Fix the DynamoDB metadata table upsert operation (#2769)

Author: KodaiD

core/src/main/java/com/scalar/db/storage/dynamo/DynamoAdmin.java

@@ -97,6 +97,7 @@ public class DynamoAdmin implements DistributedStorageAdmin {
   public static final String DEFAULT_NO_BACKUP = "false";
   public static final String DEFAULT_REQUEST_UNIT = "10";
   private static final int DEFAULT_WAITING_DURATION_SECS = 3;
+  @VisibleForTesting static final int MAX_RETRY_COUNT = 10;
 
   @VisibleForTesting static final String PARTITION_KEY = "concatenatedPartitionKey";
   @VisibleForTesting static final String CLUSTERING_KEY = "concatenatedClusteringKey";
@@ -238,15 +239,26 @@ public void createNamespace(String nonPrefixedNamespace, Map<String, String> opt
   private void upsertIntoNamespacesTable(Namespace namespace) throws ExecutionException {
     Map<String, AttributeValue> itemValues = new HashMap<>();
     itemValues.put(NAMESPACES_ATTR_NAME, AttributeValue.builder().s(namespace.prefixed()).build());
-    try {
-      client.putItem(
-          PutItemRequest.builder()
-              .tableName(ScalarDbUtils.getFullTableName(metadataNamespace, NAMESPACES_TABLE))
-              .item(itemValues)
-              .build());
-    } catch (Exception e) {
-      throw new ExecutionException(
-          "Inserting the " + namespace + " namespace into the namespaces table failed", e);
+    int retryCount = 0;
+    while (true) {
+      try {
+        client.putItem(
+            PutItemRequest.builder()
+                .tableName(ScalarDbUtils.getFullTableName(metadataNamespace, NAMESPACES_TABLE))
+                .item(itemValues)
+                .build());
+        return;
+      } catch (ResourceNotFoundException e) {
+        if (retryCount >= MAX_RETRY_COUNT) {
+          throw new ExecutionException(
+              "Inserting the " + namespace + " namespace into the namespaces table failed", e);
+        }
+        Uninterruptibles.sleepUninterruptibly(waitingDurationSecs, TimeUnit.SECONDS);
+        retryCount++;
+      } catch (Exception e) {
+        throw new ExecutionException(
+            "Inserting the " + namespace + " namespace into the namespaces table failed", e);
+      }
     }
   }
 
@@ -470,15 +482,28 @@ private void upsertTableMetadata(Namespace namespace, String table, TableMetadat
           METADATA_ATTR_SECONDARY_INDEX,
           AttributeValue.builder().ss(metadata.getSecondaryIndexNames()).build());
     }
-    try {
-      client.putItem(
-          PutItemRequest.builder()
-              .tableName(ScalarDbUtils.getFullTableName(metadataNamespace, METADATA_TABLE))
-              .item(itemValues)
-              .build());
-    } catch (Exception e) {
-      throw new ExecutionException(
-          "Adding the metadata for the " + getFullTableName(namespace, table) + " table failed", e);
+    int retryCount = 0;
+    while (true) {
+      try {
+        client.putItem(
+            PutItemRequest.builder()
+                .tableName(ScalarDbUtils.getFullTableName(metadataNamespace, METADATA_TABLE))
+                .item(itemValues)
+                .build());
+        return;
+      } catch (ResourceNotFoundException e) {
+        if (retryCount >= MAX_RETRY_COUNT) {
+          throw new ExecutionException(
+              "Adding the metadata for the " + getFullTableName(namespace, table) + " table failed",
+              e);
+        }
+        Uninterruptibles.sleepUninterruptibly(waitingDurationSecs, TimeUnit.SECONDS);
+        retryCount++;
+      } catch (Exception e) {
+        throw new ExecutionException(
+            "Adding the metadata for the " + getFullTableName(namespace, table) + " table failed",
+            e);
+      }
     }
   }
 

core/src/test/java/com/scalar/db/storage/dynamo/DynamoAdminTestBase.java

@@ -703,6 +703,38 @@ public void createTable_WhenMetadataTableExists_ShouldCreateOnlyTable()
         .isInstanceOf(IllegalArgumentException.class);
   }
 
+  @Test
+  public void createTable_WhenActualMetadataTableCreationIsDelayed_ShouldFailAfterRetries() {
+    // Arrange
+    when(client.describeTable(any(DescribeTableRequest.class))).thenReturn(tableIsActiveResponse);
+    when(client.describeContinuousBackups(any(DescribeContinuousBackupsRequest.class)))
+        .thenReturn(backupIsEnabledResponse);
+    when(client.putItem(any(PutItemRequest.class))).thenThrow(ResourceNotFoundException.class);
+    TableMetadata metadata =
+        TableMetadata.newBuilder()
+            .addPartitionKey("c1")
+            .addClusteringKey("c2", Order.DESC)
+            .addClusteringKey("c3", Order.ASC)
+            .addColumn("c1", DataType.TEXT)
+            .addColumn("c2", DataType.BIGINT)
+            .addColumn("c3", DataType.BOOLEAN)
+            .addColumn("c4", DataType.INT)
+            .addColumn("c5", DataType.BLOB)
+            .addColumn("c6", DataType.DOUBLE)
+            .addColumn("c7", DataType.FLOAT)
+            .addColumn("c8", DataType.DATE)
+            .addColumn("c9", DataType.TIME)
+            .addColumn("c10", DataType.TIMESTAMP)
+            .addColumn("c11", DataType.TIMESTAMPTZ)
+            .addSecondaryIndex("c4")
+            .build();
+
+    // Act Assert
+    assertThatThrownBy(() -> admin.createTable(NAMESPACE, TABLE, metadata))
+        .isInstanceOf(ExecutionException.class);
+    verify(client, times(DynamoAdmin.MAX_RETRY_COUNT + 1)).putItem(any(PutItemRequest.class));
+  }
+
   @Test
   public void dropTable_WithNoMetadataLeft_ShouldDropTableAndDeleteMetadata()
       throws ExecutionException {
@@ -1436,6 +1468,17 @@ public void createNamespace_WithExistingNamespacesTable_ShouldAddNamespace()
                 .build());
   }
 
+  @Test
+  public void createNamespace_WhenActualNamespaceTableCreationIsDelayed_ShouldFailAfterRetries() {
+    // Arrange
+    when(client.putItem(any(PutItemRequest.class))).thenThrow(ResourceNotFoundException.class);
+
+    // Act Assert
+    assertThatThrownBy(() -> admin.createNamespace(NAMESPACE, Collections.emptyMap()))
+        .isInstanceOf(ExecutionException.class);
+    verify(client, times(DynamoAdmin.MAX_RETRY_COUNT + 1)).putItem(any(PutItemRequest.class));
+  }
+
   @Test
   public void namespaceExists_WithExistingNamespace_ShouldReturnTrue() throws ExecutionException {
     // Arrange