fix(db/remote): Fix remote db issues

blockchain/statebackend/deprecated.go

@@ -1,6 +1,8 @@
 package statebackend
 
 import (
+	"errors"
+
 	"github.com/NethermindEth/juno/core"
 	"github.com/NethermindEth/juno/core/deprecatedstate"
 	"github.com/NethermindEth/juno/core/felt"
@@ -20,10 +22,12 @@ func (b *deprecatedStateBackend) HeadState() (core.StateReader, StateCloser, err
 	// Fail early if no block has been committed (no head state to open)
 	// only the key's existence matters, not the height itself.
 	if _, err := core.GetChainHeight(txn); err != nil {
-		return nil, nil, err
+		return nil, nil, errors.Join(err, txn.Close())
 	}
 
-	return deprecatedstate.New(txn), NoopStateCloser, nil
+	// The batch has to be closed: on a remote DB it is a gRPC stream, and the
+	// server holds a batch of its own open for as long as the stream lives.
+	return deprecatedstate.New(txn), txn.Close, nil
 }
 
 func (b *deprecatedStateBackend) StateAtBlockNumber(
@@ -38,7 +42,7 @@ func (b *deprecatedStateBackend) StateAtBlockNumber(
 	return deprecatedstate.NewHistory(
 		deprecatedstate.New(txn),
 		blockNumber,
-	), NoopStateCloser, nil
+	), txn.Close, nil
 }
 
 func (b *deprecatedStateBackend) StateAtBlockHash(
@@ -48,7 +52,7 @@ func (b *deprecatedStateBackend) StateAtBlockHash(
 	if blockHash.IsZero() {
 		memDB := memory.New()
 		txn := memDB.NewIndexedBatch()
-		return deprecatedstate.New(txn), NoopStateCloser, nil
+		return deprecatedstate.New(txn), txn.Close, nil
 	}
 
 	blockNumber, err := pruner.BlockNumberByHashIfStateRetained(b.database, blockHash)
@@ -60,7 +64,7 @@ func (b *deprecatedStateBackend) StateAtBlockHash(
 	return deprecatedstate.NewHistory(
 		deprecatedstate.New(txn),
 		blockNumber,
-	), NoopStateCloser, nil
+	), txn.Close, nil
 }
 
 func (b *deprecatedStateBackend) Store(
@@ -134,6 +138,8 @@ func (b *deprecatedStateBackend) RevertHead() error {
 func (b *deprecatedStateBackend) GetReverseStateDiff() (core.StateDiff, error) {
 	//nolint:staticcheck,nolintlint // used by old state
 	txn := b.database.NewIndexedBatch()
+	defer txn.Close()
+
 	blockNum, err := core.GetChainHeight(txn)
 	if err != nil {
 		return core.StateDiff{}, err

db/remote/db.go

@@ -11,6 +11,7 @@
 	"github.com/NethermindEth/juno/db"
 	"github.com/NethermindEth/juno/grpc/gen"
 	"github.com/NethermindEth/juno/utils/log"
+	"go.uber.org/zap"
 	"google.golang.org/grpc"
 )
 
@@ -54,12 +55,19 @@
 func (d *DB) NewTransaction(write bool) (*transaction, error) {
 	defer d.listener.OnIO(write, time.Now())
 
-	txClient, err := d.kvClient.Tx(d.ctx, grpc.MaxCallSendMsgSize(math.MaxInt), grpc.MaxCallRecvMsgSize(math.MaxInt))
+	// Every transaction owns a stream, so it needs its own context to release it.
+	ctx, cancel := context.WithCancel(d.ctx)
+	txClient, err := d.kvClient.Tx(
+		ctx,
+		grpc.MaxCallSendMsgSize(math.MaxInt),
+		grpc.MaxCallRecvMsgSize(math.MaxInt),
+	)
 	if err != nil {
+		cancel()
 		return nil, err
 	}
 
-	return &transaction{client: txClient, logger: d.logger}, nil
+	return &transaction{client: txClient, cancel: cancel, logger: d.logger}, nil
 }
 
 func (d *DB) Update(fn func(txn db.IndexedBatch) error) error {
@@ -83,10 +91,10 @@
 
 	batch := d.NewBatch()
 	if err := fn(batch); err != nil {
-		return err
+		return errors.Join(err, batch.Close())
 	}
 
-	return batch.Write()
+	return errors.Join(batch.Write(), batch.Close())
 }
 
 func (d *DB) Close() error {
@@ -110,6 +118,7 @@
 	if err != nil {
 		return err
 	}
+	defer d.discard(txn)
 
 	return txn.Get(key, cb)
 }
@@ -119,6 +128,7 @@
 	if err != nil {
 		return false, err
 	}
+	defer d.discard(txn)
 
 	return txn.Has(key)
 }
@@ -128,29 +138,25 @@
 }
 
 func (d *DB) NewBatch() db.Batch {
-	defer d.listener.OnIO(false, time.Now())
-
-	txClient, err := d.kvClient.Tx(d.ctx, grpc.MaxCallSendMsgSize(math.MaxInt), grpc.MaxCallRecvMsgSize(math.MaxInt))
+	txn, err := d.NewTransaction(false)
 	if err != nil {
 		panic(err)
 	}
 
-	return &transaction{client: txClient, logger: d.logger}
+	return txn
 }
 
 func (d *DB) NewBatchWithSize(size int) db.Batch {
 	return d.NewBatch()
 }
 
 func (d *DB) NewIndexedBatch() db.IndexedBatch {
-	defer d.listener.OnIO(true, time.Now())
-
-	txClient, err := d.kvClient.Tx(d.ctx, grpc.MaxCallSendMsgSize(math.MaxInt), grpc.MaxCallRecvMsgSize(math.MaxInt))
+	txn, err := d.NewTransaction(true)
 	if err != nil {
 		panic(err)
 	}
 
-	return &transaction{client: txClient, logger: d.logger}
+	return txn
 }
 
 func (d *DB) NewIndexedBatchWithSize(size int) db.IndexedBatch {
@@ -163,25 +169,36 @@
 		return nil, err
 	}
 
-	return txn.NewIterator(start, withUpperBound)
+	it, err := txn.NewIterator(start, withUpperBound)
+	if err != nil {
+		return nil, errors.Join(err, txn.Discard())
+	}
+
+	return &ownedIterator{Iterator: it, txn: txn}, nil
 }
 
 func (d *DB) NewSnapshot() db.Snapshot {
-	defer d.listener.OnIO(false, time.Now())
-
-	txClient, err := d.kvClient.Tx(d.ctx, grpc.MaxCallSendMsgSize(math.MaxInt), grpc.MaxCallRecvMsgSize(math.MaxInt))
+	txn, err := d.NewTransaction(false)
 	if err != nil {
 		panic(err)
 	}
 
-	return &transaction{client: txClient, logger: d.logger}
+	return txn
 }
 
 func (d *DB) WithListener(listener db.EventListener) db.KeyValueStore {
 	d.listener = listener
 	return d
 }
 
+// discard releases a transaction the DB opened for a single call. A read has
+// nothing to report on close, so the error only reaches the log.
+func (d *DB) discard(txn *transaction) {
+	if err := txn.Discard(); err != nil {
+		d.logger.Debug("Discarding remote transaction", zap.Error(err))
+	}
+}
+
 func discardTxnOnPanic(txn *transaction) {
 	p := recover()
 	if p != nil {

db/remote/db_test.go

@@ -2,6 +2,7 @@ package remote
 
 import (
 	"net"
+	"slices"
 	"testing"
 
 	"github.com/NethermindEth/juno/db"
@@ -79,6 +80,21 @@ func TestRemote(t *testing.T) {
 		assert.Equal(t, foundKeys, byte(3))
 	})
 
+	t.Run("first", func(t *testing.T) {
+		snap := remoteDB.NewSnapshot()
+		defer snap.Close()
+
+		it, err := snap.NewIterator(nil, false)
+		require.NoError(t, err)
+		defer it.Close()
+
+		require.True(t, it.First())
+		assert.Equal(t, []byte{0}, it.Key())
+		v, err := it.Value()
+		require.NoError(t, err)
+		assert.Equal(t, []byte{0}, v)
+	})
+
 	t.Run("seek", func(t *testing.T) {
 		snap := remoteDB.NewSnapshot()
 		defer snap.Close()
@@ -104,3 +120,52 @@ func TestRemote(t *testing.T) {
 	})
 	grpcSrv.GracefulStop()
 }
+
+// TestRemoteIteratorBounds guards against the bounds being dropped on the wire:
+// a key sorting before the prefix and one sorting after it must not surface.
+func TestRemoteIteratorBounds(t *testing.T) {
+	memDB := memory.New()
+	batch := memDB.NewBatch()
+	keys := [][]byte{
+		{0x00, 0xFF},
+		{0x01, 0x00},
+		{0x01, 0x01},
+		{0x01, 0x02},
+		{0x02, 0x00},
+	}
+	for _, k := range keys {
+		require.NoError(t, batch.Put(k, k))
+	}
+	require.NoError(t, batch.Write())
+
+	grpcHandler := junogrpc.New(memDB, "0.0.0")
+	grpcSrv := grpc.NewServer()
+	gen.RegisterKVServer(grpcSrv, grpcHandler)
+
+	var lc net.ListenConfig
+	l, err := lc.Listen(t.Context(), "tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+	go func() {
+		require.NoError(t, grpcSrv.Serve(l))
+	}()
+	defer grpcSrv.GracefulStop()
+
+	remoteDB, err := New(
+		l.Addr().String(),
+		t.Context(),
+		log.NewNopZapLogger(),
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+	)
+	require.NoError(t, err)
+
+	// Top-level NewIterator, not a snapshot's, so this also exercises ownedIterator.
+	it, err := remoteDB.NewIterator([]byte{0x01}, true)
+	require.NoError(t, err)
+	defer it.Close()
+
+	var found [][]byte
+	for valid := it.First(); valid; valid = it.Next() {
+		found = append(found, slices.Clone(it.Key()))
+	}
+	assert.Equal(t, [][]byte{{0x01, 0x00}, {0x01, 0x01}, {0x01, 0x02}}, found)
+}

db/remote/iterator.go

@@ -3,6 +3,7 @@ package remote
 import (
 	"slices"
 
+	"github.com/NethermindEth/juno/db"
 	"github.com/NethermindEth/juno/grpc/gen"
 	"github.com/NethermindEth/juno/utils/log"
 	"go.uber.org/zap"
@@ -89,3 +90,18 @@ func (i *iterator) Seek(key []byte) bool {
 func (i *iterator) Close() error {
 	return i.doOpAndUpdate(gen.Op_CLOSE, nil)
 }
+
+// ownedIterator holds the only reference to its transaction, so closing it has
+// to release the stream. An iterator taken from a batch or a snapshot shares
+// that stream with its owner and must leave it alone.
+type ownedIterator struct {
+	db.Iterator
+	txn *transaction
+}
+
+// Close discards the transaction, which drops the iterator on the server too.
+// It skips [gen.Op_CLOSE]: the round trip is redundant and its error would
+// surface as a failure of the scan that has already finished.
+func (i *ownedIterator) Close() error {
+	return i.txn.Discard()
+}

db/remote/transaction.go

@@ -2,6 +2,7 @@ package remote
 
 import (
 	"bytes"
+	"context"
 	"errors"
 
 	"github.com/NethermindEth/juno/db"
@@ -21,13 +22,23 @@ var (
 
 type transaction struct {
 	client gen.KV_TxClient
+	cancel context.CancelFunc
 	logger log.StructuredLogger
 }
 
-func (t *transaction) NewIterator(_ []byte, _ bool) (db.Iterator, error) {
-	err := t.client.Send(&gen.Cursor{
-		Op: gen.Op_OPEN,
-	})
+func (t *transaction) NewIterator(prefix []byte, withUpperBound bool) (db.Iterator, error) {
+	// The remote iterator has to be created with the same bounds as a local one,
+	// otherwise it scans the whole database and First returns a foreign key.
+	// BucketName carries the prefix and a non-empty V asks for the upper bound.
+	cursor := &gen.Cursor{
+		Op:         gen.Op_OPEN,
+		BucketName: prefix,
+	}
+	if withUpperBound {
+		cursor.V = []byte{1}
+	}
+
+	err := t.client.Send(cursor)
 	if err != nil {
 		return nil, err
 	}
@@ -44,8 +55,13 @@ func (t *transaction) NewIterator(_ []byte, _ bool) (db.Iterator, error) {
 	}, nil
 }
 
+// Discard releases the stream. Closing the send side lets the server return and
+// drop the batch it holds open; the cancel releases the client side, which
+// otherwise waits for the response stream to be drained to EOF.
 func (t *transaction) Discard() error {
-	return t.client.CloseSend()
+	err := t.client.CloseSend()
+	t.cancel()
+	return err
 }
 
 func (t *transaction) Commit() error {
@@ -101,4 +117,4 @@ func (t *transaction) Put(key, val []byte) error {
 func (t *transaction) Size() int    { return 0 }
 func (t *transaction) Reset()       {}
 func (t *transaction) Write() error { return nil }
-func (t *transaction) Close() error { return t.client.CloseSend() }
+func (t *transaction) Close() error { return t.Discard() }

grpc/handlers.go

@@ -68,7 +68,7 @@ func (h Handler) handleTxCursor(
 
 	// open is special case: it's the only way to receive cursor id
 	if cur.Op == gen.Op_OPEN {
-		cursorID, err := tx.newCursor()
+		cursorID, err := tx.newCursor(cur.BucketName, len(cur.V) > 0)
 		if err != nil {
 			return err
 		}
@@ -95,6 +95,11 @@ func (h Handler) handleTxCursor(
 	responsePair.CursorId = cur.Cursor
 
 	switch cur.Op {
+	case gen.Op_FIRST:
+		if it.First() {
+			responsePair.K = it.Key()
+			responsePair.V, err = it.Value()
+		}
 	case gen.Op_SEEK:
 		key := slices.Concat(cur.BucketName, cur.K)
 		if it.Seek(key) {

grpc/kv.proto

@@ -12,6 +12,7 @@ service KV {
 }
 
 // values from https://github.com/ledgerwatch/interfaces/blob/master/remote/kv.proto#L68
+// FIRST is the enum zero-value, so a Cursor with an unset op silently behaves as FIRST.
 enum Op {
   FIRST = 0;
   SEEK = 1;
@@ -28,7 +29,7 @@ message Cursor {
   bytes bucket_name = 2;
   uint32 cursor = 3;
   bytes k = 4;
-  bytes v = 5; // not used
+  bytes v = 5; // withUpperBound flag on OPEN; unused for other ops
 }
 
 message Pair {