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;
+
+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;
+
+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;
+
+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.
    *
@@ -69,8 +77,8 @@ public KafkaStreams kafkaStreams(Topology topology, Properties streamsConfig) {
    *
    * @see TracingKafkaClientSupplier
    */
-  public <K, V> ProcessorSupplier<K, V> processor(String name, Processor<K, V> processor) {
-    return new TracingProcessorSupplier<>(this, name, processor);
+  public <K, V> ProcessorSupplier<K, V> processor(String spanName, Processor<K, V> processor) {
+    return new TracingProcessorSupplier<>(this, spanName, processor);
   }
 
   /**
@@ -84,9 +92,9 @@ public <K, V> ProcessorSupplier<K, V> processor(String name, Processor<K, V> pro
    *        .to(outputTopic);
    * }</pre>
    */
-  public <K, V, R> TransformerSupplier<K, V, R> transformer(String name,
+  public <K, V, R> TransformerSupplier<K, V, R> transformer(String spanName,
       Transformer<K, V, R> transformer) {
-    return new TracingTransformerSupplier<>(this, name, transformer);
+    return new TracingTransformerSupplier<>(this, spanName, transformer);
   }
 
   /**
@@ -100,9 +108,9 @@ public <K, V, R> TransformerSupplier<K, V, R> transformer(String name,
    *        .to(outputTopic);
    * }</pre>
    */
-  public <V, VR> ValueTransformerSupplier<V, VR> valueTransformer(String name,
+  public <V, VR> ValueTransformerSupplier<V, VR> valueTransformer(String spanName,
       ValueTransformer<V, VR> valueTransformer) {
-    return new TracingValueTransformerSupplier<>(this, name, valueTransformer);
+    return new TracingValueTransformerSupplier<>(this, spanName, valueTransformer);
   }
 
   /**
@@ -116,9 +124,140 @@ public <V, VR> ValueTransformerSupplier<V, VR> valueTransformer(String name,
    *        .to(outputTopic);
    * }</pre>
    */
-  public <K, V, VR> ValueTransformerWithKeySupplier<K, V, VR> valueTransformerWithKey(String name,
+  public <K, V, VR> ValueTransformerWithKeySupplier<K, V, VR> valueTransformerWithKey(String spanName,
       ValueTransformerWithKey<K, V, VR> valueTransformerWithKey) {
-    return new TracingValueTransformerWithKeySupplier<>(this, name, valueTransformerWithKey);
+    return new TracingValueTransformerWithKeySupplier<>(this, spanName, 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(kafkaStreamsTracing.foreach("myForeach", (k, v) -> ...);
+   * }</pre>
+   */
+  public <K, V> ProcessorSupplier<K, V> foreach(String spanName, ForeachAction<K, V> action) {
+    return new TracingProcessorSupplier<>(this, spanName, 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(kafkaStreamsTracing.peek("myPeek", (k, v) -> ...)
+   *        .to(outputTopic);
+   * }</pre>
+   */
+  public <K, V> TransformerSupplier<K, V, KeyValue<K, V>> peek(String spanName, ForeachAction<K, V> action) {
+    return new TracingTransformerSupplier<>(this, spanName, 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 represent an event as part of the stream process.
+   *
+   * A common scenario for this transformer is to mark the beginning and end of a step (or set of steps)
+   * in a stream process.
+   *
+   * <p>Simple example using Kafka Streams DSL:
+   * <pre>{@code
+   * StreamsBuilder builder = new StreamsBuilder();
+   * builder.stream(inputTopic)
+   *        .transform(kafkaStreamsTracing.mark("beginning-complex-map")
+   *        .map(complexTransformation1)
+   *        .filter(predicate)
+   *        .map(complexTransformation2)
+   *        .transform(kafkaStreamsTracing.mark("end-complex-transformation")
+   *        .to(outputTopic);
+   * }</pre>
+   */
+  public <K, V> TransformerSupplier<K, V, KeyValue<K, V>> mark(String spanName) {
+    return new TracingTransformerSupplier<>(this, spanName, 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(kafkaStreamsTracing.map("myMap", (k, v) -> ...)
+   *        .to(outputTopic);
+   * }</pre>
+   */
+  public <K, V, KR, VR> TransformerSupplier<K, V, KeyValue<KR, VR>> map(String spanName,
+      KeyValueMapper<K, V, KeyValue<KR, VR>> mapper) {
+    return new TracingTransformerSupplier<>(this, spanName,
+        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(kafkaStreamsTracing.mapValues("myMapValues", (k, v) -> ...)
+   *        .to(outputTopic);
+   * }</pre>
+   */
+  public <K, V, VR> ValueTransformerWithKeySupplier<K, V, VR> mapValues(String spanName,
+      ValueMapperWithKey<K, V, VR> mapper) {
+    return new TracingValueTransformerWithKeySupplier<>(this, spanName,
+        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(kafkaStreamsTracing.mapValues("myMapValues", v -> ...)
+   *        .to(outputTopic);
+   * }</pre>
+   */
+  public <V, VR> ValueTransformerSupplier<V, VR> mapValues(String spanName, ValueMapper<V, VR> mapper) {
+    return new TracingValueTransformerSupplier<>(this, spanName, new AbstractTracingValueTransformer<V, VR>() {
+      @Override public VR transform(V value) {
+        return mapper.apply(value);
+      }
+    });
   }
 
   static void addTags(ProcessorContext processorContext, SpanCustomizer result) {

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

@@ -9,16 +9,16 @@ class TracingProcessor<K, V> implements Processor<K, V> {
 
   final KafkaStreamsTracing kafkaStreamsTracing;
   final Tracer tracer;
-  final String name;
+  final String spanName;
   final Processor<K, V> delegateProcessor;
 
   ProcessorContext processorContext;
 
   TracingProcessor(KafkaStreamsTracing kafkaStreamsTracing,
-      String name, Processor<K, V> delegateProcessor) {
+      String spanName, Processor<K, V> delegateProcessor) {
     this.kafkaStreamsTracing = kafkaStreamsTracing;
     this.tracer = kafkaStreamsTracing.tracing.tracer();
-    this.name = name;
+    this.spanName = spanName;
     this.delegateProcessor = delegateProcessor;
   }
 
@@ -32,7 +32,7 @@ public void init(ProcessorContext processorContext) {
   public void process(K k, V v) {
     Span span = kafkaStreamsTracing.nextSpan(processorContext);
     if (!span.isNoop()) {
-      span.name(name);
+      span.name(spanName);
       span.start();
     }
 

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

@@ -6,19 +6,19 @@
 class TracingProcessorSupplier<K, V> implements ProcessorSupplier<K, V> {
 
   final KafkaStreamsTracing kafkaStreamsTracing;
-  final String name;
+  final String spanName;
   final Processor<K, V> delegateProcessor;
 
   TracingProcessorSupplier(KafkaStreamsTracing kafkaStreamsTracing,
-      String name,
+      String spanName,
       Processor<K, V> delegateProcessor) {
     this.kafkaStreamsTracing = kafkaStreamsTracing;
-    this.name = name;
+    this.spanName = spanName;
     this.delegateProcessor = delegateProcessor;
   }
 
   /** This wraps process method to enable tracing. */
   @Override public Processor<K, V> get() {
-    return new TracingProcessor<>(kafkaStreamsTracing, name, delegateProcessor);
+    return new TracingProcessor<>(kafkaStreamsTracing, spanName, delegateProcessor);
   }
 }

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

@@ -9,16 +9,16 @@ class TracingTransformer<K, V, R> implements Transformer<K, V, R> {
 
   final KafkaStreamsTracing kafkaStreamsTracing;
   final Tracer tracer;
-  final String name;
+  final String spanName;
   final Transformer<K, V, R> delegateTransformer;
 
   ProcessorContext processorContext;
 
-  TracingTransformer(KafkaStreamsTracing kafkaStreamsTracing, String name,
+  TracingTransformer(KafkaStreamsTracing kafkaStreamsTracing, String spanName,
       Transformer<K, V, R> delegateTransformer) {
     this.kafkaStreamsTracing = kafkaStreamsTracing;
     this.tracer = kafkaStreamsTracing.tracing.tracer();
-    this.name = name;
+    this.spanName = spanName;
     this.delegateTransformer = delegateTransformer;
   }
 
@@ -32,7 +32,7 @@ public void init(ProcessorContext processorContext) {
   public R transform(K k, V v) {
     Span span = kafkaStreamsTracing.nextSpan(processorContext);
     if (!span.isNoop()) {
-      span.name(name);
+      span.name(spanName);
       span.start();
     }
 

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

@@ -5,19 +5,19 @@
 
 class TracingTransformerSupplier<K, V, R> implements TransformerSupplier<K, V, R> {
   final KafkaStreamsTracing kafkaStreamsTracing;
-  final String name;
+  final String spanName;
   final Transformer<K, V, R> delegateTransformer;
 
   TracingTransformerSupplier(KafkaStreamsTracing kafkaStreamsTracing,
-      String name,
+      String spanName,
       Transformer<K, V, R> delegateTransformer) {
     this.kafkaStreamsTracing = kafkaStreamsTracing;
-    this.name = name;
+    this.spanName = spanName;
     this.delegateTransformer = delegateTransformer;
   }
 
   /** This wraps transform method to enable tracing. */
   @Override public Transformer<K, V, R> get() {
-    return new TracingTransformer<>(kafkaStreamsTracing, name, delegateTransformer);
+    return new TracingTransformer<>(kafkaStreamsTracing, spanName, delegateTransformer);
  }
 }

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

@@ -9,16 +9,16 @@ class TracingValueTransformer<V, VR> implements ValueTransformer<V, VR> {
 
   final KafkaStreamsTracing kafkaStreamsTracing;
   final Tracer tracer;
-  final String name;
+  final String spanName;
   final ValueTransformer<V, VR> delegateTransformer;
 
   ProcessorContext processorContext;
 
-  TracingValueTransformer(KafkaStreamsTracing kafkaStreamsTracing, String name,
+  TracingValueTransformer(KafkaStreamsTracing kafkaStreamsTracing, String spanName,
       ValueTransformer<V, VR> delegateTransformer) {
     this.kafkaStreamsTracing = kafkaStreamsTracing;
     this.tracer = kafkaStreamsTracing.tracing.tracer();
-    this.name = name;
+    this.spanName = spanName;
     this.delegateTransformer = delegateTransformer;
   }
 
@@ -32,7 +32,7 @@ public void init(ProcessorContext processorContext) {
   public VR transform(V v) {
     Span span = kafkaStreamsTracing.nextSpan(processorContext);
     if (!span.isNoop()) {
-      span.name(name);
+      span.name(spanName);
       span.start();
     }