From 96ff7027d335467f1f601d68c41f50f235753d37 Mon Sep 17 00:00:00 2001
From: Adrian Cole <adrian@tetrate.io>
Date: Tue, 9 Jan 2024 14:22:10 +0800
Subject: [PATCH] brave: allows AsyncZipkinSpanHandler to be build with custom
 formats

This is needed for stackdriver and any future OTLP work, as well as
developing a standard proto encoder.

Signed-off-by: Adrian Cole <adrian@tetrate.io>
---
 .../brave/AsyncZipkinSpanHandler.java         | 35 +++++++++++++++++--
 1 file changed, 32 insertions(+), 3 deletions(-)

diff --git a/brave/src/main/java/zipkin2/reporter/brave/AsyncZipkinSpanHandler.java b/brave/src/main/java/zipkin2/reporter/brave/AsyncZipkinSpanHandler.java
index d9402b56..c52d7d20 100644
--- a/brave/src/main/java/zipkin2/reporter/brave/AsyncZipkinSpanHandler.java
+++ b/brave/src/main/java/zipkin2/reporter/brave/AsyncZipkinSpanHandler.java
@@ -21,6 +21,8 @@
 import java.io.Flushable;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
+import zipkin2.reporter.BytesEncoder;
+import zipkin2.reporter.Encoding;
 import zipkin2.reporter.Reporter;
 import zipkin2.reporter.ReporterMetrics;
 import zipkin2.reporter.Sender;
@@ -69,15 +71,18 @@ public Builder toBuilder() {
   /** @since 2.14 */
   public static final class Builder extends ZipkinSpanHandler.Builder {
     final AsyncReporter.Builder delegate;
+    final Encoding encoding;
 
     Builder(AsyncZipkinSpanHandler handler) {
       this.delegate = ((AsyncReporter<MutableSpan>) handler.spanReporter).toBuilder();
+      this.encoding = handler.encoding;
       this.alwaysReportSpans = handler.alwaysReportSpans;
       this.errorTag = handler.errorTag;
     }
 
     Builder(Sender sender) {
       this.delegate = AsyncReporter.newBuilder(sender);
+      this.encoding = sender.encoding();
     }
 
     /**
@@ -151,18 +156,42 @@ public Builder queuedMaxBytes(int queuedMaxBytes) {
       return (Builder) super.alwaysReportSpans(alwaysReportSpans);
     }
 
+    /**
+     * Builds an async span handler that encodes zipkin spans according to the sender's encoding.
+     */
     // AsyncZipkinSpanHandler not SpanHandler, so that Flushable and Closeable are accessible
     public AsyncZipkinSpanHandler build() {
-      return new AsyncZipkinSpanHandler(this);
+      switch (encoding) {
+        case JSON:
+          return build(new JsonV2Encoder(errorTag));
+        default:
+          throw new UnsupportedOperationException(encoding.name());
+      }
+    }
+
+    /**
+     * Builds an async span handler that encodes zipkin spans according to the encoder.
+     *
+     * <p>Note: The input encoder must use the same error tag implementation as configured by
+     * {@link #errorTag(Tag)}.
+     *
+     * @since 3.1
+     */
+    // AsyncZipkinSpanHandler not SpanHandler, so that Flushable and Closeable are accessible
+    public AsyncZipkinSpanHandler build(BytesEncoder<MutableSpan> encoder) {
+      if (encoder == null) throw new NullPointerException("encoder == null");
+      return new AsyncZipkinSpanHandler(delegate.build(encoder), this);
     }
   }
 
   final Reporter<MutableSpan> spanReporter;
+  final Encoding encoding;
   final Tag<Throwable> errorTag; // for toBuilder()
   final boolean alwaysReportSpans;
 
-  AsyncZipkinSpanHandler(Builder builder) {
-    this.spanReporter = builder.delegate.build(new JsonV2Encoder(builder.errorTag));
+  AsyncZipkinSpanHandler(AsyncReporter<MutableSpan> spanReporter, Builder builder) {
+    this.spanReporter = spanReporter;
+    this.encoding = builder.encoding;
     this.errorTag = builder.errorTag;
     this.alwaysReportSpans = builder.alwaysReportSpans;
   }