added flush memtable functionality

Author: andrewnguyen22

lib/error.go

@@ -308,6 +308,8 @@ const (
 	CodeInvalidMerkleTreeProof ErrorCode   = 11
 	CodeGarbageCollectDB       ErrorCode   = 12
 	CodeSetEntry               ErrorCode   = 13
+	CodeReadBytes              ErrorCode   = 14
+	CodeFlushMemTable          ErrorCode   = 15
 
 	RPCModule             ErrorModule = "rpc"
 	CodeRPCTimeout        ErrorCode   = 1

store/error.go

@@ -57,3 +57,11 @@ func ErrInvalidMerkleTree() lib.ErrorI {
 func ErrInvalidMerkleTreeProof() lib.ErrorI {
 	return lib.NewError(lib.CodeInvalidMerkleTreeProof, lib.StorageModule, "merkle tree proof is invalid")
 }
+
+func ErrReadBytes(err error) lib.ErrorI {
+	return lib.NewError(lib.CodeReadBytes, lib.StorageModule, fmt.Sprintf("random read bytes failed with err: %s", err.Error()))
+}
+
+func ErrFlushMemTable(err error) lib.ErrorI {
+	return lib.NewError(lib.CodeFlushMemTable, lib.StorageModule, fmt.Sprintf("flush memtable failed with err: %s", err.Error()))
+}

store/store.go

@@ -134,7 +134,6 @@ func (s *Store) NewReadOnly(queryVersion uint64) (lib.StoreI, lib.ErrorI) {
 	var useHistorical bool
 	// if the query version is older than the partition frequency
 	if queryVersion+1 < partitionHeight(s.version) {
-		fmt.Println("USING HISTORICAL")
 		useHistorical = true
 	}
 	// make a reader for the specified version
@@ -283,7 +282,7 @@ func (s *Store) Partition() {
 			// if the item is 'deleted'
 			if it.Deleted() {
 				// set the item as deleted at the partition height and discard earlier versions
-				if e := writer.SetEntryAt(newEntry(lib.Append([]byte(latestStatePrefix), k), nil, badgerDeleteBit), it.Version()); e != nil {
+				if e := writer.SetEntryAt(newEntry(lib.Append([]byte(latestStatePrefix), k), nil, badgerDeleteBit|badgerDiscardEarlierVersions), it.Version()); e != nil {
 					return ErrSetEntry(e)
 				}
 			} else {
@@ -296,7 +295,7 @@ func (s *Store) Partition() {
 					return ErrSetEntry(e)
 				}
 				// re-write the latest version with the 'discard' flag set
-				if e := writer.SetEntryAt(newEntry(lib.Append([]byte(latestStatePrefix), k), v, badgerNoDiscardBit), snapshotHeight); e != nil {
+				if e := writer.SetEntryAt(newEntry(lib.Append([]byte(latestStatePrefix), k), v, badgerDiscardEarlierVersions), snapshotHeight); e != nil {
 					return ErrSetEntry(e)
 				}
 			}
@@ -311,11 +310,14 @@ func (s *Store) Partition() {
 		}
 		// if the partition height is past the partition frequency, set the discardTs at the partition height-1
 		if snapshotHeight > partitionFrequency {
-			fmt.Println("SETTING DISCARD TS @", snapshotHeight-2)
 			sc.db.SetDiscardTs(snapshotHeight - 2)
 		}
 		// if the GC isn't already running
 		if !s.isGarbageCollecting.Swap(true) {
+			// force the mem table to flush to the LSM
+			if err = FlushMemTable(sc.db); err != nil {
+				return err
+			}
 			// run LSM compaction
 			if fe := sc.db.Flatten(32); fe != nil {
 				return ErrGarbageCollectDB(fe)

store/store_test.go

@@ -391,6 +391,71 @@ func TestPartitionIntegration(t *testing.T) {
 	require.Len(t, deletedKeys, 0)
 }
 
+func TestFlushMemtable(t *testing.T) {
+	// create temp directory for test db
+	tempDir := t.TempDir()
+	path := filepath.Join(tempDir, "test-db")
+
+	// set up db with a configuration that will trigger a GC after multiple partitions
+	db, err := badger.OpenManaged(badger.DefaultOptions(path).WithNumVersionsToKeep(math.MaxInt64).
+		WithValueThreshold(1).WithLoggingLevel(badger.INFO).WithMemTableSize(int64(units.GB)))
+	require.NoError(t, err)
+	// set up the store
+	store, err := NewStoreWithDB(db, nil, lib.NewDefaultLogger(), true)
+	require.NoError(t, err)
+	defer func() {
+		require.NoError(t, store.Close())
+	}()
+
+	// insert large dataset at various heights before partition boundary
+	var k []byte
+	keys := [][]byte{
+		bytes.Repeat([]byte("0"), 32),
+		bytes.Repeat([]byte("1"), 32),
+		bytes.Repeat([]byte("2"), 32),
+		bytes.Repeat([]byte("3"), 32),
+		bytes.Repeat([]byte("4"), 32),
+		bytes.Repeat([]byte("5"), 32),
+		bytes.Repeat([]byte("6"), 32),
+		bytes.Repeat([]byte("7"), 32),
+		bytes.Repeat([]byte("8"), 32),
+		bytes.Repeat([]byte("9"), 32),
+	}
+	iterations := 1000
+	for i := 1; i < iterations; i++ {
+		// if not the first iteration
+		if i != 0 {
+			// delete the last key in the db
+			require.NoError(t, store.Delete(k))
+		}
+		// use a key to be updated repeatedly to generate versions
+		k = keys[i%10]
+
+		// set the value in the db
+		require.NoError(t, store.Set(k, nil))
+
+		// commit regularly to create multiple versions
+		_, err = store.Commit()
+		require.NoError(t, err)
+	}
+	db.SetDiscardTs(uint64(iterations - 1))
+	require.NoError(t, FlushMemTable(db))
+	require.NoError(t, db.Flatten(32))
+
+	// use an archive iterator to iterate through the deleted keys
+	iterator, err := store.lss.ArchiveIterator(nil)
+	require.NoError(t, err)
+	it := iterator.(*Iterator)
+	defer it.Close()
+
+	numKeys := 0
+	for ; it.Valid(); it.Next() {
+		numKeys++
+	}
+
+	require.True(t, numKeys == 1)
+}
+
 func TestHistoricalPrefix(t *testing.T) {
 	tests := []struct {
 		name     string

store/wrapper_txn.go

@@ -2,6 +2,8 @@ package store
 
 import (
 	"bytes"
+	"crypto/rand"
+	"math"
 	"reflect"
 	"unsafe"
 
@@ -25,7 +27,7 @@ const (
 	badgerDiscardEarlierVersions byte = 1 << 2  // badgerDB 'discard earlier versions' flag
 	badgerDeleteBit              byte = 1 << 0  // badgerDB 'tombstoned' flag
 	badgerNoDiscardBit           byte = 1 << 3  // badgerDB 'never discard'  bit
-	badgerGCRatio                     = .000001 // the ratio when badgerDB will run the garbage collector
+	badgerGCRatio                     = .15     // the ratio when badgerDB will run the garbage collector
 	badgerSizeFieldName               = "size"  // badgerDB Txn 'size' field name
 	badgerCountFieldName              = "count" // badgerDB Txn 'count' field name
 	badgerTxnFieldName                = "txn"   // badgerDB WriteBatch 'txn' field name
@@ -206,6 +208,32 @@ func newEntry(key, value []byte, meta byte) (e *badger.Entry) {
 	return
 }
 
+// FlushMemTable() ensures badgerDB is flushing its mem table before running flatten
+// IMPORTANT - discardTs must be set before this
+func FlushMemTable(db *badger.DB) lib.ErrorI {
+	// get random 32 bytes
+	randomPrefix := make([]byte, 32)
+	if _, err := rand.Read(randomPrefix); err != nil {
+		return ErrReadBytes(err)
+	}
+	// create a new transaction to write to the database
+	tx := db.NewTransactionAt(math.MaxUint64, true)
+	// write the random prefix to the database
+	if err := tx.Set(randomPrefix, nil); err != nil {
+		return ErrSetEntry(err)
+	}
+	// commit the transaction
+	if err := tx.CommitAt(math.MaxUint64, nil); err != nil {
+		return ErrCommitDB(err)
+	}
+	// call drop prefix which triggers the mempool flush
+	// NOTE: this only works if an actual prefix exists
+	if err := db.DropPrefix(randomPrefix); err != nil {
+		return ErrFlushMemTable(err)
+	}
+	return nil
+}
+
 // setMeta() accesses the private field 'meta' of badgerDB's `Entry`
 // badger doesn't yet allow users to explicitly set keys as *do not discard*
 // https://github.com/hypermodeinc/badger/issues/2192