From b335246a306db03b55d9287dbb2522d3e0c11169 Mon Sep 17 00:00:00 2001
From: Joshua Kim <kimjos@amazon.com>
Date: Tue, 7 Apr 2020 06:15:18 -0400
Subject: [PATCH] Adding serializer for application state to lease info

---
 .../amazon/kinesis/leases/DynamoUtils.java     | 18 ++++++++++++++++++
 .../dynamodb/DynamoDBLeaseSerializer.java      | 15 +++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/amazon-kinesis-client/src/main/java/software/amazon/kinesis/leases/DynamoUtils.java b/amazon-kinesis-client/src/main/java/software/amazon/kinesis/leases/DynamoUtils.java
index 9d5f9ae2..29d6029b 100644
--- a/amazon-kinesis-client/src/main/java/software/amazon/kinesis/leases/DynamoUtils.java
+++ b/amazon-kinesis-client/src/main/java/software/amazon/kinesis/leases/DynamoUtils.java
@@ -14,6 +14,7 @@
  */
 package software.amazon.kinesis.leases;
 
+import software.amazon.awssdk.core.SdkBytes;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
 import software.amazon.kinesis.annotations.KinesisClientInternalApi;
 
@@ -36,6 +37,14 @@ public class DynamoUtils {
         return AttributeValue.builder().ss(collectionValue).build();
     }
 
+    public static AttributeValue createAttributeValue(byte[] byteBufferValue) {
+        if (byteBufferValue == null) {
+            throw new IllegalArgumentException("Byte buffer attributeValues cannot be null or empty.");
+        }
+
+        return AttributeValue.builder().b(SdkBytes.fromByteArray(byteBufferValue)).build();
+    }
+
     public static AttributeValue createAttributeValue(String stringValue) {
         if (stringValue == null || stringValue.isEmpty()) {
             throw new IllegalArgumentException("String attributeValues cannot be null or empty.");
@@ -52,6 +61,15 @@ public class DynamoUtils {
         return AttributeValue.builder().n(longValue.toString()).build();
     }
 
+    public static byte[] safeGetByteArray(Map<String, AttributeValue> dynamoRecord, String key) {
+        AttributeValue av = dynamoRecord.get(key);
+        if (av == null) {
+            return null;
+        } else {
+            return av.b().asByteArray();
+        }
+    }
+
     public static Long safeGetLong(Map<String, AttributeValue> dynamoRecord, String key) {
         AttributeValue av = dynamoRecord.get(key);
         if (av == null) {
diff --git a/amazon-kinesis-client/src/main/java/software/amazon/kinesis/leases/dynamodb/DynamoDBLeaseSerializer.java b/amazon-kinesis-client/src/main/java/software/amazon/kinesis/leases/dynamodb/DynamoDBLeaseSerializer.java
index a02e2a6e..f42bafcf 100644
--- a/amazon-kinesis-client/src/main/java/software/amazon/kinesis/leases/dynamodb/DynamoDBLeaseSerializer.java
+++ b/amazon-kinesis-client/src/main/java/software/amazon/kinesis/leases/dynamodb/DynamoDBLeaseSerializer.java
@@ -50,6 +50,7 @@ public class DynamoDBLeaseSerializer implements LeaseSerializer {
     private static final String CHECKPOINT_SUBSEQUENCE_NUMBER_KEY = "checkpointSubSequenceNumber";
     private static final String PENDING_CHECKPOINT_SEQUENCE_KEY = "pendingCheckpoint";
     private static final String PENDING_CHECKPOINT_SUBSEQUENCE_KEY = "pendingCheckpointSubSequenceNumber";
+    private static final String PENDING_CHECKPOINT_STATE_KEY = "pendingCheckpointState";
     private static final String PARENT_SHARD_ID_KEY = "parentShardId";
 
     @Override
@@ -75,6 +76,10 @@ public class DynamoDBLeaseSerializer implements LeaseSerializer {
             result.put(PENDING_CHECKPOINT_SUBSEQUENCE_KEY, DynamoUtils.createAttributeValue(lease.pendingCheckpoint().subSequenceNumber()));
         }
 
+        if (lease.pendingCheckpointState() != null) {
+            result.put(PENDING_CHECKPOINT_STATE_KEY, DynamoUtils.createAttributeValue(lease.checkpoint().subSequenceNumber()));
+        }
+
         return result;
     }
 
@@ -105,6 +110,9 @@ public class DynamoDBLeaseSerializer implements LeaseSerializer {
                             DynamoUtils.safeGetLong(dynamoRecord, PENDING_CHECKPOINT_SUBSEQUENCE_KEY))
             );
         }
+
+        leaseToUpdate.pendingCheckpointState(DynamoUtils.safeGetByteArray(dynamoRecord, PENDING_CHECKPOINT_STATE_KEY));
+
         return leaseToUpdate;
     }
 
@@ -220,6 +228,13 @@ public class DynamoDBLeaseSerializer implements LeaseSerializer {
             result.put(PENDING_CHECKPOINT_SEQUENCE_KEY, AttributeValueUpdate.builder().action(AttributeAction.DELETE).build());
             result.put(PENDING_CHECKPOINT_SUBSEQUENCE_KEY, AttributeValueUpdate.builder().action(AttributeAction.DELETE).build());
         }
+
+        if (lease.pendingCheckpointState() != null) {
+            result.put(PENDING_CHECKPOINT_STATE_KEY, putUpdate(DynamoUtils.createAttributeValue(lease.pendingCheckpointState())));
+        } else {
+            result.put(PENDING_CHECKPOINT_STATE_KEY, AttributeValueUpdate.builder().action(AttributeAction.DELETE).build());
+        }
+
         return result;
     }