cleanup partition changes

Author: andrewnguyen22

store/store.go

@@ -90,7 +90,7 @@ func NewStore(path string, metrics *lib.Metrics, log lib.LoggerI) (lib.StoreI, l
 	// single batch. It is seemingly unknown why the 10% limit is set
 	// https://discuss.dgraph.io/t/discussion-badgerdb-should-offer-arbitrarily-sized-atomic-transactions/8736
 	db, err := badger.OpenManaged(badger.DefaultOptions(path).WithNumVersionsToKeep(math.MaxInt64).
-		WithLoggingLevel(badger.DEBUG).WithMemTableSize(int64(1*units.GB + 280*units.MB)))
+		WithLoggingLevel(badger.ERROR).WithMemTableSize(int64(1*units.GB + 280*units.MB)))
 	if err != nil {
 		return nil, ErrOpenDB(err)
 	}
@@ -272,27 +272,19 @@ func (s *Store) Partition() {
 			k, v := it.Key(), it.Value()
 			// skip items that are already marked for Garbage Collection
 			if deDuplicator.Found(lib.BytesToString(k)) {
-				// set the item as prunable
-				if e := writer.SetEntryAt(newEntry(lib.Append([]byte(latestStatePrefix), k), nil, badgerDeleteBit), it.Version()); e != nil {
-					return ErrSetEntry(e)
-				}
 				continue
 			}
 			// 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|badgerDiscardEarlierVersions), it.Version()); e != nil {
+				if e := writer.SetEntryAt(newEntry(lib.Append([]byte(latestStatePrefix), k), nil, badgerDeleteBit|badgerDiscardEarlierVersions), snapshotHeight); e != nil {
 					return ErrSetEntry(e)
 				}
 			} else {
 				// set item in the historical partition
 				if e := writer.SetEntryAt(newEntry(lib.Append(partitionPrefix, k), v, badgerNoDiscardBit), snapshotHeight); e != nil {
 					return ErrSetEntry(e)
 				}
-				// set the item as prunable
-				if e := writer.SetEntryAt(newEntry(lib.Append([]byte(latestStatePrefix), k), nil, badgerDeleteBit), it.Version()); e != nil {
-					return ErrSetEntry(e)
-				}
 				// re-write the latest version with the 'discard' flag set
 				if e := writer.SetEntryAt(newEntry(lib.Append([]byte(latestStatePrefix), k), v, badgerDiscardEarlierVersions), snapshotHeight); e != nil {
 					return ErrSetEntry(e)
@@ -363,7 +355,7 @@ func (s *Store) Set(k, v []byte) lib.ErrorI {
 		return err
 	}
 	// set in the state store @ historical partition
-	if err := s.hss.SetNonPruneable(k, v); err != nil {
+	if err := s.hss.Set(k, v); err != nil {
 		return err
 	}
 	// set in the state commit store
@@ -377,7 +369,7 @@ func (s *Store) Delete(k []byte) lib.ErrorI {
 		return err
 	}
 	// delete from the state store @ historical partition
-	if err := s.hss.DeleteNonPrunable(k); err != nil {
+	if err := s.hss.Tombstone(k); err != nil {
 		return err
 	}
 	// delete from the state commit store

store/store_test.go

@@ -391,14 +391,96 @@ func TestPartitionIntegration(t *testing.T) {
 	require.Len(t, deletedKeys, 0)
 }
 
+func TestBadgerBitBehavior(t *testing.T) {
+	tests := []struct {
+		name             string
+		nonPruneable     bool
+		finalBit         byte
+		expectedKeyCount int
+	}{
+		{
+			name:             "discard earlier versions bit",
+			finalBit:         badgerDiscardEarlierVersions,
+			expectedKeyCount: 1,
+		},
+		{
+			name:             "delete bit",
+			finalBit:         badgerDeleteBit,
+			expectedKeyCount: 0,
+		},
+		{
+			name:             "discard earlier versions bit",
+			finalBit:         badgerDiscardEarlierVersions | badgerNoDiscardBit,
+			expectedKeyCount: 1000,
+		},
+		{
+			name:             "delete bit",
+			finalBit:         badgerDeleteBit | badgerNoDiscardBit,
+			expectedKeyCount: 1000,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(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).
+				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
+			k := []byte("my_key")
+			iterations := 1000
+			for i := 0; i < iterations; i++ {
+				// set the value in the db
+				require.NoError(t, store.lss.Tombstone(k)) // tombstoned but mark keep
+
+				// commit regularly to create multiple versions
+				_, err = store.Commit()
+				require.NoError(t, err)
+			}
+			tx := db.NewTransactionAt(uint64(iterations), true)
+			// set an entry with a bit that marks it as deleted and prevents it from being discarded
+			require.NoError(t, tx.SetEntry(newEntry(lib.Append([]byte(latestStatePrefix), k), nil, test.finalBit)))
+			require.NoError(t, tx.CommitAt(uint64(iterations), nil))
+
+			db.SetDiscardTs(uint64(iterations))
+
+			require.NoError(t, FlushMemTable(db))
+			require.NoError(t, db.Flatten(1))
+
+			// 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.Equal(t, test.expectedKeyCount, numKeys)
+		})
+	}
+}
+
 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)))
+		WithLoggingLevel(badger.INFO).WithMemTableSize(int64(units.GB)))
 	require.NoError(t, err)
 	// set up the store
 	store, err := NewStoreWithDB(db, nil, lib.NewDefaultLogger(), true)
@@ -453,7 +535,7 @@ func TestFlushMemtable(t *testing.T) {
 		numKeys++
 	}
 
-	require.True(t, numKeys == 1)
+	require.Equal(t, 1, numKeys)
 }
 
 func TestHistoricalPrefix(t *testing.T) {

store/wrapper_txn.go

@@ -12,16 +12,29 @@ import (
 )
 
 const (
+	// ----------------------------------------------------------------------------------------------------------------
 	// BadgerDB garbage collector behavior is not well documented leading to many open issues in their repository
 	// However, here is our current understanding based on experimentation
-	// ----------------------------------------------------------------------------------------------------------
-	// BadgerDB Garbage Collection Rules listed by priority
-	// 1. MergeOp Bit prevents garbage collection of a key regardless
-	// 2. Nothing above DiscardTs will be garbage collected
-	// 3. Deleting, expiring, and setting ‘Discard earlier versions’ signals the removal of older versions of a key
-	// 4. Deleting or expiring a key makes it eligible for GC
-	// 5. If KeepNumVersions is set to max int it will prevent automatic GC of older versions
-	// ------------------------------------------------------------------------------
+	// ----------------------------------------------------------------------------------------------------------------
+	// 1. Manual Keep (Protection)
+	//    - `badgerNoDiscardBit` prevents automatic GC of a key version.
+	//    - However, it can be manually superseded by a manual removal
+	//
+	// 2. Manual Remove (Explicit Deletion or Pruning)
+	//    - Deleting a key at a higher ts removes earlier versions once `discardTs >= ts`.
+	//    - Setting `badgerDiscardEarlierVersions` is similar, except it retains the current version.
+	//
+	// 3. Auto Remove – Tombstones
+	//    - Deleted keys (tombstoned) <= `discardTs` are automatically purged unless protected by `badgerNoDiscardBit`
+	//
+	// 4. Auto Remove – Set Entries
+	//    - For non-deleted (live) keys, Badger retains the number of versions to retain is defined by `KeepNumVersions`.
+	//    - Older versions exceeding this count are automatically eligible for GC.
+	//
+	//   Note:
+	// - The first GC pass after updating `discardTs` and flushing memtable is deterministic
+	// - Subsequent GC runs are probabilistic, depending on reclaimable space and value log thresholds
+	// ----------------------------------------------------------------------------------------------------------------
 	// Bits source: https://github.com/hypermodeinc/badger/blob/85389e88bf308c1dc271383b77b67f4ef4a85194/value.go#L37
 	badgerMetaFieldName               = "meta"  // badgerDB Entry 'meta' field name
 	badgerDiscardEarlierVersions byte = 1 << 2  // badgerDB 'discard earlier versions' flag
@@ -85,17 +98,8 @@ func (t *TxnWrapper) Delete(k []byte) lib.ErrorI {
 	return nil
 }
 
-// Set() stores the key-value pair in the BadgerDB transaction
-func (t *TxnWrapper) SetNonPruneable(k, v []byte) lib.ErrorI {
-	// set an entry with a bit that prevents it from being discarded
-	if err := t.db.SetEntry(newEntry(lib.Append(t.prefix, k), v, badgerNoDiscardBit)); err != nil {
-		return ErrStoreSet(err)
-	}
-	return nil
-}
-
-// DeleteNonPrunable() removes the key-value pair from the BadgerDB transaction but prevents it from being garbage collected
-func (t *TxnWrapper) DeleteNonPrunable(k []byte) lib.ErrorI {
+// Tombstone() removes the key-value pair from the BadgerDB transaction but prevents it from being garbage collected
+func (t *TxnWrapper) Tombstone(k []byte) lib.ErrorI {
 	// set an entry with a bit that marks it as deleted and prevents it from being discarded
 	if err := t.db.SetEntry(newEntry(lib.Append(t.prefix, k), nil, badgerDeleteBit|badgerNoDiscardBit)); err != nil {
 		return ErrStoreDelete(err)

store/wrapper_txn_test.go

@@ -28,52 +28,6 @@ func TestGetSetDelete(t *testing.T) {
 	require.Contains(t, er.Error(), badger.ErrKeyNotFound.Error())
 }
 
-func TestSetDeleteNonPruneable(t *testing.T) {
-	db, store, cleanup := newTestTxnWrapper(t)
-	defer cleanup()
-	// first set 'a' and 'b' at height 1
-	require.NoError(t, store.SetNonPruneable([]byte("a"), []byte("a")))
-	require.NoError(t, store.SetNonPruneable([]byte("b"), []byte("b")))
-	require.NoError(t, store.db.CommitAt(1, nil))
-	// delete 'b' at height 2
-	tx := db.NewTransactionAt(1, true)
-	nextStore := NewTxnWrapper(tx, lib.NewDefaultLogger(), []byte(latestStatePrefix))
-	require.NoError(t, nextStore.DeleteNonPrunable([]byte("b")))
-	require.NoError(t, nextStore.db.CommitAt(2, nil))
-	// read at height 2
-	tx2 := db.NewTransactionAt(2, true)
-	nextStore2 := NewTxnWrapper(tx2, lib.NewDefaultLogger(), []byte(latestStatePrefix))
-	defer nextStore2.Close()
-	// use an archive iterator
-	iterator, err := nextStore2.ArchiveIterator(nil)
-	require.NoError(t, err)
-	it := iterator.(*Iterator)
-	defer it.Close()
-	var count int
-	// check the expected values of each item of the archive iterator and the order
-	// including the pruning and delete bits
-	for ; it.Valid(); it.Next() {
-		require.True(t, getMeta(it.parent.Item())&badgerNoDiscardBit != 0)
-		switch count {
-		case 0:
-			require.Equal(t, it.Key(), []byte("a"))
-			require.Equal(t, it.Value(), []byte("a"))
-			require.False(t, getMeta(it.parent.Item())&badgerDeleteBit != 0)
-		case 1:
-			require.Equal(t, it.Key(), []byte("b"))
-			require.Equal(t, it.Value(), []byte(nil))
-			require.True(t, getMeta(it.parent.Item())&badgerDeleteBit != 0)
-		case 2:
-			require.Equal(t, it.Key(), []byte("b"))
-			require.Equal(t, it.Value(), []byte("b"))
-			require.False(t, getMeta(it.parent.Item())&badgerDeleteBit != 0)
-		}
-		count++
-	}
-	// ensure 3 keys
-	require.Equal(t, count, 3)
-}
-
 func TestIteratorBasic(t *testing.T) {
 	_, parent, cleanup := newTestTxnWrapper(t)
 	defer cleanup()