Add some support types for snapshot orchestration.

This patch adds a number of support types that are used to coordinate execution
of a snapshot. They correspond to the entities defined in
https://www.notion.so/responsivedev/Snapshot-Protocol-1a627549be4c8075a909cfd05adae3af

SnapshotStore defines an interface for a store for Snapshot metadata. Its essentially
a table of all the snapshots taken of an application keyed by generation.

TopicSnapshotStore implements SnapshotStore by writing snapshot metadata to a topic
and using a Kafka transaction to do transactional updates to the current snapshot.

SnapshotOrchestrator defines the interface between Kafka Streams tasks and the
snapshot orchestration process. Its how the tasks discover what the target generation
should be, and how they report their task snapshot metadata (offsets+checkpoints).

LocalSnapshotOrchestrator implements SnapshotOrchestrator for settings where the
orchestrator runs embedded in the application.

SnapshotApi defines the interface for interacting with snapshots from outside the
application.

LocalSnapshotApi implements SnapshotApi for settings where the api interacts directly
with the SnapshotStore (rather than going through another api service like we might
for responsive cloud) - e.g. in settings where we're just orchestrating snapshots
locally by storing metadata in a topic using TopicSnapshotStore.

Author: rodesai

kafka-client-examples/e2e-test/build.gradle.kts

@@ -31,7 +31,7 @@ dependencies {
     implementation(libs.guava)
     implementation(libs.slf4j.log4j2)
     implementation(libs.bundles.scylla)
-    implementation(libs.jackson)
+    implementation(libs.bundles.jackson)
     implementation(libs.mongodb.driver.core)
 
     testImplementation(testlibs.bundles.base)

kafka-client/build.gradle.kts

@@ -137,7 +137,7 @@ dependencies {
     implementation("dev.responsive:controller-api:0.16.0")
     implementation(libs.bundles.scylla)
     implementation(libs.bundles.commons)
-    implementation(libs.jackson)
+    implementation(libs.bundles.jackson)
     implementation(libs.mongodb.driver.sync)
     implementation(libs.bundles.otel)
     implementation(libs.bundles.grpc)

kafka-client/src/main/java/dev/responsive/kafka/internal/snapshot/GenerationStorage.java

@@ -0,0 +1,20 @@
+package dev.responsive.kafka.internal.snapshot;
+
+import org.apache.kafka.streams.processor.TaskId;
+
+/**
+ * Interface that abstracts away how we lookup a task's current generation.
+ * For synchronized snapshots we want to do one of the following so we can record a task's
+ * generation metadata transactionally alongside the rows with new generation markers when
+ * we bump generations:
+ * (1) store this information in the offset metadata as part of the transaction that
+ *     bumps the generation
+ * (2) if/when kafka supports 2pc store this information in another store (like rs3 or
+ *     the snapshot store)
+ *
+ * For simple uncoordinated snapshots we'll support looking up the generation by looking
+ * at the snapshot's state in the generation store.
+ */
+public interface GenerationStorage {
+  long lookupGeneration(final TaskId taskId);
+}

kafka-client/src/main/java/dev/responsive/kafka/internal/snapshot/LocalSnapshotApi.java

@@ -0,0 +1,41 @@
+package dev.responsive.kafka.internal.snapshot;
+
+import java.util.List;
+import java.util.Objects;
+
+/**
+ * Implementation of SnapshotApi that directly interacts with the Snapshot Store rather
+ * than calling out to an API service.
+ */
+public class LocalSnapshotApi implements SnapshotApi {
+  private final SnapshotStore snapshotStore;
+
+  public LocalSnapshotApi(final SnapshotStore snapshotStore) {
+    this.snapshotStore = Objects.requireNonNull(snapshotStore);
+  }
+
+  @Override
+  public Snapshot createSnapshot() {
+    return snapshotStore.updateCurrentSnapshot(snapshot -> {
+      if (snapshot.state().equals(Snapshot.State.CREATED)) {
+        throw new RuntimeException("Snapshot is currently in progress");
+      }
+      return snapshot.nextSnapshot();
+    });
+  }
+
+  @Override
+  public Snapshot getCurrentSnapshot() {
+    return snapshotStore.currentSnapshot(true);
+  }
+
+  @Override
+  public List<Snapshot> getSnapshots() {
+    return snapshotStore.listSnapshots(true);
+  }
+
+  @Override
+  public void close() {
+    snapshotStore.close();
+  }
+}

kafka-client/src/main/java/dev/responsive/kafka/internal/snapshot/LocalSnapshotOrchestrator.java

@@ -0,0 +1,99 @@
+package dev.responsive.kafka.internal.snapshot;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Objects;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import org.apache.kafka.streams.processor.TaskId;
+
+/**
+ * An implementation of the orchestrator that runs within the application and interacts
+ * directly with the snapshot store.
+ */
+public class LocalSnapshotOrchestrator implements SnapshotOrchestrator {
+  private final SnapshotStore snapshotStore;
+  private final Set<TaskId> allTasks;
+
+  public LocalSnapshotOrchestrator(
+      final SnapshotStore snapshotStore,
+      final Set<TaskId> allTasks
+  ) {
+    this.snapshotStore = Objects.requireNonNull(snapshotStore);
+    this.allTasks = Objects.requireNonNull(allTasks);
+  }
+
+  @Override
+  public long getCurrentGeneration() {
+    return snapshotStore.currentSnapshot(false).generation();
+  }
+
+  @Override
+  public void reportTaskSnapshotMetadata(
+      final long generation,
+      final List<Snapshot.TaskSnapshotMetadata> taskSnapshots
+  ) {
+    snapshotStore.updateCurrentSnapshot(snapshot -> {
+      // check that we're still working on this snapshot
+      if (snapshot.generation() != generation) {
+        throw new RuntimeException(
+            String.format("generation too old: %d > %d", snapshot.generation(), generation));
+      }
+      if (!snapshot.state().equals(Snapshot.State.CREATED)) {
+        throw new RuntimeException("Snapshot is currently completed. Cannot update");
+      }
+
+      // check that for all the specified tasks, we either haven't collected its metadata
+      // or the metadata is the same
+      final List<Snapshot.TaskSnapshotMetadata> newlyCompletedTaskSnapshots = new ArrayList<>();
+      for (final var taskSnapshot : taskSnapshots) {
+        final var found = snapshot.taskSnapshots()
+            .stream()
+            .filter(s -> s.taskId().equals(taskSnapshot.taskId()))
+            .collect(Collectors.toList());
+        if (found.size() > 1) {
+          throw new IllegalStateException(
+              "found multiple snapshots for task " + taskSnapshot.taskId());
+        }
+        if (found.isEmpty()) {
+          newlyCompletedTaskSnapshots.add(taskSnapshot);
+        } else if (!found.get(0).equals(taskSnapshot)) {
+          throw new IllegalStateException(
+              "found conflicting snapshots for task" + taskSnapshot.taskId());
+        }
+      }
+
+      // if we've collected snapshots for all tasks, mark the snapshot as completed
+      final Set<TaskId> completedTasks = Stream.concat(
+          newlyCompletedTaskSnapshots.stream().map(Snapshot.TaskSnapshotMetadata::taskId),
+          snapshot.taskSnapshots().stream().map(Snapshot.TaskSnapshotMetadata::taskId)
+      ).collect(Collectors.toSet());
+      Snapshot.State state = snapshot.state();
+      if (completedTasks.equals(this.allTasks)) {
+        state = Snapshot.State.COMPLETED;
+      }
+
+      return snapshot.withTaskSnapshots(newlyCompletedTaskSnapshots, state);
+    });
+  }
+
+  @Override
+  public void failSnapshot(long snapshotGeneration) {
+    snapshotStore.updateCurrentSnapshot(snapshot -> {
+      if (snapshot.generation() != snapshotGeneration) {
+        // todo: do something more reasonable here
+        throw new RuntimeException("Generation mismatch");
+      }
+      if (snapshot.state().equals(Snapshot.State.COMPLETED)) {
+        throw new RuntimeException("Cannot fail completed snapshot");
+      }
+      return snapshot.withStateFailed();
+    });
+  }
+
+  @Override
+  public void close() {
+    snapshotStore.close();
+  }
+}