fix(remote_storage): use cached credentials (#3128)

IMDSv2 has limits, and if we query it on every s3 interaction we are
going to go over those limits. Changes the s3_bucket client
configuration to use:
- ChainCredentialsProvider to handle env variables or imds usage
- LazyCachingCredentialsProvider to actually cache any credentials

Related: https://github.com/awslabs/aws-sdk-rust/issues/629
Possibly related: https://github.com/neondatabase/neon/issues/3118
This commit is contained in:
Joonas Koivunen
2022-12-16 13:40:01 +02:00
committed by Joonas Koivunen
parent d8f6d6fd6f
commit 73ea0a0b01

View File

@@ -4,14 +4,13 @@
//! allowing multiple api users to independently work with the same S3 bucket, if
//! their bucket prefixes are both specified and different.
use std::env::var;
use std::sync::Arc;
use std::time::Duration;
use anyhow::Context;
use aws_config::{
environment::credentials::EnvironmentVariableCredentialsProvider, imds,
imds::credentials::ImdsCredentialsProvider, meta::credentials::provide_credentials_fn,
environment::credentials::EnvironmentVariableCredentialsProvider,
imds::credentials::ImdsCredentialsProvider,
meta::credentials::{CredentialsProviderChain, LazyCachingCredentialsProvider},
};
use aws_sdk_s3::{
config::Config,
@@ -20,7 +19,6 @@ use aws_sdk_s3::{
Client, Endpoint, Region,
};
use aws_smithy_http::body::SdkBody;
use aws_types::credentials::{CredentialsError, ProvideCredentials};
use hyper::Body;
use tokio::{io, sync::Semaphore};
use tokio_util::io::ReaderStream;
@@ -31,8 +29,6 @@ use crate::{
Download, DownloadError, RemotePath, RemoteStorage, S3Config, REMOTE_STORAGE_PREFIX_SEPARATOR,
};
const DEFAULT_IMDS_TIMEOUT: Duration = Duration::from_secs(10);
pub(super) mod metrics {
use metrics::{register_int_counter_vec, IntCounterVec};
use once_cell::sync::Lazy;
@@ -122,30 +118,23 @@ impl S3Bucket {
"Creating s3 remote storage for S3 bucket {}",
aws_config.bucket_name
);
let credentials_provider = {
// uses "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"
let env_creds = EnvironmentVariableCredentialsProvider::new();
// uses imds v2
let imds = ImdsCredentialsProvider::builder().build();
// finally add caching.
// this might change in future, see https://github.com/awslabs/aws-sdk-rust/issues/629
LazyCachingCredentialsProvider::builder()
.load(CredentialsProviderChain::first_try("env", env_creds).or_else("imds", imds))
.build()
};
let mut config_builder = Config::builder()
.region(Region::new(aws_config.bucket_region.clone()))
.credentials_provider(provide_credentials_fn(|| async {
match var("AWS_ACCESS_KEY_ID").is_ok() && var("AWS_SECRET_ACCESS_KEY").is_ok() {
true => {
EnvironmentVariableCredentialsProvider::new()
.provide_credentials()
.await
}
false => {
let imds_client = imds::Client::builder()
.connect_timeout(DEFAULT_IMDS_TIMEOUT)
.read_timeout(DEFAULT_IMDS_TIMEOUT)
.build()
.await
.map_err(CredentialsError::unhandled)?;
ImdsCredentialsProvider::builder()
.imds_client(imds_client)
.build()
.provide_credentials()
.await
}
}
}));
.credentials_provider(credentials_provider);
if let Some(custom_endpoint) = aws_config.endpoint.clone() {
let endpoint = Endpoint::immutable(