Refactor pending blocks handling

Makefile

@@ -83,6 +83,7 @@ mock-gen:
 	@echo "-> Generating mocks"
 	mockery --output test/mocks --srcpkg github.com/cometbft/cometbft/rpc/client --name Client
 	mockery --output test/mocks --srcpkg github.com/cometbft/cometbft/abci/types --name Application
+	mockery --output test/mocks --srcpkg github.com/rollkit/go-da --name DA
 .PHONY: mock-gen
 
 

block/manager.go

@@ -197,6 +197,11 @@ func NewManager(
 		txsAvailableCh = nil
 	}
 
+	pendingBlocks, err := NewPendingBlocks(store, logger)
+	if err != nil {
+		return nil, err
+	}
+
 	agg := &Manager{
 		proposerKey: proposerKey,
 		conf:        conf,
@@ -219,7 +224,7 @@ func NewManager(
 		validatorSet:  &valSet,
 		txsAvailable:  txsAvailableCh,
 		buildingBlock: false,
-		pendingBlocks: NewPendingBlocks(),
+		pendingBlocks: pendingBlocks,
 		metrics:       seqMetrics,
 	}
 	return agg, nil
@@ -784,9 +789,6 @@ func (m *Manager) publishBlock(ctx context.Context) error {
 		return err
 	}
 
-	// Submit block to be published to the DA layer
-	m.pendingBlocks.addPendingBlock(block)
-
 	// Commit the new state and block which writes to disk on the proxy app
 	appHash, _, err := m.executor.Commit(ctx, newState, block, responses)
 	if err != nil {
@@ -840,7 +842,20 @@ func (m *Manager) recordMetrics(block *types.Block) {
 func (m *Manager) submitBlocksToDA(ctx context.Context) error {
 	submittedAllBlocks := false
 	backoff := initialBackoff
-	blocksToSubmit := m.pendingBlocks.getPendingBlocks()
+	blocksToSubmit, err := m.pendingBlocks.getPendingBlocks(ctx)
+	if len(blocksToSubmit) == 0 {
+		// There are no pending blocks; return because there's nothing to do, but:
+		// - it might be caused by error, then err != nil
+		// - all pending blocks are processed, then err == nil
+		// whatever the reason, error information is propagated correctly to the caller
+		return err
+	}
+	if err != nil {
+		// There are some pending blocks but also an error. It's very unlikely case - probably some error while reading
+		// blocks from the store.
+		// The error is logged and normal processing of pending blocks continues.
+		m.logger.Error("error while fetching blocks pending DA", "err", err)
+	}
 	numSubmittedBlocks := 0
 	attempt := 0
 	maxBlobSize, err := m.dalc.DA.MaxBlobSize(ctx)
@@ -863,7 +878,11 @@ func (m *Manager) submitBlocksToDA(ctx context.Context) error {
 			for _, block := range submittedBlocks {
 				m.blockCache.setDAIncluded(block.Hash().String())
 			}
-			m.pendingBlocks.removeSubmittedBlocks(submittedBlocks)
+			lastSubmittedHeight := uint64(0)
+			if l := len(submittedBlocks); l > 0 {
+				lastSubmittedHeight = submittedBlocks[l-1].Height()
+			}
+			m.pendingBlocks.setLastSubmittedHeight(ctx, lastSubmittedHeight)
 			blocksToSubmit = notSubmittedBlocks
 			// reset submission options when successful
 			// scale back gasPrice gradually

block/manager_test.go

@@ -3,9 +3,13 @@ package block
 import (
 	"bytes"
 	"context"
+	"os"
+	"strconv"
 	"testing"
 	"time"
 
+	ds "github.com/ipfs/go-datastore"
+
 	cmtypes "github.com/cometbft/cometbft/types"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -138,13 +142,20 @@ func TestSubmitBlocksToMockDA(t *testing.T) {
 	m.conf.DAMempoolTTL = 1
 	m.dalc.GasPrice = 1.0
 	m.dalc.GasMultiplier = 1.2
+	kvStore, err := store.NewDefaultInMemoryKVStore()
+	require.NoError(t, err)
+	m.store = store.New(kvStore)
 
 	t.Run("handle_tx_already_in_mempool", func(t *testing.T) {
 		var blobs [][]byte
 		block := types.GetRandomBlock(1, 5)
 		blob, err := block.MarshalBinary()
+		require.NoError(t, err)
 
+		err = m.store.SaveBlock(ctx, block, &types.Commit{})
 		require.NoError(t, err)
+		m.store.SetHeight(ctx, 1)
+
 		blobs = append(blobs, blob)
 		// Set up the mock to
 		// * throw timeout waiting for tx to be included exactly once
@@ -162,15 +173,16 @@ func TestSubmitBlocksToMockDA(t *testing.T) {
 			On("Submit", blobs, 1.0*1.2*1.2, []byte(nil)).
 			Return([][]byte{bytes.Repeat([]byte{0x00}, 8)}, nil)
 
-		m.pendingBlocks = NewPendingBlocks()
-		m.pendingBlocks.addPendingBlock(block)
+		m.pendingBlocks, err = NewPendingBlocks(m.store, m.logger)
+		require.NoError(t, err)
 		err = m.submitBlocksToDA(ctx)
 		require.NoError(t, err)
 		mockDA.AssertExpectations(t)
 	})
 }
 
 func TestSubmitBlocksToDA(t *testing.T) {
+	assert := assert.New(t)
 	require := require.New(t)
 	ctx := context.Background()
 
@@ -230,14 +242,41 @@ func TestSubmitBlocksToDA(t *testing.T) {
 	}
 
 	for _, tc := range testCases {
-		m.pendingBlocks = NewPendingBlocks()
+		// there is a limitation of value size for underlying in-memory KV store, so (temporary) on-disk store is needed
+		kvStore := getTempKVStore(t)
+		m.store = store.New(kvStore)
+		m.pendingBlocks, err = NewPendingBlocks(m.store, m.logger)
+		require.NoError(err)
 		t.Run(tc.name, func(t *testing.T) {
+			// PendingBlocks depend on store, so blocks needs to be saved and height updated
 			for _, block := range tc.blocks {
-				m.pendingBlocks.addPendingBlock(block)
+				require.NoError(m.store.SaveBlock(ctx, block, &types.Commit{}))
 			}
+			m.store.SetHeight(ctx, uint64(len(tc.blocks)))
+
 			err := m.submitBlocksToDA(ctx)
-			assert.Equal(t, tc.isErrExpected, err != nil)
-			assert.Equal(t, tc.expectedPendingBlocksLength, len(m.pendingBlocks.getPendingBlocks()))
+			assert.Equal(tc.isErrExpected, err != nil)
+			blocks, err := m.pendingBlocks.getPendingBlocks(ctx)
+			assert.NoError(err)
+			assert.Equal(tc.expectedPendingBlocksLength, len(blocks))
+
+			// ensure that metadata is updated in KV store
+			raw, err := m.store.GetMetadata(ctx, LastSubmittedHeightKey)
+			require.NoError(err)
+			lshInKV, err := strconv.ParseUint(string(raw), 10, 64)
+			require.NoError(err)
+			assert.Equal(m.store.Height(), lshInKV+uint64(tc.expectedPendingBlocksLength))
 		})
 	}
 }
+
+func getTempKVStore(t *testing.T) ds.TxnDatastore {
+	dbPath, err := os.MkdirTemp("", t.Name())
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		_ = os.RemoveAll(dbPath)
+	})
+	kvStore, err := store.NewDefaultKVStore(os.TempDir(), dbPath, t.Name())
+	require.NoError(t, err)
+	return kvStore
+}

block/pending_blocks.go

@@ -1,64 +1,112 @@
 package block
 
 import (
-	"sort"
-	"sync"
+	"context"
+	"errors"
+	"fmt"
+	"strconv"
+	"sync/atomic"
 
+	ds "github.com/ipfs/go-datastore"
+
+	"github.com/rollkit/rollkit/store"
+	"github.com/rollkit/rollkit/third_party/log"
 	"github.com/rollkit/rollkit/types"
 )
 
+// LastSubmittedHeightKey is the key used for persisting the last submitted height in store.
+const LastSubmittedHeightKey = "last submitted"
+
 // PendingBlocks maintains blocks that need to be published to DA layer
+//
+// Important assertions:
+// - blocks are safely stored in database before submission to DA
+// - blocks are always pushed to DA in order (by height)
+// - DA submission of multiple blocks is atomic - it's impossible to submit only part of a batch
+//
+// lastSubmittedHeight is updated only after receiving confirmation from DA.
+// Worst case scenario is when blocks was successfully submitted to DA, but confirmation was not received (e.g. node was
+// restarted, networking issue occurred). In this case blocks are re-submitted to DA (it's extra cost).
+// rollkit is able to skip duplicate blocks so this shouldn't affect full nodes.
+// TODO(tzdybal): we shouldn't try to push all pending blocks at once; this should depend on max blob size
 type PendingBlocks struct {
-	pendingBlocks map[uint64]*types.Block
-	mtx           *sync.RWMutex
+	store  store.Store
+	logger log.Logger
+
+	// lastSubmittedHeight holds information about last block successfully submitted to DA
+	lastSubmittedHeight atomic.Uint64
 }
 
 // NewPendingBlocks returns a new PendingBlocks struct
-func NewPendingBlocks() *PendingBlocks {
-	return &PendingBlocks{
-		pendingBlocks: make(map[uint64]*types.Block),
-		mtx:           new(sync.RWMutex),
+func NewPendingBlocks(store store.Store, logger log.Logger) (*PendingBlocks, error) {
+	pb := &PendingBlocks{
+		store:  store,
+		logger: logger,
 	}
+	if err := pb.init(); err != nil {
+		return nil, err
+	}
+	return pb, nil
 }
 
 // getPendingBlocks returns a sorted slice of pending blocks
 // that need to be published to DA layer in order of block height
-func (pb *PendingBlocks) getPendingBlocks() []*types.Block {
-	blocks := copyBlocks(pb)
-	sort.Slice(blocks, func(i, j int) bool {
-		return blocks[i].Height() < blocks[j].Height()
-	})
-	return blocks
-}
+func (pb *PendingBlocks) getPendingBlocks(ctx context.Context) ([]*types.Block, error) {
+	lastSubmitted := pb.lastSubmittedHeight.Load()
+	height := pb.store.Height()
+
+	if lastSubmitted == height {
+		return nil, nil
+	}
+	if lastSubmitted > height {
+		panic(fmt.Sprintf("height of last block submitted to DA (%d) is greater than height of last block (%d)",
+			lastSubmitted, height))
+	}
 
-// copyBlocks creates a copy of the pending blocks in a thread-safe manner.
-// It returns a slice of pointers to the copied blocks.
-func copyBlocks(pb *PendingBlocks) []*types.Block {
-	pb.mtx.RLock()
-	defer pb.mtx.RUnlock()
-	blocks := make([]*types.Block, 0, len(pb.pendingBlocks))
-	for _, block := range pb.pendingBlocks {
+	blocks := make([]*types.Block, 0, height-lastSubmitted)
+	for i := lastSubmitted + 1; i <= height; i++ {
+		block, err := pb.store.GetBlock(ctx, i)
+		if err != nil {
+			// return as much as possible + error information
+			return blocks, err
+		}
 		blocks = append(blocks, block)
 	}
-	return blocks
+	return blocks, nil
 }
 
 func (pb *PendingBlocks) isEmpty() bool {
-	pb.mtx.RLock()
-	defer pb.mtx.RUnlock()
-	return len(pb.pendingBlocks) == 0
+	return pb.store.Height() == pb.lastSubmittedHeight.Load()
 }
 
-func (pb *PendingBlocks) addPendingBlock(block *types.Block) {
-	pb.mtx.Lock()
-	defer pb.mtx.Unlock()
-	pb.pendingBlocks[block.Height()] = block
+func (pb *PendingBlocks) setLastSubmittedHeight(ctx context.Context, newLastSubmittedHeight uint64) {
+	lsh := pb.lastSubmittedHeight.Load()
+
+	if newLastSubmittedHeight > lsh && pb.lastSubmittedHeight.CompareAndSwap(lsh, newLastSubmittedHeight) {
+		err := pb.store.SetMetadata(ctx, LastSubmittedHeightKey, []byte(strconv.FormatUint(newLastSubmittedHeight, 10)))
+		if err != nil {
+			// This indicates IO error in KV store. We can't do much about this.
+			// After next successful DA submission, update will be re-attempted (with new value).
+			// If store is not updated, after node restart some blocks will be re-submitted to DA.
+			pb.logger.Error("failed to store height of latest block submitted to DA", "err", err)
+		}
+	}
 }
 
-func (pb *PendingBlocks) removeSubmittedBlocks(blocks []*types.Block) {
-	pb.mtx.Lock()
-	defer pb.mtx.Unlock()
-	for _, block := range blocks {
-		delete(pb.pendingBlocks, block.Height())
+func (pb *PendingBlocks) init() error {
+	raw, err := pb.store.GetMetadata(context.Background(), LastSubmittedHeightKey)
+	if errors.Is(err, ds.ErrNotFound) {
+		// LastSubmittedHeightKey was never used, it's special case not actual error
+		// we don't need to modify lastSubmittedHeight
+		return nil
+	}
+	if err != nil {
+		return err
+	}
+	lsh, err := strconv.ParseUint(string(raw), 10, 64)
+	if err != nil {
+		return err
 	}
+	pb.lastSubmittedHeight.CompareAndSwap(0, lsh)
+	return nil
 }