[DRAFT] Automatic async wrapping for stateful processors

kafka-client-bootstrap/build.gradle.kts

@@ -15,6 +15,10 @@ plugins {
     id("responsive.docker")
 }
 
+repositories {
+    mavenLocal()
+}
+
 application {
     mainClass.set("dev.responsive.kafka.bootstrap.main.Main")
 }

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

@@ -15,6 +15,10 @@ plugins {
     id("responsive.docker")
 }
 
+repositories {
+    mavenLocal()
+}
+
 application {
     mainClass.set("dev.responsive.examples.e2etest.Main")
 }

kafka-client-examples/e2e-test/src/main/java/dev/responsive/examples/regression/tests/KeyBatchExample.java

@@ -32,9 +32,10 @@
 import org.apache.kafka.streams.Topology;
 import org.apache.kafka.streams.kstream.Consumed;
 import org.apache.kafka.streams.kstream.Produced;
-import org.apache.kafka.streams.kstream.Transformer;
-import org.apache.kafka.streams.processor.ProcessorContext;
 import org.apache.kafka.streams.processor.PunctuationType;
+import org.apache.kafka.streams.processor.api.Processor;
+import org.apache.kafka.streams.processor.api.ProcessorContext;
+import org.apache.kafka.streams.processor.api.Record;
 import org.apache.kafka.streams.state.KeyValueIterator;
 import org.apache.kafka.streams.state.KeyValueStore;
 import org.apache.kafka.streams.state.Stores;
@@ -70,7 +71,7 @@ protected Topology buildTopology() {
     }
 
     builder.stream(ORDERS, Consumed.with(Serdes.String(), RegressionSchema.orderSerde()))
-        .transform(BatchTransformer::new, "grouped-orders-store")
+        .process(BatchTransformer::new, "grouped-orders-store")
         .peek((k, v) -> {
           if (responsive) {
             final var random = Math.abs(randomGenerator.nextLong() % 10000);
@@ -84,14 +85,13 @@ protected Topology buildTopology() {
     return builder.build();
   }
 
-  private static class BatchTransformer
-      implements Transformer<String, Order, KeyValue<String, GroupedOrder>> {
+  private static class BatchTransformer implements Processor<String, Order, String, GroupedOrder> {
 
-    private ProcessorContext context;
+    private ProcessorContext<String, GroupedOrder> context;
     private KeyValueStore<String, StoredOrder> store;
 
     @Override
-    public void init(final ProcessorContext context) {
+    public void init(final ProcessorContext<String, GroupedOrder> context) {
       this.context = context;
       this.store = context.getStateStore("grouped-orders-store");
       this.context.schedule(
@@ -102,15 +102,18 @@ public void init(final ProcessorContext context) {
     }
 
     @Override
-    public KeyValue<String, GroupedOrder> transform(final String key, final Order value) {
-      final long ts = context.timestamp();
+    public void process(final Record<String, Order> record) {
+      final long ts = record.timestamp();
 
       // first add the order to the list of orders that are stored
-      store.put(storedKey(key, ts), new StoredOrder(Optional.of(value), Optional.empty()));
+      store.put(
+          storedKey(record.key(), ts),
+          new StoredOrder(Optional.of(record.value()), Optional.empty())
+      );
 
       // next, we need to update the tracked metadata row to
       // check whether the value ought to be emitted
-      final String mKey = metaKey(key);
+      final String mKey = metaKey(record.key());
       final StoredOrder.Meta meta = Optional.ofNullable(store.get(mKey))
           .orElse(new StoredOrder(Optional.empty(), Optional.of(new StoredOrder.Meta(ts, 0, 0))))
           .meta()
@@ -122,17 +125,15 @@ public KeyValue<String, GroupedOrder> transform(final String key, final Order va
       final StoredOrder.Meta newMeta = new StoredOrder.Meta(
           ts,
           meta.count() + 1,
-          meta.size() + (long) value.amount()
+          meta.size() + (long) record.value().amount()
       );
 
       if (shouldFlush(newMeta, ts)) {
-        doFlush(key);
+        doFlush(record.key());
         store.delete(mKey);
       } else {
         store.put(mKey, new StoredOrder(Optional.empty(), Optional.of(newMeta)));
       }
-
-      return null;
     }
 
     private void flushExpired(long ts) {
@@ -174,7 +175,7 @@ private void doFlush(final String key) {
                       "Got stored order with no order! %s".formatted(value))));
         }
 
-        context.forward(key, result);
+        context.forward(new Record<>(key, result, 0L));
       }
     }
 

kafka-client-examples/simple-example/build.gradle.kts

@@ -19,12 +19,16 @@ application {
     mainClass.set("dev.responsive.examples.simpleapp.Main")
 }
 
+repositories {
+    mavenLocal()
+}
+
 dependencies {
     // todo: how to set the version here?
     implementation(project(":kafka-client"))
     implementation("com.google.guava:guava:32.1.1-jre")
-    implementation("org.apache.kafka:kafka-clients:3.4.0")
-    implementation("org.apache.kafka:kafka-streams:3.4.0")
+    implementation(libs.kafka.clients)
+    implementation(libs.kafka.streams)
     implementation("io.opentelemetry.javaagent:opentelemetry-javaagent:1.25.0")
     implementation("org.apache.logging.log4j:log4j-slf4j-impl:2.20.0")
     implementation("org.apache.commons:commons-text:1.10.0")

kafka-client/build.gradle.kts

@@ -33,6 +33,10 @@ plugins {
     id("java")
 }
 
+repositories {
+    mavenLocal()
+}
+
 /*********** Generated Resources ***********/
 
 val gitCommitId: String by lazy {

kafka-client/src/main/java/dev/responsive/kafka/api/ResponsiveKafkaStreams.java

@@ -284,16 +284,16 @@ private static Properties propsWithOverrides(
       return propsWithOverrides;
     }
 
-    final Object o = configs.originals().get(InternalConfig.INTERNAL_TASK_ASSIGNOR_CLASS);
+    final Object o = configs.originals().get(StreamsConfig.TASK_ASSIGNOR_CLASS_CONFIG);
     if (o == null) {
       propsWithOverrides.put(
-          InternalConfig.INTERNAL_TASK_ASSIGNOR_CLASS,
+          StreamsConfig.TASK_ASSIGNOR_CLASS_CONFIG,
           TASK_ASSIGNOR_CLASS_OVERRIDE
       );
     } else if (!TASK_ASSIGNOR_CLASS_OVERRIDE.equals(o.toString())) {
       final String errorMsg = String.format(
           "Invalid Streams configuration value for '%s': got %s, expected '%s'",
-          InternalConfig.INTERNAL_TASK_ASSIGNOR_CLASS,
+          StreamsConfig.TASK_ASSIGNOR_CLASS_CONFIG,
           o,
           TASK_ASSIGNOR_CLASS_OVERRIDE
       );

kafka-client/src/main/java/dev/responsive/kafka/api/async/AsyncFixedKeyProcessorSupplier.java

@@ -23,6 +23,7 @@
 import org.apache.kafka.streams.processor.api.FixedKeyProcessor;
 import org.apache.kafka.streams.processor.api.FixedKeyProcessorSupplier;
 import org.apache.kafka.streams.processor.api.ProcessorSupplier;
+import org.apache.kafka.streams.processor.api.WrappedFixedKeyProcessorSupplier;
 import org.apache.kafka.streams.state.StoreBuilder;
 
 /**
@@ -34,10 +35,10 @@
  * documentation on the async processing framework.
  */
 public class AsyncFixedKeyProcessorSupplier<KIn, VIn, VOut>
-    implements FixedKeyProcessorSupplier<KIn, VIn, VOut> {
+    implements WrappedFixedKeyProcessorSupplier<KIn, VIn, VOut> {
 
   private final FixedKeyProcessorSupplier<KIn, VIn, VOut> userProcessorSupplier;
-  private final Map<String, AbstractAsyncStoreBuilder<?, ?, ?>> asyncStoreBuilders;
+  private Map<String, AbstractAsyncStoreBuilder<?>> asyncStoreBuilders = null;
 
   /**
    * Create an AsyncProcessorSupplier that wraps a custom {@link ProcessorSupplier}
@@ -52,24 +53,31 @@ public class AsyncFixedKeyProcessorSupplier<KIn, VIn, VOut>
   public static <KIn, VIn, VOut> AsyncFixedKeyProcessorSupplier<KIn, VIn, VOut> createAsyncProcessorSupplier(
       final FixedKeyProcessorSupplier<KIn, VIn, VOut> processorSupplier
   ) {
-    return new AsyncFixedKeyProcessorSupplier<>(processorSupplier, processorSupplier.stores());
+    return new AsyncFixedKeyProcessorSupplier<>(processorSupplier);
   }
 
   private AsyncFixedKeyProcessorSupplier(
-      final FixedKeyProcessorSupplier<KIn, VIn, VOut> userProcessorSupplier,
-      final Set<StoreBuilder<?>> userStoreBuilders
+      final FixedKeyProcessorSupplier<KIn, VIn, VOut> userProcessorSupplier
   ) {
     this.userProcessorSupplier = userProcessorSupplier;
-    this.asyncStoreBuilders = initializeAsyncBuilders(userStoreBuilders);
   }
 
   @Override
   public AsyncProcessor<KIn, VIn, KIn, VOut> get() {
+    maybeInitializeAsyncStoreBuilders();
+
     return createAsyncFixedKeyProcessor(userProcessorSupplier.get(), asyncStoreBuilders);
   }
 
   @Override
   public Set<StoreBuilder<?>> stores() {
+    maybeInitializeAsyncStoreBuilders();
     return new HashSet<>(asyncStoreBuilders.values());
   }
+
+  private void maybeInitializeAsyncStoreBuilders() {
+    if (asyncStoreBuilders == null) {
+      asyncStoreBuilders = initializeAsyncBuilders(userProcessorSupplier.stores());
+    }
+  }
 }

kafka-client/src/main/java/dev/responsive/kafka/api/async/AsyncProcessorSupplier.java

@@ -28,6 +28,7 @@
 import org.apache.kafka.streams.processor.api.Processor;
 import org.apache.kafka.streams.processor.api.ProcessorContext;
 import org.apache.kafka.streams.processor.api.ProcessorSupplier;
+import org.apache.kafka.streams.processor.api.WrappedProcessorSupplier;
 import org.apache.kafka.streams.state.KeyValueStore;
 import org.apache.kafka.streams.state.StoreBuilder;
 
@@ -114,10 +115,11 @@
  *    {@link Processor} and requires at least one state store be connected.
  */
 public final class AsyncProcessorSupplier<KIn, VIn, KOut, VOut>
-    implements ProcessorSupplier<KIn, VIn, KOut, VOut> {
+    implements WrappedProcessorSupplier<KIn, VIn, KOut, VOut> {
 
   private final ProcessorSupplier<KIn, VIn, KOut, VOut> userProcessorSupplier;
-  private final Map<String, AbstractAsyncStoreBuilder<?, ?, ?>> asyncStoreBuilders;
+
+  private Map<String, AbstractAsyncStoreBuilder<?>> asyncStoreBuilders = null;
 
   /**
    * Create an AsyncProcessorSupplier that wraps a custom {@link ProcessorSupplier}
@@ -132,32 +134,30 @@ public final class AsyncProcessorSupplier<KIn, VIn, KOut, VOut>
   public static <KIn, VIn, KOut, VOut> AsyncProcessorSupplier<KIn, VIn, KOut, VOut> createAsyncProcessorSupplier(
       final ProcessorSupplier<KIn, VIn, KOut, VOut> processorSupplier
   ) {
-    return new AsyncProcessorSupplier<>(processorSupplier, processorSupplier.stores());
+    return new AsyncProcessorSupplier<>(processorSupplier);
   }
 
   private AsyncProcessorSupplier(
-      final ProcessorSupplier<KIn, VIn, KOut, VOut> userProcessorSupplier,
-      final Set<StoreBuilder<?>> userStoreBuilders
+      final ProcessorSupplier<KIn, VIn, KOut, VOut> userProcessorSupplier
   ) {
-    if (userStoreBuilders == null || userStoreBuilders.isEmpty()) {
-      throw new UnsupportedOperationException(
-          "Async processing currently requires at least one state store be "
-              + "connected to the async processor, and that stores be connected "
-              + "by implementing the #stores method in your processor supplier");
-    }
-
     this.userProcessorSupplier = userProcessorSupplier;
-    this.asyncStoreBuilders = initializeAsyncBuilders(userStoreBuilders);
   }
 
   @Override
   public AsyncProcessor<KIn, VIn, KOut, VOut> get() {
+    maybeInitializeAsyncStoreBuilders();
     return createAsyncProcessor(userProcessorSupplier.get(), asyncStoreBuilders);
   }
 
   @Override
   public Set<StoreBuilder<?>> stores() {
+    maybeInitializeAsyncStoreBuilders();
     return new HashSet<>(asyncStoreBuilders.values());
   }
 
+  private void maybeInitializeAsyncStoreBuilders() {
+    if (asyncStoreBuilders == null) {
+      asyncStoreBuilders = initializeAsyncBuilders(userProcessorSupplier.stores());
+    }
+  }
 }

kafka-client/src/main/java/dev/responsive/kafka/api/async/AsyncProcessorWrapper.java

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2024 Responsive Computing, Inc.
+ *
+ * This source code is licensed under the Responsive Business Source License Agreement v1.0
+ * available at:
+ *
+ * https://www.responsive.dev/legal/responsive-bsl-10
+ *
+ * This software requires a valid Commercial License Key for production use. Trial and commercial
+ * licenses can be obtained at https://www.responsive.dev
+ */
+
+package dev.responsive.kafka.api.async;
+
+import org.apache.kafka.streams.processor.api.FixedKeyProcessorSupplier;
+import org.apache.kafka.streams.processor.api.ProcessorSupplier;
+import org.apache.kafka.streams.processor.api.ProcessorWrapper;
+import org.apache.kafka.streams.processor.api.WrappedFixedKeyProcessorSupplier;
+import org.apache.kafka.streams.processor.api.WrappedProcessorSupplier;
+
+// TODO(sophie): only wrap when stores are Responsive (rather than eg rocksdb or in-memory)
+public class AsyncProcessorWrapper implements ProcessorWrapper {
+
+  @Override
+  public <KIn, VIn, KOut, VOut> WrappedProcessorSupplier<KIn, VIn, KOut, VOut> wrapProcessorSupplier(
+      final String processorName,
+      final ProcessorSupplier<KIn, VIn, KOut, VOut> processorSupplier
+  ) {
+    final var stores = processorSupplier.stores();
+    if (stores != null && !stores.isEmpty()) {
+      return AsyncProcessorSupplier.createAsyncProcessorSupplier(processorSupplier);
+    } else {
+      return ProcessorWrapper.asWrapped(processorSupplier);
+    }
+  }
+
+  @Override
+  public <KIn, VIn, VOut> WrappedFixedKeyProcessorSupplier<KIn, VIn, VOut> wrapFixedKeyProcessorSupplier(
+      final String processorName,
+      final FixedKeyProcessorSupplier<KIn, VIn, VOut> processorSupplier
+  ) {
+    final var stores = processorSupplier.stores();
+    if (stores != null && !stores.isEmpty()) {
+      return AsyncFixedKeyProcessorSupplier.createAsyncProcessorSupplier(processorSupplier);
+    } else {
+      return ProcessorWrapper.asWrappedFixedKey(processorSupplier);
+    }
+
+  }
+}

kafka-client/src/main/java/dev/responsive/kafka/api/async/internals/AsyncProcessor.java

@@ -71,7 +71,7 @@ public class AsyncProcessor<KIn, VIn, KOut, VOut>
   private final Processor<KIn, VIn, KOut, VOut> userProcessor;
   private final FixedKeyProcessor<KIn, VIn, VOut> userFixedKeyProcessor;
 
-  private final Map<String, AbstractAsyncStoreBuilder<?, ?, ?>> connectedStoreBuilders;
+  private final Map<String, AbstractAsyncStoreBuilder<?>> connectedStoreBuilders;
 
   // Tracks all pending events, ie from moment of creation to end of finalization/transition
   // to "DONE" state. Used to make sure all events are flushed during a commit.
@@ -115,14 +115,14 @@ public class AsyncProcessor<KIn, VIn, KOut, VOut>
 
   public static <KIn, VIn, KOut, VOut> AsyncProcessor<KIn, VIn, KOut, VOut> createAsyncProcessor(
       final Processor<KIn, VIn, KOut, VOut> userProcessor,
-      final Map<String, AbstractAsyncStoreBuilder<?, ?, ?>> connectedStoreBuilders
+      final Map<String, AbstractAsyncStoreBuilder<?>> connectedStoreBuilders
   ) {
     return new AsyncProcessor<>(userProcessor, null, connectedStoreBuilders);
   }
 
   public static <KIn, VIn, VOut> AsyncProcessor<KIn, VIn, KIn, VOut> createAsyncFixedKeyProcessor(
       final FixedKeyProcessor<KIn, VIn, VOut> userProcessor,
-      final Map<String, AbstractAsyncStoreBuilder<?, ?, ?>> connectedStoreBuilders
+      final Map<String, AbstractAsyncStoreBuilder<?>> connectedStoreBuilders
   ) {
     return new AsyncProcessor<>(null, userProcessor, connectedStoreBuilders);
   }
@@ -135,7 +135,7 @@ public static <KIn, VIn, VOut> AsyncProcessor<KIn, VIn, KIn, VOut> createAsyncFi
   private AsyncProcessor(
       final Processor<KIn, VIn, KOut, VOut> userProcessor,
       final FixedKeyProcessor<KIn, VIn, VOut> userFixedKeyProcessor,
-      final Map<String, AbstractAsyncStoreBuilder<?, ?, ?>> connectedStoreBuilders
+      final Map<String, AbstractAsyncStoreBuilder<?>> connectedStoreBuilders
   ) {
     this.userProcessor = userProcessor;
     this.userFixedKeyProcessor = userFixedKeyProcessor;
@@ -377,10 +377,10 @@ public void close() {
   private static void registerFlushListenerForStoreBuilders(
       final String streamThreadName,
       final int partition,
-      final Collection<AbstractAsyncStoreBuilder<?, ?, ?>> asyncStoreBuilders,
+      final Collection<AbstractAsyncStoreBuilder<?>> asyncStoreBuilders,
       final AsyncFlushListener flushPendingEvents
   ) {
-    for (final AbstractAsyncStoreBuilder<?, ?, ?> builder : asyncStoreBuilders) {
+    for (final AbstractAsyncStoreBuilder<?> builder : asyncStoreBuilders) {
       builder.registerFlushListenerWithAsyncStore(streamThreadName, partition, flushPendingEvents);
     }
   }
@@ -738,7 +738,7 @@ private boolean isCleared() {
    */
   private void verifyConnectedStateStores(
       final Map<String, AsyncKeyValueStore<?, ?>> accessedStores,
-      final Map<String, AbstractAsyncStoreBuilder<?, ?, ?>> connectedStores
+      final Map<String, AbstractAsyncStoreBuilder<?>> connectedStores
   ) {
     if (!accessedStores.keySet().equals(connectedStores.keySet())) {
       log.error(