fix(internal/db): fix reverse prefix iteration

Amolith and Crush created 4 days ago

Co-authored-by: Crush <[e-mail]>

Change summary

internal/db/db.go      |  10 +++
internal/db/db_test.go | 102 ++++++++++++++++++++++++++++++++++++++++++++
2 files changed, 111 insertions(+), 1 deletion(-)

Detailed changes

internal/db/db.go

@@ -270,7 +270,15 @@ func (t *Txn) Iterate(opts IterateOptions, fn func(Item) error) error {
 	defer it.Close()
 
 	if len(opts.Prefix) > 0 {
-		for it.Seek(opts.Prefix); it.ValidForPrefix(opts.Prefix); it.Next() {
+		if opts.Reverse {
+			// Seek to the end of the prefix range by appending 0xFF
+			seekKey := append(opts.Prefix, 0xFF)
+			it.Seek(seekKey)
+		} else {
+			it.Seek(opts.Prefix)
+		}
+
+		for ; it.ValidForPrefix(opts.Prefix); it.Next() {
 			if err := fn(Item{item: it.Item()}); err != nil {
 				if errors.Is(err, ErrTxnAborted) {
 					return nil

internal/db/db_test.go

@@ -589,3 +589,105 @@ func TestEncodingHelpers(t *testing.T) {
 		t.Fatalf("unexpected encodeHex: %q", hex)
 	}
 }
+
+func TestIterateReversePrefixOrder(t *testing.T) {
+	database := openTestDB(t)
+
+	keys := [][]byte{
+		[]byte("dir/abc/archived/0001/sid1"),
+		[]byte("dir/abc/archived/0002/sid2"),
+		[]byte("dir/abc/archived/0003/sid3"),
+		[]byte("dir/xyz/archived/0001/sid4"),
+		[]byte("other/key"),
+	}
+
+	err := database.Update(context.Background(), func(txn *Txn) error {
+		for _, key := range keys {
+			if err := txn.Set(key, key); err != nil {
+				return err
+			}
+		}
+		return nil
+	})
+	if err != nil {
+		t.Fatalf("populate keys: %v", err)
+	}
+
+	t.Run("ReverseWithPrefix", func(t *testing.T) {
+		var seen []string
+		err := database.View(context.Background(), func(txn *Txn) error {
+			return txn.Iterate(IterateOptions{
+				Prefix:  []byte("dir/abc/archived"),
+				Reverse: true,
+			}, func(item Item) error {
+				seen = append(seen, item.KeyString())
+				return nil
+			})
+		})
+		if err != nil {
+			t.Fatalf("reverse prefix iteration: %v", err)
+		}
+
+		expected := []string{
+			"dir/abc/archived/0003/sid3",
+			"dir/abc/archived/0002/sid2",
+			"dir/abc/archived/0001/sid1",
+		}
+		if len(seen) != len(expected) {
+			t.Fatalf("unexpected number of items: got %d want %d", len(seen), len(expected))
+		}
+		for i, want := range expected {
+			if seen[i] != want {
+				t.Fatalf("unexpected key at %d: got %q want %q", i, seen[i], want)
+			}
+		}
+	})
+
+	t.Run("ForwardWithPrefix", func(t *testing.T) {
+		var seen []string
+		err := database.View(context.Background(), func(txn *Txn) error {
+			return txn.Iterate(IterateOptions{
+				Prefix: []byte("dir/abc/archived"),
+			}, func(item Item) error {
+				seen = append(seen, item.KeyString())
+				return nil
+			})
+		})
+		if err != nil {
+			t.Fatalf("forward prefix iteration: %v", err)
+		}
+
+		expected := []string{
+			"dir/abc/archived/0001/sid1",
+			"dir/abc/archived/0002/sid2",
+			"dir/abc/archived/0003/sid3",
+		}
+		if len(seen) != len(expected) {
+			t.Fatalf("unexpected number of items: got %d want %d", len(seen), len(expected))
+		}
+		for i, want := range expected {
+			if seen[i] != want {
+				t.Fatalf("unexpected key at %d: got %q want %q", i, seen[i], want)
+			}
+		}
+	})
+
+	t.Run("EmptyPrefixReverse", func(t *testing.T) {
+		var count int
+		err := database.View(context.Background(), func(txn *Txn) error {
+			return txn.Iterate(IterateOptions{
+				Prefix:  []byte("nonexistent/prefix"),
+				Reverse: true,
+			}, func(item Item) error {
+				count++
+				return nil
+			})
+		})
+		if err != nil {
+			t.Fatalf("empty prefix iteration: %v", err)
+		}
+		if count != 0 {
+			t.Fatalf("expected 0 items for nonexistent prefix, got %d", count)
+		}
+	})
+}