Mark corruptabledb as corrupted upon an iterator error (#1829)

Author: Dan Laine

database/corruptabledb/db.go

@@ -89,6 +89,34 @@ func (db *Database) NewBatch() database.Batch {
 	}
 }
 
+func (db *Database) NewIterator() database.Iterator {
+	return &iterator{
+		Iterator: db.Database.NewIterator(),
+		db:       db,
+	}
+}
+
+func (db *Database) NewIteratorWithStart(start []byte) database.Iterator {
+	return &iterator{
+		Iterator: db.Database.NewIteratorWithStart(start),
+		db:       db,
+	}
+}
+
+func (db *Database) NewIteratorWithPrefix(prefix []byte) database.Iterator {
+	return &iterator{
+		Iterator: db.Database.NewIteratorWithPrefix(prefix),
+		db:       db,
+	}
+}
+
+func (db *Database) NewIteratorWithStartAndPrefix(start, prefix []byte) database.Iterator {
+	return &iterator{
+		Iterator: db.Database.NewIteratorWithStartAndPrefix(start, prefix),
+		db:       db,
+	}
+}
+
 func (db *Database) corrupted() error {
 	db.errorLock.RLock()
 	defer db.errorLock.RUnlock()
@@ -127,3 +155,24 @@ func (b *batch) Write() error {
 	}
 	return b.db.handleError(b.Batch.Write())
 }
+
+type iterator struct {
+	database.Iterator
+	db *Database
+}
+
+func (it *iterator) Next() bool {
+	if err := it.db.corrupted(); err != nil {
+		return false
+	}
+	val := it.Iterator.Next()
+	_ = it.db.handleError(it.Iterator.Error())
+	return val
+}
+
+func (it *iterator) Error() error {
+	if err := it.db.corrupted(); err != nil {
+		return err
+	}
+	return it.db.handleError(it.Iterator.Error())
+}

database/corruptabledb/db_test.go

@@ -10,6 +10,8 @@ import (
 
 	"github.com/stretchr/testify/require"
 
+	"go.uber.org/mock/gomock"
+
 	"github.com/ava-labs/avalanchego/database"
 	"github.com/ava-labs/avalanchego/database/memdb"
 )
@@ -75,3 +77,124 @@ func TestCorruption(t *testing.T) {
 		})
 	}
 }
+
+func TestIterator(t *testing.T) {
+	errIter := errors.New("iterator error")
+
+	type test struct {
+		name              string
+		databaseErrBefore error
+		modifyIter        func(*gomock.Controller, *iterator)
+		op                func(*require.Assertions, *iterator)
+		expectedErr       error
+	}
+
+	tests := []test{
+		{
+			name:              "corrupted database; Next",
+			databaseErrBefore: errTest,
+			expectedErr:       errTest,
+			modifyIter:        func(ctrl *gomock.Controller, iter *iterator) {},
+			op: func(require *require.Assertions, iter *iterator) {
+				require.False(iter.Next())
+			},
+		},
+		{
+			name:              "Next corrupts database",
+			databaseErrBefore: nil,
+			expectedErr:       errIter,
+			modifyIter: func(ctrl *gomock.Controller, iter *iterator) {
+				mockInnerIter := database.NewMockIterator(ctrl)
+				mockInnerIter.EXPECT().Next().Return(false)
+				mockInnerIter.EXPECT().Error().Return(errIter)
+				iter.Iterator = mockInnerIter
+			},
+			op: func(require *require.Assertions, iter *iterator) {
+				require.False(iter.Next())
+			},
+		},
+		{
+			name:              "corrupted database; Error",
+			databaseErrBefore: errTest,
+			expectedErr:       errTest,
+			modifyIter:        func(ctrl *gomock.Controller, iter *iterator) {},
+			op: func(require *require.Assertions, iter *iterator) {
+				err := iter.Error()
+				require.ErrorIs(err, errTest)
+			},
+		},
+		{
+			name:              "Error corrupts database",
+			databaseErrBefore: nil,
+			expectedErr:       errIter,
+			modifyIter: func(ctrl *gomock.Controller, iter *iterator) {
+				mockInnerIter := database.NewMockIterator(ctrl)
+				mockInnerIter.EXPECT().Error().Return(errIter)
+				iter.Iterator = mockInnerIter
+			},
+			op: func(require *require.Assertions, iter *iterator) {
+				err := iter.Error()
+				require.ErrorIs(err, errIter)
+			},
+		},
+		{
+			name:              "corrupted database; Key",
+			databaseErrBefore: errTest,
+			expectedErr:       errTest,
+			modifyIter:        func(ctrl *gomock.Controller, iter *iterator) {},
+			op: func(require *require.Assertions, iter *iterator) {
+				_ = iter.Key()
+			},
+		},
+		{
+			name:              "corrupted database; Value",
+			databaseErrBefore: errTest,
+			expectedErr:       errTest,
+			modifyIter:        func(ctrl *gomock.Controller, iter *iterator) {},
+			op: func(require *require.Assertions, iter *iterator) {
+				_ = iter.Value()
+			},
+		},
+		{
+			name:              "corrupted database; Release",
+			databaseErrBefore: errTest,
+			expectedErr:       errTest,
+			modifyIter:        func(ctrl *gomock.Controller, iter *iterator) {},
+			op: func(require *require.Assertions, iter *iterator) {
+				iter.Release()
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			require := require.New(t)
+			ctrl := gomock.NewController(t)
+
+			// Make a database
+			baseDB := memdb.New()
+			corruptableDB := New(baseDB)
+
+			// Put a key-value pair in the database.
+			require.NoError(corruptableDB.Put([]byte{0}, []byte{1}))
+
+			// Mark database as corupted, if applicable
+			_ = corruptableDB.handleError(tt.databaseErrBefore)
+
+			// Make an iterator
+			iter := &iterator{
+				Iterator: corruptableDB.NewIterator(),
+				db:       corruptableDB,
+			}
+
+			// Modify the iterator (optional)
+			tt.modifyIter(ctrl, iter)
+
+			// Do an iterator operation
+			tt.op(require, iter)
+
+			err := corruptableDB.corrupted()
+			require.ErrorIs(err, tt.expectedErr)
+		})
+	}
+}

database/mock_iterator.go

@@ -2,6 +2,7 @@ github.com/ava-labs/avalanchego/api/server=Server=api/server/mock_server.go
 github.com/ava-labs/avalanchego/chains/atomic=SharedMemory=chains/atomic/mock_shared_memory.go
 github.com/ava-labs/avalanchego/codec=Manager=codec/mock_manager.go
 github.com/ava-labs/avalanchego/database=Batch=database/mock_batch.go
+github.com/ava-labs/avalanchego/database=Iterator=database/mock_iterator.go
 github.com/ava-labs/avalanchego/message=OutboundMessage=message/mock_message.go
 github.com/ava-labs/avalanchego/message=OutboundMsgBuilder=message/mock_outbound_message_builder.go
 github.com/ava-labs/avalanchego/network/peer=GossipTracker=network/peer/mock_gossip_tracker.go