Kafka Streams helper operations

instrumentation/kafka-streams/src/main/java/brave/kafka/streams/AbstractTracingTransformer.java

@@ -0,0 +1,14 @@
+package brave.kafka.streams;
+
+import org.apache.kafka.streams.kstream.Transformer;
+import org.apache.kafka.streams.processor.ProcessorContext;
+
+public abstract class AbstractTracingTransformer<K, V, R> implements
+    Transformer<K, V, R> {
+
+  @Override public void init(ProcessorContext context) {
+  }
+
+  @Override public void close() {
+  }
+}

instrumentation/kafka-streams/src/main/java/brave/kafka/streams/AbstractTracingValueTransformer.java

@@ -0,0 +1,14 @@
+package brave.kafka.streams;
+
+import org.apache.kafka.streams.kstream.ValueTransformer;
+import org.apache.kafka.streams.processor.ProcessorContext;
+
+public abstract class AbstractTracingValueTransformer<V, VR> implements
+    ValueTransformer<V, VR> {
+
+  @Override public void init(ProcessorContext context) {
+  }
+
+  @Override public void close() {
+  }
+}

instrumentation/kafka-streams/src/main/java/brave/kafka/streams/AbstractTracingValueTransformerWithKey.java

@@ -0,0 +1,14 @@
+package brave.kafka.streams;
+
+import org.apache.kafka.streams.kstream.ValueTransformerWithKey;
+import org.apache.kafka.streams.processor.ProcessorContext;
+
+public abstract class AbstractTracingValueTransformerWithKey<K, V, VR> implements
+    ValueTransformerWithKey<K, V, VR> {
+
+  @Override public void init(ProcessorContext context) {
+  }
+
+  @Override public void close() {
+  }
+}

instrumentation/kafka-streams/src/main/java/brave/kafka/streams/KafkaStreamsTracing.java

@@ -7,16 +7,24 @@
 import brave.propagation.TraceContext;
 import brave.propagation.TraceContextOrSamplingFlags;
 import java.util.Properties;
+import java.util.function.BiConsumer;
 import org.apache.kafka.common.header.Headers;
 import org.apache.kafka.streams.KafkaClientSupplier;
 import org.apache.kafka.streams.KafkaStreams;
+import org.apache.kafka.streams.KeyValue;
 import org.apache.kafka.streams.Topology;
+import org.apache.kafka.streams.kstream.ForeachAction;
+import org.apache.kafka.streams.kstream.KStream;
+import org.apache.kafka.streams.kstream.KeyValueMapper;
 import org.apache.kafka.streams.kstream.Transformer;
 import org.apache.kafka.streams.kstream.TransformerSupplier;
+import org.apache.kafka.streams.kstream.ValueMapper;
+import org.apache.kafka.streams.kstream.ValueMapperWithKey;
 import org.apache.kafka.streams.kstream.ValueTransformer;
 import org.apache.kafka.streams.kstream.ValueTransformerSupplier;
 import org.apache.kafka.streams.kstream.ValueTransformerWithKey;
 import org.apache.kafka.streams.kstream.ValueTransformerWithKeySupplier;
+import org.apache.kafka.streams.processor.AbstractProcessor;
 import org.apache.kafka.streams.processor.Processor;
 import org.apache.kafka.streams.processor.ProcessorContext;
 import org.apache.kafka.streams.processor.ProcessorSupplier;
@@ -38,8 +46,8 @@ public static KafkaStreamsTracing create(Tracing tracing) {
 
   /**
    * Creates a {@link KafkaStreams} instance with a tracing-enabled {@link KafkaClientSupplier}. All
-   * Topology Sources and Sinks (including internal Topics) will create Spans on records
-   * processed (i.e. send or consumed).
+   * Topology Sources and Sinks (including internal Topics) will create Spans on records processed
+   * (i.e. send or consumed).
    *
    * Use this instead of {@link KafkaStreams} constructor.
    *
@@ -121,6 +129,132 @@ public <K, V, VR> ValueTransformerWithKeySupplier<K, V, VR> valueTransformerWith
     return new TracingValueTransformerWithKeySupplier<>(this, name, valueTransformerWithKey);
   }
 
+  /**
+   * Create a foreach processor, similar to {@link KStream#foreach(ForeachAction)}, where its action
+   * will be recorded in a new span with the indicated name.
+   *
+   * <p>Simple example using Kafka Streams DSL:
+   * <pre>{@code
+   * StreamsBuilder builder = new StreamsBuilder();
+   * builder.stream(inputTopic)
+   *        .process(kafkaStraemsTracing.foreach("myForeach", (k, v) -> ...);
+   * }</pre>
+   */
+  public <K, V> ProcessorSupplier<K, V> foreach(String name, ForeachAction<K, V> action) {
+    return new TracingProcessorSupplier<>(this, name, new AbstractProcessor<K, V>() {
+      @Override public void process(K key, V value) {
+        action.apply(key, value);
+      }
+    });
+  }
+
+  /**
+   * Create a peek transformer, similar to {@link KStream#peek(ForeachAction)}, where its action
+   * will be recorded in a new span with the indicated name.
+   *
+   * <p>Simple example using Kafka Streams DSL:
+   * <pre>{@code
+   * StreamsBuilder builder = new StreamsBuilder();
+   * builder.stream(inputTopic)
+   *        .transform(kafkaStraemsTracing.peek("myPeek", (k, v) -> ...)
+   *        .to(outputTopic);
+   * }</pre>
+   */
+  public <K, V> TransformerSupplier<K, V, KeyValue<K, V>> peek(String name, ForeachAction<K, V> action) {
+    return new TracingTransformerSupplier<>(this, name, new AbstractTracingTransformer<K, V, KeyValue<K, V>>() {
+      @Override public KeyValue<K, V> transform(K key, V value) {
+        action.apply(key, value);
+        return KeyValue.pair(key, value);
+      }
+    });
+  }
+
+  /**
+   * Create a mark transformer, similar to {@link KStream#peek(ForeachAction)}, but no action is executed.
+   * Instead only a span is created to mark the beginning or end of a task.
+   *
+   * <p>Simple example using Kafka Streams DSL:
+   * <pre>{@code
+   * StreamsBuilder builder = new StreamsBuilder();
+   * builder.stream(inputTopic)
+   *        .transform(kafkaStraemsTracing.mark("beginning")
+   *        .map(mapper)
+   *        .transform(kafkaStreamsTracing.mark("end")
+   *        .to(outputTopic);
+   * }</pre>
+   */
+  public <K, V> TransformerSupplier<K, V, KeyValue<K, V>> mark(String name) {
+    return new TracingTransformerSupplier<>(this, name, new AbstractTracingTransformer<K, V, KeyValue<K, V>>() {
+      @Override public KeyValue<K, V> transform(K key, V value) {
+        return KeyValue.pair(key, value);
+      }
+    });
+  }
+
+  /**
+   * Create a map transformer, similar to {@link KStream#map(KeyValueMapper)}, where its mapper action
+   * will be recorded in a new span with the indicated name.
+   *
+   * <p>Simple example using Kafka Streams DSL:
+   * <pre>{@code
+   * StreamsBuilder builder = new StreamsBuilder();
+   * builder.stream(inputTopic)
+   *        .transform(kafkaStraemsTracing.map("myMap", (k, v) -> ...)
+   *        .to(outputTopic);
+   * }</pre>
+   */
+  public <K, V, KR, VR> TransformerSupplier<K, V, KeyValue<KR, VR>> map(String name,
+      KeyValueMapper<K, V, KeyValue<KR, VR>> mapper) {
+    return new TracingTransformerSupplier<>(this, name,
+        new AbstractTracingTransformer<K, V, KeyValue<KR, VR>>() {
+          @Override public KeyValue<KR, VR> transform(K key, V value) {
+            return mapper.apply(key, value);
+          }
+        });
+  }
+
+  /**
+   * Create a peek transformer, similar to {@link KStream#mapValues(ValueMapperWithKey)}, where its mapper action
+   * will be recorded in a new span with the indicated name.
+   *
+   * <p>Simple example using Kafka Streams DSL:
+   * <pre>{@code
+   * StreamsBuilder builder = new StreamsBuilder();
+   * builder.stream(inputTopic)
+   *        .transform(kafkaStraemsTracing.mapValues("myMapValues", (k, v) -> ...)
+   *        .to(outputTopic);
+   * }</pre>
+   */
+  public <K, V, VR> ValueTransformerWithKeySupplier<K, V, VR> mapValues(String name,
+      ValueMapperWithKey<K, V, VR> mapper) {
+    return new TracingValueTransformerWithKeySupplier<>(this, name,
+        new AbstractTracingValueTransformerWithKey<K, V, VR>() {
+          @Override public VR transform(K readOnlyKey, V value) {
+            return mapper.apply(readOnlyKey, value);
+          }
+        });
+  }
+
+  /**
+   * Create a peek transformer, similar to {@link KStream#mapValues(ValueMapper)}, where its mapper action
+   * will be recorded in a new span with the indicated name.
+   *
+   * <p>Simple example using Kafka Streams DSL:
+   * <pre>{@code
+   * StreamsBuilder builder = new StreamsBuilder();
+   * builder.stream(inputTopic)
+   *        .transform(kafkaStraemsTracing.mapValues("myMapValues", v -> ...)
+   *        .to(outputTopic);
+   * }</pre>
+   */
+  public <V, VR> ValueTransformerSupplier<V, VR> mapValues(String name, ValueMapper<V, VR> mapper) {
+    return new TracingValueTransformerSupplier<>(this, name, new AbstractTracingValueTransformer<V, VR>() {
+      @Override public VR transform(V value) {
+        return mapper.apply(value);
+      }
+    });
+  }
+
   static void addTags(ProcessorContext processorContext, SpanCustomizer result) {
     result.tag(KafkaStreamsTags.KAFKA_STREAMS_APPLICATION_ID_TAG, processorContext.applicationId());
     result.tag(KafkaStreamsTags.KAFKA_STREAMS_TASK_ID_TAG, processorContext.taskId().toString());