Support enhanced fan-out feature (#90)

* Implement enhanced fan-out consumer

Signed-off-by: Ilia Cimpoes <[email protected]>

* Add test cases

Signed-off-by: Ilia Cimpoes <[email protected]>

* Small adjustments in fan-out consumer

Signed-off-by: Ilia Cimpoes <[email protected]>

Author: iliacimpoes

clientlibrary/checkpoint/dynamodb-checkpointer.go

@@ -31,14 +31,14 @@ import (
 	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
-	"github.com/aws/aws-sdk-go/aws/awserr"
 	"github.com/aws/aws-sdk-go/aws/client"
 	"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/vmware/vmware-go-kcl/clientlibrary/config"
 	par "github.com/vmware/vmware-go-kcl/clientlibrary/partition"
+	"github.com/vmware/vmware-go-kcl/clientlibrary/utils"
 	"github.com/vmware/vmware-go-kcl/logger"
 )
 
@@ -144,7 +144,7 @@ func (checkpointer *DynamoCheckpoint) GetLease(shard *par.ShardStatus, newAssign
 			return ErrLeaseNotAcquired{"current lease timeout not yet expired"}
 		}
 
-		checkpointer.log.Debugf("Attempting to get a lock for shard: %s, leaseTimeout: %s, assignedTo: %s", shard.ID, currentLeaseTimeout, assignedTo)
+		checkpointer.log.Debugf("Attempting to get a lock for shard: %s, leaseTimeout: %s, assignedTo: %s, newAssignedTo: %s", shard.ID, currentLeaseTimeout, assignedTo, newAssignTo)
 		conditionalExpression = "ShardID = :id AND AssignedTo = :assigned_to AND LeaseTimeout = :lease_timeout"
 		expressionAttributeValues = map[string]*dynamodb.AttributeValue{
 			":id": {
@@ -175,18 +175,16 @@ func (checkpointer *DynamoCheckpoint) GetLease(shard *par.ShardStatus, newAssign
 		marshalledCheckpoint[ParentShardIdKey] = &dynamodb.AttributeValue{S: aws.String(shard.ParentShardId)}
 	}
 
-	if shard.Checkpoint != "" {
+	if shard.GetCheckpoint() != "" {
 		marshalledCheckpoint[SequenceNumberKey] = &dynamodb.AttributeValue{
-			S: aws.String(shard.Checkpoint),
+			S: aws.String(shard.GetCheckpoint()),
 		}
 	}
 
 	err = checkpointer.conditionalUpdate(conditionalExpression, expressionAttributeValues, marshalledCheckpoint)
 	if err != nil {
-		if awsErr, ok := err.(awserr.Error); ok {
-			if awsErr.Code() == dynamodb.ErrCodeConditionalCheckFailedException {
-				return ErrLeaseNotAcquired{dynamodb.ErrCodeConditionalCheckFailedException}
-			}
+		if utils.AWSErrCode(err) == dynamodb.ErrCodeConditionalCheckFailedException {
+			return ErrLeaseNotAcquired{dynamodb.ErrCodeConditionalCheckFailedException}
 		}
 		return err
 	}
@@ -207,7 +205,7 @@ func (checkpointer *DynamoCheckpoint) CheckpointSequence(shard *par.ShardStatus)
 			S: aws.String(shard.ID),
 		},
 		SequenceNumberKey: {
-			S: aws.String(shard.Checkpoint),
+			S: aws.String(shard.GetCheckpoint()),
 		},
 		LeaseOwnerKey: {
 			S: aws.String(shard.AssignedTo),
@@ -236,12 +234,10 @@ func (checkpointer *DynamoCheckpoint) FetchCheckpoint(shard *par.ShardStatus) er
 		return ErrSequenceIDNotFound
 	}
 	checkpointer.log.Debugf("Retrieved Shard Iterator %s", *sequenceID.S)
-	shard.Mux.Lock()
-	defer shard.Mux.Unlock()
-	shard.Checkpoint = aws.StringValue(sequenceID.S)
+	shard.SetCheckpoint(aws.StringValue(sequenceID.S))
 
 	if assignedTo, ok := checkpoint[LeaseOwnerKey]; ok {
-		shard.AssignedTo = aws.StringValue(assignedTo.S)
+		shard.SetLeaseOwner(aws.StringValue(assignedTo.S))
 	}
 	return nil
 }

clientlibrary/checkpoint/dynamodb-checkpointer_test.go

@@ -33,12 +33,11 @@ import (
 	"testing"
 	"time"
 
-	"github.com/stretchr/testify/assert"
-
 	"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/dynamodbiface"
+	"github.com/stretchr/testify/assert"
 
 	cfg "github.com/vmware/vmware-go-kcl/clientlibrary/config"
 	par "github.com/vmware/vmware-go-kcl/clientlibrary/partition"
@@ -75,7 +74,7 @@ func TestGetLeaseNotAquired(t *testing.T) {
 	err := checkpoint.GetLease(&par.ShardStatus{
 		ID:         "0001",
 		Checkpoint: "",
-		Mux:        &sync.Mutex{},
+		Mux:        &sync.RWMutex{},
 	}, "abcd-efgh")
 	if err != nil {
 		t.Errorf("Error getting lease %s", err)
@@ -84,7 +83,7 @@ func TestGetLeaseNotAquired(t *testing.T) {
 	err = checkpoint.GetLease(&par.ShardStatus{
 		ID:         "0001",
 		Checkpoint: "",
-		Mux:        &sync.Mutex{},
+		Mux:        &sync.RWMutex{},
 	}, "ijkl-mnop")
 	if err == nil || !errors.As(err, &ErrLeaseNotAcquired{}) {
 		t.Errorf("Got a lease when it was already held by abcd-efgh: %s", err)
@@ -124,7 +123,7 @@ func TestGetLeaseAquired(t *testing.T) {
 	shard := &par.ShardStatus{
 		ID:         "0001",
 		Checkpoint: "deadbeef",
-		Mux:        &sync.Mutex{},
+		Mux:        &sync.RWMutex{},
 	}
 	err := checkpoint.GetLease(shard, "ijkl-mnop")
 
@@ -145,7 +144,7 @@ func TestGetLeaseAquired(t *testing.T) {
 
 	status := &par.ShardStatus{
 		ID:  shard.ID,
-		Mux: &sync.Mutex{},
+		Mux: &sync.RWMutex{},
 	}
 	checkpoint.FetchCheckpoint(status)
 

clientlibrary/config/config.go

@@ -41,6 +41,7 @@ import (
 
 	"github.com/aws/aws-sdk-go/aws"
 	creds "github.com/aws/aws-sdk-go/aws/credentials"
+
 	"github.com/vmware/vmware-go-kcl/clientlibrary/metrics"
 	"github.com/vmware/vmware-go-kcl/logger"
 )
@@ -169,13 +170,24 @@ type (
 		// StreamName is the name of Kinesis stream
 		StreamName string
 
+		// EnableEnhancedFanOutConsumer enables enhanced fan-out consumer
+		// See: https://docs.aws.amazon.com/streams/latest/dev/enhanced-consumers.html
+		// Either consumer name or consumer ARN must be specified when Enhanced Fan-Out is enabled.
+		EnableEnhancedFanOutConsumer bool
+
+		// EnhancedFanOutConsumerName is the name of the enhanced fan-out consumer to create.
+		EnhancedFanOutConsumerName string
+
+		// EnhancedFanOutConsumerARN is the ARN of an already created enhanced fan-out consumer, if this is set no automatic consumer creation will be attempted
+		EnhancedFanOutConsumerARN string
+
 		// WorkerID used to distinguish different workers/processes of a Kinesis application
 		WorkerID string
 
 		// InitialPositionInStream specifies the Position in the stream where a new application should start from
 		InitialPositionInStream InitialPositionInStream
 
-		// InitialPositionInStreamExtended provides actual AT_TMESTAMP value
+		// InitialPositionInStreamExtended provides actual AT_TIMESTAMP value
 		InitialPositionInStreamExtended InitialPositionInStreamExtended
 
 		// credentials to access Kinesis/Dynamo: https://docs.aws.amazon.com/sdk-for-go/api/aws/credentials/
@@ -262,18 +274,18 @@ func empty(s string) bool {
 	return len(strings.TrimSpace(s)) == 0
 }
 
-// checkIsValuePositive make sure the value is possitive.
+// checkIsValueNotEmpty makes sure the value is not empty.
 func checkIsValueNotEmpty(key string, value string) {
 	if empty(value) {
 		// There is no point to continue for incorrect configuration. Fail fast!
-		log.Panicf("Non-empty value exepected for %v, actual: %v", key, value)
+		log.Panicf("Non-empty value expected for %v, actual: %v", key, value)
 	}
 }
 
-// checkIsValuePositive make sure the value is possitive.
+// checkIsValuePositive makes sure the value is possitive.
 func checkIsValuePositive(key string, value int) {
 	if value <= 0 {
 		// There is no point to continue for incorrect configuration. Fail fast!
-		log.Panicf("Positive value exepected for %v, actual: %v", key, value)
+		log.Panicf("Positive value expected for %v, actual: %v", key, value)
 	}
 }

clientlibrary/config/config_test.go

@@ -19,10 +19,11 @@
 package config
 
 import (
-	"github.com/vmware/vmware-go-kcl/logger"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+
+	"github.com/vmware/vmware-go-kcl/logger"
 )
 
 func TestConfig(t *testing.T) {
@@ -32,13 +33,36 @@ func TestConfig(t *testing.T) {
 		WithInitialPositionInStream(TRIM_HORIZON).
 		WithIdleTimeBetweenReadsInMillis(20).
 		WithCallProcessRecordsEvenForEmptyRecordList(true).
-		WithTaskBackoffTimeMillis(10)
+		WithTaskBackoffTimeMillis(10).
+		WithEnhancedFanOutConsumer("fan-out-consumer")
 
 	assert.Equal(t, "appName", kclConfig.ApplicationName)
 	assert.Equal(t, 500, kclConfig.FailoverTimeMillis)
 	assert.Equal(t, 10, kclConfig.TaskBackoffTimeMillis)
+	assert.True(t, kclConfig.EnableEnhancedFanOutConsumer)
+	assert.Equal(t, "fan-out-consumer", kclConfig.EnhancedFanOutConsumerName)
 
 	contextLogger := kclConfig.Logger.WithFields(logger.Fields{"key1": "value1"})
 	contextLogger.Debugf("Starting with default logger")
 	contextLogger.Infof("Default logger is awesome")
 }
+
+func TestEmptyEnhancedFanOutConsumerName(t *testing.T) {
+	assert.PanicsWithValue(t, "Non-empty value expected for EnhancedFanOutConsumerName, actual: ", func() {
+		NewKinesisClientLibConfig("app", "stream", "us-west-2", "worker").WithEnhancedFanOutConsumer("")
+	})
+}
+
+func TestConfigWithEnhancedFanOutConsumerARN(t *testing.T) {
+	kclConfig := NewKinesisClientLibConfig("app", "stream", "us-west-2", "worker").
+		WithEnhancedFanOutConsumerARN("consumer:arn")
+
+	assert.True(t, kclConfig.EnableEnhancedFanOutConsumer)
+	assert.Equal(t, "consumer:arn", kclConfig.EnhancedFanOutConsumerARN)
+}
+
+func TestEmptyEnhancedFanOutConsumerARN(t *testing.T) {
+	assert.PanicsWithValue(t, "Non-empty value expected for EnhancedFanOutConsumerARN, actual: ", func() {
+		NewKinesisClientLibConfig("app", "stream", "us-west-2", "worker").WithEnhancedFanOutConsumerARN("")
+	})
+}

clientlibrary/config/kcl-config.go

@@ -37,9 +37,9 @@ import (
 	"log"
 	"time"
 
-	"github.com/vmware/vmware-go-kcl/clientlibrary/metrics"
-
 	"github.com/aws/aws-sdk-go/aws/credentials"
+
+	"github.com/vmware/vmware-go-kcl/clientlibrary/metrics"
 	"github.com/vmware/vmware-go-kcl/clientlibrary/utils"
 	"github.com/vmware/vmware-go-kcl/logger"
 )
@@ -212,3 +212,23 @@ func (c *KinesisClientLibConfiguration) WithMonitoringService(mService metrics.M
 	c.MonitoringService = mService
 	return c
 }
+
+// WithEnhancedFanOutConsumer enables enhanced fan-out consumer with the specified name
+// For more info see: https://docs.aws.amazon.com/streams/latest/dev/enhanced-consumers.html
+// Note: You can register up to twenty consumers per stream to use enhanced fan-out.
+func (c *KinesisClientLibConfiguration) WithEnhancedFanOutConsumer(consumerName string) *KinesisClientLibConfiguration {
+	checkIsValueNotEmpty("EnhancedFanOutConsumerName", consumerName)
+	c.EnhancedFanOutConsumerName = consumerName
+	c.EnableEnhancedFanOutConsumer = true
+	return c
+}
+
+// WithEnhancedFanOutConsumerARN enables enhanced fan-out consumer with the specified consumer ARN
+// For more info see: https://docs.aws.amazon.com/streams/latest/dev/enhanced-consumers.html
+// Note: You can register up to twenty consumers per stream to use enhanced fan-out.
+func (c *KinesisClientLibConfiguration) WithEnhancedFanOutConsumerARN(consumerARN string) *KinesisClientLibConfiguration {
+	checkIsValueNotEmpty("EnhancedFanOutConsumerARN", consumerARN)
+	c.EnhancedFanOutConsumerARN = consumerARN
+	c.EnableEnhancedFanOutConsumer = true
+	return c
+}

clientlibrary/partition/partition.go

@@ -37,7 +37,7 @@ type ShardStatus struct {
 	ParentShardId string
 	Checkpoint    string
 	AssignedTo    string
-	Mux           *sync.Mutex
+	Mux           *sync.RWMutex
 	LeaseTimeout  time.Time
 	// Shard Range
 	StartingSequenceNumber string
@@ -46,8 +46,8 @@ type ShardStatus struct {
 }
 
 func (ss *ShardStatus) GetLeaseOwner() string {
-	ss.Mux.Lock()
-	defer ss.Mux.Unlock()
+	ss.Mux.RLock()
+	defer ss.Mux.RUnlock()
 	return ss.AssignedTo
 }
 
@@ -56,3 +56,15 @@ func (ss *ShardStatus) SetLeaseOwner(owner string) {
 	defer ss.Mux.Unlock()
 	ss.AssignedTo = owner
 }
+
+func (ss *ShardStatus) GetCheckpoint() string {
+	ss.Mux.RLock()
+	defer ss.Mux.RUnlock()
+	return ss.Checkpoint
+}
+
+func (ss *ShardStatus) SetCheckpoint(c string) {
+	ss.Mux.Lock()
+	defer ss.Mux.Unlock()
+	ss.Checkpoint = c
+}

clientlibrary/utils/awserr.go

@@ -0,0 +1,31 @@
+/*
+ * Copyright (c) 2021 VMware, Inc.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
+ * associated documentation files (the "Software"), to deal in the Software without restriction, including
+ * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is furnished to do
+ * so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all copies or substantial
+ * portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
+ * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
+ * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
+ * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
+ * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ */
+package utils
+
+import (
+	"github.com/aws/aws-sdk-go/aws/awserr"
+)
+
+func AWSErrCode(err error) string {
+	awsErr, _ := err.(awserr.Error)
+	if awsErr != nil {
+		return awsErr.Code()
+	}
+	return ""
+}