Cleanup and unit tests.

Signed-off-by: Cody Littley <[email protected]>
cody-littley · Sep 30, 2024 · 091e6bb · 091e6bb
1 parent dfe422a
commit 091e6bb
Show file tree

Hide file tree

Showing 9 changed files with 87 additions and 25 deletions.
diff --git a/common/kvstore/batch.go b/common/kvstore/batch.go
@@ -2,7 +2,6 @@ package kvstore
 
 // Batch is a collection of operations that can be applied atomically to a TableStore.
 type Batch[T any] interface {
-	// TODO test nil value case
 	// Put stores the given key / value pair in the batch, overwriting any existing value for that key.
 	// If nil is passed as the value, a byte slice of length 0 will be stored.
 	Put(key T, value []byte)

diff --git a/common/kvstore/leveldb/leveldb_store.go b/common/kvstore/leveldb/leveldb_store.go
@@ -41,6 +41,9 @@ func NewStore(logger logging.Logger, path string) (kvstore.Store, error) {
 
 // Put stores a data in the store.
 func (store *levelDBStore) Put(key []byte, value []byte) error {
+	if value == nil {
+		value = []byte{}
+	}
 	return store.db.Put(key, value, nil)
 }
 
@@ -99,7 +102,7 @@ type levelDBBatch struct {
 
 func (m *levelDBBatch) Put(key []byte, value []byte) {
 	if value == nil {
-		value = []byte{0}
+		value = []byte{}
 	}
 	m.batch.Put(key, value)
 }

diff --git a/common/kvstore/mapstore/map_store.go b/common/kvstore/mapstore/map_store.go
@@ -27,6 +27,10 @@ func NewStore() kvstore.Store {
 
 // Put adds a key-value pair to the store.
 func (store *mapStore) Put(key []byte, value []byte) error {
+	if value == nil {
+		value = []byte{}
+	}
+
 	store.lock.Lock()
 	defer store.lock.Unlock()
 
@@ -93,7 +97,7 @@ type batch struct {
 // Put stores the given key / value pair in the batch, overwriting any existing value for that key.
 func (m *batch) Put(key []byte, value []byte) {
 	if value == nil {
-		value = []byte{0}
+		value = []byte{}
 	}
 	m.keys = append(m.keys, key)
 	m.values = append(m.values, value)

diff --git a/common/kvstore/store.go b/common/kvstore/store.go
@@ -17,6 +17,7 @@ type Store interface {
 	BatchOperator[[]byte]
 
 	// Put stores the given key / value pair in the database, overwriting any existing value for that key.
+	// If nil is passed as the value, a byte slice of length 0 will be stored.
 	Put(key []byte, value []byte) error
 
 	// Get retrieves the value for the given key. Returns a ErrNotFound error if the key does not exist.

diff --git a/common/kvstore/tablestore/table_store_test.go b/common/kvstore/tablestore/table_store_test.go
@@ -34,7 +34,7 @@ func TestTableCount(t *testing.T) {
 	assert.NoError(t, err)
 
 	base := mapstore.NewStore()
-	store, err := TableStoreWrapper(logger, base)
+	store, err := Wrapper(logger, base)
 	assert.NoError(t, err)
 
 	// table count needs to fit into 32 bytes, and two tables are reserved for internal use
@@ -69,7 +69,7 @@ func TestTableList(t *testing.T) {
 
 	base, err := leveldb.NewStore(logger, dbPath)
 	assert.NoError(t, err)
-	store, err := TableStoreWrapper(logger, base)
+	store, err := Wrapper(logger, base)
 	assert.NoError(t, err)
 
 	tables := store.GetTables()
@@ -123,7 +123,7 @@ func TestTableList(t *testing.T) {
 
 	base, err = leveldb.NewStore(logger, dbPath)
 	assert.NoError(t, err)
-	store, err = TableStoreWrapper(logger, base)
+	store, err = Wrapper(logger, base)
 	assert.NoError(t, err)
 
 	tables = store.GetTables()
@@ -152,7 +152,7 @@ func TestTableList(t *testing.T) {
 
 	base, err = leveldb.NewStore(logger, dbPath)
 	assert.NoError(t, err)
-	store, err = TableStoreWrapper(logger, base)
+	store, err = Wrapper(logger, base)
 	assert.NoError(t, err)
 
 	tables = store.GetTables()
@@ -196,7 +196,7 @@ func TestUniqueKeySpace(t *testing.T) {
 	assert.NoError(t, err)
 
 	base := mapstore.NewStore()
-	store, err := TableStoreWrapper(logger, base)
+	store, err := Wrapper(logger, base)
 	assert.NoError(t, err)
 
 	table1, err := store.GetTable("table1")
@@ -237,7 +237,7 @@ func TestBatchOperations(t *testing.T) {
 	assert.NoError(t, err)
 
 	base := mapstore.NewStore()
-	store, err := TableStoreWrapper(logger, base)
+	store, err := Wrapper(logger, base)
 	assert.NoError(t, err)
 
 	table1, err := store.GetTable("table1")
@@ -413,7 +413,7 @@ func TestDropTable(t *testing.T) {
 	assert.NoError(t, err)
 
 	base := mapstore.NewStore()
-	store, err := TableStoreWrapper(logger, base)
+	store, err := Wrapper(logger, base)
 	assert.NoError(t, err)
 
 	table1, err := store.GetTable("table1")
@@ -530,7 +530,7 @@ func TestIteration(t *testing.T) {
 	assert.NoError(t, err)
 
 	base := mapstore.NewStore()
-	store, err := TableStoreWrapper(logger, base)
+	store, err := Wrapper(logger, base)
 	assert.NoError(t, err)
 
 	table1, err := store.GetTable("table1")
@@ -685,7 +685,7 @@ func TestRestart(t *testing.T) {
 
 	base, err := leveldb.NewStore(logger, dbPath)
 	assert.NoError(t, err)
-	store, err := TableStoreWrapper(logger, base)
+	store, err := Wrapper(logger, base)
 	assert.NoError(t, err)
 
 	table1, err := store.GetTable("table1")
@@ -716,7 +716,7 @@ func TestRestart(t *testing.T) {
 
 	base, err = leveldb.NewStore(logger, dbPath)
 	assert.NoError(t, err)
-	store, err = TableStoreWrapper(logger, base)
+	store, err = Wrapper(logger, base)
 	assert.NoError(t, err)
 
 	table1, err = store.GetTable("table1")
@@ -765,7 +765,7 @@ func TestRandomOperations(t *testing.T) {
 
 	base, err := leveldb.NewStore(logger, dbPath)
 	assert.NoError(t, err)
-	store, err := TableStoreWrapper(logger, base)
+	store, err := Wrapper(logger, base)
 	assert.NoError(t, err)
 
 	tables := make(map[string]kvstore.Table)
@@ -782,7 +782,7 @@ func TestRandomOperations(t *testing.T) {
 
 			base, err = leveldb.NewStore(logger, dbPath)
 			assert.NoError(t, err)
-			store, err = TableStoreWrapper(logger, base)
+			store, err = Wrapper(logger, base)
 			assert.NoError(t, err)
 
 			for tableName := range tables {
@@ -929,7 +929,7 @@ func TestInterruptedTableDeletion(t *testing.T) {
 		deletionsRemaining: 50,
 	}
 
-	store, err := TableStoreWrapper(logger, explodingBase)
+	store, err := Wrapper(logger, explodingBase)
 	assert.NoError(t, err)
 
 	// Create a few tables
@@ -964,7 +964,7 @@ func TestInterruptedTableDeletion(t *testing.T) {
 	// Restart the store. The table should be gone by the time the method returns.
 	base, err = leveldb.NewStore(logger, dbPath)
 	assert.NoError(t, err)
-	store, err = TableStoreWrapper(logger, base)
+	store, err = Wrapper(logger, base)
 	assert.NoError(t, err)
 
 	tables := store.GetTables()

diff --git a/common/kvstore/tablestore/table_store_wrapper.go b/common/kvstore/tablestore/table_store_wrapper.go
@@ -60,12 +60,12 @@ type tableStore struct {
 	namespaceTable kvstore.Table
 }
 
-// TableStoreWrapper wraps the given Store to create a TableStore.
+// Wrapper wraps the given Store to create a TableStore.
 //
 // WARNING: it is not safe to access the wrapped store directly while the TableStore is in use. The TableStore uses
 // special key formatting, and direct access to the wrapped store may violate the TableStore's invariants, resulting
 // in undefined behavior.
-func TableStoreWrapper(logger logging.Logger, base kvstore.Store) (kvstore.TableStore, error) {
+func Wrapper(logger logging.Logger, base kvstore.Store) (kvstore.TableStore, error) {
 
 	tableIDMap := make(map[string]uint32)
 	tableIdSet := make(map[uint32]bool)
@@ -313,6 +313,9 @@ type tableStoreBatch struct {
 
 // Put adds a key-value pair to the batch.
 func (t *tableStoreBatch) Put(key kvstore.TableKey, value []byte) {
+	if value == nil {
+		value = []byte{}
+	}
 	t.batch.Put(key, value)
 }
 

diff --git a/common/kvstore/tablestore/table_view.go b/common/kvstore/tablestore/table_view.go
@@ -39,6 +39,10 @@ func (t *tableView) Name() string {
 
 // Put adds a key-value pair to the table.
 func (t *tableView) Put(key []byte, value []byte) error {
+	if value == nil {
+		value = []byte{}
+	}
+
 	k := t.TableKey(key)
 	return t.base.Put(k, value)
 }
@@ -140,6 +144,9 @@ type tableBatch struct {
 
 // Put schedules a key-value pair to be added to the table.
 func (t *tableBatch) Put(key []byte, value []byte) {
+	if value == nil {
+		value = []byte{}
+	}
 	k := t.table.TableKey(key)
 	t.batch.Put(k, value)
 }

diff --git a/common/kvstore/test/store_test.go b/common/kvstore/test/store_test.go
@@ -37,7 +37,7 @@ var storeBuilders = []func(logger logging.Logger, path string) (kvstore.Store, e
 	},
 	func(logger logging.Logger, path string) (kvstore.Store, error) {
 		store := mapstore.NewStore()
-		tableStore, err := tablestore.TableStoreWrapper(logger, store)
+		tableStore, err := tablestore.Wrapper(logger, store)
 		if err != nil {
 			return nil, err
 		}
@@ -52,7 +52,7 @@ var storeBuilders = []func(logger logging.Logger, path string) (kvstore.Store, e
 		if err != nil {
 			return nil, err
 		}
-		tableStore, err := tablestore.TableStoreWrapper(logger, store)
+		tableStore, err := tablestore.Wrapper(logger, store)
 		if err != nil {
 			return nil, err
 		}
@@ -160,8 +160,6 @@ func TestRandomOperations(t *testing.T) {
 	}
 }
 
-// TODO test mixed batch
-
 func writeBatchTest(t *testing.T, store kvstore.Store) {
 	tu.InitializeRandom()
 	deleteDBDirectory(t)
@@ -174,11 +172,22 @@ func writeBatchTest(t *testing.T, store kvstore.Store) {
 	for i := 0; i < 1000; i++ {
 		// Write a random value.
 		key := tu.RandomBytes(32)
-		value := tu.RandomBytes(32)
+
+		var value []byte
+		if i%50 == 0 {
+			// nil values are interpreted as empty slices.
+			value = nil
+		} else {
+			value = tu.RandomBytes(32)
+		}
 
 		batch.Put(key, value)
 
-		expectedData[string(key)] = value
+		if value == nil {
+			expectedData[string(key)] = []byte{}
+		} else {
+			expectedData[string(key)] = value
+		}
 
 		if i%10 == 0 {
 			// Every so often, apply the batch and check that the store matches the expected data.
@@ -421,3 +430,32 @@ func TestIterationWithPrefix(t *testing.T) {
 		iterationWithPrefixTest(t, store)
 	}
 }
+
+func putNilTest(t *testing.T, store kvstore.Store) {
+	tu.InitializeRandom()
+	deleteDBDirectory(t)
+
+	key := tu.RandomBytes(32)
+
+	err := store.Put(key, nil)
+	assert.NoError(t, err)
+
+	value, err := store.Get(key)
+	assert.NoError(t, err)
+	assert.Equal(t, []byte{}, value)
+
+	err = store.Destroy()
+	assert.NoError(t, err)
+	verifyDBIsDeleted(t)
+}
+
+func TestPutNil(t *testing.T) {
+	logger, err := common.NewLogger(common.DefaultLoggerConfig())
+	assert.NoError(t, err)
+
+	for _, builder := range storeBuilders {
+		store, err := builder(logger, dbPath)
+		assert.NoError(t, err)
+		putNilTest(t, store)
+	}
+}
diff --git a/common/kvstore/ttl/ttl_wrapper.go b/common/kvstore/ttl/ttl_wrapper.go
@@ -191,6 +191,10 @@ func (store *ttlStore) expireKeys(now time.Time) error {
 }
 
 func (store *ttlStore) Put(key []byte, value []byte) error {
+	if value == nil {
+		value = []byte{}
+	}
+
 	prefixedKey := append(keyPrefix, key...)
 	return store.store.Put(prefixedKey, value)
 }
@@ -212,6 +216,9 @@ type batch struct {
 }
 
 func (b *batch) Put(key []byte, value []byte) {
+	if value == nil {
+		value = []byte{}
+	}
 	prefixedKey := append(keyPrefix, key...)
 	b.base.Put(prefixedKey, value)
 }