More efficient zipkin traces streaming (#16)

* efficient zipkin traces streaming

* VizceralBridge: use new tracing tags, fix tests

Author: mostroverkhov

tracing-openzipkin/src/main/java/io/netifi/proteus/tracing/TracesStreamer.java

@@ -8,6 +8,8 @@
 import com.fasterxml.jackson.databind.deser.std.StdDeserializer;
 import com.fasterxml.jackson.databind.module.SimpleDeserializers;
 import com.hubspot.jackson.datatype.protobuf.ProtobufModule;
+import io.netty.handler.codec.json.JsonObjectDecoder;
+import org.reactivestreams.Publisher;
 import reactor.core.Exceptions;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
@@ -22,49 +24,64 @@
 public class TracesStreamer {
 
   private final ObjectMapper objectMapper = protoMapper();
-  private Function<Integer, Mono<InputStream>> inputSource;
+  private Function<Integer, Publisher<InputStream>> inputSource;
 
   public TracesStreamer(String zipkinUrl,
                         Mono<HttpClient> client) {
-    this.inputSource = zipkinServerStream(zipkinUrl, client);
+    this(zipkinServerStream(zipkinUrl, client));
   }
 
-  public TracesStreamer(Mono<InputStream> inputSource) {
-    this.inputSource = v -> inputSource;
+  public TracesStreamer(Publisher<InputStream> tracesSource) {
+    this(v -> tracesSource);
+  }
+
+  private TracesStreamer(Function<Integer, Publisher<InputStream>> inputSource) {
+    this.inputSource = inputSource;
   }
 
   public Flux<Trace> streamTraces(int lookbackSeconds) {
     return streamTraces(inputSource.apply(lookbackSeconds));
   }
 
-  Flux<Trace> streamTraces(Mono<InputStream> input) {
-    return input.map(is -> {
-      try {
-        return objectMapper.<Traces>readValue(
-            is,
-            new TypeReference<Traces>() {
-            });
-      } catch (IOException e) {
-        throw Exceptions.propagate(e);
-      }
-    }).flatMapIterable(Traces::getTracesList)
-        ;
+  Flux<Trace> streamTraces(Publisher<InputStream> input) {
+    return Flux.from(input)
+        .filter(is -> {
+          try {
+            return is.available() > 0;
+          } catch (IOException e) {
+            throw Exceptions.propagate(e);
+          }
+        })
+        .map(is -> {
+          try {
+            return objectMapper.readValue(
+                is,
+                new TypeReference<Trace>() {
+                });
+          } catch (IOException e) {
+            throw Exceptions.propagate(e);
+          }
+        });
   }
 
-  private static Function<Integer, Mono<InputStream>> zipkinServerStream(String zipkinUrl,
-                                                                         Mono<HttpClient> client) {
+  private static Function<Integer, Publisher<InputStream>> zipkinServerStream(String zipkinUrl,
+                                                                              Mono<HttpClient> client) {
     return lookbackSeconds -> client
-        .flatMap(c -> c
-            .get(zipkinQuery(zipkinUrl, lookbackSeconds))
-            .flatMap(resp ->
+        .flatMapMany(c -> c
+            .get(
+                zipkinQuery(zipkinUrl, lookbackSeconds),
+                req -> {
+                  req.context().addHandler(new JsonObjectDecoder(true));
+                  return Mono.empty();
+                })
+            .flatMapMany(resp ->
                 resp.receive()
-                    .aggregate()
                     .asInputStream()));
   }
 
   private static String zipkinQuery(String zipkinUrl, int lookbackSeconds) {
-    long lookbackMicros = TimeUnit.SECONDS.toMillis(lookbackSeconds);
-    return zipkinUrl + "?lookback=" + lookbackMicros + "&limit=100000";
+    long lookbackMillis = TimeUnit.SECONDS.toMillis(lookbackSeconds);
+    return zipkinUrl + "?lookback=" + lookbackMillis + "&limit=100000";
   }
 
   private ObjectMapper protoMapper() {
@@ -79,12 +96,12 @@ public static class CustomProtoModule extends ProtobufModule {
     public void setupModule(SetupContext context) {
       super.setupModule(context);
       SimpleDeserializers deser = new SimpleDeserializers();
-      deser.addDeserializer(Traces.class, new TracersDeserializer());
+      deser.addDeserializer(Trace.class, new TracersDeserializer());
       context.addDeserializers(deser);
     }
   }
 
-  public static class TracersDeserializer extends StdDeserializer<Traces> {
+  public static class TracersDeserializer extends StdDeserializer<Trace> {
 
     public TracersDeserializer() {
       this(null);
@@ -95,26 +112,12 @@ protected TracersDeserializer(Class<?> vc) {
     }
 
     @Override
-    public Traces deserialize(JsonParser p,
-                              DeserializationContext ctxt) throws IOException {
-      p.nextToken();
-      Traces.Builder tracesBuilder = Traces.newBuilder();
-      while (p.currentToken() != JsonToken.END_ARRAY) {
-        tracesBuilder.addTraces(nextTrace(p, ctxt));
-      }
-      p.nextToken();
-      return tracesBuilder.build();
-    }
-
-    private Trace nextTrace(JsonParser p,
-                            DeserializationContext ctxt) throws IOException {
+    public Trace deserialize(JsonParser p,
+                             DeserializationContext ctx) throws IOException {
       Trace.Builder traceBuilder = Trace.newBuilder();
-      p.nextToken();
-      while (p.currentToken() != JsonToken.END_ARRAY) {
-        traceBuilder.addSpans(ctxt.readValue(p, Span.class));
-        p.nextToken();
+      while (p.nextToken() != JsonToken.END_ARRAY) {
+        traceBuilder.addSpans(ctx.readValue(p, Span.class));
       }
-      p.nextToken();
       return traceBuilder.build();
     }
   }

tracing-openzipkin/src/test/java/io/netifi/proteus/tracing/ZipkinTracesStreamerTest.java

@@ -1,34 +1,93 @@
 package io.netifi.proteus.tracing;
 
+import io.netty.channel.ChannelOption;
 import org.junit.Assert;
-import org.junit.Before;
+import org.junit.Ignore;
 import org.junit.Test;
+import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
+import reactor.ipc.netty.http.client.HttpClient;
+import reactor.ipc.netty.resources.PoolResources;
 
+import java.io.ByteArrayInputStream;
 import java.io.InputStream;
+import java.nio.charset.Charset;
+import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.util.List;
+import java.util.Scanner;
+import java.util.concurrent.TimeUnit;
 
 public class ZipkinTracesStreamerTest {
 
-  private TracesStreamer tracesStreamer;
+  private static final int COUNT = 42;
 
-  @Before
-  public void setUp() {
-    tracesStreamer = new TracesStreamer(zipkinSource());
+  @Test
+  public void zipkinServerTracesStreaming() {
+    TracesStreamer tracesStreamer = new TracesStreamer(zipkinSource(COUNT));
+    List<Trace> traces = tracesStreamer
+        .streamTraces(42)
+        .collectList()
+        .block(Duration.ofSeconds(10));
+    Assert.assertEquals(COUNT, traces.size());
   }
 
   @Test
-  public void zipkinServerTracesStreaming() {
+  public void emptyResponse() {
+    TracesStreamer tracesStreamer = new TracesStreamer(emptySource());
     List<Trace> traces = tracesStreamer
         .streamTraces(42)
         .collectList()
         .block(Duration.ofSeconds(10));
-    Assert.assertFalse(traces.isEmpty());
+    Assert.assertTrue(traces.isEmpty());
+  }
+
+  private Flux<InputStream> zipkinSource(int count) {
+    return Mono.fromCallable(() -> {
+      try (InputStream trace =
+               getClass()
+                   .getClassLoader()
+                   .getResourceAsStream("zipkin_trace.json")) {
+        Charset utf8 = StandardCharsets.UTF_8;
+        try (java.util.Scanner s = new Scanner(trace, utf8.name())) {
+          return s.useDelimiter("\\A").hasNext() ? s.next() : "";
+        }
+      }
+    }).flatMapMany(trace ->
+        Flux.range(1, count)
+            .map(v -> trace)
+            .map(this::asInputStream));
+  }
+
+  private InputStream asInputStream(String s) {
+    return new ByteArrayInputStream(s.getBytes(StandardCharsets.UTF_8));
+  }
+
+  private Flux<InputStream> emptySource() {
+    return Flux.just(asInputStream(""));
+  }
+
+  @Ignore("requires local zipkin server")
+  @Test
+  public void streamerIntegrationTest() {
+    TracesStreamer streamer = new TracesStreamer("/api/v2/traces",
+        client());
+    Flux<Trace> traces = streamer.streamTraces((int) TimeUnit.SECONDS.toSeconds(10));
+    List<Trace> tracesList = traces.collectList().block();
+    Assert.assertFalse(tracesList.isEmpty());
   }
 
-  private Mono<InputStream> zipkinSource() {
-    return Mono.fromCallable(() ->
-        getClass().getClassLoader().getResourceAsStream("zipkin.json"));
+  private Mono<HttpClient> client() {
+    return Mono.just(HttpClient.builder()
+        .options(
+            builder ->
+                builder
+                    .compression(true)
+                    .poolResources(PoolResources.fixed("proteusZipkinBridge"))
+                    .option(ChannelOption.SO_KEEPALIVE, true)
+                    .option(ChannelOption.CONNECT_TIMEOUT_MILLIS, 30_000)
+                    .host("127.0.0.1")
+                    .port(9411))
+        .build());
   }
 }