chore(sinks): remove downcast and trait in sink retry code (#23543)

* Reapply "Remove downcast ref and trait usage from partial retry logic"

This reverts commit 3bd9c0d.

* Reapply "Fix make check-clippy errors"

This reverts commit abeb520.

* Use Fn instead of FnOnce

Author: thomasqueirozb

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,