From 5ee196b0782f80391d0ae525873cb0fec284600d Mon Sep 17 00:00:00 2001 From: aadityaraj7769 Date: Tue, 19 May 2026 11:08:35 +0530 Subject: [PATCH] Refactor otel metrics integration for load balancer strategies. Introduce HostStatus enum for improved health status tracking. Update Degrader and Relative load balancer strategy metrics providers to utilize HostStatus for reporting unhealthy and quarantined hosts. Enhance tests to validate new metrics behavior and ensure proper listener registration for tracker clients. Maintain backward compatibility while improving telemetry accuracy. Co-authored-by: Cursor --- .../linkedin/d2/balancer/D2ClientBuilder.java | 30 +- .../linkedin/d2/balancer/D2ClientConfig.java | 196 +++++++ .../clients/PerCallDurationListener.java | 50 ++ .../clients/PerCallDurationSemantics.java | 29 ++ .../d2/balancer/clients/TrackerClient.java | 29 ++ .../balancer/clients/TrackerClientImpl.java | 70 ++- .../simple/SimpleLoadBalancerState.java | 5 +- .../LoadBalancerStrategyFactory.java | 27 + .../d2/balancer/strategies/SchemeAware.java | 46 ++ ...DegraderLoadBalancerStrategyFactoryV3.java | 35 +- .../DegraderLoadBalancerStrategyV3.java | 123 ++++- .../RelativeLoadBalancerStrategy.java | 14 +- .../RelativeLoadBalancerStrategyFactory.java | 34 +- .../strategies/relative/StateUpdater.java | 151 ++++++ .../linkedin/d2/jmx/D2ClientJmxManager.java | 7 + ...BalancerStrategyV3OtelMetricsProvider.java | 22 + .../java/com/linkedin/d2/jmx/HostStatus.java | 32 ++ ...adBalancerStrategyOtelMetricsProvider.java | 42 ++ ...BalancerStrategyV3OtelMetricsProvider.java | 53 ++ ...adBalancerStrategyOtelMetricsProvider.java | 62 +++ ...adBalancerStrategyOtelMetricsProvider.java | 39 ++ .../d2/balancer/D2ClientBuilderTest.java | 28 + ...ClientConfigOtelForwarderWarnOnceTest.java | 95 ++++ .../clients/TrackerClientImplTest.java | 493 +++++++++++++++++- ...BalancerStrategyV3OtelIntegrationTest.java | 398 ++++++++++++++ ...adBalancerStrategyOtelIntegrationTest.java | 398 ++++++++++++++ .../AbstractRecordingOtelMetricsProvider.java | 242 +++++++++ .../d2/jmx/D2ClientJmxManagerTest.java | 25 + ...ncerStrategyV3OtelMetricsProviderTest.java | 281 ++++++++++ ...lancerStrategyOtelMetricsProviderTest.java | 212 ++++++++ ...BalancerStrategyV3OtelMetricsProvider.java | 43 ++ ...adBalancerStrategyOtelMetricsProvider.java | 49 ++ 32 files changed, 3331 insertions(+), 29 deletions(-) create mode 100644 d2/src/main/java/com/linkedin/d2/balancer/clients/PerCallDurationListener.java create mode 100644 d2/src/main/java/com/linkedin/d2/balancer/clients/PerCallDurationSemantics.java create mode 100644 d2/src/main/java/com/linkedin/d2/balancer/strategies/SchemeAware.java create mode 100644 d2/src/main/java/com/linkedin/d2/jmx/DegraderLoadBalancerStrategyV3OtelMetricsProvider.java create mode 100644 d2/src/main/java/com/linkedin/d2/jmx/HostStatus.java create mode 100644 d2/src/main/java/com/linkedin/d2/jmx/LoadBalancerStrategyOtelMetricsProvider.java create mode 100644 d2/src/main/java/com/linkedin/d2/jmx/NoOpDegraderLoadBalancerStrategyV3OtelMetricsProvider.java create mode 100644 d2/src/main/java/com/linkedin/d2/jmx/NoOpRelativeLoadBalancerStrategyOtelMetricsProvider.java create mode 100644 d2/src/main/java/com/linkedin/d2/jmx/RelativeLoadBalancerStrategyOtelMetricsProvider.java create mode 100644 d2/src/test/java/com/linkedin/d2/balancer/D2ClientConfigOtelForwarderWarnOnceTest.java create mode 100644 d2/src/test/java/com/linkedin/d2/balancer/strategies/degrader/DegraderLoadBalancerStrategyV3OtelIntegrationTest.java create mode 100644 d2/src/test/java/com/linkedin/d2/balancer/strategies/relative/RelativeLoadBalancerStrategyOtelIntegrationTest.java create mode 100644 d2/src/test/java/com/linkedin/d2/jmx/AbstractRecordingOtelMetricsProvider.java create mode 100644 d2/src/test/java/com/linkedin/d2/jmx/DegraderLoadBalancerStrategyV3OtelMetricsProviderTest.java create mode 100644 d2/src/test/java/com/linkedin/d2/jmx/RelativeLoadBalancerStrategyOtelMetricsProviderTest.java create mode 100644 d2/src/test/java/com/linkedin/d2/jmx/TestDegraderLoadBalancerStrategyV3OtelMetricsProvider.java create mode 100644 d2/src/test/java/com/linkedin/d2/jmx/TestRelativeLoadBalancerStrategyOtelMetricsProvider.java diff --git a/d2/src/main/java/com/linkedin/d2/balancer/D2ClientBuilder.java b/d2/src/main/java/com/linkedin/d2/balancer/D2ClientBuilder.java index 413587a875..8277bca312 100644 --- a/d2/src/main/java/com/linkedin/d2/balancer/D2ClientBuilder.java +++ b/d2/src/main/java/com/linkedin/d2/balancer/D2ClientBuilder.java @@ -52,11 +52,13 @@ import com.linkedin.d2.discovery.event.ServiceDiscoveryEventEmitter; import com.linkedin.d2.discovery.stores.zk.ZKPersistentConnection; import com.linkedin.d2.discovery.stores.zk.ZooKeeper; -import com.linkedin.d2.jmx.XdsServerMetricsProvider; -import com.linkedin.d2.jmx.XdsClientOtelMetricsProvider; +import com.linkedin.d2.jmx.DegraderLoadBalancerStrategyV3OtelMetricsProvider; import com.linkedin.d2.jmx.JmxManager; -import com.linkedin.d2.xds.XdsClientValidator; import com.linkedin.d2.jmx.NoOpJmxManager; +import com.linkedin.d2.jmx.RelativeLoadBalancerStrategyOtelMetricsProvider; +import com.linkedin.d2.jmx.XdsClientOtelMetricsProvider; +import com.linkedin.d2.jmx.XdsServerMetricsProvider; +import com.linkedin.d2.xds.XdsClientValidator; import com.linkedin.r2.transport.common.TransportClientFactory; import com.linkedin.r2.transport.http.client.HttpClientFactory; import com.linkedin.r2.util.NamedThreadFactory; @@ -251,7 +253,9 @@ public D2Client build() _config.d2CalleeInfoRecorder, _config.enableIndisDownstreamServicesFetcher, _config.indisDownstreamServicesFetchTimeout, - _config.xdsClientOtelMetricsProvider + _config.xdsClientOtelMetricsProvider, + _config.relativeLoadBalancerStrategyOtelMetricsProvider, + _config.degraderLoadBalancerStrategyV3OtelMetricsProvider ); final LoadBalancerWithFacilitiesFactory loadBalancerFactory = (_config.lbWithFacilitiesFactory == null) ? @@ -873,6 +877,18 @@ public D2ClientBuilder setXdsClientOtelMetricsProvider(XdsClientOtelMetricsProvi return this; } + public D2ClientBuilder setRelativeLoadBalancerStrategyOtelMetricsProvider( + RelativeLoadBalancerStrategyOtelMetricsProvider relativeLoadBalancerStrategyOtelMetricsProvider) { + _config.relativeLoadBalancerStrategyOtelMetricsProvider = relativeLoadBalancerStrategyOtelMetricsProvider; + return this; + } + + public D2ClientBuilder setDegraderLoadBalancerStrategyV3OtelMetricsProvider( + DegraderLoadBalancerStrategyV3OtelMetricsProvider degraderLoadBalancerStrategyV3OtelMetricsProvider) { + _config.degraderLoadBalancerStrategyV3OtelMetricsProvider = degraderLoadBalancerStrategyV3OtelMetricsProvider; + return this; + } + public D2ClientBuilder setLoadBalanceStreamException(boolean loadBalanceStreamException) { _config.loadBalanceStreamException = loadBalanceStreamException; return this; @@ -965,7 +981,8 @@ private Map> createDefaultLoadBalancerStr loadBalancerStrategyFactories.putIfAbsent("random", randomStrategyFactory); final DegraderLoadBalancerStrategyFactoryV3 degraderStrategyFactoryV3 = new DegraderLoadBalancerStrategyFactoryV3( - _config.healthCheckOperations, _config._executorService, _config.eventEmitter, Collections.emptyList()); + _config.healthCheckOperations, _config._executorService, _config.eventEmitter, Collections.emptyList(), + _config.degraderLoadBalancerStrategyV3OtelMetricsProvider); loadBalancerStrategyFactories.putIfAbsent("degrader", degraderStrategyFactoryV3); loadBalancerStrategyFactories.putIfAbsent("degraderV2", degraderStrategyFactoryV3); loadBalancerStrategyFactories.putIfAbsent("degraderV3", degraderStrategyFactoryV3); @@ -976,7 +993,8 @@ private Map> createDefaultLoadBalancerStr // TODO: create StateUpdater.LoadBalanceConfig and pass it to the RelativeLoadBalancerStrategyFactory final RelativeLoadBalancerStrategyFactory relativeLoadBalancerStrategyFactory = new RelativeLoadBalancerStrategyFactory( _config._executorService, _config.healthCheckOperations, Collections.emptyList(), _config.eventEmitter, - SystemClock.instance(), _config.loadBalanceStreamException, _config.enableRelativeStrategyDeferredAllocation); + SystemClock.instance(), _config.loadBalanceStreamException, _config.enableRelativeStrategyDeferredAllocation, + _config.relativeLoadBalancerStrategyOtelMetricsProvider); loadBalancerStrategyFactories.putIfAbsent(RelativeLoadBalancerStrategy.RELATIVE_LOAD_BALANCER_STRATEGY_NAME, relativeLoadBalancerStrategyFactory); } diff --git a/d2/src/main/java/com/linkedin/d2/balancer/D2ClientConfig.java b/d2/src/main/java/com/linkedin/d2/balancer/D2ClientConfig.java index 95fdbb4c28..42b68690c3 100644 --- a/d2/src/main/java/com/linkedin/d2/balancer/D2ClientConfig.java +++ b/d2/src/main/java/com/linkedin/d2/balancer/D2ClientConfig.java @@ -46,6 +46,10 @@ import com.linkedin.d2.jmx.NoOpJmxManager; import com.linkedin.d2.jmx.XdsClientOtelMetricsProvider; import com.linkedin.d2.jmx.NoOpXdsClientOtelMetricsProvider; +import com.linkedin.d2.jmx.DegraderLoadBalancerStrategyV3OtelMetricsProvider; +import com.linkedin.d2.jmx.NoOpDegraderLoadBalancerStrategyV3OtelMetricsProvider; +import com.linkedin.d2.jmx.RelativeLoadBalancerStrategyOtelMetricsProvider; +import com.linkedin.d2.jmx.NoOpRelativeLoadBalancerStrategyOtelMetricsProvider; import com.linkedin.r2.transport.common.TransportClientFactory; import io.grpc.netty.shaded.io.netty.handler.ssl.SslContext; import java.time.Duration; @@ -54,14 +58,23 @@ import java.util.concurrent.ExecutorService; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Function; import javax.net.ssl.SSLContext; import javax.net.ssl.SSLParameters; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import static com.linkedin.d2.xds.XdsClientValidator.DEFAULT_MINIMUM_JAVA_VERSION; public class D2ClientConfig { + private static final Logger LOG = LoggerFactory.getLogger(D2ClientConfig.class); + + // Guards so the deprecated-forwarder OTel-NoOp warnings are emitted at most once per JVM each. + private static final AtomicBoolean PRE_XDS_OTEL_FORWARDER_WARNED = new AtomicBoolean(false); + private static final AtomicBoolean PRE_STRATEGY_OTEL_FORWARDER_WARNED = new AtomicBoolean(false); + // default values for some configs, to be shared with other classes public static final String D2_JMX_MANAGER_PREFIX_DEFAULT = "UnknownPrefix"; public static final int DEFAULT_RETRY_LIMIT = 3; @@ -189,6 +202,19 @@ public class D2ClientConfig * Defaults to no-op implementation; can be overridden to enable metric tracking. */ public XdsClientOtelMetricsProvider xdsClientOtelMetricsProvider = new NoOpXdsClientOtelMetricsProvider(); + + /** + * Provider for OpenTelemetry metrics collection for RelativeLoadBalancerStrategy operations. + * Defaults to no-op implementation; can be overridden to enable metric tracking. + */ + public RelativeLoadBalancerStrategyOtelMetricsProvider relativeLoadBalancerStrategyOtelMetricsProvider = new NoOpRelativeLoadBalancerStrategyOtelMetricsProvider(); + + /** + * Provider for OpenTelemetry metrics collection for DegraderLoadBalancerStrategyV3 operations. + * Defaults to no-op implementation; can be overridden to enable metric tracking. + */ + public DegraderLoadBalancerStrategyV3OtelMetricsProvider degraderLoadBalancerStrategyV3OtelMetricsProvider = new NoOpDegraderLoadBalancerStrategyV3OtelMetricsProvider(); + public boolean loadBalanceStreamException = false; public boolean enablePotentialClientsCache = false; public boolean enableRelativeStrategyDeferredAllocation = false; @@ -358,8 +384,25 @@ public D2ClientConfig() enableIndisDownstreamServicesFetcher, indisDownstreamServicesFetchTimeout, new NoOpXdsClientOtelMetricsProvider()); + // Deprecated ctor: all three OTel providers defaulted to NoOp; warn once per JVM. + if (PRE_XDS_OTEL_FORWARDER_WARNED.compareAndSet(false, true)) + { + LOG.warn("Deprecated D2ClientConfig ctor: XDS, Relative LB, and Degrader LB OTel providers default to NoOp. " + + "Use the ctor with all three providers or set them on D2ClientConfig / D2ClientBuilder."); + } } + /** + * Backward-compatible overload preserving the constructor signature that took only an + * {@link XdsClientOtelMetricsProvider}, so that any in-package callers updated for that signature + * continue to compile now that the relative- and degrader-strategy OTel providers have been + * added. New callers should use the constructor that accepts all three OTel metrics providers. + * + * @deprecated Use the constructor that additionally takes + * {@link RelativeLoadBalancerStrategyOtelMetricsProvider} and + * {@link DegraderLoadBalancerStrategyV3OtelMetricsProvider}. + */ + @Deprecated D2ClientConfig(String zkHosts, String xdsServer, String hostName, @@ -445,6 +488,157 @@ public D2ClientConfig() Boolean enableIndisDownstreamServicesFetcher, Duration indisDownstreamServicesFetchTimeout, XdsClientOtelMetricsProvider xdsClientOtelMetricsProvider) + { + this(zkHosts, xdsServer, hostName, zkSessionTimeoutInMs, zkStartupTimeoutInMs, lbWaitTimeout, lbWaitUnit, + flagFile, basePath, fsBasePath, indisFsBasePath, componentFactory, clientFactories, lbWithFacilitiesFactory, + sslContext, grpcSslContext, sslParameters, isSSLEnabled, shutdownAsynchronously, isSymlinkAware, + clientServicesConfig, d2ServicePath, useNewEphemeralStoreWatcher, healthCheckOperations, executorService, + retry, restRetryEnabled, streamRetryEnabled, retryLimit, retryUpdateIntervalMs, retryAggregatedIntervalNum, + warmUp, warmUpTimeoutSeconds, indisWarmUpTimeoutSeconds, warmUpConcurrentRequests, + indisWarmUpConcurrentRequests, downstreamServicesFetcher, indisDownstreamServicesFetcher, + backupRequestsEnabled, backupRequestsStrategyStatsConsumer, + backupRequestsLatencyNotificationInterval, + backupRequestsLatencyNotificationIntervalUnit, + enableBackupRequestsClientAsync, + backupRequestsExecutorService, + emitter, + partitionAccessorRegistry, + zooKeeperDecorator, + enableSaveUriDataOnDisk, + loadBalancerStrategyFactories, + requestTimeoutHandlerEnabled, + sslSessionValidatorFactory, + zkConnection, + startUpExecutorService, + indisStartUpExecutorService, + jmxManager, + d2JmxManagerPrefix, + zookeeperReadWindowMs, + enableRelativeLoadBalancer, + deterministicSubsettingMetadataProvider, + canaryDistributionProvider, + enableClusterFailout, + failoutConfigProviderFactory, + failoutRedirectStrategy, + serviceDiscoveryEventEmitter, + dualReadStateManager, + xdsExecutorService, + xdsStreamReadyTimeout, + dualReadNewLbExecutor, + xdsChannelLoadBalancingPolicy, + xdsChannelLoadBalancingPolicyConfig, + subscribeToUriGlobCollection, + xdsServerMetricsProvider, + loadBalanceStreamException, + enablePotentialClientsCache, + xdsInitialResourceVersionsEnabled, + disableDetectLiRawD2Client, + isLiRawD2Client, + xdsStreamMaxRetryBackoffSeconds, + xdsChannelKeepAliveTimeMins, + xdsMinimumJavaVersion, + actionOnPrecheckFailure, + d2CalleeInfoRecorder, + enableIndisDownstreamServicesFetcher, + indisDownstreamServicesFetchTimeout, + xdsClientOtelMetricsProvider, + new NoOpRelativeLoadBalancerStrategyOtelMetricsProvider(), + new NoOpDegraderLoadBalancerStrategyV3OtelMetricsProvider()); + // Deprecated ctor: Relative/Degrader LB OTel providers defaulted to NoOp; warn once per JVM. + if (PRE_STRATEGY_OTEL_FORWARDER_WARNED.compareAndSet(false, true)) + { + LOG.warn("Deprecated D2ClientConfig ctor: Relative and Degrader LB OTel providers default to NoOp. " + + "Use the ctor that takes both providers or set them on D2ClientConfig / D2ClientBuilder."); + } + } + + D2ClientConfig(String zkHosts, + String xdsServer, + String hostName, + long zkSessionTimeoutInMs, + long zkStartupTimeoutInMs, + long lbWaitTimeout, + TimeUnit lbWaitUnit, + String flagFile, + String basePath, + String fsBasePath, + String indisFsBasePath, + ComponentFactory componentFactory, + Map clientFactories, + LoadBalancerWithFacilitiesFactory lbWithFacilitiesFactory, + SSLContext sslContext, + SslContext grpcSslContext, + SSLParameters sslParameters, + boolean isSSLEnabled, + boolean shutdownAsynchronously, + boolean isSymlinkAware, + Map> clientServicesConfig, + String d2ServicePath, + boolean useNewEphemeralStoreWatcher, + HealthCheckOperations healthCheckOperations, + ScheduledExecutorService executorService, + boolean retry, + boolean restRetryEnabled, + boolean streamRetryEnabled, + int retryLimit, + long retryUpdateIntervalMs, + int retryAggregatedIntervalNum, + boolean warmUp, + int warmUpTimeoutSeconds, + int indisWarmUpTimeoutSeconds, + int warmUpConcurrentRequests, + int indisWarmUpConcurrentRequests, + DownstreamServicesFetcher downstreamServicesFetcher, + DownstreamServicesFetcher indisDownstreamServicesFetcher, + boolean backupRequestsEnabled, + BackupRequestsStrategyStatsConsumer backupRequestsStrategyStatsConsumer, + long backupRequestsLatencyNotificationInterval, + TimeUnit backupRequestsLatencyNotificationIntervalUnit, + boolean enableBackupRequestsClientAsync, + ScheduledExecutorService backupRequestsExecutorService, + EventEmitter emitter, + PartitionAccessorRegistry partitionAccessorRegistry, + Function zooKeeperDecorator, + boolean enableSaveUriDataOnDisk, + Map> loadBalancerStrategyFactories, + boolean requestTimeoutHandlerEnabled, + SslSessionValidatorFactory sslSessionValidatorFactory, + ZKPersistentConnection zkConnection, + ScheduledExecutorService startUpExecutorService, + ScheduledExecutorService indisStartUpExecutorService, + JmxManager jmxManager, + String d2JmxManagerPrefix, + int zookeeperReadWindowMs, + boolean enableRelativeLoadBalancer, + DeterministicSubsettingMetadataProvider deterministicSubsettingMetadataProvider, + CanaryDistributionProvider canaryDistributionProvider, + boolean enableClusterFailout, + FailoutConfigProviderFactory failoutConfigProviderFactory, + FailoutRedirectStrategy failoutRedirectStrategy, + ServiceDiscoveryEventEmitter serviceDiscoveryEventEmitter, + DualReadStateManager dualReadStateManager, + ScheduledExecutorService xdsExecutorService, + Long xdsStreamReadyTimeout, + ExecutorService dualReadNewLbExecutor, + String xdsChannelLoadBalancingPolicy, + Map xdsChannelLoadBalancingPolicyConfig, + boolean subscribeToUriGlobCollection, + XdsServerMetricsProvider xdsServerMetricsProvider, + boolean loadBalanceStreamException, + boolean enablePotentialClientsCache, + boolean xdsInitialResourceVersionsEnabled, + boolean disableDetectLiRawD2Client, + boolean isLiRawD2Client, + Integer xdsStreamMaxRetryBackoffSeconds, + Long xdsChannelKeepAliveTimeMins, + String xdsMinimumJavaVersion, + XdsClientValidator.ActionOnPrecheckFailure actionOnPrecheckFailure, + D2CalleeInfoRecorder d2CalleeInfoRecorder, + Boolean enableIndisDownstreamServicesFetcher, + Duration indisDownstreamServicesFetchTimeout, + XdsClientOtelMetricsProvider xdsClientOtelMetricsProvider, + RelativeLoadBalancerStrategyOtelMetricsProvider relativeLoadBalancerStrategyOtelMetricsProvider, + DegraderLoadBalancerStrategyV3OtelMetricsProvider degraderLoadBalancerStrategyV3OtelMetricsProvider) { this.zkHosts = zkHosts; this.xdsServer = xdsServer; @@ -531,5 +725,7 @@ public D2ClientConfig() this.indisDownstreamServicesFetchTimeout = indisDownstreamServicesFetchTimeout; this.enableIndisDownstreamServicesFetcher = enableIndisDownstreamServicesFetcher; this.xdsClientOtelMetricsProvider = xdsClientOtelMetricsProvider; + this.relativeLoadBalancerStrategyOtelMetricsProvider = relativeLoadBalancerStrategyOtelMetricsProvider; + this.degraderLoadBalancerStrategyV3OtelMetricsProvider = degraderLoadBalancerStrategyV3OtelMetricsProvider; } } diff --git a/d2/src/main/java/com/linkedin/d2/balancer/clients/PerCallDurationListener.java b/d2/src/main/java/com/linkedin/d2/balancer/clients/PerCallDurationListener.java new file mode 100644 index 0000000000..8cbb61a98f --- /dev/null +++ b/d2/src/main/java/com/linkedin/d2/balancer/clients/PerCallDurationListener.java @@ -0,0 +1,50 @@ +/* + Copyright (c) 2026 LinkedIn Corp. + + Licensed 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 com.linkedin.d2.balancer.clients; + +/** + * Per-call duration listener for {@link TrackerClient}. + * + *

Primitive-{@code long} specialization of a {@code BiConsumer} + * to avoid boxing the duration on every completed call. Equivalent in spirit to + * {@link java.util.function.LongConsumer}, extended with a second + * {@link PerCallDurationSemantics} argument that the JDK has no built-in functional interface for. + * + *

Threading and contract. The listener fires synchronously on the thread that + * completes the underlying transport call — for the REST path it runs immediately before the + * wrapped {@code TransportCallback} is invoked, and for the streaming path it runs from the entity + * stream's {@code onDone}/{@code onError} (typically a transport / event-loop thread). Therefore + * implementations MUST NOT block (no I/O, no locks held by slow code paths, no waiting on + * other threads) and MUST NOT throw checked work onto the caller. Any blocking behaviour + * here directly stalls request completion and downstream user callbacks. + * + *

Implementations should also be cheap and allocation-free where possible; they are on the + * per-request hot path. Exceptions thrown from {@link #accept(long, PerCallDurationSemantics)} + * are caught by {@link TrackerClient} and rate-limited in the logs — they will not propagate + * to the wrapped callback, but they will cause the duration sample to be dropped. + */ +@FunctionalInterface +public interface PerCallDurationListener +{ + /** + * Records the perceived per-call duration for one completed call. + * + * @param durationMs duration in milliseconds the server was perceived to have contributed + * @param semantics what {@code durationMs} represents (full round-trip vs. streaming TTFB); + * see {@link PerCallDurationSemantics} + */ + void accept(long durationMs, PerCallDurationSemantics semantics); +} diff --git a/d2/src/main/java/com/linkedin/d2/balancer/clients/PerCallDurationSemantics.java b/d2/src/main/java/com/linkedin/d2/balancer/clients/PerCallDurationSemantics.java new file mode 100644 index 0000000000..5a132a1318 --- /dev/null +++ b/d2/src/main/java/com/linkedin/d2/balancer/clients/PerCallDurationSemantics.java @@ -0,0 +1,29 @@ +/* + Copyright (c) 2026 LinkedIn Corp. + + Licensed 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 com.linkedin.d2.balancer.clients; + +/** + * What a per-call duration measures when passed to + * {@link PerCallDurationListener#accept(long, PerCallDurationSemantics)}. + */ +public enum PerCallDurationSemantics +{ + /** REST: full callback latency. Streaming: transport failure before any response body. */ + FULL_ROUND_TRIP, + + /** Streaming only: request start through first response byte (success or mid-stream error). */ + TIME_TO_FIRST_BYTE +} diff --git a/d2/src/main/java/com/linkedin/d2/balancer/clients/TrackerClient.java b/d2/src/main/java/com/linkedin/d2/balancer/clients/TrackerClient.java index c60f569f4b..b5b9a12cd3 100644 --- a/d2/src/main/java/com/linkedin/d2/balancer/clients/TrackerClient.java +++ b/d2/src/main/java/com/linkedin/d2/balancer/clients/TrackerClient.java @@ -90,4 +90,33 @@ default double getSubsetWeight(int partitionId) * @return CallTracker. */ CallTracker getCallTracker(); + + /** + * Sets a listener that is invoked for each completed call with the duration (in ms) the server + * was perceived to have contributed, and a {@link PerCallDurationSemantics} value describing + * what that duration measures (full round trip vs. streaming TTFB). + * + *

Why this is a setter and not a constructor parameter. A {@link TrackerClient} is + * constructed by the discovery-side {@code TrackerClientFactory} before any + * {@link com.linkedin.d2.balancer.strategies.LoadBalancerStrategy} that consumes its metrics + * exists. The strategy registers its listener when it first observes the tracker client during + * a partition-state update. + * + *

Single-owner contract. A {@link TrackerClient} supports at most one duration + * listener at a time. Calling this method replaces any previously-registered listener (last + * writer wins) — listeners are not composed. + * + *

Threading. The listener fires synchronously on the transport-completion + * thread, immediately before the wrapped {@code TransportCallback} runs (REST path) or from + * the entity stream's {@code onDone}/{@code onError} (streaming path). Implementations + * MUST NOT block — any latency added here directly delays request completion + * and the downstream user callback. See {@link PerCallDurationListener} for the full contract. + * + *

Passing {@code null} resets the listener to a no-op. + * + * @param listener primitive-{@code long} duration sink; may be {@code null} + */ + default void setPerCallDurationListener(PerCallDurationListener listener) + { + } } diff --git a/d2/src/main/java/com/linkedin/d2/balancer/clients/TrackerClientImpl.java b/d2/src/main/java/com/linkedin/d2/balancer/clients/TrackerClientImpl.java index 47cbe75788..f9e43405d1 100644 --- a/d2/src/main/java/com/linkedin/d2/balancer/clients/TrackerClientImpl.java +++ b/d2/src/main/java/com/linkedin/d2/balancer/clients/TrackerClientImpl.java @@ -33,6 +33,7 @@ import com.linkedin.r2.message.stream.StreamResponse; import com.linkedin.r2.message.stream.entitystream.EntityStream; import com.linkedin.r2.message.stream.entitystream.Observer; +import com.linkedin.d2.balancer.util.RateLimitedLogger; import com.linkedin.r2.transport.common.bridge.client.TransportClient; import com.linkedin.r2.transport.common.bridge.common.TransportCallback; import com.linkedin.r2.transport.common.bridge.common.TransportResponse; @@ -70,6 +71,9 @@ public class TrackerClientImpl implements TrackerClient public static final long DEFAULT_CALL_TRACKER_INTERVAL = DegraderLoadBalancerStrategyConfig.DEFAULT_UPDATE_INTERVAL_MS; private static final Logger _log = LoggerFactory.getLogger(TrackerClient.class); + // Rate at which warnings about a misbehaving per-call duration listener are logged. Keeps a + // persistently-throwing OTel provider from flooding the logs on hot request paths. + private static final long PER_CALL_LISTENER_LOG_RATE_MS = 60_000L; private final TransportClient _transportClient; private final Map _partitionData; @@ -77,9 +81,13 @@ public class TrackerClientImpl implements TrackerClient private final Predicate _isErrorStatus; private final ConcurrentMap _subsetWeightMap; private final boolean _doNotLoadBalance; + private final Clock _clock; final CallTracker _callTracker; private boolean _doNotSlowStart; + private volatile PerCallDurationListener _perCallDurationListener = (d, s) -> { + }; + private final RateLimitedLogger _rateLimitedListenerErrorLogger; private volatile CallTracker.CallStats _latestCallStats; @@ -94,6 +102,7 @@ public TrackerClientImpl(URI uri, Map partitionDataMap, { _uri = uri; _transportClient = transportClient; + _clock = clock; _callTracker = new CallTrackerImpl(interval, clock, percentileTrackingEnabled); _isErrorStatus = isErrorStatus; _partitionData = Collections.unmodifiableMap(partitionDataMap); @@ -104,6 +113,8 @@ public TrackerClientImpl(URI uri, Map partitionDataMap, _callTracker.addStatsRolloverEventListener(event -> _latestCallStats = event.getCallStats()); + _rateLimitedListenerErrorLogger = new RateLimitedLogger(_log, PER_CALL_LISTENER_LOG_RATE_MS, clock); + debug(_log, "created tracker client: ", this); } @@ -142,13 +153,47 @@ public double getSubsetWeight(int partitionId) { return _subsetWeightMap.getOrDefault(partitionId, 1D); } + /** + * Sets a listener that is invoked once per completed call with the duration (in ms) the server + * was perceived to have contributed to the request, and with {@link PerCallDurationSemantics} + * describing that measurement. + * + *

See {@link TrackerClient#setPerCallDurationListener(PerCallDurationListener)} for the + * broader design rationale (why a setter is used instead of constructor injection, the timing + * requirement, and the silent-absence caveat for custom {@link TrackerClient} implementations). + */ + @Override + public void setPerCallDurationListener(PerCallDurationListener listener) + { + _perCallDurationListener = listener != null ? listener : (d, s) -> { + }; + } + + /** + * Invokes the per-call duration listener and swallows any {@link RuntimeException} it throws. + */ + private void safelyNotifyDurationListener(long duration, PerCallDurationSemantics semantics) + { + try + { + _perCallDurationListener.accept(duration, semantics); + } + catch (RuntimeException e) + { + _rateLimitedListenerErrorLogger.warn( + "Per-call duration listener threw an exception; latency metric not recorded for this call", e); + } + } + @Override public void restRequest(RestRequest request, RequestContext requestContext, Map wireAttrs, TransportCallback callback) { - _transportClient.restRequest(request, requestContext, wireAttrs, new TrackerClientRestCallback(callback, _callTracker.startCall())); + long startTime = _clock.currentTimeMillis(); + _transportClient.restRequest(request, requestContext, wireAttrs, + new TrackerClientRestCallback(callback, _callTracker.startCall(), startTime)); } @Override @@ -157,7 +202,9 @@ public void streamRequest(StreamRequest request, Map wireAttrs, TransportCallback callback) { - _transportClient.streamRequest(request, requestContext, wireAttrs, new TrackerClientStreamCallback(callback, _callTracker.startCall())); + long startTime = _clock.currentTimeMillis(); + _transportClient.streamRequest(request, requestContext, wireAttrs, + new TrackerClientStreamCallback(callback, _callTracker.startCall(), startTime)); } @Override @@ -182,17 +229,20 @@ private class TrackerClientRestCallback implements TransportCallback _wrappedCallback; private CallCompletion _callCompletion; + private final long _startTime; public TrackerClientRestCallback(TransportCallback wrappedCallback, - CallCompletion callCompletion) + CallCompletion callCompletion, long startTime) { _wrappedCallback = wrappedCallback; _callCompletion = callCompletion; + _startTime = startTime; } @Override public void onResponse(TransportResponse response) { + long duration = _clock.currentTimeMillis() - _startTime; if (response.hasError()) { Throwable throwable = response.getError(); @@ -202,6 +252,7 @@ public void onResponse(TransportResponse response) { _callCompletion.endCall(); } + safelyNotifyDurationListener(duration, PerCallDurationSemantics.FULL_ROUND_TRIP); _wrappedCallback.onResponse(response); } @@ -229,12 +280,14 @@ private class TrackerClientStreamCallback implements TransportCallback _wrappedCallback; private CallCompletion _callCompletion; + private final long _startTime; public TrackerClientStreamCallback(TransportCallback wrappedCallback, - CallCompletion callCompletion) + CallCompletion callCompletion, long startTime) { _wrappedCallback = wrappedCallback; _callCompletion = callCompletion; + _startTime = startTime; } @Override @@ -244,6 +297,8 @@ public void onResponse(TransportResponse response) { Throwable throwable = response.getError(); handleError(_callCompletion, throwable); + safelyNotifyDurationListener(_clock.currentTimeMillis() - _startTime, + PerCallDurationSemantics.FULL_ROUND_TRIP); } else { @@ -264,6 +319,7 @@ public void onResponse(TransportResponse response) * In this way, D2 still monitors the responsiveness of a server without the interference from the client * side events, and error counting still works as before. */ + long firstByteTime = _clock.currentTimeMillis(); _callCompletion.record(); Observer observer = new Observer() { @@ -276,12 +332,18 @@ public void onDataAvailable(ByteString data) public void onDone() { _callCompletion.endCall(); + safelyNotifyDurationListener(firstByteTime - _startTime, + PerCallDurationSemantics.TIME_TO_FIRST_BYTE); } @Override public void onError(Throwable e) { handleError(_callCompletion, e); + // Record TTFB (firstByteTime - startTime) instead of full duration up to the streaming + // error. + safelyNotifyDurationListener(firstByteTime - _startTime, + PerCallDurationSemantics.TIME_TO_FIRST_BYTE); } }; entityStream.addObserver(observer); diff --git a/d2/src/main/java/com/linkedin/d2/balancer/simple/SimpleLoadBalancerState.java b/d2/src/main/java/com/linkedin/d2/balancer/simple/SimpleLoadBalancerState.java index 4f9a0b0bd0..04fe363074 100644 --- a/d2/src/main/java/com/linkedin/d2/balancer/simple/SimpleLoadBalancerState.java +++ b/d2/src/main/java/com/linkedin/d2/balancer/simple/SimpleLoadBalancerState.java @@ -1519,7 +1519,10 @@ private Map createNewStrategies(ServiceProperties List schemes = serviceProperties.getPrioritizedSchemes(); for (String scheme : schemes) { - LoadBalancerStrategy strategy = factory.newLoadBalancer(serviceProperties); + // Pass scheme through so OTel-emitting strategies can tag metrics from the very first + // per-call listener invocation, eliminating the bootstrap window between strategy + // construction and the subsequent setScheme call from D2ClientJmxManager. + LoadBalancerStrategy strategy = factory.newLoadBalancer(serviceProperties, scheme); newStrategies.put(scheme, strategy); } diff --git a/d2/src/main/java/com/linkedin/d2/balancer/strategies/LoadBalancerStrategyFactory.java b/d2/src/main/java/com/linkedin/d2/balancer/strategies/LoadBalancerStrategyFactory.java index 3d1f0bd5b7..a320a8ca26 100644 --- a/d2/src/main/java/com/linkedin/d2/balancer/strategies/LoadBalancerStrategyFactory.java +++ b/d2/src/main/java/com/linkedin/d2/balancer/strategies/LoadBalancerStrategyFactory.java @@ -31,4 +31,31 @@ public interface LoadBalancerStrategyFactory * @return Load balancer strategy. */ T newLoadBalancer(ServiceProperties serviceProperties); + + /** + * Create new {@link LoadBalancerStrategy} for a service, with the URI scheme it will serve known + * up front. Strategies that emit OpenTelemetry metrics tagged by scheme should override this + * overload so the scheme is set at construction time, eliminating the bootstrap window between + * {@link #newLoadBalancer(ServiceProperties)} and + * {@code D2ClientJmxManager.doRegisterLoadBalancerStrategy(...)} during which per-call listener + * emissions and gauge updates would otherwise be silently dropped. + * + *

The default implementation delegates to {@link #newLoadBalancer(ServiceProperties)} for + * backward source compatibility with factories (including out-of-tree implementations) that + * have not been migrated. + * + * @param serviceProperties {@link ServiceProperties}. + * @param scheme URI scheme this strategy instance will serve (e.g. {@code "http"} or + * {@code "https"}); may be {@code null} if unknown at the call site. + * @return Load balancer strategy. + */ + default T newLoadBalancer(ServiceProperties serviceProperties, String scheme) + { + T strategy = newLoadBalancer(serviceProperties); + if (strategy instanceof SchemeAware && scheme != null) + { + ((SchemeAware) strategy).setScheme(scheme); + } + return strategy; + } } diff --git a/d2/src/main/java/com/linkedin/d2/balancer/strategies/SchemeAware.java b/d2/src/main/java/com/linkedin/d2/balancer/strategies/SchemeAware.java new file mode 100644 index 0000000000..ebb1a28135 --- /dev/null +++ b/d2/src/main/java/com/linkedin/d2/balancer/strategies/SchemeAware.java @@ -0,0 +1,46 @@ +/* + Copyright (c) 2026 LinkedIn Corp. + + Licensed 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 com.linkedin.d2.balancer.strategies; + +/** + * Mixin for {@link LoadBalancerStrategy} implementations that need to know which URI scheme + * (e.g. {@code "http"} / {@code "https"}) they are serving so they can tag per-strategy telemetry. + * + * Callers that want to set the scheme should always check first: + * + * {@code + * if (strategy instanceof SchemeAware) { + * ((SchemeAware) strategy).setScheme(scheme); + * } + * } + */ +public interface SchemeAware +{ + + String NO_VALUE = "-"; + + /** + * Sets the URI scheme this strategy is associated with. Used for OpenTelemetry metric tagging + * so the consumer can attribute samples to the correct (service, scheme) pair. + * + * Implementations should treat {@code null} and {@link #NO_VALUE} as "no change" so that + * late, repeated, or partially-initialized callers don't clobber a previously-set scheme. + * + * @param scheme the load-balancer scheme (e.g. {@code "http"}, {@code "https"}) + */ + void setScheme(String scheme); +} diff --git a/d2/src/main/java/com/linkedin/d2/balancer/strategies/degrader/DegraderLoadBalancerStrategyFactoryV3.java b/d2/src/main/java/com/linkedin/d2/balancer/strategies/degrader/DegraderLoadBalancerStrategyFactoryV3.java index 9d5d1a9654..90fda5f8ea 100644 --- a/d2/src/main/java/com/linkedin/d2/balancer/strategies/degrader/DegraderLoadBalancerStrategyFactoryV3.java +++ b/d2/src/main/java/com/linkedin/d2/balancer/strategies/degrader/DegraderLoadBalancerStrategyFactoryV3.java @@ -22,6 +22,8 @@ import com.linkedin.d2.balancer.properties.ServiceProperties; import com.linkedin.d2.balancer.strategies.LoadBalancerStrategyFactory; import com.linkedin.d2.balancer.util.healthcheck.HealthCheckOperations; +import com.linkedin.d2.jmx.DegraderLoadBalancerStrategyV3OtelMetricsProvider; +import com.linkedin.d2.jmx.NoOpDegraderLoadBalancerStrategyV3OtelMetricsProvider; import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; @@ -43,37 +45,55 @@ public class DegraderLoadBalancerStrategyFactoryV3 implements private final ScheduledExecutorService _executorService; private final EventEmitter _eventEmitter; private final List _degraderStateListenerFactories; + private final DegraderLoadBalancerStrategyV3OtelMetricsProvider _degraderLbOtelMetricsProvider; public DegraderLoadBalancerStrategyFactoryV3() { - _healthCheckOperations = null; - _executorService = null; - _eventEmitter = new NoopEventEmitter(); - _degraderStateListenerFactories = Collections.emptyList(); + this(null, null, new NoopEventEmitter(), Collections.emptyList()); } public DegraderLoadBalancerStrategyFactoryV3(HealthCheckOperations healthCheckOperations, ScheduledExecutorService executorService, EventEmitter emitter, List degraderStateListenerFactories) + { + this(healthCheckOperations, executorService, emitter, degraderStateListenerFactories, + new NoOpDegraderLoadBalancerStrategyV3OtelMetricsProvider()); + } + + public DegraderLoadBalancerStrategyFactoryV3(HealthCheckOperations healthCheckOperations, + ScheduledExecutorService executorService, EventEmitter emitter, + List degraderStateListenerFactories, + DegraderLoadBalancerStrategyV3OtelMetricsProvider degraderLbOtelMetricsProvider) { _healthCheckOperations = healthCheckOperations; _executorService = executorService; _eventEmitter = (emitter == null) ? new NoopEventEmitter() : emitter; _degraderStateListenerFactories = degraderStateListenerFactories; + _degraderLbOtelMetricsProvider = (degraderLbOtelMetricsProvider == null) + ? new NoOpDegraderLoadBalancerStrategyV3OtelMetricsProvider() + : degraderLbOtelMetricsProvider; } @Override public DegraderLoadBalancerStrategyV3 newLoadBalancer(ServiceProperties serviceProperties) + { + return newLoadBalancer(serviceProperties, null); + } + + @Override + public DegraderLoadBalancerStrategyV3 newLoadBalancer(ServiceProperties serviceProperties, String scheme) { return newLoadBalancer(serviceProperties.getServiceName(), serviceProperties.getLoadBalancerStrategyProperties(), serviceProperties.getDegraderProperties(), serviceProperties.getPath(), - serviceProperties.getClusterName()); + serviceProperties.getClusterName(), + scheme); } private DegraderLoadBalancerStrategyV3 newLoadBalancer(String serviceName, - Map strategyProperties, Map degraderProperties, String path, String clusterName) + Map strategyProperties, Map degraderProperties, String path, String clusterName, + String scheme) { debug(LOG, "created a degrader load balancer strategyV3"); @@ -95,6 +115,7 @@ private DegraderLoadBalancerStrategyV3 newLoadBalancer(String serviceName, listeners.add(new DegraderMonitorEventEmitter.Factory(serviceName)); listeners.addAll(_degraderStateListenerFactories); - return new DegraderLoadBalancerStrategyV3(config, serviceName, degraderProperties, listeners); + return new DegraderLoadBalancerStrategyV3(config, serviceName, degraderProperties, listeners, + _degraderLbOtelMetricsProvider, scheme); } } diff --git a/d2/src/main/java/com/linkedin/d2/balancer/strategies/degrader/DegraderLoadBalancerStrategyV3.java b/d2/src/main/java/com/linkedin/d2/balancer/strategies/degrader/DegraderLoadBalancerStrategyV3.java index a53a32ab11..b35d444e43 100644 --- a/d2/src/main/java/com/linkedin/d2/balancer/strategies/degrader/DegraderLoadBalancerStrategyV3.java +++ b/d2/src/main/java/com/linkedin/d2/balancer/strategies/degrader/DegraderLoadBalancerStrategyV3.java @@ -25,6 +25,7 @@ import com.linkedin.d2.balancer.strategies.DelegatingRingFactory; import com.linkedin.d2.balancer.strategies.LoadBalancerQuarantine; import com.linkedin.d2.balancer.strategies.LoadBalancerStrategy; +import com.linkedin.d2.balancer.strategies.SchemeAware; import com.linkedin.d2.balancer.util.hashing.HashFunction; import com.linkedin.d2.balancer.util.hashing.RandomHash; import com.linkedin.d2.balancer.util.hashing.Ring; @@ -32,6 +33,8 @@ import com.linkedin.d2.balancer.util.hashing.URIRegexHash; import com.linkedin.d2.balancer.util.healthcheck.HealthCheck; import com.linkedin.d2.balancer.util.healthcheck.HealthCheckClientBuilder; +import com.linkedin.d2.jmx.DegraderLoadBalancerStrategyV3OtelMetricsProvider; +import com.linkedin.d2.jmx.NoOpDegraderLoadBalancerStrategyV3OtelMetricsProvider; import com.linkedin.r2.filter.R2Constants; import com.linkedin.r2.message.Request; import com.linkedin.r2.message.RequestContext; @@ -53,7 +56,9 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.locks.Lock; +import java.util.function.Consumer; import java.util.stream.Collectors; import javax.annotation.Nonnull; import org.slf4j.Logger; @@ -71,7 +76,7 @@ * @author Oby Sumampouw (osumampouw@linkedin.com) * @author Zhenkai Zhu (zzhu@linkedin.com) */ -public class DegraderLoadBalancerStrategyV3 implements LoadBalancerStrategy +public class DegraderLoadBalancerStrategyV3 implements LoadBalancerStrategy, SchemeAware { public static final String DEGRADER_STRATEGY_NAME = "degrader"; public static final String HASH_METHOD_NONE = "none"; @@ -94,10 +99,29 @@ public class DegraderLoadBalancerStrategyV3 implements LoadBalancerStrategy private final DegraderLoadBalancerState _state; private final RateLimitedLogger _rateLimitedLogger; + private final DegraderLoadBalancerStrategyV3OtelMetricsProvider _degraderLbOtelMetricsProvider; + private volatile String _scheme; public DegraderLoadBalancerStrategyV3(DegraderLoadBalancerStrategyConfig config, String serviceName, Map degraderProperties, List degraderStateListenerFactories) + { + this(config, serviceName, degraderProperties, degraderStateListenerFactories, + new NoOpDegraderLoadBalancerStrategyV3OtelMetricsProvider()); + } + + public DegraderLoadBalancerStrategyV3(DegraderLoadBalancerStrategyConfig config, String serviceName, + Map degraderProperties, + List degraderStateListenerFactories, + DegraderLoadBalancerStrategyV3OtelMetricsProvider degraderLbOtelMetricsProvider) + { + this(config, serviceName, degraderProperties, degraderStateListenerFactories, degraderLbOtelMetricsProvider, null); + } + + public DegraderLoadBalancerStrategyV3(DegraderLoadBalancerStrategyConfig config, String serviceName, + Map degraderProperties, + List degraderStateListenerFactories, + DegraderLoadBalancerStrategyV3OtelMetricsProvider degraderLbOtelMetricsProvider, String scheme) { _updateEnabled = true; setConfig(config); @@ -107,7 +131,24 @@ public DegraderLoadBalancerStrategyV3(DegraderLoadBalancerStrategyConfig config, } _state = new DegraderLoadBalancerState(serviceName, degraderProperties, config, degraderStateListenerFactories); _rateLimitedLogger = new RateLimitedLogger(_log, config.DEFAULT_UPDATE_INTERVAL_MS, config.getClock()); + _degraderLbOtelMetricsProvider = (degraderLbOtelMetricsProvider == null) + ? new NoOpDegraderLoadBalancerStrategyV3OtelMetricsProvider() + : degraderLbOtelMetricsProvider; + _scheme = (scheme != null && !scheme.equals(NO_VALUE)) ? scheme : NO_VALUE; + } + /** + * {@link SchemeAware} implementation: stores the scheme for OTel metric tagging. Treats + * {@code null} and the {@link #NO_VALUE} sentinel as "no change" so partially-initialized + * callers cannot clobber a previously-set scheme. + */ + @Override + public void setScheme(String scheme) + { + if (scheme != null && !scheme.equals(NO_VALUE)) + { + _scheme = scheme; + } } @Override @@ -418,6 +459,29 @@ private void updatePartitionState(long clusterGenerationId, Partition partition, } } + // Register per-call OTEL latency listener for clients joining this partition for the first + // time. + // Capture serviceName outside the lambda so each per-call invocation skips a virtual call + // through _state — the value is immutable for the lifetime of this strategy. + final String serviceName = _state.getServiceName(); + Set previouslySeenClients = partitionState.getTrackerClients(); + for (DegraderTrackerClient client : trackerClients) + { + if (previouslySeenClients.contains(client)) + { + continue; + } + client.setPerCallDurationListener((duration, semantics) -> { + // skip emission if the scheme has not been initialized yet. + String scheme = _scheme; + if (NO_VALUE.equals(scheme)) + { + return; + } + _degraderLbOtelMetricsProvider.recordHostLatency(serviceName, scheme, duration, semantics); + }); + } + // doUpdatePartitionState has no side effects on _state or trackerClients. // all changes to the trackerClients would be recorded in clientUpdaters partitionState = doUpdatePartitionState(clusterGenerationId, partition.getId(), partitionState, @@ -429,6 +493,63 @@ private void updatePartitionState(long clusterGenerationId, Partition partition, { clientUpdater.update(); } + + scheduleEmitOtelMetrics(partitionState, config); + } + + /** + * Aligns with {@code StateUpdater#updateStateForPartition}, which schedules + * {@code emitOtelMetrics} on its executor. + */ + private void scheduleEmitOtelMetrics(PartitionDegraderLoadBalancerState partitionState, + DegraderLoadBalancerStrategyConfig config) + { + Runnable task = () -> { + try + { + emitOtelMetrics(partitionState); + } + catch (RuntimeException e) + { + _log.warn("OpenTelemetry degrader gauge emission failed for service " + _state.getServiceName(), e); + } + }; + ScheduledExecutorService executor = config.getExecutorService(); + if (executor != null) + { + executor.execute(task); + } + else + { + task.run(); + } + } + + /** + * Emit OpenTelemetry metrics for the current partition state. + * + * Skips emission entirely if the scheme has not been initialized yet (i.e. before + * {@link #setScheme(String)} runs from D2ClientJmxManager). + * + * Invoked from {@link #scheduleEmitOtelMetrics} (executor thread or inline when no executor). + */ + private void emitOtelMetrics(PartitionDegraderLoadBalancerState partitionState) + { + String scheme = _scheme; + if (NO_VALUE.equals(scheme)) + { + return; + } + + _degraderLbOtelMetricsProvider.updateOverrideClusterDropRate(_state.getServiceName(), scheme, + partitionState.getCurrentOverrideDropRate()); + + int totalPoints = 0; + for (Integer points : partitionState.getPointsMap().values()) + { + totalPoints += points; + } + _degraderLbOtelMetricsProvider.updateTotalPointsInHashRing(_state.getServiceName(), scheme, totalPoints); } diff --git a/d2/src/main/java/com/linkedin/d2/balancer/strategies/relative/RelativeLoadBalancerStrategy.java b/d2/src/main/java/com/linkedin/d2/balancer/strategies/relative/RelativeLoadBalancerStrategy.java index d58935199e..1dad7f070e 100644 --- a/d2/src/main/java/com/linkedin/d2/balancer/strategies/relative/RelativeLoadBalancerStrategy.java +++ b/d2/src/main/java/com/linkedin/d2/balancer/strategies/relative/RelativeLoadBalancerStrategy.java @@ -18,6 +18,7 @@ import com.linkedin.d2.balancer.clients.TrackerClient; import com.linkedin.d2.balancer.strategies.LoadBalancerStrategy; +import com.linkedin.d2.balancer.strategies.SchemeAware; import com.linkedin.d2.balancer.util.hashing.HashFunction; import com.linkedin.d2.balancer.util.hashing.Ring; import com.linkedin.r2.message.Request; @@ -41,7 +42,7 @@ * * @see com.linkedin.d2.D2RelativeStrategyProperties */ -public class RelativeLoadBalancerStrategy implements LoadBalancerStrategy +public class RelativeLoadBalancerStrategy implements LoadBalancerStrategy, SchemeAware { private static final Logger LOG = LoggerFactory.getLogger(RelativeLoadBalancerStrategy.class); public static final String RELATIVE_LOAD_BALANCER_STRATEGY_NAME = "relative"; @@ -147,6 +148,17 @@ public int getTotalHostsInAllPartitions() return _stateUpdater.getTotalHostsInAllPartitions(); } + /** + * {@link SchemeAware} implementation: forwards to the underlying {@link StateUpdater} so OTel + * metrics can be tagged with the scheme. Treats {@code null} / {@code "-"} as "no change" via + * {@code StateUpdater#setScheme} for partially-initialized callers. + */ + @Override + public void setScheme(String scheme) + { + _stateUpdater.setScheme(scheme); + } + /** * Exposed for testings */ diff --git a/d2/src/main/java/com/linkedin/d2/balancer/strategies/relative/RelativeLoadBalancerStrategyFactory.java b/d2/src/main/java/com/linkedin/d2/balancer/strategies/relative/RelativeLoadBalancerStrategyFactory.java index a40fe097c6..ae74f8ed4f 100644 --- a/d2/src/main/java/com/linkedin/d2/balancer/strategies/relative/RelativeLoadBalancerStrategyFactory.java +++ b/d2/src/main/java/com/linkedin/d2/balancer/strategies/relative/RelativeLoadBalancerStrategyFactory.java @@ -34,6 +34,8 @@ import com.linkedin.d2.balancer.util.hashing.RandomHash; import com.linkedin.d2.balancer.util.hashing.URIRegexHash; import com.linkedin.d2.balancer.util.healthcheck.HealthCheckOperations; +import com.linkedin.d2.jmx.NoOpRelativeLoadBalancerStrategyOtelMetricsProvider; +import com.linkedin.d2.jmx.RelativeLoadBalancerStrategyOtelMetricsProvider; import com.linkedin.r2.message.Request; import com.linkedin.util.clock.Clock; import java.util.ArrayList; @@ -78,24 +80,35 @@ public class RelativeLoadBalancerStrategyFactory implements LoadBalancerStrategy private final Clock _clock; private final boolean _loadBalanceStreamException; private final boolean _enableRelativeStrategyDeferredAllocation; + private final RelativeLoadBalancerStrategyOtelMetricsProvider _relativeLbOtelMetricsProvider; public RelativeLoadBalancerStrategyFactory(ScheduledExecutorService executorService, HealthCheckOperations healthCheckOperations, List> stateListenerFactories, EventEmitter eventEmitter, Clock clock) { - this(executorService, healthCheckOperations, stateListenerFactories, eventEmitter, clock, false, false); + this(executorService, healthCheckOperations, stateListenerFactories, eventEmitter, clock, false, false, + new NoOpRelativeLoadBalancerStrategyOtelMetricsProvider()); } public RelativeLoadBalancerStrategyFactory(ScheduledExecutorService executorService, HealthCheckOperations healthCheckOperations, List> stateListenerFactories, EventEmitter eventEmitter, Clock clock, boolean loadBalanceStreamException) { - this(executorService, healthCheckOperations, stateListenerFactories, eventEmitter, clock, - loadBalanceStreamException, false); + this(executorService, healthCheckOperations, stateListenerFactories, eventEmitter, clock, loadBalanceStreamException, + false, new NoOpRelativeLoadBalancerStrategyOtelMetricsProvider()); } public RelativeLoadBalancerStrategyFactory(ScheduledExecutorService executorService, HealthCheckOperations healthCheckOperations, List> stateListenerFactories, EventEmitter eventEmitter, Clock clock, boolean loadBalanceStreamException, boolean enableRelativeStrategyDeferredAllocation) + { + this(executorService, healthCheckOperations, stateListenerFactories, eventEmitter, clock, loadBalanceStreamException, + enableRelativeStrategyDeferredAllocation, new NoOpRelativeLoadBalancerStrategyOtelMetricsProvider()); + } + + public RelativeLoadBalancerStrategyFactory(ScheduledExecutorService executorService, HealthCheckOperations healthCheckOperations, + List> stateListenerFactories, EventEmitter eventEmitter, Clock clock, + boolean loadBalanceStreamException, boolean enableRelativeStrategyDeferredAllocation, + RelativeLoadBalancerStrategyOtelMetricsProvider relativeLbOtelMetricsProvider) { _executorService = executorService; _healthCheckOperations = healthCheckOperations; @@ -104,11 +117,20 @@ public RelativeLoadBalancerStrategyFactory(ScheduledExecutorService executorServ _clock = clock; _loadBalanceStreamException = loadBalanceStreamException; _enableRelativeStrategyDeferredAllocation = enableRelativeStrategyDeferredAllocation; + _relativeLbOtelMetricsProvider = (relativeLbOtelMetricsProvider == null) + ? new NoOpRelativeLoadBalancerStrategyOtelMetricsProvider() + : relativeLbOtelMetricsProvider; } @Override public RelativeLoadBalancerStrategy newLoadBalancer(ServiceProperties serviceProperties) + { + return newLoadBalancer(serviceProperties, null); + } + + @Override + public RelativeLoadBalancerStrategy newLoadBalancer(ServiceProperties serviceProperties, String scheme) { D2RelativeStrategyProperties relativeStrategyProperties = RelativeStrategyPropertiesConverter .toProperties(serviceProperties.getRelativeStrategyProperties()); @@ -116,12 +138,12 @@ public RelativeLoadBalancerStrategy newLoadBalancer(ServiceProperties servicePro return new RelativeLoadBalancerStrategy(getRelativeStateUpdater(relativeStrategyProperties, serviceProperties.getServiceName(), serviceProperties.getClusterName(), - serviceProperties.getPath()), getClientSelector(relativeStrategyProperties), + serviceProperties.getPath(), scheme), getClientSelector(relativeStrategyProperties), _enableRelativeStrategyDeferredAllocation); } private StateUpdater getRelativeStateUpdater(D2RelativeStrategyProperties relativeStrategyProperties, - String serviceName, String clusterName, String servicePath) + String serviceName, String clusterName, String servicePath, String scheme) { QuarantineManager quarantineManager = getQuarantineManager(relativeStrategyProperties, serviceName, servicePath); final List> listenerFactories = new ArrayList<>(); @@ -133,7 +155,7 @@ private StateUpdater getRelativeStateUpdater(D2RelativeStrategyProperties relati listenerFactories.addAll(_stateListenerFactories); } return new StateUpdater(relativeStrategyProperties, quarantineManager, _executorService, listenerFactories, - serviceName, _loadBalanceStreamException); + serviceName, _loadBalanceStreamException, _relativeLbOtelMetricsProvider, scheme); } private ClientSelector getClientSelector(D2RelativeStrategyProperties relativeStrategyProperties) diff --git a/d2/src/main/java/com/linkedin/d2/balancer/strategies/relative/StateUpdater.java b/d2/src/main/java/com/linkedin/d2/balancer/strategies/relative/StateUpdater.java index 32b8ee974f..d7d2143df5 100644 --- a/d2/src/main/java/com/linkedin/d2/balancer/strategies/relative/StateUpdater.java +++ b/d2/src/main/java/com/linkedin/d2/balancer/strategies/relative/StateUpdater.java @@ -19,9 +19,14 @@ import com.google.common.annotations.VisibleForTesting; import com.linkedin.d2.D2RelativeStrategyProperties; import com.linkedin.d2.balancer.clients.TrackerClient; +import com.linkedin.d2.balancer.strategies.LoadBalancerQuarantine; import com.linkedin.d2.balancer.strategies.PartitionStateUpdateListener; import com.linkedin.d2.balancer.strategies.DelegatingRingFactory; +import com.linkedin.d2.balancer.strategies.SchemeAware; import com.linkedin.d2.balancer.util.hashing.Ring; +import com.linkedin.d2.jmx.HostStatus; +import com.linkedin.d2.jmx.NoOpRelativeLoadBalancerStrategyOtelMetricsProvider; +import com.linkedin.d2.jmx.RelativeLoadBalancerStrategyOtelMetricsProvider; import com.linkedin.util.degrader.CallTracker; import com.linkedin.util.degrader.ErrorType; import java.net.URI; @@ -65,10 +70,12 @@ public class StateUpdater private final Lock _lock; private final List> _listenerFactories; private final String _serviceName; + private volatile String _scheme; private final ScheduledFuture scheduledFuture; private ConcurrentMap _partitionLoadBalancerStateMap; private int _firstPartitionId = -1; private final boolean _loadBalanceStreamException; + private final RelativeLoadBalancerStrategyOtelMetricsProvider _relativeLbOtelMetricsProvider; @Deprecated StateUpdater(D2RelativeStrategyProperties relativeStrategyProperties, @@ -91,6 +98,29 @@ public class StateUpdater serviceName, loadBalanceStreamException); } + StateUpdater(D2RelativeStrategyProperties relativeStrategyProperties, + QuarantineManager quarantineManager, + ScheduledExecutorService executorService, + List> listenerFactories, + String serviceName, boolean loadBalanceStreamException, + RelativeLoadBalancerStrategyOtelMetricsProvider relativeLbOtelMetricsProvider) + { + this(relativeStrategyProperties, quarantineManager, executorService, new ConcurrentHashMap<>(), listenerFactories, + serviceName, loadBalanceStreamException, relativeLbOtelMetricsProvider); + } + + StateUpdater(D2RelativeStrategyProperties relativeStrategyProperties, + QuarantineManager quarantineManager, + ScheduledExecutorService executorService, + List> listenerFactories, + String serviceName, boolean loadBalanceStreamException, + RelativeLoadBalancerStrategyOtelMetricsProvider relativeLbOtelMetricsProvider, + String scheme) + { + this(relativeStrategyProperties, quarantineManager, executorService, new ConcurrentHashMap<>(), listenerFactories, + serviceName, loadBalanceStreamException, relativeLbOtelMetricsProvider, scheme); + } + StateUpdater(D2RelativeStrategyProperties relativeStrategyProperties, QuarantineManager quarantineManager, ScheduledExecutorService executorService, @@ -108,6 +138,31 @@ public class StateUpdater ConcurrentMap partitionLoadBalancerStateMap, List> listenerFactories, String serviceName, boolean loadBalanceStreamException) + { + this(relativeStrategyProperties, quarantineManager, executorService, partitionLoadBalancerStateMap, + listenerFactories, serviceName, loadBalanceStreamException, new NoOpRelativeLoadBalancerStrategyOtelMetricsProvider()); + } + + StateUpdater(D2RelativeStrategyProperties relativeStrategyProperties, + QuarantineManager quarantineManager, + ScheduledExecutorService executorService, + ConcurrentMap partitionLoadBalancerStateMap, + List> listenerFactories, + String serviceName, boolean loadBalanceStreamException, + RelativeLoadBalancerStrategyOtelMetricsProvider relativeLbOtelMetricsProvider) + { + this(relativeStrategyProperties, quarantineManager, executorService, partitionLoadBalancerStateMap, + listenerFactories, serviceName, loadBalanceStreamException, relativeLbOtelMetricsProvider, null); + } + + StateUpdater(D2RelativeStrategyProperties relativeStrategyProperties, + QuarantineManager quarantineManager, + ScheduledExecutorService executorService, + ConcurrentMap partitionLoadBalancerStateMap, + List> listenerFactories, + String serviceName, boolean loadBalanceStreamException, + RelativeLoadBalancerStrategyOtelMetricsProvider relativeLbOtelMetricsProvider, + String scheme) { _relativeStrategyProperties = relativeStrategyProperties; _quarantineManager = quarantineManager; @@ -116,6 +171,10 @@ public class StateUpdater _partitionLoadBalancerStateMap = partitionLoadBalancerStateMap; _lock = new ReentrantLock(); _serviceName = serviceName; + _scheme = (scheme != null && !scheme.equals(SchemeAware.NO_VALUE)) ? scheme : SchemeAware.NO_VALUE; + _relativeLbOtelMetricsProvider = (relativeLbOtelMetricsProvider == null) + ? new NoOpRelativeLoadBalancerStrategyOtelMetricsProvider() + : relativeLbOtelMetricsProvider; scheduledFuture = executorService.scheduleWithFixedDelay(this::updateState, EXECUTOR_INITIAL_DELAY, _relativeStrategyProperties.getUpdateIntervalMs(), @@ -123,6 +182,20 @@ public class StateUpdater _loadBalanceStreamException = loadBalanceStreamException; } + /** + * Sets the scheme for this state updater. Used for OTEL metrics tagging. + * This is called after strategy creation when the scheme becomes available. + * + * @param scheme the load balancer scheme (e.g., "http", "https") + */ + public void setScheme(String scheme) + { + if (scheme != null && !scheme.equals(SchemeAware.NO_VALUE)) + { + _scheme = scheme; + } + } + /** * Update the state of the partition if necessary * This update is triggered by the request. If the cluster is not initialized or the uris changed, we will update the state. @@ -238,6 +311,24 @@ void updateStateForPartition(Collection trackerClients, int parti LOG.debug("Updating for partition: " + partitionId + ", state: " + oldPartitionState); PartitionState newPartitionState = new PartitionState(oldPartitionState); + // Register per-call OTel latency listener for clients joining this partition for the first time. + Map oldStateMap = oldPartitionState.getTrackerClientStateMap(); + for (TrackerClient trackerClient : trackerClients) + { + if (!oldStateMap.containsKey(trackerClient) && !trackerClient.doNotLoadBalance()) + { + trackerClient.setPerCallDurationListener((duration, semantics) -> { + // skip emission if the scheme has not been initialized yet. + String scheme = _scheme; + if (SchemeAware.NO_VALUE.equals(scheme)) + { + return; + } + _relativeLbOtelMetricsProvider.recordHostLatency(_serviceName, scheme, duration, semantics); + }); + } + } + // Step 1: Update the base health scores for each {@link TrackerClient} in the cluster Map latestCallStatsMap = new HashMap<>(); long avgClusterLatency = getAvgClusterLatency(trackerClients, latestCallStatsMap); @@ -257,6 +348,14 @@ void updateStateForPartition(Collection trackerClients, int parti // Step 4: Log and emit monitor event _executorService.execute(() -> { logState(oldPartitionState, newPartitionState, partitionId); + try + { + emitOtelMetrics(newPartitionState); + } + catch (RuntimeException e) + { + LOG.warn("OpenTelemetry relative gauge emission failed for service " + _serviceName, e); + } notifyPartitionStateUpdateListener(newPartitionState); }); } @@ -464,6 +563,58 @@ private void notifyPartitionStateUpdateListener(PartitionState state) state.getListeners().forEach(listener -> listener.onUpdate(state)); } + /** + * Emit OpenTelemetry metrics for the current partition state. + * Host latencies are emitted per-call via the listener registered in + * {@link #calculateBaseHealthScore}. + * + *

Skips emission entirely if the scheme has not been initialized yet (i.e. before + * {@link #setScheme(String)} runs from D2ClientJmxManager). + */ + private void emitOtelMetrics(PartitionState partitionState) + { + String scheme = _scheme; + if (SchemeAware.NO_VALUE.equals(scheme)) + { + return; + } + + Map trackerClientStateMap = partitionState.getTrackerClientStateMap(); + Map quarantineMap = partitionState.getQuarantineMap(); + Map pointsMap = partitionState.getPointsMap(); + + // Single traversal over the tracker-client state map: count unhealthy hosts and, in the same + // pass, count hosts that are currently quarantined (joining via map.get instead of streaming + // the quarantine map separately). + int unhealthyCount = 0; + int quarantineCount = 0; + for (Map.Entry entry : trackerClientStateMap.entrySet()) + { + if (entry.getValue().isUnhealthy()) + { + unhealthyCount++; + } + LoadBalancerQuarantine quarantine = quarantineMap.get(entry.getKey()); + if (quarantine != null && quarantine.isInQuarantine()) + { + quarantineCount++; + } + } + + // Single traversal for total ring points; primitive sum, no boxing. + int totalPoints = 0; + for (Integer points : pointsMap.values()) + { + totalPoints += points; + } + + String serviceName = _serviceName; + _relativeLbOtelMetricsProvider.updateTotalHostsInAllPartitionsCount(serviceName, scheme, getTotalHostsInAllPartitions()); + _relativeLbOtelMetricsProvider.updateDegradedHostsCount(serviceName, scheme, HostStatus.UNHEALTHY, unhealthyCount); + _relativeLbOtelMetricsProvider.updateDegradedHostsCount(serviceName, scheme, HostStatus.QUARANTINED, quarantineCount); + _relativeLbOtelMetricsProvider.updateTotalPointsInHashRing(serviceName, scheme, totalPoints); + } + @VisibleForTesting double getErrorRate(Map errorTypeCounts, int callCount) { diff --git a/d2/src/main/java/com/linkedin/d2/jmx/D2ClientJmxManager.java b/d2/src/main/java/com/linkedin/d2/jmx/D2ClientJmxManager.java index 9635c7d5cf..5ec1ba4686 100644 --- a/d2/src/main/java/com/linkedin/d2/jmx/D2ClientJmxManager.java +++ b/d2/src/main/java/com/linkedin/d2/jmx/D2ClientJmxManager.java @@ -29,6 +29,7 @@ import com.linkedin.d2.balancer.simple.SimpleLoadBalancerState; import com.linkedin.d2.balancer.simple.SimpleLoadBalancerState.SimpleLoadBalancerStateListener; import com.linkedin.d2.balancer.strategies.LoadBalancerStrategy; +import com.linkedin.d2.balancer.strategies.SchemeAware; import com.linkedin.d2.discovery.stores.file.FileStore; import com.linkedin.d2.discovery.stores.zk.ZooKeeperEphemeralStore; import com.linkedin.d2.discovery.stores.zk.ZooKeeperPermanentStore; @@ -207,6 +208,12 @@ public void onServicePropertiesRemoval(LoadBalancerStateItem private void doRegisterLoadBalancerStrategy(String serviceName, String scheme, LoadBalancerStrategy strategy, @Nullable DualReadModeProvider.DualReadMode mode) { + // Only strategies that opt-in via the SchemeAware mixin care about scheme tagging + // (currently the OTel-emitting ones). + if (strategy instanceof SchemeAware) + { + ((SchemeAware) strategy).setScheme(scheme); + } String jmxName = getLoadBalancerStrategyJmxName(serviceName, scheme, mode); _jmxManager.registerLoadBalancerStrategy(jmxName, strategy); } diff --git a/d2/src/main/java/com/linkedin/d2/jmx/DegraderLoadBalancerStrategyV3OtelMetricsProvider.java b/d2/src/main/java/com/linkedin/d2/jmx/DegraderLoadBalancerStrategyV3OtelMetricsProvider.java new file mode 100644 index 0000000000..0451e95716 --- /dev/null +++ b/d2/src/main/java/com/linkedin/d2/jmx/DegraderLoadBalancerStrategyV3OtelMetricsProvider.java @@ -0,0 +1,22 @@ +/* + Copyright (c) 2026 LinkedIn Corp. + + Licensed 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 com.linkedin.d2.jmx; + +/** OTel metrics for {@code DegraderLoadBalancerStrategyV3}. */ +public interface DegraderLoadBalancerStrategyV3OtelMetricsProvider extends LoadBalancerStrategyOtelMetricsProvider +{ + void updateOverrideClusterDropRate(String serviceName, String scheme, double overrideClusterDropRate); +} diff --git a/d2/src/main/java/com/linkedin/d2/jmx/HostStatus.java b/d2/src/main/java/com/linkedin/d2/jmx/HostStatus.java new file mode 100644 index 0000000000..af76c93b62 --- /dev/null +++ b/d2/src/main/java/com/linkedin/d2/jmx/HostStatus.java @@ -0,0 +1,32 @@ +/* + Copyright (c) 2026 LinkedIn Corp. + + Licensed 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 com.linkedin.d2.jmx; + +/** Host status for {@code D2.RelativeLb.DegradedHostsCount} gauge attributes. */ +public enum HostStatus +{ + /** + * Hosts whose health score has been reduced due to high latency or error rate. Still receive + * traffic at a reduced weight; not yet quarantined. + */ + UNHEALTHY, + + /** + * Hosts currently in quarantine. They are not receiving production traffic and are pending + * health-check recovery before being re-admitted. + */ + QUARANTINED +} diff --git a/d2/src/main/java/com/linkedin/d2/jmx/LoadBalancerStrategyOtelMetricsProvider.java b/d2/src/main/java/com/linkedin/d2/jmx/LoadBalancerStrategyOtelMetricsProvider.java new file mode 100644 index 0000000000..d7476483c1 --- /dev/null +++ b/d2/src/main/java/com/linkedin/d2/jmx/LoadBalancerStrategyOtelMetricsProvider.java @@ -0,0 +1,42 @@ +/* + Copyright (c) 2026 LinkedIn Corp. + + Licensed 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 com.linkedin.d2.jmx; + +import com.linkedin.d2.balancer.clients.PerCallDurationSemantics; + +/** Shared OTel metrics for D2 load balancer strategies. */ +public interface LoadBalancerStrategyOtelMetricsProvider +{ + /** + * Records a per-call host latency sample. + * + * @param serviceName the name of the service + * @param scheme the load balancer scheme (e.g., "http", "https") + * @param hostLatencyMs the duration in milliseconds + * @param semantics the semantics of the latency + */ + void recordHostLatency(String serviceName, String scheme, long hostLatencyMs, + PerCallDurationSemantics semantics); + + /** + * Updates the total number of points across all hosts in the consistent hash ring. + * + * @param serviceName the name of the service + * @param scheme the load balancer scheme (e.g., "http", "https") + * @param totalPointsInHashRing the total number of points in the hash ring + */ + void updateTotalPointsInHashRing(String serviceName, String scheme, int totalPointsInHashRing); +} diff --git a/d2/src/main/java/com/linkedin/d2/jmx/NoOpDegraderLoadBalancerStrategyV3OtelMetricsProvider.java b/d2/src/main/java/com/linkedin/d2/jmx/NoOpDegraderLoadBalancerStrategyV3OtelMetricsProvider.java new file mode 100644 index 0000000000..f0c3b33300 --- /dev/null +++ b/d2/src/main/java/com/linkedin/d2/jmx/NoOpDegraderLoadBalancerStrategyV3OtelMetricsProvider.java @@ -0,0 +1,53 @@ +/* + Copyright (c) 2026 LinkedIn Corp. + + Licensed 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 com.linkedin.d2.jmx; + +import com.linkedin.d2.balancer.clients.PerCallDurationSemantics; + +/** + * No-Op implementation of {@link DegraderLoadBalancerStrategyV3OtelMetricsProvider}. + * Used when OpenTelemetry metrics are disabled. + */ +public class NoOpDegraderLoadBalancerStrategyV3OtelMetricsProvider implements DegraderLoadBalancerStrategyV3OtelMetricsProvider +{ + /** + * {@inheritDoc} + */ + @Override + public void recordHostLatency(String serviceName, String scheme, long hostLatencyMs, + PerCallDurationSemantics semantics) + { + // No-op + } + + /** + * {@inheritDoc} + */ + @Override + public void updateOverrideClusterDropRate(String serviceName, String scheme, double overrideClusterDropRate) + { + // No-op + } + + /** + * {@inheritDoc} + */ + @Override + public void updateTotalPointsInHashRing(String serviceName, String scheme, int totalPointsInHashRing) + { + // No-op + } +} diff --git a/d2/src/main/java/com/linkedin/d2/jmx/NoOpRelativeLoadBalancerStrategyOtelMetricsProvider.java b/d2/src/main/java/com/linkedin/d2/jmx/NoOpRelativeLoadBalancerStrategyOtelMetricsProvider.java new file mode 100644 index 0000000000..6d68f77625 --- /dev/null +++ b/d2/src/main/java/com/linkedin/d2/jmx/NoOpRelativeLoadBalancerStrategyOtelMetricsProvider.java @@ -0,0 +1,62 @@ +/* + Copyright (c) 2026 LinkedIn Corp. + + Licensed 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 com.linkedin.d2.jmx; + +import com.linkedin.d2.balancer.clients.PerCallDurationSemantics; + +/** + * No-Op implementation of {@link RelativeLoadBalancerStrategyOtelMetricsProvider}. + * Used when OpenTelemetry metrics are disabled. + */ +public class NoOpRelativeLoadBalancerStrategyOtelMetricsProvider implements RelativeLoadBalancerStrategyOtelMetricsProvider +{ + /** + * {@inheritDoc} + */ + @Override + public void recordHostLatency(String serviceName, String scheme, long hostLatencyMs, + PerCallDurationSemantics semantics) + { + // No-op + } + + /** + * {@inheritDoc} + */ + @Override + public void updateTotalHostsInAllPartitionsCount(String serviceName, String scheme, int totalHostsInAllPartitionsCount) + { + // No-op + } + + /** + * {@inheritDoc} + */ + @Override + public void updateDegradedHostsCount(String serviceName, String scheme, HostStatus status, int count) + { + // No-op + } + + /** + * {@inheritDoc} + */ + @Override + public void updateTotalPointsInHashRing(String serviceName, String scheme, int totalPointsInHashRing) + { + // No-op + } +} diff --git a/d2/src/main/java/com/linkedin/d2/jmx/RelativeLoadBalancerStrategyOtelMetricsProvider.java b/d2/src/main/java/com/linkedin/d2/jmx/RelativeLoadBalancerStrategyOtelMetricsProvider.java new file mode 100644 index 0000000000..6e4b3db833 --- /dev/null +++ b/d2/src/main/java/com/linkedin/d2/jmx/RelativeLoadBalancerStrategyOtelMetricsProvider.java @@ -0,0 +1,39 @@ +/* + Copyright (c) 2026 LinkedIn Corp. + + Licensed 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 com.linkedin.d2.jmx; + +/** OTel metrics for {@code RelativeLoadBalancerStrategy}. */ +public interface RelativeLoadBalancerStrategyOtelMetricsProvider extends LoadBalancerStrategyOtelMetricsProvider +{ + /** + * Updates the number of total hosts in all partitions regardless of their status. + * + * @param serviceName the name of the service + * @param scheme the load balancer scheme (e.g., "http", "https") + * @param totalHostsInAllPartitionsCount the count of total hosts in all partitions + */ + void updateTotalHostsInAllPartitionsCount(String serviceName, String scheme, int totalHostsInAllPartitionsCount); + + /** + * Updates the number of hosts currently in a degraded state, broken down by status. + * + * @param serviceName the name of the service + * @param scheme the load balancer scheme (e.g., "http", "https") + * @param status the degraded-host bucket this count belongs to + * @param count the number of hosts currently in {@code status} + */ + void updateDegradedHostsCount(String serviceName, String scheme, HostStatus status, int count); +} diff --git a/d2/src/test/java/com/linkedin/d2/balancer/D2ClientBuilderTest.java b/d2/src/test/java/com/linkedin/d2/balancer/D2ClientBuilderTest.java index 3598179640..7f07fe43bd 100644 --- a/d2/src/test/java/com/linkedin/d2/balancer/D2ClientBuilderTest.java +++ b/d2/src/test/java/com/linkedin/d2/balancer/D2ClientBuilderTest.java @@ -47,4 +47,32 @@ void testD2ServicePathNotNull(String d2ServicePath, String expectedD2ServicePath return Mockito.mock(LoadBalancerWithFacilities.class); }); } + + @Test + void testSetRelativeLoadBalancerStrategyOtelMetricsProviderNullForwardedToConfig() + { + D2ClientBuilder d2ClientBuilder = new D2ClientBuilder(); + d2ClientBuilder.setRelativeLoadBalancerStrategyOtelMetricsProvider(null); + + d2ClientBuilder.setLoadBalancerWithFacilitiesFactory(config -> { + Assert.assertNull(config.relativeLoadBalancerStrategyOtelMetricsProvider, + "Builder must forward null to the config field; NoOp-fallback is the strategy's " + + "responsibility, not the builder's"); + return Mockito.mock(LoadBalancerWithFacilities.class); + }); + } + + @Test + void testSetDegraderLoadBalancerStrategyV3OtelMetricsProviderNullForwardedToConfig() + { + D2ClientBuilder d2ClientBuilder = new D2ClientBuilder(); + d2ClientBuilder.setDegraderLoadBalancerStrategyV3OtelMetricsProvider(null); + + d2ClientBuilder.setLoadBalancerWithFacilitiesFactory(config -> { + Assert.assertNull(config.degraderLoadBalancerStrategyV3OtelMetricsProvider, + "Builder must forward null to the config field; NoOp-fallback is the strategy's " + + "responsibility, not the builder's"); + return Mockito.mock(LoadBalancerWithFacilities.class); + }); + } } diff --git a/d2/src/test/java/com/linkedin/d2/balancer/D2ClientConfigOtelForwarderWarnOnceTest.java b/d2/src/test/java/com/linkedin/d2/balancer/D2ClientConfigOtelForwarderWarnOnceTest.java new file mode 100644 index 0000000000..5e184e19e2 --- /dev/null +++ b/d2/src/test/java/com/linkedin/d2/balancer/D2ClientConfigOtelForwarderWarnOnceTest.java @@ -0,0 +1,95 @@ +/* + Copyright (c) 2026 LinkedIn Corp. + + Licensed 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 com.linkedin.d2.balancer; + +import java.lang.reflect.Field; +import java.lang.reflect.Modifier; +import java.util.concurrent.atomic.AtomicBoolean; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.fail; + + +/** Tests {@code D2ClientConfig.PRE_STRATEGY_OTEL_FORWARDER_WARNED} warn-once guard. */ +public class D2ClientConfigOtelForwarderWarnOnceTest +{ + private static final String GUARD_FIELD_NAME = "PRE_STRATEGY_OTEL_FORWARDER_WARNED"; + + @Test + public void testGuardFieldIsPrivateStaticFinalAtomicBoolean() throws NoSuchFieldException + { + Field guard = D2ClientConfig.class.getDeclaredField(GUARD_FIELD_NAME); + int modifiers = guard.getModifiers(); + assertTrue(Modifier.isPrivate(modifiers), + GUARD_FIELD_NAME + " must remain private; otherwise external code can flip the guard " + + "and suppress the deprecation warning entirely"); + assertTrue(Modifier.isStatic(modifiers), + GUARD_FIELD_NAME + " must remain static; the warn-once contract is per-JVM, not " + + "per-instance"); + assertTrue(Modifier.isFinal(modifiers), + GUARD_FIELD_NAME + " must remain final so the reference can't be reseated after " + + "JVM startup"); + assertEquals(guard.getType(), AtomicBoolean.class, + GUARD_FIELD_NAME + " must remain an AtomicBoolean — a plain boolean is not safe under " + + "concurrent first-time callers and would let the warning fire twice in a race"); + } + + @Test + public void testGuardEmitsExactlyOnceAcrossRepeatedCompareAndSet() throws Exception + { + Field guardField = D2ClientConfig.class.getDeclaredField(GUARD_FIELD_NAME); + guardField.setAccessible(true); + AtomicBoolean guard = (AtomicBoolean) guardField.get(null); + + boolean originalState = guard.get(); + try + { + guard.set(false); + + assertTrue(guard.compareAndSet(false, true), + "First compareAndSet on a fresh guard must return true so the deprecation warning is " + + "emitted exactly once"); + assertFalse(guard.compareAndSet(false, true), + "Subsequent compareAndSet must return false so the deprecation warning is suppressed"); + assertFalse(guard.compareAndSet(false, true), + "Subsequent compareAndSet must remain false; the once-per-JVM contract is sticky"); + } + finally + { + guard.set(originalState); + } + } + + @Test + public void testGuardIsNotAnInstanceField() throws NoSuchFieldException + { + Field guard; + try + { + guard = D2ClientConfig.class.getDeclaredField(GUARD_FIELD_NAME); + } + catch (NoSuchFieldException expected) + { + fail(GUARD_FIELD_NAME + " must exist on D2ClientConfig"); + return; + } + assertTrue(Modifier.isStatic(guard.getModifiers()), + GUARD_FIELD_NAME + " must be static so the warn-once contract is per-JVM"); + } +} diff --git a/d2/src/test/java/com/linkedin/d2/balancer/clients/TrackerClientImplTest.java b/d2/src/test/java/com/linkedin/d2/balancer/clients/TrackerClientImplTest.java index 83e9fdf699..3f5738b70a 100644 --- a/d2/src/test/java/com/linkedin/d2/balancer/clients/TrackerClientImplTest.java +++ b/d2/src/test/java/com/linkedin/d2/balancer/clients/TrackerClientImplTest.java @@ -1,20 +1,68 @@ package com.linkedin.d2.balancer.clients; +import com.linkedin.common.callback.Callback; +import com.linkedin.common.util.None; +import com.linkedin.r2.RemoteInvocationException; +import com.linkedin.r2.message.RequestContext; +import com.linkedin.r2.message.rest.RestRequest; +import com.linkedin.r2.message.rest.RestRequestBuilder; +import com.linkedin.r2.message.rest.RestResponse; +import com.linkedin.r2.message.rest.RestResponseBuilder; +import com.linkedin.r2.message.stream.StreamRequest; +import com.linkedin.r2.message.stream.StreamRequestBuilder; +import com.linkedin.r2.message.stream.StreamResponse; +import com.linkedin.r2.message.stream.StreamResponseBuilder; +import com.linkedin.r2.message.stream.entitystream.EntityStream; +import com.linkedin.r2.message.stream.entitystream.EntityStreams; +import com.linkedin.r2.message.stream.entitystream.Observer; +import com.linkedin.r2.message.stream.entitystream.Reader; +import com.linkedin.r2.transport.common.bridge.client.TransportClient; +import com.linkedin.r2.transport.common.bridge.common.TransportCallback; +import com.linkedin.r2.transport.common.bridge.common.TransportResponse; +import com.linkedin.r2.transport.common.bridge.common.TransportResponseImpl; +import com.linkedin.util.clock.SettableClock; +import com.linkedin.util.clock.SystemClock; import java.net.URI; import java.util.HashMap; - -import com.linkedin.util.clock.SystemClock; - +import java.util.Map; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; import org.junit.Assert; +import org.testng.annotations.DataProvider; import org.testng.annotations.Test; + /** * Tests {@link TrackerClientImpl}. + * + *

In addition to coverage of the basic configuration accessors, this class verifies the + * per-call duration instrumentation introduced for OTel: + *

    + *
  • {@link TrackerClientImpl#setPerCallDurationListener(PerCallDurationListener)} stores + * the listener and tolerates {@code null}.
    • + *
  • The listener is invoked with the correct duration on every callback path: + * REST success, REST error, stream transport error, stream {@code onDone}, and stream + * {@code onError} during streaming.
    • + *
  • Optional {@link org.testng.annotations.DataProvider}-driven tests exercise the same + * scenarios as compact matrix rows (see {@code *_Parametrized} methods).
    • + *
*/ public class TrackerClientImplTest { + private static final URI URI_FOO = URI.create("http://foo.example:1234/svc"); + private static final long DEFAULT_INTERVAL_MS = 1000L; + private TrackerClientImpl _trackerClient; + private static PerCallDurationListener durationSink( + AtomicLong durationOut, AtomicReference semanticsOut) + { + return (d, s) -> { + durationOut.set(d); + semanticsOut.set(s); + }; + } + @Test public void testDoNotLoadBalance() { @@ -28,4 +76,443 @@ public void testDoNotLoadBalance() Assert.assertEquals(_trackerClient.doNotLoadBalance(), doNotLoadBalance); } + + // --------------------------------------------------------------------------- + // setPerCallDurationListener + // --------------------------------------------------------------------------- + + @Test + public void testSetPerCallDurationListenerNullIsTolerated() + { + SettableClock clock = new SettableClock(); + RecordingTransportClient transport = new RecordingTransportClient(); + TrackerClientImpl client = newClient(clock, transport); + + // Listener defaults to a no-op; setting null should fall back to that no-op without throwing. + client.setPerCallDurationListener(null); + + // Simulate a successful REST call to make sure no NPE is raised when the listener fires. + client.restRequest(new RestRequestBuilder(URI_FOO).build(), new RequestContext(), new HashMap<>(), + new CapturingTransportCallback<>()); + clock.addDuration(50); + transport.restCallback.onResponse(TransportResponseImpl.success(new RestResponseBuilder().build())); + // Reaching this line without an exception is the assertion. + } + + // --------------------------------------------------------------------------- + // REST callback paths (line 224 in TrackerClientImpl) + // --------------------------------------------------------------------------- + + @Test + public void testRestSuccessInvokesListenerWithMeasuredDuration() + { + SettableClock clock = new SettableClock(); + RecordingTransportClient transport = new RecordingTransportClient(); + TrackerClientImpl client = newClient(clock, transport); + + AtomicLong observedDuration = new AtomicLong(-1); + AtomicReference observedSemantics = new AtomicReference<>(); + client.setPerCallDurationListener(durationSink(observedDuration, observedSemantics)); + + client.restRequest(new RestRequestBuilder(URI_FOO).build(), new RequestContext(), new HashMap<>(), + new CapturingTransportCallback<>()); + + long expectedDuration = 137L; + clock.addDuration(expectedDuration); + transport.restCallback.onResponse(TransportResponseImpl.success(new RestResponseBuilder().build())); + + Assert.assertEquals("REST success path should record the measured duration", expectedDuration, + observedDuration.get()); + Assert.assertEquals(PerCallDurationSemantics.FULL_ROUND_TRIP, observedSemantics.get()); + } + + @Test + public void testRestErrorInvokesListenerWithMeasuredDuration() + { + SettableClock clock = new SettableClock(); + RecordingTransportClient transport = new RecordingTransportClient(); + TrackerClientImpl client = newClient(clock, transport); + + AtomicLong observedDuration = new AtomicLong(-1); + AtomicReference observedSemantics = new AtomicReference<>(); + client.setPerCallDurationListener(durationSink(observedDuration, observedSemantics)); + + client.restRequest(new RestRequestBuilder(URI_FOO).build(), new RequestContext(), new HashMap<>(), + new CapturingTransportCallback<>()); + + long expectedDuration = 250L; + clock.addDuration(expectedDuration); + transport.restCallback.onResponse( + TransportResponseImpl.error(new RemoteInvocationException("simulated REST error"))); + + Assert.assertEquals("REST error path should record the measured duration", expectedDuration, + observedDuration.get()); + Assert.assertEquals(PerCallDurationSemantics.FULL_ROUND_TRIP, observedSemantics.get()); + } + + // --------------------------------------------------------------------------- + // Stream callback paths + // --------------------------------------------------------------------------- + + /** + * Verifies the immediate stream-transport-error path (line 269 in TrackerClientImpl): when the + * transport callback fires with an error, the listener is invoked using the current clock time + * minus the start time. + */ + @Test + public void testStreamTransportErrorInvokesListenerWithMeasuredDuration() + { + SettableClock clock = new SettableClock(); + RecordingTransportClient transport = new RecordingTransportClient(); + TrackerClientImpl client = newClient(clock, transport); + + AtomicLong observedDuration = new AtomicLong(-1); + AtomicReference observedSemantics = new AtomicReference<>(); + client.setPerCallDurationListener(durationSink(observedDuration, observedSemantics)); + + StreamRequest streamRequest = new StreamRequestBuilder(URI_FOO).build(EntityStreams.emptyStream()); + client.streamRequest(streamRequest, new RequestContext(), new HashMap<>(), new CapturingTransportCallback<>()); + + long expectedDuration = 73L; + clock.addDuration(expectedDuration); + transport.streamCallback.onResponse( + TransportResponseImpl.error(new RemoteInvocationException("simulated stream transport error"))); + + Assert.assertEquals("Stream transport-error path should record the measured duration", + expectedDuration, observedDuration.get()); + Assert.assertEquals(PerCallDurationSemantics.FULL_ROUND_TRIP, observedSemantics.get()); + } + + /** + * Verifies the stream success path (line 303 in TrackerClientImpl): the listener should receive + * {@code firstByteTime - startTime}, NOT the time at {@code onDone}. This is the documented + * D2 behavior (avoid penalizing servers for client-side back-pressure during streaming). + */ + @Test + public void testStreamSuccessOnDoneInvokesListenerWithFirstByteDuration() + { + SettableClock clock = new SettableClock(); + RecordingTransportClient transport = new RecordingTransportClient(); + TrackerClientImpl client = newClient(clock, transport); + + AtomicLong observedDuration = new AtomicLong(-1); + AtomicReference observedSemantics = new AtomicReference<>(); + client.setPerCallDurationListener(durationSink(observedDuration, observedSemantics)); + + StreamRequest streamRequest = new StreamRequestBuilder(URI_FOO).build(EntityStreams.emptyStream()); + client.streamRequest(streamRequest, new RequestContext(), new HashMap<>(), new CapturingTransportCallback<>()); + + // Time until first byte arrives. + long firstByteOffset = 90L; + clock.addDuration(firstByteOffset); + + // Build a stream response backed by a stream that captures the observer the strategy adds. + CapturingEntityStream entityStream = new CapturingEntityStream(); + StreamResponse response = new StreamResponseBuilder().build(entityStream); + transport.streamCallback.onResponse(TransportResponseImpl.success(response)); + + Assert.assertNotNull("Strategy must add an observer to the entity stream", + entityStream.observer); + + // Simulate a slow client consuming the body (back-pressure). The listener must NOT include + // this delay -- it should only reflect time-to-first-byte. + clock.addDuration(500L); + entityStream.observer.onDone(); + + Assert.assertEquals("Stream success path should record firstByteTime - startTime, not the " + + "onDone time", firstByteOffset, observedDuration.get()); + Assert.assertEquals(PerCallDurationSemantics.TIME_TO_FIRST_BYTE, observedSemantics.get()); + } + + /** + * Verifies the stream error-during-streaming path: when the stream observer fires + * {@code onError} after the first byte has arrived, the listener should receive + * {@code firstByteTime - startTime} (TTFB), aligned with the success {@code onDone} path. This + * avoids penalizing the server for time spent streaming the body — which can be dominated by + * client-side back-pressure rather than server responsiveness. + */ + @Test + public void testStreamErrorAfterFirstByteRecordsTtfbNotFullDuration() + { + SettableClock clock = new SettableClock(); + RecordingTransportClient transport = new RecordingTransportClient(); + TrackerClientImpl client = newClient(clock, transport); + + AtomicLong observedDuration = new AtomicLong(-1); + AtomicReference observedSemantics = new AtomicReference<>(); + client.setPerCallDurationListener(durationSink(observedDuration, observedSemantics)); + + StreamRequest streamRequest = new StreamRequestBuilder(URI_FOO).build(EntityStreams.emptyStream()); + client.streamRequest(streamRequest, new RequestContext(), new HashMap<>(), new CapturingTransportCallback<>()); + + long firstByteOffset = 60L; + clock.addDuration(firstByteOffset); + + CapturingEntityStream entityStream = new CapturingEntityStream(); + StreamResponse response = new StreamResponseBuilder().build(entityStream); + transport.streamCallback.onResponse(TransportResponseImpl.success(response)); + + // Stream errors mid-flight after additional time. This extra time should NOT be reflected in + // the recorded duration -- the listener should only see firstByteTime - startTime, matching + // the onDone path for cross-path histogram consistency. + long extraStreamingTime = 200L; + clock.addDuration(extraStreamingTime); + entityStream.observer.onError(new RemoteInvocationException("simulated streaming error")); + + Assert.assertEquals("Stream onError path should record firstByteTime - startTime (TTFB), " + + "ignoring streaming-time after first byte", firstByteOffset, observedDuration.get()); + Assert.assertEquals(PerCallDurationSemantics.TIME_TO_FIRST_BYTE, observedSemantics.get()); + } + + // --------------------------------------------------------------------------- + // Parametrized variants (same scenarios as explicit tests above; matrix-style coverage) + // --------------------------------------------------------------------------- + + @DataProvider(name = "doNotLoadBalanceFlags") + public Object[][] doNotLoadBalanceFlags() + { + return new Object[][] {{true}, {false}}; + } + + @Test(dataProvider = "doNotLoadBalanceFlags") + public void testDoNotLoadBalance_Parametrized(boolean doNotLoadBalance) + { + _trackerClient = new TrackerClientImpl(URI.create("uri"), new HashMap<>(), null, SystemClock.instance(), 1000, + (test) -> false, false, false, doNotLoadBalance); + + Assert.assertEquals(_trackerClient.doNotLoadBalance(), doNotLoadBalance); + } + + @DataProvider(name = "restDurationPaths") + public Object[][] restDurationPaths() + { + return new Object[][] { + {true, 137L, "REST success path should record the measured duration"}, + {false, 250L, "REST error path should record the measured duration"} + }; + } + + @Test(dataProvider = "restDurationPaths") + public void testRestInvokesListenerWithMeasuredDuration_Parametrized(boolean restSuccess, long expectedDuration, + String assertionMessage) + { + SettableClock clock = new SettableClock(); + RecordingTransportClient transport = new RecordingTransportClient(); + TrackerClientImpl client = newClient(clock, transport); + + AtomicLong observedDuration = new AtomicLong(-1); + AtomicReference observedSemantics = new AtomicReference<>(); + client.setPerCallDurationListener(durationSink(observedDuration, observedSemantics)); + + client.restRequest(new RestRequestBuilder(URI_FOO).build(), new RequestContext(), new HashMap<>(), + new CapturingTransportCallback<>()); + + clock.addDuration(expectedDuration); + if (restSuccess) + { + transport.restCallback.onResponse(TransportResponseImpl.success(new RestResponseBuilder().build())); + } + else + { + transport.restCallback.onResponse( + TransportResponseImpl.error(new RemoteInvocationException("simulated REST error"))); + } + + Assert.assertEquals(assertionMessage, expectedDuration, observedDuration.get()); + Assert.assertEquals(PerCallDurationSemantics.FULL_ROUND_TRIP, observedSemantics.get()); + } + + /** + * Parametrized counterpart to {@link #testStreamSuccessOnDoneInvokesListenerWithFirstByteDuration()} + * and {@link #testStreamErrorAfterFirstByteRecordsTtfbNotFullDuration()}. + */ + @DataProvider(name = "streamTtfbAfterFirstBytePaths") + public Object[][] streamTtfbAfterFirstBytePaths() + { + return new Object[][] { + { + true, + 90L, + 500L, + "Stream success path should record firstByteTime - startTime, not the onDone time" + }, + { + false, + 60L, + 200L, + "Stream onError path should record firstByteTime - startTime (TTFB), " + + "ignoring streaming-time after first byte" + } + }; + } + + @Test(dataProvider = "streamTtfbAfterFirstBytePaths") + public void testStreamRecordsTtfbAfterFirstByte_Parametrized(boolean completeWithOnDone, long firstByteOffset, + long delayAfterFirstByte, String assertionMessage) + { + SettableClock clock = new SettableClock(); + RecordingTransportClient transport = new RecordingTransportClient(); + TrackerClientImpl client = newClient(clock, transport); + + AtomicLong observedDuration = new AtomicLong(-1); + AtomicReference observedSemantics = new AtomicReference<>(); + client.setPerCallDurationListener(durationSink(observedDuration, observedSemantics)); + + StreamRequest streamRequest = new StreamRequestBuilder(URI_FOO).build(EntityStreams.emptyStream()); + client.streamRequest(streamRequest, new RequestContext(), new HashMap<>(), new CapturingTransportCallback<>()); + + clock.addDuration(firstByteOffset); + + CapturingEntityStream entityStream = new CapturingEntityStream(); + StreamResponse response = new StreamResponseBuilder().build(entityStream); + transport.streamCallback.onResponse(TransportResponseImpl.success(response)); + + Assert.assertNotNull("Strategy must add an observer to the entity stream", entityStream.observer); + + clock.addDuration(delayAfterFirstByte); + if (completeWithOnDone) + { + entityStream.observer.onDone(); + } + else + { + entityStream.observer.onError(new RemoteInvocationException("simulated streaming error")); + } + + Assert.assertEquals(assertionMessage, firstByteOffset, observedDuration.get()); + Assert.assertEquals(PerCallDurationSemantics.TIME_TO_FIRST_BYTE, observedSemantics.get()); + } + + /** + * Verifies graceful degradation: if the per-call listener throws, the wrapped transport + * callback must still be invoked so the application's request is not stuck. The exception must + * not escape from {@link com.linkedin.r2.transport.common.bridge.common.TransportCallback#onResponse}. + */ + @Test + public void testListenerExceptionDoesNotBlockWrappedCallback() + { + SettableClock clock = new SettableClock(); + RecordingTransportClient transport = new RecordingTransportClient(); + TrackerClientImpl client = newClient(clock, transport); + + AtomicLong listenerInvocations = new AtomicLong(); + client.setPerCallDurationListener((duration, semantics) -> { + listenerInvocations.incrementAndGet(); + throw new RuntimeException("simulated OTel SDK failure"); + }); + + AtomicLong wrappedCallbackInvocations = new AtomicLong(); + TransportCallback wrappedCallback = response -> wrappedCallbackInvocations.incrementAndGet(); + + // The exception from the listener must NOT propagate out of onResponse; if it did, the test + // would fail with an uncaught RuntimeException here. + client.restRequest(new RestRequestBuilder(URI_FOO).build(), new RequestContext(), new HashMap<>(), + wrappedCallback); + clock.addDuration(50); + transport.restCallback.onResponse(TransportResponseImpl.success(new RestResponseBuilder().build())); + + Assert.assertEquals("Listener should still have been invoked once", 1, listenerInvocations.get()); + Assert.assertEquals("Wrapped callback must fire even when the listener throws", 1, + wrappedCallbackInvocations.get()); + } + + /** + * Sanity check: replacing the listener with another non-null listener swaps the destination. + */ + @Test + public void testSetPerCallDurationListenerReplacesPreviousListener() + { + SettableClock clock = new SettableClock(); + RecordingTransportClient transport = new RecordingTransportClient(); + TrackerClientImpl client = newClient(clock, transport); + + AtomicReference firstSink = new AtomicReference<>(); + AtomicReference secondSink = new AtomicReference<>(); + client.setPerCallDurationListener((d, s) -> firstSink.set(d)); + client.setPerCallDurationListener((d, s) -> secondSink.set(d)); + + client.restRequest(new RestRequestBuilder(URI_FOO).build(), new RequestContext(), new HashMap<>(), + new CapturingTransportCallback<>()); + clock.addDuration(42); + transport.restCallback.onResponse(TransportResponseImpl.success(new RestResponseBuilder().build())); + + Assert.assertNull("First listener must not receive any duration after being replaced", + firstSink.get()); + Assert.assertNotNull("Second listener must receive the duration", secondSink.get()); + Assert.assertEquals(42L, secondSink.get().longValue()); + } + + // --------------------------------------------------------------------------- + // Helpers + // --------------------------------------------------------------------------- + + private static TrackerClientImpl newClient(SettableClock clock, TransportClient transport) + { + Map partitionData = new HashMap<>(); + partitionData.put(0, new com.linkedin.d2.balancer.properties.PartitionData(1)); + return new TrackerClientImpl(URI_FOO, partitionData, transport, clock, DEFAULT_INTERVAL_MS, + status -> false, false, false, false); + } + + /** + * Captures the most recently passed-in transport callbacks so the test can drive them. + */ + private static final class RecordingTransportClient implements TransportClient + { + volatile TransportCallback restCallback; + volatile TransportCallback streamCallback; + + @Override + public void restRequest(RestRequest request, RequestContext requestContext, Map wireAttrs, + TransportCallback callback) + { + restCallback = callback; + } + + @Override + public void streamRequest(StreamRequest request, RequestContext requestContext, Map wireAttrs, + TransportCallback callback) + { + streamCallback = callback; + } + + @Override + public void shutdown(Callback callback) + { + callback.onSuccess(None.none()); + } + } + + /** + * Captures the {@link Observer} the strategy adds to the entity stream so the test can drive + * {@code onDone} / {@code onError} directly without relying on the real reader/writer flow. + */ + private static final class CapturingEntityStream implements EntityStream + { + volatile Observer observer; + + @Override + public void addObserver(Observer o) + { + observer = o; + } + + @Override + public void setReader(Reader r) + { + // Not used by these tests; production code only adds an observer. + } + } + + /** + * No-op {@link TransportCallback} used to satisfy the wrappedCallback invariant. The tests + * inspect side effects on the per-call duration listener, not on the wrapped callback. + */ + private static final class CapturingTransportCallback implements TransportCallback + { + @Override + public void onResponse(TransportResponse response) + { + } + } } diff --git a/d2/src/test/java/com/linkedin/d2/balancer/strategies/degrader/DegraderLoadBalancerStrategyV3OtelIntegrationTest.java b/d2/src/test/java/com/linkedin/d2/balancer/strategies/degrader/DegraderLoadBalancerStrategyV3OtelIntegrationTest.java new file mode 100644 index 0000000000..b03e96831c --- /dev/null +++ b/d2/src/test/java/com/linkedin/d2/balancer/strategies/degrader/DegraderLoadBalancerStrategyV3OtelIntegrationTest.java @@ -0,0 +1,398 @@ +/* + Copyright (c) 2026 LinkedIn Corp. + + Licensed 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 com.linkedin.d2.balancer.strategies.degrader; + +import com.linkedin.common.callback.Callback; +import com.linkedin.common.util.None; +import com.linkedin.d2.balancer.clients.DegraderTrackerClient; +import com.linkedin.d2.balancer.clients.PerCallDurationListener; +import com.linkedin.d2.balancer.clients.PerCallDurationSemantics; +import com.linkedin.d2.balancer.clients.DegraderTrackerClientImpl; +import com.linkedin.d2.balancer.clients.TrackerClient; +import com.linkedin.d2.balancer.properties.PartitionData; +import com.linkedin.d2.balancer.util.hashing.Ring; +import com.linkedin.d2.balancer.util.partitions.DefaultPartitionAccessor; +import com.linkedin.d2.jmx.DegraderLoadBalancerStrategyV3OtelMetricsProvider; +import com.linkedin.d2.jmx.TestDegraderLoadBalancerStrategyV3OtelMetricsProvider; +import com.linkedin.r2.message.RequestContext; +import com.linkedin.r2.message.rest.RestRequest; +import com.linkedin.r2.message.rest.RestResponse; +import com.linkedin.r2.message.stream.StreamRequest; +import com.linkedin.r2.message.stream.StreamResponse; +import com.linkedin.r2.transport.common.bridge.client.TransportClient; +import com.linkedin.r2.transport.common.bridge.common.TransportCallback; +import com.linkedin.util.clock.SystemClock; +import java.net.URI; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertTrue; + + +/** Integration tests for {@link DegraderLoadBalancerStrategyV3} OTel metrics wiring. */ +public class DegraderLoadBalancerStrategyV3OtelIntegrationTest +{ + private static final String SERVICE_NAME = "integration-test-service"; + private static final String SCHEME = "http"; + private static final long CLUSTER_GENERATION_ID = 1L; + private static final int PARTITION_ID = DefaultPartitionAccessor.DEFAULT_PARTITION_ID; + private static final List NO_LISTENERS = + Collections.emptyList(); + + private TestDegraderLoadBalancerStrategyV3OtelMetricsProvider _provider; + + @BeforeMethod + public void setUp() + { + _provider = new TestDegraderLoadBalancerStrategyV3OtelMetricsProvider(); + } + + @Test + public void testGaugeMetricsEmittedOnStateUpdate() + { + DegraderLoadBalancerStrategyV3 strategy = newStrategy(); + strategy.setScheme(SCHEME); + + Map trackerClients = newTrackerClientMap(3); + + strategy.getRing(CLUSTER_GENERATION_ID, PARTITION_ID, trackerClients); + + assertEquals(_provider.getCallCount("updateOverrideClusterDropRate"), 1, + "updateOverrideClusterDropRate should be invoked exactly once after a state update"); + assertEquals(_provider.getCallCount("updateTotalPointsInHashRing"), 1, + "updateTotalPointsInHashRing should be invoked exactly once after a state update"); + assertEquals(_provider.getLastServiceName("updateOverrideClusterDropRate"), SERVICE_NAME); + assertEquals(_provider.getLastScheme("updateOverrideClusterDropRate"), SCHEME); + assertEquals(_provider.getLastServiceName("updateTotalPointsInHashRing"), SERVICE_NAME); + assertEquals(_provider.getLastScheme("updateTotalPointsInHashRing"), SCHEME); + + assertEquals(_provider.getLastIntValue("updateTotalPointsInHashRing").intValue(), 300, + "Total points in hash ring should reflect 3 healthy hosts at 100 points each"); + } + + @Test + public void testPerCallLatencyListenerRegisteredAndForwardsToProvider() + { + DegraderLoadBalancerStrategyV3 strategy = newStrategy(); + strategy.setScheme(SCHEME); + + List capturingClients = new ArrayList<>(); + Map trackerClients = new HashMap<>(); + for (int i = 0; i < 2; i++) + { + ListenerCapturingTrackerClient c = newListenerCapturingTrackerClient(URI.create("http://host" + i + ":1234")); + capturingClients.add(c); + trackerClients.put(c.getUri(), c); + } + + strategy.getRing(CLUSTER_GENERATION_ID, PARTITION_ID, trackerClients); + + long simulatedLatencyMs = 123L; + int invokedListeners = 0; + for (ListenerCapturingTrackerClient client : capturingClients) + { + if (client.capturedListener != null) + { + client.capturedListener.accept(simulatedLatencyMs, PerCallDurationSemantics.FULL_ROUND_TRIP); + invokedListeners++; + } + } + assertTrue(invokedListeners > 0, "Expected at least one tracker client to have a per-call listener registered"); + + assertEquals(_provider.getCallCount("recordHostLatency"), invokedListeners, + "recordHostLatency should be invoked once per simulated per-call listener invocation"); + assertEquals(_provider.getLastLongValue("recordHostLatency").longValue(), simulatedLatencyMs); + assertEquals(_provider.getLastServiceName("recordHostLatency"), SERVICE_NAME); + assertEquals(_provider.getLastScheme("recordHostLatency"), SCHEME); + assertEquals(_provider.getLastPerCallDurationSemantics("recordHostLatency"), + PerCallDurationSemantics.FULL_ROUND_TRIP); + } + + @Test + public void testPerCallListenerNotReRegisteredForExistingClientOnSubsequentCycle() + { + DegraderLoadBalancerStrategyV3 strategy = newStrategy(); + strategy.setScheme(SCHEME); + + URI uri = URI.create("http://host0:1234"); + CountingListenerCapturingTrackerClient client = newCountingListenerCapturingTrackerClient(uri); + Map trackerClients = new HashMap<>(); + trackerClients.put(uri, client); + + strategy.getRing(CLUSTER_GENERATION_ID, PARTITION_ID, trackerClients); + assertEquals(client.setListenerCallCount, 1, + "Listener must be registered exactly once on the first state update for a new client"); + + strategy.getRing(CLUSTER_GENERATION_ID + 1, PARTITION_ID, trackerClients); + assertEquals(client.setListenerCallCount, 1, + "Listener must not be re-registered for a client already known to the partition"); + } + + @Test + public void testPerCallListenerReRegisteredAfterClusterRegenWithSameUris() + { + DegraderLoadBalancerStrategyV3 strategy = newStrategy(); + strategy.setScheme(SCHEME); + + URI uri = URI.create("http://host0:1234"); + ListenerCapturingTrackerClient firstGeneration = newListenerCapturingTrackerClient(uri); + Map firstClients = new HashMap<>(); + firstClients.put(uri, firstGeneration); + + strategy.getRing(CLUSTER_GENERATION_ID, PARTITION_ID, firstClients); + assertTrue(firstGeneration.capturedListener != null, "First generation client should get a listener"); + + ListenerCapturingTrackerClient secondGeneration = newListenerCapturingTrackerClient(uri); + Map secondClients = new HashMap<>(); + secondClients.put(uri, secondGeneration); + + strategy.getRing(CLUSTER_GENERATION_ID + 1, PARTITION_ID, secondClients); + + assertTrue(secondGeneration.capturedListener != null, + "New tracker client instance at an existing URI must receive a per-call listener after cluster regen"); + secondGeneration.capturedListener.accept(99L, PerCallDurationSemantics.FULL_ROUND_TRIP); + assertEquals(_provider.getCallCount("recordHostLatency"), 1); + assertEquals(_provider.getLastLongValue("recordHostLatency").longValue(), 99L); + } + + @Test + public void testStateUpdateBeforeSetSchemeSkipsEmission() + { + DegraderLoadBalancerStrategyV3 strategy = newStrategy(); + strategy.getRing(CLUSTER_GENERATION_ID, PARTITION_ID, newTrackerClientMap(1)); + + assertEquals(_provider.getCallCount("updateOverrideClusterDropRate"), 0, + "Gauge metrics should not be emitted before setScheme is called"); + assertEquals(_provider.getCallCount("updateTotalPointsInHashRing"), 0, + "Gauge metrics should not be emitted before setScheme is called"); + } + + @Test + public void testPerCallListenerSkipsEmissionBeforeSetScheme() + { + DegraderLoadBalancerStrategyV3 strategy = newStrategy(); + ListenerCapturingTrackerClient client = newListenerCapturingTrackerClient(URI.create("http://host:1234")); + Map trackerClients = new HashMap<>(); + trackerClients.put(client.getUri(), client); + + strategy.getRing(CLUSTER_GENERATION_ID, PARTITION_ID, trackerClients); + + assertTrue(client.capturedListener != null, "Listener must still be registered on the tracker client"); + + client.capturedListener.accept(123L, PerCallDurationSemantics.FULL_ROUND_TRIP); + assertEquals(_provider.getCallCount("recordHostLatency"), 0, + "Per-call latency must not be recorded before setScheme is called"); + + strategy.setScheme(SCHEME); + client.capturedListener.accept(321L, PerCallDurationSemantics.FULL_ROUND_TRIP); + assertEquals(_provider.getCallCount("recordHostLatency"), 1, + "Per-call latency should be recorded once scheme is initialized"); + assertEquals(_provider.getLastLongValue("recordHostLatency").longValue(), 321L); + assertEquals(_provider.getLastScheme("recordHostLatency"), SCHEME); + } + + @Test + public void testSetSchemeNullPreservesExistingScheme() + { + DegraderLoadBalancerStrategyV3 strategy = newStrategy(); + strategy.setScheme(SCHEME); + strategy.setScheme(null); + + strategy.getRing(CLUSTER_GENERATION_ID, PARTITION_ID, newTrackerClientMap(1)); + + assertEquals(_provider.getCallCount("updateOverrideClusterDropRate"), 1, + "Metrics must still be emitted with the previously-set scheme"); + assertEquals(_provider.getLastScheme("updateOverrideClusterDropRate"), SCHEME, + "setScheme(null) must not change the scheme"); + } + + @Test + public void testSetSchemeWithPlaceholderPreservesExistingScheme() + { + DegraderLoadBalancerStrategyV3 strategy = newStrategy(); + strategy.setScheme(SCHEME); + strategy.setScheme("-"); + + strategy.getRing(CLUSTER_GENERATION_ID, PARTITION_ID, newTrackerClientMap(1)); + + assertEquals(_provider.getCallCount("updateOverrideClusterDropRate"), 1); + assertEquals(_provider.getLastScheme("updateOverrideClusterDropRate"), SCHEME, + "setScheme(\"-\") must not change the scheme"); + } + + @Test + public void testGaugeEmissionRuntimeExceptionDoesNotAbortRingUpdate() + { + DegraderLoadBalancerStrategyV3OtelMetricsProvider provider = new TestDegraderLoadBalancerStrategyV3OtelMetricsProvider() + { + @Override + public void updateOverrideClusterDropRate(String serviceName, String scheme, double overrideClusterDropRate) + { + throw new RuntimeException("simulated OTel gauge failure"); + } + }; + DegraderLoadBalancerStrategyV3 strategy = new DegraderLoadBalancerStrategyV3( + new DegraderLoadBalancerStrategyConfig(5000), + SERVICE_NAME, + null, + NO_LISTENERS, + provider); + strategy.setScheme(SCHEME); + + Ring ring = strategy.getRing(CLUSTER_GENERATION_ID, PARTITION_ID, newTrackerClientMap(1)); + assertNotNull(ring, "Ring must be returned even when gauge emission throws"); + + Ring secondRing = strategy.getRing(CLUSTER_GENERATION_ID + 1, PARTITION_ID, newTrackerClientMap(2)); + assertNotNull(secondRing, + "Subsequent ring must still be returned; provider exceptions on gauges must not halt updates"); + } + + @Test + public void testNullProviderIsCoalescedToNoOpAndDoesNotNpeOnEmission() + { + DegraderLoadBalancerStrategyV3 strategy = new DegraderLoadBalancerStrategyV3( + new DegraderLoadBalancerStrategyConfig(5000), + SERVICE_NAME, + null, + NO_LISTENERS, + null); + strategy.setScheme(SCHEME); + + Ring ring = strategy.getRing(CLUSTER_GENERATION_ID, PARTITION_ID, newTrackerClientMap(2)); + assertNotNull(ring, + "Ring must be returned when provider is null (constructor must coalesce to NoOp)"); + } + + @Test + public void testEmptyTrackerClientsDoesNotEmitMetrics() + { + DegraderLoadBalancerStrategyV3 strategy = newStrategy(); + strategy.setScheme(SCHEME); + + strategy.getRing(CLUSTER_GENERATION_ID, PARTITION_ID, Collections.emptyMap()); + + assertEquals(_provider.getCallCount("updateOverrideClusterDropRate"), 0, + "No metrics should be emitted when there are no tracker clients"); + assertEquals(_provider.getCallCount("updateTotalPointsInHashRing"), 0, + "No metrics should be emitted when there are no tracker clients"); + assertEquals(_provider.getCallCount("recordHostLatency"), 0, + "No latencies should be recorded when there are no tracker clients"); + } + + private DegraderLoadBalancerStrategyV3 newStrategy() + { + return new DegraderLoadBalancerStrategyV3( + new DegraderLoadBalancerStrategyConfig(5000), + SERVICE_NAME, + null, + NO_LISTENERS, + _provider); + } + + private static Map newTrackerClientMap(int count) + { + Map clients = new HashMap<>(); + for (int i = 0; i < count; i++) + { + DegraderTrackerClient client = newDegraderTrackerClient(URI.create("http://host" + i + ":1234")); + clients.put(client.getUri(), client); + } + return clients; + } + + private static DegraderTrackerClient newDegraderTrackerClient(URI uri) + { + Map partitionDataMap = new HashMap<>(); + partitionDataMap.put(DefaultPartitionAccessor.DEFAULT_PARTITION_ID, new PartitionData(1)); + return new DegraderTrackerClientImpl(uri, partitionDataMap, new NoopTransportClient(), + SystemClock.instance(), null); + } + + private static ListenerCapturingTrackerClient newListenerCapturingTrackerClient(URI uri) + { + Map partitionDataMap = new HashMap<>(); + partitionDataMap.put(DefaultPartitionAccessor.DEFAULT_PARTITION_ID, new PartitionData(1)); + return new ListenerCapturingTrackerClient(uri, partitionDataMap, new NoopTransportClient()); + } + + private static CountingListenerCapturingTrackerClient newCountingListenerCapturingTrackerClient(URI uri) + { + Map partitionDataMap = new HashMap<>(); + partitionDataMap.put(DefaultPartitionAccessor.DEFAULT_PARTITION_ID, new PartitionData(1)); + return new CountingListenerCapturingTrackerClient(uri, partitionDataMap, new NoopTransportClient()); + } + + private static final class ListenerCapturingTrackerClient extends DegraderTrackerClientImpl + { + volatile PerCallDurationListener capturedListener; + + ListenerCapturingTrackerClient(URI uri, Map partitionDataMap, TransportClient client) + { + super(uri, partitionDataMap, client, SystemClock.instance(), null); + } + + @Override + public void setPerCallDurationListener(PerCallDurationListener listener) + { + capturedListener = listener; + super.setPerCallDurationListener(listener); + } + } + + private static final class CountingListenerCapturingTrackerClient extends DegraderTrackerClientImpl + { + volatile int setListenerCallCount; + + CountingListenerCapturingTrackerClient(URI uri, Map partitionDataMap, TransportClient client) + { + super(uri, partitionDataMap, client, SystemClock.instance(), null); + } + + @Override + public void setPerCallDurationListener(PerCallDurationListener listener) + { + setListenerCallCount++; + super.setPerCallDurationListener(listener); + } + } + + private static final class NoopTransportClient implements TransportClient + { + @Override + public void restRequest(RestRequest request, RequestContext requestContext, Map wireAttrs, + TransportCallback callback) + { + } + + @Override + public void streamRequest(StreamRequest request, RequestContext requestContext, Map wireAttrs, + TransportCallback callback) + { + } + + @Override + public void shutdown(Callback callback) + { + } + } +} diff --git a/d2/src/test/java/com/linkedin/d2/balancer/strategies/relative/RelativeLoadBalancerStrategyOtelIntegrationTest.java b/d2/src/test/java/com/linkedin/d2/balancer/strategies/relative/RelativeLoadBalancerStrategyOtelIntegrationTest.java new file mode 100644 index 0000000000..2514f96e34 --- /dev/null +++ b/d2/src/test/java/com/linkedin/d2/balancer/strategies/relative/RelativeLoadBalancerStrategyOtelIntegrationTest.java @@ -0,0 +1,398 @@ +/* + Copyright (c) 2026 LinkedIn Corp. + + Licensed 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 com.linkedin.d2.balancer.strategies.relative; + +import com.linkedin.d2.D2RelativeStrategyProperties; +import com.linkedin.d2.balancer.clients.PerCallDurationListener; +import com.linkedin.d2.balancer.clients.PerCallDurationSemantics; +import com.linkedin.d2.balancer.clients.TrackerClient; +import com.linkedin.d2.balancer.strategies.PartitionStateUpdateListener; +import com.linkedin.d2.jmx.HostStatus; +import com.linkedin.d2.jmx.RelativeLoadBalancerStrategyOtelMetricsProvider; +import com.linkedin.d2.jmx.TestRelativeLoadBalancerStrategyOtelMetricsProvider; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ScheduledExecutorService; +import org.mockito.ArgumentCaptor; +import org.mockito.Mockito; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; + + +/** Integration tests for {@link StateUpdater} OTel metrics wiring. */ +public class RelativeLoadBalancerStrategyOtelIntegrationTest +{ + private static final int DEFAULT_PARTITION_ID = 0; + private static final long DEFAULT_CLUSTER_GENERATION_ID = 0; + private static final String SERVICE_NAME = "integration-test-service"; + private static final String SCHEME = "http"; + + private TestRelativeLoadBalancerStrategyOtelMetricsProvider _provider; + private QuarantineManager _quarantineManager; + private ScheduledExecutorService _executorService; + + @BeforeMethod + public void setUp() + { + _provider = new TestRelativeLoadBalancerStrategyOtelMetricsProvider(); + _quarantineManager = Mockito.mock(QuarantineManager.class); + _executorService = Mockito.mock(ScheduledExecutorService.class); + Mockito.doAnswer(invocation -> { + ((Runnable) invocation.getArguments()[0]).run(); + return null; + }).when(_executorService).execute(Mockito.any(Runnable.class)); + } + + @Test + public void testGaugeMetricsEmittedAfterStateUpdate() + { + StateUpdater stateUpdater = newStateUpdater(); + stateUpdater.setScheme(SCHEME); + + List trackerClients = TrackerClientMockHelper.mockTrackerClients( + 3, + Arrays.asList(20, 20, 20), + Arrays.asList(10, 10, 10), + Arrays.asList(200L, 200L, 200L), + Arrays.asList(100L, 100L, 100L), + Arrays.asList(0, 0, 0)); + + stateUpdater.updateState(new HashSet<>(trackerClients), DEFAULT_PARTITION_ID, + DEFAULT_CLUSTER_GENERATION_ID, false); + + assertEquals(_provider.getCallCount("updateTotalHostsInAllPartitionsCount"), 1, + "updateTotalHostsInAllPartitionsCount should be invoked once after a state update"); + assertEquals(_provider.getCallCountForHostStatus("updateDegradedHostsCount", HostStatus.UNHEALTHY), 1, + "updateDegradedHostsCount(UNHEALTHY) should be invoked once after a state update"); + assertEquals(_provider.getCallCountForHostStatus("updateDegradedHostsCount", HostStatus.QUARANTINED), 1, + "updateDegradedHostsCount(QUARANTINED) should be invoked once after a state update"); + assertEquals(_provider.getCallCount("updateTotalPointsInHashRing"), 1, + "updateTotalPointsInHashRing should be invoked once after a state update"); + + assertEquals(_provider.getLastServiceName("updateTotalHostsInAllPartitionsCount"), SERVICE_NAME); + assertEquals(_provider.getLastScheme("updateTotalHostsInAllPartitionsCount"), SCHEME); + assertEquals(_provider.getLastServiceName("updateTotalPointsInHashRing"), SERVICE_NAME); + assertEquals(_provider.getLastScheme("updateTotalPointsInHashRing"), SCHEME); + + assertEquals(_provider.getLastIntValue("updateTotalHostsInAllPartitionsCount").intValue(), 3, + "Total hosts gauge should match the number of tracker clients"); + + assertEquals(_provider.getLastIntValueForHostStatus("updateDegradedHostsCount", HostStatus.UNHEALTHY).intValue(), 0, + "All hosts in the fixture are healthy; unhealthy gauge must be 0"); + assertEquals(_provider.getLastIntValueForHostStatus("updateDegradedHostsCount", HostStatus.QUARANTINED).intValue(), 0, + "No host in the fixture is quarantined; quarantine gauge must be 0"); + assertEquals(_provider.getLastServiceName("updateDegradedHostsCount"), SERVICE_NAME); + assertEquals(_provider.getLastScheme("updateDegradedHostsCount"), SCHEME); + assertTrue(_provider.getLastIntValue("updateTotalPointsInHashRing") > 0, + "Total ring points should be strictly positive when there are healthy hosts"); + } + + @Test + public void testFourTupleGaugeSnapshotIsInternallyConsistent() + { + StateUpdater stateUpdater = newStateUpdater(); + stateUpdater.setScheme(SCHEME); + + List trackerClients = TrackerClientMockHelper.mockTrackerClients( + 3, + Arrays.asList(20, 20, 20), + Arrays.asList(10, 10, 10), + Arrays.asList(200L, 200L, 200L), + Arrays.asList(100L, 100L, 100L), + Arrays.asList(0, 0, 0)); + + stateUpdater.updateState(new HashSet<>(trackerClients), DEFAULT_PARTITION_ID, + DEFAULT_CLUSTER_GENERATION_ID, false); + + assertEquals(_provider.getCallCount("updateTotalHostsInAllPartitionsCount"), 1, + "snapshot: total-hosts gauge must fire exactly once per cycle"); + assertEquals(_provider.getCallCountForHostStatus("updateDegradedHostsCount", HostStatus.UNHEALTHY), 1, + "snapshot: degraded(UNHEALTHY) gauge must fire exactly once per cycle"); + assertEquals(_provider.getCallCountForHostStatus("updateDegradedHostsCount", HostStatus.QUARANTINED), 1, + "snapshot: degraded(QUARANTINED) gauge must fire exactly once per cycle"); + assertEquals(_provider.getCallCount("updateTotalPointsInHashRing"), 1, + "snapshot: ring-points gauge must fire exactly once per cycle"); + + int total = _provider.getLastIntValue("updateTotalHostsInAllPartitionsCount"); + int unhealthy = + _provider.getLastIntValueForHostStatus("updateDegradedHostsCount", HostStatus.UNHEALTHY); + int quarantined = + _provider.getLastIntValueForHostStatus("updateDegradedHostsCount", HostStatus.QUARANTINED); + int ringPoints = _provider.getLastIntValue("updateTotalPointsInHashRing"); + + assertEquals(_provider.getLastServiceName("updateTotalHostsInAllPartitionsCount"), SERVICE_NAME); + assertEquals(_provider.getLastScheme("updateTotalHostsInAllPartitionsCount"), SCHEME); + assertEquals(_provider.getLastServiceName("updateDegradedHostsCount"), SERVICE_NAME); + assertEquals(_provider.getLastScheme("updateDegradedHostsCount"), SCHEME); + assertEquals(_provider.getLastServiceName("updateTotalPointsInHashRing"), SERVICE_NAME); + assertEquals(_provider.getLastScheme("updateTotalPointsInHashRing"), SCHEME); + + assertEquals(total, 3, "total-hosts must match the tracker-client fixture size"); + assertTrue(unhealthy >= 0, "unhealthy count must be non-negative"); + assertTrue(quarantined >= 0, "quarantined count must be non-negative"); + assertTrue(unhealthy + quarantined <= total, + "degraded counts cannot exceed total hosts: unhealthy=" + unhealthy + + ", quarantined=" + quarantined + ", total=" + total); + assertTrue(ringPoints > 0, + "ring points must be strictly positive when the cluster has any healthy hosts"); + } + + @Test + public void testPerCallLatencyListenerRegisteredAndForwardsToProvider() + { + StateUpdater stateUpdater = newStateUpdater(); + stateUpdater.setScheme(SCHEME); + + List trackerClients = TrackerClientMockHelper.mockTrackerClients( + 1, + Collections.singletonList(20), + Collections.singletonList(10), + Collections.singletonList(200L), + Collections.singletonList(100L), + Collections.singletonList(0)); + + stateUpdater.updateState(new HashSet<>(trackerClients), DEFAULT_PARTITION_ID, + DEFAULT_CLUSTER_GENERATION_ID, false); + + ArgumentCaptor listenerCaptor = + ArgumentCaptor.forClass(PerCallDurationListener.class); + Mockito.verify(trackerClients.get(0)).setPerCallDurationListener(listenerCaptor.capture()); + + long simulatedLatencyMs = 321L; + PerCallDurationListener capturedListener = listenerCaptor.getValue(); + capturedListener.accept(simulatedLatencyMs, PerCallDurationSemantics.FULL_ROUND_TRIP); + + assertEquals(_provider.getCallCount("recordHostLatency"), 1, + "Listener should forward exactly one recordHostLatency invocation"); + assertEquals(_provider.getLastLongValue("recordHostLatency").longValue(), simulatedLatencyMs); + assertEquals(_provider.getLastServiceName("recordHostLatency"), SERVICE_NAME); + assertEquals(_provider.getLastScheme("recordHostLatency"), SCHEME); + assertEquals(_provider.getLastPerCallDurationSemantics("recordHostLatency"), + PerCallDurationSemantics.FULL_ROUND_TRIP); + } + + @Test + public void testGaugeMetricsSkippedBeforeSetScheme() + { + StateUpdater stateUpdater = newStateUpdater(); + List trackerClients = TrackerClientMockHelper.mockTrackerClients( + 2, + Arrays.asList(20, 20), + Arrays.asList(10, 10), + Arrays.asList(200L, 200L), + Arrays.asList(100L, 100L), + Arrays.asList(0, 0)); + + stateUpdater.updateState(new HashSet<>(trackerClients), DEFAULT_PARTITION_ID, + DEFAULT_CLUSTER_GENERATION_ID, false); + + assertEquals(_provider.getCallCount("updateTotalHostsInAllPartitionsCount"), 0, + "Gauge metrics should not be emitted before setScheme is called"); + assertEquals(_provider.getCallCount("updateDegradedHostsCount"), 0, + "Gauge metrics should not be emitted before setScheme is called"); + assertEquals(_provider.getCallCount("updateTotalPointsInHashRing"), 0, + "Gauge metrics should not be emitted before setScheme is called"); + } + + @Test + public void testPerCallListenerSkipsEmissionBeforeSetScheme() + { + StateUpdater stateUpdater = newStateUpdater(); + List trackerClients = TrackerClientMockHelper.mockTrackerClients( + 1, + Collections.singletonList(20), + Collections.singletonList(10), + Collections.singletonList(200L), + Collections.singletonList(100L), + Collections.singletonList(0)); + + stateUpdater.updateState(new HashSet<>(trackerClients), DEFAULT_PARTITION_ID, + DEFAULT_CLUSTER_GENERATION_ID, false); + + ArgumentCaptor listenerCaptor = + ArgumentCaptor.forClass(PerCallDurationListener.class); + Mockito.verify(trackerClients.get(0)).setPerCallDurationListener(listenerCaptor.capture()); + + PerCallDurationListener capturedListener = listenerCaptor.getValue(); + + capturedListener.accept(100L, PerCallDurationSemantics.FULL_ROUND_TRIP); + assertEquals(_provider.getCallCount("recordHostLatency"), 0, + "Per-call latency must not be recorded before setScheme is called"); + + stateUpdater.setScheme(SCHEME); + capturedListener.accept(200L, PerCallDurationSemantics.FULL_ROUND_TRIP); + assertEquals(_provider.getCallCount("recordHostLatency"), 1, + "Per-call latency should be recorded once scheme is initialized"); + assertEquals(_provider.getLastLongValue("recordHostLatency").longValue(), 200L); + assertEquals(_provider.getLastScheme("recordHostLatency"), SCHEME); + } + + @Test + public void testSetSchemeNullPreservesExistingScheme() + { + StateUpdater stateUpdater = newStateUpdater(); + stateUpdater.setScheme(SCHEME); + stateUpdater.setScheme(null); + + List trackerClients = TrackerClientMockHelper.mockTrackerClients( + 1, + Collections.singletonList(20), + Collections.singletonList(10), + Collections.singletonList(200L), + Collections.singletonList(100L), + Collections.singletonList(0)); + + stateUpdater.updateState(new HashSet<>(trackerClients), DEFAULT_PARTITION_ID, + DEFAULT_CLUSTER_GENERATION_ID, false); + + assertEquals(_provider.getCallCount("updateTotalPointsInHashRing"), 1, + "Metrics must still be emitted with the previously-set scheme"); + assertEquals(_provider.getLastScheme("updateTotalPointsInHashRing"), SCHEME, + "setScheme(null) must not change the scheme"); + } + + @Test + public void testSetSchemeWithPlaceholderPreservesExistingScheme() + { + StateUpdater stateUpdater = newStateUpdater(); + stateUpdater.setScheme(SCHEME); + stateUpdater.setScheme("-"); + + List trackerClients = TrackerClientMockHelper.mockTrackerClients( + 1, + Collections.singletonList(20), + Collections.singletonList(10), + Collections.singletonList(200L), + Collections.singletonList(100L), + Collections.singletonList(0)); + + stateUpdater.updateState(new HashSet<>(trackerClients), DEFAULT_PARTITION_ID, + DEFAULT_CLUSTER_GENERATION_ID, false); + + assertEquals(_provider.getCallCount("updateTotalPointsInHashRing"), 1); + assertEquals(_provider.getLastScheme("updateTotalPointsInHashRing"), SCHEME, + "setScheme(\"-\") must not change the scheme"); + } + + @Test + public void testStateUpdateAdvancesWhenProviderThrowsOnGauge() + { + _provider = new TestRelativeLoadBalancerStrategyOtelMetricsProvider() + { + @Override + public void updateTotalPointsInHashRing(String serviceName, String scheme, int totalPointsInHashRing) + { + throw new RuntimeException("simulated OTel SDK gauge failure"); + } + }; + + StateUpdater stateUpdater = newStateUpdater(); + stateUpdater.setScheme(SCHEME); + + List trackerClients = TrackerClientMockHelper.mockTrackerClients( + 2, + Arrays.asList(20, 20), + Arrays.asList(10, 10), + Arrays.asList(200L, 200L), + Arrays.asList(100L, 100L), + Arrays.asList(0, 0)); + + stateUpdater.updateState(new HashSet<>(trackerClients), DEFAULT_PARTITION_ID, + DEFAULT_CLUSTER_GENERATION_ID, false); + + stateUpdater.updateState(new HashSet<>(trackerClients), DEFAULT_PARTITION_ID, + DEFAULT_CLUSTER_GENERATION_ID + 1, true); + + assertTrue(_provider.getCallCount("updateTotalHostsInAllPartitionsCount") >= 1, + "Earlier gauges must still be recorded even when a later gauge throws"); + } + + @Test + public void testListenerNotRegisteredForDoNotLoadBalanceClients() + { + StateUpdater stateUpdater = newStateUpdater(); + stateUpdater.setScheme(SCHEME); + + List trackerClients = TrackerClientMockHelper.mockTrackerClients( + 2, + Arrays.asList(20, 20), + Arrays.asList(10, 10), + Arrays.asList(200L, 200L), + Arrays.asList(100L, 100L), + Arrays.asList(0, 0), + false, + Arrays.asList(true, true)); + + stateUpdater.updateState(new HashSet<>(trackerClients), DEFAULT_PARTITION_ID, + DEFAULT_CLUSTER_GENERATION_ID, false); + + for (TrackerClient client : trackerClients) + { + Mockito.verify(client, Mockito.never()).setPerCallDurationListener(Mockito.any()); + } + assertTrue(_provider.getCallCount("updateTotalPointsInHashRing") >= 1, + "Gauge metrics should still be emitted regardless of per-call listener registration"); + } + + @Test + public void testNullProviderIsCoalescedToNoOpAndDoesNotNpeOnEmission() + { + D2RelativeStrategyProperties props = new D2RelativeStrategyProperties(); + RelativeLoadBalancerStrategyFactory.putDefaultValues(props); + StateUpdater stateUpdater = new StateUpdater( + props, + _quarantineManager, + _executorService, + new ConcurrentHashMap<>(), + Collections.>emptyList(), + SERVICE_NAME, + false, + null); + stateUpdater.setScheme(SCHEME); + + List trackerClients = TrackerClientMockHelper.mockTrackerClients( + 2, + Arrays.asList(20, 20), + Arrays.asList(10, 10), + Arrays.asList(200L, 200L), + Arrays.asList(100L, 100L), + Arrays.asList(0, 0)); + + stateUpdater.updateState(new HashSet<>(trackerClients), DEFAULT_PARTITION_ID, + DEFAULT_CLUSTER_GENERATION_ID, false); + } + + private StateUpdater newStateUpdater() + { + D2RelativeStrategyProperties props = new D2RelativeStrategyProperties(); + RelativeLoadBalancerStrategyFactory.putDefaultValues(props); + return new StateUpdater( + props, + _quarantineManager, + _executorService, + new ConcurrentHashMap<>(), + Collections.>emptyList(), + SERVICE_NAME, + false, + _provider); + } +} diff --git a/d2/src/test/java/com/linkedin/d2/jmx/AbstractRecordingOtelMetricsProvider.java b/d2/src/test/java/com/linkedin/d2/jmx/AbstractRecordingOtelMetricsProvider.java new file mode 100644 index 0000000000..85f15d06da --- /dev/null +++ b/d2/src/test/java/com/linkedin/d2/jmx/AbstractRecordingOtelMetricsProvider.java @@ -0,0 +1,242 @@ +/* + Copyright (c) 2026 LinkedIn Corp. + + Licensed 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 com.linkedin.d2.jmx; + +import com.linkedin.d2.balancer.clients.PerCallDurationSemantics; +import java.util.ArrayList; +import java.util.List; + + +/** Shared recording base for strategy OTel metrics test doubles. */ +abstract class AbstractRecordingOtelMetricsProvider +{ + private final List _calls = new ArrayList<>(); + + protected void recordLong(String methodName, String serviceName, String scheme, long longValue) + { + _calls.add(new MetricsInvocation(methodName, serviceName, scheme, longValue)); + } + + protected void recordLong(String methodName, String serviceName, String scheme, long longValue, + PerCallDurationSemantics semantics) + { + _calls.add(new MetricsInvocation(methodName, serviceName, scheme, longValue, semantics)); + } + + protected void recordInt(String methodName, String serviceName, String scheme, int intValue) + { + _calls.add(new MetricsInvocation(methodName, serviceName, scheme, intValue)); + } + + protected void recordInt(String methodName, String serviceName, String scheme, int intValue, + HostStatus hostStatus) + { + _calls.add(new MetricsInvocation(methodName, serviceName, scheme, intValue, hostStatus)); + } + + protected void recordDouble(String methodName, String serviceName, String scheme, double doubleValue) + { + _calls.add(new MetricsInvocation(methodName, serviceName, scheme, doubleValue)); + } + + public int getCallCount(String methodName) + { + int count = 0; + for (MetricsInvocation call : _calls) + { + if (call.methodName.equals(methodName)) + { + count++; + } + } + return count; + } + + public String getLastServiceName(String methodName) + { + MetricsInvocation last = lastFor(methodName); + return last == null ? null : last.serviceName; + } + + public String getLastScheme(String methodName) + { + MetricsInvocation last = lastFor(methodName); + return last == null ? null : last.scheme; + } + + public Long getLastLongValue(String methodName) + { + MetricsInvocation last = lastFor(methodName); + return last == null ? null : last.longValue; + } + + public Integer getLastIntValue(String methodName) + { + MetricsInvocation last = lastFor(methodName); + return last == null ? null : last.intValue; + } + + public Double getLastDoubleValue(String methodName) + { + MetricsInvocation last = lastFor(methodName); + return last == null ? null : last.doubleValue; + } + + public PerCallDurationSemantics getLastPerCallDurationSemantics(String methodName) + { + MetricsInvocation last = lastFor(methodName); + return last == null ? null : last.perCallDurationSemantics; + } + + /** + * Returns the number of invocations of {@code methodName} whose recorded {@link HostStatus} + * tag matched {@code status}. Used to assert that a single attribute-dimensioned gauge has + * been emitted once per status value in the same emission cycle. + */ + public int getCallCountForHostStatus(String methodName, HostStatus status) + { + int count = 0; + for (MetricsInvocation call : _calls) + { + if (call.methodName.equals(methodName) && call.hostStatus == status) + { + count++; + } + } + return count; + } + + /** + * Returns the most recent integer value recorded for {@code methodName} with the given + * {@link HostStatus} attribute, or {@code null} if no such call was recorded. + */ + public Integer getLastIntValueForHostStatus(String methodName, HostStatus status) + { + for (int i = _calls.size() - 1; i >= 0; i--) + { + MetricsInvocation call = _calls.get(i); + if (call.methodName.equals(methodName) && call.hostStatus == status) + { + return call.intValue; + } + } + return null; + } + + /** + * Returns all recorded latency values for a given service name and scheme, in invocation order. + * Useful for verifying histogram data points. + * + * @param serviceName the service name to filter by, or {@code null} for all services + * @param scheme the scheme to filter by, or {@code null} for all schemes + * @return list of recorded latency values + */ + public List getAllLatencyValues(String serviceName, String scheme) + { + List latencies = new ArrayList<>(); + for (MetricsInvocation call : _calls) + { + if (call.methodName.equals("recordHostLatency") + && (serviceName == null || serviceName.equals(call.serviceName)) + && (scheme == null || scheme.equals(call.scheme))) + { + latencies.add(call.longValue); + } + } + return latencies; + } + + /** + * Clears all recorded calls. Useful for resetting state between tests. + */ + public void reset() + { + _calls.clear(); + } + + private MetricsInvocation lastFor(String methodName) + { + for (int i = _calls.size() - 1; i >= 0; i--) + { + MetricsInvocation call = _calls.get(i); + if (call.methodName.equals(methodName)) + { + return call; + } + } + return null; + } + + // ------------------------------------------------------------------------- + // Single record type — only the relevant numeric slot is populated per call. + // ------------------------------------------------------------------------- + + private static final class MetricsInvocation + { + final String methodName; + final String serviceName; + final String scheme; + Long longValue; + Integer intValue; + Double doubleValue; + PerCallDurationSemantics perCallDurationSemantics; + HostStatus hostStatus; + + MetricsInvocation(String methodName, String serviceName, String scheme, long longValue) + { + this.methodName = methodName; + this.serviceName = serviceName; + this.scheme = scheme; + this.longValue = longValue; + } + + MetricsInvocation(String methodName, String serviceName, String scheme, long longValue, + PerCallDurationSemantics semantics) + { + this.methodName = methodName; + this.serviceName = serviceName; + this.scheme = scheme; + this.longValue = longValue; + this.perCallDurationSemantics = semantics; + } + + MetricsInvocation(String methodName, String serviceName, String scheme, int intValue) + { + this.methodName = methodName; + this.serviceName = serviceName; + this.scheme = scheme; + this.intValue = intValue; + } + + MetricsInvocation(String methodName, String serviceName, String scheme, int intValue, + HostStatus hostStatus) + { + this.methodName = methodName; + this.serviceName = serviceName; + this.scheme = scheme; + this.intValue = intValue; + this.hostStatus = hostStatus; + } + + MetricsInvocation(String methodName, String serviceName, String scheme, double doubleValue) + { + this.methodName = methodName; + this.serviceName = serviceName; + this.scheme = scheme; + this.doubleValue = doubleValue; + } + } +} diff --git a/d2/src/test/java/com/linkedin/d2/jmx/D2ClientJmxManagerTest.java b/d2/src/test/java/com/linkedin/d2/jmx/D2ClientJmxManagerTest.java index 7f5747bad5..136028ad3c 100644 --- a/d2/src/test/java/com/linkedin/d2/jmx/D2ClientJmxManagerTest.java +++ b/d2/src/test/java/com/linkedin/d2/jmx/D2ClientJmxManagerTest.java @@ -373,6 +373,31 @@ public Object[][] sourceTypeAndDualReadModeForDualReadModeSwitch() DualReadModeProvider.DualReadMode.OLD_LB_ONLY, true, false} }; } + /** + * Verifies that {@code D2ClientJmxManager}'s {@code SimpleLoadBalancerStateListener#onStrategyAdded} + * propagates the scheme to a {@link com.linkedin.d2.balancer.strategies.SchemeAware} strategy via + * {@link com.linkedin.d2.balancer.strategies.SchemeAware#setScheme(String)}. + * Since the scheme is a dimension on every per-strategy OTel metric, this wiring is critical: + * a regression that drops the call (e.g. accidentally removing the {@code instanceof SchemeAware} + * dispatch in {@code doRegisterLoadBalancerStrategy}) would leave all strategies tagged with their + * default placeholder. + */ + @Test + public void testOnStrategyAddedPropagatesSchemeToStrategy() + { + D2ClientJmxManagerFixture fixture = new D2ClientJmxManagerFixture(); + // Use the simplest configuration (null source, single-read) so we don't have to reason about + // dual-read JMX-name prefixing. The unit under test here is scheme propagation, not naming. + D2ClientJmxManager d2ClientJmxManager = fixture.getD2ClientJmxManager("Foo", null, false); + d2ClientJmxManager.setSimpleLoadBalancerState(fixture._simpleLoadBalancerState); + SimpleLoadBalancerState.SimpleLoadBalancerStateListener lbStateListener = + fixture._simpleLoadBalancerStateListenerCaptor.getValue(); + + lbStateListener.onStrategyAdded("S_Foo", "https", fixture._relativeLoadBalancerStrategy); + + Mockito.verify(fixture._relativeLoadBalancerStrategy).setScheme("https"); + } + @Test(dataProvider = "sourceTypeAndDualReadModeForLixSwitch") public void testJmxNamesOnDualReadModeSwitch(D2ClientJmxManager.DiscoverySourceType sourceType, DualReadModeProvider.DualReadMode oldMode, DualReadModeProvider.DualReadMode newMode, boolean isPrimaryBefore, boolean isPrimaryAfter) diff --git a/d2/src/test/java/com/linkedin/d2/jmx/DegraderLoadBalancerStrategyV3OtelMetricsProviderTest.java b/d2/src/test/java/com/linkedin/d2/jmx/DegraderLoadBalancerStrategyV3OtelMetricsProviderTest.java new file mode 100644 index 0000000000..d699c7cfb3 --- /dev/null +++ b/d2/src/test/java/com/linkedin/d2/jmx/DegraderLoadBalancerStrategyV3OtelMetricsProviderTest.java @@ -0,0 +1,281 @@ +/* + Copyright (c) 2026 LinkedIn Corp. + + Licensed 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 com.linkedin.d2.jmx; + +import com.linkedin.d2.balancer.clients.PerCallDurationSemantics; + +import java.util.Arrays; +import java.util.List; + +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; + +/** Tests for {@link DegraderLoadBalancerStrategyV3OtelMetricsProvider}. */ +public class DegraderLoadBalancerStrategyV3OtelMetricsProviderTest +{ + private TestDegraderLoadBalancerStrategyV3OtelMetricsProvider _testProvider; + + @BeforeMethod + public void setUp() + { + _testProvider = new TestDegraderLoadBalancerStrategyV3OtelMetricsProvider(); + } + + @Test + public void testRecordHostLatency() + { + String serviceName = "test-service-latency"; + String scheme = "http"; + long latencyMs = 150L; + + _testProvider.recordHostLatency(serviceName, scheme, latencyMs, PerCallDurationSemantics.FULL_ROUND_TRIP); + + assertEquals(_testProvider.getCallCount("recordHostLatency"), 1); + assertEquals(_testProvider.getLastServiceName("recordHostLatency"), serviceName); + assertEquals(_testProvider.getLastScheme("recordHostLatency"), scheme); + assertEquals(_testProvider.getLastLongValue("recordHostLatency").longValue(), latencyMs); + assertEquals(_testProvider.getLastPerCallDurationSemantics("recordHostLatency"), + PerCallDurationSemantics.FULL_ROUND_TRIP); + } + + @Test + public void testRecordHostLatencySemanticsDimensionLastWriteWins() + { + _testProvider.recordHostLatency("svc", "http", 10L, PerCallDurationSemantics.FULL_ROUND_TRIP); + assertEquals(_testProvider.getLastPerCallDurationSemantics("recordHostLatency"), + PerCallDurationSemantics.FULL_ROUND_TRIP); + _testProvider.recordHostLatency("svc", "http", 20L, PerCallDurationSemantics.TIME_TO_FIRST_BYTE); + assertEquals(_testProvider.getLastLongValue("recordHostLatency").longValue(), 20L); + assertEquals(_testProvider.getLastPerCallDurationSemantics("recordHostLatency"), + PerCallDurationSemantics.TIME_TO_FIRST_BYTE); + } + + @Test + public void testRecordMultipleHostLatencies() + { + String serviceName = "test-service-multi-latency"; + String scheme = "https"; + + _testProvider.recordHostLatency(serviceName, scheme, 100L, PerCallDurationSemantics.FULL_ROUND_TRIP); + _testProvider.recordHostLatency(serviceName, scheme, 150L, PerCallDurationSemantics.FULL_ROUND_TRIP); + _testProvider.recordHostLatency(serviceName, scheme, 200L, PerCallDurationSemantics.FULL_ROUND_TRIP); + _testProvider.recordHostLatency(serviceName, scheme, 120L, PerCallDurationSemantics.FULL_ROUND_TRIP); + _testProvider.recordHostLatency(serviceName, scheme, 180L, PerCallDurationSemantics.FULL_ROUND_TRIP); + + assertEquals(_testProvider.getCallCount("recordHostLatency"), 5); + assertEquals(_testProvider.getLastServiceName("recordHostLatency"), serviceName); + assertEquals(_testProvider.getLastScheme("recordHostLatency"), scheme); + assertEquals(_testProvider.getLastLongValue("recordHostLatency").longValue(), 180L); + + List allLatencies = _testProvider.getAllLatencyValues(serviceName, scheme); + assertEquals(allLatencies.size(), 5); + assertEquals(allLatencies, Arrays.asList(100L, 150L, 200L, 120L, 180L)); + } + + @Test + public void testRecordHostLatencyZero() + { + _testProvider.recordHostLatency("test-service", "http", 0L, PerCallDurationSemantics.FULL_ROUND_TRIP); + + assertEquals(_testProvider.getCallCount("recordHostLatency"), 1); + assertEquals(_testProvider.getLastLongValue("recordHostLatency").longValue(), 0L); + } + + @Test + public void testUpdateOverrideClusterDropRate() + { + String serviceName = "test-service-drop-rate"; + String scheme = "http"; + double dropRate = 0.25; + + _testProvider.updateOverrideClusterDropRate(serviceName, scheme, dropRate); + + assertEquals(_testProvider.getCallCount("updateOverrideClusterDropRate"), 1); + assertEquals(_testProvider.getLastServiceName("updateOverrideClusterDropRate"), serviceName); + assertEquals(_testProvider.getLastScheme("updateOverrideClusterDropRate"), scheme); + assertEquals(_testProvider.getLastDoubleValue("updateOverrideClusterDropRate"), dropRate, 1e-9); + } + + @Test + public void testUpdateOverrideClusterDropRateZero() + { + _testProvider.updateOverrideClusterDropRate("test-service", "http", 0.0); + + assertEquals(_testProvider.getCallCount("updateOverrideClusterDropRate"), 1); + assertEquals(_testProvider.getLastDoubleValue("updateOverrideClusterDropRate"), 0.0, 1e-9); + } + + @Test + public void testUpdateOverrideClusterDropRateFull() + { + _testProvider.updateOverrideClusterDropRate("test-service", "https", 1.0); + + assertEquals(_testProvider.getCallCount("updateOverrideClusterDropRate"), 1); + assertEquals(_testProvider.getLastDoubleValue("updateOverrideClusterDropRate"), 1.0, 1e-9); + } + + @Test + public void testUpdateTotalPointsInHashRing() + { + String serviceName = "test-service-ring"; + String scheme = "https"; + int totalPoints = 1000; + + _testProvider.updateTotalPointsInHashRing(serviceName, scheme, totalPoints); + + assertEquals(_testProvider.getCallCount("updateTotalPointsInHashRing"), 1); + assertEquals(_testProvider.getLastServiceName("updateTotalPointsInHashRing"), serviceName); + assertEquals(_testProvider.getLastScheme("updateTotalPointsInHashRing"), scheme); + assertEquals(_testProvider.getLastIntValue("updateTotalPointsInHashRing").intValue(), totalPoints); + } + + @Test + public void testUpdateTotalPointsInHashRingZero() + { + _testProvider.updateTotalPointsInHashRing("test-service", "http", 0); + + assertEquals(_testProvider.getCallCount("updateTotalPointsInHashRing"), 1); + assertEquals(_testProvider.getLastIntValue("updateTotalPointsInHashRing").intValue(), 0); + } + + @Test + public void testDifferentServiceNames() + { + _testProvider.updateTotalPointsInHashRing("service-A", "http", 100); + _testProvider.updateTotalPointsInHashRing("service-B", "https", 200); + _testProvider.updateTotalPointsInHashRing("service-C", "http", 300); + + assertEquals(_testProvider.getCallCount("updateTotalPointsInHashRing"), 3); + assertEquals(_testProvider.getLastServiceName("updateTotalPointsInHashRing"), "service-C"); + assertEquals(_testProvider.getLastScheme("updateTotalPointsInHashRing"), "http"); + assertEquals(_testProvider.getLastIntValue("updateTotalPointsInHashRing").intValue(), 300); + } + + @Test + public void testDifferentSchemes() + { + _testProvider.updateOverrideClusterDropRate("my-service", "http", 0.1); + _testProvider.updateOverrideClusterDropRate("my-service", "https", 0.2); + + assertEquals(_testProvider.getCallCount("updateOverrideClusterDropRate"), 2); + assertEquals(_testProvider.getLastScheme("updateOverrideClusterDropRate"), "https"); + assertEquals(_testProvider.getLastDoubleValue("updateOverrideClusterDropRate"), 0.2, 1e-9); + } + + @Test + public void testLatencyIsolatedByServiceAndScheme() + { + _testProvider.recordHostLatency("service-A", "http", 100L, PerCallDurationSemantics.FULL_ROUND_TRIP); + _testProvider.recordHostLatency("service-A", "https", 200L, PerCallDurationSemantics.FULL_ROUND_TRIP); + _testProvider.recordHostLatency("service-B", "http", 300L, PerCallDurationSemantics.FULL_ROUND_TRIP); + + List serviceAHttp = _testProvider.getAllLatencyValues("service-A", "http"); + List serviceAHttps = _testProvider.getAllLatencyValues("service-A", "https"); + List serviceBHttp = _testProvider.getAllLatencyValues("service-B", "http"); + + assertEquals(serviceAHttp, Arrays.asList(100L)); + assertEquals(serviceAHttps, Arrays.asList(200L)); + assertEquals(serviceBHttp, Arrays.asList(300L)); + } + + @Test + public void testAllMethodsCalled() + { + String serviceName = "comprehensive-service"; + String scheme = "https"; + + _testProvider.recordHostLatency(serviceName, scheme, 100L, PerCallDurationSemantics.FULL_ROUND_TRIP); + _testProvider.recordHostLatency(serviceName, scheme, 150L, PerCallDurationSemantics.FULL_ROUND_TRIP); + _testProvider.recordHostLatency(serviceName, scheme, 200L, PerCallDurationSemantics.FULL_ROUND_TRIP); + + _testProvider.updateOverrideClusterDropRate(serviceName, scheme, 0.0); + _testProvider.updateTotalPointsInHashRing(serviceName, scheme, 1000); + + assertEquals(_testProvider.getCallCount("recordHostLatency"), 3); + assertEquals(_testProvider.getCallCount("updateOverrideClusterDropRate"), 1); + assertEquals(_testProvider.getCallCount("updateTotalPointsInHashRing"), 1); + } + + @Test + public void testDropRateIncreasedDuringDegradation() + { + String serviceName = "degrading-service"; + String scheme = "http"; + + _testProvider.updateOverrideClusterDropRate(serviceName, scheme, 0.0); + _testProvider.updateOverrideClusterDropRate(serviceName, scheme, 0.1); + _testProvider.updateOverrideClusterDropRate(serviceName, scheme, 0.3); + _testProvider.updateOverrideClusterDropRate(serviceName, scheme, 0.5); + + assertEquals(_testProvider.getCallCount("updateOverrideClusterDropRate"), 4); + assertEquals(_testProvider.getLastDoubleValue("updateOverrideClusterDropRate"), 0.5, 1e-9); + } + + @Test + public void testReset() + { + String serviceName = "test-service-reset"; + String scheme = "http"; + + _testProvider.recordHostLatency(serviceName, scheme, 100L, PerCallDurationSemantics.FULL_ROUND_TRIP); + _testProvider.updateOverrideClusterDropRate(serviceName, scheme, 0.2); + _testProvider.updateTotalPointsInHashRing(serviceName, scheme, 500); + + assertEquals(_testProvider.getCallCount("recordHostLatency"), 1); + assertEquals(_testProvider.getCallCount("updateOverrideClusterDropRate"), 1); + assertEquals(_testProvider.getCallCount("updateTotalPointsInHashRing"), 1); + + _testProvider.reset(); + + assertEquals(_testProvider.getCallCount("recordHostLatency"), 0); + assertEquals(_testProvider.getCallCount("updateOverrideClusterDropRate"), 0); + assertEquals(_testProvider.getCallCount("updateTotalPointsInHashRing"), 0); + } + + @Test + public void testNoOpProviderDoesNotThrow() + { + NoOpDegraderLoadBalancerStrategyV3OtelMetricsProvider noOpProvider = + new NoOpDegraderLoadBalancerStrategyV3OtelMetricsProvider(); + + noOpProvider.recordHostLatency("service", "http", 150L, PerCallDurationSemantics.FULL_ROUND_TRIP); + + noOpProvider.updateOverrideClusterDropRate("service", "https", 0.25); + noOpProvider.updateTotalPointsInHashRing("service", "http", 1000); + } + + @Test + public void testLatencyHistogramDistribution() + { + String serviceName = "production-service"; + String scheme = "https"; + + long[] callLatencies = {50L, 55L, 60L, 65L, 70L, 75L, 100L, 150L, 200L, 500L}; + for (long latency : callLatencies) { + _testProvider.recordHostLatency(serviceName, scheme, latency, PerCallDurationSemantics.FULL_ROUND_TRIP); + } + + List recorded = _testProvider.getAllLatencyValues(serviceName, scheme); + assertEquals(recorded.size(), callLatencies.length); + + assertNotNull(_testProvider.getLastLongValue("recordHostLatency")); + assertEquals(_testProvider.getLastLongValue("recordHostLatency").longValue(), 500L); + } +} diff --git a/d2/src/test/java/com/linkedin/d2/jmx/RelativeLoadBalancerStrategyOtelMetricsProviderTest.java b/d2/src/test/java/com/linkedin/d2/jmx/RelativeLoadBalancerStrategyOtelMetricsProviderTest.java new file mode 100644 index 0000000000..566348630a --- /dev/null +++ b/d2/src/test/java/com/linkedin/d2/jmx/RelativeLoadBalancerStrategyOtelMetricsProviderTest.java @@ -0,0 +1,212 @@ +/* + Copyright (c) 2026 LinkedIn Corp. + + Licensed 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 com.linkedin.d2.jmx; + +import com.linkedin.d2.balancer.clients.PerCallDurationSemantics; + +import java.util.Arrays; +import java.util.List; + +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; + +/** Tests for {@link RelativeLoadBalancerStrategyOtelMetricsProvider}. */ +public class RelativeLoadBalancerStrategyOtelMetricsProviderTest +{ + private TestRelativeLoadBalancerStrategyOtelMetricsProvider _testProvider; + + @BeforeMethod + public void setUp() + { + _testProvider = new TestRelativeLoadBalancerStrategyOtelMetricsProvider(); + } + + @DataProvider(name = "intMethodProvider") + public Object[][] intMethodProvider() + { + return new Object[][] { + {"updateTotalHostsInAllPartitionsCount", 100}, + {"updateTotalPointsInHashRing", 1000} + }; + } + + @Test + public void testRecordHostLatency() + { + String serviceName = "test-service-latency"; + String scheme = "http"; + long latencyMs = 150L; + + _testProvider.recordHostLatency(serviceName, scheme, latencyMs, PerCallDurationSemantics.FULL_ROUND_TRIP); + + assertEquals(_testProvider.getCallCount("recordHostLatency"), 1); + assertEquals(_testProvider.getLastServiceName("recordHostLatency"), serviceName); + assertEquals(_testProvider.getLastLongValue("recordHostLatency").longValue(), latencyMs); + } + + @Test + public void testRecordMultipleHostLatencies() + { + String serviceName = "test-service-multi-latency"; + String scheme = "https"; + + _testProvider.recordHostLatency(serviceName, scheme, 100L, PerCallDurationSemantics.FULL_ROUND_TRIP); + _testProvider.recordHostLatency(serviceName, scheme, 150L, PerCallDurationSemantics.FULL_ROUND_TRIP); + _testProvider.recordHostLatency(serviceName, scheme, 200L, PerCallDurationSemantics.FULL_ROUND_TRIP); + _testProvider.recordHostLatency(serviceName, scheme, 120L, PerCallDurationSemantics.FULL_ROUND_TRIP); + _testProvider.recordHostLatency(serviceName, scheme, 180L, PerCallDurationSemantics.FULL_ROUND_TRIP); + + assertEquals(_testProvider.getCallCount("recordHostLatency"), 5); + assertEquals(_testProvider.getLastServiceName("recordHostLatency"), serviceName); + assertEquals(_testProvider.getLastScheme("recordHostLatency"), scheme); + assertEquals(_testProvider.getLastLongValue("recordHostLatency").longValue(), 180L); + + List allLatencies = _testProvider.getAllLatencyValues(serviceName, scheme); + assertEquals(allLatencies.size(), 5); + assertEquals(allLatencies, Arrays.asList(100L, 150L, 200L, 120L, 180L)); + } + + @Test(dataProvider = "intMethodProvider") + public void testIntUpdateMethods(String methodName, int value) + { + String serviceName = "test-service-" + methodName; + String scheme = "http"; + + switch (methodName) + { + case "updateTotalHostsInAllPartitionsCount": + _testProvider.updateTotalHostsInAllPartitionsCount(serviceName, scheme, value); + break; + case "updateTotalPointsInHashRing": + _testProvider.updateTotalPointsInHashRing(serviceName, scheme, value); + break; + default: + throw new IllegalArgumentException("Unknown method: " + methodName); + } + + assertEquals(_testProvider.getCallCount(methodName), 1); + assertEquals(_testProvider.getLastServiceName(methodName), serviceName); + assertEquals(_testProvider.getLastScheme(methodName), scheme); + assertEquals(_testProvider.getLastIntValue(methodName).intValue(), value); + } + + @Test + public void testDifferentServiceNames() + { + _testProvider.updateTotalHostsInAllPartitionsCount("service-A", "http", 10); + _testProvider.updateTotalHostsInAllPartitionsCount("service-B", "https", 20); + _testProvider.updateTotalHostsInAllPartitionsCount("service-C", "http", 30); + + assertEquals(_testProvider.getCallCount("updateTotalHostsInAllPartitionsCount"), 3); + assertEquals(_testProvider.getLastServiceName("updateTotalHostsInAllPartitionsCount"), "service-C"); + assertEquals(_testProvider.getLastScheme("updateTotalHostsInAllPartitionsCount"), "http"); + assertEquals(_testProvider.getLastIntValue("updateTotalHostsInAllPartitionsCount").intValue(), 30); + } + + @Test + public void testAllMethodsCalled() + { + String serviceName = "comprehensive-service"; + String scheme = "https"; + + _testProvider.recordHostLatency(serviceName, scheme, 100L, PerCallDurationSemantics.FULL_ROUND_TRIP); + _testProvider.recordHostLatency(serviceName, scheme, 150L, PerCallDurationSemantics.FULL_ROUND_TRIP); + _testProvider.recordHostLatency(serviceName, scheme, 200L, PerCallDurationSemantics.FULL_ROUND_TRIP); + + _testProvider.updateTotalHostsInAllPartitionsCount(serviceName, scheme, 100); + _testProvider.updateDegradedHostsCount(serviceName, scheme, HostStatus.UNHEALTHY, 3); + _testProvider.updateDegradedHostsCount(serviceName, scheme, HostStatus.QUARANTINED, 2); + _testProvider.updateTotalPointsInHashRing(serviceName, scheme, 1000); + + assertEquals(_testProvider.getCallCount("recordHostLatency"), 3); + + assertEquals(_testProvider.getCallCount("updateTotalHostsInAllPartitionsCount"), 1); + assertEquals(_testProvider.getCallCount("updateDegradedHostsCount"), 2); + assertEquals(_testProvider.getCallCountForHostStatus("updateDegradedHostsCount", HostStatus.UNHEALTHY), 1); + assertEquals(_testProvider.getCallCountForHostStatus("updateDegradedHostsCount", HostStatus.QUARANTINED), 1); + assertEquals(_testProvider.getLastIntValueForHostStatus("updateDegradedHostsCount", HostStatus.UNHEALTHY).intValue(), 3); + assertEquals(_testProvider.getLastIntValueForHostStatus("updateDegradedHostsCount", HostStatus.QUARANTINED).intValue(), 2); + assertEquals(_testProvider.getCallCount("updateTotalPointsInHashRing"), 1); + } + + @Test + public void testReset() + { + String serviceName = "test-service-reset"; + String scheme = "http"; + + _testProvider.recordHostLatency(serviceName, scheme, 100L, PerCallDurationSemantics.FULL_ROUND_TRIP); + _testProvider.updateTotalHostsInAllPartitionsCount(serviceName, scheme, 10); + + assertEquals(_testProvider.getCallCount("recordHostLatency"), 1); + assertEquals(_testProvider.getCallCount("updateTotalHostsInAllPartitionsCount"), 1); + + _testProvider.reset(); + + assertEquals(_testProvider.getCallCount("recordHostLatency"), 0); + assertEquals(_testProvider.getCallCount("updateTotalHostsInAllPartitionsCount"), 0); + } + + @Test + public void testZeroValues() + { + String serviceName = "test-service-zero"; + String scheme = "http"; + + _testProvider.recordHostLatency(serviceName, scheme, 0L, PerCallDurationSemantics.FULL_ROUND_TRIP); + _testProvider.updateTotalHostsInAllPartitionsCount(serviceName, scheme, 0); + _testProvider.updateDegradedHostsCount(serviceName, scheme, HostStatus.UNHEALTHY, 0); + _testProvider.updateDegradedHostsCount(serviceName, scheme, HostStatus.QUARANTINED, 0); + + assertEquals(_testProvider.getLastLongValue("recordHostLatency").longValue(), 0L); + assertEquals(_testProvider.getLastIntValue("updateTotalHostsInAllPartitionsCount").intValue(), 0); + assertEquals(_testProvider.getLastIntValueForHostStatus("updateDegradedHostsCount", HostStatus.UNHEALTHY).intValue(), 0); + assertEquals(_testProvider.getLastIntValueForHostStatus("updateDegradedHostsCount", HostStatus.QUARANTINED).intValue(), 0); + } + + @Test + public void testNoOpProviderDoesNotThrow() + { + NoOpRelativeLoadBalancerStrategyOtelMetricsProvider noOpProvider = + new NoOpRelativeLoadBalancerStrategyOtelMetricsProvider(); + + noOpProvider.recordHostLatency("service", "http", 150L, PerCallDurationSemantics.FULL_ROUND_TRIP); + + noOpProvider.updateTotalHostsInAllPartitionsCount("service", "https", 100); + noOpProvider.updateDegradedHostsCount("service", "http", HostStatus.UNHEALTHY, 3); + noOpProvider.updateDegradedHostsCount("service", "https", HostStatus.QUARANTINED, 2); + noOpProvider.updateTotalPointsInHashRing("service", "http", 1000); + } + + @Test + public void testLatencyHistogramDistribution() + { + String serviceName = "production-service"; + String scheme = "https"; + + long[] hostLatencies = {50L, 55L, 60L, 65L, 70L, 75L, 100L, 150L, 200L, 500L}; + for (long latency : hostLatencies) + { + _testProvider.recordHostLatency(serviceName, scheme, latency, PerCallDurationSemantics.FULL_ROUND_TRIP); + } + + List recordedLatencies = _testProvider.getAllLatencyValues(serviceName, scheme); + assertEquals(recordedLatencies.size(), 10); + } +} diff --git a/d2/src/test/java/com/linkedin/d2/jmx/TestDegraderLoadBalancerStrategyV3OtelMetricsProvider.java b/d2/src/test/java/com/linkedin/d2/jmx/TestDegraderLoadBalancerStrategyV3OtelMetricsProvider.java new file mode 100644 index 0000000000..f0f235ac9d --- /dev/null +++ b/d2/src/test/java/com/linkedin/d2/jmx/TestDegraderLoadBalancerStrategyV3OtelMetricsProvider.java @@ -0,0 +1,43 @@ +/* + Copyright (c) 2026 LinkedIn Corp. + + Licensed 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 com.linkedin.d2.jmx; + +import com.linkedin.d2.balancer.clients.PerCallDurationSemantics; + + +/** Recording test double for {@link DegraderLoadBalancerStrategyV3OtelMetricsProvider}. */ +public class TestDegraderLoadBalancerStrategyV3OtelMetricsProvider extends AbstractRecordingOtelMetricsProvider + implements DegraderLoadBalancerStrategyV3OtelMetricsProvider +{ + @Override + public void recordHostLatency(String serviceName, String scheme, long hostLatencyMs, + PerCallDurationSemantics semantics) + { + recordLong("recordHostLatency", serviceName, scheme, hostLatencyMs, semantics); + } + + @Override + public void updateOverrideClusterDropRate(String serviceName, String scheme, double overrideClusterDropRate) + { + recordDouble("updateOverrideClusterDropRate", serviceName, scheme, overrideClusterDropRate); + } + + @Override + public void updateTotalPointsInHashRing(String serviceName, String scheme, int totalPointsInHashRing) + { + recordInt("updateTotalPointsInHashRing", serviceName, scheme, totalPointsInHashRing); + } +} diff --git a/d2/src/test/java/com/linkedin/d2/jmx/TestRelativeLoadBalancerStrategyOtelMetricsProvider.java b/d2/src/test/java/com/linkedin/d2/jmx/TestRelativeLoadBalancerStrategyOtelMetricsProvider.java new file mode 100644 index 0000000000..2b53036a8c --- /dev/null +++ b/d2/src/test/java/com/linkedin/d2/jmx/TestRelativeLoadBalancerStrategyOtelMetricsProvider.java @@ -0,0 +1,49 @@ +/* + Copyright (c) 2026 LinkedIn Corp. + + Licensed 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 com.linkedin.d2.jmx; + +import com.linkedin.d2.balancer.clients.PerCallDurationSemantics; + + +/** Recording test double for {@link RelativeLoadBalancerStrategyOtelMetricsProvider}. */ +public class TestRelativeLoadBalancerStrategyOtelMetricsProvider extends AbstractRecordingOtelMetricsProvider + implements RelativeLoadBalancerStrategyOtelMetricsProvider +{ + @Override + public void recordHostLatency(String serviceName, String scheme, long hostLatencyMs, + PerCallDurationSemantics semantics) + { + recordLong("recordHostLatency", serviceName, scheme, hostLatencyMs, semantics); + } + + @Override + public void updateTotalHostsInAllPartitionsCount(String serviceName, String scheme, int totalHostsInAllPartitionsCount) + { + recordInt("updateTotalHostsInAllPartitionsCount", serviceName, scheme, totalHostsInAllPartitionsCount); + } + + @Override + public void updateDegradedHostsCount(String serviceName, String scheme, HostStatus status, int count) + { + recordInt("updateDegradedHostsCount", serviceName, scheme, count, status); + } + + @Override + public void updateTotalPointsInHashRing(String serviceName, String scheme, int totalPointsInHashRing) + { + recordInt("updateTotalPointsInHashRing", serviceName, scheme, totalPointsInHashRing); + } +}