chore: apply suggestions

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

Author: WenyXu

src/client/controller.rs

@@ -71,7 +71,8 @@ impl ControllerClient {
             let response = broker
                 .request(request)
                 .await
-                .map_err(|e| ErrorOrThrottle::Error((e.into(), Some(r#gen))))?;
+                .map_err(|e| ErrorOrThrottle::Error((e.into(), Some(r#gen))))?
+                .response;
 
             maybe_throttle(response.throttle_time_ms)?;
 
@@ -122,7 +123,8 @@ impl ControllerClient {
             let response = broker
                 .request(request)
                 .await
-                .map_err(|e| ErrorOrThrottle::Error((e.into(), Some(r#gen))))?;
+                .map_err(|e| ErrorOrThrottle::Error((e.into(), Some(r#gen))))?
+                .response;
 
             maybe_throttle(response.throttle_time_ms)?;
 

src/client/partition.rs

@@ -222,7 +222,7 @@ impl PartitionClient {
                     response,
                     encoded_request_size,
                 } = broker
-                    .request_with_metadata(&request)
+                    .request(&request)
                     .await
                     .map_err(|e| ErrorOrThrottle::Error((e.into(), Some(r#gen))))?;
                 maybe_throttle(response.throttle_time_ms)?;
@@ -267,7 +267,8 @@ impl PartitionClient {
                 let response = broker
                     .request(&request)
                     .await
-                    .map_err(|e| ErrorOrThrottle::Error((e.into(), Some(r#gen))))?;
+                    .map_err(|e| ErrorOrThrottle::Error((e.into(), Some(r#gen))))?
+                    .response;
                 maybe_throttle(response.throttle_time_ms)?;
                 process_fetch_response(self.partition, &self.topic, response, offset)
                     .map_err(|e| ErrorOrThrottle::Error((e, Some(r#gen))))
@@ -303,7 +304,8 @@ impl PartitionClient {
                 let response = broker
                     .request(&request)
                     .await
-                    .map_err(|e| ErrorOrThrottle::Error((e.into(), Some(r#gen))))?;
+                    .map_err(|e| ErrorOrThrottle::Error((e.into(), Some(r#gen))))?
+                    .response;
                 maybe_throttle(response.throttle_time_ms)?;
                 process_list_offsets_response(self.partition, &self.topic, response)
                     .map_err(|e| ErrorOrThrottle::Error((e, Some(r#gen))))
@@ -336,7 +338,8 @@ impl PartitionClient {
                 let response = broker
                     .request(&request)
                     .await
-                    .map_err(|e| ErrorOrThrottle::Error((e.into(), Some(r#gen))))?;
+                    .map_err(|e| ErrorOrThrottle::Error((e.into(), Some(r#gen))))?
+                    .response;
                 maybe_throttle(Some(response.throttle_time_ms))?;
                 process_delete_records_response(&self.topic, self.partition, response)
                     .map_err(|e| ErrorOrThrottle::Error((e, Some(r#gen))))

src/client/producer.rs

@@ -300,8 +300,9 @@ impl BatchProducerBuilder {
     }
 }
 
-/// The result of the produce call.
+/// The result of the [produce](ProducerClient::produce) call.
 #[derive(Debug, Default)]
+#[non_exhaustive]
 pub struct ProduceResult {
     /// The offsets of the produced records.
     pub offsets: Vec<i64>,

src/connection.rs

@@ -392,7 +392,7 @@ impl RequestHandler for MessengerTransport {
         &self,
         request_params: &MetadataRequest,
     ) -> Result<MetadataResponse, RequestError> {
-        self.request(request_params).await
+        self.request(request_params).await.map(|r| r.response)
     }
 }
 

src/messenger.rs

@@ -317,17 +317,7 @@ where
         self.version_ranges = ranges;
     }
 
-    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>(
+    pub async fn request<R>(
         &self,
         msg: R,
     ) -> Result<ResponseBodyWithMetadata<R::ResponseBody>, RequestError>
@@ -580,7 +570,7 @@ where
         auth_bytes: Vec<u8>,
     ) -> Result<SaslAuthenticateResponse, SaslError> {
         let req = SaslAuthenticateRequest::new(auth_bytes);
-        let resp = self.request(req).await?;
+        let resp = self.request(req).await?.response;
         if let Some(err) = resp.error_code {
             if let Some(s) = resp.error_message.0 {
                 debug!("Sasl auth error message: {s}");
@@ -593,7 +583,7 @@ where
 
     async fn sasl_handshake(&self, mechanism: &str) -> Result<SaslHandshakeResponse, SaslError> {
         let req = SaslHandshakeRequest::new(mechanism);
-        let resp = self.request(req).await?;
+        let resp = self.request(req).await?.response;
         if let Some(err) = resp.error_code {
             return Err(SaslError::ApiError(err));
         }
@@ -1204,7 +1194,8 @@ mod tests {
                 tagged_fields: Some(TaggedFields::default()),
             })
             .await
-            .unwrap();
+            .unwrap()
+            .response;
         assert_eq!(actual, resp);
     }
 

src/protocol/messages/mod.rs

@@ -86,7 +86,7 @@ impl<W: Write, T: WriteVersionedType<W>> WriteVersionedType<W> for &T {
     }
 }
 
-/// A response body with metadata about the request.
+/// A response body with metadata about the request & response.
 #[derive(Debug)]
 pub struct ResponseBodyWithMetadata<R> {
     /// The response body.

tests/client.rs

@@ -6,6 +6,7 @@ use rskafka::{
         ClientBuilder,
         error::{Error as ClientError, ProtocolError, ServerErrorResponse},
         partition::{Compression, OffsetAt, UnknownTopicHandling},
+        producer::ProduceResult,
     },
     record::{Record, RecordAndOffset},
 };
@@ -390,6 +391,55 @@ async fn test_consume_offset_out_of_range() {
     );
 }
 
+#[tokio::test]
+async fn test_produce_metadata() {
+    maybe_start_logging();
+
+    let test_cfg = maybe_skip_kafka_integration!();
+    let topic_name = random_topic_name();
+    let n_partitions = 1;
+
+    let client = ClientBuilder::new(test_cfg.bootstrap_brokers.clone())
+        .build()
+        .await
+        .unwrap();
+    let controller_client = client.controller_client().unwrap();
+    controller_client
+        .create_topic(&topic_name, n_partitions, 1, 5_000)
+        .await
+        .unwrap();
+
+    let partition_client = client
+        .partition_client(topic_name.clone(), 0, UnknownTopicHandling::Retry)
+        .await
+        .unwrap();
+
+    let record1 = record(b"");
+    let ProduceResult {
+        offsets,
+        encoded_request_size,
+        ..
+    } = partition_client
+        .produce(vec![record1.clone()], Compression::NoCompression)
+        .await
+        .unwrap();
+    assert_eq!(offsets[0], 0);
+    assert_eq!(encoded_request_size, 177);
+
+    // Produce a record with repeated string
+    let record2 = record(b"value value value value value value value");
+    let ProduceResult {
+        offsets,
+        encoded_request_size,
+        ..
+    } = partition_client
+        .produce(vec![record2.clone()], Compression::NoCompression)
+        .await
+        .unwrap();
+    assert_eq!(offsets[0], 1);
+    assert_eq!(encoded_request_size, 219);
+}
+
 #[tokio::test]
 async fn test_get_offset() {
     maybe_start_logging();