[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/docker/scripts/run-e2etest

@@ -7,4 +7,6 @@ JMX_OPTS="-Dcom.sun.management.jmxremote=true
 -Djava.rmi.server.hostname=${POD_IP}
 "
 
+
+cp
 java ${JAVA_OPTS} ${EXPORTER_OPTS} ${JMX_OPTS}  -cp "/usr/share/java/responsive-e2etest/*" dev.responsive.examples.e2etest.Main ${CONFIG_PATH}

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

@@ -14,15 +14,16 @@
 
 import static dev.responsive.kafka.api.async.internals.AsyncProcessor.createAsyncFixedKeyProcessor;
 import static dev.responsive.kafka.api.async.internals.AsyncUtils.initializeAsyncBuilders;
+import static dev.responsive.kafka.internal.utils.Utils.isExecutingOnStreamThread;
 
-import dev.responsive.kafka.api.async.internals.AsyncProcessor;
 import dev.responsive.kafka.api.async.internals.stores.AbstractAsyncStoreBuilder;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 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,45 @@ 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() {
-    return createAsyncFixedKeyProcessor(userProcessorSupplier.get(), asyncStoreBuilders);
+  public FixedKeyProcessor<KIn, VIn, VOut> get() {
+    if (maybeInitializeAsyncStoreBuilders()) {
+      return createAsyncFixedKeyProcessor(userProcessorSupplier.get(), asyncStoreBuilders);
+    } else {
+      return userProcessorSupplier.get();
+    }
   }
 
   @Override
   public Set<StoreBuilder<?>> stores() {
-    return new HashSet<>(asyncStoreBuilders.values());
+    if (maybeInitializeAsyncStoreBuilders()) {
+      return new HashSet<>(asyncStoreBuilders.values());
+    } else {
+      return userProcessorSupplier.stores();
+    }
+  }
+
+  private boolean maybeInitializeAsyncStoreBuilders() {
+    if (!isExecutingOnStreamThread()) {
+      // short circuit and delay the actual initialization until we detect that an actual
+      // StreamThread is calling this, since the application will actually run through the
+      // entire topology and build everything once when it first starts up, then throws it away
+      return false;
+    }
+
+    if (asyncStoreBuilders == null) {
+      asyncStoreBuilders = initializeAsyncBuilders(userProcessorSupplier.stores());
+    }
+
+    return true;
   }
 }

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

@@ -15,7 +15,6 @@
 import static dev.responsive.kafka.api.async.internals.AsyncProcessor.createAsyncProcessor;
 import static dev.responsive.kafka.api.async.internals.AsyncUtils.initializeAsyncBuilders;
 
-import dev.responsive.kafka.api.async.internals.AsyncProcessor;
 import dev.responsive.kafka.api.async.internals.stores.AbstractAsyncStoreBuilder;
 import java.util.HashSet;
 import java.util.Map;
@@ -28,6 +27,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 +114,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 +133,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() {
+  public Processor<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());
+    }
+  }
 }