parthchandra commented on code in PR #6025: URL: https://github.com/apache/datafusion-comet/pull/6025#discussion_r4051355771
########## 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: Fixed — a failed refresh is now remembered briefly so a burst collapses to one STS call. - `SharedEntry` gains `last_failure_at: RwLock<Option<Instant>>`. On a refresh error we record the time instead of leaving the cache empty. - `credentials()` checks `in_failure_cooldown()` both before and after taking the refresh lock. Within `FAILURE_COOLDOWN` (1s) waiters return the error without calling STS; a successful refresh clears the marker. - The cooldown is short and bounded on purpose: the SDK has already spent its retry budget by the time we record a failure, and we still want to recover quickly once the throttle clears. Test added: `concurrent_failed_refresh_is_coalesced` — 8 concurrent readers against a persistently failing provider now produce exactly **1** provider call (was 8), and all receive an error (no downgrade). ########## 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: Good catch, fixed. The take-over now stands aside when explicit static credentials are set in the environment, so it never shadows a higher-priority source. - Added `explicit_env_credentials()` (non-empty `AWS_ACCESS_KEY_ID` + `AWS_SECRET_ACCESS_KEY`) and short-circuit it in the shared `take_over_if_irsa`, so both the Parquet and Iceberg paths keep the default-chain precedence (Environment -> Profile -> WebIdentity). - When env creds are present we return `None` and let the existing default chain run, exactly as before this change. Test added: `explicit_env_credentials_keep_precedence_over_irsa` — with IRSA present and no env creds it takes over; with `AWS_ACCESS_KEY_ID` / `AWS_SECRET_ACCESS_KEY` set it stands aside. Scope note: I guard the environment case you described. Profile credentials aren't specially detected (reliably probing profile presence is messier and uncommon on executors); that user can name the provider explicitly. Happy to extend to profiles if you think it's worth it. ########## 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: Fixed — the resolved settings are now part of the cache key, so configuration is honored regardless of initialization order. - `EntryKey` now includes `max_attempts`, `min_ttl`, and `max_jitter` in addition to `(role_arn, token_file, region)`. - A catalog configured with `comet.s3.credentials.webIdentity.maxAttempts=8` gets its own entry with 8 attempts even if a default (5-attempt) entry for the same identity was created first; the reverse order no longer changes the first catalog's behavior either. - Two callers with the same identity **and** settings still share one entry (and one STS call), so the sharing benefit is preserved for the common case. Applies across the Parquet and Iceberg paths. Test added: `entry_key_includes_resolved_settings` — same identity with different `maxAttempts` produces distinct keys; identical identity + settings produce equal keys. -- 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]
