Repository: genjidb/genji
Pull request #368 (2021-03-11 07:17:21)

Fix index state with UPDATE queries
This fixes how UPDATE queries on indexed tables would cause some indexes to be in an incorrect state.
This was caused by the "trash bin" mechanism implemented in both bolt and memory engines, where nodes would get marked as deleted without modifying the btree and then actually removed when the transaction is committed.
The error occurred when a key was marked for deletion and then overwritten (which happens a lot when we run UPDATE on indexed fields that are not modified), the cleanup mechanism would delete them even when it's not supposed to.
This PR also fixes a race condition occurring in the memory engine.

Fixes https://github.com/genjidb/genji/issues/355


Patch:
diff --git a/database/table.go b/database/table.go
index ff0afb69..1d57e91a 100644
--- a/database/table.go
+++ b/database/table.go
@@ -181,7 +181,7 @@ func (t *Table) replace(indexes []*Index, key []byte, d document.Document) error
 	for _, idx := range indexes {
 		v, err := idx.Info.Path.GetValueFromDocument(old)
 		if err != nil {
-			return err
+			v = document.NewNullValue()
 		}
 
 		err = idx.Delete(v, key)
@@ -209,16 +209,20 @@ func (t *Table) replace(indexes []*Index, key []byte, d document.Document) error
 	for _, idx := range indexes {
 		v, err := idx.Info.Path.GetValueFromDocument(d)
 		if err != nil {
-			continue
+			v = document.NewNullValue()
 		}
 
 		err = idx.Set(v, key)
 		if err != nil {
+			if err == ErrIndexDuplicateValue {
+				return ErrDuplicateDocument
+			}
+
 			return err
 		}
 	}
 
-	return err
+	return nil
 }
 
 type documentWithKey struct {
diff --git a/database/transaction.go b/database/transaction.go
index a1d95045..1d719cfb 100644
--- a/database/transaction.go
+++ b/database/transaction.go
@@ -67,7 +67,6 @@ func (tx *Transaction) Rollback() error {
 // Commit the transaction. Calling this method on read-only transactions
 // will return an error.
 func (tx *Transaction) Commit() error {
-
 	err := tx.tx.Commit()
 	if err != nil {
 		return err
diff --git a/engine/boltengine/engine.go b/engine/boltengine/engine.go
index 1ec0f29e..938cd048 100644
--- a/engine/boltengine/engine.go
+++ b/engine/boltengine/engine.go
@@ -214,9 +214,13 @@ func (t *Transaction) cleanupBinBucket() error {
 			buckets[string(bucketName)] = b
 		}
 
-		err := b.Delete(key)
-		if err != nil {
-			return err
+		// if the key has been rewritten during the lifecycle of the transaction
+		// do not delete it
+		if b.Get(key) == nil {
+			err := b.Delete(key)
+			if err != nil {
+				return err
+			}
 		}
 	}
 
diff --git a/engine/memoryengine/engine.go b/engine/memoryengine/engine.go
index 87100029..035ffaec 100644
--- a/engine/memoryengine/engine.go
+++ b/engine/memoryengine/engine.go
@@ -16,17 +16,17 @@ import (
 const btreeDegree = 12
 
 // Engine is a simple memory engine implementation that stores data in
-// an in-memory Btree. It allows multiple readers and one single writer.
+// an in-memory Btree. It is not thread safe.
 type Engine struct {
 	closed    bool
-	stores    map[string]*btree.BTree
+	stores    map[string]*tree
 	sequences map[string]uint64
 }
 
 // NewEngine creates an in-memory engine.
 func NewEngine() *Engine {
 	return &Engine{
-		stores:    make(map[string]*btree.BTree),
+		stores:    make(map[string]*tree),
 		sequences: make(map[string]uint64),
 	}
 }
@@ -158,7 +158,7 @@ func (tx *transaction) CreateStore(name []byte) error {
 
 	tr := btree.New(btreeDegree)
 
-	tx.ng.stores[string(name)] = tr
+	tx.ng.stores[string(name)] = &tree{bt: tr}
 
 	// on rollback, remove the btree from the list of stores
 	tx.onRollback = append(tx.onRollback, func() {
@@ -193,3 +193,81 @@ func (tx *transaction) DropStore(name []byte) error {
 
 	return nil
 }
+
+// tree is a thread safe wrapper aroung BTree.
+// It prevents modifying and rebalancing the btree while other
+// routines are reading it.
+type tree struct {
+	bt *btree.BTree
+
+	m sync.RWMutex
+}
+
+func (t *tree) Get(key btree.Item) btree.Item {
+	t.m.RLock()
+	defer t.m.RUnlock()
+
+	return t.bt.Get(key)
+}
+
+func (t *tree) Delete(key btree.Item) btree.Item {
+	t.m.Lock()
+	defer t.m.Unlock()
+
+	return t.bt.Delete(key)
+}
+
+func (t *tree) ReplaceOrInsert(key btree.Item) btree.Item {
+	t.m.Lock()
+	defer t.m.Unlock()
+
+	return t.bt.ReplaceOrInsert(key)
+}
+
+func (t *tree) Ascend(iterator btree.ItemIterator) {
+	t.m.RLock()
+	defer t.m.RUnlock()
+
+	t.bt.Ascend(func(i btree.Item) bool {
+		t.m.RUnlock()
+		defer t.m.RLock()
+
+		return iterator(i)
+	})
+}
+
+func (t *tree) AscendGreaterOrEqual(pivot btree.Item, iterator btree.ItemIterator) {
+	t.m.RLock()
+	defer t.m.RUnlock()
+
+	t.bt.AscendGreaterOrEqual(pivot, func(i btree.Item) bool {
+		t.m.RUnlock()
+		defer t.m.RLock()
+
+		return iterator(i)
+	})
+}
+
+func (t *tree) Descend(iterator btree.ItemIterator) {
+	t.m.RLock()
+	defer t.m.RUnlock()
+
+	t.bt.Descend(func(i btree.Item) bool {
+		t.m.RUnlock()
+		defer t.m.RLock()
+
+		return iterator(i)
+	})
+}
+
+func (t *tree) DescendLessOrEqual(pivot btree.Item, iterator btree.ItemIterator) {
+	t.m.RLock()
+	defer t.m.RUnlock()
+
+	t.bt.DescendLessOrEqual(pivot, func(i btree.Item) bool {
+		t.m.RUnlock()
+		defer t.m.RLock()
+
+		return iterator(i)
+	})
+}
diff --git a/engine/memoryengine/store.go b/engine/memoryengine/store.go
index 398ab360..803618c5 100644
--- a/engine/memoryengine/store.go
+++ b/engine/memoryengine/store.go
@@ -37,7 +37,7 @@ func (i *item) Less(than btree.Item) bool {
 
 // storeTx implements an engine.Store.
 type storeTx struct {
-	tr   *btree.BTree
+	tr   *tree
 	tx   *transaction
 	name string
 }
@@ -158,7 +158,9 @@ func (s *storeTx) Delete(k []byte) error {
 
 	// on commit, remove the item from the tree.
 	s.tx.onCommit = append(s.tx.onCommit, func() {
-		s.tr.Delete(i)
+		if i.deleted {
+			s.tr.Delete(i)
+		}
 	})
 	return nil
 }
@@ -178,7 +180,7 @@ func (s *storeTx) Truncate() error {
 	}
 
 	old := s.tr
-	s.tr = btree.New(btreeDegree)
+	s.tr = &tree{bt: btree.New(btreeDegree)}
 
 	// on rollback replace the new tree by the old one.
 	s.tx.onRollback = append(s.tx.onRollback, func() {
@@ -220,7 +222,7 @@ func (s *storeTx) Iterator(opts engine.IteratorOptions) engine.Iterator {
 type iterator struct {
 	tx      *transaction
 	reverse bool
-	tr      *btree.BTree
+	tr      *tree
 	item    *item // current item
 	ch      chan *item
 	closed  chan struct{} // closed by the goroutine when it's shutdown
@@ -234,11 +236,11 @@ func (it *iterator) Seek(pivot []byte) {
 	// is closed before creating a new one
 	if it.cancel != nil {
 		it.cancel()
-		it.ch = make(chan *item)
 		<-it.closed
-		it.closed = make(chan struct{})
 	}
 
+	it.ch = make(chan *item)
+	it.closed = make(chan struct{})
 	it.ctx, it.cancel = context.WithCancel(it.tx.ctx)
 
 	it.runIterator(pivot)
@@ -252,7 +254,7 @@ func (it *iterator) Seek(pivot []byte) {
 func (it *iterator) runIterator(pivot []byte) {
 	it.tx.wg.Add(1)
 
-	go func(ctx context.Context, ch chan *item, tr *btree.BTree) {
+	go func(ctx context.Context, ch chan *item, tr *tree) {
 		defer it.tx.wg.Done()
 		defer close(ch)
 		defer close(it.closed)


Test patch:
diff --git a/database/table_test.go b/database/table_test.go
index d8203e49..28c84423 100644
--- a/database/table_test.go
+++ b/database/table_test.go
@@ -13,6 +13,7 @@ import (
 	"github.com/genjidb/genji/document/encoding/msgpack"
 	"github.com/genjidb/genji/engine/memoryengine"
 	"github.com/genjidb/genji/sql/parser"
+	"github.com/genjidb/genji/testutil"
 	"github.com/stretchr/testify/require"
 )
 
@@ -633,6 +634,52 @@ func TestTableReplace(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, "c", f.V.(string))
 	})
+
+	t.Run("Should update indexes", func(t *testing.T) {
+		_, tx, cleanup := newTestTx(t)
+		defer cleanup()
+
+		err := tx.CreateTable("test", nil)
+		require.NoError(t, err)
+
+		err = tx.CreateIndex(database.IndexInfo{
+			Path:      document.NewPath("a"),
+			Unique:    true,
+			TableName: "test",
+			IndexName: "idx_foo_a",
+		})
+		require.NoError(t, err)
+
+		tb, err := tx.GetTable("test")
+		require.NoError(t, err)
+
+		// insert two different documents
+		d1, err := tb.Insert(testutil.MakeDocument(t, `{"a": 1, "b": 1}`))
+		require.NoError(t, err)
+		d2, err := tb.Insert(testutil.MakeDocument(t, `{"a": 2, "b": 2}`))
+		require.NoError(t, err)
+
+		before := testutil.GetIndexContent(t, tx, "idx_foo_a")
+
+		// replace doc 1 without modifying indexed key
+		err = tb.Replace(d1.(document.Keyer).RawKey(), testutil.MakeDocument(t, `{"a": 1, "b": 3}`))
+		require.NoError(t, err)
+		// index should be the same as before
+		require.Equal(t, before, testutil.GetIndexContent(t, tx, "idx_foo_a"))
+
+		// replace doc 2 and modify indexed key
+		err = tb.Replace(d2.(document.Keyer).RawKey(), testutil.MakeDocument(t, `{"a": 3, "b": 3}`))
+		require.NoError(t, err)
+		// index should be different for doc 2
+		got := testutil.GetIndexContent(t, tx, "idx_foo_a")
+		require.Equal(t, before[0], got[0])
+		require.NotEqual(t, before[1], got[1])
+
+		// replace doc 1 with duplicate indexed key
+		err = tb.Replace(d1.(document.Keyer).RawKey(), testutil.MakeDocument(t, `{"a": 3, "b": 3}`))
+		// index should be the same as before
+		require.Equal(t, database.ErrDuplicateDocument, err)
+	})
 }
 
 // TestTableTruncate verifies Truncate behaviour.
diff --git a/engine/enginetest/testing.go b/engine/enginetest/testing.go
index 9e2229b9..00b4f48e 100644
--- a/engine/enginetest/testing.go
+++ b/engine/enginetest/testing.go
@@ -857,6 +857,28 @@ func TestStoreIterator(t *testing.T, builder Builder) {
 		require.True(t, it.Valid())
 		require.Equal(t, it.Item().Key(), k)
 	})
+
+	t.Run("Iterating while deleting current key should work", func(t *testing.T) {
+		st, cleanup := storeBuilder(t, builder)
+		defer cleanup()
+
+		for i := 0; i < 50; i++ {
+			err := st.Put([]byte{byte(i)}, []byte{byte(i)})
+			require.NoError(t, err)
+		}
+
+		i := 0
+		it := st.Iterator(engine.IteratorOptions{})
+		defer it.Close()
+
+		for it.Seek(nil); it.Valid() && i < 50; it.Next() {
+			require.Equal(t, []byte{byte(i)}, it.Item().Key())
+
+			err := st.Delete([]byte{byte(i)})
+			require.NoError(t, err)
+			i++
+		}
+	})
 }
 
 // TestStorePut verifies Put behaviour.
@@ -1015,6 +1037,48 @@ func TestStoreDelete(t *testing.T, builder Builder) {
 		require.Equal(t, 1, i)
 	})
 
+	t.Run("Should not rollback document if deleted then put", func(t *testing.T) {
+		ng, cleanup := builder()
+		defer cleanup()
+
+		tx, err := ng.Begin(context.Background(), engine.TxOptions{Writable: true})
+		require.NoError(t, err)
+		err = tx.CreateStore([]byte("test"))
+		require.NoError(t, err)
+		st, err := tx.GetStore([]byte("test"))
+		require.NoError(t, err)
+
+		err = st.Put([]byte("foo"), []byte("FOO"))
+		require.NoError(t, err)
+
+		// delete the key
+		err = st.Delete([]byte("foo"))
+		require.NoError(t, err)
+
+		v, err := st.Get([]byte("foo"))
+		require.Equal(t, engine.ErrKeyNotFound, err)
+
+		err = st.Put([]byte("foo"), []byte("bar"))
+		require.NoError(t, err)
+
+		v, err = st.Get([]byte("foo"))
+		require.Equal(t, []byte("bar"), v)
+
+		// commit and reopen a transaction
+		err = tx.Commit()
+		require.NoError(t, err)
+
+		tx, err = ng.Begin(context.Background(), engine.TxOptions{Writable: false})
+		require.NoError(t, err)
+		defer tx.Rollback()
+
+		st, err = tx.GetStore([]byte("test"))
+		require.NoError(t, err)
+
+		v, err = st.Get([]byte("foo"))
+		require.Equal(t, []byte("bar"), v)
+	})
+
 	t.Run("Should fail if context canceled", func(t *testing.T) {
 		ctx, cancel := context.WithCancel(context.Background())
 		defer cancel()
diff --git a/query/update_test.go b/query/update_test.go
index 3554497e..21188596 100644
--- a/query/update_test.go
+++ b/query/update_test.go
@@ -44,35 +44,46 @@ func TestUpdateStmt(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			db, err := genji.Open(":memory:")
-			require.NoError(t, err)
-			defer db.Close()
-
-			err = db.Exec("CREATE TABLE test (a text not null)")
-			require.NoError(t, err)
-			err = db.Exec("INSERT INTO test (a, b, c) VALUES ('foo1', 'bar1', 'baz1')")
-			require.NoError(t, err)
-			err = db.Exec("INSERT INTO test (a, b) VALUES ('foo2', 'bar2')")
-			require.NoError(t, err)
-			err = db.Exec("INSERT INTO test (a, d, e) VALUES ('foo3', 'bar3', 'baz3')")
-			require.NoError(t, err)
-
-			err = db.Exec(test.query, test.params...)
-			if test.fails {
-				require.Error(t, err)
-				return
-			}
-			require.NoError(t, err)
+			runTest := func(indexed bool) {
+				db, err := genji.Open(":memory:")
+				require.NoError(t, err)
+				defer db.Close()
 
-			st, err := db.Query("SELECT * FROM test")
-			require.NoError(t, err)
-			defer st.Close()
+				err = db.Exec("CREATE TABLE test (a text not null)")
+				require.NoError(t, err)
 
-			var buf bytes.Buffer
+				if indexed {
+					err = db.Exec("CREATE INDEX idx_test_a ON test(a)")
+					require.NoError(t, err)
+				}
 
-			err = document.IteratorToJSONArray(&buf, st)
-			require.NoError(t, err)
-			require.JSONEq(t, test.expected, buf.String())
+				err = db.Exec("INSERT INTO test (a, b, c) VALUES ('foo1', 'bar1', 'baz1')")
+				require.NoError(t, err)
+				err = db.Exec("INSERT INTO test (a, b) VALUES ('foo2', 'bar2')")
+				require.NoError(t, err)
+				err = db.Exec("INSERT INTO test (a, d, e) VALUES ('foo3', 'bar3', 'baz3')")
+				require.NoError(t, err)
+
+				err = db.Exec(test.query, test.params...)
+				if test.fails {
+					require.Error(t, err)
+					return
+				}
+				require.NoError(t, err)
+
+				st, err := db.Query("SELECT * FROM test")
+				require.NoError(t, err)
+				defer st.Close()
+
+				var buf bytes.Buffer
+
+				err = document.IteratorToJSONArray(&buf, st)
+				require.NoError(t, err)
+				require.JSONEq(t, test.expected, buf.String())
+			}
+
+			runTest(false)
+			runTest(true)
 		})
 	}
 
@@ -123,7 +134,6 @@ func TestUpdateStmt(t *testing.T) {
 				require.NoError(t, err)
 				require.JSONEq(t, tt.expected, buf.String())
 			})
-
 		}
 	})
 }
diff --git a/testutil/document.go b/testutil/document.go
index 54d81668..8c88c544 100644
--- a/testutil/document.go
+++ b/testutil/document.go
@@ -1,7 +1,9 @@
 package testutil
 
 import (
+	"encoding/json"
 	"fmt"
+	"os"
 	"testing"
 
 	"github.com/genjidb/genji/document"
@@ -54,3 +56,13 @@ func (docs Docs) RequireEqual(t testing.TB, others Docs) {
 		}
 	}
 }
+
+// Dump a json representation of v to os.Stdout.
+func Dump(t testing.TB, v interface{}) {
+	t.Helper()
+
+	enc := json.NewEncoder(os.Stdout)
+	enc.SetIndent("", "  ")
+	err := enc.Encode(v)
+	require.NoError(t, err)
+}
diff --git a/testutil/index.go b/testutil/index.go
new file mode 100644
index 00000000..6bd251e9
--- /dev/null
+++ b/testutil/index.go
@@ -0,0 +1,29 @@
+package testutil
+
+import (
+	"testing"
+
+	"github.com/genjidb/genji/database"
+	"github.com/genjidb/genji/document"
+	"github.com/stretchr/testify/require"
+)
+
+// GetIndexContent iterates over the entire index and returns all the key-value pairs in order.
+func GetIndexContent(t testing.TB, tx *database.Transaction, indexName string) []KV {
+	t.Helper()
+
+	idx, err := tx.GetIndex(indexName)
+	require.NoError(t, err)
+
+	var content []KV
+	err = idx.AscendGreaterOrEqual(document.Value{}, func(val, key []byte) error {
+		content = append(content, KV{
+			Key:   append([]byte{}, val...),
+			Value: append([]byte{}, key...),
+		})
+		return nil
+	})
+	require.NoError(t, err)
+
+	return content
+}
diff --git a/testutil/misc.go b/testutil/misc.go
new file mode 100644
index 00000000..41cc0529
--- /dev/null
+++ b/testutil/misc.go
@@ -0,0 +1,6 @@
+package testutil
+
+// KV is used to represent key-value pairs stored in tables or indexes.
+type KV struct {
+	Key, Value []byte
+}