diff --git a/astra/src/main/java/com/slack/astra/server/Astra.java b/astra/src/main/java/com/slack/astra/server/Astra.java index 90e6f011d1..c6c47c3f87 100644 --- 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 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 index 94cf1be2fc..01d12e5065 100644 --- 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 index 05158e2d6f..0256dad8f9 100644 --- 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 index d00b42a1be..605122f971 100644 --- 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 index 3324868c01..4d9890ca3a 100644 --- 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 index 9154d28a7f..bf19dc80ca 100644 --- a/astra/src/test/java/com/slack/astra/zipkinApi/ZipkinServiceTest.java +++ b/astra/src/test/java/com/slack/astra/zipkinApi/ZipkinServiceTest.java @@ -24,6 +24,7 @@ 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 { @@ -31,10 +32,12 @@ public class ZipkinServiceTest { 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 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 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 index 5587141347..9356fb25f9 100644 --- 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 index b1b39d87ae..991b39194b 100644 --- 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 index f2c9fb3a0a..310bd85065 100644 --- 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 index 2a125590e5..d4b465d7ad 100644 --- a/docs/topics/Config-options.md +++ b/docs/topics/Config-options.md @@ -345,6 +345,16 @@ seconds higher than the queryCon to allow for aggregation post-processing to occur. +### 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: