Merge pull request coinbase#55 from coinbase/ian-lai/dynamodb-locking

Migrate from S3 to DynamoDB for locking

Author: Graham Jenson

aws/aws.go

@@ -6,6 +6,8 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
 	"github.com/aws/aws-sdk-go/aws/session"
+	"github.com/aws/aws-sdk-go/service/dynamodb"
+	"github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface"
 	"github.com/aws/aws-sdk-go/service/lambda"
 	"github.com/aws/aws-sdk-go/service/lambda/lambdaiface"
 	"github.com/aws/aws-sdk-go/service/s3"
@@ -21,11 +23,13 @@ import (
 type S3API s3iface.S3API
 type LambdaAPI lambdaiface.LambdaAPI
 type SFNAPI sfniface.SFNAPI
+type DynamoDBAPI dynamodbiface.DynamoDBAPI
 
 type AwsClients interface {
 	S3Client(region *string, account_id *string, role *string) S3API
 	LambdaClient(region *string, account_id *string, role *string) LambdaAPI
 	SFNClient(region *string, account_id *string, role *string) SFNAPI
+	DynamoDBClient(region *string, account_id *string, role *string) DynamoDBAPI
 }
 
 ////////////
@@ -106,3 +110,7 @@ func (c *Clients) LambdaClient(region *string, account_id *string, role *string)
 func (c *Clients) SFNClient(region *string, account_id *string, role *string) SFNAPI {
 	return sfn.New(c.Session(), c.Config(region, account_id, role))
 }
+
+func (c *Clients) DynamoDBClient(region *string, account_id *string, role *string) DynamoDBAPI {
+	return dynamodb.New(c.Session(), c.Config(region, account_id, role))
+}

aws/dynamodb/lock.go

@@ -0,0 +1,99 @@
+package dynamodb
+
+import (
+	"fmt"
+	"time"
+
+	awssdk "github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/awserr"
+	"github.com/aws/aws-sdk-go/service/dynamodb"
+	"github.com/aws/aws-sdk-go/service/dynamodb/expression"
+
+	stepaws "github.com/coinbase/step/aws"
+)
+
+var (
+	columnKey  = "key"
+	columnId   = "id"
+	columnTime = "time"
+)
+
+type DynamoDBLocker struct {
+	client stepaws.DynamoDBAPI
+}
+
+func NewDynamoDBLocker(client stepaws.DynamoDBAPI) *DynamoDBLocker {
+	return &DynamoDBLocker{client}
+}
+
+func (l *DynamoDBLocker) GrabLock(namespace string, lockPath string, uuid string, reason string) (bool, error) {
+	// Construct a conditional expression such that we only allow a new lock
+	// to be created if there is not already one for the same key.
+	condExp := expression.Name(columnKey).AttributeNotExists()
+	condExp = condExp.Or(expression.Name(columnId).Equal(expression.Value(uuid)))
+
+	expr, err := expression.NewBuilder().WithCondition(condExp).Build()
+	if err != nil {
+		return false, err
+	}
+
+	// Attempt to create a lock
+	_, err = l.client.PutItem(&dynamodb.PutItemInput{
+		TableName:                 awssdk.String(namespace),
+		ConditionExpression:       expr.Condition(),
+		ExpressionAttributeNames:  expr.Names(),
+		ExpressionAttributeValues: expr.Values(),
+		Item: map[string]*dynamodb.AttributeValue{
+			columnKey: {
+				S: awssdk.String(lockPath),
+			},
+			columnId: {
+				S: awssdk.String(uuid),
+			},
+			columnTime: {
+				S: awssdk.String(time.Now().Format(time.RFC3339)),
+			},
+		},
+	})
+
+	if err != nil {
+		awsErr, ok := err.(awserr.Error)
+		// A lock already exists for the same key.
+		if ok && awsErr.Code() == dynamodb.ErrCodeConditionalCheckFailedException {
+			return false, nil
+		}
+
+		return false, err
+	}
+
+	return true, nil
+}
+
+func (l *DynamoDBLocker) ReleaseLock(namespace string, lockPath string, uuid string) error {
+	// Construct a condition expression such that we only allow a lock
+	// to be deleted if the key, and the UUID aligns.
+	condExp := expression.Name(columnId).Equal(expression.Value(uuid))
+	expr, err := expression.NewBuilder().WithCondition(condExp).Build()
+	if err != nil {
+		return err
+	}
+
+	// Attempt to delete lock
+	_, err = l.client.DeleteItem(&dynamodb.DeleteItemInput{
+		TableName:                 awssdk.String(namespace),
+		ConditionExpression:       expr.Condition(),
+		ExpressionAttributeNames:  expr.Names(),
+		ExpressionAttributeValues: expr.Values(),
+		Key: map[string]*dynamodb.AttributeValue{
+			columnKey: {
+				S: awssdk.String(lockPath),
+			},
+		},
+	})
+
+	if awsErr, ok := err.(awserr.Error); ok && awsErr.Code() == dynamodb.ErrCodeConditionalCheckFailedException {
+		// A lock already exists, but with a different UUID.
+		return fmt.Errorf("Lock was stolen for release with UUID(%v)", uuid)
+	}
+	return err
+}

aws/dynamodb/lock_test.go

@@ -0,0 +1,94 @@
+package dynamodb
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/aws/aws-sdk-go/aws/awserr"
+	"github.com/aws/aws-sdk-go/service/dynamodb"
+	"github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface"
+	"github.com/stretchr/testify/assert"
+)
+
+type MockDynamoDBClient struct {
+	dynamodbiface.DynamoDBAPI
+	putItemCallback    func(input *dynamodb.PutItemInput) (*dynamodb.PutItemOutput, error)
+	deleteItemCallback func(input *dynamodb.DeleteItemInput) (*dynamodb.DeleteItemOutput, error)
+}
+
+func (c *MockDynamoDBClient) PutItem(input *dynamodb.PutItemInput) (*dynamodb.PutItemOutput, error) {
+	return c.putItemCallback(input)
+}
+
+func (c *MockDynamoDBClient) DeleteItem(input *dynamodb.DeleteItemInput) (*dynamodb.DeleteItemOutput, error) {
+	return c.deleteItemCallback(input)
+}
+
+func TestLock(t *testing.T) {
+	t.Run("lock failure", func(t *testing.T) {
+		client := &MockDynamoDBClient{}
+		client.putItemCallback = func(input *dynamodb.PutItemInput) (*dynamodb.PutItemOutput, error) {
+			return nil, awserr.New(dynamodb.ErrCodeConditionalCheckFailedException, "The conditional request failed.", errors.New("fake error"))
+		}
+
+		locker := &DynamoDBLocker{client}
+
+		grabbed, err := locker.GrabLock("tableName", "lockPath", "uuid", "testing")
+		assert.NoError(t, err)
+		assert.False(t, grabbed)
+	})
+
+	t.Run("lock acquired successfully", func(t *testing.T) {
+		client := &MockDynamoDBClient{}
+		client.putItemCallback = func(input *dynamodb.PutItemInput) (*dynamodb.PutItemOutput, error) {
+			assert.Equal(t, "tableName", *input.TableName)
+			assert.Equal(t, "lockPath", *input.Item[columnKey].S)
+			assert.Equal(t, "uuid", *input.Item[columnId].S)
+			assert.Equal(t, "(attribute_not_exists (#0)) OR (#1 = :0)", *input.ConditionExpression)
+
+			assert.Equal(t, "key", *input.ExpressionAttributeNames["#0"])
+			assert.Equal(t, "id", *input.ExpressionAttributeNames["#1"])
+			assert.Equal(t, "uuid", *input.ExpressionAttributeValues[":0"].S)
+
+			return &dynamodb.PutItemOutput{}, nil
+		}
+
+		locker := &DynamoDBLocker{client}
+
+		grabbed, err := locker.GrabLock("tableName", "lockPath", "uuid", "testing")
+		assert.NoError(t, err)
+		assert.True(t, grabbed)
+	})
+}
+
+func TestUnlock(t *testing.T) {
+	t.Run("unlock failure", func(t *testing.T) {
+		client := &MockDynamoDBClient{}
+		client.deleteItemCallback = func(input *dynamodb.DeleteItemInput) (*dynamodb.DeleteItemOutput, error) {
+			return nil, awserr.New(dynamodb.ErrCodeConditionalCheckFailedException, "The conditional request failed.", errors.New("fake error"))
+		}
+
+		locker := &DynamoDBLocker{client}
+
+		err := locker.ReleaseLock("tableName", "lockPath", "uuid")
+		assert.Error(t, err)
+	})
+
+	t.Run("unlock released", func(t *testing.T) {
+		client := &MockDynamoDBClient{}
+		client.deleteItemCallback = func(input *dynamodb.DeleteItemInput) (*dynamodb.DeleteItemOutput, error) {
+			assert.Equal(t, "tableName", *input.TableName)
+
+			assert.Equal(t, "id", *input.ExpressionAttributeNames["#0"])
+			assert.Equal(t, "uuid", *input.ExpressionAttributeValues[":0"].S)
+			assert.Equal(t, "lockPath", *input.Key[columnKey].S)
+
+			return &dynamodb.DeleteItemOutput{}, nil
+		}
+
+		locker := &DynamoDBLocker{client}
+
+		err := locker.ReleaseLock("tableName", "lockPath", "uuid")
+		assert.NoError(t, err)
+	})
+}

aws/mocks/mock_dynamodb.go

@@ -0,0 +1,33 @@
+package mocks
+
+import (
+	"github.com/aws/aws-sdk-go/service/dynamodb"
+	"github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface"
+)
+
+type MockDynamoDBClient struct {
+	dynamodbiface.DynamoDBAPI
+
+	PutItemInputs    []*dynamodb.PutItemInput
+	DeleteItemInputs []*dynamodb.DeleteItemInput
+}
+
+func (m *MockDynamoDBClient) init() {
+	if m.PutItemInputs == nil {
+		m.PutItemInputs = []*dynamodb.PutItemInput{}
+	}
+
+	if m.DeleteItemInputs == nil {
+		m.DeleteItemInputs = []*dynamodb.DeleteItemInput{}
+	}
+}
+
+func (m *MockDynamoDBClient) PutItem(input *dynamodb.PutItemInput) (*dynamodb.PutItemOutput, error) {
+	m.PutItemInputs = append(m.PutItemInputs, input)
+	return &dynamodb.PutItemOutput{}, nil
+}
+
+func (m *MockDynamoDBClient) DeleteItem(input *dynamodb.DeleteItemInput) (*dynamodb.DeleteItemOutput, error) {
+	m.DeleteItemInputs = append(m.DeleteItemInputs, input)
+	return &dynamodb.DeleteItemOutput{}, nil
+}

aws/mocks/mocks.go

@@ -3,9 +3,10 @@ package mocks
 import "github.com/coinbase/step/aws"
 
 type MockClients struct {
-	S3     *MockS3Client
-	Lambda *MockLambdaClient
-	SFN    *MockSFNClient
+	S3       *MockS3Client
+	Lambda   *MockLambdaClient
+	SFN      *MockSFNClient
+	DynamoDB *MockDynamoDBClient
 }
 
 func (awsc *MockClients) S3Client(*string, *string, *string) aws.S3API {
@@ -20,10 +21,15 @@ func (awsc *MockClients) SFNClient(*string, *string, *string) aws.SFNAPI {
 	return awsc.SFN
 }
 
+func (awsc *MockClients) DynamoDBClient(*string, *string, *string) aws.DynamoDBAPI {
+	return awsc.DynamoDB
+}
+
 func MockAwsClients() *MockClients {
 	return &MockClients{
 		&MockS3Client{},
 		&MockLambdaClient{},
 		&MockSFNClient{},
+		&MockDynamoDBClient{},
 	}
 }

aws/s3/s3.go

@@ -157,7 +157,7 @@ func PutWithTypeAndCacheControl(s3c aws.S3API, bucket *string, path *string, con
 		Key:          path,
 		Body:         bytes.NewReader(*content),
 		ACL:          to.Strp("private"),
-		ContentType: contentType,
+		ContentType:  contentType,
 		CacheControl: cacheControl,
 	})
 }

bifrost/inmemory_locker.go

@@ -0,0 +1,88 @@
+package bifrost
+
+import (
+	"fmt"
+	"sync"
+)
+
+type Lock struct {
+	lockPath string
+	uuid     string
+	reason   string
+}
+
+type InMemoryLocker struct {
+	mu    sync.RWMutex
+	locks map[string][]*Lock
+}
+
+func NewInMemoryLocker() *InMemoryLocker {
+	return &InMemoryLocker{
+		locks: make(map[string][]*Lock),
+	}
+}
+
+func (l *InMemoryLocker) GrabLock(namespace string, lockPath string, uuid string, reason string) (bool, error) {
+	existingLock := l.GetLockByPath(namespace, lockPath)
+	if existingLock != nil {
+		return existingLock.uuid == uuid, nil
+	}
+
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	l.locks[namespace] = append(l.locks[namespace], &Lock{
+		lockPath: lockPath,
+		uuid:     uuid,
+		reason:   reason,
+	})
+
+	return true, nil
+}
+
+func (l *InMemoryLocker) ReleaseLock(namespace string, lockPath string, uuid string) error {
+	existingLock := l.GetLockByPath(namespace, lockPath)
+	if existingLock != nil && existingLock.uuid != uuid {
+		return fmt.Errorf("failed to release lock: %s is currently held by UUID(%v)", lockPath, existingLock.uuid)
+	}
+
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	var updatedLocks []*Lock
+	for _, lock := range l.locks[namespace] {
+		if lock.uuid == uuid {
+			continue
+		}
+		updatedLocks = append(updatedLocks, lock)
+	}
+
+	l.locks[namespace] = updatedLocks
+
+	return nil
+}
+
+func (l *InMemoryLocker) GetLockByNamespace(namespace string) []*Lock {
+	l.mu.RLock()
+	defer l.mu.RUnlock()
+
+	locks, found := l.locks[namespace]
+	if !found {
+		return []*Lock{}
+	}
+
+	return locks
+}
+
+func (l *InMemoryLocker) GetLockByPath(namespace string, lockPath string) *Lock {
+	l.mu.RLock()
+	defer l.mu.RUnlock()
+
+	for _, lock := range l.GetLockByNamespace(namespace) {
+		if lock.lockPath == lockPath {
+			return lock
+		}
+	}
+
+	return nil
+}