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); + } +}