parthchandra commented on code in PR #6025:
URL: https://github.com/apache/datafusion-comet/pull/6025#discussion_r4065106262
##########
native/core/src/parquet/objectstore/s3.rs:
##########
@@ -100,9 +101,27 @@ pub fn create_store(
builder.with_credentials(Arc::new(bridge))
}
None => {
- match get_runtime().block_on(build_credential_provider(configs,
bucket, min_ttl))? {
+ // IRSA take-over. When no explicit `aws.credentials.provider` is
configured, the
+ // default AWS chain places IMDS/instance-role after web-identity,
so a throttled
+ // AssumeRoleWithWebIdentity silently downgrades to the node role
-> hard S3 403. On
+ // EKS/IRSA use the Comet web-identity provider instead (retries
the throttle, no
+ // node-role fallback, shared jittered cache). Any explicit
provider config is
+ // respected -- it takes the normal `build_credential_provider`
path below.
+ let explicit_provider = get_config_trimmed(configs, bucket,
"aws.credentials.provider")
+ .is_some_and(|s| !s.is_empty());
+ let web_identity = take_over_if_irsa(explicit_provider, |key| {
+ get_config_trimmed(configs, bucket, key).map(|s| s.to_string())
+ });
Review Comment:
I've dropped my earlier "profiles out of scope" position.
`take_over_if_irsa`
now stands aside when a profile is configured, in addition to the
environment-key case:
- `AWS_PROFILE` set, or
- a shared credentials file present (`AWS_SHARED_CREDENTIALS_FILE`, else
`~/.aws/credentials`).
The default chain ranks a working profile ahead of web-identity, so the
take-over defers to
it and the previous principal keeps being used. The check is deliberately
conservative: a
present-but-empty file just makes us fall back to the default chain, which
is the
pre-existing behavior, so it is never worse than before.
Test added: `configured_profile_keeps_precedence_over_irsa` — with
`AWS_PROFILE` set the
take-over returns `None`; on a clean IRSA baseline it takes over.
##########
native/core/src/cloud/s3/web_identity.rs:
##########
@@ -0,0 +1,918 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! IRSA (EKS "IAM Roles for Service Accounts") web-identity credential
provider for the native S3
+//! paths.
+//!
+//! Why this exists: on EKS with IRSA the native reader assumes the app role
by calling STS
+//! `AssumeRoleWithWebIdentity`. Under a concurrent burst (many executors x
many cores starting
+//! together) STS throttles that call. opendal's default reqsign chain (used
by the Iceberg path
+//! when no Comet provider class is set) does NOT retry the throttle and
silently downgrades to the
+//! EC2/EKS node instance role, which lacks bucket access -> every read then
fails with a hard S3
+//! 403. See docs/source/contributor-guide/s3-credential-provider-design.md.
+//!
+//! This provider fixes all three parts of that failure:
+//! 1. Retry on throttle. It uses the AWS SDK
`WebIdentityTokenCredentialsProvider`, whose STS
+//! client retries throttling with exponential backoff + jitter.
`max_attempts` is
+//! configurable (default higher than the SDK's default of 3).
+//! 2. No silent downgrade. The provider is web-identity ONLY -- there is no
IMDS/instance-role
+//! fallback -- so a transient throttle surfaces as a retryable error
instead of a
+//! wrong-identity credential.
+//! 3. Shared, jittered cache. One assumed-role credential is cached per
process, keyed by
+//! identity (role_arn, token_file, region) and the resolved
retry/refresh settings, and shared
+//! across all reader threads and scans that resolve to the same key.
Refresh fires ahead of
+//! expiry by `min_ttl` plus a per-process random jitter so cluster-wide
refreshes do not
+//! synchronize into another burst; a failed refresh is briefly
remembered so a throttled burst
+//! costs one STS call rather than one per reader.
+//!
+//! The same struct is exposed as both `object_store::CredentialProvider` (raw
Parquet path) and
+//! reqsign's `ProvideCredential` (Iceberg via opendal /
`CustomAwsCredentialLoader`), mirroring
+//! `credential_bridge::CometS3CredentialBridge`.
+
+use std::collections::HashMap;
+use std::sync::{Arc, OnceLock, RwLock};
+use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
+
+use async_trait::async_trait;
+use aws_config::provider_config::ProviderConfig;
+use aws_config::retry::RetryConfig;
+use aws_config::web_identity_token::WebIdentityTokenCredentialsProvider;
+use aws_credential_types::provider::ProvideCredentials;
+use aws_credential_types::Credentials;
+use iceberg_storage_opendal::AwsCredential as IcebergAwsCredential;
+use object_store::aws::AwsCredential;
+use object_store::CredentialProvider;
+use rand::RngExt;
+use reqsign_core::time::Timestamp;
+use reqsign_core::{
+ Context, Error as ReqsignError, ErrorKind as ReqsignErrorKind,
+ ProvideCredential as IcebergProvideCredential,
+};
+
+use crate::cloud::s3::credential_bridge::DEFAULT_EXPIRY_WHEN_UNKNOWN;
+
+/// EKS-projected env vars that signal IRSA is in effect. Both must be present.
+const ENV_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE";
+const ENV_ROLE_ARN: &str = "AWS_ROLE_ARN";
+
+/// Config keys read from the Iceberg catalog property bag. A
non-`s3.`/`client.` prefix keeps them
+/// from being forwarded into opendal's S3 config (see
`iceberg_common::STORAGE_PROPERTY_PREFIXES`).
+const KEY_ENABLED: &str = "comet.s3.credentials.webIdentity.enabled";
+const KEY_MAX_ATTEMPTS: &str = "comet.s3.credentials.webIdentity.maxAttempts";
+const KEY_MIN_TTL_SECS: &str =
"comet.s3.credentials.webIdentity.minTtlSeconds";
+const KEY_JITTER_SECS: &str =
"comet.s3.credentials.webIdentity.refreshJitterSeconds";
+
+const DEFAULT_ENABLED: bool = true;
+const DEFAULT_MAX_ATTEMPTS: u32 = 5;
+const DEFAULT_MIN_TTL_SECS: u64 = 300;
+const DEFAULT_JITTER_SECS: u64 = 60;
+
+/// After a refresh exhausts its STS retries and fails, waiters within this
window get the failure
+/// without each firing their own assume-role call. Bounds STS pressure during
a sustained throttle
+/// (one call per entry per window instead of one per reader) while still
letting the credential
+/// recover shortly after. Kept short: the SDK has already spent its retry
budget by the time we
+/// record a failure.
+const FAILURE_COOLDOWN: Duration = Duration::from_secs(1);
+
+/// Detected IRSA identity plus the resolved tuning knobs. Cheap to clone; the
expensive AWS SDK
+/// provider lives in the process-wide `SharedEntry` keyed by `entry_key`.
+#[derive(Clone, Debug)]
+pub struct WebIdentityConfig {
+ role_arn: String,
+ token_file: String,
+ /// From `AWS_REGION` / `AWS_DEFAULT_REGION`; only part of the cache key.
The actual region
+ /// resolution is done by the AWS SDK's `with_default_region`.
+ region: Option<String>,
+ max_attempts: u32,
+ min_ttl: Duration,
+ max_jitter: Duration,
+}
+
+impl WebIdentityConfig {
+ /// Returns a config only when IRSA is in effect (both env vars present)
and the feature is
+ /// enabled. `resolve` looks up a bare setting key (e.g.
`KEY_MAX_ATTEMPTS`) in whichever config
+ /// bag the caller owns -- the Iceberg catalog bag or the Parquet
`fs.s3a.*` bag -- so the two
+ /// scan paths share one detection routine without sharing a config-key
scheme. Returns `None`
+ /// when IRSA is not detected or the feature is disabled.
+ pub fn detect_with<F>(resolve: F) -> Option<Self>
+ where
+ F: Fn(&str) -> Option<String>,
+ {
+ let token_file = non_empty_env(ENV_TOKEN_FILE)?;
+ let role_arn = non_empty_env(ENV_ROLE_ARN)?;
+ if !parse_setting(resolve(KEY_ENABLED), DEFAULT_ENABLED) {
+ return None;
+ }
+ Some(Self {
+ role_arn,
+ token_file,
+ region: non_empty_env("AWS_REGION").or_else(||
non_empty_env("AWS_DEFAULT_REGION")),
+ max_attempts: parse_u32(resolve(KEY_MAX_ATTEMPTS),
DEFAULT_MAX_ATTEMPTS),
+ min_ttl: Duration::from_secs(parse_setting(
+ resolve(KEY_MIN_TTL_SECS),
+ DEFAULT_MIN_TTL_SECS,
+ )),
+ max_jitter: Duration::from_secs(parse_setting(
+ resolve(KEY_JITTER_SECS),
+ DEFAULT_JITTER_SECS,
+ )),
+ })
+ }
+
+ fn entry_key(&self) -> EntryKey {
+ EntryKey {
+ role_arn: self.role_arn.clone(),
+ token_file: self.token_file.clone(),
+ region: self.region.clone(),
+ max_attempts: self.max_attempts,
+ min_ttl: self.min_ttl,
+ max_jitter: self.max_jitter,
+ }
+ }
+}
+
+/// Process-wide cache key. A credential is shared per distinct identity AND
resolved settings, so a
+/// catalog that configures its own retry/refresh knobs gets its own entry
with its own
+/// configuration honored -- independent of which scan initializes first. Two
callers with the same
+/// identity and the same settings still share one entry (and one STS call).
+#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+struct EntryKey {
+ role_arn: String,
+ token_file: String,
+ region: Option<String>,
+ max_attempts: u32,
+ min_ttl: Duration,
+ max_jitter: Duration,
+}
+
+/// The shared, cached credential for one identity. `provider` is the AWS SDK
web-identity provider
+/// built once; `cached` holds the last credential; `refresh_jitter` is drawn
once per process so
+/// each executor refreshes at a slightly different time. `last_failure_at`
coalesces a burst of
+/// readers that hit a persistent throttle into a single STS call.
+#[derive(Debug)]
+struct SharedEntry {
+ provider: Arc<dyn ProvideCredentials>,
+ cached: RwLock<Option<Credentials>>,
+ /// Single-flights refreshes so a burst of readers triggers exactly one
STS call.
+ refresh_lock: tokio::sync::Mutex<()>,
+ /// When the last refresh failed. Waiters within `FAILURE_COOLDOWN` of
this get the failure
+ /// without re-calling STS, so a failed burst costs one call rather than
one per reader.
+ last_failure_at: RwLock<Option<Instant>>,
+ min_ttl: Duration,
+ refresh_jitter: Duration,
+}
+
+impl SharedEntry {
+ /// Returns the cached credential if it is still fresh, i.e. it does not
expire within
+ /// `min_ttl + refresh_jitter`.
+ fn fresh(&self) -> Option<Credentials> {
+ let guard = self.cached.read().unwrap();
+ let cred = guard.as_ref()?;
+ match cred.expiry() {
+ Some(expiry) => {
+ if expiry <= SystemTime::now() + self.min_ttl +
self.refresh_jitter {
+ None
+ } else {
+ Some(cred.clone())
+ }
+ }
+ // No expiry reported: keep it. Web-identity credentials normally
carry one.
+ None => Some(cred.clone()),
+ }
+ }
+
+ /// `Some(error)` if a refresh failed within the last `FAILURE_COOLDOWN`,
so callers can bail out
+ /// instead of piling another assume-role call onto a throttled STS.
+ fn in_failure_cooldown(&self) -> Option<String> {
+ let at = (*self.last_failure_at.read().unwrap())?;
+ (at.elapsed() < FAILURE_COOLDOWN).then(|| {
+ "web-identity credential refresh failed recently; backing off
before retrying STS"
+ .to_string()
+ })
+ }
+
+ /// Fetches a fresh credential, refreshing from STS at most once at a
time. On a refresh error
+ /// the error propagates -- we never fall back to a lower-privilege
identity -- and is briefly
+ /// remembered so concurrent waiters do not each re-issue the same
throttled call.
+ async fn credentials(&self) -> Result<Credentials, String> {
+ if let Some(cred) = self.fresh() {
+ return Ok(cred);
+ }
+ if let Some(err) = self.in_failure_cooldown() {
+ return Err(err);
+ }
+ let _guard = self.refresh_lock.lock().await;
+ // Re-check: another task may have refreshed (or just failed) while we
waited on the lock.
+ if let Some(cred) = self.fresh() {
+ return Ok(cred);
+ }
+ if let Some(err) = self.in_failure_cooldown() {
+ return Err(err);
+ }
+ match self.provider.provide_credentials().await {
+ Ok(cred) => {
+ *self.cached.write().unwrap() = Some(cred.clone());
+ *self.last_failure_at.write().unwrap() = None;
+ Ok(cred)
+ }
+ Err(e) => {
+ *self.last_failure_at.write().unwrap() = Some(Instant::now());
+ Err(format!("web-identity assume-role failed: {e}"))
+ }
+ }
+ }
+}
+
+/// Registry of shared credential entries, one per identity, for the lifetime
of the process.
+///
+/// Process lifetime is the right scope for the same reason as the region
cache in `s3.rs`: each
+/// executor is dedicated to one Spark application, and there is a bounded set
of assumed roles per
+/// job. Entries are never evicted; the map stays proportional to the number
of distinct roles.
+fn registry() -> &'static std::sync::Mutex<HashMap<EntryKey,
Arc<SharedEntry>>> {
+ static REGISTRY: OnceLock<std::sync::Mutex<HashMap<EntryKey,
Arc<SharedEntry>>>> =
+ OnceLock::new();
+ REGISTRY.get_or_init(|| std::sync::Mutex::new(HashMap::new()))
+}
+
+/// Returns the shared entry for `cfg`, building the AWS SDK provider once if
needed. The provider
+/// is built outside the registry lock (it is async); a concurrent builder
just loses the insert
+/// race, which is harmless.
+async fn shared_entry(cfg: &WebIdentityConfig) -> Arc<SharedEntry> {
+ let key = cfg.entry_key();
+ if let Some(entry) = registry().lock().unwrap().get(&key).cloned() {
+ return entry;
+ }
+
+ let provider = build_provider(cfg).await;
+ // Draw the refresh jitter once. subsec_nanos at build time differs across
processes, so this
+ // seeds a per-executor offset even before rand is consulted.
+ let jitter = if cfg.max_jitter.is_zero() {
+ Duration::ZERO
+ } else {
+
Duration::from_secs(rand::rng().random_range(0..=cfg.max_jitter.as_secs()))
+ };
+ let entry = Arc::new(SharedEntry {
+ provider,
+ cached: RwLock::new(None),
+ refresh_lock: tokio::sync::Mutex::new(()),
+ last_failure_at: RwLock::new(None),
+ min_ttl: cfg.min_ttl,
+ refresh_jitter: jitter,
+ });
+
+ let mut map = registry().lock().unwrap();
+ Arc::clone(map.entry(key).or_insert(entry))
+}
+
+/// Builds a web-identity-only AWS SDK credential provider with STS retry
raised to
+/// `cfg.max_attempts`. No IMDS/instance-role fallback is wired in, so a
throttle that outlasts the
+/// retries errors instead of downgrading.
+async fn build_provider(cfg: &WebIdentityConfig) -> Arc<dyn
ProvideCredentials> {
+ let provider_config = base_provider_config(cfg).await;
+ let provider = WebIdentityTokenCredentialsProvider::builder()
+ .configure(&provider_config)
+ .build();
+ Arc::new(provider)
+}
+
+/// The `ProviderConfig` shared by the production build and the tests: default
region resolution
+/// plus the raised STS retry budget. Tests attach an in-memory HTTP client to
this so the retry
+/// path can be exercised without a socket.
+async fn base_provider_config(cfg: &WebIdentityConfig) -> ProviderConfig {
+ ProviderConfig::with_default_region()
+ .await
+
.with_retry_config(RetryConfig::standard().with_max_attempts(cfg.max_attempts))
Review Comment:
Fixed. `base_provider_config` no longer uses a bare `with_default_region()`.
It now loads a
full `SdkConfig` (`aws_config::defaults(BehaviorVersion::latest()).load()`)
and forwards the
resolved `region`, `use_fips`, and `use_dual_stack` onto the web-identity
`ProviderConfig`.
Since `ProviderConfig::client_config()` feeds those into the STS client, the
call now
resolves the FIPS / dual-stack endpoint whenever `AWS_USE_FIPS_ENDPOINT` /
`AWS_USE_DUALSTACK_ENDPOINT` (or the profile equivalents) are set — matching
what the
default chain did. Loading `defaults()` also means the SDK applies the usual
environment-over-profile precedence for both settings.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
