diff --git a/src/sources/aws_s3/mod.rs b/src/sources/aws_s3/mod.rs index 7794f6612..ba613ab9f 100644 --- a/src/sources/aws_s3/mod.rs +++ b/src/sources/aws_s3/mod.rs @@ -17,8 +17,8 @@ use vrl::value::{kind::Collection, Kind}; use super::util::MultilineConfig; use crate::codecs::DecodingConfig; use crate::{ - aws::{auth::AwsAuthentication, create_client, create_client_and_region, RegionOrEndpoint}, - common::{s3::S3ClientBuilder, sqs::SqsClientBuilder}, + aws::{auth::AwsAuthentication, create_client_and_region, RegionOrEndpoint}, + common::sqs::SqsClientBuilder, config::{ ProxyConfig, SourceAcknowledgementsConfig, SourceConfig, SourceContext, SourceOutput, }, @@ -207,20 +207,6 @@ impl AwsS3Config { ) -> crate::Result { let region = self.region.region(); let endpoint = self.region.endpoint(); - let force_path_style_value: bool = true; - - let s3_client = create_client::( - &S3ClientBuilder { - force_path_style: Some(force_path_style_value), - }, - &self.auth, - region.clone(), - endpoint.clone(), - proxy, - self.tls_options.as_ref(), - None, - ) - .await?; let decoder = DecodingConfig::new(self.framing.clone(), self.decoding.clone(), log_namespace) @@ -232,7 +218,7 @@ impl AwsS3Config { &SqsClientBuilder {}, &self.auth, region.clone(), - endpoint, + endpoint.clone(), proxy, sqs.tls_options.as_ref(), sqs.timeout.as_ref(), @@ -242,11 +228,14 @@ impl AwsS3Config { let ingestor = sqs::Ingestor::new( region, sqs_client, - s3_client, sqs.clone(), self.compression, multiline, decoder, + self.auth.clone(), + endpoint, + proxy.clone(), + self.tls_options.clone(), ) .await?; @@ -390,8 +379,9 @@ mod integration_tests { config::{ProxyConfig, SourceConfig, SourceContext}, event::EventStatus::{self, *}, line_agg, + common::s3::S3ClientBuilder, sources::{ - aws_s3::{sqs::S3Event, S3ClientBuilder}, + aws_s3::sqs::S3Event, util::MultilineConfig, }, test_util::{ diff --git a/src/sources/aws_s3/sqs.rs b/src/sources/aws_s3/sqs.rs index 479c7174a..beb574af0 100644 --- a/src/sources/aws_s3/sqs.rs +++ b/src/sources/aws_s3/sqs.rs @@ -1,4 +1,5 @@ -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; +use std::sync::RwLock; use std::{future::ready, num::NonZeroUsize, panic, sync::Arc, sync::LazyLock}; use std::time::Duration; use aws_sdk_s3::operation::get_object::GetObjectError; @@ -31,8 +32,9 @@ use vector_lib::internal_event::{ use crate::codecs::Decoder; use crate::event::{Event, LogEvent}; use crate::{ - aws::AwsTimeout, - config::{SourceAcknowledgementsConfig, SourceContext}, + aws::{auth::AwsAuthentication, create_client, AwsTimeout}, + common::s3::S3ClientBuilder, + config::{ProxyConfig, SourceAcknowledgementsConfig, SourceContext}, event::{BatchNotifier, BatchStatus, EstimatedJsonEncodedSizeOf}, internal_events::{ EventsReceived, SqsMessageDeleteBatchError, SqsMessageDeletePartialError, @@ -222,16 +224,10 @@ pub enum ProcessingError { bucket: String, key: String, }, - #[snafu(display( - "Object notification for s3://{}/{} is a bucket in another region: {}", - bucket, - key, - region - ))] - WrongRegion { + #[snafu(display("Failed to create S3 client for region {}: {}", region, source))] + CreateS3Client { + source: Box, region: String, - bucket: String, - key: String, }, #[snafu(display("Unsupported S3 event version: {}.", version,))] UnsupportedS3EventVersion { version: semver::Version }, @@ -250,10 +246,14 @@ pub enum ProcessingError { pub struct State { region: Region, - - s3_client: S3Client, sqs_client: SqsClient, + auth: AwsAuthentication, + endpoint: Option, + proxy: ProxyConfig, + tls_options: Option, + s3_client_cache: RwLock>, + multiline: Option, compression: super::Compression, @@ -268,6 +268,34 @@ pub struct State { pub message_format: SqsMessageFormat, } +impl State { + async fn get_s3_client(&self, region_str: &str) -> Result { + if let Some(client) = self.s3_client_cache.read().unwrap().get(region_str) { + return Ok(client.clone()); + } + + let client = create_client::( + &S3ClientBuilder { + force_path_style: Some(true), + }, + &self.auth, + Some(Region::new(region_str.to_owned())), + self.endpoint.clone(), + &self.proxy, + self.tls_options.as_ref(), + None, + ) + .await + .map_err(|e| ProcessingError::CreateS3Client { + source: e.into(), + region: region_str.to_owned(), + })?; + + let mut cache = self.s3_client_cache.write().unwrap(); + Ok(cache.entry(region_str.to_owned()).or_insert(client).clone()) + } +} + pub(super) struct Ingestor { state: Arc, } @@ -276,23 +304,32 @@ impl Ingestor { pub(super) async fn new( region: Region, sqs_client: SqsClient, - s3_client: S3Client, config: Config, compression: super::Compression, multiline: Option, decoder: Decoder, + auth: AwsAuthentication, + endpoint: Option, + proxy: ProxyConfig, + tls_options: Option, ) -> Result { if config.max_number_of_messages < 1 || config.max_number_of_messages > 10 { return Err(IngestorNewError::InvalidNumberOfMessages { messages: config.max_number_of_messages, }); } + let s3_client_cache = RwLock::new(BTreeMap::new()); + let state = Arc::new(State { region, - - s3_client, sqs_client, + auth, + endpoint, + proxy, + tls_options, + s3_client_cache, + compression, multiline, @@ -570,19 +607,9 @@ impl IngestorProcess { return Ok(()); } - // S3 has to send notifications to a queue in the same region so I don't think this will - // actually ever be hit unless messages are being forwarded from one queue to another - if self.state.region.as_ref() != s3_event.aws_region.as_str() { - return Err(ProcessingError::WrongRegion { - bucket: s3_event.s3.bucket.name.clone(), - key: s3_event.s3.object.key.clone(), - region: s3_event.aws_region, - }); - } + let s3_client = self.state.get_s3_client(&s3_event.aws_region).await?; - let object_result = self - .state - .s3_client + let object_result = s3_client .get_object() .bucket(s3_event.s3.bucket.name.clone()) .key(s3_event.s3.object.key.clone())