Redis locking in Go with redigo #golang

lock.go

package main

import (
	"github.com/garyburd/redigo/redis"
)

var ErrLockMismatch = errors.New("key is locked with a different secret")

const lockScript = `
local v = redis.call("GET", KEYS[1])
if v == false or v == ARGV[1]
then
	return redis.call("SET", KEYS[1], ARGV[1], "EX", ARGV[2]) and 1
else
	return 0
end
`

const unlockScript = `
local v = redis.call("GET",KEYS[1])
if v == false then
	return 1
elseif v == ARGV[1] then
	return redis.call("DEL",KEYS[1])
else
	return 0
end
`

// writeLock attempts to grab a redis lock. The error returned is safe to ignore
// if all you care about is whether or not the lock was acquired successfully.
func writeLock(name, secret string, ttl uint64) (bool, error) {
	rc := redisPool.Get()
	defer rc.Close()

	script := redis.NewScript(1, lockScript)
	resp, err := redis.Int(script.Do(rc, name, secret, int64(ttl)))
	if err != nil {
		return false, err
	}
	if resp == 0 {
		return false, ErrLockMismatch
	}
	return true, nil
}

// writeLock releases the redis lock
func releaseLock(name, secret string) (bool, error) {
	rc := redisPool.Get()
	defer rc.Close()

	script := redis.NewScript(1, unlockScript)
	resp, err := redis.Int(script.Do(rc, name, secret))
	if err != nil {
		return false, err
	}
	if resp == 0 {
		return false, ErrLockMismatch
	}
	return true, nil
}

lock_test.go

package main

import (
	"github.com/garyburd/redigo/redis"
	"os"
	"testing"
)

func requireRedis(t *testing.T) {
	if os.Getenv("REDIS_URL") == "" {
		t.Fatalf("aggregate tests skipped due to missing REDIS_URL")
	}
	if redisPool == nil {
		e := initRedis(os.Getenv("REDIS_URL"))
		if e != nil {
			t.Fatalf("error initializing redis: %q", e.Error())
		}
	}
}

func TestLocking(t *testing.T) {
	requireRedis(t)
	// attempt to lock
	locked, err := writeLock("mykey", "secret", uint64(10))
	if err != nil {
		t.Fatalf("error acquiring lock: %q", err.Error())
	}
	if !locked {
		t.Fatal("expected writeLock to return true")
	}

	// attempt to re-lock (should succeed)
	locked, err = writeLock("mykey", "secret", uint64(5))
	if !locked || err != nil {
		t.Fatalf("expected re-lock attempt to succeed, got locked=%t error=%q", locked, err)
	}

	// attempt to re-lock with different secret (should return error)
	locked, err = writeLock("mykey", "differentsecret", uint64(5))
	if locked || err != ErrLockMismatch {
		t.Fatalf("expected re-lock attempt to fail, got locked=%t error=%q", locked, err)
	}

	// attempt to unlock w/ bad secret
	unlocked, err := releaseLock("mykey", "wrongsecret")
	if unlocked || err != ErrLockMismatch {
		t.Fatalf("expected unlock w/ bad secret to fail, got unlocked=%t error=%q", unlocked, err)
	}

	// attempt to unlock w/ correct secret
	unlocked, err = releaseLock("mykey", "secret")
	if !unlocked || err != nil {
		t.Fatalf("expected unlock to succeed, got unlocked=%t error=%q", unlocked, err)
	}

	// attempt to unlock again (should return true, nil)
	unlocked, err = releaseLock("mykey", "secret")
	if !unlocked || err != nil {
		t.Fatalf("expected repeat unlock to succeed, got unlocked=%t error=%q", unlocked, err)
	}
}