Added get multiple api for badger

2026-07-27 06:54:45 +00:00 · 2025-02-18 14:25:28 +05:30
parent 57da589029
commit f225f6150b
6 changed files with 325 additions and 0 deletions
@@ -736,6 +736,8 @@ func (db *DB) getMemTables() ([]*memTable, func()) {
 // get returns the value in memtable or disk for given key.
 // Note that value will include meta byte.
 //
+// getBatch would return the values of list of keys in order
+//
 // IMPORTANT: We should never write an entry with an older timestamp for the same key, We need to
 // maintain this invariant to search for the latest value of a key, or else we need to search in all
 // tables and find the max version among them.  To maintain this invariant, we also need to ensure
@@ -747,7 +749,54 @@ func (db *DB) getMemTables() ([]*memTable, func()) {
 // do that. For every get("fooX") call where X is the version, we will search
 // for "fooX" in all the levels of the LSM tree. This is expensive but it
 // removes the overhead of handling move keys completely.
+func (db *DB) getBatch(keys [][]byte, done []bool) ([]y.ValueStruct, error) {
+	if db.IsClosed() {
+		return []y.ValueStruct{}, ErrDBClosed
+	}
+	tables, decr := db.getMemTables() // Lock should be released.
+	defer decr()
+
+	maxVs := make([]y.ValueStruct, len(keys))
+
+	y.NumGetsAdd(db.opt.MetricsEnabled, 1)
+	// For memtable, we need to check every memtable each time
+	for j, key := range keys {
+		if done[j] {
+			continue
+		}
+		version := y.ParseTs(key)
+		for i := 0; i < len(tables); i++ {
+			vs := tables[i].sl.Get(key)
+			y.NumMemtableGetsAdd(db.opt.MetricsEnabled, 1)
+			if vs.Meta == 0 && vs.Value == nil {
+				continue
+			}
+			// Found the required version of the key, mark as done, no need to process
+			// it further
+			if vs.Version == version {
+				y.NumGetsWithResultsAdd(db.opt.MetricsEnabled, 1)
+				maxVs[j] = vs
+				done[j] = true
+				break
+			}
+			if maxVs[j].Version < vs.Version {
+				maxVs[j] = vs
+			}
+		}
+	}
+	return db.lc.getBatch(keys, maxVs, 0, done)
+}
+
 func (db *DB) get(key []byte) (y.ValueStruct, error) {
+	if db.opt.useGetBatch {
+		done := make([]bool, 1)
+		vals, err := db.getBatch([][]byte{key}, done)
+		if len(vals) != 0 {
+			return vals[0], err
+		}
+		return y.ValueStruct{}, err
+	}
+
 	if db.IsClosed() {
 		return y.ValueStruct{}, ErrDBClosed
 	}
@@ -262,6 +262,113 @@ func (s *levelHandler) getTableForKey(key []byte) ([]*table.Table, func() error)
 	return []*table.Table{tbl}, tbl.DecrRef
 }

+func (s *levelHandler) getBatch(keys [][]byte, done []bool) ([]y.ValueStruct, error) {
+	// Find the table for which the key is in, and then seek it
+	getForKey := func(key []byte) (y.ValueStruct, func() error, []*table.Iterator) {
+		tables, decr := s.getTableForKey(key)
+		keyNoTs := y.ParseKey(key)
+		itrs := make([]*table.Iterator, 0)
+
+		hash := y.Hash(keyNoTs)
+		var maxVs y.ValueStruct
+		for _, th := range tables {
+			if th.DoesNotHave(hash) {
+				y.NumLSMBloomHitsAdd(s.db.opt.MetricsEnabled, s.strLevel, 1)
+				continue
+			}
+
+			it := th.NewIterator(0)
+			itrs = append(itrs, it)
+
+			y.NumLSMGetsAdd(s.db.opt.MetricsEnabled, s.strLevel, 1)
+			it.Seek(key)
+			if !it.Valid() {
+				continue
+			}
+			if y.SameKey(key, it.Key()) {
+				if version := y.ParseTs(it.Key()); maxVs.Version < version {
+					maxVs = it.ValueCopy()
+					maxVs.Version = version
+				}
+			}
+		}
+
+		return maxVs, decr, itrs
+	}
+
+	// Use old results from getForKey and find in those tables.
+	findInIter := func(key []byte, itrs []*table.Iterator) y.ValueStruct {
+		var maxVs y.ValueStruct
+
+		for _, it := range itrs {
+			it.Seek(key)
+			if !it.Valid() {
+				continue
+			}
+			if y.SameKey(key, it.Key()) {
+				if version := y.ParseTs(it.Key()); maxVs.Version < version {
+					maxVs = it.ValueCopy()
+					maxVs.Version = version
+				}
+			}
+		}
+
+		return maxVs
+	}
+
+	results := make([]y.ValueStruct, len(keys))
+	// For L0, we need to search all tables each time, so we can just call get() as required
+	if s.level == 0 {
+		var err error
+		for i, key := range keys {
+			if done[i] {
+				continue
+			}
+			results[i], err = s.get(key)
+			if err != nil {
+				return results, err
+			}
+		}
+		return results, nil
+	} else {
+		decr := func() error { return nil }
+		var itrs []*table.Iterator
+
+		started := false
+		for i := 0; i < len(keys); i++ {
+			if done[i] {
+				continue
+			}
+			if !started {
+				var maxVs y.ValueStruct
+				maxVs, decr, itrs = getForKey(keys[0])
+				results[i] = maxVs
+				started = true
+			} else {
+				results[i] = findInIter(keys[i], itrs)
+				// If we can't find in the current tables, maybe the
+				// data is there in other tables
+				if len(results[i].Value) == 0 {
+					for i := 0; i < len(itrs); i++ {
+						itrs[i].Close()
+					}
+					err := decr()
+					if err != nil {
+						return nil, err
+					}
+					results[i], decr, itrs = getForKey(keys[i])
+				}
+			}
+		}
+
+		for i := 0; i < len(itrs); i++ {
+			itrs[i].Close()
+		}
+		return results, decr()
+	}
+
+}
+
 // get returns value for a given key or the key after that. If not found, return nil.
 func (s *levelHandler) get(key []byte) (y.ValueStruct, error) {
 	tables, decr := s.getTableForKey(key)
@@ -1594,6 +1594,55 @@ func (s *levelsController) close() error {
 	return y.Wrap(err, "levelsController.Close")
 }

+func (s *levelsController) getBatch(keys [][]byte, maxVs []y.ValueStruct, startLevel int, done []bool) (
+	[]y.ValueStruct, error) {
+	if s.kv.IsClosed() {
+		return []y.ValueStruct{}, ErrDBClosed
+	}
+	// It's important that we iterate the levels from 0 on upward. The reason is, if we iterated
+	// in opposite order, or in parallel (naively calling all the h.RLock() in some order) we could
+	// read level L's tables post-compaction and level L+1's tables pre-compaction. (If we do
+	// parallelize this, we will need to call the h.RLock() function by increasing order of level
+	// number.)
+	for _, h := range s.levels {
+		// Ignore all levels below startLevel. This is useful for GC when L0 is kept in memory.
+		if h.level < startLevel {
+			continue
+		}
+		vs, err := h.getBatch(keys, done) // Calls h.RLock() and h.RUnlock().
+		if err != nil {
+			return []y.ValueStruct{}, y.Wrapf(err, "get keys: %q", keys)
+		}
+
+		for i, v := range vs {
+			// Done is only update by this function or one in db. levelhandler will
+			// not update done. No need to do anything is done is set.
+			if done[i] {
+				continue
+			}
+			if v.Value == nil && v.Meta == 0 {
+				continue
+			}
+			y.NumBytesReadsLSMAdd(s.kv.opt.MetricsEnabled, int64(len(v.Value)))
+			version := y.ParseTs(keys[i])
+			if v.Version == version {
+				maxVs[i] = v
+				done[i] = true
+			}
+			if maxVs[i].Version < v.Version {
+				maxVs[i] = v
+			}
+		}
+	}
+
+	for i := 0; i < len(maxVs); i++ {
+		if len(maxVs[i].Value) > 0 {
+			y.NumGetsWithResultsAdd(s.kv.opt.MetricsEnabled, 1)
+		}
+	}
+	return maxVs, nil
+}
+
 // get searches for a given key in all the levels of the LSM tree. It returns
 // key version <= the expected version (version in key). If not found,
 // it returns an empty y.ValueStruct.
@@ -118,6 +118,9 @@ type Options struct {
 	maxBatchSize  int64 // max batch size in bytes

 	maxValueThreshold float64
+
+	// This would let you use get batch instead of get, an experimental api instead
+	useGetBatch bool
 }

 // DefaultOptions sets a list of recommended options for good performance.
@@ -176,6 +179,7 @@ func DefaultOptions(path string) Options {
 		EncryptionKeyRotationDuration: 10 * 24 * time.Hour, // Default 10 days.
 		DetectConflicts:               true,
 		NamespaceOffset:               -1,
+		useGetBatch:                   true,
 	}
 }

@@ -429,6 +429,83 @@ func (txn *Txn) Delete(key []byte) error {
 	return txn.modify(e)
 }

+func (txn *Txn) GetBatch(keys [][]byte) (items []*Item, rerr error) {
+	if txn.discarded {
+		return nil, ErrDiscardedTxn
+	}
+
+	for _, key := range keys {
+		if len(key) == 0 {
+			return nil, ErrEmptyKey
+		}
+		if err := txn.db.isBanned(key); err != nil {
+			return nil, err
+		}
+	}
+
+	items = make([]*Item, len(keys))
+	done := make([]bool, len(keys))
+
+	if txn.update {
+		doneAll := 0
+		for i, key := range keys {
+			item := items[i]
+			if e, has := txn.pendingWrites[string(key)]; has && bytes.Equal(key, e.Key) {
+				if isDeletedOrExpired(e.meta, e.ExpiresAt) {
+					return nil, ErrKeyNotFound
+				}
+				// Fulfill from cache.
+				item.meta = e.meta
+				item.val = e.Value
+				item.userMeta = e.UserMeta
+				item.key = key
+				item.status = prefetched
+				item.version = txn.readTs
+				item.expiresAt = e.ExpiresAt
+				// We probably don't need to set db on item here.
+				done[i] = true
+				doneAll += 1
+			}
+			// Only track reads if this is update txn. No need to track read if txn serviced it
+			// internally.
+			txn.addReadKey(key)
+		}
+		if doneAll == len(keys) {
+			return items, nil
+		}
+	}
+
+	seeks := make([][]byte, len(keys))
+	for i, key := range keys {
+		seeks[i] = y.KeyWithTs(key, txn.readTs)
+	}
+	vss, err := txn.db.getBatch(seeks, done)
+	if err != nil {
+		return nil, y.Wrapf(err, "DB::Get keys: %q", keys)
+	}
+
+	for i, vs := range vss {
+		if vs.Value == nil && vs.Meta == 0 {
+			items[i] = nil
+		}
+		if isDeletedOrExpired(vs.Meta, vs.ExpiresAt) {
+			items[i] = nil
+		}
+
+		items[i] = new(Item)
+		items[i].key = keys[i]
+		items[i].version = vs.Version
+		items[i].meta = vs.Meta
+		items[i].userMeta = vs.UserMeta
+		items[i].vptr = y.SafeCopy(items[i].vptr, vs.Value)
+		items[i].txn = txn
+		items[i].expiresAt = vs.ExpiresAt
+	}
+
+	return items, nil
+
+}
+
 // Get looks for key and returns corresponding Item.
 // If key is not found, ErrKeyNotFound is returned.
 func (txn *Txn) Get(key []byte) (item *Item, rerr error) {
@@ -21,6 +21,45 @@ import (
 	"github.com/dgraph-io/ristretto/v2/z"
 )

+func TestTxnSimpleTsRead(t *testing.T) {
+	dir, err := os.MkdirTemp("", "badger-test")
+	require.NoError(t, err)
+	defer removeDir(dir)
+	opts := getTestOptions(dir)
+	opts.Dir = dir
+	opts.ValueDir = dir
+
+	opts.managedTxns = true
+
+	db, err := Open(opts)
+	require.NoError(t, err)
+	defer func() {
+		require.NoError(t, db.Close())
+	}()
+
+	for i := 0; i < 10; i++ {
+		txn := db.NewTransactionAt(uint64(i)+1, true)
+		k := []byte(fmt.Sprintf("key=%d", 1))
+		v := []byte(fmt.Sprintf("val=%d", i))
+		require.NoError(t, txn.SetEntry(NewEntry(k, v)))
+		err = txn.CommitAt(uint64(i)*3+1, nil)
+		require.NoError(t, err)
+	}
+
+	for i := 7; i < 10; i++ {
+		txn := db.NewTransactionAt(uint64(i), false)
+		item, err := txn.Get([]byte("key=1"))
+		require.NoError(t, err)
+
+		require.NoError(t, item.Value(func(val []byte) error {
+			require.Equal(t, []byte("val=2"), val)
+			return nil
+		}))
+
+		txn.Discard()
+	}
+}
+
 func TestTxnSimple(t *testing.T) {
 	runBadgerTest(t, nil, func(t *testing.T, db *DB) {
 		txn := db.NewTransaction(true)