diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/authenticator/AzureFederatedCallbackHandler.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/authenticator/AzureFederatedCallbackHandler.java index 53e5eee03a..9281ab488f 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/authenticator/AzureFederatedCallbackHandler.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/authenticator/AzureFederatedCallbackHandler.java @@ -56,9 +56,14 @@ public void configure(final Map configs, final String saslMechanism, final Map stsHeaderOverrides = decodeStsHeaderOverrides(options.get(OPT_STS_HEADER_OVERRIDES)); final AwsCredentialsSupplier awsCredentialsSupplier = AwsCredentialsSupplierProvider.getInstance().getAwsCredentialsSupplier(); + // Nullable: null when no source published a metrics object (e.g. buffer/admin paths), mirroring + // how the AWS credentials supplier is read from its own process-global singleton above. + final KafkaSourceAuthMetrics authMetrics = + KafkaSourceAuthMetricsProvider.getInstance().getAuthMetrics(); this.tokenProvider = new AzureFederatedTokenProvider( options.get(OPT_REGION), options.get(OPT_STS_ROLE_ARN), options.get(OPT_TOKEN_ENDPOINT), - options.get(OPT_CLIENT_ID), options.get(OPT_SCOPE), stsHeaderOverrides, awsCredentialsSupplier); + options.get(OPT_CLIENT_ID), options.get(OPT_SCOPE), stsHeaderOverrides, awsCredentialsSupplier, + authMetrics); } private static Map decodeStsHeaderOverrides(final String encoded) { diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/authenticator/AzureFederatedTokenProvider.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/authenticator/AzureFederatedTokenProvider.java index 218a3127a3..006846b515 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/authenticator/AzureFederatedTokenProvider.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/authenticator/AzureFederatedTokenProvider.java @@ -36,8 +36,6 @@ import java.util.function.Function; import java.util.function.Supplier; -import static org.opensearch.dataprepper.logging.DataPrepperMarkers.SENSITIVE; - public class AzureFederatedTokenProvider { private static final Logger LOG = LoggerFactory.getLogger(AzureFederatedTokenProvider.class); @@ -57,6 +55,8 @@ public class AzureFederatedTokenProvider { private final String scope; private final Supplier stsClientSupplier; private final HttpClient httpClient; + // Null on paths without a PluginMetrics handle (buffer/admin); metric sites are then no-ops. + private final KafkaSourceAuthMetrics authMetrics; private final ReentrantLock lock = new ReentrantLock(); private volatile CachedToken cached; @@ -66,12 +66,21 @@ public AzureFederatedTokenProvider(final String region, final String stsRoleArn, final String clientId, final String scope, final Map stsHeaderOverrides, final AwsCredentialsSupplier awsCredentialsSupplier) { + this(region, stsRoleArn, tokenEndpoint, clientId, scope, stsHeaderOverrides, awsCredentialsSupplier, null); + } + + public AzureFederatedTokenProvider(final String region, final String stsRoleArn, final String tokenEndpoint, + final String clientId, final String scope, + final Map stsHeaderOverrides, + final AwsCredentialsSupplier awsCredentialsSupplier, + final KafkaSourceAuthMetrics authMetrics) { this(region, stsRoleArn, tokenEndpoint, clientId, scope, stsHeaderOverrides, awsCredentialsSupplier, provider -> StsClient.builder() .region(Region.of(region)) .credentialsProvider(provider) .build(), - HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(10)).build()); + HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(10)).build(), + authMetrics); } AzureFederatedTokenProvider(final String region, final String stsRoleArn, final String tokenEndpoint, @@ -80,6 +89,17 @@ public AzureFederatedTokenProvider(final String region, final String stsRoleArn, final AwsCredentialsSupplier awsCredentialsSupplier, final Function stsClientFactory, final HttpClient httpClient) { + this(region, stsRoleArn, tokenEndpoint, clientId, scope, stsHeaderOverrides, awsCredentialsSupplier, + stsClientFactory, httpClient, null); + } + + AzureFederatedTokenProvider(final String region, final String stsRoleArn, final String tokenEndpoint, + final String clientId, final String scope, + final Map stsHeaderOverrides, + final AwsCredentialsSupplier awsCredentialsSupplier, + final Function stsClientFactory, + final HttpClient httpClient, + final KafkaSourceAuthMetrics authMetrics) { this.region = region; this.stsRoleArn = stsRoleArn; this.tokenEndpoint = tokenEndpoint; @@ -88,6 +108,7 @@ public AzureFederatedTokenProvider(final String region, final String stsRoleArn, this.stsClientSupplier = () -> stsClientFactory.apply(baseCredentials(region, stsRoleArn, stsHeaderOverrides, awsCredentialsSupplier)); this.httpClient = httpClient; + this.authMetrics = authMetrics; } String getRegion() { @@ -120,8 +141,20 @@ public AzureFederatedOAuthBearerToken getToken() { if (cached != null && cached.isValid()) { return cached.newTokenInstance(); } - final AzureFederatedOAuthBearerToken token = exchange(); + final AzureFederatedOAuthBearerToken token; + try { + token = exchange(); + } catch (final RuntimeException e) { + if (authMetrics != null) { + authMetrics.recordFailure(classifyFailure(e)); + } + throw e; + } cached = new CachedToken(token); + if (authMetrics != null) { + // Headroom to the proactive re-mint deadline (expiry minus skew), not raw expiry. + authMetrics.recordRefresh(token.lifetimeMs() - SKEW_BUFFER_MS); + } return token; } finally { lock.unlock(); @@ -132,6 +165,33 @@ private AzureFederatedOAuthBearerToken exchange() { return postToAzure(mintAwsJwt()); } + // Walk the cause chain to a bounded cause tag; anything unmatched is CAUSE_OTHER. + private static String classifyFailure(final RuntimeException e) { + for (Throwable t = e; t != null; t = t.getCause()) { + if (t instanceof OutboundWebIdentityFederationDisabledException) { + return KafkaSourceAuthMetrics.CAUSE_AWS_OUTBOUND_FEDERATION_DISABLED; + } + if (t instanceof StsException) { + final StsException stsException = (StsException) t; + final String errorCode = errorCodeOf(stsException); + return errorCode != null && errorCode.contains("AccessDenied") + ? KafkaSourceAuthMetrics.CAUSE_AWS_STS_ACCESS_DENIED + : KafkaSourceAuthMetrics.CAUSE_AWS_STS_ERROR; + } + if (t instanceof AzureTokenExchangeException) { + return KafkaSourceAuthMetrics.CAUSE_AZURE_TOKEN_EXCHANGE_REJECTED; + } + if (t instanceof IOException || t instanceof InterruptedException) { + return KafkaSourceAuthMetrics.CAUSE_NETWORK; + } + } + return KafkaSourceAuthMetrics.CAUSE_OTHER; + } + + private static String errorCodeOf(final StsException e) { + return e.awsErrorDetails() != null ? e.awsErrorDetails().errorCode() : null; + } + private StsClient stsClient() { if (stsClient == null) { stsClient = stsClientSupplier.get(); @@ -172,26 +232,20 @@ private String mintAwsJwt() { .durationSeconds(JWT_DURATION_SECONDS)); return response.webIdentityToken(); } catch (final OutboundWebIdentityFederationDisabledException e) { - LOG.error(SENSITIVE, "AWS Outbound Identity Federation is not enabled for {}", identity()); - throw new RuntimeException("AWS Outbound Identity Federation is not enabled for " - + identity() + ". Enable it once per account.", e); + LOG.error("The outbound web identity federation feature is not enabled for the account.", e); + throw new RuntimeException("The outbound web identity federation feature is not enabled for the account.", e); } catch (final StsException e) { - if (e.statusCode() == 403) { - LOG.error(SENSITIVE, "{} lacks sts:GetWebIdentityToken", identity()); - throw new RuntimeException(identity() - + " lacks sts:GetWebIdentityToken (AccessDenied). Add it to the role policy.", e); - } - throw new RuntimeException("STS GetWebIdentityToken failed for " + identity(), e); + final String errorCode = errorCodeOf(e); + LOG.error("STS request failed during azure_federated authentication (HTTP {}, errorCode {}).", + e.statusCode(), errorCode, e); + throw new RuntimeException("STS request failed during azure_federated authentication (HTTP " + + e.statusCode() + ", errorCode " + errorCode + ")", e); } catch (final Exception e) { - throw new RuntimeException("AWS credential resolution failed for " + identity() - + " (verify the AWS credentials/role configuration)", e); + LOG.error("Unexpected failure during azure_federated authentication.", e); + throw new RuntimeException("Unexpected failure during azure_federated authentication", e); } } - private String identity() { - return stsRoleArn != null ? "Role " + stsRoleArn : "the resolved AWS credentials identity"; - } - private AzureFederatedOAuthBearerToken postToAzure(final String awsJwt) { final String form = "grant_type=" + encode(GRANT_TYPE) + "&client_id=" + encode(clientId) @@ -207,8 +261,8 @@ private AzureFederatedOAuthBearerToken postToAzure(final String awsJwt) { final HttpResponse response = httpClient.send(request, HttpResponse.BodyHandlers.ofString()); if (response.statusCode() != 200) { final String aadsts = extractAadstsError(response.body()); - LOG.error(SENSITIVE, "Azure token exchange failed: HTTP {} {}", response.statusCode(), aadsts); - throw new RuntimeException("Azure token exchange failed: HTTP " + response.statusCode() + " " + aadsts); + LOG.error("Azure token exchange was rejected (HTTP {}, {}).", response.statusCode(), aadsts); + throw new AzureTokenExchangeException(response.statusCode(), aadsts); } final JsonNode json = OBJECT_MAPPER.readTree(response.body()); final String accessToken = json.get("access_token").asText(); @@ -237,6 +291,14 @@ private static String encode(final String value) { return URLEncoder.encode(value, StandardCharsets.UTF_8); } + // Dedicated type so the failure metric classifies the Azure non-200 path by instanceof rather than + // by message text (which is then free to change). The AADSTS code is surfaced, not interpreted. + static final class AzureTokenExchangeException extends RuntimeException { + AzureTokenExchangeException(final int statusCode, final String aadsts) { + super("Azure token exchange was rejected (HTTP " + statusCode + ", " + aadsts + ")"); + } + } + private final class CachedToken { private final AzureFederatedOAuthBearerToken token; private final long validUntilMs; diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/authenticator/KafkaSourceAuthMetrics.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/authenticator/KafkaSourceAuthMetrics.java new file mode 100644 index 0000000000..0dcf94f935 --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/authenticator/KafkaSourceAuthMetrics.java @@ -0,0 +1,87 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + */ + +package org.opensearch.dataprepper.plugins.kafka.authenticator; + +import io.micrometer.core.instrument.Counter; +import org.opensearch.dataprepper.metrics.PluginMetrics; + +import java.util.function.LongSupplier; + +/** + * Operator-facing metrics for the azure_federated OAUTHBEARER token lifecycle. The success counter + * and the refresh-headroom gauge are registered in the constructor; the per-cause failure counters are + * registered on first use (the registry dedupes by name+tags). This keeps + * {@link AzureFederatedTokenProvider} free of any direct Micrometer usage (mirroring the pattern + * used by KafkaTopicConsumerMetrics). These are internal/operator signals only - they are not + * vended to customers. + */ +public class KafkaSourceAuthMetrics { + static final String TOKEN_REFRESH_SUCCESS = "tokenRefreshSuccess"; + static final String TOKEN_REFRESH_FAILURES = "tokenRefreshFailures"; + static final String TIME_TO_TOKEN_REFRESH = "timeToTokenRefresh"; + static final String ERROR_TYPE_TAG = "errorType"; + + // Failure cause values (bounded, low cardinality). The token provider already distinguishes these. + public static final String CAUSE_AWS_STS_ACCESS_DENIED = "aws_sts_access_denied"; + public static final String CAUSE_AWS_STS_ERROR = "aws_sts_error"; + public static final String CAUSE_AWS_OUTBOUND_FEDERATION_DISABLED = "aws_outbound_federation_disabled"; + public static final String CAUSE_AZURE_TOKEN_EXCHANGE_REJECTED = "azure_token_exchange_rejected"; + public static final String CAUSE_NETWORK = "network"; + public static final String CAUSE_OTHER = "other"; + + private final PluginMetrics pluginMetrics; + private final Counter tokenRefreshSuccess; + private final LongSupplier clock; + + // Absolute epoch millis of the refresh deadline; 0 = no token minted yet. Written on the refresh + // thread, read on the scrape thread, hence volatile. + private volatile long refreshDeadlineMs = 0L; + + public KafkaSourceAuthMetrics(final PluginMetrics pluginMetrics) { + this(pluginMetrics, System::currentTimeMillis); + } + + KafkaSourceAuthMetrics(final PluginMetrics pluginMetrics, final LongSupplier clock) { + this.pluginMetrics = pluginMetrics; + this.clock = clock; + this.tokenRefreshSuccess = pluginMetrics.counter(TOKEN_REFRESH_SUCCESS); + // Headroom is computed at scrape time so it counts down between refreshes. The gauge holds only + // a weak ref to this object, so KafkaSourceAuthMetricsProvider must keep it strongly referenced. + pluginMetrics.gauge(TIME_TO_TOKEN_REFRESH, this, KafkaSourceAuthMetrics::currentRefreshHeadroomSeconds); + } + + double currentRefreshHeadroomSeconds() { + final long deadline = refreshDeadlineMs; + if (deadline == 0L) { + return 0.0; + } + return Math.max(0.0, (deadline - clock.getAsLong()) / 1000.0); + } + + /** + * Record a successful token (re)mint. {@code refreshDeadlineMs} is the absolute epoch-millis + * instant at which the provider will proactively re-mint (token expiry minus the skew buffer); + * the timeToTokenRefresh gauge reports headroom to this deadline. + */ + public void recordRefresh(final long refreshDeadlineMs) { + this.refreshDeadlineMs = refreshDeadlineMs; + tokenRefreshSuccess.increment(); + } + + /** + * Record a failed token (re)mint, tagged by {@code cause} (one of the {@code CAUSE_*} constants) + * so an operator can route by remediation path (role policy vs Azure app registration vs network). + * The registry dedupes by name+tags, so repeated causes resolve to the same counter. + */ + public void recordFailure(final String cause) { + pluginMetrics.counterWithTags(TOKEN_REFRESH_FAILURES, ERROR_TYPE_TAG, cause).increment(); + } +} diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/authenticator/KafkaSourceAuthMetricsProvider.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/authenticator/KafkaSourceAuthMetricsProvider.java new file mode 100644 index 0000000000..a23978700a --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/authenticator/KafkaSourceAuthMetricsProvider.java @@ -0,0 +1,44 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + */ + +package org.opensearch.dataprepper.plugins.kafka.authenticator; + +/** + * Process-global handoff for the azure_federated auth metrics object, mirroring + * {@link AwsCredentialsSupplierProvider}. {@link AzureFederatedCallbackHandler} is instantiated by + * kafka-clients via reflection (no-arg constructor), so a {@code PluginMetrics}-backed object cannot + * be passed down the constructor chain; the source publishes it here before Kafka client construction + * and the handler reads it in {@code configure()}. + * + *

Like {@link AwsCredentialsSupplierProvider}, this is a plain last-writer-wins singleton. If two + * federated Kafka sources ever ran in the same JVM, the second {@code set()} would overwrite the + * first and its refreshes would be attributed to the second source's metric prefix. This does not + * occur in the target deployment (one Kafka source per container), and it is the same limitation the + * sibling {@link AwsCredentialsSupplierProvider} already accepts. + */ +public class KafkaSourceAuthMetricsProvider { + private static final KafkaSourceAuthMetricsProvider singleton = new KafkaSourceAuthMetricsProvider(); + + public static KafkaSourceAuthMetricsProvider getInstance() { + return singleton; + } + + private volatile KafkaSourceAuthMetrics authMetrics; + + protected KafkaSourceAuthMetricsProvider() {} + + public KafkaSourceAuthMetrics getAuthMetrics() { + return authMetrics; + } + + public void set(final KafkaSourceAuthMetrics authMetrics) { + this.authMetrics = authMetrics; + } +} diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumer.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumer.java index 74aa88520e..c3e076232b 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumer.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumer.java @@ -235,6 +235,8 @@ void consumeRecords() throws Exception { resetAuthBackoff(); LOG.debug("Consumed records with count {}", records.count()); if (Objects.nonNull(records) && !records.isEmpty() && records.count() > 0) { + // Time poll() return through buffer handoff, once per non-empty batch. + final long processingStartMillis = System.currentTimeMillis(); Map offsets = new HashMap<>(); AcknowledgementSet acknowledgementSet = null; if (acknowledgementsEnabled) { @@ -250,6 +252,7 @@ void consumeRecords() throws Exception { acknowledgementSet.complete(); numberOfAcksPending.incrementAndGet(); } + topicMetrics.recordProcessingLatency(System.currentTimeMillis() - processingStartMillis); } } catch (AuthenticationException e) { authFailureAttempts++; @@ -620,6 +623,7 @@ private void iterateRecordPartitions(ConsumerRecords records, fin } public void closeConsumer(){ + topicMetrics.deregister(consumer); consumer.close(); } @@ -629,6 +633,8 @@ public void shutdownConsumer(){ @Override public void onPartitionsAssigned(Collection partitions) { + // Monotonic count of assignments; unlike joinRate it cannot miss brief events between scrapes. + topicMetrics.getNumberOfRebalances().increment(); synchronized(this) { final long epoch = getCurrentTimeNanos(); @@ -649,6 +655,7 @@ public void onPartitionsAssigned(Collection partitions) { @Override public void onPartitionsRevoked(Collection partitions) { + topicMetrics.getNumberOfPartitionsRevoked().increment(partitions.size()); synchronized(this) { commitOffsets(true); for (TopicPartition topicPartition : partitions) { @@ -659,6 +666,7 @@ public void onPartitionsRevoked(Collection partitions) { LOG.info("Revoked partition {}", topicPartition); ownedPartitionsEpoch.remove(topicPartition); partitionCommitTrackerMap.remove(topicPartition.partition()); + topicMetrics.clearPartitionLag(topicPartition.partition()); } if (paused) { consumer.pause(consumer.assignment()); diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSource.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSource.java index 3265e21ec1..edd392f251 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSource.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/source/KafkaSource.java @@ -44,6 +44,8 @@ import org.opensearch.dataprepper.plugins.kafka.consumer.KafkaCustomConsumer; import org.opensearch.dataprepper.plugins.kafka.consumer.KafkaCustomConsumerFactory; import org.opensearch.dataprepper.plugins.kafka.consumer.PauseConsumePredicate; +import org.opensearch.dataprepper.plugins.kafka.authenticator.KafkaSourceAuthMetrics; +import org.opensearch.dataprepper.plugins.kafka.authenticator.KafkaSourceAuthMetricsProvider; import org.opensearch.dataprepper.plugins.kafka.extension.KafkaClusterConfigSupplier; import org.opensearch.dataprepper.plugins.kafka.util.ClientDNSLookupType; import org.opensearch.dataprepper.plugins.kafka.util.KafkaSecurityConfigurer; @@ -130,10 +132,15 @@ public void start(Buffer> buffer) { setMdc(); Properties authProperties = new Properties(); KafkaSecurityConfigurer.setDynamicSaslClientCallbackHandler(authProperties, sourceConfig, pluginConfigObservable); + // Hand off the auth metrics before client construction (see KafkaSourceAuthMetricsProvider). + // azure_federated only: other mechanisms mint no token, so the meters would stay at zero. + if (usesAzureFederatedAuth()) { + KafkaSourceAuthMetricsProvider.getInstance().set(new KafkaSourceAuthMetrics(pluginMetrics)); + } KafkaSecurityConfigurer.setAuthProperties(authProperties, sourceConfig, awsCredentialsSupplier, LOG); sourceConfig.getTopics().forEach(topic -> { consumerGroupID = topic.getGroupId(); - KafkaTopicConsumerMetrics topicMetrics = new KafkaTopicConsumerMetrics(topic.getName(), pluginMetrics, true); + KafkaTopicConsumerMetrics topicMetrics = new KafkaTopicConsumerMetrics(topic.getName(), pluginMetrics, true, topic.getWorkers()); Properties consumerProperties = getConsumerProperties(topic, authProperties); MessageFormat schema = MessageFormat.getByMessageFormatByName(schemaType); try { @@ -351,6 +358,13 @@ private void setConsumerTopicProperties(Properties properties, TopicConsumerConf KafkaCustomConsumerFactory.setConsumerTopicProperties(properties, topicConfig, consumerGroupID); } + private boolean usesAzureFederatedAuth() { + final AuthConfig authConfig = sourceConfig.getAuthConfig(); + return authConfig != null + && authConfig.getSaslAuthConfig() != null + && authConfig.getSaslAuthConfig().getAzureFederatedAuthConfig() != null; + } + private void setPropertiesForSchemaRegistryConnectivity(Properties properties) { AuthConfig authConfig = sourceConfig.getAuthConfig(); String schemaRegistryApiKey = sourceConfig.getSchemaConfig().getSchemaRegistryApiKey(); diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicConsumerMetrics.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicConsumerMetrics.java index aba517177e..8072743fa4 100644 --- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicConsumerMetrics.java +++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicConsumerMetrics.java @@ -6,6 +6,7 @@ package org.opensearch.dataprepper.plugins.kafka.util; import io.micrometer.core.instrument.Counter; +import io.micrometer.core.instrument.Tag; import io.micrometer.core.instrument.Timer; import org.opensearch.dataprepper.metrics.PluginMetrics; import org.apache.kafka.common.Metric; @@ -15,10 +16,14 @@ import org.slf4j.LoggerFactory; import java.time.Instant; +import java.util.List; import java.util.Objects; import java.util.Map; import java.util.HashMap; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.TimeUnit; +import java.util.function.DoubleUnaryOperator; public class KafkaTopicConsumerMetrics { private static final Logger LOG = LoggerFactory.getLogger(KafkaTopicConsumerMetrics.class); @@ -33,6 +38,15 @@ public class KafkaTopicConsumerMetrics { static final String NUMBER_OF_RECORDS_CONSUMED = "numberOfRecordsConsumed"; static final String NUMBER_OF_BYTES_CONSUMED = "numberOfBytesConsumed"; static final String ACTUAL_POLL_INTERVAL = "actualPollInterval"; + static final String NUMBER_OF_ASSIGNED_PARTITIONS = "numberOfAssignedPartitions"; + static final String NUMBER_OF_ACTIVE_READERS = "numberOfActiveReaders"; + static final String NUMBER_OF_CONFIGURED_WORKERS = "numberOfConfiguredWorkers"; + static final String NUMBER_OF_REBALANCES = "numberOfRebalances"; + static final String NUMBER_OF_PARTITIONS_REVOKED = "numberOfPartitionsRevoked"; + static final String RECORDS_LAG_PER_PARTITION = "recordsLagPerPartition"; + static final String RECORDS_PROCESSING_LATENCY = "recordsProcessingLatency"; + // kafka-clients metric key; stored raw, inverted at read time (see update()). + private static final String ASSIGNED_PARTITIONS = "assigned-partitions"; private final String topicName; private long updateTime; @@ -52,12 +66,29 @@ public class KafkaTopicConsumerMetrics { private final Timer timeBetweenPollCalls; private Instant lastPollTime; + private final Counter numberOfRebalances; + private final Counter numberOfPartitionsRevoked; + private final Timer recordsProcessingLatency; + private final int configuredWorkers; + + private final Map perPartitionLag = new ConcurrentHashMap<>(); + private final Set registeredLagGauges = ConcurrentHashMap.newKeySet(); + public KafkaTopicConsumerMetrics(final String topicName, final PluginMetrics pluginMetrics, final boolean topicNameInMetrics) { + this(topicName, pluginMetrics, topicNameInMetrics, 0); + } + + public KafkaTopicConsumerMetrics(final String topicName, final PluginMetrics pluginMetrics, + final boolean topicNameInMetrics, final int configuredWorkers) { this.pluginMetrics = pluginMetrics; this.topicName = topicName; + this.configuredWorkers = configuredWorkers; this.updateTime = Instant.now().getEpochSecond(); - this.metricValues = new HashMap<>(); + // ConcurrentHashMap: mutated by every worker thread in update()/register() while the metric + // scrape thread iterates it inside the gauge lambdas. A plain HashMap here is a + // ConcurrentModificationException hazard, especially with the new scaling gauges below. + this.metricValues = new ConcurrentHashMap<>(); initializeMetricNamesMap(topicNameInMetrics); this.numberOfRecordsConsumed = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_RECORDS_CONSUMED, topicNameInMetrics)); this.numberOfBytesConsumed = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_BYTES_CONSUMED, topicNameInMetrics)); @@ -70,9 +101,37 @@ public KafkaTopicConsumerMetrics(final String topicName, final PluginMetrics plu this.numberOfPositiveAcknowledgements = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_POSITIVE_ACKNOWLEDGEMENTS, topicNameInMetrics)); this.numberOfNegativeAcknowledgements = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_NEGATIVE_ACKNOWLEDGEMENTS, topicNameInMetrics)); this.timeBetweenPollCalls = pluginMetrics.timer(getTopicMetricName(ACTUAL_POLL_INTERVAL, topicNameInMetrics)); + this.numberOfRebalances = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_REBALANCES, topicNameInMetrics)); + this.numberOfPartitionsRevoked = pluginMetrics.counter(getTopicMetricName(NUMBER_OF_PARTITIONS_REVOKED, topicNameInMetrics)); + this.recordsProcessingLatency = pluginMetrics.timer(getTopicMetricName(RECORDS_PROCESSING_LATENCY, topicNameInMetrics)); + registerScalingGauges(topicNameInMetrics); lastPollTime = Instant.now(); } + private void registerScalingGauges(final boolean topicNameInMetrics) { + pluginMetrics.gauge(getTopicMetricName(NUMBER_OF_ASSIGNED_PARTITIONS, topicNameInMetrics), metricValues, + mv -> foldAssignedPartitions(mv, raw -> raw)); + pluginMetrics.gauge(getTopicMetricName(NUMBER_OF_ACTIVE_READERS, topicNameInMetrics), metricValues, + mv -> foldAssignedPartitions(mv, raw -> raw > 0 ? 1.0 : 0.0)); + // Register against the long-lived metricValues object, NOT gauge(name, Integer) whose weak ref + // to the boxed Integer is GC'd to NaN. + if (configuredWorkers > 0) { + pluginMetrics.gauge(getTopicMetricName(NUMBER_OF_CONFIGURED_WORKERS, topicNameInMetrics), metricValues, mv -> (double) configuredWorkers); + } + } + + private static double foldAssignedPartitions(final Map> mv, + final DoubleUnaryOperator perConsumer) { + double total = 0; + for (final Map.Entry> entry : mv.entrySet()) { + final Map consumerMetrics = entry.getValue(); + synchronized (consumerMetrics) { + total += perConsumer.applyAsDouble(consumerMetrics.getOrDefault(ASSIGNED_PARTITIONS, 0.0)); + } + } + return total; + } + private void initializeMetricNamesMap(final boolean topicNameInMetrics) { this.metricsNameMap = new HashMap<>(); this.metricsNameMap.put("bytes-consumed-total", "bytesConsumedTotal"); @@ -112,6 +171,10 @@ private void initializeMetricNamesMap(final boolean topicNameInMetrics) { } return min; }); + } else if (metricName.equals("assigned-partitions")) { + // numberOfNonConsumers = count of consumers with ZERO assigned partitions (read-time inversion). + pluginMetrics.gauge(getTopicMetricName(camelCaseName, topicNameInMetrics), metricValues, + mv -> foldAssignedPartitions(mv, raw -> raw == 0.0 ? 1.0 : 0.0)); } else if (!metricName.contains("-total")) { pluginMetrics.gauge(getTopicMetricName(camelCaseName, topicNameInMetrics), metricValues, metricValues -> { double sum = 0; @@ -175,6 +238,28 @@ public Counter getNumberOfPositiveAcknowledgements() { return numberOfPositiveAcknowledgements; } + public Counter getNumberOfRebalances() { + return numberOfRebalances; + } + + public Counter getNumberOfPartitionsRevoked() { + return numberOfPartitionsRevoked; + } + + public void recordProcessingLatency(final long durationMillis) { + recordsProcessingLatency.record(durationMillis, TimeUnit.MILLISECONDS); + } + + // Drop a closed consumer's entries so they stop skewing the aggregate gauges. + public void deregister(final KafkaConsumer consumer) { + metricValues.remove(consumer); + } + + // Zero (not remove) a revoked partition's lag: the gauge meter stays registered, only the value resets. + public void clearPartitionLag(final int partition) { + perPartitionLag.put(String.valueOf(partition), 0.0); + } + public void recordTimeBetweenPolls() { final long timeBetweenPolls = Instant.now().toEpochMilli() - lastPollTime.toEpochMilli(); timeBetweenPollCalls.record(timeBetweenPolls, TimeUnit.MILLISECONDS); @@ -211,6 +296,21 @@ public void update(final KafkaConsumer consumer) { if (Objects.nonNull(metricsNameMap.get(metricName))) { if (metric.tags().containsKey("partition") && (metricName.equals("records-lag-max") || metricName.equals("records-lead-min"))) { + if (metricName.equals("records-lag-max")) { + final String partition = metric.tags().get("partition"); + final double lag = (Double) value.metricValue(); + if (!Double.isNaN(lag) && !Double.isInfinite(lag)) { + // plain put (latest value), NOT merge(Math::max): a monotonic max would turn + // this into a high-water-mark rather than a live lag gauge. + perPartitionLag.put(partition, lag); + if (registeredLagGauges.add(partition)) { + // topic pinned in the name regardless of topicNameInMetrics: partition tags collide across topics otherwise. + pluginMetrics.gaugeWithTags(getTopicMetricName(RECORDS_LAG_PER_PARTITION, true), + List.of(Tag.of("partition", partition)), perPartitionLag, + m -> m.getOrDefault(partition, 0.0)); + } + } + } continue; } @@ -236,10 +336,9 @@ public void update(final KafkaConsumer consumer) { numberOfBytesConsumed.increment(newValue - prevValue); } } - // Keep the count of number of consumers without any assigned partitions. This value can go up or down. So, it is made as Guage metric - if (metricName.equals("assigned-partitions")) { - newValue = (newValue == 0.0) ? 1.0 : 0.0; - } + // Store the RAW assigned-partitions value; numberOfNonConsumers, numberOfAssignedPartitions + // and numberOfActiveReaders all derive from it at read time in their gauge lambdas. This + // replaced a destructive write-time inversion that overwrote the raw value. synchronized(consumerMetrics) { consumerMetrics.put(metricName, newValue); } diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/authenticator/AzureFederatedCallbackHandlerTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/authenticator/AzureFederatedCallbackHandlerTest.java index 0de52f247e..399d7100da 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/authenticator/AzureFederatedCallbackHandlerTest.java +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/authenticator/AzureFederatedCallbackHandlerTest.java @@ -11,6 +11,9 @@ package org.opensearch.dataprepper.plugins.kafka.authenticator; import com.fasterxml.jackson.databind.ObjectMapper; +import io.micrometer.core.instrument.Counter; +import io.micrometer.core.instrument.Metrics; +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; import org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule; import org.apache.kafka.common.security.oauthbearer.OAuthBearerTokenCallback; import org.junit.jupiter.api.AfterEach; @@ -18,6 +21,7 @@ import org.mockito.ArgumentCaptor; import org.opensearch.dataprepper.aws.api.AwsCredentialsOptions; import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; +import org.opensearch.dataprepper.metrics.PluginMetrics; import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; import javax.security.auth.callback.Callback; @@ -33,6 +37,9 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.containsString; import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.instanceOf; +import static org.hamcrest.Matchers.not; +import static org.hamcrest.Matchers.notNullValue; import static org.hamcrest.Matchers.sameInstance; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; @@ -68,6 +75,7 @@ private List jaasEntries(final Map option @AfterEach void resetSupplier() { AwsCredentialsSupplierProvider.getInstance().set(null); + KafkaSourceAuthMetricsProvider.getInstance().set(null); } @Test @@ -101,6 +109,58 @@ void configure_readsSupplierFromSingleton_andPassesItToProvider() { verify(supplier).getProvider(any(AwsCredentialsOptions.class)); } + @Test + void configure_readsAuthMetricsFromSingleton_andProviderEmitsFailureMetric() { + // End-to-end wiring: object published on the singleton -> read in configure() -> provider + // instruments a failure. Guards the silent-metrics-loss path (handler not reading the object). + final SimpleMeterRegistry meterRegistry = new SimpleMeterRegistry(); + Metrics.addRegistry(meterRegistry); + // Unique pipeline scope: PluginMetrics writes to the static global registry, so a fixed name + // could collide with meters registered by other test classes in the same JVM. + final String pipelineName = "handler-test-pipeline-" + java.util.UUID.randomUUID(); + try { + final AwsCredentialsSupplier supplier = mock(AwsCredentialsSupplier.class); + when(supplier.getProvider(any(AwsCredentialsOptions.class))) + .thenThrow(new RuntimeException("credential resolution failed")); + AwsCredentialsSupplierProvider.getInstance().set(supplier); + KafkaSourceAuthMetricsProvider.getInstance() + .set(new KafkaSourceAuthMetrics(PluginMetrics.fromNames("kafka", pipelineName))); + final AzureFederatedCallbackHandler handler = new AzureFederatedCallbackHandler(); + + handler.configure(Map.of(), OAuthBearerLoginModule.OAUTHBEARER_MECHANISM, jaasEntries(jaasOptions())); + assertThrows(RuntimeException.class, () -> handler.getTokenProvider().getToken()); + + // The credential-resolution failure classifies as CAUSE_OTHER; the point is a tagged + // tokenRefreshFailures counter exists, proving the metrics object reached the provider. + final Counter failure = meterRegistry.find(pipelineName + ".kafka." + KafkaSourceAuthMetrics.TOKEN_REFRESH_FAILURES) + .tag(KafkaSourceAuthMetrics.ERROR_TYPE_TAG, KafkaSourceAuthMetrics.CAUSE_OTHER).counter(); + assertThat(failure, notNullValue()); + assertThat(failure.count(), equalTo(1.0)); + } finally { + Metrics.removeRegistry(meterRegistry); + meterRegistry.clear(); + meterRegistry.close(); + } + } + + @Test + void configure_withoutAuthMetricsSet_providerHasNoMetricsAndDoesNotThrow() { + // No metrics object published => get() returns null => provider null-guards, no NPE. + final AwsCredentialsSupplier supplier = mock(AwsCredentialsSupplier.class); + when(supplier.getProvider(any(AwsCredentialsOptions.class))) + .thenThrow(new RuntimeException("credential resolution failed")); + AwsCredentialsSupplierProvider.getInstance().set(supplier); + final AzureFederatedCallbackHandler handler = new AzureFederatedCallbackHandler(); + + handler.configure(Map.of(), OAuthBearerLoginModule.OAUTHBEARER_MECHANISM, jaasEntries(jaasOptions())); + + // The original credential-resolution failure must survive; a null-metrics dereference would + // also be a RuntimeException, so assertThrows(RuntimeException) alone would miss that regression. + final RuntimeException thrown = + assertThrows(RuntimeException.class, () -> handler.getTokenProvider().getToken()); + assertThat(thrown, not(instanceOf(NullPointerException.class))); + } + @Test void configure_withStsHeaderOverridesOption_decodesAndPassesThemToProvider() throws Exception { final Map stsHeaderOverrides = Map.of( diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/authenticator/AzureFederatedTokenProviderTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/authenticator/AzureFederatedTokenProviderTest.java index dfe8ab0c6d..01cb74185a 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/authenticator/AzureFederatedTokenProviderTest.java +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/authenticator/AzureFederatedTokenProviderTest.java @@ -10,6 +10,10 @@ package org.opensearch.dataprepper.plugins.kafka.authenticator; +import io.micrometer.core.instrument.Counter; +import io.micrometer.core.instrument.Metrics; +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; @@ -19,7 +23,9 @@ import org.mockito.quality.Strictness; import org.opensearch.dataprepper.aws.api.AwsCredentialsOptions; import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; +import org.opensearch.dataprepper.metrics.PluginMetrics; import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; +import software.amazon.awssdk.awscore.exception.AwsErrorDetails; import software.amazon.awssdk.services.sts.StsClient; import software.amazon.awssdk.services.sts.model.GetWebIdentityTokenResponse; import software.amazon.awssdk.services.sts.model.OutboundWebIdentityFederationDisabledException; @@ -43,7 +49,9 @@ import static org.hamcrest.Matchers.containsString; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.greaterThan; +import static org.hamcrest.Matchers.instanceOf; import static org.hamcrest.Matchers.not; +import static org.hamcrest.Matchers.notNullValue; import static org.hamcrest.Matchers.sameInstance; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; @@ -67,19 +75,58 @@ class AzureFederatedTokenProviderTest { private StsClient stsClient; private HttpClient httpClient; + private SimpleMeterRegistry meterRegistry; + private String pipelineName; + private String metricPrefix; @BeforeEach void setUp() { stsClient = mock(StsClient.class); httpClient = mock(HttpClient.class); + meterRegistry = new SimpleMeterRegistry(); + Metrics.addRegistry(meterRegistry); + // Globally-unique pipeline name: PluginMetrics writes to the static Metrics.globalRegistry + // whose meters outlive an add/removeRegistry cycle, so a name shared with another test class + // would resolve a stale removed meter. + pipelineName = "test-pipeline-" + UUID.randomUUID(); + metricPrefix = pipelineName + ".kafka."; + } + + @AfterEach + void tearDown() { + Metrics.removeRegistry(meterRegistry); + meterRegistry.clear(); + meterRegistry.close(); } private AzureFederatedTokenProvider providerWith(final StsClient sts, final HttpClient http) { + return providerWith(sts, http, null); + } + + private AzureFederatedTokenProvider providerWith(final StsClient sts, final HttpClient http, + final KafkaSourceAuthMetrics authMetrics) { final AwsCredentialsSupplier supplier = mock(AwsCredentialsSupplier.class); when(supplier.getProvider(any(AwsCredentialsOptions.class))) .thenReturn(mock(AwsCredentialsProvider.class)); return new AzureFederatedTokenProvider(REGION, STS_ROLE_ARN, TOKEN_ENDPOINT, CLIENT_ID, SCOPE, - Collections.emptyMap(), supplier, provider -> sts, http); + Collections.emptyMap(), supplier, provider -> sts, http, authMetrics); + } + + private KafkaSourceAuthMetrics newAuthMetrics() { + return new KafkaSourceAuthMetrics(PluginMetrics.fromNames("kafka", pipelineName)); + } + + private double failureCount(final String cause) { + final Counter counter = meterRegistry.find(metricPrefix + KafkaSourceAuthMetrics.TOKEN_REFRESH_FAILURES) + .tag(KafkaSourceAuthMetrics.ERROR_TYPE_TAG, cause).counter(); + return counter == null ? 0.0 : counter.count(); + } + + // Dynamic token body (random access_token + expiry) so tests prove behavior rather than pinning + // to constant literals, per the project's dynamic-test-values convention. + private String successBody() { + return "{\"access_token\":\"" + UUID.randomUUID() + + "\",\"expires_in\":" + ThreadLocalRandom.current().nextInt(60, 86_400) + "}"; } private void stubWebIdentityTokenSuccess() { @@ -126,26 +173,29 @@ void getToken_onNon200_throwsWithHttpStatusAndAadstsCode() throws IOException, I } @Test - void getToken_whenGetWebIdentityToken403_throwsAccessDeniedMessage() { + void getToken_whenStsRequestFails_throwsGenericStsMessageWithStatusAndErrorCode() { when(stsClient.getWebIdentityToken(any(java.util.function.Consumer.class))) - .thenThrow(StsException.builder().statusCode(403).message("AccessDenied").build()); + .thenThrow(StsException.builder().statusCode(403) + .awsErrorDetails(AwsErrorDetails.builder().errorCode("AccessDenied").build()) + .message("AccessDenied").build()); final RuntimeException e = assertThrows(RuntimeException.class, () -> providerWith(stsClient, httpClient).getToken()); - assertThat(e.getMessage(), containsString(STS_ROLE_ARN)); - assertThat(e.getMessage(), containsString("sts:GetWebIdentityToken")); + assertThat(e.getMessage(), containsString("STS request failed during azure_federated authentication")); + assertThat(e.getMessage(), containsString("403")); + assertThat(e.getMessage(), containsString("AccessDenied")); } @Test - void getToken_whenOutboundFederationDisabled_throwsEnableFederationMessage() { + void getToken_whenOutboundFederationDisabled_throwsFederationDisabledMessage() { when(stsClient.getWebIdentityToken(any(java.util.function.Consumer.class))) .thenThrow(OutboundWebIdentityFederationDisabledException.builder().message("disabled").build()); final RuntimeException e = assertThrows(RuntimeException.class, () -> providerWith(stsClient, httpClient).getToken()); - assertThat(e.getMessage(), containsString("Outbound")); + assertThat(e.getMessage(), containsString("outbound web identity federation feature is not enabled")); } @Test @@ -234,7 +284,7 @@ void getToken_whileCached_returnsDistinctInstancesWithSameValueAndLifetime() thr @Test void getToken_resolvesCredentialsThroughSupplier_withRoleAndRegion() throws Exception { stubWebIdentityTokenSuccess(); - final HttpResponse response = httpResponse(200, "{\"access_token\":\"t\",\"expires_in\":3599}"); + final HttpResponse response = httpResponse(200, successBody()); when(httpClient.send(any(HttpRequest.class), any(HttpResponse.BodyHandler.class))) .thenReturn(response); final AwsCredentialsSupplier supplier = mock(AwsCredentialsSupplier.class); @@ -254,7 +304,7 @@ void getToken_resolvesCredentialsThroughSupplier_withRoleAndRegion() throws Exce @Test void getToken_passesStsHeaderOverridesToSupplier() throws Exception { stubWebIdentityTokenSuccess(); - final HttpResponse response = httpResponse(200, "{\"access_token\":\"t\",\"expires_in\":3599}"); + final HttpResponse response = httpResponse(200, successBody()); when(httpClient.send(any(HttpRequest.class), any(HttpResponse.BodyHandler.class))) .thenReturn(response); final AwsCredentialsSupplier supplier = mock(AwsCredentialsSupplier.class); @@ -276,7 +326,9 @@ void getToken_passesStsHeaderOverridesToSupplier() throws Exception { @Test void getToken_withNullRole_stillResolvesThroughSupplierAndMints() throws Exception { stubWebIdentityTokenSuccess(); - final HttpResponse response = httpResponse(200, "{\"access_token\":\"t\",\"expires_in\":3599}"); + final String accessToken = UUID.randomUUID().toString(); + final HttpResponse response = + httpResponse(200, "{\"access_token\":\"" + accessToken + "\",\"expires_in\":3599}"); when(httpClient.send(any(HttpRequest.class), any(HttpResponse.BodyHandler.class))) .thenReturn(response); final AwsCredentialsSupplier supplier = mock(AwsCredentialsSupplier.class); @@ -289,13 +341,15 @@ void getToken_withNullRole_stillResolvesThroughSupplierAndMints() throws Excepti provider -> stsClient, httpClient) .getToken(); - assertThat(token.value(), equalTo("t")); + assertThat(token.value(), equalTo(accessToken)); } @Test - void getToken_whenGetWebIdentityToken403WithNullRole_usesStableIdentifierNotNullLiteral() { + void getToken_whenStsRequestFailsWithNullRole_doesNotLeakNullLiteral() { when(stsClient.getWebIdentityToken(any(java.util.function.Consumer.class))) - .thenThrow(StsException.builder().statusCode(403).message("AccessDenied").build()); + .thenThrow(StsException.builder().statusCode(403) + .awsErrorDetails(AwsErrorDetails.builder().errorCode("AccessDenied").build()) + .message("AccessDenied").build()); final AwsCredentialsSupplier supplier = mock(AwsCredentialsSupplier.class); when(supplier.getProvider(any(AwsCredentialsOptions.class))) .thenReturn(mock(AwsCredentialsProvider.class)); @@ -305,11 +359,11 @@ void getToken_whenGetWebIdentityToken403WithNullRole_usesStableIdentifierNotNull Collections.emptyMap(), supplier, provider -> stsClient, httpClient).getToken()); assertThat(e.getMessage(), not(containsString("null"))); - assertThat(e.getMessage(), containsString("sts:GetWebIdentityToken")); + assertThat(e.getMessage(), containsString("STS request failed during azure_federated authentication")); } @Test - void getToken_whenCredentialResolutionThrowsNonSdkException_surfacesActionableMessage() { + void getToken_whenUnrecognizedExceptionThrown_surfacesGenericMessage() { // simulate a non-SDK failure surfacing at the credential-materialization point when(stsClient.getWebIdentityToken(any(java.util.function.Consumer.class))) .thenThrow(new RuntimeException("com.amazonaws delegation assume failed")); @@ -321,7 +375,7 @@ void getToken_whenCredentialResolutionThrowsNonSdkException_surfacesActionableMe new AzureFederatedTokenProvider(REGION, STS_ROLE_ARN, TOKEN_ENDPOINT, CLIENT_ID, SCOPE, Collections.emptyMap(), supplier, provider -> stsClient, httpClient).getToken()); - assertThat(e.getMessage(), containsString("AWS credential resolution failed")); + assertThat(e.getMessage(), containsString("Unexpected failure during azure_federated authentication")); } @Test @@ -356,4 +410,121 @@ void close_whenStsClientNeverCreated_isNoOp() { verify(stsClient, never()).close(); } + + @Test + void getToken_onSuccess_incrementsRefreshCountAndSetsExpiry() throws IOException, InterruptedException { + stubWebIdentityTokenSuccess(); + final HttpResponse response = httpResponse(200, successBody()); + when(httpClient.send(any(HttpRequest.class), any(HttpResponse.BodyHandler.class))).thenReturn(response); + final KafkaSourceAuthMetrics authMetrics = newAuthMetrics(); + + providerWith(stsClient, httpClient, authMetrics).getToken(); + + final Counter refreshCount = + meterRegistry.find(metricPrefix + KafkaSourceAuthMetrics.TOKEN_REFRESH_SUCCESS).counter(); + assertThat(refreshCount, notNullValue()); + assertThat(refreshCount.count(), equalTo(1.0)); + final io.micrometer.core.instrument.Gauge expiry = + meterRegistry.find(metricPrefix + KafkaSourceAuthMetrics.TIME_TO_TOKEN_REFRESH).gauge(); + assertThat(expiry.value(), greaterThan(0.0)); + } + + @Test + void getToken_whileCached_incrementsRefreshCountOnlyOnActualMint() throws IOException, InterruptedException { + stubWebIdentityTokenSuccess(); + final HttpResponse response = httpResponse(200, successBody()); + when(httpClient.send(any(HttpRequest.class), any(HttpResponse.BodyHandler.class))).thenReturn(response); + final KafkaSourceAuthMetrics authMetrics = newAuthMetrics(); + final AzureFederatedTokenProvider provider = providerWith(stsClient, httpClient, authMetrics); + + provider.getToken(); + provider.getToken(); // served from cache, must NOT be counted as a refresh + + assertThat(meterRegistry.find(metricPrefix + KafkaSourceAuthMetrics.TOKEN_REFRESH_SUCCESS).counter().count(), + equalTo(1.0)); + } + + @Test + void getToken_on403_incrementsFailuresTaggedAccessDenied() { + when(stsClient.getWebIdentityToken(any(java.util.function.Consumer.class))) + .thenThrow(StsException.builder().statusCode(403) + .awsErrorDetails(AwsErrorDetails.builder().errorCode("AccessDenied").build()) + .message("AccessDenied").build()); + final KafkaSourceAuthMetrics authMetrics = newAuthMetrics(); + + assertThrows(RuntimeException.class, () -> providerWith(stsClient, httpClient, authMetrics).getToken()); + + assertThat(failureCount(KafkaSourceAuthMetrics.CAUSE_AWS_STS_ACCESS_DENIED), equalTo(1.0)); + } + + @Test + void getToken_onNonAccessDeniedStsError_incrementsFailuresTaggedStsError() { + when(stsClient.getWebIdentityToken(any(java.util.function.Consumer.class))) + .thenThrow(StsException.builder().statusCode(500) + .awsErrorDetails(AwsErrorDetails.builder().errorCode("InternalServerError").build()) + .message("STS internal error").build()); + final KafkaSourceAuthMetrics authMetrics = newAuthMetrics(); + + assertThrows(RuntimeException.class, () -> providerWith(stsClient, httpClient, authMetrics).getToken()); + + assertThat(failureCount(KafkaSourceAuthMetrics.CAUSE_AWS_STS_ERROR), equalTo(1.0)); + } + + @Test + void getToken_whenFederationDisabled_incrementsFailuresTaggedFederationDisabled() { + when(stsClient.getWebIdentityToken(any(java.util.function.Consumer.class))) + .thenThrow(OutboundWebIdentityFederationDisabledException.builder().message("disabled").build()); + final KafkaSourceAuthMetrics authMetrics = newAuthMetrics(); + + assertThrows(RuntimeException.class, () -> providerWith(stsClient, httpClient, authMetrics).getToken()); + + assertThat(failureCount(KafkaSourceAuthMetrics.CAUSE_AWS_OUTBOUND_FEDERATION_DISABLED), equalTo(1.0)); + } + + @Test + void getToken_onNon200_incrementsFailuresTaggedAadsts() throws IOException, InterruptedException { + stubWebIdentityTokenSuccess(); + final HttpResponse response = httpResponse(401, + "{\"error\":\"invalid_client\",\"error_description\":\"AADSTS700016: not found\"}"); + when(httpClient.send(any(HttpRequest.class), any(HttpResponse.BodyHandler.class))).thenReturn(response); + final KafkaSourceAuthMetrics authMetrics = newAuthMetrics(); + + assertThrows(RuntimeException.class, () -> providerWith(stsClient, httpClient, authMetrics).getToken()); + + assertThat(failureCount(KafkaSourceAuthMetrics.CAUSE_AZURE_TOKEN_EXCHANGE_REJECTED), equalTo(1.0)); + } + + @Test + void getToken_onHttpIoException_incrementsFailuresTaggedNetwork() throws IOException, InterruptedException { + stubWebIdentityTokenSuccess(); + when(httpClient.send(any(HttpRequest.class), any(HttpResponse.BodyHandler.class))) + .thenThrow(new IOException("connection reset")); + final KafkaSourceAuthMetrics authMetrics = newAuthMetrics(); + + assertThrows(RuntimeException.class, () -> providerWith(stsClient, httpClient, authMetrics).getToken()); + + assertThat(failureCount(KafkaSourceAuthMetrics.CAUSE_NETWORK), equalTo(1.0)); + } + + @Test + void getToken_withNullMetrics_doesNotThrowOnSuccessOrFailure() throws IOException, InterruptedException { + // success path with null metrics + stubWebIdentityTokenSuccess(); + final HttpResponse response = httpResponse(200, successBody()); + when(httpClient.send(any(HttpRequest.class), any(HttpResponse.BodyHandler.class))).thenReturn(response); + providerWith(stsClient, httpClient, null).getToken(); + + // failure path with null metrics: the original actionable error must survive, i.e. the null + // metrics must NOT be dereferenced (an NPE from a missing null-guard also extends + // RuntimeException, so assertThrows(RuntimeException) alone would not catch that regression). + final StsClient failing = mock(StsClient.class); + when(failing.getWebIdentityToken(any(java.util.function.Consumer.class))) + .thenThrow(StsException.builder().statusCode(403) + .awsErrorDetails(AwsErrorDetails.builder().errorCode("AccessDenied").build()) + .message("AccessDenied").build()); + final RuntimeException thrown = assertThrows(RuntimeException.class, + () -> providerWith(failing, httpClient, null).getToken()); + assertThat(thrown, not(instanceOf(NullPointerException.class))); + assertThat(thrown.getMessage(), containsString("STS request failed during azure_federated authentication")); + } } diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/authenticator/KafkaSourceAuthMetricsProviderTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/authenticator/KafkaSourceAuthMetricsProviderTest.java new file mode 100644 index 0000000000..85dd2a4a12 --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/authenticator/KafkaSourceAuthMetricsProviderTest.java @@ -0,0 +1,52 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + */ + +package org.opensearch.dataprepper.plugins.kafka.authenticator; + +import org.junit.jupiter.api.Test; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.nullValue; +import static org.hamcrest.Matchers.sameInstance; +import static org.mockito.Mockito.mock; + +class KafkaSourceAuthMetricsProviderTest { + + @Test + void getInstance_returnsSameSingleton() { + assertThat(KafkaSourceAuthMetricsProvider.getInstance(), + sameInstance(KafkaSourceAuthMetricsProvider.getInstance())); + } + + @Test + void get_beforeSet_isNull() { + // fresh instance via the protected test ctor to avoid cross-test static leakage + final KafkaSourceAuthMetricsProvider provider = new KafkaSourceAuthMetricsProvider() {}; + assertThat(provider.getAuthMetrics(), nullValue()); + } + + @Test + void set_thenGet_returnsSuppliedInstance() { + final KafkaSourceAuthMetricsProvider provider = new KafkaSourceAuthMetricsProvider() {}; + final KafkaSourceAuthMetrics metrics = mock(KafkaSourceAuthMetrics.class); + provider.set(metrics); + assertThat(provider.getAuthMetrics(), sameInstance(metrics)); + } + + @Test + void set_overwritesPrevious_lastWriterWins() { + final KafkaSourceAuthMetricsProvider provider = new KafkaSourceAuthMetricsProvider() {}; + final KafkaSourceAuthMetrics first = mock(KafkaSourceAuthMetrics.class); + final KafkaSourceAuthMetrics second = mock(KafkaSourceAuthMetrics.class); + provider.set(first); + provider.set(second); + assertThat(provider.getAuthMetrics(), sameInstance(second)); + } +} diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/authenticator/KafkaSourceAuthMetricsTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/authenticator/KafkaSourceAuthMetricsTest.java new file mode 100644 index 0000000000..8ed9e74032 --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/authenticator/KafkaSourceAuthMetricsTest.java @@ -0,0 +1,135 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + */ + +package org.opensearch.dataprepper.plugins.kafka.authenticator; + +import io.micrometer.core.instrument.Counter; +import io.micrometer.core.instrument.Gauge; +import io.micrometer.core.instrument.Metrics; +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.metrics.PluginMetrics; + +import java.util.UUID; +import java.util.concurrent.ThreadLocalRandom; +import java.util.concurrent.atomic.AtomicLong; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.closeTo; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.greaterThan; +import static org.hamcrest.Matchers.notNullValue; +import static org.hamcrest.Matchers.nullValue; + +class KafkaSourceAuthMetricsTest { + + private String prefix; + private SimpleMeterRegistry meterRegistry; + private AtomicLong now; + private KafkaSourceAuthMetrics authMetrics; + + @BeforeEach + void setUp() { + // PluginMetrics writes to the static Metrics.globalRegistry, whose meters outlive an + // add/removeRegistry cycle. A globally-unique pipeline name (not just per-class-unique) avoids + // cross-test meter-id collisions that otherwise make find() resolve a stale, removed meter. + final String pipelineName = "test-pipeline-" + UUID.randomUUID(); + prefix = pipelineName + ".kafka."; + meterRegistry = new SimpleMeterRegistry(); + Metrics.addRegistry(meterRegistry); + // Controllable clock so the pull-model gauge is exercised deterministically (advance time + // between refresh and read without sleeping). + now = new AtomicLong(1_000_000_000_000L); + authMetrics = new KafkaSourceAuthMetrics(PluginMetrics.fromNames("kafka", pipelineName), now::get); + } + + @AfterEach + void tearDown() { + Metrics.removeRegistry(meterRegistry); + meterRegistry.clear(); + meterRegistry.close(); + } + + private Gauge expiryGauge() { + return meterRegistry.find(prefix + KafkaSourceAuthMetrics.TIME_TO_TOKEN_REFRESH).gauge(); + } + + @Test + void recordRefresh_incrementsCountAndSetsExpiryHeadroom() { + final long headroomSeconds = ThreadLocalRandom.current().nextInt(60, 7200); + authMetrics.recordRefresh(now.get() + headroomSeconds * 1000L); + + final Counter counter = meterRegistry.find(prefix + KafkaSourceAuthMetrics.TOKEN_REFRESH_SUCCESS).counter(); + assertThat(counter, notNullValue()); + assertThat(counter.count(), equalTo(1.0)); + + assertThat(expiryGauge(), notNullValue()); + assertThat(expiryGauge().value(), closeTo(headroomSeconds, 0.001)); + } + + @Test + void timeToTokenRefresh_beforeAnyToken_isZero() { + assertThat(expiryGauge(), notNullValue()); + assertThat(expiryGauge().value(), equalTo(0.0)); + } + + @Test + void timeToTokenRefresh_isPullModel_countsDownAsClockAdvancesForSameToken() { + // Guards against a snapshot-at-refresh regression: for a SINGLE refresh, the gauge value must + // strictly decrease as the clock advances (it is computed at read time, not frozen at refresh). + authMetrics.recordRefresh(now.get() + 3_600_000L); + final double atMint = expiryGauge().value(); + + now.addAndGet(600_000L); // 10 minutes pass, no new refresh + final double tenMinutesLater = expiryGauge().value(); + + assertThat(atMint, closeTo(3600.0, 0.001)); + assertThat(atMint, greaterThan(tenMinutesLater)); + assertThat(tenMinutesLater, closeTo(3000.0, 0.001)); + } + + @Test + void timeToTokenRefresh_pastDeadline_flooredAtZero() { + authMetrics.recordRefresh(now.get() - 10_000L); + assertThat(expiryGauge().value(), equalTo(0.0)); + } + + @Test + void recordFailure_incrementsCounterTaggedByCause() { + authMetrics.recordFailure(KafkaSourceAuthMetrics.CAUSE_AWS_STS_ACCESS_DENIED); + authMetrics.recordFailure(KafkaSourceAuthMetrics.CAUSE_AWS_STS_ACCESS_DENIED); + authMetrics.recordFailure(KafkaSourceAuthMetrics.CAUSE_AZURE_TOKEN_EXCHANGE_REJECTED); + + final Counter accessDenied = meterRegistry.find(prefix + KafkaSourceAuthMetrics.TOKEN_REFRESH_FAILURES) + .tag(KafkaSourceAuthMetrics.ERROR_TYPE_TAG, KafkaSourceAuthMetrics.CAUSE_AWS_STS_ACCESS_DENIED).counter(); + final Counter aadsts = meterRegistry.find(prefix + KafkaSourceAuthMetrics.TOKEN_REFRESH_FAILURES) + .tag(KafkaSourceAuthMetrics.ERROR_TYPE_TAG, KafkaSourceAuthMetrics.CAUSE_AZURE_TOKEN_EXCHANGE_REJECTED).counter(); + + assertThat(accessDenied, notNullValue()); + assertThat(accessDenied.count(), equalTo(2.0)); + assertThat(aadsts, notNullValue()); + assertThat(aadsts.count(), equalTo(1.0)); + } + + @Test + void recordFailure_distinctCauses_produceDistinctCounters() { + authMetrics.recordFailure(KafkaSourceAuthMetrics.CAUSE_NETWORK); + final Counter network = meterRegistry.find(prefix + KafkaSourceAuthMetrics.TOKEN_REFRESH_FAILURES) + .tag(KafkaSourceAuthMetrics.ERROR_TYPE_TAG, KafkaSourceAuthMetrics.CAUSE_NETWORK).counter(); + final Counter other = meterRegistry.find(prefix + KafkaSourceAuthMetrics.TOKEN_REFRESH_FAILURES) + .tag(KafkaSourceAuthMetrics.ERROR_TYPE_TAG, KafkaSourceAuthMetrics.CAUSE_OTHER).counter(); + assertThat(network, notNullValue()); + assertThat(network.count(), equalTo(1.0)); + // a cause never recorded has no counter registered + assertThat(other, nullValue()); + } +} diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumerTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumerTest.java index e1f5b020de..e4c2e54d89 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumerTest.java +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumerTest.java @@ -86,6 +86,7 @@ import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.times; +import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.spy; import static org.mockito.Mockito.doNothing; @@ -176,6 +177,8 @@ public void setUp() throws JsonProcessingException { when(topicMetrics.getNumberOfDeserializationErrors()).thenReturn(counter); when(topicMetrics.getNumberOfInvalidTimeStamps()).thenReturn(counter); when(topicMetrics.getNumberOfPollAuthErrors()).thenReturn(counter); + lenient().when(topicMetrics.getNumberOfRebalances()).thenReturn(counter); + lenient().when(topicMetrics.getNumberOfPartitionsRevoked()).thenReturn(counter); when(topicConfig.getThreadWaitingTime()).thenReturn(Duration.ofSeconds(1)); when(topicConfig.getSerdeFormat()).thenReturn(MessageFormat.PLAINTEXT); when(topicConfig.getAutoCommit()).thenReturn(false); @@ -348,6 +351,43 @@ public void testPlainTextConsumeRecords() throws InterruptedException { assertNotNull(event.getEventHandle().getExternalOriginationTime()); verify(topicMetrics).recordTimeBetweenPolls(); + verify(topicMetrics).recordProcessingLatency(anyLong()); + } + + @Test + public void onPartitionsAssigned_incrementsRebalanceCount() { + final Counter rebalanceCounter = mock(Counter.class); + when(topicMetrics.getNumberOfRebalances()).thenReturn(rebalanceCounter); + final String topic = topicConfig.getName(); + consumer = createObjectUnderTest("plaintext", false); + + consumer.onPartitionsAssigned(List.of(new TopicPartition(topic, testPartition))); + + verify(rebalanceCounter).increment(); + } + + @Test + public void onPartitionsRevoked_incrementsRevokedCountAndClearsPartitionLag() { + final Counter revokedCounter = mock(Counter.class); + when(topicMetrics.getNumberOfPartitionsRevoked()).thenReturn(revokedCounter); + final String topic = topicConfig.getName(); + final TopicPartition topicPartition = new TopicPartition(topic, testPartition); + consumer = createObjectUnderTest("plaintext", false); + consumer.onPartitionsAssigned(List.of(topicPartition)); + + consumer.onPartitionsRevoked(List.of(topicPartition)); + + verify(revokedCounter).increment(1.0); + verify(topicMetrics).clearPartitionLag(testPartition); + } + + @Test + public void closeConsumer_deregistersFromMetrics() { + consumer = createObjectUnderTest("plaintext", false); + + consumer.closeConsumer(); + + verify(topicMetrics).deregister(kafkaConsumer); } @Test diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicConsumerMetricsLagGaugeTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicConsumerMetricsLagGaugeTest.java new file mode 100644 index 0000000000..2c26bd8bfe --- /dev/null +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicConsumerMetricsLagGaugeTest.java @@ -0,0 +1,164 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + */ + +package org.opensearch.dataprepper.plugins.kafka.util; + +import io.micrometer.core.instrument.Metrics; +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; +import org.apache.commons.lang3.RandomStringUtils; +import org.apache.kafka.clients.consumer.KafkaConsumer; +import org.apache.kafka.common.Metric; +import org.apache.kafka.common.MetricName; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.metrics.PluginMetrics; + +import java.util.HashMap; +import java.util.Map; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.closeTo; +import static org.hamcrest.Matchers.notNullValue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +/** + * Verifies that {@link KafkaTopicConsumerMetrics} emits an always-on {@code recordsLagPerPartition} + * gauge tagged by partition, reflecting the consumer's per-partition {@code records-lag-max}. + */ +public class KafkaTopicConsumerMetricsLagGaugeTest { + + private static final class KafkaTestMetric implements Metric { + private final Object value; + private final MetricName name; + + private KafkaTestMetric(final Object value, final MetricName name) { + this.value = value; + this.name = name; + } + + @Override + public MetricName metricName() { + return name; + } + + @Override + public Object metricValue() { + return value; + } + } + + private SimpleMeterRegistry meterRegistry; + private String topicName; + + @BeforeEach + void setUp() { + topicName = RandomStringUtils.randomAlphabetic(8); + meterRegistry = new SimpleMeterRegistry(); + Metrics.addRegistry(meterRegistry); + } + + @AfterEach + void tearDown() { + Metrics.removeRegistry(meterRegistry); + meterRegistry.clear(); + meterRegistry.close(); + } + + private KafkaTopicConsumerMetrics newTopicMetrics() { + return new KafkaTopicConsumerMetrics(topicName, PluginMetrics.fromNames("kafka", "test-pipeline"), true); + } + + private MetricName lagMetricName(final String partition) { + final Map tags = new HashMap<>(); + tags.put("topic", topicName); + tags.put("partition", partition); + return new MetricName("records-lag-max", "consumer-fetch-manager-metrics", "per-partition lag", tags); + } + + private void putLag(final Map metrics, final String partition, final double lag) { + final MetricName name = lagMetricName(partition); + metrics.put(name, new KafkaTestMetric(lag, name)); + } + + @SuppressWarnings({"unchecked", "rawtypes"}) + private KafkaConsumer consumerWithPartitionLag(final String partition, final double lag) { + final KafkaConsumer consumer = mock(KafkaConsumer.class); + final Map metrics = new HashMap<>(); + putLag(metrics, partition, lag); + when(consumer.metrics()).thenReturn(metrics); + return consumer; + } + + private double lagGaugeValue(final String partition) { + final String gaugeName = "test-pipeline.kafka.topic." + topicName + ".recordsLagPerPartition"; + return meterRegistry.find(gaugeName).tag("partition", partition).gauge().value(); + } + + @Test + public void update_registersPerPartitionLagGaugeTaggedByPartition() { + final KafkaTopicConsumerMetrics topicMetrics = newTopicMetrics(); + final KafkaConsumer consumer = consumerWithPartitionLag("3", 42.0); + topicMetrics.register(consumer); + + topicMetrics.update(consumer); + + final String gaugeName = "test-pipeline.kafka.topic." + topicName + ".recordsLagPerPartition"; + assertThat("recordsLagPerPartition gauge for partition=3 should be registered", + meterRegistry.find(gaugeName).tag("partition", "3").gauge(), notNullValue()); + assertThat(lagGaugeValue("3"), closeTo(42.0, 0.001)); + } + + @Test + public void update_multiplePartitions_registersDistinctTaggedGauges() { + final KafkaTopicConsumerMetrics topicMetrics = newTopicMetrics(); + final KafkaConsumer consumer = mock(KafkaConsumer.class); + final Map metrics = new HashMap<>(); + putLag(metrics, "3", 10.0); + putLag(metrics, "7", 99.0); + when(consumer.metrics()).thenReturn(metrics); + topicMetrics.register(consumer); + + topicMetrics.update(consumer); + + assertThat(lagGaugeValue("3"), closeTo(10.0, 0.001)); + assertThat(lagGaugeValue("7"), closeTo(99.0, 0.001)); + } + + @Test + public void update_reportsLatestLagNotHighWaterMark() { + final KafkaTopicConsumerMetrics topicMetrics = newTopicMetrics(); + final KafkaConsumer consumer = mock(KafkaConsumer.class); + final Map metrics = new HashMap<>(); + when(consumer.metrics()).thenReturn(metrics); + topicMetrics.register(consumer); + + putLag(metrics, "3", 500.0); + topicMetrics.update(consumer); + putLag(metrics, "3", 20.0); + topicMetrics.update(consumer); + + // A live gauge must fall to the latest lag, not stay at the earlier high-water mark. + assertThat(lagGaugeValue("3"), closeTo(20.0, 0.001)); + } + + @Test + public void clearPartitionLag_zeroesTheGaugeForARevokedPartition() { + final KafkaTopicConsumerMetrics topicMetrics = newTopicMetrics(); + final KafkaConsumer consumer = consumerWithPartitionLag("3", 42.0); + topicMetrics.register(consumer); + topicMetrics.update(consumer); + + topicMetrics.clearPartitionLag(3); + + assertThat(lagGaugeValue("3"), closeTo(0.0, 0.001)); + } +} diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicConsumerMetricsTests.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicConsumerMetricsTests.java index 9bf2f3e130..1715de2588 100644 --- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicConsumerMetricsTests.java +++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/util/KafkaTopicConsumerMetricsTests.java @@ -252,6 +252,11 @@ public void KafkaTopicMetricTest_checkMetricUpdates(int numConsumers) { } else if (k.contains("numberOfNonConsumers")) { int expectedValue = numConsumers/2; assertThat(result, equalTo((double)expectedValue)); + } else if (k.contains("numberOfAssignedPartitions") || k.contains("numberOfActiveReaders") + || k.contains("numberOfConfiguredWorkers")) { + // New always-on capacity gauges register real doubles; not part of this + // rate-oriented assertion set, so just confirm they produce a finite value. + assertThat(Double.isNaN(result), equalTo(false)); } else { assertThat(result, equalTo(k+": Unknown Metric")); } @@ -375,6 +380,36 @@ public void update_withNaNMetricValue_andNopriorUpdate_preservesInitialZero() { assertEquals(0.0, lagResult, 0.01d); } + @Test + public void scalingGauges_reportConfiguredWorkersAndAssignmentCounts() { + final int configuredWorkers = 4; + topicMetrics = new KafkaTopicConsumerMetrics(topicName, pluginMetrics, true, configuredWorkers); + + // Three consumers: two each owning a partition (active readers), one idle (non-consumer). + final double[] assignedPerConsumer = {1.0, 1.0, 0.0}; + for (final double assigned : assignedPerConsumer) { + final KafkaConsumer kafkaConsumer = mock(KafkaConsumer.class); + topicMetrics.register(kafkaConsumer); + final Map metrics = new HashMap<>(); + when(kafkaConsumer.metrics()).thenReturn(metrics); + final KafkaTestMetric metric = getMetric("assigned-partitions", assigned, new HashMap<>()); + metrics.put(metric.metricName(), metric); + topicMetrics.update(kafkaConsumer); + } + + final Map> mv = topicMetrics.getMetricValues(); + // numberOfConfiguredWorkers only registers when configuredWorkers > 0 (the 4-arg constructor + // that KafkaSource now uses); assert it reports the configured count. + assertEquals(configuredWorkers, + pluginMetricsMap.get("topic." + topicName + ".numberOfConfiguredWorkers").applyAsDouble(mv), 0.01d); + assertEquals(2.0, + pluginMetricsMap.get("topic." + topicName + ".numberOfAssignedPartitions").applyAsDouble(mv), 0.01d); + assertEquals(2.0, + pluginMetricsMap.get("topic." + topicName + ".numberOfActiveReaders").applyAsDouble(mv), 0.01d); + assertEquals(1.0, + pluginMetricsMap.get("topic." + topicName + ".numberOfNonConsumers").applyAsDouble(mv), 0.01d); + } + @Test void recordTimeBetweenPolls_records_metric_correctly() throws InterruptedException { topicMetrics = createObjectUnderTest();