iterator_test.go

iterator_test.go - Overview

This file contains unit tests for the iterator functionality in BadgerDB, focusing on prefix-based iteration, timestamp-based iteration, and read-only mode handling.

Detailed Documentation

`tableMock` struct

type tableMock struct {
	left, right []byte
}

Purpose: This struct is a mock implementation of the table.Table interface, used for testing the pickTable function. It stores the smallest and largest keys of a table.
Fields:
- left: The smallest key in the table.
- right: The largest key in the table.

`(*tableMock) Smallest() []byte`

func (tm *tableMock) Smallest() []byte { return tm.left }

Purpose: Returns the smallest key in the mock table.
Returns: The left field (smallest key) as a []byte.

`(*tableMock) Biggest() []byte`

func (tm *tableMock) Biggest() []byte { return tm.right }

Purpose: Returns the largest key in the mock table.
Returns: The right field (largest key) as a []byte.

`(*tableMock) DoesNotHave(hash uint32) bool`

func (tm *tableMock) DoesNotHave(hash uint32) bool { return false }

Purpose: A mock implementation that always returns false.
Parameters:
- hash: uint32 hash.
Returns: Always returns false.

`(*tableMock) MaxVersion() uint64`

func (tm *tableMock) MaxVersion() uint64 { return math.MaxUint64 }

Purpose: Returns the maximum possible unsigned 64-bit integer, representing the maximum version.
Returns: math.MaxUint64.

`TestPickTables(t *testing.T)`

func TestPickTables(t *testing.T) {
	opt := DefaultIteratorOptions

	within := func(prefix, left, right []byte) {
		opt.Prefix = prefix
		// PickTable expects smallest and biggest to contain timestamps.
		tm := &tableMock{left: y.KeyWithTs(left, 1), right: y.KeyWithTs(right, 1)}
		require.True(t, opt.pickTable(tm), "within failed for %b %b %b\n", prefix, left, right)
	}
	outside := func(prefix, left, right string) {
		opt.Prefix = []byte(prefix)
		// PickTable expects smallest and biggest to contain timestamps.
		tm := &tableMock{left: y.KeyWithTs([]byte(left), 1), right: y.KeyWithTs([]byte(right), 1)}
		require.False(t, opt.pickTable(tm), "outside failed for %b %b %b", prefix, left, right)
	}
	within([]byte("abc"), []byte("ab"), []byte("ad"))
	within([]byte("abc"), []byte("abc"), []byte("ad"))
	within([]byte("abc"), []byte("abb123"), []byte("ad"))
	within([]byte("abc"), []byte("abc123"), []byte("abd234"))
	within([]byte("abc"), []byte("abc123"), []byte("abc456"))
	// Regression test for https://github.com/hypermodeinc/badger/issues/992
	within([]byte{0, 0, 1}, []byte{0}, []byte{0, 0, 1})

	outside("abd", "abe", "ad")
	outside("abd", "ac", "ad")
	outside("abd", "b", "e")
	outside("abd", "a", "ab")
	outside("abd", "ab", "abc")
	outside("abd", "ab", "abc123")
}

Purpose: Tests the pickTable method of IteratorOptions to ensure it correctly determines whether a table falls within a given prefix.
Parameters:
- t: A pointer to the testing.T struct for test management.
Logic: Defines helper functions within and outside to assert that pickTable returns true or false, respectively, for different prefixes and key ranges. It then calls these functions with various test cases.

`TestPickSortTables(t *testing.T)`

func TestPickSortTables(t *testing.T) {
	type MockKeys struct {
		small string
		large string
	}
	genTables := func(mks ...MockKeys) []*table.Table {
		out := make([]*table.Table, 0)
		for _, mk := range mks {
			opts := table.Options{ChkMode: options.OnTableAndBlockRead}
			tbl := buildTable(t, [][]string{{mk.small, "some value"},
				{mk.large, "some value"}}, opts)
			defer func() { require.NoError(t, tbl.DecrRef()) }()
			out = append(out, tbl)
		}
		return out
	}
	tables := genTables(MockKeys{small: "a", large: "abc"},
		MockKeys{small: "abcd", large: "cde"},
		MockKeys{small: "cge", large: "chf"},
		MockKeys{small: "glr", large: "gyup"})
	opt := DefaultIteratorOptions
	opt.prefixIsKey = false
	opt.Prefix = []byte("c")
	filtered := opt.pickTables(tables)
	require.Equal(t, 2, len(filtered))
	// build table adds time stamp so removing tailing bytes.
	require.Equal(t, filtered[0].Smallest()[:4], []byte("abcd"))
	require.Equal(t, filtered[1].Smallest()[:3], []byte("cge"))
	tables = genTables(MockKeys{small: "a", large: "abc"},
		MockKeys{small: "abcd", large: "ade"},
		MockKeys{small: "cge", large: "chf"},
		MockKeys{small: "glr", large: "gyup"})
	filtered = opt.pickTables(tables)
	require.Equal(t, 1, len(filtered))
	require.Equal(t, filtered[0].Smallest()[:3], []byte("cge"))
	tables = genTables(MockKeys{small: "a", large: "abc"},
		MockKeys{small: "abcd", large: "ade"},
		MockKeys{small: "cge", large: "chf"},
		MockKeys{small: "ckr", large: "cyup"},
		MockKeys{small: "csfr", large: "gyup"})
	filtered = opt.pickTables(tables)
	require.Equal(t, 3, len(filtered))
	require.Equal(t, filtered[0].Smallest()[:3], []byte("cge"))
	require.Equal(t, filtered[1].Smallest()[:3], []byte("ckr"))
	require.Equal(t, filtered[2].Smallest()[:4], []byte("csfr"))

	opt.Prefix = []byte("aa")
	filtered = opt.pickTables(tables)
	require.Equal(t, y.ParseKey(filtered[0].Smallest()), []byte("a"))
	require.Equal(t, y.ParseKey(filtered[0].Biggest()), []byte("abc"))
}

Purpose: Tests the pickTables method of IteratorOptions, ensuring that it correctly filters and sorts a slice of table.Table based on a given prefix.
Parameters:
- t: A pointer to the testing.T struct for test management.
Logic:
- Defines a MockKeys struct to represent the smallest and largest keys in a mock table.
- Defines a genTables function to generate a slice of table.Table from a slice of MockKeys.
- Creates several slices of table.Table with different key ranges.
- Sets the Prefix field of DefaultIteratorOptions to different values.
- Calls the pickTables method with each slice of tables and the iterator options.
- Asserts that the returned slice of tables contains the expected tables in the correct order.

`TestIterateSinceTs(t *testing.T)`

func TestIterateSinceTs(t *testing.T) {
	bkey := func(i int) []byte {
		return []byte(fmt.Sprintf("%04d", i))
	}
	val := []byte("OK")
	n := 100000

	runBadgerTest(t, nil, func(t *testing.T, db *DB) {
		batch := db.NewWriteBatch()
		for i := 0; i < n; i++ {
			if (i % 10000) == 0 {
				t.Logf("Put i=%d\n", i)
			}
			require.NoError(t, batch.Set(bkey(i), val))
		}
		require.NoError(t, batch.Flush())

		maxVs := db.MaxVersion()
		sinceTs := maxVs - maxVs/10
		iopt := DefaultIteratorOptions
		iopt.SinceTs = sinceTs

		require.NoError(t, db.View(func(txn *Txn) error {
			it := txn.NewIterator(iopt)
			defer it.Close()

			for it.Rewind(); it.Valid(); it.Next() {
				i := it.Item()
				require.GreaterOrEqual(t, i.Version(), sinceTs)
			}
			return nil
		}))

	})
}

Purpose: Tests that the iterator only returns keys with a version greater than or equal to a specified timestamp (SinceTs).
Parameters:
- t: A pointer to the testing.T struct for test management.
Logic:
- Writes a large number of keys to the database with a write batch.
- Calculates a sinceTs value based on the maximum version in the database.
- Creates an iterator with the SinceTs option set to the calculated value.
- Iterates over the keys in the database and asserts that each key's version is greater than or equal to sinceTs.

`TestIterateSinceTsWithPendingWrites(t *testing.T)`

func TestIterateSinceTsWithPendingWrites(t *testing.T) {
	// The pending entries still have version=0. Even IteratorOptions.SinceTs is 0, the entries
	// should be visible.
	runBadgerTest(t, nil, func(t *testing.T, db *DB) {
		txn := db.NewTransaction(true)
		defer txn.Discard()
		require.NoError(t, txn.Set([]byte("key1"), []byte("value1")))
		require.NoError(t, txn.Set([]byte("key2"), []byte("value2")))
		itr := txn.NewIterator(DefaultIteratorOptions)
		defer itr.Close()
		count := 0
		for itr.Rewind(); itr.Valid(); itr.Next() {
			count++
		}
		require.Equal(t, 2, count)
	})
}

Purpose: Tests that the iterator includes pending writes (writes that have not yet been flushed to disk) even when SinceTs is set to 0.
Parameters:
- t: A pointer to the testing.T struct for test management.
Logic:
- Creates a new transaction and sets two key-value pairs.
- Creates an iterator with the default options (including SinceTs set to 0).
- Iterates over the keys in the transaction and asserts that both key-value pairs are returned.

`TestIteratePrefix(t *testing.T)`

func TestIteratePrefix(t *testing.T) {
	if !*manual {
		t.Skip("Skipping test meant to be run manually.")
		return
	}
	testIteratorPrefix := func(t *testing.T, db *DB) {
		bkey := func(i int) []byte {
			return []byte(fmt.Sprintf("%04d", i))
		}
		val := []byte("OK")
		n := 10000

		batch := db.NewWriteBatch()
		for i := 0; i &lt; n; i++ {
			if (i % 1000) == 0 {
				t.Logf("Put i=%d\n", i)
			}
			require.NoError(t, batch.Set(bkey(i), val))
		}
		require.NoError(t, batch.Flush())

		countKeys := func(prefix string) int {
			t.Logf("Testing with prefix: %s", prefix)
			var count int
			opt := DefaultIteratorOptions
			opt.Prefix = []byte(prefix)
			err := db.View(func(txn *Txn) error {
				itr := txn.NewIterator(opt)
				defer itr.Close()
				for itr.Rewind(); itr.Valid(); itr.Next() {
					item := itr.Item()
					err := item.Value(func(v []byte) error {
						require.Equal(t, val, v)
						return nil
					})
					require.NoError(t, err)
					require.True(t, bytes.HasPrefix(item.Key(), opt.Prefix))
					count++
				}
				return nil
			})
			require.NoError(t, err)
			return count
		}

		countOneKey := func(key []byte) int {
			var count int
			err := db.View(func(txn *Txn) error {
				itr := txn.NewKeyIterator(key, DefaultIteratorOptions)
				defer itr.Close()
				for itr.Rewind(); itr.Valid(); itr.Next() {
					item := itr.Item()
					err := item.Value(func(v []byte) error {
						require.Equal(t, val, v)
						return nil
					})
					require.NoError(t, err)
					require.Equal(t, key, item.Key())
					count++
				}
				return nil
			})
			require.NoError(t, err)
			return count
		}

		for i := 0; i &lt;= 9; i++ {
			require.Equal(t, 1, countKeys(fmt.Sprintf("%d%d%d%d", i, i, i, i)))
			require.Equal(t, 10, countKeys(fmt.Sprintf("%d%d%d", i, i, i)))
			require.Equal(t, 100, countKeys(fmt.Sprintf("%d%d", i, i)))
			require.Equal(t, 1000, countKeys(fmt.Sprintf("%d", i)))
		}
		require.Equal(t, 10000, countKeys(""))

		t.Logf("Testing each key with key iterator")
		for i := 0; i &lt; n; i++ {
			require.Equal(t, 1, countOneKey(bkey(i)))
		}
	}

	t.Run("With Default options", func(t *testing.T) {
		runBadgerTest(t, nil, func(t *testing.T, db *DB) {
			testIteratorPrefix(t, db)
		})
	})

	t.Run("With Block Offsets in Cache", func(t *testing.T) {
		opts := getTestOptions("")
		opts.IndexCacheSize = 100 &lt;&lt; 20
		runBadgerTest(t, &opts, func(t *testing.T, db *DB) {
			testIteratorPrefix(t, db)
		})
	})

	t.Run("With Block Offsets and Blocks in Cache", func(t *testing.T) {
		opts := getTestOptions("")
		opts.BlockCacheSize = 100 &lt;&lt; 20
		opts.IndexCacheSize = 100 &lt;&lt; 20
		runBadgerTest(t, &opts, func(t *testing.T, db *DB) {
			testIteratorPrefix(t, db)
		})
	})

	t.Run("With Blocks in Cache", func(t *testing.T) {
		opts := getTestOptions("")
		opts.BlockCacheSize = 100 &lt;&lt; 20
		runBadgerTest(t, &opts, func(t *testing.T, db *DB) {
			testIteratorPrefix(t, db)
		})
	})

}

Purpose: Tests the iterator's prefix-based iteration functionality.
Parameters:
- t: A pointer to the testing.T struct for test management.
Logic:
- Inserts a set of keys with a specific format into the database.
- Defines two helper functions:
  - countKeys: Counts the number of keys that match a given prefix using a regular iterator.
  - countOneKey: Counts the number of keys that match a given key using a key iterator.
- Calls countKeys and countOneKey with various prefixes and keys, and asserts that the counts are correct.
- The test is run with different cache configurations to test the iterator's performance under different conditions.

`TestIteratorReadOnlyWithNoData(t *testing.T)`

func TestIteratorReadOnlyWithNoData(t *testing.T) {
	dir, err := os.MkdirTemp(".", "badger-test")
	y.Check(err)
	defer removeDir(dir)
	opts := getTestOptions(dir)
	db, err := Open(opts)
	require.NoError(t, err)
	require.NoError(t, db.Close())

	opts.ReadOnly = true
	db, err = Open(opts)
	require.NoError(t, err)
	defer func() {
		require.NoError(t, db.Close())
	}()

	require.NoError(t, db.View(func(txn *Txn) error {
		iopts := DefaultIteratorOptions
		iopts.Prefix = []byte("xxx")
		itr := txn.NewIterator(iopts)
		defer itr.Close()
		return nil
	}))
}

Purpose: Tests that the iterator does not crash the database in read-only mode when the database is empty.
Parameters:
- t: A pointer to the testing.T struct for test management.
Logic:
- Creates a temporary directory and opens a BadgerDB instance in that directory.
- Closes the database and reopens it in read-only mode.
- Creates a new transaction and an iterator with a prefix.
- Closes the iterator and the transaction.
- Asserts that no errors occurred during the process.

`BenchmarkIteratePrefixSingleKey(b *testing.B)`

func BenchmarkIteratePrefixSingleKey(b *testing.B) {
	dir, err := os.MkdirTemp(".", "badger-test")
	y.Check(err)
	defer removeDir(dir)
	opts := getTestOptions(dir)
	db, err := Open(opts)
	y.Check(err)
	defer db.Close()

	N := 100000 // Should generate around 80 SSTables.
	val := []byte("OK")
	bkey := func(i int) []byte {
		return []byte(fmt.Sprintf("%06d", i))
	}

	batch := db.NewWriteBatch()
	for i := 0; i &lt; N; i++ {
		y.Check(batch.Set(bkey(i), val))
	}
	y.Check(batch.Flush())
	var lsmFiles int
	err = filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if strings.HasSuffix(path, ".sst") {
			lsmFiles++
		}
		if err != nil {
			return err
		}
		return nil
	})
	y.Check(err)
	b.Logf("LSM files: %d", lsmFiles)
	b.Logf("Key splits: %v", db.Ranges(nil, 10000))
	b.Logf("Key splits with prefix: %v", db.Ranges([]byte("09"), 10000))

	b.Logf("Outer b.N: %d", b.N)
	b.Run("Key lookups", func(b *testing.B) {
		b.Logf("Inner b.N: %d", b.N)
		for i := 0; i &lt; b.N; i++ {
			key := bkey(rand.Intn(N))
			err := db.View(func(txn *Txn) error {
				opt := DefaultIteratorOptions
				// NOTE: Comment opt.Prefix out here to compare the performance
				// difference between providing Prefix as an option, v/s not. I
				// see a 20% improvement when there are ~80 SSTables.
				opt.Prefix = key
				opt.AllVersions = true

				itr := txn.NewIterator(opt)
				defer itr.Close()

				var count int
				for itr.Seek(key); itr.ValidForPrefix(key); itr.Next() {
					count++
				}
				if count != 1 {
					b.Fatalf("Count must be one key: %s. Found: %d", key, count)
				}
				return nil
			})
			if err != nil {
				b.Fatalf("Error while View: %v", err)
			}
		}
	})
}

Purpose: Benchmarks the performance of iterating over a single key with a prefix.
Parameters:
- b: A pointer to the testing.B struct for benchmark management.
Logic:
- Creates a temporary directory and opens a BadgerDB instance in that directory.
- Writes a large number of keys to the database with a write batch.
- Runs a sub-benchmark called "Key lookups" that iterates over a single key with a prefix.
- The sub-benchmark measures the time it takes to perform a large number of key lookups.
- The opt.Prefix is set to key, and the performance difference between providing Prefix as an option vs not is compared.

Code Examples

None.

iterator_test.go - Overview​

Detailed Documentation​

tableMock struct​

(*tableMock) Smallest() []byte​

(*tableMock) Biggest() []byte​

(*tableMock) DoesNotHave(hash uint32) bool​

(*tableMock) MaxVersion() uint64​

TestPickTables(t *testing.T)​

TestPickSortTables(t *testing.T)​

TestIterateSinceTs(t *testing.T)​

TestIterateSinceTsWithPendingWrites(t *testing.T)​

TestIteratePrefix(t *testing.T)​

TestIteratorReadOnlyWithNoData(t *testing.T)​

BenchmarkIteratePrefixSingleKey(b *testing.B)​

Code Examples​

Getting Started Relevance​

iterator_test.go - Overview

Detailed Documentation

`tableMock` struct

`(*tableMock) Smallest() []byte`

`(*tableMock) Biggest() []byte`

`(*tableMock) DoesNotHave(hash uint32) bool`

`(*tableMock) MaxVersion() uint64`

`TestPickTables(t *testing.T)`

`TestPickSortTables(t *testing.T)`

`TestIterateSinceTs(t *testing.T)`

`TestIterateSinceTsWithPendingWrites(t *testing.T)`

`TestIteratePrefix(t *testing.T)`

`TestIteratorReadOnlyWithNoData(t *testing.T)`

`BenchmarkIteratePrefixSingleKey(b *testing.B)`

Code Examples

Getting Started Relevance