Merge pull request #19 from wpaulino/height-hint-sweep

account+order: update account height hints

account/interface.go

@@ -269,6 +269,14 @@ func OutPointModifier(op wire.OutPoint) Modifier {
 	}
 }
 
+// HeightHintModifier is a functional option that modifies the height hint of an
+// account.
+func HeightHintModifier(heightHint uint32) Modifier {
+	return func(account *Account) {
+		account.HeightHint = heightHint
+	}
+}
+
 // CloseTxModifier is a functional option that modifies the closing transaction
 // of an account.
 func CloseTxModifier(tx *wire.MsgTx) Modifier {

account/manager.go

@@ -598,8 +598,11 @@ func (m *Manager) handleAccountConf(traderKey *btcec.PublicKey,
 		traderKey.SerializeCompressed(), confDetails.BlockHeight)
 
 	// Mark the account as open and proceed with the rest of the flow.
-	err = m.cfg.Store.UpdateAccount(account, StateModifier(StateOpen))
-	if err != nil {
+	mods := []Modifier{
+		StateModifier(StateOpen),
+		HeightHintModifier(confDetails.BlockHeight),
+	}
+	if err := m.cfg.Store.UpdateAccount(account, mods...); err != nil {
 		return err
 	}
 
@@ -700,7 +703,9 @@ func (m *Manager) handleAccountSpend(traderKey *btcec.PublicKey,
 	// Write the spending transaction once again in case the one we
 	// previously broadcast was replaced with a higher fee one.
 	return m.cfg.Store.UpdateAccount(
-		account, StateModifier(StateClosed), CloseTxModifier(spendTx),
+		account, StateModifier(StateClosed),
+		HeightHintModifier(uint32(spendDetails.SpendingHeight)),
+		CloseTxModifier(spendTx),
 	)
 }
 
@@ -934,6 +939,9 @@ func (m *Manager) spendAccount(ctx context.Context, account *Account,
 		return nil, nil, err
 	}
 
+	// Update the account's height hint.
+	modifiers = append(modifiers, HeightHintModifier(bestHeight))
+
 	// With the transaction crafted, update our on-disk state and broadcast
 	// the transaction. We'll need some additional modifiers based on
 	// whether the account is being closed or not.

account/manager_test.go

@@ -152,11 +152,15 @@ func (h *testHarness) openAccount(value btcutil.Amount, expiry uint32, // nolint
 	h.assertAccountNotSubscribed(account.TraderKey.PubKey)
 
 	// Notify the confirmation of the account.
-	h.notifier.confChan <- &chainntnfs.TxConfirmation{}
+	confHeight := bestHeight + 6
+	h.notifier.confChan <- &chainntnfs.TxConfirmation{
+		BlockHeight: confHeight,
+	}
 
 	// This should prompt the account to now be in a StateOpen state and
 	// the subscription for updates should now be realized.
 	account.State = StateOpen
+	account.HeightHint = confHeight
 	h.assertAccountExists(account)
 	h.assertAccountSubscribed(account.TraderKey.PubKey)
 
@@ -195,14 +199,20 @@ func (h *testHarness) closeAccount(account *Account, outputs []*wire.TxOut,
 	closeTx := h.assertSpendTxBroadcast(account, nil, nil, nil)
 
 	account.State = StatePendingClosed
+	account.HeightHint = bestHeight
 	account.CloseTx = closeTx
 	h.assertAccountExists(account)
 
 	// Notify the transaction as a spend of the account.
-	h.notifier.spendChan <- &chainntnfs.SpendDetail{SpendingTx: closeTx}
+	spendHeight := bestHeight + 6
+	h.notifier.spendChan <- &chainntnfs.SpendDetail{
+		SpendingTx:     closeTx,
+		SpendingHeight: int32(spendHeight),
+	}
 
 	// This should prompt the account to now be in a StateClosed state.
 	account.State = StateClosed
+	account.HeightHint = spendHeight
 	h.assertAccountExists(account)
 
 	return closeTx
@@ -371,7 +381,8 @@ func (h *testHarness) assertAuctioneerReceived(inputs []*lnwallet.Utxo,
 // StateOpen.
 func (h *testHarness) assertAccountModification(account *Account,
 	inputs []*lnwallet.Utxo, outputs []*wire.TxOut,
-	newAccountValue btcutil.Amount, accountInputIdx, accountOutputIdx uint32) {
+	newAccountValue btcutil.Amount, accountInputIdx,
+	accountOutputIdx uint32, broadcastHeight uint32) {
 
 	h.t.Helper()
 
@@ -395,6 +406,7 @@ func (h *testHarness) assertAccountModification(account *Account,
 			Hash:  spendTx.TxHash(),
 			Index: accountOutputIdx,
 		}),
+		HeightHintModifier(broadcastHeight),
 		IncrementBatchKey(),
 	}
 	for _, mod := range mods {
@@ -413,8 +425,13 @@ func (h *testHarness) assertAccountModification(account *Account,
 
 	// Notify the confirmation, causing the account to transition back to
 	// StateOpen.
-	h.notifier.confChan <- &chainntnfs.TxConfirmation{Tx: spendTx}
+	confHeight := broadcastHeight + 6
+	h.notifier.confChan <- &chainntnfs.TxConfirmation{
+		Tx:          spendTx,
+		BlockHeight: confHeight,
+	}
 	StateModifier(StateOpen)(account)
+	HeightHintModifier(confHeight)(account)
 	h.assertAccountExists(account)
 }
 
@@ -564,10 +581,14 @@ func TestResumeAccountAfterRestart(t *testing.T) {
 	h.assertAccountExists(account)
 
 	// Notify the confirmation of the account.
-	h.notifier.confChan <- &chainntnfs.TxConfirmation{}
+	confHeight := uint32(bestHeight + 6)
+	h.notifier.confChan <- &chainntnfs.TxConfirmation{
+		BlockHeight: confHeight,
+	}
 
 	// This should prompt the account to now be in a Confirmed state.
 	account.State = StateOpen
+	account.HeightHint = confHeight
 	h.assertAccountExists(account)
 }
 
@@ -709,7 +730,7 @@ func TestAccountWithdrawal(t *testing.T) {
 	withdrawOutputSum := valuePerOutput * btcutil.Amount(len(outputs))
 	valueAfterWithdrawal := account.Value - withdrawOutputSum - expectedFee
 	h.assertAccountModification(
-		account, nil, outputs, valueAfterWithdrawal, 0, 0,
+		account, nil, outputs, valueAfterWithdrawal, 0, 0, bestHeight,
 	)
 
 	// Finally, close the account to ensure we can process another spend
@@ -804,6 +825,7 @@ func TestAccountDeposit(t *testing.T) {
 	h.assertAccountModification(
 		account, h.wallet.utxos, []*wire.TxOut{changeOutput},
 		valueAfterDeposit, accountInputIdx, accountOutputIdx,
+		bestHeight,
 	)
 
 	// Finally, close the account to ensure we can process another spend

order/batch.go

@@ -221,7 +221,7 @@ type BatchSigner interface {
 type BatchStorer interface {
 	// StorePendingBatch makes sure all changes executed by a pending batch
 	// are correctly and atomically stored to the database.
-	StorePendingBatch(*Batch) error
+	StorePendingBatch(_ *Batch, bestHeight uint32) error
 
 	// MarkBatchComplete marks a pending batch as complete, allowing a
 	// trader to participate in a new batch.

order/batch_storer.go

@@ -9,6 +9,13 @@ import (
 	"github.com/lightninglabs/llm/clmrpc"
 )
 
+const (
+	// heightHintPadding is the padding we add to our best known height to
+	// avoid any discrepancies in block propagation between us and the
+	// auctioneer.
+	heightHintPadding = -3
+)
+
 // batchStorer is a type that implements BatchStorer and can persist a batch to
 // the local trader database.
 type batchStorer struct {
@@ -23,7 +30,7 @@ type batchStorer struct {
 // modifications will be applied atomically as a result of MarkBatchComplete.
 //
 // NOTE: This method is part of the BatchStorer interface.
-func (s *batchStorer) StorePendingBatch(batch *Batch) error {
+func (s *batchStorer) StorePendingBatch(batch *Batch, bestHeight uint32) error {
 	// Prepare the order modifications first.
 	orders := make([]Nonce, len(batch.MatchedOrders))
 	orderModifiers := make([][]Modifier, len(orders))
@@ -64,6 +71,13 @@ func (s *batchStorer) StorePendingBatch(batch *Batch) error {
 	// Next create our account modifiers.
 	accounts := make([]*account.Account, len(batch.AccountDiffs))
 	accountModifiers := make([][]account.Modifier, len(accounts))
+
+	// Each account will have the same height hint applied.
+	heightHint := int64(bestHeight) + heightHintPadding
+	if heightHint < 0 {
+		heightHint = 0
+	}
+
 	for idx, diff := range batch.AccountDiffs {
 		// Get the current state of the account first so we can create
 		// a proper diff.
@@ -107,10 +121,14 @@ func (s *batchStorer) StorePendingBatch(batch *Batch) error {
 				diff.EndingState)
 		}
 
-		// Finally update the account value and expiry.
-		accountModifiers[idx] = append(
+		// Finally update the account value and height hint.
+		modifiers = append(
 			modifiers, account.ValueModifier(diff.EndingBalance),
 		)
+		modifiers = append(
+			modifiers, account.HeightHintModifier(uint32(heightHint)),
+		)
+		accountModifiers[idx] = modifiers
 	}
 
 	// Everything is ready to be persisted now.

order/batch_storer_test.go

@@ -15,6 +15,7 @@ import (
 func TestBatchStorer(t *testing.T) {
 	t.Parallel()
 
+	const bestHeight = 1337
 	var (
 		storeMock = newMockStore()
 		storer    = &batchStorer{
@@ -119,7 +120,7 @@ func TestBatchStorer(t *testing.T) {
 	}
 
 	// Pass the assembled batch to the storer now.
-	err := storer.StorePendingBatch(batch)
+	err := storer.StorePendingBatch(batch, bestHeight)
 	if err != nil {
 		t.Fatalf("error storing batch: %v", err)
 	}
@@ -165,6 +166,11 @@ func TestBatchStorer(t *testing.T) {
 		t.Fatalf("invalid account expiry, got %d wanted %d",
 			smallAcct.Value, 144)
 	}
+	heightHint := uint32(bestHeight + heightHintPadding)
+	if smallAcct.HeightHint != heightHint {
+		t.Fatalf("invalid account height hint, got %d wanted %d",
+			smallAcct.Value, heightHint)
+	}
 
 	if bigAcct.State != account.StatePendingUpdate {
 		t.Fatalf("invalid account state, got %d wanted %d",
@@ -178,6 +184,10 @@ func TestBatchStorer(t *testing.T) {
 		t.Fatalf("invalid account expiry, got %d wanted %d",
 			bigAcct.Value, 144)
 	}
+	if bigAcct.HeightHint != heightHint {
+		t.Fatalf("invalid account height hint, got %d wanted %d",
+			bigAcct.Value, heightHint)
+	}
 }
 
 func newKit(nonce Nonce, units SupplyUnit) Kit {

order/manager.go

@@ -239,13 +239,14 @@ func (m *Manager) PendingBatch() *Batch {
 // belong to the trader. Before sending off the signature to the auctioneer,
 // we'll also persist the batch to disk as pending to ensure we can recover
 // after a crash.
-func (m *Manager) BatchSign() (BatchSignature, error) {
+func (m *Manager) BatchSign(bestHeight uint32) (BatchSignature, error) {
 	sig, err := m.batchSigner.Sign(m.pendingBatch)
 	if err != nil {
 		return nil, err
 	}
 
-	if err := m.batchStorer.StorePendingBatch(m.pendingBatch); err != nil {
+	err = m.batchStorer.StorePendingBatch(m.pendingBatch, bestHeight)
+	if err != nil {
 		return nil, fmt.Errorf("unable to store batch: %v", err)
 	}
 

rpcserver.go

@@ -663,7 +663,8 @@ func (s *rpcServer) handleServerMessage(rpcMsg *clmrpc.ServerAuctionMessage) err
 			"num_orders=%v", batch.ID[:], len(batch.MatchedOrders))
 
 		// Sign for the accounts in the batch.
-		sigs, err := s.orderManager.BatchSign()
+		bestHeight := atomic.LoadUint32(&s.bestHeight)
+		sigs, err := s.orderManager.BatchSign(bestHeight)
 		if err != nil {
 			rpcLog.Errorf("Error signing batch: %v", err)
 			return s.sendRejectBatch(batch, err)