support MultiGet operation

db.go

@@ -495,13 +495,44 @@ func (db *DB) get(key []byte) (y.ValueStruct, error) {
 	for i := 0; i < len(tables); i++ {
 		vs := tables[i].Get(key)
 		y.NumMemtableGets.Add(1)
-		if vs.Meta != 0 || vs.Value != nil {
+		if vs.Valid() {
 			return vs, nil
 		}
 	}
 	return db.lc.get(key)
 }
 
+func (db *DB) multiGet(pairs []keyValuePair) error {
+	tables := db.getMemTables() // Lock should be released.
+	defer func() {
+		for _, tbl := range tables {
+			tbl.DecrRef()
+		}
+	}()
+	y.NumGets.Add(int64(len(pairs)))
+	var foundCount int
+	for i := 0; i < len(tables); i++ {
+		table := tables[i]
+		for j := range pairs {
+			pair := &pairs[j]
+			if pair.found {
+				continue
+			}
+			val := table.Get(pair.key)
+			y.NumMemtableGets.Add(1)
+			if val.Valid() {
+				pair.val = val
+				pair.found = true
+				foundCount++
+			}
+		}
+	}
+	if foundCount == len(pairs) {
+		return nil
+	}
+	return db.lc.multiGet(pairs)
+}
+
 func (db *DB) updateOffset(ptrs []valuePointer) {
 	var ptr valuePointer
 	for i := len(ptrs) - 1; i >= 0; i-- {

db_test.go

@@ -387,6 +387,23 @@ func TestGetMore(t *testing.T) {
 			txn.Discard()
 		}
 
+		// MultiGet
+		var multiGetKeys [][]byte
+		var expectedValues []string
+		for i := 0; i < n; i += 100 {
+			multiGetKeys = append(multiGetKeys, data(i))
+			expectedValues = append(expectedValues, fmt.Sprintf("zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz%9d", i))
+		}
+		txn1 := db.NewTransaction(false)
+		items, err := txn1.MultiGet(multiGetKeys)
+		require.NoError(t, err)
+		for i, item := range items {
+			val, err1 := item.Value()
+			require.NoError(t, err1)
+			require.Equal(t, expectedValues[i], string(val))
+		}
+		txn1.Discard()
+
 		// "Delete" key.
 		for i := 0; i < n; i += m {
 			if (i % 10000) == 0 {

levels.go

@@ -860,14 +860,33 @@ func (s *levelsController) get(key []byte) (y.ValueStruct, error) {
 		if err != nil {
 			return y.ValueStruct{}, errors.Wrapf(err, "get key: %q", key)
 		}
-		if vs.Value == nil && vs.Meta == 0 {
-			continue
+		if vs.Valid() {
+			return vs, nil
 		}
-		return vs, nil
 	}
 	return y.ValueStruct{}, nil
 }
 
+func (s *levelsController) multiGet(pairs []keyValuePair) error {
+	for _, h := range s.levels {
+		for i := range pairs {
+			pair := &pairs[i]
+			if pair.found {
+				continue
+			}
+			val, err := h.get(pair.key) // Calls h.RLock() and h.RUnlock().
+			if err != nil {
+				return errors.Wrapf(err, "get key: %q", pair.key)
+			}
+			if val.Valid() {
+				pair.val = val
+				pair.found = true
+			}
+		}
+	}
+	return nil
+}
+
 func appendIteratorsReversed(out []y.Iterator, th []*table.Table, reversed bool) []y.Iterator {
 	for i := len(th) - 1; i >= 0; i-- {
 		// This will increment the reference of the table handler.

transaction.go

@@ -349,7 +349,7 @@ func (txn *Txn) Get(key []byte) (item *Item, rerr error) {
 	if err != nil {
 		return nil, errors.Wrapf(err, "DB::Get key: %q", key)
 	}
-	if vs.Value == nil && vs.Meta == 0 {
+	if !vs.Valid() {
 		return nil, ErrKeyNotFound
 	}
 	if isDeleted(vs.Meta) {
@@ -366,6 +366,49 @@ func (txn *Txn) Get(key []byte) (item *Item, rerr error) {
 	return item, nil
 }
 
+type keyValuePair struct {
+	key   []byte
+	found bool
+	val   y.ValueStruct
+}
+
+// MultiGet gets items for keys, if not found, the corresponding item will be nil.
+// It only supports read-only transaction for simplicity.
+func (txn *Txn) MultiGet(keys [][]byte) (items []*Item, err error) {
+	if txn.update {
+		return nil, errors.New("not supported")
+	}
+	if txn.discarded {
+		return nil, ErrDiscardedTxn
+	}
+	keyValuePairs := make([]keyValuePair, len(keys))
+	for i, key := range keys {
+		if len(key) == 0 {
+			return nil, ErrEmptyKey
+		}
+		keyValuePairs[i].key = y.KeyWithTs(key, txn.readTs)
+	}
+	err = txn.db.multiGet(keyValuePairs)
+	if err != nil {
+		return nil, err
+	}
+	items = make([]*Item, len(keys))
+	for i, pair := range keyValuePairs {
+		if pair.found {
+			items[i] = &Item{
+				key:      keys[i],
+				version:  pair.val.Version,
+				meta:     pair.val.Meta,
+				userMeta: pair.val.UserMeta,
+				db:       txn.db,
+				vptr:     pair.val.Value,
+				txn:      txn,
+			}
+		}
+	}
+	return items, nil
+}
+
 // Discard discards a created transaction. This method is very important and must be called. Commit
 // method calls this internally, however, calling this multiple times doesn't cause any issues. So,
 // this can safely be called via a defer right when transaction is created.

transaction_test.go

@@ -838,3 +838,29 @@ func TestArmV7Issue311Fix(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+func TestTxnMultiGet(t *testing.T) {
+	runBadgerTest(t, nil, func(t *testing.T, db *DB) {
+		txn := db.NewTransaction(true)
+		keys := make([][]byte, 10)
+		vals := make([][]byte, 10)
+		for i := 0; i < 10; i++ {
+			k := []byte(fmt.Sprintf("key=%d", i))
+			keys[i] = k
+			v := []byte(fmt.Sprintf("val=%d", i))
+			vals[i] = v
+			txn.Set(k, v)
+		}
+		require.NoError(t, txn.Commit())
+
+		txn = db.NewTransaction(false)
+		items, err := txn.MultiGet(keys)
+		require.NoError(t, err)
+		for i, item := range items {
+			val, err := item.Value()
+			require.NoError(t, err)
+			require.Equal(t, vals[i], val)
+		}
+		txn.Discard()
+	})
+}

y/iterator.go

@@ -52,6 +52,11 @@ func (v *ValueStruct) Encode(b []byte) {
 	copy(b[2+len(v.UserMeta):], v.Value)
 }
 
+// Valid checks if the ValueStruct is valid.
+func (v *ValueStruct) Valid() bool {
+	return v.Meta != 0 || v.Value != nil
+}
+
 // EncodeTo should be kept in sync with the Encode function above. The reason
 // this function exists is to avoid creating byte arrays per key-value pair in
 // table/builder.go.