rocketmq-client-rust 1.0.0

// Copyright 2023 The RocketMQ Rust Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicI64;
use std::sync::atomic::Ordering;
use std::sync::Arc;

use rocketmq_model::common::message::message_queue::MessageQueue;
use tokio::sync::RwLock;

use crate::consumer::consumer_impl::process_queue::ProcessQueue;

/// Represents a message queue with its associated state for lite pull consumer.
/// Contains process queue reference, pause flag, and offset tracking for pull, consume, and seek
/// operations.
#[derive(Clone)]
pub(crate) struct AssignedQueue {
    pub(crate) message_queue: MessageQueue,
    pub(crate) process_queue: Arc<ProcessQueue>,
    pub(crate) paused: Arc<AtomicBool>,
    pub(crate) pull_offset: Arc<AtomicI64>,
    pub(crate) consume_offset: Arc<AtomicI64>,
    pub(crate) seek_offset: Arc<AtomicI64>,
}

impl AssignedQueue {
    pub fn new(message_queue: MessageQueue) -> Self {
        Self {
            message_queue,
            process_queue: Arc::new(ProcessQueue::new()),
            paused: Arc::new(AtomicBool::new(false)),
            pull_offset: Arc::new(AtomicI64::new(-1)),
            consume_offset: Arc::new(AtomicI64::new(-1)),
            seek_offset: Arc::new(AtomicI64::new(-1)),
        }
    }
}

/// Manages assigned message queues and their process queues for lite pull consumer.
#[derive(Clone)]
pub struct AssignedMessageQueue {
    queue_map: Arc<RwLock<HashMap<MessageQueue, AssignedQueue>>>,
}

impl AssignedMessageQueue {
    /// Creates a new empty AssignedMessageQueue.
    pub fn new() -> Self {
        Self {
            queue_map: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Assigns a new message queue.
    pub async fn put(&self, mq: MessageQueue) {
        let mut map = self.queue_map.write().await;
        map.entry(mq.clone()).or_insert_with(|| AssignedQueue::new(mq));
    }

    /// Updates all assigned message queues, dropping queues that are no longer assigned.
    pub async fn update_assigned_message_queue<I>(&self, assigned: I)
    where
        I: IntoIterator<Item = MessageQueue>,
    {
        let assigned: HashSet<MessageQueue> = assigned.into_iter().collect();
        let mut map = self.queue_map.write().await;

        map.retain(|mq, aq| {
            if assigned.contains(mq) {
                true
            } else {
                aq.process_queue.set_dropped(true);
                false
            }
        });

        for mq in assigned {
            map.entry(mq.clone()).or_insert_with(|| AssignedQueue::new(mq));
        }
    }

    /// Updates assigned queues for one topic, leaving assignments of other topics untouched.
    pub async fn update_assigned_message_queue_for_topic<I>(&self, topic: &str, assigned: I)
    where
        I: IntoIterator<Item = MessageQueue>,
    {
        let assigned: HashSet<MessageQueue> = assigned.into_iter().collect();
        let mut map = self.queue_map.write().await;

        map.retain(|mq, aq| {
            if mq.topic() == topic && !assigned.contains(mq) {
                aq.process_queue.set_dropped(true);
                false
            } else {
                true
            }
        });

        for mq in assigned {
            map.entry(mq.clone()).or_insert_with(|| AssignedQueue::new(mq));
        }
    }

    /// Removes a message queue and returns its process queue.
    pub async fn remove(&self, mq: &MessageQueue) -> Option<Arc<ProcessQueue>> {
        let mut map = self.queue_map.write().await;
        map.remove(mq).map(|aq| {
            aq.process_queue.set_dropped(true);
            aq.process_queue
        })
    }

    /// Removes all message queues for the specified topic and returns their process queues.
    pub async fn remove_by_topic(&self, topic: &str) -> Vec<Arc<ProcessQueue>> {
        let mut map = self.queue_map.write().await;
        let mut removed_pqs = Vec::new();
        map.retain(|mq, aq| {
            if mq.topic() == topic {
                aq.process_queue.set_dropped(true);
                removed_pqs.push(aq.process_queue.clone());
                false
            } else {
                true
            }
        });
        removed_pqs
    }

    /// Java-compatible alias for removing all assigned queues of a topic.
    pub async fn remove_assigned_message_queue(&self, topic: &str) -> Vec<Arc<ProcessQueue>> {
        self.remove_by_topic(topic).await
    }

    /// Returns the process queue for the specified message queue.
    pub async fn get_process_queue(&self, mq: &MessageQueue) -> Option<Arc<ProcessQueue>> {
        let map = self.queue_map.read().await;
        map.get(mq).map(|aq| aq.process_queue.clone())
    }

    /// Checks if the specified queue is paused.
    pub async fn is_paused(&self, mq: &MessageQueue) -> bool {
        let map = self.queue_map.read().await;
        map.get(mq).map(|aq| aq.paused.load(Ordering::Acquire)).unwrap_or(true)
    }

    /// Sets the pause state for the specified queue.
    pub async fn set_paused(&self, mq: &MessageQueue, paused: bool) {
        let map = self.queue_map.read().await;
        if let Some(aq) = map.get(mq) {
            aq.paused.store(paused, Ordering::Release);
        }
    }

    /// Returns the seek offset for the specified queue (-1 if not seeking).
    pub async fn get_seek_offset(&self, mq: &MessageQueue) -> i64 {
        let map = self.queue_map.read().await;
        map.get(mq)
            .map(|aq| aq.seek_offset.load(Ordering::Acquire))
            .unwrap_or(-1)
    }

    /// Sets the seek offset for the specified queue.
    pub async fn set_seek_offset(&self, mq: &MessageQueue, offset: i64) {
        let map = self.queue_map.read().await;
        if let Some(aq) = map.get(mq) {
            aq.seek_offset.store(offset, Ordering::Release);
        }
    }

    /// Clears the seek offset for the specified queue (sets to -1).
    pub async fn clear_seek_offset(&self, mq: &MessageQueue) {
        self.set_seek_offset(mq, -1).await;
    }

    /// Returns the pull offset for the specified queue.
    pub async fn get_pull_offset(&self, mq: &MessageQueue) -> i64 {
        let map = self.queue_map.read().await;
        map.get(mq)
            .map(|aq| aq.pull_offset.load(Ordering::Acquire))
            .unwrap_or(-1)
    }

    /// Updates the pull offset for the specified queue.
    pub async fn update_pull_offset(&self, mq: &MessageQueue, offset: i64, process_queue: &Arc<ProcessQueue>) {
        let map = self.queue_map.read().await;
        if let Some(aq) = map.get(mq) {
            // Only update if the process queue instance matches to prevent race conditions
            if Arc::ptr_eq(&aq.process_queue, process_queue) {
                aq.pull_offset.store(offset, Ordering::Release);
            }
        }
    }

    /// Returns the consume offset for the specified queue.
    pub async fn get_consume_offset(&self, mq: &MessageQueue) -> i64 {
        let map = self.queue_map.read().await;
        map.get(mq)
            .map(|aq| aq.consume_offset.load(Ordering::Acquire))
            .unwrap_or(-1)
    }

    /// Updates the consume offset for the specified queue.
    pub async fn update_consume_offset(&self, mq: &MessageQueue, offset: i64) {
        let map = self.queue_map.read().await;
        if let Some(aq) = map.get(mq) {
            aq.consume_offset.store(offset, Ordering::Release);
        }
    }

    /// Returns all assigned message queues.
    pub async fn message_queues(&self) -> HashSet<MessageQueue> {
        let map = self.queue_map.read().await;
        map.keys().cloned().collect()
    }

    /// Java-compatible alias for all assigned message queues.
    pub async fn assigned_message_queues(&self) -> HashSet<MessageQueue> {
        self.message_queues().await
    }

    /// Returns the number of assigned queues.
    pub async fn size(&self) -> usize {
        let map = self.queue_map.read().await;
        map.len()
    }

    /// Clears all assigned queues.
    pub async fn clear(&self) {
        let mut map = self.queue_map.write().await;
        for aq in map.values() {
            aq.process_queue.set_dropped(true);
        }
        map.clear();
    }

    /// Returns the total cached message count across all queues.
    pub async fn total_msg_count(&self) -> i64 {
        let map = self.queue_map.read().await;
        map.values().map(|aq| aq.process_queue.msg_count() as i64).sum()
    }

    /// Returns the total cached message size in MiB across all queues.
    pub async fn total_msg_size_in_mib(&self) -> i64 {
        let map = self.queue_map.read().await;
        map.values()
            .map(|aq| (aq.process_queue.msg_size() / (1024 * 1024)) as i64)
            .sum()
    }
}

impl Default for AssignedMessageQueue {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rocketmq_model::common::message::message_queue::MessageQueue;

    fn mock_mq(topic: &str, queue_id: i32) -> MessageQueue {
        MessageQueue::from_parts(topic, "mock_broker", queue_id)
    }

    #[tokio::test]
    async fn test_put_and_size() {
        let amq = AssignedMessageQueue::new();
        let mq1 = mock_mq("topic_a", 0);
        let mq2 = mock_mq("topic_a", 1);

        assert_eq!(amq.size().await, 0);

        amq.put(mq1.clone()).await;
        assert_eq!(amq.size().await, 1);

        amq.put(mq2.clone()).await;
        assert_eq!(amq.size().await, 2);

        amq.put(mq1.clone()).await;
        assert_eq!(amq.size().await, 2);

        let queues = amq.message_queues().await;
        assert!(queues.contains(&mq1));
        assert!(queues.contains(&mq2));
    }

    #[tokio::test]
    async fn test_remove() {
        let amq = AssignedMessageQueue::new();
        let mq1 = mock_mq("topic_a", 0);

        amq.put(mq1.clone()).await;
        assert_eq!(amq.size().await, 1);

        let removed_pq = amq.remove(&mq1).await;
        assert!(removed_pq.is_some());
        assert_eq!(amq.size().await, 0);

        let mq2 = mock_mq("topic_b", 0);
        let removed_none = amq.remove(&mq2).await;
        assert!(removed_none.is_none());
    }

    #[tokio::test]
    async fn test_remove_by_topic() {
        let amq = AssignedMessageQueue::new();
        let mq1 = mock_mq("topic_a", 0);
        let mq2 = mock_mq("topic_a", 1);
        let mq3 = mock_mq("topic_b", 0);

        amq.put(mq1.clone()).await;
        amq.put(mq2.clone()).await;
        amq.put(mq3.clone()).await;
        assert_eq!(amq.size().await, 3);

        let removed_pqs = amq.remove_by_topic("topic_a").await;
        assert_eq!(removed_pqs.len(), 2);
        assert_eq!(amq.size().await, 1);
    }

    #[tokio::test]
    async fn update_assigned_message_queue_removes_unassigned_queues_like_java_assign() {
        let amq = AssignedMessageQueue::new();
        let retained = mock_mq("topic_a", 0);
        let removed = mock_mq("topic_a", 1);

        amq.put(retained.clone()).await;
        amq.put(removed.clone()).await;
        let removed_pq = amq.get_process_queue(&removed).await.unwrap();

        amq.update_assigned_message_queue(vec![retained.clone()]).await;

        let queues = amq.assigned_message_queues().await;
        assert_eq!(queues.len(), 1);
        assert!(queues.contains(&retained));
        assert!(amq.get_process_queue(&removed).await.is_none());
        assert!(removed_pq.is_dropped());
    }

    #[tokio::test]
    async fn update_assigned_message_queue_for_topic_keeps_other_topics() {
        let amq = AssignedMessageQueue::new();
        let retained = mock_mq("topic_a", 0);
        let removed_same_topic = mock_mq("topic_a", 1);
        let other_topic = mock_mq("topic_b", 0);

        amq.put(retained.clone()).await;
        amq.put(removed_same_topic.clone()).await;
        amq.put(other_topic.clone()).await;
        let removed_pq = amq.get_process_queue(&removed_same_topic).await.unwrap();

        amq.update_assigned_message_queue_for_topic("topic_a", vec![retained.clone()])
            .await;

        let queues = amq.message_queues().await;
        assert_eq!(queues.len(), 2);
        assert!(queues.contains(&retained));
        assert!(queues.contains(&other_topic));
        assert!(!queues.contains(&removed_same_topic));
        assert!(removed_pq.is_dropped());
    }

    #[tokio::test]
    async fn test_pause_and_resume() {
        let amq = AssignedMessageQueue::new();
        let mq = mock_mq("topic_a", 0);
        let unassigned_mq = mock_mq("topic_a", 1);

        amq.put(mq.clone()).await;

        assert!(!amq.is_paused(&mq).await);
        assert!(amq.is_paused(&unassigned_mq).await);

        amq.set_paused(&mq, true).await;
        assert!(amq.is_paused(&mq).await);

        amq.set_paused(&mq, false).await;
        assert!(!amq.is_paused(&mq).await);
    }

    #[tokio::test]
    async fn test_seek_offset() {
        let amq = AssignedMessageQueue::new();
        let mq = mock_mq("topic_a", 0);

        amq.put(mq.clone()).await;

        assert_eq!(amq.get_seek_offset(&mq).await, -1);

        amq.set_seek_offset(&mq, 1024).await;
        assert_eq!(amq.get_seek_offset(&mq).await, 1024);

        amq.clear_seek_offset(&mq).await;
        assert_eq!(amq.get_seek_offset(&mq).await, -1);
    }

    #[tokio::test]
    async fn test_pull_and_consume_offset() {
        let amq = AssignedMessageQueue::new();
        let mq = mock_mq("topic_a", 0);

        amq.put(mq.clone()).await;

        let pq = amq.get_process_queue(&mq).await.unwrap();

        assert_eq!(amq.get_pull_offset(&mq).await, -1);
        amq.update_pull_offset(&mq, 500, &pq).await;
        assert_eq!(amq.get_pull_offset(&mq).await, 500);

        assert_eq!(amq.get_consume_offset(&mq).await, -1);
        amq.update_consume_offset(&mq, 300).await;
        assert_eq!(amq.get_consume_offset(&mq).await, 300);
    }

    #[tokio::test]
    async fn test_clear() {
        let amq = AssignedMessageQueue::new();
        amq.put(mock_mq("topic_a", 0)).await;
        amq.put(mock_mq("topic_b", 1)).await;

        assert_eq!(amq.size().await, 2);

        amq.clear().await;
        assert_eq!(amq.size().await, 0);
    }
}