From ef3163324cf701b0cea2650eba922c7b986a6ba6 Mon Sep 17 00:00:00 2001
From: Mahmoud Abdelsalam <mahmoud.abdelsalam@gmail.com>
Date: Mon, 1 Apr 2019 14:19:18 -0700
Subject: [PATCH] Fix DynamoDB HA race issue

---
 physical/dynamodb/dynamodb.go | 32 ++++++++++++++++++++------------
 1 file changed, 20 insertions(+), 12 deletions(-)

diff --git a/physical/dynamodb/dynamodb.go b/physical/dynamodb/dynamodb.go
index 55f9a9ed70a9..47eba9669b18 100644
--- a/physical/dynamodb/dynamodb.go
+++ b/physical/dynamodb/dynamodb.go
@@ -606,7 +606,7 @@ func (l *DynamoDBLock) tryToLock(stop, success chan struct{}, errors chan error)
 		case <-stop:
 			ticker.Stop()
 		case <-ticker.C:
-			err := l.writeItem()
+			err := l.updateItem(true)
 			if err != nil {
 				if err, ok := err.(awserr.Error); ok {
 					// Don't report a condition check failure, this means that the lock
@@ -633,7 +633,8 @@ func (l *DynamoDBLock) periodicallyRenewLock(done chan struct{}) {
 	for {
 		select {
 		case <-ticker.C:
-			l.writeItem()
+			// This should not renew the lock if the lock was deleted from under you.
+			l.updateItem(false)
 		case <-done:
 			ticker.Stop()
 			return
@@ -643,9 +644,24 @@ func (l *DynamoDBLock) periodicallyRenewLock(done chan struct{}) {
 
 // Attempts to put/update the dynamodb item using condition expressions to
 // evaluate the TTL.
-func (l *DynamoDBLock) writeItem() error {
+func (l *DynamoDBLock) updateItem(createIfMissing bool) error {
 	now := time.Now()
 
+	conditionExpression := ""
+	if createIfMissing {
+		conditionExpression += "attribute_not_exists(#path) or " +
+			"attribute_not_exists(#key) or "
+	} else {
+		conditionExpression += "attribute_exists(#path) and " +
+			"attribute_exists(#key) and "
+	}
+
+	// To work when upgrading from older versions that did not include the
+	// Identity attribute, we first check if the attr doesn't exist, and if
+	// it does, then we check if the identity is equal to our own.
+	// We also write if the lock expired.
+	conditionExpression += "(attribute_not_exists(#identity) or #identity = :identity or #expires <= :now)"
+
 	_, err := l.backend.client.UpdateItem(&dynamodb.UpdateItemInput{
 		TableName: aws.String(l.backend.table),
 		Key: map[string]*dynamodb.AttributeValue{
@@ -657,15 +673,7 @@ func (l *DynamoDBLock) writeItem() error {
 		// A. identity is equal to our identity (or the identity doesn't exist)
 		// or
 		// B. The ttl on the item is <= to the current time
-		ConditionExpression: aws.String(
-			"attribute_not_exists(#path) or " +
-				"attribute_not_exists(#key) or " +
-				// To work when upgrading from older versions that did not include the
-				// Identity attribute, we first check if the attr doesn't exist, and if
-				// it does, then we check if the identity is equal to our own.
-				"(attribute_not_exists(#identity) or #identity = :identity) or " +
-				"#expires <= :now",
-		),
+		ConditionExpression: aws.String(conditionExpression),
 		ExpressionAttributeNames: map[string]*string{
 			"#path":     aws.String("Path"),
 			"#key":      aws.String("Key"),