add ctx to blocksItr

Add context.Context to the block iterator so that waitForBlock respects
context cancellation. This fixes goroutine leaks in callers like
deliver.go where cursor.Next() blocks indefinitely even after the gRPC
stream's context is cancelled.

Signed-off-by: Senthilnathan <cendhu@gmail.com>

Author: cendhu

common/deliver/deliver.go

@@ -261,7 +261,7 @@ func (h *Handler) deliverBlocks(ctx context.Context, srv *Server, envelope *cb.E
 
 	logger.Debugf("[channel: %s] Received seekInfo (%p) %v from %s", chdr.ChannelId, seekInfo, seekInfo, addr)
 
-	cursor, number := chain.Reader().Iterator(seekInfo.Start)
+	cursor, number := chain.Reader().Iterator(ctx, seekInfo.Start)
 	defer cursor.Close()
 	var stopNum uint64
 	switch stop := seekInfo.Stop.Type.(type) {

common/deliver/deliver_test.go

@@ -341,7 +341,7 @@ var _ = ginkgo.Describe("Deliver", func() {
 			gomega.Expect(err).NotTo(gomega.HaveOccurred())
 
 			gomega.Expect(fakeBlockReader.IteratorCallCount()).To(gomega.Equal(1))
-			startPosition := fakeBlockReader.IteratorArgsForCall(0)
+			_, startPosition := fakeBlockReader.IteratorArgsForCall(0)
 			gomega.Expect(startPosition).To(test.ProtoEqual(seekInfo.Start))
 		})
 
@@ -423,7 +423,7 @@ var _ = ginkgo.Describe("Deliver", func() {
 				gomega.Expect(err).NotTo(gomega.HaveOccurred())
 
 				gomega.Expect(fakeBlockReader.IteratorCallCount()).To(gomega.Equal(1))
-				start := fakeBlockReader.IteratorArgsForCall(0)
+				_, start := fakeBlockReader.IteratorArgsForCall(0)
 				gomega.Expect(start).To(test.ProtoEqual(&ab.SeekPosition{}))
 				gomega.Expect(fakeBlockIterator.NextCallCount()).To(gomega.Equal(1))
 
@@ -453,7 +453,7 @@ var _ = ginkgo.Describe("Deliver", func() {
 				gomega.Expect(err).NotTo(gomega.HaveOccurred())
 
 				gomega.Expect(fakeBlockReader.IteratorCallCount()).To(gomega.Equal(1))
-				start := fakeBlockReader.IteratorArgsForCall(0)
+				_, start := fakeBlockReader.IteratorArgsForCall(0)
 				gomega.Expect(start).To(test.ProtoEqual(&ab.SeekPosition{}))
 
 				gomega.Expect(fakeBlockIterator.NextCallCount()).To(gomega.Equal(2))
@@ -525,7 +525,7 @@ var _ = ginkgo.Describe("Deliver", func() {
 				gomega.Expect(err).NotTo(gomega.HaveOccurred())
 
 				gomega.Expect(fakeBlockReader.IteratorCallCount()).To(gomega.Equal(1))
-				start := fakeBlockReader.IteratorArgsForCall(0)
+				_, start := fakeBlockReader.IteratorArgsForCall(0)
 				gomega.Expect(start).To(test.ProtoEqual(&ab.SeekPosition{}))
 
 				gomega.Expect(fakeBlockIterator.NextCallCount()).To(gomega.Equal(2))
@@ -592,7 +592,7 @@ var _ = ginkgo.Describe("Deliver", func() {
 					err := handler.Handle(context.Background(), server)
 					gomega.Expect(err).NotTo(gomega.HaveOccurred())
 					gomega.Expect(fakeBlockReader.IteratorCallCount()).To(gomega.Equal(1))
-					start := fakeBlockReader.IteratorArgsForCall(0)
+					_, start := fakeBlockReader.IteratorArgsForCall(0)
 					gomega.Expect(start).To(test.ProtoEqual(&ab.SeekPosition{}))
 					gomega.Expect(fakeBlockIterator.NextCallCount()).To(gomega.Equal(3))
 					gomega.Expect(fakeResponseSender.SendBlockResponseCallCount()).To(gomega.Equal(3))

common/deliver/mock/block_reader.go

@@ -8,6 +8,7 @@ package blkstorage
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"math"
 	"sync"
@@ -609,14 +610,14 @@ func (mgr *blockfileMgr) retrieveBlockHeaderByNumber(blockNum uint64) (*common.B
 	return info.blockHeader, nil
 }
 
-func (mgr *blockfileMgr) retrieveBlocks(startNum uint64) (*blocksItr, error) {
+func (mgr *blockfileMgr) retrieveBlocks(ctx context.Context, startNum uint64) (*blocksItr, error) {
 	if startNum < mgr.firstPossibleBlockNumberInBlockFiles() {
 		return nil, errors.Errorf(
 			"cannot serve block [%d]. The ledger is bootstrapped from a snapshot. First available block = [%d]",
 			startNum, mgr.firstPossibleBlockNumberInBlockFiles(),
 		)
 	}
-	return newBlockItr(mgr, startNum), nil
+	return newBlockItr(ctx, mgr, startNum), nil
 }
 
 func (mgr *blockfileMgr) txIDExists(txID string) (bool, error) {

common/ledger/blkstorage/blockfile_mgr.go

@@ -141,7 +141,7 @@ func TestBlockfileMgrBlockIterator(t *testing.T) {
 func testBlockfileMgrBlockIterator(t *testing.T, blockfileMgr *blockfileMgr,
 	firstBlockNum, lastBlockNum int, expectedBlocks []*common.Block,
 ) {
-	itr, err := blockfileMgr.retrieveBlocks(uint64(firstBlockNum))
+	itr, err := blockfileMgr.retrieveBlocks(t.Context(), uint64(firstBlockNum))
 	require.NoError(t, err, "Error while getting blocks iterator")
 	defer itr.Close()
 	numBlocksItrated := 0

common/ledger/blkstorage/blockfile_mgr_test.go

@@ -7,6 +7,7 @@ SPDX-License-Identifier: Apache-2.0
 package blkstorage
 
 import (
+	"context"
 	"sync"
 
 	"github.com/hyperledger/fabric-x-common/common/ledger"
@@ -15,23 +16,36 @@ import (
 // blocksItr - an iterator for iterating over a sequence of blocks
 type blocksItr struct {
 	mgr                  *blockfileMgr
+	ctx                  context.Context
 	maxBlockNumAvailable uint64
 	blockNumToRetrieve   uint64
 	stream               *blockStream
 	closeMarker          bool
 	closeMarkerLock      *sync.Mutex
 }
 
-func newBlockItr(mgr *blockfileMgr, startBlockNum uint64) *blocksItr {
+func newBlockItr(ctx context.Context, mgr *blockfileMgr, startBlockNum uint64) *blocksItr {
 	mgr.blkfilesInfoCond.L.Lock()
 	defer mgr.blkfilesInfoCond.L.Unlock()
-	return &blocksItr{mgr, mgr.blockfilesInfo.lastPersistedBlock, startBlockNum, nil, false, &sync.Mutex{}}
+	return &blocksItr{
+		mgr:                  mgr,
+		ctx:                  ctx,
+		maxBlockNumAvailable: mgr.blockfilesInfo.lastPersistedBlock,
+		blockNumToRetrieve:   startBlockNum,
+		closeMarkerLock:      &sync.Mutex{},
+	}
 }
 
 func (itr *blocksItr) waitForBlock(blockNum uint64) uint64 {
 	itr.mgr.blkfilesInfoCond.L.Lock()
 	defer itr.mgr.blkfilesInfoCond.L.Unlock()
-	for itr.mgr.blockfilesInfo.lastPersistedBlock < blockNum && !itr.shouldClose() {
+
+	stop := context.AfterFunc(itr.ctx, func() {
+		itr.mgr.blkfilesInfoCond.Broadcast()
+	})
+	defer stop()
+
+	for itr.mgr.blockfilesInfo.lastPersistedBlock < blockNum && !itr.shouldClose() && itr.ctx.Err() == nil {
 		logger.Debugf("Going to wait for newer blocks. maxAvailaBlockNumber=[%d], waitForBlockNum=[%d]",
 			itr.mgr.blockfilesInfo.lastPersistedBlock, blockNum)
 		itr.mgr.blkfilesInfoCond.Wait()
@@ -68,6 +82,9 @@ func (itr *blocksItr) Next() (ledger.QueryResult, error) {
 	if itr.closeMarker {
 		return nil, nil
 	}
+	if err := itr.ctx.Err(); err != nil {
+		return nil, err
+	}
 
 	cachedBlock, existsInCache := itr.mgr.cache.get(itr.blockNumToRetrieve)
 	if existsInCache {

common/ledger/blkstorage/blocks_itr.go

@@ -27,7 +27,7 @@ func TestBlocksItrBlockingNext(t *testing.T) {
 	blocks := testutil.ConstructTestBlocks(t, 10)
 	blkfileMgrWrapper.addBlocks(blocks[:5])
 
-	itr, err := blkfileMgr.retrieveBlocks(1)
+	itr, err := blkfileMgr.retrieveBlocks(t.Context(), 1)
 	require.NoError(t, err)
 	defer itr.Close()
 	readyChan := make(chan struct{})
@@ -51,7 +51,7 @@ func TestBlockItrClose(t *testing.T) {
 	blocks := testutil.ConstructTestBlocks(t, 5)
 	blkfileMgrWrapper.addBlocks(blocks)
 
-	itr, err := blkfileMgr.retrieveBlocks(1)
+	itr, err := blkfileMgr.retrieveBlocks(t.Context(), 1)
 	require.NoError(t, err)
 
 	bh, _ := itr.Next()
@@ -74,7 +74,7 @@ func TestRaceToDeadlock(t *testing.T) {
 	blkfileMgrWrapper.addBlocks(blocks)
 
 	for i := 0; i < 1000; i++ {
-		itr, err := blkfileMgr.retrieveBlocks(5)
+		itr, err := blkfileMgr.retrieveBlocks(t.Context(), 5)
 		if err != nil {
 			panic(err)
 		}
@@ -85,7 +85,7 @@ func TestRaceToDeadlock(t *testing.T) {
 	}
 
 	for i := 0; i < 1000; i++ {
-		itr, err := blkfileMgr.retrieveBlocks(5)
+		itr, err := blkfileMgr.retrieveBlocks(t.Context(), 5)
 		if err != nil {
 			panic(err)
 		}
@@ -105,7 +105,7 @@ func TestBlockItrCloseWithoutRetrieve(t *testing.T) {
 	blocks := testutil.ConstructTestBlocks(t, 5)
 	blkfileMgrWrapper.addBlocks(blocks)
 
-	itr, err := blkfileMgr.retrieveBlocks(2)
+	itr, err := blkfileMgr.retrieveBlocks(t.Context(), 2)
 	require.NoError(t, err)
 	itr.Close()
 }
@@ -121,12 +121,12 @@ func TestCloseMultipleItrsWaitForFutureBlock(t *testing.T) {
 
 	wg := &sync.WaitGroup{}
 	wg.Add(2)
-	itr1, err := blkfileMgr.retrieveBlocks(7)
+	itr1, err := blkfileMgr.retrieveBlocks(t.Context(), 7)
 	require.NoError(t, err)
 	// itr1 does not retrieve any block because it closes before new blocks are added
 	go iterateInBackground(t, itr1, 9, wg, []uint64{})
 
-	itr2, err := blkfileMgr.retrieveBlocks(8)
+	itr2, err := blkfileMgr.retrieveBlocks(t.Context(), 8)
 	require.NoError(t, err)
 	// itr2 retrieves two blocks 8 and 9. Because it started waiting for 8 and quits at 9
 	go iterateInBackground(t, itr2, 9, wg, []uint64{8, 9})

common/ledger/blkstorage/blocks_itr_test.go

@@ -7,6 +7,7 @@ SPDX-License-Identifier: Apache-2.0
 package blkstorage
 
 import (
+	"context"
 	"time"
 
 	"github.com/hyperledger/fabric-protos-go-apiv2/common"
@@ -77,8 +78,8 @@ func (store *BlockStore) GetBlockchainInfo() (*common.BlockchainInfo, error) {
 }
 
 // RetrieveBlocks returns an iterator that can be used for iterating over a range of blocks
-func (store *BlockStore) RetrieveBlocks(startNum uint64) (ledger.ResultsIterator, error) {
-	return store.fileMgr.retrieveBlocks(startNum)
+func (store *BlockStore) RetrieveBlocks(ctx context.Context, startNum uint64) (ledger.ResultsIterator, error) {
+	return store.fileMgr.retrieveBlocks(ctx, startNum)
 }
 
 // RetrieveBlockByHash returns the block for given block-hash

common/ledger/blkstorage/blockstore.go

@@ -7,13 +7,16 @@ SPDX-License-Identifier: Apache-2.0
 package blkstorage
 
 import (
+	"context"
 	"fmt"
 	"testing"
+	"time"
 
 	"github.com/hyperledger/fabric-protos-go-apiv2/common"
 	"github.com/hyperledger/fabric-protos-go-apiv2/peer"
 	"github.com/stretchr/testify/require"
 
+	"github.com/hyperledger/fabric-x-common/common/ledger"
 	"github.com/hyperledger/fabric-x-common/common/ledger/testutil"
 	"github.com/hyperledger/fabric-x-common/protoutil"
 	"github.com/hyperledger/fabric-x-common/tools/pkg/txflags"
@@ -99,7 +102,7 @@ func checkBlocks(t *testing.T, expectedBlocks []*common.Block, store *BlockStore
 	require.Equal(t, uint64(len(expectedBlocks)), bcInfo.Height)
 	require.Equal(t, protoutil.BlockHeaderHash(expectedBlocks[len(expectedBlocks)-1].GetHeader()), bcInfo.CurrentBlockHash)
 
-	itr, _ := store.RetrieveBlocks(0)
+	itr, _ := store.RetrieveBlocks(t.Context(), 0)
 	for i := 0; i < len(expectedBlocks); i++ {
 		block, _ := itr.Next()
 		require.Equal(t, expectedBlocks[i], block)
@@ -246,6 +249,46 @@ func TestDrop(t *testing.T) {
 	require.EqualError(t, provider.Drop("ledger2"), "internal leveldb error while obtaining db iterator: leveldb: closed")
 }
 
+func TestBlocksItrContextCancel(t *testing.T) {
+	env := newTestEnv(t, NewConf(t.TempDir(), 0))
+	t.Cleanup(env.Cleanup)
+	defer env.Cleanup()
+	store, err := env.provider.Open("testledger-ctx")
+	require.NoError(t, err)
+	t.Cleanup(store.Shutdown)
+
+	// Add one block so the store is initialized
+	blocks := testutil.ConstructTestBlocks(t, 1)
+	require.NoError(t, store.AddBlock(blocks[0]))
+
+	// Create a cancellable context and an iterator starting at block 1 (not yet available)
+	ctx, cancel := context.WithCancel(t.Context())
+	itr, err := store.RetrieveBlocks(ctx, 1)
+	require.NoError(t, err)
+	t.Cleanup(itr.Close)
+
+	// Next() should block waiting for block 1. Cancel the context to unblock it.
+	done := make(chan any)
+	var result ledger.QueryResult
+	var nextErr error
+	go func() {
+		result, nextErr = itr.Next()
+		close(done)
+	}()
+
+	// Cancel the context — this should wake up waitForBlock
+	cancel()
+
+	select {
+	case <-done:
+	case <-time.After(5 * time.Second):
+		t.Fatal("Next() did not return after context cancellation")
+	}
+
+	require.Nil(t, result)
+	require.ErrorIs(t, nextErr, context.Canceled)
+}
+
 func constructLedgerid(id int) string {
 	return fmt.Sprintf("ledger_%d", id)
 }

common/ledger/blkstorage/blockstore_provider_test.go

@@ -465,7 +465,7 @@ func verifyQueriesOnBlocksPriorToSnapshot(
 		_, err := bootstrappedBlockStore.RetrieveBlockByNumber(blockNum)
 		require.EqualError(t, err, expectedErrStr)
 
-		_, err = bootstrappedBlockStore.RetrieveBlocks(blockNum)
+		_, err = bootstrappedBlockStore.RetrieveBlocks(t.Context(), blockNum)
 		require.EqualError(t, err, expectedErrStr)
 
 		_, err = bootstrappedBlockStore.RetrieveTxByBlockNumTranNum(blockNum, 0)
@@ -517,7 +517,7 @@ func verifyQueriesOnBlocksAddedAfterBootstrapping(t *testing.T,
 		require.NoError(t, err)
 		require.Equal(t, b, retrievedBlock)
 
-		itr, err := bootstrappedBlockStore.RetrieveBlocks(b.Header.Number)
+		itr, err := bootstrappedBlockStore.RetrieveBlocks(t.Context(), b.Header.Number)
 		require.NoError(t, err)
 		blk, err := itr.Next()
 		require.NoError(t, err)