This is an automated email from the ASF dual-hosted git repository. zqr10159 pushed a commit to branch 2.0.0 in repository https://gitbox.apache.org/repos/asf/hertzbeat.git
commit d40787953d2f2d1926ddcf0d03cb84b413feceed Author: Logic <[email protected]> AuthorDate: Wed Aug 26 15:49:06 2026 +0800 feat: isolate observability signal queries --- .../manager/controller/TopologyController.java | 13 +- .../manager/support/GlobalExceptionHandler.java | 15 ++ .../manager/controller/TopologyControllerTest.java | 7 + ...obalExceptionHandlerObservabilityQueryTest.java | 61 +++++++ .../controller/OtlpIngestionController.java | 29 ++-- .../logs/controller/LogQueryController.java | 36 ++-- .../traces/controller/TraceQueryController.java | 35 ++-- .../controller/OtlpIngestionControllerTest.java | 10 +- .../logs/controller/LogQueryControllerTest.java | 16 +- .../LogQueryWorkspaceControllerContractTest.java | 6 +- .../controller/TraceQueryControllerTest.java | 6 +- .../src/main/resources/application.yml | 13 ++ .../ObservabilityQueryAdmissionException.java | 52 ++++++ .../ObservabilityQueryAdmissionService.java | 133 +++++++++++++++ .../ObservabilityQueryAdmissionServiceTest.java | 182 +++++++++++++++++++++ 15 files changed, 563 insertions(+), 51 deletions(-) diff --git a/hertzbeat-manager/src/main/java/org/apache/hertzbeat/manager/controller/TopologyController.java b/hertzbeat-manager/src/main/java/org/apache/hertzbeat/manager/controller/TopologyController.java index 7ef40ddf1b..9ecbf076d0 100644 --- a/hertzbeat-manager/src/main/java/org/apache/hertzbeat/manager/controller/TopologyController.java +++ b/hertzbeat-manager/src/main/java/org/apache/hertzbeat/manager/controller/TopologyController.java @@ -28,6 +28,7 @@ import org.apache.hertzbeat.common.observability.gateway.AuthTokenScopes; import org.apache.hertzbeat.common.support.exception.CommonException; import org.apache.hertzbeat.manager.pojo.dto.EntityTopologyGraphInfo; import org.apache.hertzbeat.manager.service.entity.EntityTopologyQueryService; +import org.apache.hertzbeat.warehouse.query.admission.ObservabilityQueryAdmissionService; import org.springframework.http.ResponseEntity; import org.springframework.util.StringUtils; import org.springframework.web.bind.annotation.GetMapping; @@ -44,9 +45,12 @@ import org.springframework.web.bind.annotation.RestController; public class TopologyController { private final EntityTopologyQueryService entityTopologyQueryService; + private final ObservabilityQueryAdmissionService queryAdmissionService; - public TopologyController(EntityTopologyQueryService entityTopologyQueryService) { + public TopologyController(EntityTopologyQueryService entityTopologyQueryService, + ObservabilityQueryAdmissionService queryAdmissionService) { this.entityTopologyQueryService = entityTopologyQueryService; + this.queryAdmissionService = queryAdmissionService; } @GetMapping @@ -77,9 +81,10 @@ public class TopologyController { throw new CommonException("topology_workspace_unavailable"); } String trustedWorkspaceId = AuthTokenScopes.normalizeWorkspaceId(workspaceId); - EntityTopologyGraphInfo graph = entityTopologyQueryService.buildFocusedTopology( - trustedWorkspaceId, focusEntityId, depth, environment, sourceKind, start, end, - relationType, hideInternal, pageIndex, pageSize); + EntityTopologyGraphInfo graph = queryAdmissionService.execute("topology", + () -> entityTopologyQueryService.buildFocusedTopology( + trustedWorkspaceId, focusEntityId, depth, environment, sourceKind, start, end, + relationType, hideInternal, pageIndex, pageSize)); return ResponseEntity.ok(Message.success(graph)); } } diff --git a/hertzbeat-manager/src/main/java/org/apache/hertzbeat/manager/support/GlobalExceptionHandler.java b/hertzbeat-manager/src/main/java/org/apache/hertzbeat/manager/support/GlobalExceptionHandler.java index f68611abbf..1dafb467cf 100644 --- a/hertzbeat-manager/src/main/java/org/apache/hertzbeat/manager/support/GlobalExceptionHandler.java +++ b/hertzbeat-manager/src/main/java/org/apache/hertzbeat/manager/support/GlobalExceptionHandler.java @@ -33,7 +33,9 @@ import org.apache.hertzbeat.alert.notice.AlertNoticeException; import org.apache.hertzbeat.manager.support.exception.MonitorDatabaseException; import org.apache.hertzbeat.manager.support.exception.MonitorDetectException; import org.apache.hertzbeat.manager.support.exception.MonitorMetricsException; +import org.apache.hertzbeat.warehouse.query.admission.ObservabilityQueryAdmissionException; import org.springframework.dao.DataAccessException; +import org.springframework.http.HttpHeaders; import org.springframework.http.HttpStatus; import org.springframework.http.ResponseEntity; import org.springframework.http.converter.HttpMessageNotReadableException; @@ -58,6 +60,19 @@ public class GlobalExceptionHandler { private static final String UNKNOWN_ERROR_MESSAGE = "unknown error happen"; private static final String TELEMETRY_STORAGE_UNAVAILABLE_MESSAGE = "telemetry storage unavailable"; + /** Return an explicit retryable response when a bounded signal-query lane cannot accept work. */ + @ExceptionHandler(ObservabilityQueryAdmissionException.class) + @ResponseBody + ResponseEntity<Message<Void>> handleObservabilityQueryAdmissionException( + ObservabilityQueryAdmissionException exception) { + HttpStatus status = exception.getReason() == ObservabilityQueryAdmissionException.Reason.OVERLOADED + ? HttpStatus.TOO_MANY_REQUESTS + : HttpStatus.SERVICE_UNAVAILABLE; + return ResponseEntity.status(status) + .header(HttpHeaders.RETRY_AFTER, "1") + .body(Message.fail(FAIL_CODE, exception.getMessage())); + } + /** * Preserve the difference between a valid empty telemetry query and an unavailable store. * diff --git a/hertzbeat-manager/src/test/java/org/apache/hertzbeat/manager/controller/TopologyControllerTest.java b/hertzbeat-manager/src/test/java/org/apache/hertzbeat/manager/controller/TopologyControllerTest.java index 349fb2f44e..4f028ca1b7 100644 --- a/hertzbeat-manager/src/test/java/org/apache/hertzbeat/manager/controller/TopologyControllerTest.java +++ b/hertzbeat-manager/src/test/java/org/apache/hertzbeat/manager/controller/TopologyControllerTest.java @@ -26,6 +26,7 @@ import static org.springframework.test.web.servlet.result.MockMvcResultMatchers. import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.jsonPath; import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status; +import java.time.Duration; import java.util.List; import java.util.Set; import org.apache.hertzbeat.common.constants.CommonConstants; @@ -34,12 +35,14 @@ import org.apache.hertzbeat.common.support.exception.CommonException; import org.apache.hertzbeat.manager.pojo.dto.EntityTopologyGraphInfo; import org.apache.hertzbeat.manager.service.entity.EntityTopologyQueryService; import org.apache.hertzbeat.manager.support.GlobalExceptionHandler; +import org.apache.hertzbeat.warehouse.query.admission.ObservabilityQueryAdmissionService; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.InjectMocks; import org.mockito.Mock; +import org.mockito.Spy; import org.mockito.junit.jupiter.MockitoExtension; import org.springframework.test.web.servlet.MockMvc; import org.springframework.test.web.servlet.request.MockMvcRequestBuilders; @@ -59,6 +62,10 @@ class TopologyControllerTest { @Mock private EntityTopologyQueryService entityTopologyQueryService; + @Spy + private ObservabilityQueryAdmissionService queryAdmissionService = + new ObservabilityQueryAdmissionService(8, 8, 8, 4, 8, Duration.ofMillis(100)); + @BeforeEach void setUp() { AuthTokenRequestContext.bindWorkspaceId("default"); diff --git a/hertzbeat-manager/src/test/java/org/apache/hertzbeat/manager/support/GlobalExceptionHandlerObservabilityQueryTest.java b/hertzbeat-manager/src/test/java/org/apache/hertzbeat/manager/support/GlobalExceptionHandlerObservabilityQueryTest.java new file mode 100644 index 0000000000..b1842f3937 --- /dev/null +++ b/hertzbeat-manager/src/test/java/org/apache/hertzbeat/manager/support/GlobalExceptionHandlerObservabilityQueryTest.java @@ -0,0 +1,61 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hertzbeat.manager.support; + +import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.get; +import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.header; +import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.jsonPath; +import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status; + +import org.apache.hertzbeat.warehouse.query.admission.ObservabilityQueryAdmissionException; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.springframework.test.web.servlet.MockMvc; +import org.springframework.test.web.servlet.setup.MockMvcBuilders; +import org.springframework.web.bind.annotation.GetMapping; +import org.springframework.web.bind.annotation.RestController; + +class GlobalExceptionHandlerObservabilityQueryTest { + + private MockMvc mockMvc; + + @BeforeEach + void setUp() { + mockMvc = MockMvcBuilders.standaloneSetup(new FailingQueryController()) + .setControllerAdvice(new GlobalExceptionHandler()) + .build(); + } + + @Test + void mapsOverloadToTooManyRequestsWithRetryHint() throws Exception { + mockMvc.perform(get("/test/query/overloaded")) + .andExpect(status().isTooManyRequests()) + .andExpect(header().string("Retry-After", "1")) + .andExpect(jsonPath("$.msg").value("Observability metrics query capacity is exhausted.")); + } + + @RestController + private static class FailingQueryController { + + @GetMapping("/test/query/overloaded") + void overloaded() { + throw new ObservabilityQueryAdmissionException( + "metrics", ObservabilityQueryAdmissionException.Reason.OVERLOADED, null); + } + } +} diff --git a/hertzbeat-observability/src/main/java/org/apache/hertzbeat/observability/ingestion/controller/OtlpIngestionController.java b/hertzbeat-observability/src/main/java/org/apache/hertzbeat/observability/ingestion/controller/OtlpIngestionController.java index d75b28bdf2..c52cd033dd 100644 --- a/hertzbeat-observability/src/main/java/org/apache/hertzbeat/observability/ingestion/controller/OtlpIngestionController.java +++ b/hertzbeat-observability/src/main/java/org/apache/hertzbeat/observability/ingestion/controller/OtlpIngestionController.java @@ -33,6 +33,7 @@ import org.apache.hertzbeat.common.observability.dto.metrics.OtlpRelatedMetricsD import org.apache.hertzbeat.observability.ingestion.red.OtlpIngestionRedSummaryService; import org.apache.hertzbeat.observability.ingestion.service.OtlpIngestionWorkspaceService; import org.apache.hertzbeat.observability.metrics.service.CollectorScopedMetricsQueryService; +import org.apache.hertzbeat.warehouse.query.admission.ObservabilityQueryAdmissionService; import org.springframework.http.ResponseEntity; import org.springframework.util.StringUtils; import org.springframework.web.bind.annotation.GetMapping; @@ -52,6 +53,7 @@ public class OtlpIngestionController { private final OtlpIngestionWorkspaceService otlpIngestionWorkspaceService; private final OtlpIngestionRedSummaryService otlpIngestionRedSummaryService; private final CollectorScopedMetricsQueryService collectorScopedMetricsQueryService; + private final ObservabilityQueryAdmissionService queryAdmissionService; @GetMapping("/overview") @Operation(summary = "Unified OTLP ingestion overview") @@ -103,11 +105,12 @@ public class OtlpIngestionController { @RequestParam(value = "limit", required = false) String limit, @RequestParam(value = "operationName", required = false) String operationName) { String workspaceId = currentWorkspaceId(); - return ResponseEntity.ok(Message.success(collectorScopedMetricsQueryService.query( - new CollectorScopedMetricsQueryService.Request( - workspaceId, entityId, entityType, start, end, serviceName, serviceNamespace, environment, collectorId, - instance, endpoint, query, filter, groupBy, aggregation, temporalAggregation, step, limit, - operationName)))); + return ResponseEntity.ok(Message.success(queryAdmissionService.execute("metrics", + () -> collectorScopedMetricsQueryService.query( + new CollectorScopedMetricsQueryService.Request( + workspaceId, entityId, entityType, start, end, serviceName, serviceNamespace, + environment, collectorId, instance, endpoint, query, filter, groupBy, aggregation, + temporalAggregation, step, limit, operationName))))); } @GetMapping("/metrics/inventory") @@ -125,10 +128,11 @@ public class OtlpIngestionController { @RequestParam(value = "endpoint", required = false) String endpoint, @RequestParam(value = "limit", required = false) String limit) { String workspaceId = currentWorkspaceId(); - return ResponseEntity.ok(Message.success(collectorScopedMetricsQueryService.inventory( - new CollectorScopedMetricsQueryService.InventoryRequest( - workspaceId, entityId, entityType, start, end, serviceName, serviceNamespace, environment, - collectorId, instance, endpoint, limit)))); + return ResponseEntity.ok(Message.success(queryAdmissionService.execute("metrics", + () -> collectorScopedMetricsQueryService.inventory( + new CollectorScopedMetricsQueryService.InventoryRequest( + workspaceId, entityId, entityType, start, end, serviceName, serviceNamespace, + environment, collectorId, instance, endpoint, limit))))); } @GetMapping("/metrics/related") @@ -145,9 +149,10 @@ public class OtlpIngestionController { @RequestParam(value = "operationName", required = false) String operationName, @RequestParam(value = "limit", required = false) String limit) { String workspaceId = currentWorkspaceId(); - return ResponseEntity.ok(Message.success(otlpIngestionWorkspaceService.getRelatedMetrics( - workspaceId, entityId, entityType, start, end, serviceName, serviceNamespace, environment, filter, - operationName, limit))); + return ResponseEntity.ok(Message.success(queryAdmissionService.execute("metrics", + () -> otlpIngestionWorkspaceService.getRelatedMetrics( + workspaceId, entityId, entityType, start, end, serviceName, serviceNamespace, environment, + filter, operationName, limit)))); } private String currentWorkspaceId() { diff --git a/hertzbeat-observability/src/main/java/org/apache/hertzbeat/observability/logs/controller/LogQueryController.java b/hertzbeat-observability/src/main/java/org/apache/hertzbeat/observability/logs/controller/LogQueryController.java index 800bd3f6cf..644d0c2cb8 100644 --- a/hertzbeat-observability/src/main/java/org/apache/hertzbeat/observability/logs/controller/LogQueryController.java +++ b/hertzbeat-observability/src/main/java/org/apache/hertzbeat/observability/logs/controller/LogQueryController.java @@ -21,6 +21,7 @@ import io.swagger.v3.oas.annotations.Operation; import io.swagger.v3.oas.annotations.Parameter; import io.swagger.v3.oas.annotations.tags.Tag; import java.util.Map; +import java.util.function.Supplier; import org.apache.hertzbeat.common.entity.dto.Message; import org.apache.hertzbeat.common.entity.dto.PageResponse; import org.apache.hertzbeat.common.entity.log.LogEntry; @@ -31,6 +32,7 @@ import org.apache.hertzbeat.observability.ingestion.semantic.OtlpResourceSemanti import org.apache.hertzbeat.observability.logs.service.LogQueryService; import org.apache.hertzbeat.observability.shared.query.CollectorResourceScope; import org.apache.hertzbeat.observability.shared.query.TelemetryQueryContextScope; +import org.apache.hertzbeat.warehouse.query.admission.ObservabilityQueryAdmissionService; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.data.domain.Page; import org.springframework.http.ResponseEntity; @@ -49,10 +51,13 @@ import org.springframework.web.bind.annotation.RestController; public class LogQueryController { private final LogQueryService logQueryService; + private final ObservabilityQueryAdmissionService queryAdmissionService; @Autowired - public LogQueryController(LogQueryService logQueryService) { + public LogQueryController(LogQueryService logQueryService, + ObservabilityQueryAdmissionService queryAdmissionService) { this.logQueryService = logQueryService; + this.queryAdmissionService = queryAdmissionService; } @GetMapping("/list") @@ -101,10 +106,11 @@ public class LogQueryController { String workspaceId = trustedWorkspaceId(); ScopedFilters scopedFilters = scopeFilters( entityId, entityType, collectorId, instance, endpoint, resourceFilter, attributeFilter); - Page<LogEntry> result = logQueryService.list(workspaceId, entityId, start, end, traceId, spanId, + Page<LogEntry> result = executeQuery(() -> logQueryService.list( + workspaceId, entityId, start, end, traceId, spanId, severityNumber, severityText, search, serviceName, serviceNamespace, environment, scopedFilters.resourceFilter(), scopedFilters.attributeFilter(), - pageIndex, pageSize, hideInternal, hideNoise); + pageIndex, pageSize, hideInternal, hideNoise)); return ResponseEntity.ok(Message.success(PageResponse.from(result))); } @@ -148,10 +154,10 @@ public class LogQueryController { String workspaceId = trustedWorkspaceId(); ScopedFilters scopedFilters = scopeFilters( entityId, entityType, collectorId, instance, endpoint, resourceFilter, attributeFilter); - return ResponseEntity.ok(Message.success(logQueryService.context( + return ResponseEntity.ok(Message.success(executeQuery(() -> logQueryService.context( workspaceId, entityId, logTimeUnixNano, start, end, serviceName, serviceNamespace, environment, scopedFilters.resourceFilter(), scopedFilters.attributeFilter(), limit, direction, - cursorLogTimeUnixNano, hideInternal, hideNoise))); + cursorLogTimeUnixNano, hideInternal, hideNoise)))); } @GetMapping("/stats/overview") @@ -196,10 +202,10 @@ public class LogQueryController { String workspaceId = trustedWorkspaceId(); ScopedFilters scopedFilters = scopeFilters( entityId, entityType, collectorId, instance, endpoint, resourceFilter, attributeFilter); - return ResponseEntity.ok(Message.success(logQueryService.overviewStats( + return ResponseEntity.ok(Message.success(executeQuery(() -> logQueryService.overviewStats( workspaceId, entityId, start, end, traceId, spanId, severityNumber, severityText, search, serviceName, serviceNamespace, environment, scopedFilters.resourceFilter(), scopedFilters.attributeFilter(), - hideInternal, hideNoise))); + hideInternal, hideNoise)))); } @GetMapping("/stats/trace-coverage") @@ -244,10 +250,10 @@ public class LogQueryController { String workspaceId = trustedWorkspaceId(); ScopedFilters scopedFilters = scopeFilters( entityId, entityType, collectorId, instance, endpoint, resourceFilter, attributeFilter); - return ResponseEntity.ok(Message.success(logQueryService.traceCoverageStats( + return ResponseEntity.ok(Message.success(executeQuery(() -> logQueryService.traceCoverageStats( workspaceId, entityId, start, end, traceId, spanId, severityNumber, severityText, search, serviceName, serviceNamespace, environment, scopedFilters.resourceFilter(), scopedFilters.attributeFilter(), - hideInternal, hideNoise))); + hideInternal, hideNoise)))); } @GetMapping("/stats/trend") @@ -292,10 +298,10 @@ public class LogQueryController { String workspaceId = trustedWorkspaceId(); ScopedFilters scopedFilters = scopeFilters( entityId, entityType, collectorId, instance, endpoint, resourceFilter, attributeFilter); - return ResponseEntity.ok(Message.success(logQueryService.trendStats( + return ResponseEntity.ok(Message.success(executeQuery(() -> logQueryService.trendStats( workspaceId, entityId, start, end, traceId, spanId, severityNumber, severityText, search, serviceName, serviceNamespace, environment, scopedFilters.resourceFilter(), scopedFilters.attributeFilter(), - hideInternal, hideNoise))); + hideInternal, hideNoise)))); } @GetMapping("/stats/group-by") @@ -348,11 +354,15 @@ public class LogQueryController { String workspaceId = trustedWorkspaceId(); ScopedFilters scopedFilters = scopeFilters( entityId, entityType, collectorId, instance, endpoint, resourceFilter, attributeFilter); - return ResponseEntity.ok(Message.success(logQueryService.groupByStats( + return ResponseEntity.ok(Message.success(executeQuery(() -> logQueryService.groupByStats( workspaceId, entityId, start, end, traceId, spanId, severityNumber, severityText, search, serviceName, serviceNamespace, environment, scopedFilters.resourceFilter(), scopedFilters.attributeFilter(), groupBy, - limit, orderBy, minCount, hideInternal, hideNoise))); + limit, orderBy, minCount, hideInternal, hideNoise)))); + } + + private <T> T executeQuery(Supplier<T> query) { + return queryAdmissionService.execute("logs", query); } private String mergeEntityContextResourceFilter(Long entityId, String entityType, String resourceFilter) { diff --git a/hertzbeat-observability/src/main/java/org/apache/hertzbeat/observability/traces/controller/TraceQueryController.java b/hertzbeat-observability/src/main/java/org/apache/hertzbeat/observability/traces/controller/TraceQueryController.java index 8ea5950d24..968351c59c 100644 --- a/hertzbeat-observability/src/main/java/org/apache/hertzbeat/observability/traces/controller/TraceQueryController.java +++ b/hertzbeat-observability/src/main/java/org/apache/hertzbeat/observability/traces/controller/TraceQueryController.java @@ -35,6 +35,7 @@ import org.apache.hertzbeat.observability.shared.query.CollectorResourceScope; import org.apache.hertzbeat.observability.shared.query.TelemetryQueryContextScope; import org.apache.hertzbeat.observability.traces.service.EntityTraceQueryService; import org.apache.hertzbeat.observability.traces.service.EntityTraceQueryService.TraceDetailQuery; +import org.apache.hertzbeat.warehouse.query.admission.ObservabilityQueryAdmissionService; import org.springframework.data.domain.Page; import org.springframework.http.ResponseEntity; import org.springframework.util.StringUtils; @@ -54,6 +55,7 @@ import org.springframework.web.bind.annotation.RestController; public class TraceQueryController { private final EntityTraceQueryService entityTraceQueryService; + private final ObservabilityQueryAdmissionService queryAdmissionService; @GetMapping("/list") @Operation(summary = "Query traces with entity context and canonical resource filters") @@ -82,10 +84,11 @@ public class TraceQueryController { String workspaceId = trustedWorkspaceId(); ScopedFilters scopedFilters = scopeFilters( entityId, entityType, collectorId, instance, endpoint, resourceFilter, attributeFilter); - Page<TraceListItemDto> page = entityTraceQueryService.queryTraceList( - workspaceId, entityId, start, end, traceId, errorOnly, serviceName, serviceNamespace, environment, - scopedFilters.resourceFilter(), operationName, minDurationMs, maxDurationMs, pageIndex, pageSize, - hideInternal, spanScope, scopedFilters.attributeFilter()); + Page<TraceListItemDto> page = queryAdmissionService.execute("traces", + () -> entityTraceQueryService.queryTraceList( + workspaceId, entityId, start, end, traceId, errorOnly, serviceName, serviceNamespace, environment, + scopedFilters.resourceFilter(), operationName, minDurationMs, maxDurationMs, pageIndex, pageSize, + hideInternal, spanScope, scopedFilters.attributeFilter())); return ResponseEntity.ok(Message.success(page)); } @@ -114,10 +117,11 @@ public class TraceQueryController { String workspaceId = trustedWorkspaceId(); ScopedFilters scopedFilters = scopeFilters( entityId, entityType, collectorId, instance, endpoint, resourceFilter, attributeFilter); - return ResponseEntity.ok(Message.success(entityTraceQueryService.getTraceOverview( - workspaceId, entityId, start, end, traceId, errorOnly, serviceName, serviceNamespace, environment, - scopedFilters.resourceFilter(), operationName, minDurationMs, maxDurationMs, hideInternal, spanScope, - scopedFilters.attributeFilter()))); + return ResponseEntity.ok(Message.success(queryAdmissionService.execute("traces", + () -> entityTraceQueryService.getTraceOverview( + workspaceId, entityId, start, end, traceId, errorOnly, serviceName, serviceNamespace, environment, + scopedFilters.resourceFilter(), operationName, minDurationMs, maxDurationMs, hideInternal, + spanScope, scopedFilters.attributeFilter())))); } @GetMapping("/stats/group-by") @@ -149,10 +153,11 @@ public class TraceQueryController { String workspaceId = trustedWorkspaceId(); ScopedFilters scopedFilters = scopeFilters( entityId, entityType, collectorId, instance, endpoint, resourceFilter, attributeFilter); - return ResponseEntity.ok(Message.success(entityTraceQueryService.getTraceGroupByStats( - workspaceId, entityId, start, end, traceId, errorOnly, serviceName, serviceNamespace, environment, - scopedFilters.resourceFilter(), operationName, minDurationMs, maxDurationMs, groupBy, limit, orderBy, - minCount, hideInternal, spanScope, scopedFilters.attributeFilter()))); + return ResponseEntity.ok(Message.success(queryAdmissionService.execute("traces", + () -> entityTraceQueryService.getTraceGroupByStats( + workspaceId, entityId, start, end, traceId, errorOnly, serviceName, serviceNamespace, environment, + scopedFilters.resourceFilter(), operationName, minDurationMs, maxDurationMs, groupBy, limit, + orderBy, minCount, hideInternal, spanScope, scopedFilters.attributeFilter())))); } @GetMapping("/{traceId}") @@ -186,7 +191,8 @@ public class TraceQueryController { TraceDetailQuery query = detailQuery( entityId, traceId, spanId, start, end, serviceName, serviceNamespace, environment, collectorId, instance, endpoint, resourceFilter, attributeFilter, minDurationMs, maxDurationMs); - return ResponseEntity.ok(Message.success(entityTraceQueryService.getTraceDetail(workspaceId, query))); + return ResponseEntity.ok(Message.success(queryAdmissionService.execute("traces", + () -> entityTraceQueryService.getTraceDetail(workspaceId, query)))); } @GetMapping("/{traceId}/spans") @@ -224,7 +230,8 @@ public class TraceQueryController { TraceDetailQuery query = detailQuery( entityId, traceId, spanId, start, end, serviceName, serviceNamespace, environment, collectorId, instance, endpoint, resourceFilter, attributeFilter, minDurationMs, maxDurationMs); - TraceDetailDto detail = entityTraceQueryService.getTraceDetail(workspaceId, query); + TraceDetailDto detail = queryAdmissionService.execute("traces", + () -> entityTraceQueryService.getTraceDetail(workspaceId, query)); return ResponseEntity.ok(Message.success(detail == null ? List.of() : detail.getSpans())); } diff --git a/hertzbeat-observability/src/test/java/org/apache/hertzbeat/observability/ingestion/controller/OtlpIngestionControllerTest.java b/hertzbeat-observability/src/test/java/org/apache/hertzbeat/observability/ingestion/controller/OtlpIngestionControllerTest.java index 948bcd60d9..1f7df02e61 100644 --- a/hertzbeat-observability/src/test/java/org/apache/hertzbeat/observability/ingestion/controller/OtlpIngestionControllerTest.java +++ b/hertzbeat-observability/src/test/java/org/apache/hertzbeat/observability/ingestion/controller/OtlpIngestionControllerTest.java @@ -20,6 +20,8 @@ package org.apache.hertzbeat.observability.ingestion.controller; import static org.mockito.ArgumentMatchers.argThat; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; + +import java.time.Duration; import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.get; import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.jsonPath; import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status; @@ -38,6 +40,7 @@ import org.apache.hertzbeat.common.observability.gateway.AuthTokenRequestContext import org.apache.hertzbeat.observability.ingestion.red.OtlpIngestionRedSummaryService; import org.apache.hertzbeat.observability.ingestion.service.OtlpIngestionWorkspaceService; import org.apache.hertzbeat.observability.metrics.service.CollectorScopedMetricsQueryService; +import org.apache.hertzbeat.warehouse.query.admission.ObservabilityQueryAdmissionService; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -65,7 +68,8 @@ class OtlpIngestionControllerTest { void setUp() { AuthTokenRequestContext.bindWorkspaceId("team-a"); OtlpIngestionController controller = new OtlpIngestionController( - otlpIngestionWorkspaceService, otlpIngestionRedSummaryService, collectorScopedMetricsQueryService); + otlpIngestionWorkspaceService, otlpIngestionRedSummaryService, collectorScopedMetricsQueryService, + queryAdmissionService()); this.mockMvc = MockMvcBuilders.standaloneSetup(controller).build(); } @@ -74,6 +78,10 @@ class OtlpIngestionControllerTest { AuthTokenRequestContext.clear(); } + private ObservabilityQueryAdmissionService queryAdmissionService() { + return new ObservabilityQueryAdmissionService(8, 8, 8, 4, 8, Duration.ofMillis(100)); + } + @Test void shouldReturnWrappedOverviewPayload() throws Exception { OtlpIngestionOverviewDto overview = new OtlpIngestionOverviewDto( diff --git a/hertzbeat-observability/src/test/java/org/apache/hertzbeat/observability/logs/controller/LogQueryControllerTest.java b/hertzbeat-observability/src/test/java/org/apache/hertzbeat/observability/logs/controller/LogQueryControllerTest.java index 6a469713b9..3599a5b5ba 100644 --- a/hertzbeat-observability/src/test/java/org/apache/hertzbeat/observability/logs/controller/LogQueryControllerTest.java +++ b/hertzbeat-observability/src/test/java/org/apache/hertzbeat/observability/logs/controller/LogQueryControllerTest.java @@ -33,6 +33,7 @@ import static org.springframework.test.web.servlet.result.MockMvcResultMatchers. import java.lang.reflect.InvocationTargetException; import java.lang.reflect.Method; import java.lang.reflect.Proxy; +import java.time.Duration; import java.time.Instant; import java.time.LocalDateTime; import java.time.ZoneId; @@ -50,6 +51,7 @@ import org.apache.hertzbeat.common.entity.manager.ObserveEntity; import org.apache.hertzbeat.common.observability.gateway.AuthTokenRequestContext; import org.apache.hertzbeat.common.observability.gateway.ObservabilityWorkspaceQueryGateway; import org.apache.hertzbeat.observability.logs.service.impl.LogQueryServiceImpl; +import org.apache.hertzbeat.warehouse.query.admission.ObservabilityQueryAdmissionService; import org.apache.hertzbeat.warehouse.store.history.tsdb.HistoryDataReader; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; @@ -81,7 +83,8 @@ class LogQueryControllerTest { void setUp() { AuthTokenRequestContext.bindWorkspaceId("default"); this.logQueryController = new LogQueryController( - new LogQueryServiceImpl(List.of(workspaceCompatibleReader(historyDataReader)))); + new LogQueryServiceImpl(List.of(workspaceCompatibleReader(historyDataReader))), + queryAdmissionService()); this.mockMvc = MockMvcBuilders.standaloneSetup(logQueryController).build(); } @@ -325,7 +328,7 @@ class LogQueryControllerTest { ObservabilityWorkspaceQueryGateway workspaceQueryGateway = org.mockito.Mockito.mock(ObservabilityWorkspaceQueryGateway.class); this.logQueryController = new LogQueryController( new LogQueryServiceImpl(List.of(workspaceCompatibleReader(historyDataReader)), - Optional.of(workspaceQueryGateway))); + Optional.of(workspaceQueryGateway)), queryAdmissionService()); this.mockMvc = MockMvcBuilders.standaloneSetup(logQueryController).build(); EntityIdentity serviceName = EntityIdentity.builder() .entityId(42L) @@ -400,7 +403,7 @@ class LogQueryControllerTest { ObservabilityWorkspaceQueryGateway workspaceQueryGateway = org.mockito.Mockito.mock(ObservabilityWorkspaceQueryGateway.class); this.logQueryController = new LogQueryController( new LogQueryServiceImpl(List.of(workspaceCompatibleReader(historyDataReader)), - Optional.of(workspaceQueryGateway))); + Optional.of(workspaceQueryGateway)), queryAdmissionService()); this.mockMvc = MockMvcBuilders.standaloneSetup(logQueryController).build(); when(workspaceQueryGateway.findEntityById("default", 42L)).thenReturn(Optional.of(ObserveEntity.builder() .id(42L) @@ -653,7 +656,7 @@ class LogQueryControllerTest { ObservabilityWorkspaceQueryGateway workspaceQueryGateway = org.mockito.Mockito.mock(ObservabilityWorkspaceQueryGateway.class); this.logQueryController = new LogQueryController( new LogQueryServiceImpl(List.of(workspaceCompatibleReader(historyDataReader)), - Optional.of(workspaceQueryGateway))); + Optional.of(workspaceQueryGateway)), queryAdmissionService()); this.mockMvc = MockMvcBuilders.standaloneSetup(logQueryController).build(); when(workspaceQueryGateway.findEntityById("default", 42L)).thenReturn(Optional.of(ObserveEntity.builder() .id(42L) @@ -1520,7 +1523,7 @@ class LogQueryControllerTest { MockMvc fallbackMockMvc = MockMvcBuilders .standaloneSetup(new LogQueryController(new LogQueryServiceImpl(List.of( workspaceCompatibleReader(historyDataReader), - workspaceCompatibleReader(secondaryHistoryDataReader))))) + workspaceCompatibleReader(secondaryHistoryDataReader))), queryAdmissionService())) .build(); when(historyDataReader.countLogsByMultipleConditions(any(), any(), any(), any(), any(), any(), any(), @@ -1958,4 +1961,7 @@ class LogQueryControllerTest { org.mockito.ArgumentMatchers.<Map<String, String>>any(), eq("resource:service.version")); } + private ObservabilityQueryAdmissionService queryAdmissionService() { + return new ObservabilityQueryAdmissionService(8, 8, 8, 4, 8, Duration.ofMillis(100)); + } } diff --git a/hertzbeat-observability/src/test/java/org/apache/hertzbeat/observability/logs/controller/LogQueryWorkspaceControllerContractTest.java b/hertzbeat-observability/src/test/java/org/apache/hertzbeat/observability/logs/controller/LogQueryWorkspaceControllerContractTest.java index 8829515b58..74d9629a2e 100644 --- a/hertzbeat-observability/src/test/java/org/apache/hertzbeat/observability/logs/controller/LogQueryWorkspaceControllerContractTest.java +++ b/hertzbeat-observability/src/test/java/org/apache/hertzbeat/observability/logs/controller/LogQueryWorkspaceControllerContractTest.java @@ -19,6 +19,8 @@ package org.apache.hertzbeat.observability.logs.controller; import static org.assertj.core.api.Assertions.assertThat; import static org.mockito.Mockito.mock; + +import java.time.Duration; import static org.mockito.Mockito.mockingDetails; import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status; @@ -27,6 +29,7 @@ import java.util.Map; import org.apache.hertzbeat.common.observability.gateway.AuthTokenRequestContext; import org.apache.hertzbeat.common.support.exception.TelemetryStorageUnavailableException; import org.apache.hertzbeat.observability.logs.service.LogQueryService; +import org.apache.hertzbeat.warehouse.query.admission.ObservabilityQueryAdmissionService; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -56,7 +59,8 @@ class LogQueryWorkspaceControllerContractTest { return null; }); mockMvc = org.springframework.test.web.servlet.setup.MockMvcBuilders - .standaloneSetup(new LogQueryController(logQueryService)) + .standaloneSetup(new LogQueryController(logQueryService, + new ObservabilityQueryAdmissionService(8, 8, 8, 4, 8, Duration.ofMillis(100)))) .setControllerAdvice(new UnavailableAdvice()) .build(); } diff --git a/hertzbeat-observability/src/test/java/org/apache/hertzbeat/observability/traces/controller/TraceQueryControllerTest.java b/hertzbeat-observability/src/test/java/org/apache/hertzbeat/observability/traces/controller/TraceQueryControllerTest.java index 183eca98c4..380ca0e7d8 100644 --- a/hertzbeat-observability/src/test/java/org/apache/hertzbeat/observability/traces/controller/TraceQueryControllerTest.java +++ b/hertzbeat-observability/src/test/java/org/apache/hertzbeat/observability/traces/controller/TraceQueryControllerTest.java @@ -22,6 +22,8 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoInteractions; import static org.mockito.Mockito.when; + +import java.time.Duration; import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.get; import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.jsonPath; import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status; @@ -34,6 +36,7 @@ import org.apache.hertzbeat.common.observability.gateway.AuthTokenRequestContext import org.apache.hertzbeat.common.support.exception.TelemetryStorageUnavailableException; import org.apache.hertzbeat.observability.traces.service.EntityTraceQueryService; import org.apache.hertzbeat.observability.traces.service.EntityTraceQueryService.TraceDetailQuery; +import org.apache.hertzbeat.warehouse.query.admission.ObservabilityQueryAdmissionService; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -56,7 +59,8 @@ class TraceQueryControllerTest { @BeforeEach void setUp() { AuthTokenRequestContext.bindWorkspaceId("team-a"); - TraceQueryController controller = new TraceQueryController(entityTraceQueryService); + TraceQueryController controller = new TraceQueryController(entityTraceQueryService, + new ObservabilityQueryAdmissionService(8, 8, 8, 4, 8, Duration.ofMillis(100))); this.mockMvc = MockMvcBuilders.standaloneSetup(controller).build(); } diff --git a/hertzbeat-startup/src/main/resources/application.yml b/hertzbeat-startup/src/main/resources/application.yml index 6a1e0b87da..2be40287c9 100644 --- a/hertzbeat-startup/src/main/resources/application.yml +++ b/hertzbeat-startup/src/main/resources/application.yml @@ -368,6 +368,19 @@ hertzbeat: max-concurrent-requests: ${HERTZBEAT_OTLP_MAX_CONCURRENT_REQUESTS:32} queue-capacity: ${HERTZBEAT_OTLP_QUEUE_CAPACITY:32} max-queue-wait: ${HERTZBEAT_OTLP_MAX_QUEUE_WAIT:100ms} + observability: + query: + admission: + metrics: + max-concurrent-requests: ${HERTZBEAT_METRICS_QUERY_MAX_CONCURRENT_REQUESTS:8} + logs: + max-concurrent-requests: ${HERTZBEAT_LOGS_QUERY_MAX_CONCURRENT_REQUESTS:8} + traces: + max-concurrent-requests: ${HERTZBEAT_TRACES_QUERY_MAX_CONCURRENT_REQUESTS:8} + topology: + max-concurrent-requests: ${HERTZBEAT_TOPOLOGY_QUERY_MAX_CONCURRENT_REQUESTS:4} + queue-capacity: ${HERTZBEAT_QUERY_QUEUE_CAPACITY:8} + max-queue-wait: ${HERTZBEAT_QUERY_MAX_QUEUE_WAIT:100ms} # Optional virtual-thread overrides. Remove this whole block to use built-in defaults. vthreads: enabled: true diff --git a/hertzbeat-warehouse/src/main/java/org/apache/hertzbeat/warehouse/query/admission/ObservabilityQueryAdmissionException.java b/hertzbeat-warehouse/src/main/java/org/apache/hertzbeat/warehouse/query/admission/ObservabilityQueryAdmissionException.java new file mode 100644 index 0000000000..d4a1dd1e17 --- /dev/null +++ b/hertzbeat-warehouse/src/main/java/org/apache/hertzbeat/warehouse/query/admission/ObservabilityQueryAdmissionException.java @@ -0,0 +1,52 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hertzbeat.warehouse.query.admission; + +/** Signals an overloaded or cancelled observability query before execution. */ +public class ObservabilityQueryAdmissionException extends RuntimeException { + + private final String signal; + private final Reason reason; + + public ObservabilityQueryAdmissionException(String signal, Reason reason, Throwable cause) { + super(message(signal, reason), cause); + this.signal = signal; + this.reason = reason; + } + + public String getSignal() { + return signal; + } + + public Reason getReason() { + return reason; + } + + private static String message(String signal, Reason reason) { + return switch (reason) { + case OVERLOADED -> "Observability " + signal + " query capacity is exhausted."; + case CANCELLED -> "Observability " + signal + " query admission was cancelled."; + }; + } + + /** Query admission failure classification used by the HTTP exception boundary. */ + public enum Reason { + OVERLOADED, + CANCELLED + } +} diff --git a/hertzbeat-warehouse/src/main/java/org/apache/hertzbeat/warehouse/query/admission/ObservabilityQueryAdmissionService.java b/hertzbeat-warehouse/src/main/java/org/apache/hertzbeat/warehouse/query/admission/ObservabilityQueryAdmissionService.java new file mode 100644 index 0000000000..1f20a78449 --- /dev/null +++ b/hertzbeat-warehouse/src/main/java/org/apache/hertzbeat/warehouse/query/admission/ObservabilityQueryAdmissionService.java @@ -0,0 +1,133 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hertzbeat.warehouse.query.admission; + +import java.time.Duration; +import java.util.Map; +import java.util.Objects; +import java.util.concurrent.Semaphore; +import java.util.concurrent.TimeUnit; +import java.util.function.Supplier; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.beans.factory.annotation.Value; +import org.springframework.stereotype.Service; + +/** + * Isolates observability reads in bounded per-signal lanes without moving work off the request thread. + * + * <p>Executing inline preserves request-scoped security context and avoids a second executor and timeout layer. + * Backend deadlines remain owned by the storage-specific query guard.</p> + */ +@Service +public class ObservabilityQueryAdmissionService { + + private final Map<String, Lane> lanes; + + @Autowired + public ObservabilityQueryAdmissionService( + @Value("${hertzbeat.observability.query.admission.metrics.max-concurrent-requests:8}") + int metricsMaxConcurrentRequests, + @Value("${hertzbeat.observability.query.admission.logs.max-concurrent-requests:8}") + int logsMaxConcurrentRequests, + @Value("${hertzbeat.observability.query.admission.traces.max-concurrent-requests:8}") + int tracesMaxConcurrentRequests, + @Value("${hertzbeat.observability.query.admission.topology.max-concurrent-requests:4}") + int topologyMaxConcurrentRequests, + @Value("${hertzbeat.observability.query.admission.queue-capacity:8}") + int queueCapacity, + @Value("${hertzbeat.observability.query.admission.max-queue-wait:100ms}") + Duration maxQueueWait) { + this.lanes = Map.of( + "metrics", new Lane("metrics", metricsMaxConcurrentRequests, queueCapacity, maxQueueWait), + "logs", new Lane("logs", logsMaxConcurrentRequests, queueCapacity, maxQueueWait), + "traces", new Lane("traces", tracesMaxConcurrentRequests, queueCapacity, maxQueueWait), + "topology", new Lane("topology", topologyMaxConcurrentRequests, queueCapacity, maxQueueWait)); + } + + /** Execute a query in the lane dedicated to its signal. */ + public <T> T execute(String signal, Supplier<T> operation) { + Lane lane = lanes.get(signal); + if (lane == null) { + throw new IllegalArgumentException("Unsupported observability query signal: " + signal); + } + return lane.execute(Objects.requireNonNull(operation, "operation")); + } + + private static final class Lane { + + private final String signal; + private final Semaphore admittedRequests; + private final Semaphore activeRequests; + private final long maxQueueWaitNanos; + + private Lane(String signal, int maxConcurrentRequests, int queueCapacity, Duration maxQueueWait) { + if (maxConcurrentRequests <= 0) { + throw new IllegalArgumentException("maxConcurrentRequests must be greater than zero"); + } + if (queueCapacity < 0) { + throw new IllegalArgumentException("queueCapacity must not be negative"); + } + long admissionCapacity = (long) maxConcurrentRequests + queueCapacity; + if (admissionCapacity > Integer.MAX_VALUE) { + throw new IllegalArgumentException("Combined admission capacity exceeds the supported limit"); + } + this.signal = signal; + this.admittedRequests = new Semaphore((int) admissionCapacity, true); + this.activeRequests = new Semaphore(maxConcurrentRequests, true); + this.maxQueueWaitNanos = durationNanos(maxQueueWait); + } + + private <T> T execute(Supplier<T> operation) { + if (!admittedRequests.tryAcquire()) { + throw failure(ObservabilityQueryAdmissionException.Reason.OVERLOADED, null); + } + boolean active = false; + try { + if (!activeRequests.tryAcquire(maxQueueWaitNanos, TimeUnit.NANOSECONDS)) { + throw failure(ObservabilityQueryAdmissionException.Reason.OVERLOADED, null); + } + active = true; + return operation.get(); + } catch (InterruptedException exception) { + Thread.currentThread().interrupt(); + throw failure(ObservabilityQueryAdmissionException.Reason.CANCELLED, exception); + } finally { + if (active) { + activeRequests.release(); + } + admittedRequests.release(); + } + } + + private ObservabilityQueryAdmissionException failure( + ObservabilityQueryAdmissionException.Reason reason, Throwable cause) { + return new ObservabilityQueryAdmissionException(signal, reason, cause); + } + + private static long durationNanos(Duration duration) { + if (duration == null || duration.isNegative()) { + throw new IllegalArgumentException("maxQueueWait must be zero or positive"); + } + try { + return duration.toNanos(); + } catch (ArithmeticException ignored) { + return Long.MAX_VALUE; + } + } + } +} diff --git a/hertzbeat-warehouse/src/test/java/org/apache/hertzbeat/warehouse/query/admission/ObservabilityQueryAdmissionServiceTest.java b/hertzbeat-warehouse/src/test/java/org/apache/hertzbeat/warehouse/query/admission/ObservabilityQueryAdmissionServiceTest.java new file mode 100644 index 0000000000..6dea4c3f04 --- /dev/null +++ b/hertzbeat-warehouse/src/test/java/org/apache/hertzbeat/warehouse/query/admission/ObservabilityQueryAdmissionServiceTest.java @@ -0,0 +1,182 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hertzbeat.warehouse.query.admission; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.time.Duration; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import org.apache.hertzbeat.common.observability.gateway.AuthTokenRequestContext; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +class ObservabilityQueryAdmissionServiceTest { + + @AfterEach + void clearContext() { + AuthTokenRequestContext.clear(); + } + + @Test + void preservesAllRequestContextOnTheCallingThread() { + ObservabilityQueryAdmissionService service = service(1, 0, Duration.ZERO); + AuthTokenRequestContext.bindWorkspaceId("workspace-a"); + AuthTokenRequestContext.bindAuthenticatedWorkspaceId("workspace-authenticated"); + AuthTokenRequestContext.bindCollectorId("collector-a"); + Thread requestThread = Thread.currentThread(); + + assertEquals("workspace-a|workspace-authenticated|collector-a|true", service.execute("logs", () -> + AuthTokenRequestContext.currentWorkspaceId() + "|" + + AuthTokenRequestContext.currentAuthenticatedWorkspaceId() + "|" + + AuthTokenRequestContext.currentCollectorId() + "|" + + (Thread.currentThread() == requestThread))); + } + + @Test + void rejectsFullLaneAndRecoversAfterActiveQueryCompletes() throws Exception { + ObservabilityQueryAdmissionService service = service(1, 0, Duration.ZERO); + CountDownLatch queryStarted = new CountDownLatch(1); + CountDownLatch releaseQuery = new CountDownLatch(1); + Thread activeRequest = startBlockingQuery(service, "metrics", queryStarted, releaseQuery); + assertTrue(queryStarted.await(1, TimeUnit.SECONDS)); + + ObservabilityQueryAdmissionException exception = assertThrows( + ObservabilityQueryAdmissionException.class, + () -> service.execute("metrics", () -> "rejected")); + + assertEquals(ObservabilityQueryAdmissionException.Reason.OVERLOADED, exception.getReason()); + assertEquals("metrics", exception.getSignal()); + releaseQuery.countDown(); + activeRequest.join(1_000); + assertFalse(activeRequest.isAlive()); + assertEquals("recovered", service.execute("metrics", () -> "recovered")); + } + + @Test + void boundsQueueWaitAndNeverRunsExpiredOperation() throws Exception { + ObservabilityQueryAdmissionService service = service(1, 1, Duration.ofMillis(50)); + CountDownLatch queryStarted = new CountDownLatch(1); + CountDownLatch releaseQuery = new CountDownLatch(1); + Thread activeRequest = startBlockingQuery(service, "logs", queryStarted, releaseQuery); + assertTrue(queryStarted.await(1, TimeUnit.SECONDS)); + AtomicReference<String> result = new AtomicReference<>(); + long startedAt = System.nanoTime(); + + ObservabilityQueryAdmissionException exception = assertThrows( + ObservabilityQueryAdmissionException.class, + () -> service.execute("logs", () -> { + result.set("ran"); + return "late"; + })); + long elapsedMillis = Duration.ofNanos(System.nanoTime() - startedAt).toMillis(); + + assertEquals(ObservabilityQueryAdmissionException.Reason.OVERLOADED, exception.getReason()); + assertTrue(elapsedMillis >= 25 && elapsedMillis < 500); + assertNull(result.get()); + releaseQuery.countDown(); + activeRequest.join(1_000); + } + + @Test + void isolatesMetricsLogsTracesAndTopologyLanes() throws Exception { + ObservabilityQueryAdmissionService service = service(1, 0, Duration.ZERO); + CountDownLatch queryStarted = new CountDownLatch(1); + CountDownLatch releaseQuery = new CountDownLatch(1); + Thread activeRequest = startBlockingQuery(service, "metrics", queryStarted, releaseQuery); + assertTrue(queryStarted.await(1, TimeUnit.SECONDS)); + + assertEquals("logs", service.execute("logs", () -> "logs")); + assertEquals("traces", service.execute("traces", () -> "traces")); + assertEquals("topology", service.execute("topology", () -> "topology")); + + releaseQuery.countDown(); + activeRequest.join(1_000); + } + + @Test + void releasesCapacityWhenQueryFails() { + ObservabilityQueryAdmissionService service = service(1, 0, Duration.ZERO); + + assertThrows(IllegalStateException.class, + () -> service.execute("traces", () -> { + throw new IllegalStateException("query failed"); + })); + + assertEquals("recovered", service.execute("traces", () -> "recovered")); + } + + @Test + void interruptionWhileWaitingPreservesInterruptAndReleasesQueueCapacity() throws Exception { + ObservabilityQueryAdmissionService service = service(1, 1, Duration.ofSeconds(5)); + CountDownLatch queryStarted = new CountDownLatch(1); + CountDownLatch releaseQuery = new CountDownLatch(1); + Thread activeRequest = startBlockingQuery(service, "topology", queryStarted, releaseQuery); + assertTrue(queryStarted.await(1, TimeUnit.SECONDS)); + AtomicReference<Throwable> failure = new AtomicReference<>(); + AtomicReference<Boolean> interrupted = new AtomicReference<>(false); + Thread queuedRequest = Thread.ofVirtual().start(() -> { + try { + service.execute("topology", () -> "unexpected"); + } catch (Throwable throwable) { + failure.set(throwable); + interrupted.set(Thread.currentThread().isInterrupted()); + } + }); + + queuedRequest.interrupt(); + queuedRequest.join(1_000); + + ObservabilityQueryAdmissionException exception = + (ObservabilityQueryAdmissionException) failure.get(); + assertEquals(ObservabilityQueryAdmissionException.Reason.CANCELLED, exception.getReason()); + assertTrue(interrupted.get()); + releaseQuery.countDown(); + activeRequest.join(1_000); + assertEquals("recovered", service.execute("topology", () -> "recovered")); + } + + private ObservabilityQueryAdmissionService service( + int maxConcurrentRequests, int queueCapacity, Duration maxQueueWait) { + return new ObservabilityQueryAdmissionService( + maxConcurrentRequests, maxConcurrentRequests, maxConcurrentRequests, maxConcurrentRequests, + queueCapacity, maxQueueWait); + } + + private Thread startBlockingQuery( + ObservabilityQueryAdmissionService service, + String signal, + CountDownLatch queryStarted, + CountDownLatch releaseQuery) { + return Thread.ofVirtual().start(() -> service.execute(signal, () -> { + queryStarted.countDown(); + try { + releaseQuery.await(); + } catch (InterruptedException exception) { + Thread.currentThread().interrupt(); + throw new IllegalStateException("Test query interrupted", exception); + } + return "completed"; + })); + } +} --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
