Merge pull request #223 from YilunAllenChen/list-offsets

list offsets

Author: johnward

src/client/mod.rs

@@ -15,7 +15,9 @@ use std::time::{Duration, Instant};
 
 // pub re-export
 pub use crate::compression::Compression;
+use crate::protocol::list_offset::ListOffsetVersion;
 pub use crate::utils::PartitionOffset;
+use crate::utils::TimestampedPartitionOffset;
 
 #[cfg(feature = "security")]
 pub use self::network::SecurityConfig;
@@ -930,7 +932,111 @@ impl KafkaClient {
         Ok(res)
     }
 
-    /// Takes ownership back from the given `HashMap` Entry.
+    /// Fetch offsets for a list of topics
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// use kafka::client::KafkaClient;
+    ///
+    /// let mut client = KafkaClient::new(vec!["localhost:9092".to_owned()]);
+    /// client.load_metadata_all().unwrap();
+    /// let topics = vec!["test-topic".to_string()];
+    /// let offsets = client.list_offsets(&s, FetchOffset::ByTime(1698425676797));
+    /// ```
+    ///
+    /// Returns a mapping of topic name to `TimestampedPartitionOffset`s.
+    /// Each entry in the vector represents the timestamp, and the corresponding offset,
+    /// that Kafka finds to be the first message with timestamp *later* than the passed in
+    /// FetchOffset parameter.
+    /// example: Ok({"test-topic": [TimestampedPartitionOffset { offset: 20, partition: 0, time: 1698425676798 } ]
+    /// Note that the message might not be exactly at the given timestamp.
+    pub fn list_offsets<T: AsRef<str>>(
+        &mut self,
+        topics: &[T],
+        offset: FetchOffset,
+    ) -> Result<HashMap<String, Vec<TimestampedPartitionOffset>>> {
+        let api_ver = ListOffsetVersion::V1;
+        let time = offset.to_kafka_value();
+        let n_topics = topics.len();
+
+        let state = &mut self.state;
+        let correlation = state.next_correlation_id();
+
+        // Map topic and partition to the corresponding broker
+        let config = &self.config;
+        let mut reqs: HashMap<&str, protocol::ListOffsetsRequest<'_>> =
+            HashMap::with_capacity(n_topics);
+        for topic in topics {
+            let topic = topic.as_ref();
+            if let Some(ps) = state.partitions_for(topic) {
+                for (id, host) in ps
+                    .iter()
+                    .filter_map(|(id, p)| p.broker(state).map(|b| (id, b.host())))
+                {
+                    let entry = reqs.entry(host).or_insert_with(|| {
+                        protocol::ListOffsetsRequest::new(correlation, api_ver, &config.client_id)
+                    });
+                    entry.add(topic, id, time);
+                }
+            }
+        }
+
+        // Call each broker with the request formed earlier
+        let now = Instant::now();
+        let mut res: HashMap<String, Vec<TimestampedPartitionOffset>> =
+            HashMap::with_capacity(n_topics);
+        for (host, req) in reqs {
+            let resp = __send_receive::<_, protocol::ListOffsetsResponse>(
+                &mut self.conn_pool,
+                host,
+                now,
+                req,
+            )?;
+            for tp in resp.topics {
+                let mut entry = res.entry(tp.topic);
+                let mut new_resp_offsets = None;
+                let mut err = None;
+                // Use an explicit scope here to allow insertion into a vacant entry
+                // below
+                {
+                    let resp_offsets = match entry {
+                        hash_map::Entry::Occupied(ref mut e) => e.get_mut(),
+                        hash_map::Entry::Vacant(_) => {
+                            new_resp_offsets = Some(Vec::new());
+                            new_resp_offsets.as_mut().unwrap()
+                        }
+                    };
+                    for p in tp.partitions {
+                        let pto: TimestampedPartitionOffset = match p.to_offset() {
+                            Ok(po) => po,
+                            Err(code) => {
+                                err = Some((p.partition, code));
+                                break;
+                            }
+                        };
+                        resp_offsets.push(pto);
+                    }
+                }
+                if let Some((partition, code)) = err {
+                    let topic = KafkaClient::get_key_from_entry(entry);
+                    return Err(Error::TopicPartitionError {
+                        topic_name: topic,
+                        partition_id: partition,
+                        error_code: code,
+                    });
+                }
+                if let hash_map::Entry::Vacant(e) = entry {
+                    // unwrap is ok because if it is Vacant, it would have
+                    // been made into a Some above
+                    e.insert(new_resp_offsets.unwrap());
+                }
+            }
+        }
+        Ok(res)
+    }
+
+    /// Takes ownership back from the given HashMap Entry.
     fn get_key_from_entry<'a, K: 'a, V: 'a>(entry: hash_map::Entry<'a, K, V>) -> K {
         match entry {
             hash_map::Entry::Occupied(e) => e.remove_entry().0,

src/protocol/list_offset.rs

@@ -0,0 +1,168 @@
+use std::io::{Read, Write};
+
+use crate::codecs::{FromByte, ToByte};
+use crate::error::{KafkaCode, Result};
+use crate::utils::TimestampedPartitionOffset;
+
+use super::{HeaderRequest, HeaderResponse, API_KEY_OFFSET};
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub enum ListOffsetVersion {
+    // currently only support 1
+    V1 = 1,
+}
+
+
+/// https://kafka.apache.org/protocol.html#The_Messages_ListOffsets
+#[derive(Debug)]
+pub struct ListOffsetsRequest<'a> {
+    pub header: HeaderRequest<'a>,
+    pub replica: i32,
+    pub topics: Vec<TopicListOffsetsRequest<'a>>,
+}
+
+#[derive(Debug)]
+pub struct TopicListOffsetsRequest<'a> {
+    pub topic: &'a str,
+    pub partitions: Vec<PartitionListOffsetsRequest>,
+}
+
+#[derive(Default, Debug)]
+pub struct PartitionListOffsetsRequest {
+    pub partition: i32,
+    pub time: i64,
+}
+
+impl<'a> ListOffsetsRequest<'a> {
+    pub fn new(
+        correlation_id: i32,
+        version: ListOffsetVersion,
+        client_id: &'a str,
+    ) -> ListOffsetsRequest<'a> {
+        ListOffsetsRequest {
+            header: HeaderRequest::new(API_KEY_OFFSET, version as i16, correlation_id, client_id),
+            replica: -1,
+            topics: vec![],
+        }
+    }
+
+    pub fn add(&mut self, topic: &'a str, partition: i32, time: i64) {
+        for tp in &mut self.topics {
+            if tp.topic == topic {
+                tp.add(partition, time);
+                return;
+            }
+        }
+        let mut tp = TopicListOffsetsRequest::new(topic);
+        tp.add(partition, time);
+        self.topics.push(tp);
+    }
+}
+
+impl<'a> TopicListOffsetsRequest<'a> {
+    fn new(topic: &'a str) -> TopicListOffsetsRequest<'a> {
+        TopicListOffsetsRequest {
+            topic,
+            partitions: vec![],
+        }
+    }
+    fn add(&mut self, partition: i32, time: i64) {
+        self.partitions
+            .push(PartitionListOffsetsRequest::new(partition, time));
+    }
+}
+
+impl PartitionListOffsetsRequest {
+    fn new(partition: i32, time: i64) -> PartitionListOffsetsRequest {
+        PartitionListOffsetsRequest { partition, time }
+    }
+}
+
+impl<'a> ToByte for ListOffsetsRequest<'a> {
+    fn encode<T: Write>(&self, buffer: &mut T) -> Result<()> {
+        try_multi!(
+            self.header.encode(buffer),
+            self.replica.encode(buffer),
+            self.topics.encode(buffer)
+        )
+    }
+}
+
+impl<'a> ToByte for TopicListOffsetsRequest<'a> {
+    fn encode<T: Write>(&self, buffer: &mut T) -> Result<()> {
+        try_multi!(self.topic.encode(buffer), self.partitions.encode(buffer))
+    }
+}
+
+impl ToByte for PartitionListOffsetsRequest {
+    fn encode<T: Write>(&self, buffer: &mut T) -> Result<()> {
+        try_multi!(self.partition.encode(buffer), self.time.encode(buffer))
+    }
+}
+
+// -------------------------------------
+
+#[derive(Default, Debug)]
+pub struct ListOffsetsResponse {
+    pub header: HeaderResponse,
+    pub topics: Vec<TopicListOffsetsResponse>,
+}
+
+#[derive(Default, Debug)]
+pub struct TopicListOffsetsResponse {
+    pub topic: String,
+    pub partitions: Vec<TimestampedPartitionOffsetListOffsetsResponse>,
+}
+
+#[derive(Default, Debug)]
+pub struct TimestampedPartitionOffsetListOffsetsResponse {
+    pub partition: i32,
+    pub error_code: i16,
+    pub timestamp: i64,
+    pub offset: i64,
+}
+
+impl TimestampedPartitionOffsetListOffsetsResponse {
+    pub fn to_offset(&self) -> std::result::Result<TimestampedPartitionOffset, KafkaCode> {
+        match KafkaCode::from_protocol(self.error_code) {
+            Some(code) => Err(code),
+            None => Ok(TimestampedPartitionOffset {
+                partition: self.partition,
+                offset: self.offset,
+                time: self.timestamp,
+            }),
+        }
+    }
+}
+
+impl FromByte for ListOffsetsResponse {
+    type R = ListOffsetsResponse;
+
+    #[allow(unused_must_use)]
+    fn decode<T: Read>(&mut self, buffer: &mut T) -> Result<()> {
+        try_multi!(self.header.decode(buffer), self.topics.decode(buffer))
+    }
+}
+
+impl FromByte for TopicListOffsetsResponse {
+    type R = TopicListOffsetsResponse;
+
+    #[allow(unused_must_use)]
+    fn decode<T: Read>(&mut self, buffer: &mut T) -> Result<()> {
+        try_multi!(self.topic.decode(buffer), self.partitions.decode(buffer))
+    }
+}
+
+impl FromByte for TimestampedPartitionOffsetListOffsetsResponse {
+    type R = TimestampedPartitionOffsetListOffsetsResponse;
+
+    #[allow(unused_must_use)]
+    fn decode<T: Read>(&mut self, buffer: &mut T) -> Result<()> {
+        try_multi!(
+            self.partition.decode(buffer),
+            self.error_code.decode(buffer),
+            self.timestamp.decode(buffer),
+            self.offset.decode(buffer)
+        )
+    }
+}

src/protocol/mod.rs

@@ -18,6 +18,7 @@ pub mod consumer;
 pub mod metadata;
 pub mod offset;
 pub mod produce;
+pub mod list_offset;
 
 pub mod fetch;
 mod zreader;
@@ -32,6 +33,7 @@ pub use self::fetch::FetchRequest;
 pub use self::metadata::{MetadataRequest, MetadataResponse};
 pub use self::offset::{OffsetRequest, OffsetResponse};
 pub use self::produce::{ProduceRequest, ProduceResponse};
+pub use self::list_offset::{ListOffsetsRequest, ListOffsetsResponse};
 
 // --------------------------------------------------------------------
 

src/utils.rs

@@ -9,3 +9,12 @@ pub struct PartitionOffset {
     pub offset: i64,
     pub partition: i32,
 }
+
+/// A retrieved offset for a particular partition in the context of an
+/// already known topic, specific to a timestamp.
+#[derive(Debug, Hash, PartialEq, Eq)]
+pub struct TimestampedPartitionOffset {
+    pub offset: i64,
+    pub partition: i32,
+    pub time: i64
+}