opensearch-project · graytaylor0 · Aug 20, 2026 · Aug 17, 2026
@@ -56,9 +56,14 @@ public void configure(final Map<String, ?> configs, final String saslMechanism,
         final Map<String, String> 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<String, String> decodeStsHeaderOverrides(final String encoded) {

@@ -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<StsClient> 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<String, String> 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<String, String> 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<AwsCredentialsProvider, StsClient> 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<String, String> stsHeaderOverrides,
+                                final AwsCredentialsSupplier awsCredentialsSupplier,
+                                final Function<AwsCredentialsProvider, StsClient> 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<String> 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;

@@ -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();
+    }
+}
@@ -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()}.
+ *
+ * <p>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;
+    }
+}
@@ -235,6 +235,8 @@ <T> 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<TopicPartition, CommitOffsetRange> offsets = new HashMap<>();
                 AcknowledgementSet acknowledgementSet = null;
                 if (acknowledgementsEnabled) {
@@ -250,6 +252,7 @@ <T> void consumeRecords() throws Exception {
                     acknowledgementSet.complete();
                     numberOfAcksPending.incrementAndGet();
                 }
+                topicMetrics.recordProcessingLatency(System.currentTimeMillis() - processingStartMillis);
             }
         } catch (AuthenticationException e) {
             authFailureAttempts++;
@@ -620,6 +623,7 @@ private <T> void iterateRecordPartitions(ConsumerRecords<String, T> records, fin
     }
 
     public void closeConsumer(){
+        topicMetrics.deregister(consumer);
         consumer.close();
     }
 
@@ -629,6 +633,8 @@ public void shutdownConsumer(){
 
     @Override
     public void onPartitionsAssigned(Collection<TopicPartition> 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<TopicPartition> partitions) {
 
     @Override
     public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
+        topicMetrics.getNumberOfPartitionsRevoked().increment(partitions.size());
         synchronized(this) {
             commitOffsets(true);
             for (TopicPartition topicPartition : partitions) {
@@ -659,6 +666,7 @@ public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
                 LOG.info("Revoked partition {}", topicPartition);
                 ownedPartitionsEpoch.remove(topicPartition);
                 partitionCommitTrackerMap.remove(topicPartition.partition());
+                topicMetrics.clearPartitionLag(topicPartition.partition());
             }
             if (paused) {
                 consumer.pause(consumer.assignment());