Make KVItem.Value() method asynchronous, and return error.

KVItem.Value() now takes a func that is called with the slice containing
the bytes for the value. This slice directly references the memory map
for the value log. It allows us to avoid copying bytes from the mmap-ed
value log.

Other related changes:

* Add a field to KVItem to hold a reference to the corresponding KV.

* Add a field to KVItem to track whether the value has been prefetched.

* Add a field to KVItem to record any errors during prefetching.

* Change the semantics of iteration. We don’t fetch values by design
anymore. So FetchValues does not apply normally. It only applies in the
case of prefetching. So renamed it to PrefetchValues. PrefetchSize now
only applies if PrefetchValues is true.

* If PrefetchValues is true, we copy values from mmap and set item.val.

* Remove KV.FillValue() method. We don’t fetch value synchronously
anymore, this method is not relevant.

* Refactor some private methods in value.go

* Fixes to tests. Added a helper method to asynchronously retrieve the
value and assign it to a variable.

Author: deepakjois

doc_test.go

@@ -41,16 +41,38 @@ func Example() {
 	var item badger.KVItem
 	if err := kv.Get(key, &item); err != nil {
 		fmt.Printf("Error while getting key: %q", key)
+		return
 	}
-	fmt.Printf("GET %s %s\n", key, item.Value())
+	var val []byte
+	err := item.Value(func(v []byte) {
+		val = make([]byte, len(v))
+		copy(val, v)
+	})
+	if err != nil {
+		fmt.Printf("Error while getting value for key: %q", key)
+		return
+	}
+
+	fmt.Printf("GET %s %s\n", key, val)
 
 	if err := kv.CompareAndSet(key, []byte("venus"), 100); err != nil {
 		fmt.Println("CAS counter mismatch")
 	} else {
 		if err = kv.Get(key, &item); err != nil {
 			fmt.Printf("Error while getting key: %q", key)
 		}
-		fmt.Printf("Set to %s\n", item.Value())
+
+		err := item.Value(func(v []byte) {
+			val = make([]byte, len(v))
+			copy(val, v)
+		})
+
+		if err != nil {
+			fmt.Printf("Error while getting value for key: %q", key)
+			return
+		}
+
+		fmt.Printf("Set to %s\n", val)
 	}
 	if err := kv.CompareAndSet(key, []byte("mars"), item.Counter()); err == nil {
 		fmt.Println("Set to mars")

iterator.go

@@ -23,10 +23,20 @@ import (
 	"github.com/dgraph-io/badger/y"
 )
 
+type prefetchStatus uint8
+
+const (
+	empty prefetchStatus = iota
+	prefetched
+)
+
 // KVItem is returned during iteration. Both the Key() and Value() output is only valid until
 // iterator.Next() is called.
 type KVItem struct {
+	status     prefetchStatus
+	err        error
 	wg         sync.WaitGroup
+	kv         *KV
 	key        []byte
 	vptr       []byte
 	meta       byte
@@ -42,12 +52,21 @@ func (item *KVItem) Key() []byte {
 	return item.key
 }
 
-// Value returns the value, generally fetched from the value log. This call can block while the
-// value is populated asynchronously via a disk read. Remember to parse or copy it if you need to
-// reuse it. DO NOT modify or append to this slice; it would result in internal data overwrite.
-func (item *KVItem) Value() []byte {
+// Value retrieves the value of the item from the value log. It calls the
+// consumer function with a slice argument representing the value. In case
+// of error, the consumer function is not called
+//
+// Remember to parse or copy it if you need to reuse it. DO NOT modify or
+// append to this slice; it would result in a panic.
+func (item *KVItem) Value(consumer func([]byte)) error {
 	item.wg.Wait()
-	return item.val
+	if item.status == prefetched {
+		if item.err != nil {
+			return item.err
+		}
+		consumer(item.val)
+	}
+	return item.kv.yieldItemValue(item, consumer)
 }
 
 func (item *KVItem) hasValue() bool {
@@ -62,9 +81,25 @@ func (item *KVItem) hasValue() bool {
 	return true
 }
 
-// EstimatedSize returns approximate size of the key-value pair.  This can be called with
-// FetchValues=false, to quickly iterate through and estimate the size of a range of key-value
-// pairs (without fetching the corresponding values).
+func (item *KVItem) prefetchValue() {
+	item.err = item.kv.yieldItemValue(item, func(val []byte) {
+		if val == nil {
+			return
+		}
+		buf := item.slice.Resize(len(val))
+		// FIXME in case of non-mmaped read buf and val might be the same location, in
+		// which case this is redundant. Not sure if this is a no-op in that case.
+		copy(buf, val)
+		item.val = buf
+		item.status = prefetched
+	})
+}
+
+// EstimatedSize returns approximate size of the key-value pair.
+//
+// This can be called while iterating through a store to quickly estimate the
+// size of a range of key-value pairs (without fetching the corresponding
+// values).
 func (item *KVItem) EstimatedSize() int64 {
 	if !item.hasValue() {
 		return 0
@@ -121,16 +156,18 @@ func (l *list) pop() *KVItem {
 
 // IteratorOptions is used to set options when iterating over Badger key-value stores.
 type IteratorOptions struct {
-	PrefetchSize int  // How many KV pairs to prefetch while iterating.
-	FetchValues  bool // Controls whether the values should be fetched from the value log.
+	// Indicates whether we should prefetch values during iteration and store them.
+	PrefetchValues bool
+	// How many KV pairs to prefetch while iterating. Valid only if PrefetchValues is true.
+	PrefetchSize int
 	Reverse      bool // Direction of iteration. False is forward, true is backward.
 }
 
 // DefaultIteratorOptions contains default options when iterating over Badger key-value stores.
 var DefaultIteratorOptions = IteratorOptions{
-	PrefetchSize: 100,
-	FetchValues:  true,
-	Reverse:      false,
+	PrefetchValues: false,
+	PrefetchSize:   100,
+	Reverse:        false,
 }
 
 // Iterator helps iterating over the KV pairs in a lexicographically sorted order.
@@ -147,7 +184,7 @@ type Iterator struct {
 func (it *Iterator) newItem() *KVItem {
 	item := it.waste.pop()
 	if item == nil {
-		item = &KVItem{slice: new(y.Slice)}
+		item = &KVItem{slice: new(y.Slice), kv: it.kv}
 	}
 	return item
 }
@@ -206,20 +243,20 @@ func (it *Iterator) fill(item *KVItem) {
 	item.key = y.Safecopy(item.key, it.iitr.Key())
 	item.vptr = y.Safecopy(item.vptr, vs.Value)
 	item.val = nil
-	if it.opt.FetchValues {
+	if it.opt.PrefetchValues {
 		item.wg.Add(1)
 		go func() {
-			it.kv.fillItem(item)
+			// FIXME we are not handling errors here.
+			item.prefetchValue()
 			item.wg.Done()
 		}()
 	}
 }
 
 func (it *Iterator) prefetch() {
-	prefetchSize := it.opt.PrefetchSize
-	if it.opt.PrefetchSize <= 1 {
-		// Try prefetching atleast the first two items to put into it.item and it.data.
-		prefetchSize = 2
+	prefetchSize := 2
+	if it.opt.PrefetchValues && it.opt.PrefetchSize > 1 {
+		prefetchSize = it.opt.PrefetchSize
 	}
 
 	i := it.iitr

kv.go

@@ -180,7 +180,16 @@ func NewKV(optParam *Options) (out *KV, err error) {
 	if err := out.Get(head, &item); err != nil {
 		return nil, errors.Wrap(err, "Retrieving head")
 	}
-	val := item.Value()
+
+	var val []byte
+	err = item.Value(func(v []byte) {
+		val = make([]byte, len(v))
+		copy(val, v)
+	})
+
+	if err != nil {
+		return nil, errors.Wrap(err, "Retrieving head value")
+	}
 	// lastUsedCasCounter will either be the value stored in !badger!head, or some subsequently
 	// written value log entry that we replay.  (Subsequent value log entries might be _less_
 	// than lastUsedCasCounter, if there was value log gc so we have to max() values while
@@ -389,53 +398,29 @@ func (s *KV) getMemTables() ([]*skl.Skiplist, func()) {
 	}
 }
 
-// FillValue populates item with a value.
-//
-// item must be a valid KVItem returned by Badger during iteration. This method
-// could be used to fetch values explicitly during a key-only iteration
-// (FetchValues is set to false). It is useful for example, if values are
-// required for some keys only.
-//
-// This method should not be called when iteration is performed with
-// FetchValues set to true, as it will cause additional copying.
-//
-// Multiple calls to this method will result in multiple copies from the value
-// log. It is the caller’s responsibility to make sure they don’t call this
-// method more than once.
-func (s *KV) FillValue(item *KVItem) error {
-	// Wait for any pending fill operations to finish.
-	item.wg.Wait()
-	item.wg.Add(1)
-	defer item.wg.Done()
-	return s.fillItem(item)
-}
-
-func (s *KV) fillItem(item *KVItem) error {
+func (s *KV) yieldItemValue(item *KVItem, consumer func([]byte)) error {
 	if !item.hasValue() {
-		item.val = nil
+		consumer(nil)
 		return nil
 	}
 
 	if item.slice == nil {
 		item.slice = new(y.Slice)
 	}
+
 	if (item.meta & BitValuePointer) == 0 {
-		item.val = item.slice.Resize(len(item.vptr))
-		copy(item.val, item.vptr)
+		val := item.slice.Resize(len(item.vptr))
+		copy(val, item.vptr)
+		consumer(val)
 		return nil
 	}
 
 	var vp valuePointer
 	vp.Decode(item.vptr)
-	entry, err := s.vlog.Read(vp, item.slice)
+	err := s.vlog.Read(vp, item.slice, consumer)
 	if err != nil {
-		return errors.Wrapf(err, "Unable to read from value log: %+v", vp)
-	}
-	if (entry.Meta & BitDelete) != 0 { // Is a tombstone.
-		item.val = nil
-		return nil
+		return err
 	}
-	item.val = entry.Value
 	return nil
 }
 
@@ -463,18 +448,14 @@ func (s *KV) Get(key []byte, item *KVItem) error {
 	if err != nil {
 		return errors.Wrapf(err, "KV::Get key: %q", key)
 	}
-	if item.slice == nil {
-		item.slice = new(y.Slice)
-	}
+
 	item.meta = vs.Meta
 	item.userMeta = vs.UserMeta
 	item.casCounter = vs.CASCounter
 	item.key = key
+	item.kv = s
 	item.vptr = vs.Value
 
-	if err := s.fillItem(item); err != nil {
-		return errors.Wrapf(err, "KV::Get key: %q", key)
-	}
 	return nil
 }