Fix Dynamo records not uploading after DynamoDbPublisher close() is called

src/main/java/com/example/submission/FooPojoSparkPublisherApp.java

@@ -8,6 +8,7 @@
 import java.util.List;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.TimeUnit;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
@@ -27,6 +28,7 @@ public final class FooPojoSparkPublisherApp {
   private FooPojoSparkPublisherApp() {
   }
 
+  public static ConcurrentLinkedQueue<FooPojo> allFooPojos = new ConcurrentLinkedQueue<>();
   public static void main(String[] args) {
     int partitions = Integer.parseInt(System.getProperty("foo.partitions", "50"));
     int numRecords = Integer.parseInt(System.getenv().getOrDefault("RECORDS_TO_WRITE", "10000000"));
@@ -39,7 +41,6 @@ public static void main(String[] args) {
         .coalesce(partitions)
         .mapPartitions((MapPartitionsFunction<Long, FooPojo>) iterator -> {
 
-
           List<FooPojo> fooPojos = new ArrayList<>();
           // 1. Create FooPojos.
           for (int i = 0; i < recordsPerPartition; i++) {
@@ -49,15 +50,13 @@ public static void main(String[] args) {
 
             // 2. Non-blocking publish (unless queue full, then block).
             fooPojos.add(fooPojo);
+            allFooPojos.add(fooPojo);
             DynamoPublisherCreator.getInstance().publish(fooPojo);
           }
 
           return fooPojos.iterator();
-
         }, Encoders.bean(FooPojo.class))
-
         // 3. Await all async publishers.
-
         .mapPartitions( (MapPartitionsFunction<FooPojo, FooPojo>) iterator -> {
           try {
             CompletableFuture.allOf(DynamoPublisherCreator.getInstance().close())
@@ -66,11 +65,13 @@ public static void main(String[] args) {
           } catch (Exception e) {
             throw new RuntimeException(e);
           }
-        }, Encoders.bean(FooPojo.class) )
+        }, Encoders.bean(FooPojo.class))
 
         // 4. Write everything out to an s3 file or local
         .write()
         .json("fooPojos" + UUID.randomUUID());
 
+    System.out.println("Size of Foo Pojos: " + allFooPojos.size());
+    System.out.println(SPARK);
   }
 }

src/main/java/com/example/submission/dynamo/DynamoPublisher.java

@@ -43,7 +43,8 @@ public class DynamoPublisher<T>  {
   private final BlockingDeque<T> recordQueue;
   private final Semaphore semaphore;
   private final int maxInFlight;
-  private final CompletableFuture<Void> publisherTask;
+  private volatile CompletableFuture<Void> publisherTask;
+  private volatile boolean closed = false;
 
   private static final java.util.concurrent.ExecutorService HANDOFF_EXECUTOR =
       Executors.newFixedThreadPool(1, r -> {
@@ -77,22 +78,30 @@ public DynamoPublisher(
     this.recordQueue = new LinkedBlockingDeque<>(maxQueueSize);
     this.maxInFlight = Math.max(1, maxInFlight);
     this.semaphore = new Semaphore(this.maxInFlight);
-    this.publisherTask = CompletableFuture.runAsync(this::publisherLoop, HANDOFF_EXECUTOR);
+    startPublisherTaskIfNeeded();
   }
 
   /**
    * Enqueues a record for asynchronous publishing. Blocks when the queue is full.
    */
   @SneakyThrows
   public void publish(@NonNull T record) {
+    startPublisherTaskIfNeeded();
     recordQueue.put(record);
   }
 
   public CompletableFuture<Void> close() {
     shutdownRequested = true;
-    publisherTask.join();
-    dynamoDbAsyncClient.close();
-    return publisherTask;
+    CompletableFuture<Void> task = publisherTask;
+    if (task != null) {
+      task.join();
+    }
+    // Do NOT close the Dynamo client here; allow restarting the loop later.
+    return task == null ? CompletableFuture.completedFuture(null) : task;
+  }
+
+  public boolean isClosed() {
+    return closed;
   }
 
   @SneakyThrows
@@ -110,6 +119,19 @@ private void publisherLoop() {
       semaphore.acquire();
       createBatchWriteFuture(request, 0);
     }
+    // Intentionally sleep after draining to help reproduce stall scenarios in tests.
+    try {
+      Thread.sleep(Duration.ofSeconds(1).toMillis());
+    } catch (InterruptedException ie) {
+      Thread.currentThread().interrupt();
+    }
+  }
+
+  private synchronized void startPublisherTaskIfNeeded() {
+    if (publisherTask == null || publisherTask.isDone()) {
+      shutdownRequested = false;
+      publisherTask = CompletableFuture.runAsync(this::publisherLoop, HANDOFF_EXECUTOR);
+    }
   }
 
   @SneakyThrows

src/main/java/com/example/submission/samples/FooPojoTableCreator.java

@@ -0,0 +1,105 @@
+package com.example.submission.samples;
+
+import com.example.submission.pojo.FooPojo;
+import java.util.Objects;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClient;
+import software.amazon.awssdk.enhanced.dynamodb.DynamoDbTable;
+import software.amazon.awssdk.enhanced.dynamodb.TableSchema;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
+import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition;
+import software.amazon.awssdk.services.dynamodb.model.BillingMode;
+import software.amazon.awssdk.services.dynamodb.model.CreateTableRequest;
+import software.amazon.awssdk.services.dynamodb.model.DescribeTableRequest;
+import software.amazon.awssdk.services.dynamodb.model.KeySchemaElement;
+import software.amazon.awssdk.services.dynamodb.model.KeyType;
+import software.amazon.awssdk.services.dynamodb.model.ResourceNotFoundException;
+import software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType;
+
+/**
+ * Small helper to create a DynamoDB table for {@link FooPojo} with partition key {@code pojoId}.
+ *
+ * Usage:
+ *   ./gradlew run --args="--create-table --table FooPojoDummy --region us-east-1"
+ * or run via the IntelliJ run configuration added for convenience.
+ */
+public final class FooPojoTableCreator {
+
+  private static final Logger LOGGER = LogManager.getLogger(FooPojoTableCreator.class);
+
+  private FooPojoTableCreator() {}
+
+  public static void main(String[] args) {
+    Config cfg = Config.from(args);
+
+    try (DynamoDbClient dynamo = DynamoDbClient.builder()
+        .region(cfg.region)
+        .build()) {
+
+      if (tableExists(dynamo, cfg.tableName)) {
+        LOGGER.info("Table '{}' already exists in region {}", cfg.tableName, cfg.region.id());
+        return;
+      }
+
+      // Use on-demand billing for simplicity.
+      CreateTableRequest request = CreateTableRequest.builder()
+          .tableName(cfg.tableName)
+          .billingMode(BillingMode.PAY_PER_REQUEST)
+          .keySchema(KeySchemaElement.builder()
+              .attributeName("pojoId")
+              .keyType(KeyType.HASH)
+              .build())
+          .attributeDefinitions(AttributeDefinition.builder()
+              .attributeName("pojoId")
+              .attributeType(ScalarAttributeType.S)
+              .build())
+          .build();
+
+      dynamo.createTable(request);
+      LOGGER.info("Created table '{}' (billing=PAY_PER_REQUEST)", cfg.tableName);
+
+      // Optional: validate mapping using the enhanced client
+      DynamoDbEnhancedClient enhanced = DynamoDbEnhancedClient.builder().dynamoDbClient(dynamo).build();
+      DynamoDbTable<FooPojo> table = enhanced.table(cfg.tableName, TableSchema.fromBean(FooPojo.class));
+      LOGGER.info("Validated FooPojo schema mapping for table '{}': {}", cfg.tableName, table.tableName());
+    }
+  }
+
+  private static boolean tableExists(DynamoDbClient dynamo, String table) {
+    try {
+      dynamo.describeTable(DescribeTableRequest.builder().tableName(table).build());
+      return true;
+    } catch (ResourceNotFoundException rnfe) {
+      return false;
+    }
+  }
+
+  private record Config(String tableName, Region region) {
+    private static Config from(String[] args) {
+      String table = null;
+      Region region = null;
+      for (int i = 0; i < args.length; i++) {
+        switch (args[i]) {
+          case "--table" -> table = readValue(args, ++i, "--table");
+          case "--region" -> region = Region.of(readValue(args, ++i, "--region"));
+          default -> {}
+        }
+      }
+      String resolvedTable = Objects.requireNonNullElseGet(
+          table, () -> System.getenv().getOrDefault("FOO_TABLE_NAME", "FooPojoDummy"));
+      Region resolvedRegion = Objects.requireNonNullElseGet(
+          region, () -> Region.of(System.getenv().getOrDefault("AWS_REGION", "us-east-1")));
+      return new Config(resolvedTable, resolvedRegion);
+    }
+
+    private static String readValue(String[] args, int idx, String flag) {
+      if (idx >= args.length) {
+        throw new IllegalArgumentException("Missing value after " + flag);
+      }
+      return args[idx];
+    }
+  }
+}
+