sunchao commented on code in PR #6025:
URL: https://github.com/apache/datafusion-comet/pull/6025#discussion_r4050898438
##########
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:
### Correctness
[P2] Could this preserve the higher-priority environment credential source
before installing the web-identity-only provider? With the native default chain
selected (no nonempty `aws.credentials.provider`) and both IRSA variables
present, a valid `AWS_ACCESS_KEY_ID` / `AWS_SECRET_ACCESS_KEY` pair previously
wins through the AWS SDK's `Environment -> Profile -> WebIdentityToken` chain.
This branch now ignores that pair and assumes `AWS_ROLE_ARN` instead. An
executor that has explicit credentials for a second account will therefore read
as the service-account role and can get `AccessDenied`, even though its
existing credentials still work. Please keep the original precedence for
explicit environment credentials and add a synthetic mixed-environment case.
##########
native/core/src/cloud/s3/web_identity.rs:
##########
@@ -0,0 +1,756 @@
+// 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
+//! (role_arn, token_file, region), and shared across all reader threads
and scans. Refresh
+//! fires ahead of expiry by `min_ttl` plus a per-process random jitter
so cluster-wide
+//! refreshes do not synchronize into another burst.
+//!
+//! 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, 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;
+
+/// 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(),
+ }
+ }
+}
+
+/// Process-wide cache key. One assumed-role credential is shared per distinct
identity.
+#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+struct EntryKey {
+ role_arn: String,
+ token_file: String,
+ region: Option<String>,
+}
+
+/// 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.
+#[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<()>,
+ 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()),
+ }
+ }
+
+ /// 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.
+ async fn credentials(&self) -> Result<Credentials, String> {
+ if let Some(cred) = self.fresh() {
+ return Ok(cred);
+ }
+ let _guard = self.refresh_lock.lock().await;
+ // Re-check: another task may have refreshed while we waited on the
lock.
+ if let Some(cred) = self.fresh() {
+ return Ok(cred);
+ }
+ let cred = self
+ .provider
+ .provide_credentials()
+ .await
+ .map_err(|e| format!("web-identity assume-role failed: {e}"))?;
Review Comment:
### Performance
[P2] Could the shared entry also coalesce a failed refresh, or apply a
bounded cooldown before another waiter retries? When `provide_credentials()`
exhausts its STS retries, `?` releases the mutex without recording any result.
Every reader already waiting on that mutex then sees the same empty or stale
cache and calls the provider again. An isolated harness using these exact
`SharedEntry` methods produces eight provider calls for eight concurrent
readers on persistent failure, versus one call on success. Thus the throttle
case this change targets still scales STS work with the number of readers and
serializes repeated refresh attempts. Please cover a concurrent failed burst as
well as the successful burst.
##########
native/core/src/cloud/s3/web_identity.rs:
##########
@@ -0,0 +1,756 @@
+// 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
+//! (role_arn, token_file, region), and shared across all reader threads
and scans. Refresh
+//! fires ahead of expiry by `min_ttl` plus a per-process random jitter
so cluster-wide
+//! refreshes do not synchronize into another burst.
+//!
+//! 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, 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;
+
+/// 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(),
+ }
+ }
+}
+
+/// Process-wide cache key. One assumed-role credential is shared per distinct
identity.
+#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+struct EntryKey {
+ role_arn: String,
+ token_file: String,
+ region: Option<String>,
+}
+
+/// 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.
+#[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<()>,
+ 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()),
+ }
+ }
+
+ /// 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.
+ async fn credentials(&self) -> Result<Credentials, String> {
+ if let Some(cred) = self.fresh() {
+ return Ok(cred);
+ }
+ let _guard = self.refresh_lock.lock().await;
+ // Re-check: another task may have refreshed while we waited on the
lock.
+ if let Some(cred) = self.fresh() {
+ return Ok(cred);
+ }
+ let cred = self
+ .provider
+ .provide_credentials()
+ .await
+ .map_err(|e| format!("web-identity assume-role failed: {e}"))?;
+ *self.cached.write().unwrap() = Some(cred.clone());
+ Ok(cred)
+ }
+}
+
+/// 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;
Review Comment:
### Correctness
[P2] Could entry reuse account for the resolved retry and refresh settings?
`EntryKey` contains only the role ARN, token file and region, but the stored
provider, `min_ttl` and jitter are created from the first caller's
configuration. For example, if a scan first creates the default five-attempt
entry, a second catalog configured with
`comet.s3.credentials.webIdentity.maxAttempts=8` immediately returns that entry
and still gets five attempts. The reverse initialization order also changes the
first catalog's behavior. The same issue affects both refresh knobs and sharing
between Parquet and Iceberg. Please make the effective settings deterministic
and honor the documented per-catalog configuration.
--
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]
