From 7f945c8fffd35e73f278fc3f0bce9961398dc1e1 Mon Sep 17 00:00:00 2001
From: Pawan Rawal <pawan@dgraph.io>
Date: Thu, 28 Sep 2017 12:11:23 +1000
Subject: [PATCH] Fix #243, make writableLogOffset atomic

---
 kv_test.go | 57 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 value.go   | 16 +++++++++------
 2 files changed, 67 insertions(+), 6 deletions(-)

diff --git a/kv_test.go b/kv_test.go
index d8e1b48f0..b75655aae 100644
--- a/kv_test.go
+++ b/kv_test.go
@@ -18,6 +18,7 @@ package badger
 
 import (
 	"bytes"
+	"crypto/rand"
 	"fmt"
 	"io/ioutil"
 	"os"
@@ -940,3 +941,59 @@ func TestSetIfAbsentAsync(t *testing.T) {
 	require.Equal(t, n, count)
 	require.NoError(t, kv.Close())
 }
+
+func TestGetSetRace(t *testing.T) {
+	dir, err := ioutil.TempDir("", "badger")
+	require.NoError(t, err)
+	defer os.RemoveAll(dir)
+	kv, _ := NewKV(getTestOptions(dir))
+
+	data := make([]byte, 4096)
+	_, err = rand.Read(data)
+	require.NoError(t, err)
+
+	var (
+		numOp = 100
+		wg    sync.WaitGroup
+		keyCh = make(chan string)
+	)
+
+	// writer
+	wg.Add(1)
+	go func() {
+		defer func() {
+			wg.Done()
+			close(keyCh)
+		}()
+
+		for i := 0; i < numOp; i++ {
+			key := fmt.Sprintf("%d", i)
+			err = kv.Set([]byte(key), data, 0x00)
+			require.NoError(t, err)
+			keyCh <- key
+		}
+	}()
+
+	// reader
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+
+		for key := range keyCh {
+			var item KVItem
+
+			err := kv.Get([]byte(key), &item)
+			require.NoError(t, err)
+
+			var val []byte
+			err = item.Value(func(v []byte) error {
+				val = make([]byte, len(v))
+				copy(val, v)
+				return nil
+			})
+			require.NoError(t, err)
+		}
+	}()
+
+	wg.Wait()
+}
diff --git a/value.go b/value.go
index d61993601..dd0f8e3fb 100644
--- a/value.go
+++ b/value.go
@@ -768,7 +768,7 @@ func (vlog *valueLog) Replay(ptr valuePointer, fn logEntry) error {
 	var err error
 	last := vlog.filesMap[vlog.maxFid]
 	lastOffset, err := last.fd.Seek(0, io.SeekEnd)
-	vlog.writableLogOffset = uint32(lastOffset)
+	atomic.AddUint32(&vlog.writableLogOffset, uint32(lastOffset))
 	return errors.Wrapf(err, "Unable to seek to end of value log: %q", last.path)
 }
 
@@ -807,6 +807,10 @@ func (vlog *valueLog) sync() error {
 	return err
 }
 
+func (vlog *valueLog) writableOffset() uint32 {
+	return atomic.LoadUint32(&vlog.writableLogOffset)
+}
+
 // write is thread-unsafe by design and should not be called concurrently.
 func (vlog *valueLog) write(reqs []*request) error {
 	vlog.filesLock.RLock()
@@ -825,10 +829,10 @@ func (vlog *valueLog) write(reqs []*request) error {
 		y.NumWrites.Add(1)
 		y.NumBytesWritten.Add(int64(n))
 		vlog.elog.Printf("Done")
-		vlog.writableLogOffset += uint32(n)
+		atomic.AddUint32(&vlog.writableLogOffset, uint32(n))
 		vlog.buf.Reset()
 
-		if vlog.writableLogOffset > uint32(vlog.opt.ValueLogFileSize) {
+		if vlog.writableOffset() > uint32(vlog.opt.ValueLogFileSize) {
 			var err error
 			if err = curlf.doneWriting(vlog.writableLogOffset); err != nil {
 				return err
@@ -865,7 +869,7 @@ func (vlog *valueLog) write(reqs []*request) error {
 
 			p.Fid = curlf.fid
 			// Use the offset including buffer length so far.
-			p.Offset = vlog.writableLogOffset + uint32(vlog.buf.Len())
+			p.Offset = vlog.writableOffset() + uint32(vlog.buf.Len())
 			plen, err := encodeEntry(e, &vlog.buf) // Now encode the entry into buffer.
 			if err != nil {
 				return err
@@ -903,10 +907,10 @@ func (vlog *valueLog) getFileRLocked(fid uint32) (*logFile, error) {
 // Read reads the value log at a given location.
 func (vlog *valueLog) Read(vp valuePointer, consumer func([]byte) error) error {
 	// Check for valid offset if we are reading to writable log.
-	if vp.Fid == vlog.maxFid && vp.Offset >= vlog.writableLogOffset {
+	if vp.Fid == vlog.maxFid && vp.Offset >= vlog.writableOffset() {
 		return errors.Errorf(
 			"Invalid value pointer offset: %d greater than current offset: %d",
-			vp.Offset, vlog.writableLogOffset)
+			vp.Offset, vlog.writableOffset())
 	}
 
 	fn := func(buf []byte) error {