make default maxSpans for zipkin endpoint configurable (#19)

airbnb · Dec 6, 2024 · 471edf9 · 471edf9
1 parent 16d8f56
commit 471edf9
Show file tree

Hide file tree

Showing 10 changed files with 81 additions and 5 deletions.
diff --git a/astra/src/main/java/com/slack/astra/server/Astra.java b/astra/src/main/java/com/slack/astra/server/Astra.java
@@ -223,7 +223,10 @@ private static Set<Service> getServices(
               .withRequestTimeout(requestTimeout)
               .withTracing(astraConfig.getTracingConfig())
               .withAnnotatedService(new ElasticsearchApiService(astraDistributedQueryService))
-              .withAnnotatedService(new ZipkinService(astraDistributedQueryService))
+              .withAnnotatedService(
+                  new ZipkinService(
+                      astraDistributedQueryService,
+                      astraConfig.getQueryConfig().getZipkinDefaultMaxSpans()))
               .withGrpcService(astraDistributedQueryService)
               .build();
       services.add(armeriaService);

diff --git a/astra/src/main/java/com/slack/astra/server/ValidateAstraConfig.java b/astra/src/main/java/com/slack/astra/server/ValidateAstraConfig.java
@@ -50,6 +50,9 @@ private static void validateQueryConfig(AstraConfigs.QueryServiceConfig queryCon
         queryConfig.getServerConfig().getRequestTimeoutMs()
             > queryConfig.getDefaultQueryTimeoutMs(),
         "QueryConfig requestTimeoutMs must be higher than defaultQueryTimeoutMs");
+    checkArgument(
+        queryConfig.getZipkinDefaultMaxSpans() >= 1000,
+        "QueryConfig zipkinDefaultMaxSpans cannot less than 1000");
   }
 
   private static void validateCacheConfig(AstraConfigs.CacheConfig cacheConfig) {

diff --git a/astra/src/main/java/com/slack/astra/zipkinApi/ZipkinService.java b/astra/src/main/java/com/slack/astra/zipkinApi/ZipkinService.java
@@ -148,7 +148,7 @@ protected static long convertToMicroSeconds(Instant instant) {
   private static final Logger LOG = LoggerFactory.getLogger(ZipkinService.class);
   private static long LOOKBACK_MINS = 60 * 24 * 7;
 
-  private static final int MAX_SPANS = 20_000;
+  private final int defaultMaxSpans;
 
   private final AstraQueryServiceBase searcher;
 
@@ -160,8 +160,9 @@ protected static long convertToMicroSeconds(Instant instant) {
           .serializationInclusion(JsonInclude.Include.NON_EMPTY)
           .build();
 
-  public ZipkinService(AstraQueryServiceBase searcher) {
+  public ZipkinService(AstraQueryServiceBase searcher, int defaultMaxSpans) {
     this.searcher = searcher;
+    this.defaultMaxSpans = defaultMaxSpans;
   }
 
   @Get
@@ -216,7 +217,7 @@ public HttpResponse getTraceByTraceId(
     long endTime =
         endTimeEpochMs.orElseGet(
             () -> Instant.now().plus(LOOKBACK_MINS, ChronoUnit.MINUTES).toEpochMilli());
-    int howMany = maxSpans.orElse(MAX_SPANS);
+    int howMany = maxSpans.orElse(this.defaultMaxSpans);
 
     brave.Span span = Tracing.currentTracer().currentSpan();
     span.tag("startTimeEpochMs", String.valueOf(startTime));

diff --git a/astra/src/main/proto/astra_configs.proto b/astra/src/main/proto/astra_configs.proto
@@ -85,6 +85,7 @@ message QueryServiceConfig {
   ServerConfig server_config = 1;
   int32 default_query_timeout_ms = 2;
   string managerConnectString = 3;
+  int32 zipkin_default_max_spans = 4;
 }
 
 enum KafkaOffsetLocation {

diff --git a/astra/src/test/java/com/slack/astra/testlib/AstraConfigUtil.java b/astra/src/test/java/com/slack/astra/testlib/AstraConfigUtil.java
@@ -81,6 +81,7 @@ public static AstraConfigs.AstraConfig makeAstraConfig(
                     .setRequestTimeoutMs(5000)
                     .build())
             .setDefaultQueryTimeoutMs(3000)
+            .setZipkinDefaultMaxSpans(20000)
             .build();
 
     AstraConfigs.RecoveryConfig recoveryConfig =

diff --git a/astra/src/test/java/com/slack/astra/zipkinApi/ZipkinServiceTest.java b/astra/src/test/java/com/slack/astra/zipkinApi/ZipkinServiceTest.java
@@ -24,17 +24,20 @@
 import org.junit.jupiter.api.Test;
 import org.mockito.Mock;
 import org.mockito.MockedStatic;
+import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 
 public class ZipkinServiceTest {
   @Mock private AstraQueryServiceBase searcher;
   private ZipkinService zipkinService;
   private AstraSearch.SearchResult mockSearchResult;
 
+  private static final int defaultMaxSpans = 2000;
+
   @BeforeEach
   public void setup() throws IOException {
     MockitoAnnotations.openMocks(this);
-    zipkinService = spy(new ZipkinService(searcher));
+    zipkinService = spy(new ZipkinService(searcher, defaultMaxSpans));
     // Build mockSearchResult
     ObjectMapper objectMapper = new ObjectMapper();
     JsonNode jsonNode =
@@ -68,4 +71,55 @@ public void testGetTraceByTraceId_onlyTraceIdProvided() throws Exception {
       assertEquals(HttpStatus.OK, aggregatedResponse.status());
     }
   }
+
+  @Test
+  public void testGetTraceByTraceId_respectsDefaultMaxSpans() throws Exception {
+    try (MockedStatic<Tracing> mockedTracing = mockStatic(Tracing.class)) {
+      // Mocking Tracing and Span
+      brave.Tracer mockTracer = mock(brave.Tracer.class);
+      Span mockSpan = mock(Span.class);
+
+      mockedTracing.when(Tracing::currentTracer).thenReturn(mockTracer);
+      when(mockTracer.currentSpan()).thenReturn(mockSpan);
+
+      String traceId = "test_trace_2";
+      when(searcher.doSearch(any())).thenReturn(mockSearchResult);
+
+      zipkinService.getTraceByTraceId(
+          traceId, Optional.empty(), Optional.empty(), Optional.empty());
+
+      Mockito.verify(searcher)
+          .doSearch(
+              Mockito.argThat(
+                  request ->
+                      request.getHowMany() == defaultMaxSpans
+                          && request.getQuery().contains("\"trace_id\":\"" + traceId + "\"")));
+    }
+  }
+
+  @Test
+  public void testGetTraceByTraceId_respectsMaxSpans() throws Exception {
+    try (MockedStatic<Tracing> mockedTracing = mockStatic(Tracing.class)) {
+      // Mocking Tracing and Span
+      brave.Tracer mockTracer = mock(brave.Tracer.class);
+      Span mockSpan = mock(Span.class);
+
+      mockedTracing.when(Tracing::currentTracer).thenReturn(mockTracer);
+      when(mockTracer.currentSpan()).thenReturn(mockSpan);
+
+      String traceId = "test_trace_2";
+      when(searcher.doSearch(any())).thenReturn(mockSearchResult);
+      int maxSpansParam = 10000;
+
+      zipkinService.getTraceByTraceId(
+          traceId, Optional.empty(), Optional.empty(), Optional.of(maxSpansParam));
+
+      Mockito.verify(searcher)
+          .doSearch(
+              Mockito.argThat(
+                  request ->
+                      request.getHowMany() == maxSpansParam
+                          && request.getQuery().contains("\"trace_id\":\"" + traceId + "\"")));
+    }
+  }
 }
diff --git a/astra/src/test/resources/test_config.json b/astra/src/test/resources/test_config.json
@@ -57,6 +57,7 @@
       "requestTimeoutMs": 3000
     },
     "defaultQueryTimeoutMs": 1500,
+    "zipkinDefaultMaxSpans": 20000,
     "managerConnectString": "localhost:8083"
   },
   "metadataStoreConfig": {

diff --git a/astra/src/test/resources/test_config.yaml b/astra/src/test/resources/test_config.yaml
@@ -33,6 +33,7 @@ queryConfig:
     serverAddress: "1.2.3.4"
     requestTimeoutMs: 3000
   defaultQueryTimeoutMs: 2500
+  zipkinDefaultMaxSpans: 20000
   managerConnectString: localhost:8083
 
 s3Config:

diff --git a/config/config.yaml b/config/config.yaml
@@ -50,6 +50,7 @@ queryConfig:
     serverAddress: ${ASTRA_QUERY_SERVER_ADDRESS:-localhost}
     requestTimeoutMs: ${ASTRA_QUERY_REQUEST_TIMEOUT_MS:-5000}
   defaultQueryTimeoutMs: ${ASTRA_QUERY_DEFAULT_QUERY_TIMEOUT_MS:-3000}
+  zipkinDefaultMaxSpans: ${ASTRA_QUERY_ZIPKIN_DEFAULT_MAX_SPANS:-20000}
   managerConnectString: ${ASTRA_MANAGER_CONNECTION_STRING:-localhost:8083}
 
 metadataStoreConfig:

diff --git a/docs/topics/Config-options.md b/docs/topics/Config-options.md
@@ -345,6 +345,16 @@ seconds higher than the <a href="Config-options.md#defaultquerytimeout">queryCon
 to allow for aggregation post-processing to occur.
 </tip>
 
+### zipkinDefaultMaxSpans
+
+```yaml
+queryConfig:
+  zipkinDefaultMaxSpans: 25000
+```
+
+Max spans that the zipkin endpoint will return when the API call does not include `maxSpans`. A trace with more than 
+this amount of spans will be truncated.
+
 ### managerConnectString
 ```yaml
 queryConfig: