From f4b7918e9336b883d78accd297e3fe749d105ef2 Mon Sep 17 00:00:00 2001
From: James Mills <1290234+prologic@users.noreply.github.com>
Date: Wed, 13 Mar 2019 20:21:15 +1000
Subject: [PATCH] Add flock on database Open()/Close() to prevent multiple
 concurrent processes write access. Fixes #2

---
 bitcask.go      | 21 ++++++++++++++++++++-
 bitcask_test.go | 15 +++++++++++++++
 go.mod          |  1 +
 go.sum          |  2 ++
 4 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/bitcask.go b/bitcask.go
index b26bd55..17211d2 100644
--- a/bitcask.go
+++ b/bitcask.go
@@ -11,6 +11,8 @@ import (
 	"strconv"
 	"strings"
 	"time"
+
+	"github.com/gofrs/flock"
 )
 
 const (
@@ -18,10 +20,13 @@ 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 {
+	*flock.Flock
+
 	path      string
 	curr      *Datafile
 	keydir    *Keydir
@@ -31,6 +36,10 @@ type Bitcask struct {
 }
 
 func (b *Bitcask) Close() error {
+	defer func() {
+		b.Flock.Unlock()
+	}()
+
 	for _, df := range b.datafiles {
 		df.Close()
 	}
@@ -343,6 +352,7 @@ func Open(path string, options ...func(*Bitcask) error) (*Bitcask, error) {
 	}
 
 	bitcask := &Bitcask{
+		Flock:     flock.New(filepath.Join(path, "lock")),
 		path:      path,
 		curr:      curr,
 		keydir:    keydir,
@@ -358,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
 }
diff --git a/bitcask_test.go b/bitcask_test.go
index b7bbf8e..aed1f9e 100644
--- a/bitcask_test.go
+++ b/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 {
diff --git a/go.mod b/go.mod
index 96256e1..1360dc1 100644
--- a/go.mod
+++ b/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
diff --git a/go.sum b/go.sum
index d5643a3..e4127ac 100644
--- a/go.sum
+++ b/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=