feat(elasticsearch sink): fix partial retry logic

changelog.d/elasticsearch_sink_skip_retrying_succeeded_documents.feature.md

@@ -0,0 +1,3 @@
+The `request_retry_partial` behavior for the `elasticsearch` was changed. Now only the failed retriable requests in a bulk will be retried (instead of all requests in the bulk).
+
+authors: Serendo

src/sinks/aws_cloudwatch_logs/retry.rs

@@ -9,28 +9,34 @@ use crate::aws::is_retriable_error;
 use crate::sinks::{aws_cloudwatch_logs::service::CloudwatchError, util::retries::RetryLogic};
 
 #[derive(Debug)]
-pub struct CloudwatchRetryLogic<T> {
-    phantom: PhantomData<T>,
+pub struct CloudwatchRetryLogic<Request, Response> {
+    request: PhantomData<Request>,
+    response: PhantomData<Response>,
 }
-impl<T> CloudwatchRetryLogic<T> {
-    pub const fn new() -> CloudwatchRetryLogic<T> {
+impl<Request, Response> CloudwatchRetryLogic<Request, Response> {
+    pub const fn new() -> CloudwatchRetryLogic<Request, Response> {
         CloudwatchRetryLogic {
-            phantom: PhantomData,
+            request: PhantomData,
+            response: PhantomData,
         }
     }
 }
 
-impl<T> Clone for CloudwatchRetryLogic<T> {
+impl<Request, Response> Clone for CloudwatchRetryLogic<Request, Response> {
     fn clone(&self) -> Self {
         CloudwatchRetryLogic {
-            phantom: PhantomData,
+            request: PhantomData,
+            response: PhantomData,
         }
     }
 }
 
-impl<T: Send + Sync + 'static> RetryLogic for CloudwatchRetryLogic<T> {
+impl<Request: Send + Sync + 'static, Response: Send + Sync + 'static> RetryLogic
+    for CloudwatchRetryLogic<Request, Response>
+{
     type Error = CloudwatchError;
-    type Response = T;
+    type Request = Request;
+    type Response = Response;
 
     // TODO this match may not be necessary given the logic in `is_retriable_error()`
     #[allow(clippy::cognitive_complexity)] // long, but just a hair over our limit
@@ -84,7 +90,7 @@ mod test {
 
     #[test]
     fn test_throttle_retry() {
-        let retry_logic: CloudwatchRetryLogic<()> = CloudwatchRetryLogic::new();
+        let retry_logic: CloudwatchRetryLogic<(), ()> = CloudwatchRetryLogic::new();
 
         let meta_err = aws_smithy_types::error::ErrorMetadata::builder()
             .code("ThrottlingException")

src/sinks/aws_cloudwatch_logs/service.rs

@@ -50,7 +50,7 @@ type Svc = Buffer<
     <ConcurrencyLimit<
         RateLimit<
             Retry<
-                FibonacciRetryPolicy<CloudwatchRetryLogic<()>>,
+                FibonacciRetryPolicy<CloudwatchRetryLogic<Vec<InputLogEvent>, ()>>,
                 Buffer<
                     Vec<InputLogEvent>,
                     <Timeout<CloudwatchLogsSvc> as Service<Vec<InputLogEvent>>>::Future,

src/sinks/aws_cloudwatch_metrics/mod.rs

@@ -201,6 +201,7 @@ struct CloudWatchMetricsRetryLogic;
 
 impl RetryLogic for CloudWatchMetricsRetryLogic {
     type Error = SdkError<PutMetricDataError>;
+    type Request = PartitionInnerBuffer<Vec<Metric>, String>;
     type Response = ();
 
     fn is_retriable_error(&self, error: &Self::Error) -> bool {

src/sinks/aws_kinesis/config.rs

@@ -99,7 +99,7 @@ where
     R: Send + 'static,
     RR: Record + Record<T = R> + Clone + Send + Sync + Unpin + 'static,
     E: Send + 'static,
-    RT: RetryLogic<Response = KinesisResponse> + Default,
+    RT: RetryLogic<Request = BatchKinesisRequest<RR>, Response = KinesisResponse> + Default,
 {
     let request_limits = config.request.into_settings();
 

src/sinks/aws_kinesis/firehose/config.rs

@@ -16,6 +16,7 @@ use crate::{
     },
 };
 
+use super::sink::BatchKinesisRequest;
 use super::{
     build_sink,
     record::{KinesisFirehoseClient, KinesisFirehoseRecord},
@@ -173,6 +174,7 @@ struct KinesisRetryLogic {
 
 impl RetryLogic for KinesisRetryLogic {
     type Error = SdkError<KinesisError, HttpResponse>;
+    type Request = BatchKinesisRequest<KinesisFirehoseRecord>;
     type Response = KinesisResponse;
 
     fn is_retriable_error(&self, error: &Self::Error) -> bool {
@@ -187,7 +189,7 @@ impl RetryLogic for KinesisRetryLogic {
         is_retriable_error(error)
     }
 
-    fn should_retry_response(&self, response: &Self::Response) -> RetryAction {
+    fn should_retry_response(&self, response: &Self::Response) -> RetryAction<Self::Request> {
         if response.failure_count > 0 && self.retry_partial {
             let msg = format!("partial error count {}", response.failure_count);
             RetryAction::Retry(msg.into())

src/sinks/aws_kinesis/streams/config.rs

@@ -15,6 +15,7 @@ use crate::{
     },
 };
 
+use super::sink::BatchKinesisRequest;
 use super::{
     build_sink,
     record::{KinesisStreamClient, KinesisStreamRecord},
@@ -170,6 +171,7 @@ struct KinesisRetryLogic {
 
 impl RetryLogic for KinesisRetryLogic {
     type Error = SdkError<KinesisError, HttpResponse>;
+    type Request = BatchKinesisRequest<KinesisStreamRecord>;
     type Response = KinesisResponse;
 
     fn is_retriable_error(&self, error: &Self::Error) -> bool {
@@ -190,7 +192,7 @@ impl RetryLogic for KinesisRetryLogic {
         is_retriable_error(error)
     }
 
-    fn should_retry_response(&self, response: &Self::Response) -> RetryAction {
+    fn should_retry_response(&self, response: &Self::Response) -> RetryAction<Self::Request> {
         if response.failure_count > 0 && self.retry_partial {
             let msg = format!("partial error count {}", response.failure_count);
             RetryAction::Retry(msg.into())

src/sinks/aws_s_s/retry.rs

@@ -2,7 +2,7 @@ use std::marker::PhantomData;
 
 use aws_smithy_runtime_api::client::{orchestrator::HttpResponse, result::SdkError};
 
-use super::service::SendMessageResponse;
+use super::{request_builder::SendMessageEntry, service::SendMessageResponse};
 use crate::{aws::is_retriable_error, sinks::util::retries::RetryLogic};
 
 #[derive(Debug)]
@@ -26,6 +26,7 @@ where
     E: std::fmt::Debug + std::fmt::Display + std::error::Error + Sync + Send + 'static,
 {
     type Error = SdkError<E, HttpResponse>;
+    type Request = SendMessageEntry;
     type Response = SendMessageResponse;
 
     fn is_retriable_error(&self, error: &Self::Error) -> bool {

src/sinks/azure_common/config.rs

@@ -57,6 +57,7 @@ pub struct AzureBlobRetryLogic;
 
 impl RetryLogic for AzureBlobRetryLogic {
     type Error = HttpError;
+    type Request = AzureBlobRequest;
     type Response = AzureBlobResponse;
 
     fn is_retriable_error(&self, error: &Self::Error) -> bool {

src/sinks/clickhouse/service.rs

@@ -23,18 +23,19 @@ use snafu::ResultExt;
 
 #[derive(Debug, Default, Clone)]
 pub struct ClickhouseRetryLogic {
-    inner: HttpRetryLogic,
+    inner: HttpRetryLogic<HttpRequest<PartitionKey>>,
 }
 
 impl RetryLogic for ClickhouseRetryLogic {
     type Error = HttpError;
+    type Request = HttpRequest<PartitionKey>;
     type Response = HttpResponse;
 
     fn is_retriable_error(&self, error: &Self::Error) -> bool {
         self.inner.is_retriable_error(error)
     }
 
-    fn should_retry_response(&self, response: &Self::Response) -> RetryAction {
+    fn should_retry_response(&self, response: &Self::Response) -> RetryAction<Self::Request> {
         match response.http_response.status() {
             StatusCode::INTERNAL_SERVER_ERROR => {
                 let body = response.http_response.body();

src/sinks/clickhouse/sink.rs

@@ -83,7 +83,7 @@ where
 
 /// PartitionKey used to partition events by (database, table) pair.
 #[derive(Hash, Eq, PartialEq, Clone, Debug)]
-pub(super) struct PartitionKey {
+pub struct PartitionKey {
     pub database: String,
     pub table: String,
     pub format: Format,

src/sinks/databend/service.rs

@@ -23,6 +23,7 @@ pub struct DatabendRetryLogic;
 
 impl RetryLogic for DatabendRetryLogic {
     type Error = DatabendError;
+    type Request = DatabendRequest;
     type Response = DatabendResponse;
 
     fn is_retriable_error(&self, error: &Self::Error) -> bool {

src/sinks/datadog/logs/service.rs

@@ -29,6 +29,7 @@ pub struct LogApiRetry;
 
 impl RetryLogic for LogApiRetry {
     type Error = DatadogApiError;
+    type Request = LogApiRequest;
     type Response = LogApiResponse;
 
     fn is_retriable_error(&self, error: &Self::Error) -> bool {

src/sinks/datadog/metrics/service.rs

@@ -26,6 +26,7 @@ pub struct DatadogMetricsRetryLogic;
 
 impl RetryLogic for DatadogMetricsRetryLogic {
     type Error = DatadogApiError;
+    type Request = DatadogMetricsRequest;
     type Response = DatadogMetricsResponse;
 
     fn is_retriable_error(&self, error: &Self::Error) -> bool {

src/sinks/datadog/traces/service.rs

@@ -23,13 +23,14 @@ pub struct TraceApiRetry;
 
 impl RetryLogic for TraceApiRetry {
     type Error = HttpError;
+    type Request = TraceApiRequest;
     type Response = TraceApiResponse;
 
     fn is_retriable_error(&self, _error: &Self::Error) -> bool {
         true
     }
 
-    fn should_retry_response(&self, response: &Self::Response) -> RetryAction {
+    fn should_retry_response(&self, response: &Self::Response) -> RetryAction<Self::Request> {
         let status = response.status_code;
         match status {
             // Use the same status code/retry policy as the Trace agent, additionally retrying

src/sinks/elasticsearch/encoder.rs

@@ -19,7 +19,7 @@ use crate::{
     },
 };
 
-#[derive(Serialize)]
+#[derive(Serialize, Clone, Debug)]
 pub enum DocumentVersionType {
     External,
     ExternalGte,
@@ -34,20 +34,20 @@ impl DocumentVersionType {
     }
 }
 
-#[derive(Serialize)]
+#[derive(Serialize, Clone, Debug)]
 pub struct DocumentVersion {
     pub kind: DocumentVersionType,
     pub value: u64,
 }
 
-#[derive(Serialize)]
+#[derive(Serialize, Clone, Debug)]
 pub enum DocumentMetadata {
     WithoutId,
     Id(String),
     IdAndVersion(String, DocumentVersion),
 }
 
-#[derive(Serialize)]
+#[derive(Serialize, Clone, Debug)]
 pub struct ProcessedEvent {
     pub index: String,
     pub bulk_action: BulkAction,

src/sinks/elasticsearch/request_builder.rs

@@ -26,6 +26,7 @@ pub struct Metadata {
     finalizers: EventFinalizers,
     batch_size: usize,
     events_byte_size: JsonSize,
+    original_events: Vec<ProcessedEvent>,
 }
 
 impl RequestBuilder<Vec<ProcessedEvent>> for ElasticsearchRequestBuilder {
@@ -60,6 +61,7 @@ impl RequestBuilder<Vec<ProcessedEvent>> for ElasticsearchRequestBuilder {
             finalizers: events.take_finalizers(),
             batch_size: events.len(),
             events_byte_size,
+            original_events: events.clone(),
         };
         (es_metadata, metadata_builder, events)
     }
@@ -76,6 +78,8 @@ impl RequestBuilder<Vec<ProcessedEvent>> for ElasticsearchRequestBuilder {
             batch_size: es_metadata.batch_size,
             events_byte_size: es_metadata.events_byte_size,
             metadata,
+            original_events: es_metadata.original_events,
+            elasticsearch_request_builder: self.clone(),
         }
     }
 }