Add Consul-based locking (closes #17)

.travis.yml

@@ -2,6 +2,11 @@ language: go
 go:
   - 1.8
 
+before_script:
+  - wget https://releases.hashicorp.com/consul/0.8.3/consul_0.8.3_linux_amd64.zip
+  - unzip consul_0.8.3_linux_amd64.zip
+  - export PATH=$PATH:$PWD
+
 script:
   - make all
   - make goveralls

src/lock.go

@@ -14,9 +14,14 @@
 package main
 
 import (
+	"errors"
+	"os"
 	"path/filepath"
+	"strconv"
 
-	"github.com/nightlyone/lockfile"
+	"io/ioutil"
+
+	"github.com/hashicorp/consul/api"
 )
 
 // Lock interface abstracting over file-based or consul-based locks
@@ -27,7 +32,7 @@ type Lock interface {
 
 // FileLock is for file-based locks
 type FileLock struct {
-	lock lockfile.Lockfile
+	path string
 }
 
 // InitFileLock builds a FileLock at the path speicifed by name
@@ -36,19 +41,81 @@ func InitFileLock(name string) (Lock, error) {
 	if err != nil {
 		return nil, err
 	}
-	lock, err := lockfile.New(path)
+	return &FileLock{path: path}, nil
+}
+
+// TryLock tries to acquire a lock on a file
+func (fl FileLock) TryLock() error {
+	// need to check that the file doesn't exist since we support locks surviving process shutdown
+	if _, err := os.Stat(fl.path); err == nil {
+		return errors.New("Lock currently held at " + fl.path)
+	}
+
+	pid := os.Getppid()
+	err := ioutil.WriteFile(fl.path, []byte(strconv.Itoa(pid)+"\n"), 0666)
+	return err
+}
+
+// Unlock tries to release the lock on a file
+func (fl FileLock) Unlock() error {
+	return os.Remove(fl.path)
+}
+
+// ConsulLock is for Consul-based locks
+type ConsulLock struct {
+	kv  *api.KV
+	key string
+}
+
+// InitConsulLock builds a ConsulLock (a KV pair in Consul) with the name argument as key
+func InitConsulLock(consulAddress, name string) (Lock, error) {
+	client, err := api.NewClient(&api.Config{Address: consulAddress})
 	if err != nil {
 		return nil, err
 	}
-	return &FileLock{lock: lock}, nil
+
+	kv := client.KV()
+	return &ConsulLock{kv: kv, key: name}, nil
+}
+
+// TryLock tries to acquire a lock from Consul
+func (cl ConsulLock) TryLock() error {
+	p, _, err := cl.kv.Get(cl.key, nil)
+	if err != nil {
+		return err
+	}
+	if p != nil {
+		return errors.New("Lock currently held at " + cl.key)
+	}
+	pid := os.Getppid()
+	_, err = cl.kv.Put(&api.KVPair{Key: cl.key, Value: []byte(strconv.Itoa(pid))}, nil)
+	return err
 }
 
-// TryLock tries to lock the file
-func (fl FileLock) TryLock() error {
-	return fl.lock.TryLock()
+// Unlock tries to release the lock from Consul
+func (cl ConsulLock) Unlock() error {
+	p, _, err := cl.kv.Get(cl.key, nil)
+	if err != nil {
+		return err
+	}
+	if p == nil {
+		return errors.New("Lock not held")
+	}
+	_, err = cl.kv.Delete(cl.key, nil)
+	return err
 }
 
-// Unlock tries to unlock the file
-func (fl FileLock) Unlock() error {
-	return fl.lock.Unlock()
+// GetLock builds a file-based or consul-based lock depending on the consul varialbe
+func GetLock(lock, consul string) (Lock, error) {
+	var l Lock
+	var err error
+	if consul != "" {
+		l, err = InitConsulLock(consul, lock)
+	} else {
+		l, err = InitFileLock(lock)
+	}
+	if err != nil {
+		return nil, err
+	}
+	return l, nil
 }

src/lock_test.go

@@ -19,10 +19,74 @@ import (
 	"strconv"
 	"testing"
 
-	"github.com/nightlyone/lockfile"
+	"github.com/hashicorp/consul/api"
+	"github.com/hashicorp/consul/testutil"
 	"github.com/stretchr/testify/assert"
 )
 
+func makeClient(t *testing.T) (*api.Client, *testutil.TestServer) {
+	// Make client config
+	conf := api.DefaultConfig()
+	// Create server
+	server, err := testutil.NewTestServerConfigT(t, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	conf.Address = server.HTTPAddr
+
+	// Create client
+	client, err := api.NewClient(conf)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	return client, server
+}
+
+func TestConsulLock(t *testing.T) {
+	assert := assert.New(t)
+
+	c, s := makeClient(t)
+	assert.NotNil(c)
+	assert.NotNil(s)
+	defer s.Stop()
+
+	lockName := "lock"
+
+	cl, err := InitConsulLock("some://faulty.address", lockName)
+	assert.NotNil(err)
+	assert.Nil(cl)
+	assert.Equal("Unknown protocol scheme: some", err.Error())
+
+	cl, err = InitConsulLock(s.HTTPAddr, lockName)
+	assert.Nil(err)
+	assert.NotNil(cl)
+
+	err = cl.TryLock()
+	assert.Nil(err)
+
+	// fail if already locked
+	err = cl.TryLock()
+	assert.NotNil(err)
+	assert.Equal("Lock currently held at "+lockName, err.Error())
+
+	// fail if already locked by another lock
+	ocl, err := InitConsulLock(s.HTTPAddr, lockName)
+	assert.Nil(err)
+	assert.NotNil(ocl)
+	err = ocl.TryLock()
+	assert.NotNil(err)
+	assert.Equal("Lock currently held at "+lockName, err.Error())
+
+	// fail if already unlocked
+	err = cl.Unlock()
+	assert.Nil(err)
+
+	err = cl.Unlock()
+	assert.NotNil(err)
+	assert.Equal(api.ErrLockNotHeld, err)
+}
+
 func TestFileLock(t *testing.T) {
 	assert := assert.New(t)
 
@@ -31,7 +95,7 @@ func TestFileLock(t *testing.T) {
 	fl, err := InitFileLock(lockPath)
 	assert.NotNil(fl)
 	assert.Nil(err)
-	assert.Equal(fl, &FileLock{lock: lockfile.Lockfile(lockPath)})
+	assert.Equal(fl, &FileLock{path: lockPath})
 
 	// write to the file so that we can't get a lock on it
 	pid := os.Getppid()
@@ -40,7 +104,7 @@ func TestFileLock(t *testing.T) {
 
 	err = fl.TryLock()
 	assert.NotNil(err)
-	assert.Equal(lockfile.ErrBusy, err)
+	assert.Equal("Lock currently held at "+lockPath, err.Error())
 
 	// cleanup
 	err = os.Remove(lockPath)
@@ -54,5 +118,36 @@ func TestFileLock(t *testing.T) {
 
 	err = fl.Unlock()
 	assert.NotNil(err)
-	assert.Equal(lockfile.ErrRogueDeletion, err)
+	assert.Equal("remove "+lockPath+": no such file or directory", err.Error())
+}
+
+func TestGetLock(t *testing.T) {
+	assert := assert.New(t)
+
+	c, s := makeClient(t)
+	assert.NotNil(c)
+	assert.NotNil(s)
+	defer s.Stop()
+
+	lockName := "/tmp/lock"
+
+	// FileLock if consul == ""
+	lock, err := GetLock(lockName, "")
+	assert.NotNil(lock)
+	assert.Nil(err)
+	assert.Equal(lock, &FileLock{path: lockName})
+
+	// ConsulLock if consul != ""
+	lock, err = GetLock(lockName, s.HTTPAddr)
+	assert.NotNil(lock)
+	assert.Nil(err)
+	cl, ok := lock.(*ConsulLock)
+	assert.NotNil(cl)
+	assert.Equal(true, ok)
+
+	// error otherwise
+	lock, err = GetLock(lockName, "some://faulty.address")
+	assert.Nil(lock)
+	assert.NotNil(err)
+	assert.Equal("Unknown protocol scheme: some", err.Error())
 }