special case handling to look up traceId in both base64 and hex

Author: autata

astra/pom.xml

@@ -152,6 +152,12 @@
             <version>${curator.version}</version>
         </dependency>
 
+      <dependency>
+        <groupId>commons-codec</groupId>
+        <artifactId>commons-codec</artifactId>
+        <version>1.17.0</version>
+      </dependency>
+
         <!-- Lucene dependencies -->
         <dependency>
             <groupId>org.apache.lucene</groupId>

astra/src/main/java/com/slack/astra/zipkinApi/TraceFetcher.java

@@ -22,6 +22,7 @@
 import java.time.Instant;
 import java.time.temporal.ChronoUnit;
 import java.util.ArrayList;
+import java.util.Base64;
 import java.util.Comparator;
 import java.util.HashMap;
 import java.util.List;
@@ -30,6 +31,8 @@
 import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 import java.util.regex.Pattern;
+import org.apache.commons.codec.binary.Hex;
+import org.json.JSONArray;
 import org.json.JSONObject;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -88,20 +91,94 @@ private static boolean isDDTraceId(String s) {
     return s != null && DIGITS.matcher(s).matches();
   }
 
+  private record TraceIds(String hex, String base64) {}
+
+  private static TraceIds convertTraceId(String traceId) {
+    if (traceId == null || traceId.isEmpty()) return null;
+
+    String hex = null;
+    String base64Url = null;
+
+    // If input is hex → convert to Base64 URL-safe
+    if (traceId.matches("^[0-9a-fA-F]+$") && traceId.length() % 2 == 0) {
+      try {
+        byte[] bytes = Hex.decodeHex(traceId.toCharArray());
+        hex = traceId.toLowerCase();
+        base64Url = Base64.getUrlEncoder().encodeToString(bytes);
+        return new TraceIds(hex, base64Url);
+      } catch (Exception ignored) {
+        return null;
+      }
+    }
+
+    // Otherwise, assume Base64-URL → convert to hex
+    try {
+      byte[] bytes = Base64.getUrlDecoder().decode(traceId);
+      hex = Hex.encodeHexString(bytes);
+      base64Url = traceId;
+      return new TraceIds(hex, base64Url);
+    } catch (Exception ignored) {
+      return null;
+    }
+  }
+
+  private static JSONObject singleTermQuery(String traceFieldName, String traceId) {
+    JSONObject traceObject = new JSONObject();
+    traceObject.put(traceFieldName, traceId);
+    JSONObject queryJson = new JSONObject();
+    queryJson.put("term", traceObject);
+    return queryJson;
+  }
+
+  private static JSONObject buildTraceIdQuery(
+      String traceFieldName, String traceId, String convertedId) {
+    // When there is no converted ID, return the original simple term query
+    if (convertedId == null) {
+      return singleTermQuery(traceFieldName, traceId);
+    }
+
+    // When we have a convertedId, do traceId OR convertedId
+    JSONArray shouldArray = new JSONArray();
+    // term for original traceId
+    shouldArray.put(singleTermQuery(traceFieldName, traceId));
+    // term for convertedId
+    shouldArray.put(singleTermQuery(traceFieldName, convertedId));
+    JSONObject boolQuery = new JSONObject();
+    boolQuery.put("should", shouldArray);
+    boolQuery.put("minimum_should_match", 1);
+
+    JSONObject queryJson = new JSONObject();
+    queryJson.put("bool", boolQuery);
+
+    return queryJson;
+  }
+
   private Result fetchTraceResult(
       String traceId,
       Optional<Long> startTimeEpochMs,
       Optional<Long> endTimeEpochMs,
       Optional<Integer> maxSpans,
       Optional<Boolean> userRequest,
       Optional<Long> dataFreshnessInSeconds,
-      Optional<Boolean> ddTraceIdEnabled)
+      Optional<Boolean> ddTraceIdEnabled,
+      Optional<Boolean> searchByHexAndBase64)
       throws IOException {
 
     String traceFieldName = "trace_id";
-    // if trace id looks like dd_trace_id, then use dd_trace_id field to search
+    String convertedId = null;
+    // if trace id looks like dd_trace_id, then use dd_trace_id field to search. If true, ignores
+    // the hex/base64 flag
     if (ddTraceIdEnabled.isPresent() && ddTraceIdEnabled.get() && isDDTraceId(traceId)) {
       traceFieldName = "dd_trace_id";
+    } else if (searchByHexAndBase64.isPresent() && searchByHexAndBase64.get()) {
+      TraceIds traceIds = convertTraceId(traceId);
+      // to ensure we only cache the same trace once, we use the base64 version as the traceId.
+      // TraceIds are null
+      // in the case of being unable to convert, fallback to original traceId
+      if (traceIds != null) {
+        traceId = traceIds.base64;
+        convertedId = traceIds.hex;
+      }
     }
 
     // Log the custom header userRequest value if present
@@ -117,11 +194,9 @@ private Result fetchTraceResult(
       }
     }
 
-    JSONObject traceObject = new JSONObject();
-    traceObject.put(traceFieldName, traceId);
-    JSONObject queryJson = new JSONObject();
-    queryJson.put("term", traceObject);
+    JSONObject queryJson = buildTraceIdQuery(traceFieldName, traceId, convertedId);
     String queryString = queryJson.toString();
+    LOG.debug("Querying with queryString={}", queryString);
 
     long startTime =
         startTimeEpochMs.orElseGet(
@@ -180,7 +255,8 @@ public List<ZipkinSpanResponse> getSpansByTraceId(
       Optional<Integer> maxSpans,
       Optional<Boolean> userRequest,
       Optional<Long> dataFreshnessInSeconds,
-      Optional<Boolean> ddTraceIdEnabled)
+      Optional<Boolean> ddTraceIdEnabled,
+      Optional<Boolean> searchByHexAndBase64)
       throws IOException {
     Result result =
         fetchTraceResult(
@@ -190,7 +266,8 @@ public List<ZipkinSpanResponse> getSpansByTraceId(
             maxSpans,
             userRequest,
             dataFreshnessInSeconds,
-            ddTraceIdEnabled);
+            ddTraceIdEnabled,
+            searchByHexAndBase64);
 
     if (result.isCached()) {
       return objectMapper.readValue(result.rawJson, new TypeReference<>() {});
@@ -205,7 +282,8 @@ public String getByTraceId(
       Optional<Integer> maxSpans,
       Optional<Boolean> userRequest,
       Optional<Long> dataFreshnessInSeconds,
-      Optional<Boolean> ddTraceIdEnabled)
+      Optional<Boolean> ddTraceIdEnabled,
+      Optional<Boolean> searchByHexAndBase64)
       throws IOException {
     Result result =
         fetchTraceResult(
@@ -215,7 +293,8 @@ public String getByTraceId(
             maxSpans,
             userRequest,
             dataFreshnessInSeconds,
-            ddTraceIdEnabled);
+            ddTraceIdEnabled,
+            searchByHexAndBase64);
 
     if (result.isCached()) {
       return result.rawJson;

astra/src/main/java/com/slack/astra/zipkinApi/ZipkinService.java

@@ -71,7 +71,8 @@ public HttpResponse getTraceByTraceId(
       @Param("maxSpans") Optional<Integer> maxSpans,
       @Header("X-User-Request") Optional<Boolean> userRequest,
       @Header("X-Data-Freshness-In-Seconds") Optional<Long> dataFreshnessInSeconds,
-      @Header("X-DD-TRACE-ID") Optional<Boolean> ddTraceIdEnabled)
+      @Header("X-DD-TRACE-ID") Optional<Boolean> ddTraceIdEnabled,
+      @Header("X-E2E-hex-base64") Optional<Boolean> searchByHexAndBase64)
       throws IOException {
     String output =
         this.traceFetcher.getByTraceId(
@@ -81,7 +82,8 @@ public HttpResponse getTraceByTraceId(
             maxSpans,
             userRequest,
             dataFreshnessInSeconds,
-            ddTraceIdEnabled);
+            ddTraceIdEnabled,
+            searchByHexAndBase64);
     return HttpResponse.of(HttpStatus.OK, MediaType.JSON_UTF_8, output);
   }
 }