feat: add ProduceResult with request size metadata

Signed-off-by: WenyXu <[email protected]>

Author: WenyXu

src/client/partition.rs

@@ -1,6 +1,9 @@
 use crate::{
     backoff::{Backoff, BackoffConfig, ErrorOrThrottle},
-    client::error::{Error, RequestContext, Result},
+    client::{
+        error::{Error, RequestContext, Result},
+        producer::ProduceResult,
+    },
     connection::{
         BrokerCache, BrokerCacheGeneration, BrokerConnection, BrokerConnector, MessengerTransport,
         MetadataLookupMode,
@@ -15,6 +18,7 @@ use crate::{
             ListOffsetsRequest, ListOffsetsRequestPartition, ListOffsetsRequestTopic,
             ListOffsetsResponse, ListOffsetsResponsePartition, NORMAL_CONSUMER, ProduceRequest,
             ProduceRequestPartitionData, ProduceRequestTopicData, ProduceResponse,
+            ResponseBodyWithMetadata,
         },
         primitives::*,
         record::{Record as ProtocolRecord, *},
@@ -195,10 +199,10 @@ impl PartitionClient {
         &self,
         records: Vec<Record>,
         compression: Compression,
-    ) -> Result<Vec<i64>> {
+    ) -> Result<ProduceResult> {
         // skip request entirely if `records` is empty
         if records.is_empty() {
-            return Ok(vec![]);
+            return Ok(ProduceResult::default());
         }
 
         let n = records.len() as i64;
@@ -214,13 +218,20 @@ impl PartitionClient {
                     .get()
                     .await
                     .map_err(|e| ErrorOrThrottle::Error((e, None)))?;
-                let response = broker
-                    .request(&request)
+                let ResponseBodyWithMetadata {
+                    response,
+                    encoded_request_size,
+                } = broker
+                    .request_with_metadata(&request)
                     .await
                     .map_err(|e| ErrorOrThrottle::Error((e.into(), Some(r#gen))))?;
                 maybe_throttle(response.throttle_time_ms)?;
                 process_produce_response(self.partition, &self.topic, n, response)
                     .map_err(|e| ErrorOrThrottle::Error((e, Some(r#gen))))
+                    .map(|offsets| ProduceResult {
+                        offsets,
+                        encoded_request_size,
+                    })
             },
         )
         .await

src/client/producer.rs

@@ -207,7 +207,7 @@
 use std::sync::Arc;
 use std::time::Duration;
 
-use futures::future::BoxFuture;
+use futures::{FutureExt, TryFutureExt, future::BoxFuture};
 use thiserror::Error;
 use tokio::task::JoinHandle;
 use tracing::*;
@@ -300,6 +300,16 @@ impl BatchProducerBuilder {
     }
 }
 
+/// The result of the produce call.
+#[derive(Debug, Default)]
+pub struct ProduceResult {
+    /// The offsets of the produced records.
+    pub offsets: Vec<i64>,
+
+    /// The size of the request encoded in bytes.
+    pub encoded_request_size: usize,
+}
+
 /// The [`ProducerClient`] provides an abstraction over a Kafka client than can
 /// produce a record.
 ///
@@ -316,15 +326,15 @@ pub trait ProducerClient: std::fmt::Debug + Send + Sync {
         &self,
         records: Vec<Record>,
         compression: Compression,
-    ) -> BoxFuture<'_, Result<Vec<i64>, ClientError>>;
+    ) -> BoxFuture<'_, Result<ProduceResult, ClientError>>;
 }
 
 impl ProducerClient for PartitionClient {
     fn produce(
         &self,
         records: Vec<Record>,
         compression: Compression,
-    ) -> BoxFuture<'_, Result<Vec<i64>, ClientError>> {
+    ) -> BoxFuture<'_, Result<ProduceResult, ClientError>> {
         Box::pin(self.produce(records, compression))
     }
 }
@@ -703,7 +713,7 @@ mod tests {
             &self,
             records: Vec<Record>,
             _compression: Compression,
-        ) -> BoxFuture<'_, Result<Vec<i64>, ClientError>> {
+        ) -> BoxFuture<'_, Result<ProduceResult, ClientError>> {
             Box::pin(async move {
                 tokio::time::sleep(self.delay).await;
 
@@ -727,7 +737,12 @@ mod tests {
                     .map(|x| (x + offset_base) as i64)
                     .collect();
                 batch_sizes.push(records.len());
-                Ok(offsets)
+                let record_size = records.iter().map(|r| r.approximate_size()).sum::<usize>();
+                Ok(ProduceResult {
+                    offsets,
+                    // Uses the approximate size of the records to estimate the size of the request.
+                    encoded_request_size: record_size,
+                })
             })
         }
     }

src/client/producer/batch.rs

@@ -158,7 +158,7 @@ where
             async move {
                 let res = match client.produce(batch, compression).await {
                     Ok(status) => Ok(Arc::new(AggregatedStatus {
-                        aggregated_status: status,
+                        aggregated_status: status.offsets,
                         status_deagg,
                     })),
                     Err(e) => {

src/messenger.rs

@@ -27,7 +27,10 @@ use tokio::{
 };
 use tracing::{debug, info, warn};
 
-use crate::protocol::{messages::ApiVersionsRequest, traits::ReadType};
+use crate::protocol::{
+    messages::{ApiVersionsRequest, ResponseBodyWithMetadata},
+    traits::ReadType,
+};
 use crate::{
     backoff::ErrorOrThrottle,
     protocol::{
@@ -315,6 +318,19 @@ where
     }
 
     pub async fn request<R>(&self, msg: R) -> Result<R::ResponseBody, RequestError>
+    where
+        R: RequestBody + Send + WriteVersionedType<Vec<u8>>,
+        R::ResponseBody: ReadVersionedType<Cursor<Vec<u8>>>,
+    {
+        self.request_with_version_ranges(msg, &self.version_ranges)
+            .await
+            .map(|body| body.response)
+    }
+
+    pub async fn request_with_metadata<R>(
+        &self,
+        msg: R,
+    ) -> Result<ResponseBodyWithMetadata<R::ResponseBody>, RequestError>
     where
         R: RequestBody + Send + WriteVersionedType<Vec<u8>>,
         R::ResponseBody: ReadVersionedType<Cursor<Vec<u8>>>,
@@ -327,7 +343,7 @@ where
         &self,
         msg: R,
         version_ranges: &HashMap<ApiKey, ApiVersionRange>,
-    ) -> Result<R::ResponseBody, RequestError>
+    ) -> Result<ResponseBodyWithMetadata<R::ResponseBody>, RequestError>
     where
         R: RequestBody + Send + WriteVersionedType<Vec<u8>>,
         R::ResponseBody: ReadVersionedType<Cursor<Vec<u8>>>,
@@ -371,6 +387,7 @@ where
             .write_versioned(&mut buf, header_version)
             .expect("Writing header to buffer should always work");
         msg.write_versioned(&mut buf, body_api_version)?;
+        let encoded_size = buf.len();
 
         let (tx, rx) = channel();
 
@@ -412,7 +429,10 @@ where
             });
         }
 
-        Ok(body)
+        Ok(ResponseBodyWithMetadata {
+            response: body,
+            encoded_request_size: encoded_size,
+        })
     }
 
     async fn send_message(&self, msg: Vec<u8>) -> Result<(), RequestError> {
@@ -468,7 +488,7 @@ where
                     .request_with_version_ranges(&body, &version_ranges)
                     .await
                 {
-                    Ok(response) => {
+                    Ok(ResponseBodyWithMetadata { response, .. }) => {
                         if let Err(ErrorOrThrottle::Throttle(throttle)) =
                             maybe_throttle::<SyncVersionsError>(response.throttle_time_ms)
                         {

src/protocol/messages/mod.rs

@@ -86,6 +86,16 @@ impl<W: Write, T: WriteVersionedType<W>> WriteVersionedType<W> for &T {
     }
 }
 
+/// A response body with metadata about the request.
+#[derive(Debug)]
+pub struct ResponseBodyWithMetadata<R> {
+    /// The response body.
+    pub response: R,
+
+    /// The size of the request encoded in bytes.
+    pub encoded_request_size: usize,
+}
+
 /// Specifies a request body.
 pub trait RequestBody {
     /// The response type that will follow when issuing this request.