Merge pull request #323 from lightninglabs/caretaker_atomic_batch_state

caretaker: make batch state atomic

rpcserver.go

@@ -1881,7 +1881,9 @@ func marshalMintingBatch(batch *tapgarden.MintingBatch) (*mintrpc.MintingBatch,
 func marshalBatchState(batch *tapgarden.MintingBatch) (mintrpc.BatchState,
 	error) {
 
-	switch batch.BatchState {
+	currentBatchState := batch.State()
+
+	switch currentBatchState {
 	case tapgarden.BatchStatePending:
 		return mintrpc.BatchState_BATCH_STATE_PEDNING, nil
 
@@ -1907,8 +1909,8 @@ func marshalBatchState(batch *tapgarden.MintingBatch) (mintrpc.BatchState,
 		return mintrpc.BatchState_BATCH_STATE_SPROUT_CANCELLED, nil
 
 	default:
-		return 0, fmt.Errorf("unknown batch state: %d",
-			batch.BatchState)
+		return 0, fmt.Errorf("unknown batch state: %v",
+			currentBatchState.String())
 	}
 }
 

tapdb/asset_minting.go

@@ -832,9 +832,6 @@ func marshalMintingBatch(ctx context.Context, q PendingAssetStore,
 	// For each batch, we'll assemble an intermediate batch struct, then
 	// fill in all the seedlings with another sub-query.
 	batch := &tapgarden.MintingBatch{
-		BatchState: tapgarden.BatchState(
-			dbBatch.BatchState,
-		),
 		BatchKey: keychain.KeyDescriptor{
 			KeyLocator: keychain.KeyLocator{
 				Family: keychain.KeyFamily(
@@ -848,6 +845,13 @@ func marshalMintingBatch(ctx context.Context, q PendingAssetStore,
 		CreationTime: dbBatch.CreationTimeUnix.UTC(),
 	}
 
+	batchState, err := tapgarden.NewBatchState(uint8(dbBatch.BatchState))
+	if err != nil {
+		return nil, err
+	}
+
+	batch.UpdateState(batchState)
+
 	if dbBatch.MintingTxPsbt != nil {
 		genesisPkt, err := psbt.NewFromRawBytes(
 			bytes.NewReader(dbBatch.MintingTxPsbt), false,
@@ -867,7 +871,7 @@ func marshalMintingBatch(ctx context.Context, q PendingAssetStore,
 	// either fetch the set of seedlings (asset
 	// descriptions w/ no real assets), or the set of
 	// sprouts (full defined assets, but not yet mined).
-	switch batch.BatchState {
+	switch batchState {
 	case tapgarden.BatchStatePending,
 		tapgarden.BatchStateFrozen,
 		tapgarden.BatchStateSeedlingCancelled:

tapdb/asset_minting_test.go

@@ -53,14 +53,14 @@ func newAssetStore(t *testing.T) (*AssetMintingStore, *AssetStore,
 func assertBatchState(t *testing.T, batch *tapgarden.MintingBatch,
 	state tapgarden.BatchState) {
 
-	require.Equal(t, state, batch.BatchState)
+	require.Equal(t, state, batch.State())
 }
 
 func assertBatchEqual(t *testing.T, a, b *tapgarden.MintingBatch) {
 	t.Helper()
 
 	require.Equal(t, a.CreationTime.Unix(), b.CreationTime.Unix())
-	require.Equal(t, a.BatchState, b.BatchState)
+	require.Equal(t, a.State(), b.State())
 	require.Equal(t, a.BatchKey, b.BatchKey)
 	require.Equal(t, a.Seedlings, b.Seedlings)
 	require.Equal(t, a.GenesisPacket, b.GenesisPacket)

tapgarden/batch.go

@@ -2,6 +2,7 @@ package tapgarden
 
 import (
 	"fmt"
+	"sync/atomic"
 	"time"
 
 	"github.com/btcsuite/btcd/btcec/v2"
@@ -30,8 +31,8 @@ type MintingBatch struct {
 	// transaction.
 	HeightHint uint32
 
-	// BatchState is the state of the batch.
-	BatchState BatchState
+	// batchState is the state of the batch.
+	batchState atomic.Uint32
 
 	// BatchKey is the unique identifier for a batch.
 	BatchKey keychain.KeyDescriptor
@@ -132,3 +133,19 @@ func (m *MintingBatch) genesisScript() ([]byte, error) {
 
 	return tapscript.PayToTaprootScript(mintingOutputKey)
 }
+
+// State returns the private state of the batch.
+func (m *MintingBatch) State() BatchState {
+	currentBatchState := m.batchState.Load()
+
+	// Drop the error when converting the stored state to a BatchState, as
+	// we verify the batch state before storing it.
+	batchStateCopy, _ := NewBatchState(uint8(currentBatchState))
+	return batchStateCopy
+}
+
+// UpdateState updates the state of a batch to a value that has been verified to
+// be a valid batch state.
+func (m *MintingBatch) UpdateState(state BatchState) {
+	m.batchState.Store(uint32(state))
+}

tapgarden/caretaker.go

@@ -146,7 +146,7 @@ func (b *BatchCaretaker) Cancel() CancelResp {
 	defer cancel()
 
 	batchKey := b.cfg.Batch.BatchKey.PubKey.SerializeCompressed()
-	batchState := b.cfg.Batch.BatchState
+	batchState := b.cfg.Batch.State()
 	// This function can only be called before the caretaker state stepping
 	// function, so the batch state read is the next state that has not yet
 	// been executed. Seedlings are converted to asset sprouts in the Frozen
@@ -235,7 +235,7 @@ func (b *BatchCaretaker) advanceStateUntil(currentState,
 
 		currentState = nextState
 
-		b.cfg.Batch.BatchState = currentState
+		b.cfg.Batch.UpdateState(currentState)
 	}
 
 	return currentState, nil
@@ -248,9 +248,10 @@ func (b *BatchCaretaker) advanceStateUntil(currentState,
 func (b *BatchCaretaker) assetCultivator() {
 	defer b.Wg.Done()
 
+	currentBatchState := b.cfg.Batch.State()
 	// If the batch is already marked as confirmed, then we just need to
 	// advance it one more level to be finalized.
-	if b.cfg.Batch.BatchState == BatchStateConfirmed {
+	if currentBatchState == BatchStateConfirmed {
 		log.Infof("MintingBatch(%x): already confirmed!", b.batchKey[:])
 
 		_, err := b.advanceStateUntil(
@@ -271,7 +272,7 @@ func (b *BatchCaretaker) assetCultivator() {
 	// confirmation notification, which'll let us advance to the final
 	// state.
 	_, err := b.advanceStateUntil(
-		b.cfg.Batch.BatchState, BatchStateBroadcast,
+		currentBatchState, BatchStateBroadcast,
 	)
 	if err != nil {
 		log.Errorf("unable to advance state machine: %v", err)
@@ -295,11 +296,12 @@ func (b *BatchCaretaker) assetCultivator() {
 				confInfo.BlockHash, confInfo.BlockHeight)
 
 			b.confInfo = confInfo
-			b.cfg.Batch.BatchState = BatchStateConfirmed
+			b.cfg.Batch.UpdateState(BatchStateConfirmed)
+			currentBatchState = b.cfg.Batch.State()
 
 			// TODO(roasbeef): use a "trigger" here instead?
 			_, err = b.advanceStateUntil(
-				b.cfg.Batch.BatchState, BatchStateFinalized,
+				currentBatchState, BatchStateFinalized,
 			)
 			if err != nil {
 				log.Error(err)

tapgarden/interface.go

@@ -126,7 +126,41 @@ func (b BatchState) String() string {
 		return "BatchStateSproutCancelled"
 
 	default:
-		return fmt.Sprintf("UnknownState(%v)", int(b))
+		return fmt.Sprintf("UnknownState(%d)", b)
+	}
+}
+
+// NewBatchState creates a BatchState from a uint8, returning an error if the
+// input value does not map to a valid BatchState.
+func NewBatchState(state uint8) (BatchState, error) {
+	switch BatchState(state) {
+	case BatchStatePending:
+		return BatchStatePending, nil
+
+	case BatchStateFrozen:
+		return BatchStateFrozen, nil
+
+	case BatchStateCommitted:
+		return BatchStateCommitted, nil
+
+	case BatchStateBroadcast:
+		return BatchStateBroadcast, nil
+
+	case BatchStateConfirmed:
+		return BatchStateConfirmed, nil
+
+	case BatchStateFinalized:
+		return BatchStateFinalized, nil
+
+	case BatchStateSeedlingCancelled:
+		return BatchStateSeedlingCancelled, nil
+
+	case BatchStateSproutCancelled:
+		return BatchStateSproutCancelled, nil
+
+	default:
+		return BatchStateSproutCancelled,
+			fmt.Errorf("unknown batch state: %v", state)
 	}
 }
 

tapgarden/planter.go

@@ -254,8 +254,10 @@ func (c *ChainPlanter) Start() error {
 		// caretaker which'll handle progressing each batch to
 		// completion. We'll skip batches that were cancelled.
 		for _, batch := range nonFinalBatches {
-			if batch.BatchState == BatchStateSeedlingCancelled ||
-				batch.BatchState == BatchStateSproutCancelled {
+			batchState := batch.State()
+
+			if batchState == BatchStateSeedlingCancelled ||
+				batchState == BatchStateSproutCancelled {
 
 				continue
 			}
@@ -748,13 +750,13 @@ func (c *ChainPlanter) prepAssetSeedling(ctx context.Context,
 		newBatch := &MintingBatch{
 			CreationTime: time.Now(),
 			HeightHint:   currentHeight,
-			BatchState:   BatchStatePending,
 			BatchKey:     newInternalKey,
 			Seedlings: map[string]*Seedling{
 				req.AssetName: req,
 			},
 			AssetMetas: make(AssetMetas),
 		}
+		newBatch.UpdateState(BatchStatePending)
 		ctx, cancel = c.WithCtxQuit()
 		defer cancel()
 		err = c.cfg.Log.CommitMintingBatch(ctx, newBatch)

tapgarden/planter_test.go

@@ -363,8 +363,9 @@ func (t *mintingTestHarness) assertSeedlingsExist(
 }
 
 func isCancelledBatch(batch *tapgarden.MintingBatch) bool {
-	return batch.BatchState == tapgarden.BatchStateSeedlingCancelled ||
-		batch.BatchState == tapgarden.BatchStateSproutCancelled
+	batchState := batch.State()
+	return batchState == tapgarden.BatchStateSeedlingCancelled ||
+		batchState == tapgarden.BatchStateSproutCancelled
 }
 
 func (t *mintingTestHarness) assertBatchState(batchKey *btcec.PublicKey,
@@ -377,7 +378,7 @@ func (t *mintingTestHarness) assertBatchState(batchKey *btcec.PublicKey,
 	require.Len(t, batches, 1)
 
 	batch := batches[0]
-	require.Equal(t, batchState, batch.BatchState)
+	require.Equal(t, batchState, batch.State())
 }
 
 // assertSeedlingsMatchSprouts asserts that the seedlings were properly matched
@@ -398,7 +399,7 @@ func (t *mintingTestHarness) assertSeedlingsMatchSprouts(
 
 		// Filter out any cancelled batches.
 		isCommittedBatch := func(batch *tapgarden.MintingBatch) bool {
-			return batch.BatchState == tapgarden.BatchStateCommitted
+			return batch.State() == tapgarden.BatchStateCommitted
 		}
 		batch, err := chanutils.First(pendingBatches, isCommittedBatch)
 		if err != nil {