Async processor (#70)

Author: avinashkolluru

kafka-streams-framework/src/main/java/org/hypertrace/core/kafkastreams/framework/async/AsyncProcessor.java

@@ -1,9 +1,99 @@
 package org.hypertrace.core.kafkastreams.framework.async;
 
+import com.google.common.util.concurrent.RateLimiter;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ArrayBlockingQueue;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executor;
+import java.util.function.Supplier;
+import lombok.SneakyThrows;
+import lombok.extern.slf4j.Slf4j;
 import org.apache.kafka.streams.processor.api.Processor;
+import org.apache.kafka.streams.processor.api.ProcessorContext;
 import org.apache.kafka.streams.processor.api.Record;
 
-public class AsyncProcessor implements Processor {
+/**
+ * Async version of the {@link Processor} interface. Offloads the entire processing to a different
+ * executor(thread pool), collects the records to forward to the next stage of the topology and
+ * flushes them after the processing is complete
+ *
+ * @param <K> the type of input keys
+ * @param <V> the type of input values
+ * @param <KOUT> the type of output keys
+ * @param <VOUT> the type of output values
+ */
+@Slf4j
+@SuppressWarnings("UnstableApiUsage")
+public abstract class AsyncProcessor<K, V, KOUT, VOUT> implements Processor<K, V, KOUT, VOUT> {
+
+  private final Executor executor;
+  private final BlockingQueue<CompletableFuture<List<RecordToForward<KOUT, VOUT>>>> pendingFutures;
+  private final RateLimiter rateLimiter;
+  private ProcessorContext<KOUT, VOUT> context;
+
+  public AsyncProcessor(
+      Supplier<Executor> executorSupplier, AsyncProcessorConfig asyncProcessorConfig) {
+    this.executor = executorSupplier.get();
+    this.pendingFutures = new ArrayBlockingQueue<>(asyncProcessorConfig.getMaxBatchSize());
+    this.rateLimiter =
+        RateLimiter.create(1000.0 / asyncProcessorConfig.getCommitIntervalMs().toMillis());
+    log.info(
+        "async processor config. maxBatchSize: {}, commit rate: {}",
+        this.pendingFutures.remainingCapacity(),
+        this.rateLimiter.getRate());
+    // warmup to prevent commit on first message
+    rateLimiter.tryAcquire();
+  }
+
   @Override
-  public void process(Record record) {}
+  public final void init(ProcessorContext<KOUT, VOUT> context) {
+    this.context = context;
+    doInit(context.appConfigs());
+  }
+
+  protected abstract void doInit(Map<String, Object> appConfigs);
+
+  public abstract List<RecordToForward<KOUT, VOUT>> asyncProcess(K key, V value);
+
+  @SneakyThrows
+  @Override
+  public void process(Record<K, V> record) {
+    CompletableFuture<List<RecordToForward<KOUT, VOUT>>> future =
+        CompletableFuture.supplyAsync(() -> asyncProcess(record.key(), record.value()), executor);
+    // with put, thread gets blocked when queue is full. queue consumer runs in this same thread.
+    pendingFutures.put(future);
+
+    // Flush based on size or time - whichever occurs first
+    // once the queue is full, flush the queue.
+    if (pendingFutures.remainingCapacity() == 0 || rateLimiter.tryAcquire()) {
+      log.debug("flush start - queue size before flush: {}", pendingFutures.size());
+      processResults();
+      log.debug("flush end - queue size after flush: {}", pendingFutures.size());
+    }
+  }
+
+  @SneakyThrows
+  private void processResults() {
+    while (!pendingFutures.isEmpty()) {
+      CompletableFuture<List<RecordToForward<KOUT, VOUT>>> future = pendingFutures.poll();
+      // makes sure processing is complete
+      future.join();
+      // another join is needed to make sure downstream forward is also complete
+      future
+          .thenAccept(
+              result -> {
+                if (result != null) {
+                  result.forEach(
+                      recordToForward ->
+                          context.forward(
+                              recordToForward.getRecord(), recordToForward.getChildName()));
+                }
+              })
+          .join();
+    }
+    // commit once per batch
+    context.commit();
+  }
 }

kafka-streams-framework/src/main/java/org/hypertrace/core/kafkastreams/framework/async/AsyncProcessorConfig.java

@@ -0,0 +1,42 @@
+package org.hypertrace.core.kafkastreams.framework.async;
+
+import static org.hypertrace.core.kafkastreams.framework.async.Constants.DEFAULT_ASYNC_PROCESSOR_BATCH_SIZE;
+import static org.hypertrace.core.kafkastreams.framework.async.Constants.DEFAULT_ASYNC_PROCESSOR_COMMIT_INTERVAL;
+
+import com.typesafe.config.Config;
+import java.time.Duration;
+import lombok.Getter;
+
+@Getter
+public class AsyncProcessorConfig {
+  private static final String COMMIT_INTERVAL_CONFIG_KEY = "commitIntervalMs";
+  private static final String MAX_BATCH_SIZE_CONFIG_KEY = "maxBatchSize";
+  private static final String PROCESSORS_CONFIG_KEY = "async.processors";
+  private final int maxBatchSize;
+  private final Duration commitIntervalMs;
+
+  AsyncProcessorConfig(int maxBatchSize, int commitIntervalMs) {
+    this.maxBatchSize = maxBatchSize;
+    this.commitIntervalMs = Duration.ofMillis(commitIntervalMs);
+  }
+
+  public static AsyncProcessorConfig buildWith(Config config, String processorName) {
+    Config processorsConfig =
+        config.hasPath(PROCESSORS_CONFIG_KEY) ? config.getConfig(PROCESSORS_CONFIG_KEY) : null;
+    if (processorsConfig != null && processorsConfig.hasPath(processorName)) {
+      Config processorConfig = processorsConfig.getConfig(processorName);
+      int batchSize =
+          processorConfig.hasPath(MAX_BATCH_SIZE_CONFIG_KEY)
+              ? processorConfig.getInt(MAX_BATCH_SIZE_CONFIG_KEY)
+              : DEFAULT_ASYNC_PROCESSOR_BATCH_SIZE;
+      int commitInterval =
+          processorConfig.hasPath(COMMIT_INTERVAL_CONFIG_KEY)
+              ? processorConfig.getInt(COMMIT_INTERVAL_CONFIG_KEY)
+              : DEFAULT_ASYNC_PROCESSOR_COMMIT_INTERVAL;
+      return new AsyncProcessorConfig(batchSize, commitInterval);
+    }
+
+    return new AsyncProcessorConfig(
+        DEFAULT_ASYNC_PROCESSOR_BATCH_SIZE, DEFAULT_ASYNC_PROCESSOR_COMMIT_INTERVAL);
+  }
+}

kafka-streams-framework/src/main/java/org/hypertrace/core/kafkastreams/framework/async/AsyncTransformer.java

@@ -2,7 +2,7 @@
 
 public class Constants {
   public static int DEFAULT_ASYNC_EXECUTOR_POOL_SIZE = 16;
-  public static int DEFAULT_ASYNC_TRANSFORMER_BATCH_SIZE = 5120;
-  public static int DEFAULT_ASYNC_TRANSFORMER_COMMIT_INTERVAL = 5000;
+  public static int DEFAULT_ASYNC_PROCESSOR_BATCH_SIZE = 5120;
+  public static int DEFAULT_ASYNC_PROCESSOR_COMMIT_INTERVAL = 5000;
   public static String ASYNC_EXECUTOR_POOL_SIZE_KEY = "async.executors.maxPoolSize";
 }

kafka-streams-framework/src/main/java/org/hypertrace/core/kafkastreams/framework/async/AsyncTransformerConfig.java

@@ -0,0 +1,20 @@
+package org.hypertrace.core.kafkastreams.framework.async;
+
+import javax.annotation.Nullable;
+import lombok.AllArgsConstructor;
+import lombok.Getter;
+import org.apache.kafka.streams.processor.api.Record;
+
+/**
+ * Encapsulates the {@link Record record} and the {@link String childName} used by the {@link
+ * AsyncProcessor} to forward messages to the downstream processor
+ *
+ * @param <K> type of the record key
+ * @param <V> type of the record value
+ */
+@AllArgsConstructor
+@Getter
+public class RecordToForward<K, V> {
+  @Nullable private String childName;
+  private Record<K, V> record;
+}

kafka-streams-framework/src/main/java/org/hypertrace/core/kafkastreams/framework/async/Constants.java

@@ -7,14 +7,14 @@
 import java.util.function.Supplier;
 import lombok.SneakyThrows;
 import lombok.extern.slf4j.Slf4j;
-import org.apache.kafka.streams.KeyValue;
 import org.apache.kafka.streams.StreamsBuilder;
 import org.apache.kafka.streams.kstream.KStream;
 import org.apache.kafka.streams.processor.api.Processor;
 import org.apache.kafka.streams.processor.api.Record;
-import org.hypertrace.core.kafkastreams.framework.async.AsyncTransformer;
-import org.hypertrace.core.kafkastreams.framework.async.AsyncTransformerConfig;
+import org.hypertrace.core.kafkastreams.framework.async.AsyncProcessor;
+import org.hypertrace.core.kafkastreams.framework.async.AsyncProcessorConfig;
 import org.hypertrace.core.kafkastreams.framework.async.ExecutorFactory;
+import org.hypertrace.core.kafkastreams.framework.async.RecordToForward;
 import org.hypertrace.core.kafkastreams.framework.constants.KafkaStreamsAppConstants;
 import org.hypertrace.core.serviceframework.config.ConfigClient;
 
@@ -42,11 +42,11 @@ public StreamsBuilder buildTopology(
 
     Config kafkaStreamsConfig = configClient.getConfig().getConfig(KAFKA_STREAMS_CONFIG_KEY);
     KStream<String, String> transform =
-        stream.transform(
+        stream.process(
             () ->
-                new SlowTransformer(
+                new SlowProcessor(
                     ExecutorFactory.getExecutorSupplier(kafkaStreamsConfig),
-                    AsyncTransformerConfig.buildWith(kafkaStreamsConfig, "slow.transformer")));
+                    AsyncProcessorConfig.buildWith(kafkaStreamsConfig, "slow.processor")));
     transform.process(LoggingProcessor::new);
     transform.to(OUTPUT_TOPIC);
     return streamsBuilder;
@@ -64,11 +64,10 @@ public String getServiceName() {
 }
 
 @Slf4j
-class SlowTransformer extends AsyncTransformer<String, String, String, String> {
-
-  public SlowTransformer(
-      Supplier<Executor> executorSupplier, AsyncTransformerConfig asyncTransformerConfig) {
-    super(executorSupplier, asyncTransformerConfig);
+class SlowProcessor extends AsyncProcessor<String, String, String, String> {
+  public SlowProcessor(
+      Supplier<Executor> executorSupplier, AsyncProcessorConfig asyncProcessorConfig) {
+    super(executorSupplier, asyncProcessorConfig);
   }
 
   @Override
@@ -78,13 +77,15 @@ protected void doInit(Map<String, Object> appConfigs) {
 
   @SneakyThrows
   @Override
-  public List<KeyValue<String, String>> asyncTransform(String key, String value) {
+  public List<RecordToForward<String, String>> asyncProcess(String key, String value) {
     if (!key.startsWith("key")) {
       return null;
     }
-    log.info("transforming - key: {}, value: {}", key, value);
+    log.info("processing - key: {}, value: {}", key, value);
     Thread.sleep(25);
-    return List.of(KeyValue.pair("out:" + key, "out:" + value));
+    return List.of(
+        new RecordToForward<>(
+            null, new Record<>("out:" + key, "out:" + value, System.currentTimeMillis())));
   }
 }
 

kafka-streams-framework/src/main/java/org/hypertrace/core/kafkastreams/framework/async/RecordToForward.java

@@ -10,8 +10,8 @@ kafka.streams.config = {
     executors {
       maxPoolSize = 16
     }
-    transformers {
-      slow.transformer = {
+    processors {
+      slow.processor = {
         maxBatchSize = 32
         commitIntervalMs = 100
       }