Initial pass for adding multiple trace lookup

llinder · Adrian Cole · commit f26f72098b9c · 2018-07-05T18:09:38.000+07:00
diff --git a/zipkin2/src/main/java/zipkin2/storage/InMemoryStorage.java b/zipkin2/src/main/java/zipkin2/storage/InMemoryStorage.java
@@ -307,6 +307,30 @@ public synchronized Call<List<Span>> getTrace(String traceId) {
     return Call.create(filtered);
   }
 
+  @Override
+  public synchronized Call<List<List<Span>>> getTraces(List<String> traceIds) {
+    List<List<Span>> traces = new ArrayList<>();
+    for (String traceId : traceIds) {
+      traceId = Span.normalizeTraceId(traceId);
+      List<Span> spans = spansByTraceId(lowTraceId(traceId));
+      if (spans == null || spans.isEmpty()) continue;
+      if (!strictTraceId) {
+        traces.add(spans);
+      } else {
+        List<Span> filtered = new ArrayList<>(spans);
+        Iterator<Span> iterator = filtered.iterator();
+        while (iterator.hasNext()) {
+          if (!iterator.next().traceId().equals(traceId)) {
+            iterator.remove();
+          }
+        }
+        traces.add(filtered);
+      }
+    }
+
+    return Call.create(traces);
+  }
+
   @Override
   public synchronized Call<List<String>> getServiceNames() {
     if (!searchEnabled) return Call.emptyList();
diff --git a/zipkin2/src/main/java/zipkin2/storage/SpanStore.java b/zipkin2/src/main/java/zipkin2/storage/SpanStore.java
@@ -49,6 +49,19 @@ public interface SpanStore {
    */
   Call<List<Span>> getTrace(String traceId);
 
+  /**
+   * Retrieves list of spans that share a 128-bit trace id with no ordering expectation or empty if none are
+   * found.
+   *
+   * <p>When strict trace ID is disabled, spans with the same right-most 16 characters are returned
+   * even if the characters to the left are not.
+   *
+   * <p>Implementations should use {@link Span#normalizeTraceId(String)} to ensure consistency.
+   *
+   * @param traceIds a list of {@link Span#traceId() trace ID}
+   */
+  Call<List<List<Span>>> getTraces(List<String> traceIds);
+
   /**
    * Retrieves all {@link Span#localEndpoint() local} and {@link Span#remoteEndpoint() remote}
    * {@link Endpoint#serviceName service names}, sorted lexicographically.
diff --git a/zipkin2/src/test/java/zipkin2/storage/InMemoryStorageTest.java b/zipkin2/src/test/java/zipkin2/storage/InMemoryStorageTest.java
@@ -14,9 +14,7 @@
 package zipkin2.storage;
 
 import java.io.IOException;
-import java.util.Collections;
-import java.util.List;
-import java.util.Map;
+import java.util.*;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import org.junit.Test;
@@ -122,4 +120,32 @@ public class InMemoryStorageTest {
       "root"
     );
   }
+
+  @Test public void getTraces_byTraceIds() throws IOException {
+    Span trace1Span1 = Span.newBuilder().traceId("1").id("1").name("root")
+      .localEndpoint(Endpoint.newBuilder().serviceName("app").build())
+      .timestamp(TODAY * 1000)
+      .build();
+    Span trace1Span2 = Span.newBuilder().traceId("1").parentId("1").id("2")
+      .localEndpoint(Endpoint.newBuilder().serviceName("app").build())
+      .timestamp(TODAY * 1000)
+      .build();
+
+    Span trace2Span1 = Span.newBuilder().traceId("2").id("1").name("root")
+      .localEndpoint(Endpoint.newBuilder().serviceName("app").build())
+      .timestamp(TODAY * 1000)
+      .build();
+    Span trace2Span2 = Span.newBuilder().traceId("2").parentId("1").id("2")
+      .localEndpoint(Endpoint.newBuilder().serviceName("app").build())
+      .timestamp(TODAY * 1000)
+      .build();
+
+    storage.accept(asList(trace1Span1, trace1Span2, trace2Span1, trace2Span2));
+
+
+    assertThat(storage.getTraces(asList("1", "2")).execute()).contains(
+      asList(trace1Span1, trace1Span2),
+      asList(trace2Span1, trace2Span2)
+    );
+  }
 }
