Spikes ES support

Author: Adrian Cole

zipkin-storage/elasticsearch-http/src/main/java/zipkin/storage/elasticsearch/http/ElasticsearchHttpSpanConsumer.java

@@ -13,25 +13,19 @@
  */
 package zipkin.storage.elasticsearch.http;
 
-import com.squareup.moshi.JsonWriter;
-import java.io.IOException;
-import java.util.LinkedHashMap;
-import java.util.LinkedHashSet;
 import java.util.List;
-import java.util.Map;
-import java.util.Set;
 import java.util.concurrent.TimeUnit;
-import okio.Buffer;
-import zipkin.Codec;
 import zipkin.Span;
-import zipkin.internal.Pair;
+import zipkin.internal.Span2;
+import zipkin.internal.Span2Codec;
+import zipkin.internal.Span2Converter;
 import zipkin.storage.AsyncSpanConsumer;
 import zipkin.storage.Callback;
 
 import static zipkin.internal.ApplyTimestampAndDuration.guessTimestamp;
 import static zipkin.internal.Util.UTF_8;
 import static zipkin.internal.Util.propagateIfFatal;
-import static zipkin.storage.elasticsearch.http.ElasticsearchHttpSpanStore.SERVICE_SPAN;
+import static zipkin.storage.elasticsearch.http.ElasticsearchHttpSpanStore.SPAN;
 
 class ElasticsearchHttpSpanConsumer implements AsyncSpanConsumer { // not final for testing
 
@@ -50,27 +44,22 @@ class ElasticsearchHttpSpanConsumer implements AsyncSpanConsumer { // not final
     }
     try {
       HttpBulkIndexer indexer = new HttpBulkIndexer("index-span", es);
-      Map<String, Set<Pair<String>>> indexToServiceSpans = indexSpans(indexer, spans);
-      if (!indexToServiceSpans.isEmpty()) {
-        indexNames(indexer, indexToServiceSpans);
-      }
+      indexSpans(indexer, spans);
       indexer.execute(callback);
     } catch (Throwable t) {
       propagateIfFatal(t);
       callback.onError(t);
     }
   }
 
-  /** Indexes spans and returns a mapping of indexes that may need a names update */
-  Map<String, Set<Pair<String>>> indexSpans(HttpBulkIndexer indexer, List<Span> spans) {
-    Map<String, Set<Pair<String>>> indexToServiceSpans = new LinkedHashMap<>();
+  void indexSpans(HttpBulkIndexer indexer, List<Span> spans) {
     for (Span span : spans) {
       Long timestamp = guessTimestamp(span);
       Long timestampMillis;
       String index; // which index to store this span into
       if (timestamp != null) {
         timestampMillis = TimeUnit.MICROSECONDS.toMillis(timestamp);
-        index = indexNameFormatter.indexNameForTimestamp(timestampMillis);
+        index = indexNameFormatter.indexNameForTimestamp(SPAN, timestampMillis);
       } else {
         timestampMillis = null;
         // guessTimestamp is made for determining the span's authoritative timestamp. When choosing
@@ -81,42 +70,14 @@ Map<String, Set<Pair<String>>> indexSpans(HttpBulkIndexer indexer, List<Span> sp
           break;
         }
         if (indexTimestamp == null) indexTimestamp = System.currentTimeMillis();
-        index = indexNameFormatter.indexNameForTimestamp(indexTimestamp);
+        index = indexNameFormatter.indexNameForTimestamp(SPAN, indexTimestamp);
       }
-      if (!span.name.isEmpty()) putServiceSpans(indexToServiceSpans, index, span);
-      byte[] document = Codec.JSON.writeSpan(span);
-      if (timestampMillis != null) document = prefixWithTimestampMillis(document, timestampMillis);
-      indexer.add(index, ElasticsearchHttpSpanStore.SPAN, document, null /* Allow ES to choose an ID */);
-    }
-    return indexToServiceSpans;
-  }
-
-  void putServiceSpans(Map<String, Set<Pair<String>>> indexToServiceSpans, String index, Span s) {
-    Set<Pair<String>> serviceSpans = indexToServiceSpans.get(index);
-    if (serviceSpans == null) indexToServiceSpans.put(index, serviceSpans = new LinkedHashSet<>());
-    for (String serviceName : s.serviceNames()) {
-      serviceSpans.add(Pair.create(serviceName, s.name));
-    }
-  }
-
-  /**
-   * Adds service and span names to the pending batch. The id is "serviceName|spanName" to prevent
-   * a large order of duplicates ending up in the daily index. This also means queries do not need
-   * to deduplicate.
-   */
-  void indexNames(HttpBulkIndexer indexer, Map<String, Set<Pair<String>>> indexToServiceSpans)
-      throws IOException {
-    Buffer buffer = new Buffer();
-    for (Map.Entry<String, Set<Pair<String>>> entry : indexToServiceSpans.entrySet()) {
-      String index = entry.getKey();
-      for (Pair<String> serviceSpan : entry.getValue()) {
-        JsonWriter writer = JsonWriter.of(buffer);
-        writer.beginObject();
-        writer.name("serviceName").value(serviceSpan._1);
-        writer.name("spanName").value(serviceSpan._2);
-        writer.endObject();
-        byte[] document = buffer.readByteArray();
-        indexer.add(index, SERVICE_SPAN, document, serviceSpan._1 + "|" + serviceSpan._2);
+      for (Span2 span2 : Span2Converter.fromSpan(span)) {
+        byte[] document = Span2Codec.JSON.writeSpan(span2);
+        if (timestampMillis != null) {
+          document = prefixWithTimestampMillis(document, timestampMillis);
+        }
+        indexer.add(index, SPAN, document, null /* Allow ES to choose an ID */);
       }
     }
   }

zipkin-storage/elasticsearch-http/src/main/java/zipkin/storage/elasticsearch/http/ElasticsearchHttpSpanStore.java

@@ -15,7 +15,6 @@
 
 import java.util.Collections;
 import java.util.Iterator;
-import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
@@ -38,19 +37,18 @@
 
 final class ElasticsearchHttpSpanStore implements AsyncSpanStore {
 
+  static final String DEPENDENCY = "dependency";
   static final String SPAN = "span";
-  static final String DEPENDENCY_LINK = "dependencylink";
-  static final String SERVICE_SPAN = "servicespan";
 
   final SearchCallFactory search;
-  final String[] allIndices;
+  final String[] allSpanIndices;
   final IndexNameFormatter indexNameFormatter;
   final boolean strictTraceId;
   final int namesLookback;
 
   ElasticsearchHttpSpanStore(ElasticsearchHttpStorage es) {
     this.search = new SearchCallFactory(es.http());
-    this.allIndices = new String[] {es.indexNameFormatter().allIndices()};
+    this.allSpanIndices = new String[] {es.indexNameFormatter().allSpanIndices()};
     this.indexNameFormatter = es.indexNameFormatter();
     this.strictTraceId = es.strictTraceId();
     this.namesLookback = es.namesLookback();
@@ -63,39 +61,21 @@ final class ElasticsearchHttpSpanStore implements AsyncSpanStore {
     SearchRequest.Filters filters = new SearchRequest.Filters();
     filters.addRange("timestamp_millis", beginMillis, endMillis);
     if (request.serviceName != null) {
-      filters.addNestedTerms(asList(
-          "annotations.endpoint.serviceName",
-          "binaryAnnotations.endpoint.serviceName"
-      ), request.serviceName);
+      filters.addTerm("localEndpoint.serviceName", request.serviceName);
     }
 
     if (request.spanName != null) {
       filters.addTerm("name", request.spanName);
     }
 
     for (String annotation : request.annotations) {
-      Map<String, String> annotationValues = new LinkedHashMap<>();
-      annotationValues.put("annotations.value", annotation);
-      Map<String, String> binaryAnnotationKeys = new LinkedHashMap<>();
-      binaryAnnotationKeys.put("binaryAnnotations.key", annotation);
-      if (request.serviceName != null) {
-        annotationValues.put("annotations.endpoint.serviceName", request.serviceName);
-        binaryAnnotationKeys.put("binaryAnnotations.endpoint.serviceName", request.serviceName);
-      }
-      filters.addNestedTerms(annotationValues, binaryAnnotationKeys);
+      filters.should()
+        .addTerm("annotations.value", annotation)
+        .addExists("tags." + annotation);
     }
 
     for (Map.Entry<String, String> kv : request.binaryAnnotations.entrySet()) {
-      // In our index template, we make sure the binaryAnnotation value is indexed as string,
-      // meaning non-string values won't even be indexed at all. This means that we can only
-      // match string values here, which happens to be exactly what we want.
-      Map<String, String> nestedTerms = new LinkedHashMap<>();
-      nestedTerms.put("binaryAnnotations.key", kv.getKey());
-      nestedTerms.put("binaryAnnotations.value", kv.getValue());
-      if (request.serviceName != null) {
-        nestedTerms.put("binaryAnnotations.endpoint.serviceName", request.serviceName);
-      }
-      filters.addNestedTerms(nestedTerms);
+      filters.addTerm("tags." + kv.getKey(), kv.getValue());
     }
 
     if (request.minDuration != null) {
@@ -113,8 +93,9 @@ final class ElasticsearchHttpSpanStore implements AsyncSpanStore {
         .addSubAggregation(Aggregation.min("timestamp_millis"))
         .orderBy("timestamp_millis", "desc");
 
-    List<String> indices = indexNameFormatter.indexNamePatternsForRange(beginMillis, endMillis);
-    SearchRequest esRequest = SearchRequest.forIndicesAndType(indices, SPAN)
+    List<String> indices =
+      indexNameFormatter.indexNamePatternsForRange(SPAN, beginMillis, endMillis);
+    SearchRequest esRequest = SearchRequest.create(indices)
         .filters(filters).addAggregation(traceIdTimestamp);
 
     HttpCall<List<String>> traceIdsCall = search.newCall(esRequest, BodyConverters.SORTED_KEYS);
@@ -146,7 +127,7 @@ final class ElasticsearchHttpSpanStore implements AsyncSpanStore {
           callback.onSuccess(Collections.emptyList());
           return;
         }
-        SearchRequest request = SearchRequest.forIndicesAndType(indices, SPAN)
+        SearchRequest request = SearchRequest.create(indices)
             .terms("traceId", traceIds);
         search.newCall(request, BodyConverters.SPANS).submit(successCallback);
       }
@@ -182,7 +163,7 @@ final class ElasticsearchHttpSpanStore implements AsyncSpanStore {
   public void getRawTrace(long traceIdHigh, long traceIdLow, Callback<List<Span>> callback) {
     String traceIdHex = Util.toLowerHex(strictTraceId ? traceIdHigh : 0L, traceIdLow);
 
-    SearchRequest request = SearchRequest.forIndicesAndType(asList(allIndices), SPAN)
+    SearchRequest request = SearchRequest.create(asList(allSpanIndices))
         .term("traceId", traceIdHex);
 
     search.newCall(request, BodyConverters.NULLABLE_SPANS).submit(callback);
@@ -192,29 +173,17 @@ public void getRawTrace(long traceIdHigh, long traceIdLow, Callback<List<Span>>
     long endMillis =  System.currentTimeMillis();
     long beginMillis =  endMillis - namesLookback;
 
-    List<String> indices = indexNameFormatter.indexNamePatternsForRange(beginMillis, endMillis);
-    SearchRequest request = SearchRequest.forIndicesAndType(indices, SERVICE_SPAN)
-        .addAggregation(Aggregation.terms("serviceName", Integer.MAX_VALUE));
-
-    search.newCall(request, BodyConverters.SORTED_KEYS).submit(new Callback<List<String>>() {
-      @Override public void onSuccess(List<String> value) {
-        if (!value.isEmpty()) callback.onSuccess(value);
-
-        // Special cased code until sites update their collectors. What this does is do a more
-        // expensive nested query to get service names when the servicespan type returns nothing.
-        SearchRequest.Filters filters = new SearchRequest.Filters();
-        filters.addRange("timestamp_millis", beginMillis, endMillis);
-        SearchRequest request = SearchRequest.forIndicesAndType(indices, SPAN)
-            .filters(filters)
-            .addAggregation(Aggregation.nestedTerms("annotations.endpoint.serviceName"))
-            .addAggregation(Aggregation.nestedTerms("binaryAnnotations.endpoint.serviceName"));
-        search.newCall(request, BodyConverters.SORTED_KEYS).submit(callback);
-      }
-
-      @Override public void onError(Throwable t) {
-        callback.onError(t);
-      }
-    });
+    List<String> indices =
+      indexNameFormatter.indexNamePatternsForRange(SPAN, beginMillis, endMillis);
+    // Service name queries include both local and remote endpoints. This is different than
+    // Span name, as a span name can only be on a local endpoint.
+    SearchRequest.Filters filters = new SearchRequest.Filters();
+    filters.addRange("timestamp_millis", beginMillis, endMillis);
+    SearchRequest request = SearchRequest.create(indices)
+      .filters(filters)
+      .addAggregation(Aggregation.terms("localEndpoint.serviceName", Integer.MAX_VALUE))
+      .addAggregation(Aggregation.terms("remoteEndpoint.serviceName", Integer.MAX_VALUE));
+    search.newCall(request, BodyConverters.SORTED_KEYS).submit(callback);
   }
 
   @Override public void getSpanNames(String serviceName, Callback<List<String>> callback) {
@@ -226,34 +195,18 @@ public void getRawTrace(long traceIdHigh, long traceIdLow, Callback<List<Span>>
     long endMillis =  System.currentTimeMillis();
     long beginMillis =  endMillis - namesLookback;
 
-    List<String> indices = indexNameFormatter.indexNamePatternsForRange(beginMillis, endMillis);
-
-    SearchRequest request = SearchRequest.forIndicesAndType(indices, SERVICE_SPAN)
-        .term("serviceName", serviceName.toLowerCase(Locale.ROOT))
-        .addAggregation(Aggregation.terms("spanName", Integer.MAX_VALUE));
-
-    search.newCall(request, BodyConverters.SORTED_KEYS).submit(new Callback<List<String>>() {
-      @Override public void onSuccess(List<String> value) {
-        if (!value.isEmpty()) callback.onSuccess(value);
-
-        // Special cased code until sites update their collectors. What this does is do a more
-        // expensive nested query to get span names when the servicespan type returns nothing.
-        SearchRequest.Filters filters = new SearchRequest.Filters();
-        filters.addRange("timestamp_millis", beginMillis, endMillis);
-        filters.addNestedTerms(asList(
-            "annotations.endpoint.serviceName",
-            "binaryAnnotations.endpoint.serviceName"
-        ), serviceName.toLowerCase(Locale.ROOT));
-        SearchRequest request = SearchRequest.forIndicesAndType(indices, SPAN)
-            .filters(filters)
-            .addAggregation(Aggregation.terms("name", Integer.MAX_VALUE));
-        search.newCall(request, BodyConverters.SORTED_KEYS).submit(callback);
-      }
+    List<String> indices =
+      indexNameFormatter.indexNamePatternsForRange(SPAN, beginMillis, endMillis);
 
-      @Override public void onError(Throwable t) {
-        callback.onError(t);
-      }
-    });
+    // A span name is only valid on a local endpoint, as a span name is defined locally
+    SearchRequest.Filters filters = new SearchRequest.Filters()
+      .addRange("timestamp_millis", beginMillis, endMillis)
+      .addTerm("localEndpoint.serviceName", serviceName.toLowerCase(Locale.ROOT));
+
+    SearchRequest request = SearchRequest.create(indices)
+      .filters(filters)
+      .addAggregation(Aggregation.terms("name", Integer.MAX_VALUE));
+    search.newCall(request, BodyConverters.SORTED_KEYS).submit(callback);
   }
 
   @Override public void getDependencies(long endTs, @Nullable Long lookback,
@@ -262,13 +215,8 @@ public void getRawTrace(long traceIdHigh, long traceIdLow, Callback<List<Span>>
     long beginMillis = lookback != null ? endTs - lookback : 0;
     // We just return all dependencies in the days that fall within endTs and lookback as
     // dependency links themselves don't have timestamps.
-    List<String> indices = indexNameFormatter.indexNamePatternsForRange(beginMillis, endTs);
-    getDependencies(indices, callback);
-  }
-
-  void getDependencies(List<String> indices, Callback<List<DependencyLink>> callback) {
-    SearchRequest request = SearchRequest.forIndicesAndType(indices, DEPENDENCY_LINK);
-
-    search.newCall(request, BodyConverters.DEPENDENCY_LINKS).submit(callback);
+    List<String> indices =
+      indexNameFormatter.indexNamePatternsForRange(DEPENDENCY, beginMillis, endTs);
+    search.newCall(SearchRequest.create(indices), BodyConverters.DEPENDENCY_LINKS).submit(callback);
   }
 }