Add flock on database Open()/Close() to prevent multiple concurrent processes write access. Fixes #2

bitcask.go

@@ -11,6 +11,8 @@ import (
 	"strconv"
 	"strings"
 	"time"
+
+	"github.com/gofrs/flock"
 )
 
 const (
@@ -18,18 +20,29 @@ const (
 )
 
 var (
-	ErrKeyNotFound = errors.New("error: key not found")
+	ErrKeyNotFound       = errors.New("error: key not found")
+	ErrCannotAcquireLock = errors.New("error: cannot acquire lock")
 )
 
 type Bitcask struct {
-	path   string
-	curr   *Datafile
-	keydir *Keydir
+	*flock.Flock
+
+	path      string
+	curr      *Datafile
+	keydir    *Keydir
+	datafiles []*Datafile
 
 	maxDatafileSize int64
 }
 
 func (b *Bitcask) Close() error {
+	defer func() {
+		b.Flock.Unlock()
+	}()
+
+	for _, df := range b.datafiles {
+		df.Close()
+	}
 	return b.curr.Close()
 }
 
@@ -38,26 +51,17 @@ func (b *Bitcask) Sync() error {
 }
 
 func (b *Bitcask) Get(key string) ([]byte, error) {
+	var df *Datafile
+
 	item, ok := b.keydir.Get(key)
 	if !ok {
 		return nil, ErrKeyNotFound
 	}
 
-	var (
-		df  *Datafile
-		err error
-	)
-
-	// Optimization
 	if item.FileID == b.curr.id {
 		df = b.curr
 	} else {
-		// TODO: Pre-open non-active Datafiles and cache the file pointers?
-		df, err = NewDatafile(b.path, item.FileID, true)
-		if err != nil {
-			return nil, err
-		}
-		defer df.Close()
+		df = b.datafiles[item.FileID]
 	}
 
 	e, err := df.ReadAt(item.Index)
@@ -111,6 +115,9 @@ func (b *Bitcask) put(key string, value []byte) (int64, error) {
 			return -1, err
 		}
 
+		df, err := NewDatafile(b.path, b.curr.id, true)
+		b.datafiles = append(b.datafiles, df)
+
 		id := b.curr.id + 1
 		curr, err := NewDatafile(b.path, id, false)
 		if err != nil {
@@ -277,8 +284,6 @@ func Open(path string, options ...func(*Bitcask) error) (*Bitcask, error) {
 		return nil, err
 	}
 
-	keydir := NewKeydir()
-
 	fns, err := getDatafiles(path)
 	if err != nil {
 		return nil, err
@@ -289,7 +294,16 @@ func Open(path string, options ...func(*Bitcask) error) (*Bitcask, error) {
 		return nil, err
 	}
 
+	keydir := NewKeydir()
+	var datafiles []*Datafile
+
 	for i, fn := range fns {
+		df, err := NewDatafile(path, ids[i], true)
+		if err != nil {
+			return nil, err
+		}
+		datafiles = append(datafiles, df)
+
 		if filepath.Ext(fn) == ".hint" {
 			f, err := os.Open(filepath.Join(path, fn))
 			if err != nil {
@@ -307,11 +321,6 @@ func Open(path string, options ...func(*Bitcask) error) (*Bitcask, error) {
 				keydir.Add(key, item.FileID, item.Index, item.Timestamp)
 			}
 		} else {
-			df, err := NewDatafile(path, ids[i], true)
-			if err != nil {
-				return nil, err
-			}
-
 			for {
 				e, err := df.Read()
 				if err != nil {
@@ -336,15 +345,18 @@ func Open(path string, options ...func(*Bitcask) error) (*Bitcask, error) {
 	if len(ids) > 0 {
 		id = ids[(len(ids) - 1)]
 	}
+
 	curr, err := NewDatafile(path, id, false)
 	if err != nil {
 		return nil, err
 	}
 
 	bitcask := &Bitcask{
-		path:   path,
-		curr:   curr,
-		keydir: keydir,
+		Flock:     flock.New(filepath.Join(path, "lock")),
+		path:      path,
+		curr:      curr,
+		keydir:    keydir,
+		datafiles: datafiles,
 
 		maxDatafileSize: DefaultMaxDatafileSize,
 	}
@@ -356,5 +368,14 @@ func Open(path string, options ...func(*Bitcask) error) (*Bitcask, error) {
 		}
 	}
 
+	locked, err := bitcask.Flock.TryLock()
+	if err != nil {
+		return nil, err
+	}
+
+	if !locked {
+		return nil, ErrCannotAcquireLock
+	}
+
 	return bitcask, nil
 }

bitcask_test.go

@@ -198,6 +198,21 @@ func TestMerge(t *testing.T) {
 	})
 }
 
+func TestLocking(t *testing.T) {
+	assert := assert.New(t)
+
+	testdir, err := ioutil.TempDir("", "bitcask")
+	assert.NoError(err)
+
+	db, err := Open(testdir)
+	assert.NoError(err)
+	defer db.Close()
+
+	_, err = Open(testdir)
+	assert.Error(err)
+	assert.Equal("error: cannot acquire lock", err.Error())
+}
+
 func BenchmarkGet(b *testing.B) {
 	testdir, err := ioutil.TempDir("", "bitcask")
 	if err != nil {

go.mod

@@ -1,6 +1,7 @@
 module github.com/prologic/bitcask
 
 require (
+	github.com/gofrs/flock v0.7.1
 	github.com/gogo/protobuf v1.2.1
 	github.com/golang/protobuf v1.2.0
 	github.com/gorilla/websocket v1.4.0 // indirect

go.sum

@@ -10,6 +10,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/fsnotify/fsnotify v1.4.7 h1:IXs+QLmnXW2CcXuY+8Mzv/fWEsPGWxqefPtCP5CnV9I=
 github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
+github.com/gofrs/flock v0.7.1 h1:DP+LD/t0njgoPBvT5MJLeliUIVQR03hiKR6vezdwHlc=
+github.com/gofrs/flock v0.7.1/go.mod h1:F1TvTiK9OcQqauNUHlbJvyl9Qa1QvF/gOUDKA14jxHU=
 github.com/gogo/protobuf v1.2.1 h1:/s5zKNz0uPFCZ5hddgPdo2TK2TVrUNMn0OOX8/aZMTE=
 github.com/gogo/protobuf v1.2.1/go.mod h1:hp+jE20tsWTFYpLwKvXlhS1hjn+gTNwPg2I6zVXpSg4=
 github.com/golang/protobuf v1.2.0 h1:P3YflyNX/ehuJFLhxviNdFxQPkGK5cDcApsge1SqnvM=