Repository: genjidb/genji
Pull request #336 (2021-01-04 14:28:25)

Soft delete store keys before commit
This PR makes the engine store soft delete keys and cleanup on commit. This is done to prevent tree rebalancing when iterating and deleting at the same time, which was a problem with Bolt.
It also prevents inserting empty values in stores.


Patch:
diff --git a/engine/badgerengine/store.go b/engine/badgerengine/store.go
index bc29a3bb..72344e3b 100644
--- a/engine/badgerengine/store.go
+++ b/engine/badgerengine/store.go
@@ -50,6 +50,10 @@ func (s *Store) Put(k, v []byte) error {
 		return errors.New("cannot store empty key")
 	}
 
+	if len(v) == 0 {
+		return errors.New("cannot store empty value")
+	}
+
 	return s.tx.Set(buildKey(s.prefix, k), v)
 }
 
diff --git a/engine/boltengine/engine.go b/engine/boltengine/engine.go
index 606168ae..f344b1d4 100644
--- a/engine/boltengine/engine.go
+++ b/engine/boltengine/engine.go
@@ -3,6 +3,7 @@ package boltengine
 
 import (
 	"context"
+	"encoding/binary"
 	"os"
 
 	"github.com/genjidb/genji/engine"
@@ -11,6 +12,8 @@ import (
 
 const (
 	separator byte = 0x1F
+	// name of the bucket used to mark keys for deletion
+	binBucket = "__bin"
 )
 
 // Engine represents a BoltDB engine. Each store is stored in a dedicated bucket.
@@ -60,6 +63,9 @@ type Transaction struct {
 	ctx      context.Context
 	tx       *bolt.Tx
 	writable bool
+	// if set to true,
+	// the __bin bucket will be cleanup on commit.
+	cleanupBin bool
 }
 
 // Rollback the transaction. Can be used safely after commit.
@@ -87,6 +93,14 @@ func (t *Transaction) Commit() error {
 	default:
 	}
 
+	// remove keys marked for deletion
+	if t.cleanupBin {
+		err := t.cleanupBinBucket()
+		if err != nil {
+			return err
+		}
+	}
+
 	return t.tx.Commit()
 }
 
@@ -106,6 +120,7 @@ func (t *Transaction) GetStore(name []byte) (engine.Store, error) {
 	return &Store{
 		bucket: b,
 		tx:     t.tx,
+		ngTx:   t,
 		name:   name,
 		ctx:    t.ctx,
 	}, nil
@@ -151,3 +166,53 @@ func (t *Transaction) DropStore(name []byte) error {
 
 	return err
 }
+
+func (t *Transaction) markForDeletion(bucketName, key []byte) error {
+	// create the bin bucket
+	bb, err := t.tx.CreateBucketIfNotExists([]byte(binBucket))
+	if err != nil {
+		return err
+	}
+
+	// store the key in the bin bucket.
+	// store the offset of the key in the value.
+	var buf [binary.MaxVarintLen64]byte
+	n := binary.PutUvarint(buf[:], uint64(len(bucketName)))
+	err = bb.Put(append(bucketName, key...), buf[:n])
+	if err != nil {
+		return err
+	}
+
+	// tell the transaction to cleanup on commit
+	t.cleanupBin = true
+	return nil
+}
+
+func (t *Transaction) cleanupBinBucket() error {
+	buckets := make(map[string]*bolt.Bucket)
+
+	c := t.tx.Bucket([]byte(binBucket)).Cursor()
+	for k, v := c.Seek(nil); k != nil; k, v = c.Next() {
+		offset, _ := binary.Uvarint(v)
+		bucketName, key := k[:int(offset)], k[int(offset):]
+		b, ok := buckets[string(bucketName)]
+		if !ok {
+			b = t.tx.Bucket(bucketName)
+			// if b is nil, the bucket must have been deleted during this transaction
+			// after having deleted some of its keys, we can ignore it.
+			if b == nil {
+				continue
+			}
+
+			buckets[string(bucketName)] = b
+		}
+
+		err := b.Delete(key)
+		if err != nil {
+			return err
+		}
+	}
+
+	// we can now drop the bin bucket
+	return t.tx.DeleteBucket([]byte(binBucket))
+}
diff --git a/engine/boltengine/store.go b/engine/boltengine/store.go
index 653cfbf6..665391fb 100644
--- a/engine/boltengine/store.go
+++ b/engine/boltengine/store.go
@@ -3,6 +3,7 @@ package boltengine
 import (
 	"bytes"
 	"context"
+	"errors"
 
 	"github.com/genjidb/genji/engine"
 	bolt "go.etcd.io/bbolt"
@@ -12,6 +13,7 @@ import (
 type Store struct {
 	bucket *bolt.Bucket
 	tx     *bolt.Tx
+	ngTx   *Transaction
 	name   []byte
 	ctx    context.Context
 }
@@ -28,6 +30,10 @@ func (s *Store) Put(k, v []byte) error {
 		return engine.ErrTransactionReadOnly
 	}
 
+	if len(k) == 0 || len(v) == 0 {
+		return errors.New("empty key or value")
+	}
+
 	return s.bucket.Put(k, v)
 }
 
@@ -48,6 +54,9 @@ func (s *Store) Get(k []byte) ([]byte, error) {
 }
 
 // Delete a record by key. If not found, returns table.ErrDocumentNotFound.
+// It hides the key without deleting the actual node from the tree, to tree rebalancing during iterations.
+// It then adds it to a sub bucket containing the list of keys to delete when the transaction
+// is committed.
 func (s *Store) Delete(k []byte) error {
 	select {
 	case <-s.ctx.Done():
@@ -64,7 +73,15 @@ func (s *Store) Delete(k []byte) error {
 		return engine.ErrKeyNotFound
 	}
 
-	return s.bucket.Delete(k)
+	// setting the value to nil hides the key
+	// without deleting the actual node from the tree.
+	err := s.bucket.Put(k, nil)
+	if err != nil {
+		return err
+	}
+
+	// mark the key for deletion on commit
+	return s.ngTx.markForDeletion(s.name, k)
 }
 
 // Truncate deletes all the records of the store.
@@ -130,17 +147,23 @@ func (it *iterator) Seek(pivot []byte) {
 
 	if !it.reverse {
 		it.item.k, it.item.v = it.c.Seek(pivot)
+		if it.item.v == nil {
+			it.getKey(it.c.Next)
+		}
 		return
 	}
 
 	if len(pivot) == 0 {
 		it.item.k, it.item.v = it.c.Last()
+		if it.item.v == nil {
+			it.getKey(it.c.Prev)
+		}
 		return
 	}
 
 	it.item.k, it.item.v = it.c.Seek(pivot)
 	if it.item.k != nil {
-		for bytes.Compare(it.item.k, pivot) > 0 {
+		for bytes.Compare(it.item.k, pivot) > 0 || (len(it.item.k) > 0 && len(it.item.v) == 0) {
 			it.item.k, it.item.v = it.c.Prev()
 		}
 	}
@@ -152,9 +175,21 @@ func (it *iterator) Valid() bool {
 
 func (it *iterator) Next() {
 	if it.reverse {
-		it.item.k, it.item.v = it.c.Prev()
+		it.getKey(it.c.Prev)
 	} else {
-		it.item.k, it.item.v = it.c.Next()
+		it.getKey(it.c.Next)
+	}
+}
+
+func (it *iterator) getKey(fn func() (key []byte, value []byte)) {
+	// skip items with nil value.
+	// these items have been soft-deleted and will be removed on commit.
+	for {
+		it.item.k, it.item.v = fn()
+
+		if it.item.k == nil || len(it.item.v) != 0 {
+			break
+		}
 	}
 }
 
diff --git a/engine/engine.go b/engine/engine.go
index 5531015c..cb9d5be3 100644
--- a/engine/engine.go
+++ b/engine/engine.go
@@ -66,6 +66,7 @@ type Store interface {
 	// Get returns a value associated with the given key. If no key is not found, it returns ErrKeyNotFound.
 	Get(k []byte) ([]byte, error)
 	// Put stores a key value pair. If it already exists, it overrides it.
+	// Both k and v must be not nil.
 	Put(k, v []byte) error
 	// Delete a key value pair. If the key is not found, returns ErrKeyNotFound.
 	Delete(k []byte) error
diff --git a/engine/memoryengine/store.go b/engine/memoryengine/store.go
index d47957f7..398ab360 100644
--- a/engine/memoryengine/store.go
+++ b/engine/memoryengine/store.go
@@ -57,6 +57,10 @@ func (s *storeTx) Put(k, v []byte) error {
 		return errors.New("empty keys are forbidden")
 	}
 
+	if len(v) == 0 {
+		return errors.New("empty values are forbidden")
+	}
+
 	it := &item{k: k}
 	// if there is an existing value, fetch it
 	// and overwrite it directly using the pointer.
@@ -290,6 +294,15 @@ func (it *iterator) runIterator(pivot []byte) {
 }
 
 func (it *iterator) Valid() bool {
+	if it.err != nil {
+		return false
+	}
+	select {
+	case <-it.tx.ctx.Done():
+		it.err = it.tx.ctx.Err()
+	default:
+	}
+
 	return it.item != nil && it.err == nil
 }
 


Test patch:
diff --git a/engine/enginetest/testing.go b/engine/enginetest/testing.go
index c21c38a9..3de20e6a 100644
--- a/engine/enginetest/testing.go
+++ b/engine/enginetest/testing.go
@@ -899,15 +899,15 @@ func TestStorePut(t *testing.T, builder Builder) {
 		require.Error(t, err)
 	})
 
-	t.Run("Should succeed when value is nil or empty", func(t *testing.T) {
+	t.Run("Should fail when value is nil or empty", func(t *testing.T) {
 		st, cleanup := storeBuilder(t, builder)
 		defer cleanup()
 
 		err := st.Put([]byte("foo"), nil)
-		require.NoError(t, err)
+		require.Error(t, err)
 
 		err = st.Put([]byte("foo"), []byte(""))
-		require.NoError(t, err)
+		require.Error(t, err)
 	})
 
 	t.Run("Should fail if context canceled", func(t *testing.T) {
@@ -992,17 +992,27 @@ func TestStoreDelete(t *testing.T, builder Builder) {
 		require.Equal(t, []byte("FOO"), v)
 
 		// delete the key
-		err = st.Delete([]byte("foo"))
+		err = st.Delete([]byte("bar"))
 		require.NoError(t, err)
 
 		// try again, should fail
-		err = st.Delete([]byte("foo"))
+		err = st.Delete([]byte("bar"))
 		require.Equal(t, engine.ErrKeyNotFound, err)
 
 		// make sure it didn't also delete the other one
-		v, err = st.Get([]byte("bar"))
+		v, err = st.Get([]byte("foo"))
 		require.NoError(t, err)
-		require.Equal(t, []byte("BAR"), v)
+		require.Equal(t, []byte("FOO"), v)
+
+		// the deleted key must not appear on iteration
+		it := st.Iterator(engine.IteratorOptions{})
+		defer it.Close()
+		i := 0
+		for it.Seek(nil); it.Valid(); it.Next() {
+			require.Equal(t, []byte("foo"), it.Item().Key())
+			i++
+		}
+		require.Equal(t, 1, i)
 	})
 
 	t.Run("Should fail if context canceled", func(t *testing.T) {