feat(common): add deferred poll timeout support

core/binary_protocol/src/requests/messages/poll_messages.rs

@@ -17,7 +17,7 @@
 
 use crate::WireError;
 use crate::WireIdentifier;
-use crate::codec::{WireDecode, WireEncode, read_u8, read_u32_le};
+use crate::codec::{WireDecode, WireEncode, read_u8, read_u32_le, read_u64_le};
 use crate::primitives::consumer::WireConsumer;
 use crate::primitives::polling_strategy::WirePollingStrategy;
 use bytes::{BufMut, BytesMut};
@@ -27,7 +27,7 @@ use bytes::{BufMut, BytesMut};
 /// Wire format:
 /// ```text
 /// [consumer][stream_id][topic_id][partition_flag:1][partition_id:4 LE]
-/// [strategy:9][count:4 LE][auto_commit:1]
+/// [strategy:9][count:4 LE][auto_commit:1][wait_timeout_us:8 LE]
 /// ```
 ///
 /// `partition_id` encoding: a u8 flag (1=Some, 0=None) followed by 4 bytes
@@ -41,13 +41,15 @@ pub struct PollMessagesRequest {
     pub strategy: WirePollingStrategy,
     pub count: u32,
     pub auto_commit: bool,
+    pub wait_timeout_us: u64,
 }
 
 const PARTITION_FLAG_SIZE: usize = 1;
 const PARTITION_VALUE_SIZE: usize = 4;
 const STRATEGY_SIZE: usize = 9;
 const COUNT_SIZE: usize = 4;
 const AUTO_COMMIT_SIZE: usize = 1;
+const WAIT_TIMEOUT_SIZE: usize = 8;
 
 impl WireEncode for PollMessagesRequest {
     fn encoded_size(&self) -> usize {
@@ -59,6 +61,7 @@ impl WireEncode for PollMessagesRequest {
             + STRATEGY_SIZE
             + COUNT_SIZE
             + AUTO_COMMIT_SIZE
+            + WAIT_TIMEOUT_SIZE
     }
 
     fn encode(&self, buf: &mut BytesMut) {
@@ -75,6 +78,7 @@ impl WireEncode for PollMessagesRequest {
         self.strategy.encode(buf);
         buf.put_u32_le(self.count);
         buf.put_u8(u8::from(self.auto_commit));
+        buf.put_u64_le(self.wait_timeout_us);
     }
 }
 
@@ -104,6 +108,13 @@ impl WireDecode for PollMessagesRequest {
         pos += 4;
         let auto_commit = read_u8(buf, pos)? != 0;
         pos += 1;
+        let wait_timeout_us = if buf.len() == pos {
+            0
+        } else {
+            let wait_timeout_us = read_u64_le(buf, pos)?;
+            pos += WAIT_TIMEOUT_SIZE;
+            wait_timeout_us
+        };
 
         Ok((
             Self {
@@ -114,6 +125,7 @@ impl WireDecode for PollMessagesRequest {
                 strategy,
                 count,
                 auto_commit,
+                wait_timeout_us,
             },
             pos,
         ))
@@ -124,17 +136,31 @@ impl WireDecode for PollMessagesRequest {
 mod tests {
     use super::*;
 
-    #[test]
-    fn roundtrip_with_partition() {
-        let req = PollMessagesRequest {
+    fn request(wait_timeout_us: u64) -> PollMessagesRequest {
+        PollMessagesRequest {
             consumer: WireConsumer::consumer(WireIdentifier::numeric(1)),
             stream_id: WireIdentifier::numeric(10),
             topic_id: WireIdentifier::numeric(20),
             partition_id: Some(5),
             strategy: WirePollingStrategy::offset(100),
             count: 50,
             auto_commit: true,
-        };
+            wait_timeout_us,
+        }
+    }
+
+    #[test]
+    fn roundtrip_with_zero_wait_timeout() {
+        let req = request(0);
+        let bytes = req.to_bytes();
+        let (decoded, consumed) = PollMessagesRequest::decode(&bytes).unwrap();
+        assert_eq!(consumed, bytes.len());
+        assert_eq!(decoded, req);
+    }
+
+    #[test]
+    fn roundtrip_with_non_zero_wait_timeout() {
+        let req = request(250_000);
         let bytes = req.to_bytes();
         let (decoded, consumed) = PollMessagesRequest::decode(&bytes).unwrap();
         assert_eq!(consumed, bytes.len());
@@ -151,6 +177,7 @@ mod tests {
             strategy: WirePollingStrategy::first(),
             count: 10,
             auto_commit: false,
+            wait_timeout_us: 0,
         };
         let bytes = req.to_bytes();
         let (decoded, consumed) = PollMessagesRequest::decode(&bytes).unwrap();
@@ -168,13 +195,62 @@ mod tests {
             strategy: WirePollingStrategy::offset(0),
             count: 1,
             auto_commit: false,
+            wait_timeout_us: 1,
         };
         let bytes = req.to_bytes();
         let (decoded, consumed) = PollMessagesRequest::decode(&bytes).unwrap();
         assert_eq!(consumed, bytes.len());
         assert_eq!(decoded, req);
     }
 
+    #[test]
+    fn legacy_request_without_wait_timeout_decodes_as_zero() {
+        let req = request(0);
+        let bytes = req.to_bytes();
+        let legacy_len = bytes.len() - WAIT_TIMEOUT_SIZE;
+        let (decoded, consumed) = PollMessagesRequest::decode(&bytes[..legacy_len]).unwrap();
+
+        assert_eq!(consumed, legacy_len);
+        assert_eq!(decoded.wait_timeout_us, 0);
+        assert_eq!(decoded, req);
+    }
+
+    #[test]
+    fn partial_trailing_wait_timeout_returns_error() {
+        let req = request(0);
+        let bytes = req.to_bytes();
+        let legacy_len = bytes.len() - WAIT_TIMEOUT_SIZE;
+
+        for timeout_bytes in 1..WAIT_TIMEOUT_SIZE {
+            assert!(
+                PollMessagesRequest::decode(&bytes[..legacy_len + timeout_bytes]).is_err(),
+                "expected error with {timeout_bytes} trailing timeout bytes"
+            );
+        }
+    }
+
+    #[test]
+    fn encoded_size_includes_wait_timeout() {
+        let req = request(42_000);
+        let bytes = req.to_bytes();
+        let legacy_size = req.consumer.encoded_size()
+            + req.stream_id.encoded_size()
+            + req.topic_id.encoded_size()
+            + PARTITION_FLAG_SIZE
+            + PARTITION_VALUE_SIZE
+            + STRATEGY_SIZE
+            + COUNT_SIZE
+            + AUTO_COMMIT_SIZE;
+        let wait_timeout_bytes = req.wait_timeout_us.to_le_bytes();
+
+        assert_eq!(req.encoded_size(), legacy_size + WAIT_TIMEOUT_SIZE);
+        assert_eq!(bytes.len(), req.encoded_size());
+        assert_eq!(
+            &bytes[legacy_size..legacy_size + WAIT_TIMEOUT_SIZE],
+            wait_timeout_bytes.as_slice()
+        );
+    }
+
     #[test]
     fn partition_none_encodes_zero_bytes() {
         let req = PollMessagesRequest {
@@ -185,6 +261,7 @@ mod tests {
             strategy: WirePollingStrategy::first(),
             count: 1,
             auto_commit: false,
+            wait_timeout_us: 0,
         };
         let bytes = req.to_bytes();
         // After consumer(7) + stream_id(6) + topic_id(6) = offset 19
@@ -199,18 +276,11 @@ mod tests {
     }
 
     #[test]
-    fn truncated_returns_error() {
-        let req = PollMessagesRequest {
-            consumer: WireConsumer::consumer(WireIdentifier::numeric(1)),
-            stream_id: WireIdentifier::numeric(1),
-            topic_id: WireIdentifier::numeric(1),
-            partition_id: Some(1),
-            strategy: WirePollingStrategy::offset(0),
-            count: 1,
-            auto_commit: false,
-        };
+    fn truncated_required_fields_return_error() {
+        let req = request(0);
         let bytes = req.to_bytes();
-        for i in 0..bytes.len() {
+        let legacy_len = bytes.len() - WAIT_TIMEOUT_SIZE;
+        for i in 0..legacy_len {
             assert!(
                 PollMessagesRequest::decode(&bytes[..i]).is_err(),
                 "expected error for truncation at byte {i}"

core/common/src/traits/binary_impls/messages.rs

@@ -37,11 +37,27 @@ use iggy_binary_protocol::requests::messages::{
     FlushUnsavedBufferRequest, PollMessagesRequest, RawMessage, SendMessagesEncoder,
 };
 use iggy_binary_protocol::responses::consumer_groups::SyncConsumerGroupResponse;
+use std::time::Duration;
 
 /// Max attempts to resolve a fenced consumer-group poll: one re-sync after the
 /// coordinator rejects a stale assignment, then retry once.
 const GROUP_POLL_MAX_ATTEMPTS: usize = 2;
 
+struct PollGroupOptions<'a> {
+    consumer: &'a Consumer,
+    strategy: &'a PollingStrategy,
+    count: u32,
+    auto_commit: bool,
+    wait_timeout_us: u64,
+}
+
+fn duration_to_wait_timeout_us(wait_timeout: Duration) -> Result<u64, IggyError> {
+    wait_timeout
+        .as_micros()
+        .try_into()
+        .map_err(|_| IggyError::InvalidNumberValue)
+}
+
 fn group_cache_key(stream_id: &Identifier, topic_id: &Identifier, group_id: &Identifier) -> String {
     format!("{stream_id}|{topic_id}|{group_id}")
 }
@@ -172,11 +188,15 @@ async fn poll_group_messages<B: BinaryClient>(
     client: &B,
     stream_id: &Identifier,
     topic_id: &Identifier,
-    consumer: &Consumer,
-    strategy: &PollingStrategy,
-    count: u32,
-    auto_commit: bool,
+    options: PollGroupOptions<'_>,
 ) -> Result<PolledMessages, IggyError> {
+    let PollGroupOptions {
+        consumer,
+        strategy,
+        count,
+        auto_commit,
+        wait_timeout_us,
+    } = options;
     let key = group_cache_key(stream_id, topic_id, &consumer.id);
     if !client.consumer_group_state().has_assignment(&key) {
         sync_group_assignment(client, stream_id, topic_id, &consumer.id).await?;
@@ -205,6 +225,7 @@ async fn poll_group_messages<B: BinaryClient>(
             strategy: polling_strategy_to_wire(strategy),
             count,
             auto_commit,
+            wait_timeout_us,
         };
         match client
             .send_raw_with_response(POLL_MESSAGES_CODE, request.to_bytes())
@@ -283,8 +304,33 @@ impl<B: BinaryClient> MessageClient for B {
         strategy: &PollingStrategy,
         count: u32,
         auto_commit: bool,
+    ) -> Result<PolledMessages, IggyError> {
+        self.poll_messages_with_timeout(
+            stream_id,
+            topic_id,
+            partition_id,
+            consumer,
+            strategy,
+            count,
+            auto_commit,
+            Duration::ZERO,
+        )
+        .await
+    }
+
+    async fn poll_messages_with_timeout(
+        &self,
+        stream_id: &Identifier,
+        topic_id: &Identifier,
+        partition_id: Option<u32>,
+        consumer: &Consumer,
+        strategy: &PollingStrategy,
+        count: u32,
+        auto_commit: bool,
+        wait_timeout: Duration,
     ) -> Result<PolledMessages, IggyError> {
         fail_if_not_authenticated(self).await?;
+        let wait_timeout_us = duration_to_wait_timeout_us(wait_timeout)?;
         // VSR: a consumer-group poll without an explicit partition is resolved
         // client-side from the member's cached assignment (the broker routes
         // explicit partitions only).
@@ -293,10 +339,13 @@ impl<B: BinaryClient> MessageClient for B {
                 self,
                 stream_id,
                 topic_id,
-                consumer,
-                strategy,
-                count,
-                auto_commit,
+                PollGroupOptions {
+                    consumer,
+                    strategy,
+                    count,
+                    auto_commit,
+                    wait_timeout_us,
+                },
             )
             .await;
         }
@@ -308,6 +357,7 @@ impl<B: BinaryClient> MessageClient for B {
             strategy: polling_strategy_to_wire(strategy),
             count,
             auto_commit,
+            wait_timeout_us,
         };
         let response = self
             .send_raw_with_response(POLL_MESSAGES_CODE, req.to_bytes())
@@ -402,9 +452,12 @@ impl<B: BinaryClient> MessageClient for B {
 
 #[cfg(test)]
 mod tests {
-    use super::{committed_send_confirmations, decode_send_confirmations};
+    use super::{
+        committed_send_confirmations, decode_send_confirmations, duration_to_wait_timeout_us,
+    };
     use crate::{IggyError, SendMessagesConfirmationResponse, SendMessagesResponse};
     use iggy_binary_protocol::codec::WireEncode;
+    use std::time::Duration;
 
     fn response() -> SendMessagesResponse {
         SendMessagesResponse {
@@ -494,4 +547,24 @@ mod tests {
             );
         }
     }
+
+    #[test]
+    fn wait_timeout_uses_microseconds() {
+        assert_eq!(
+            duration_to_wait_timeout_us(Duration::from_millis(25)).unwrap(),
+            25_000
+        );
+        assert_eq!(
+            duration_to_wait_timeout_us(Duration::from_nanos(999)).unwrap(),
+            0
+        );
+    }
+
+    #[test]
+    fn wait_timeout_overflow_is_rejected() {
+        assert_eq!(
+            duration_to_wait_timeout_us(Duration::new(u64::MAX, 0)),
+            Err(IggyError::InvalidNumberValue)
+        );
+    }
 }