vector/sinks/aws_kinesis/streams/
config.rs

1use aws_sdk_kinesis::operation::{
2    describe_stream::DescribeStreamError, put_records::PutRecordsError,
3};
4use aws_smithy_runtime_api::client::{orchestrator::HttpResponse, result::SdkError};
5use futures::FutureExt;
6use snafu::Snafu;
7use vector_lib::configurable::{component::GenerateConfig, configurable_component};
8
9use super::{
10    KinesisClient, KinesisError, KinesisRecord, KinesisResponse, KinesisSinkBaseConfig, build_sink,
11    record::{KinesisStreamClient, KinesisStreamRecord},
12    sink::BatchKinesisRequest,
13};
14use crate::{
15    aws::{ClientBuilder, create_client, is_retriable_error},
16    config::{
17        AcknowledgementsConfig, DynValidatedSink, Input, ProxyConfig, SinkConfig, SinkContext,
18        ValidatedSink,
19    },
20    sinks::{
21        Healthcheck, VectorSink,
22        prelude::*,
23        util::{
24            BatchConfig, SinkBatchSettings,
25            retries::{RetryAction, RetryLogic},
26        },
27    },
28};
29
30#[allow(clippy::large_enum_variant)]
31#[derive(Debug, Snafu)]
32enum HealthcheckError {
33    #[snafu(display("DescribeStream failed: {}", source))]
34    DescribeStreamFailed {
35        source: SdkError<DescribeStreamError, HttpResponse>,
36    },
37    #[snafu(display("Stream names do not match, got {}, expected {}", name, stream_name))]
38    StreamNamesMismatch { name: String, stream_name: String },
39}
40
41pub struct KinesisClientBuilder;
42
43impl ClientBuilder for KinesisClientBuilder {
44    type Client = KinesisClient;
45
46    fn build(&self, config: &aws_types::SdkConfig) -> Self::Client {
47        KinesisClient::new(config)
48    }
49}
50
51pub const MAX_PAYLOAD_SIZE: usize = 5_000_000;
52pub const MAX_PAYLOAD_EVENTS: usize = 500;
53
54#[derive(Clone, Copy, Debug, Default)]
55pub struct KinesisDefaultBatchSettings;
56
57impl SinkBatchSettings for KinesisDefaultBatchSettings {
58    const MAX_EVENTS: Option<usize> = Some(MAX_PAYLOAD_EVENTS);
59    const MAX_BYTES: Option<usize> = Some(MAX_PAYLOAD_SIZE);
60    const TIMEOUT_SECS: f64 = 1.0;
61}
62
63/// Configuration for the `aws_kinesis_streams` sink.
64#[configurable_component(sink(
65    "aws_kinesis_streams",
66    "Publish logs to AWS Kinesis Streams topics."
67))]
68#[derive(Clone, Debug)]
69pub struct KinesisStreamsSinkConfig {
70    #[serde(flatten)]
71    pub base: KinesisSinkBaseConfig,
72
73    #[configurable(derived)]
74    #[serde(default)]
75    pub batch: BatchConfig<KinesisDefaultBatchSettings>,
76}
77
78impl KinesisStreamsSinkConfig {
79    async fn healthcheck(self, client: KinesisClient) -> crate::Result<()> {
80        let stream_name = self.base.stream_name;
81
82        let describe_result = client
83            .describe_stream()
84            .stream_name(stream_name.clone())
85            .set_exclusive_start_shard_id(None)
86            .limit(1)
87            .send()
88            .await;
89
90        match describe_result {
91            Ok(resp) => {
92                let name = resp
93                    .stream_description
94                    .map(|x| x.stream_name)
95                    .unwrap_or_default();
96                if name == stream_name {
97                    Ok(())
98                } else {
99                    Err(HealthcheckError::StreamNamesMismatch { name, stream_name }.into())
100                }
101            }
102            Err(source) => Err(HealthcheckError::DescribeStreamFailed { source }.into()),
103        }
104    }
105
106    pub async fn create_client(&self, proxy: &ProxyConfig) -> crate::Result<KinesisClient> {
107        create_client::<KinesisClientBuilder>(
108            &KinesisClientBuilder {},
109            &self.base.auth,
110            self.base.region.region(),
111            self.base.region.endpoint(),
112            proxy,
113            self.base.tls.as_ref(),
114            None,
115        )
116        .await
117    }
118}
119
120#[async_trait::async_trait]
121#[typetag::serde(name = "aws_kinesis_streams")]
122impl SinkConfig for KinesisStreamsSinkConfig {
123    fn input(&self) -> Input {
124        self.base.input()
125    }
126
127    fn acknowledgements(&self) -> &AcknowledgementsConfig {
128        self.base.acknowledgements()
129    }
130
131    fn as_dyn_validated(&self) -> Option<&dyn DynValidatedSink> {
132        Some(self)
133    }
134}
135
136#[derive(Clone, Debug)]
137pub struct ValidatedKinesisStreams {
138    batch_settings: BatcherSettings,
139}
140
141#[async_trait::async_trait]
142impl ValidatedSink for KinesisStreamsSinkConfig {
143    type Validated = ValidatedKinesisStreams;
144
145    fn validate(&self) -> crate::Result<ValidatedKinesisStreams> {
146        let batch_settings = self
147            .batch
148            .validate()?
149            .limit_max_bytes(MAX_PAYLOAD_SIZE)?
150            .limit_max_events(MAX_PAYLOAD_EVENTS)?
151            .into_batcher_settings()?;
152
153        Ok(ValidatedKinesisStreams { batch_settings })
154    }
155
156    async fn build(
157        &self,
158        validated: &ValidatedKinesisStreams,
159        cx: SinkContext,
160    ) -> crate::Result<(VectorSink, Healthcheck)> {
161        let client = self.create_client(&cx.proxy).await?;
162        let healthcheck = self.clone().healthcheck(client.clone()).boxed();
163
164        let sink = build_sink::<
165            KinesisStreamClient,
166            KinesisRecord,
167            KinesisStreamRecord,
168            KinesisError,
169            KinesisRetryLogic,
170        >(
171            &self.base,
172            self.base.partition_key_field.clone(),
173            validated.batch_settings,
174            KinesisStreamClient { client },
175            KinesisRetryLogic {
176                retry_partial: self.base.request_retry_partial,
177            },
178        )?;
179
180        Ok((sink, healthcheck))
181    }
182}
183
184impl GenerateConfig for KinesisStreamsSinkConfig {
185    fn generate_config() -> serde_json::Value {
186        serde_yaml::from_str(indoc::indoc! {
187            r#"partition_key_field: foo
188            stream_name: my-stream
189            encoding:
190              codec: json"#,
191        })
192        .unwrap()
193    }
194}
195#[derive(Default, Clone)]
196struct KinesisRetryLogic {
197    retry_partial: bool,
198}
199
200impl RetryLogic for KinesisRetryLogic {
201    type Error = SdkError<KinesisError, HttpResponse>;
202    type Request = BatchKinesisRequest<KinesisStreamRecord>;
203    type Response = KinesisResponse;
204
205    fn is_retriable_error(&self, error: &Self::Error) -> bool {
206        if let SdkError::ServiceError(inner) = error {
207            // Note that if the request partially fails (records sent to one
208            // partition fail but the others do not, for example), Vector
209            // does not retry. This line only covers a failure for the entire
210            // request.
211            //
212            // https://github.com/vectordotdev/vector/issues/359
213            if matches!(
214                inner.err(),
215                PutRecordsError::ProvisionedThroughputExceededException(_)
216            ) {
217                return true;
218            }
219        }
220        is_retriable_error(error)
221    }
222
223    fn should_retry_response(&self, response: &Self::Response) -> RetryAction<Self::Request> {
224        if response.failure_count > 0 && self.retry_partial && !response.failed_records.is_empty() {
225            let failed_records = response.failed_records.clone();
226            RetryAction::RetryPartial(Box::new(move |original_request| {
227                let failed_events: Vec<_> = failed_records
228                    .iter()
229                    .filter_map(|r| original_request.events.get(r.index).cloned())
230                    .collect();
231
232                let metadata = RequestMetadata::from_batch(
233                    failed_events.iter().map(|req| req.get_metadata().clone()),
234                );
235
236                BatchKinesisRequest {
237                    events: failed_events,
238                    metadata,
239                }
240            }))
241        } else {
242            RetryAction::Successful
243        }
244    }
245}
246
247#[cfg(test)]
248mod tests {
249    use super::*;
250
251    #[test]
252    fn generate_config() {
253        crate::test_util::test_generate_config::<KinesisStreamsSinkConfig>();
254    }
255
256    #[test]
257    fn validate_produces_batch_settings() {
258        let config = KinesisStreamsSinkConfig {
259            batch: BatchConfig::<KinesisDefaultBatchSettings>::default(),
260            base: KinesisSinkBaseConfig {
261                stream_name: String::from("test"),
262                region: crate::aws::RegionOrEndpoint::with_both(
263                    "us-east-1",
264                    "http://localhost:4566",
265                ),
266                encoding: vector_lib::codecs::JsonSerializerConfig::default().into(),
267                compression: Compression::None,
268                request: Default::default(),
269                tls: None,
270                auth: Default::default(),
271                request_retry_partial: false,
272                acknowledgements: Default::default(),
273                partition_key_field: None,
274            },
275        };
276
277        let validated = config.validate().expect("validation should succeed");
278        assert_eq!(validated.batch_settings.item_limit, MAX_PAYLOAD_EVENTS);
279        assert_eq!(validated.batch_settings.size_limit, MAX_PAYLOAD_SIZE);
280    }
281}
vector/sinks/aws_kinesis/streams/config.rs

vector/sinks/aws_kinesis/streams/
config.rs
