use std::fmt::Display;
use std::hash::Hash;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicI32;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;
use std::sync::atomic::Ordering::Acquire;
use std::sync::atomic::Ordering::Relaxed;
use std::sync::Arc;
use rocketmq_protocol::protocol::body::pop_process_queue_info::PopProcessQueueInfo;
use rocketmq_runtime::common::time_utils::current_millis;
use crate::consumer::consumer_impl::PULL_MAX_IDLE_TIME;
#[derive(Clone)]
pub struct PopProcessQueue {
last_pop_timestamp: Arc<AtomicU64>,
wait_ack_counter: Arc<AtomicI32>,
dropped: Arc<AtomicBool>,
}
impl Hash for PopProcessQueue {
#[inline]
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
self.last_pop_timestamp.load(Relaxed).hash(state);
self.wait_ack_counter.load(Relaxed).hash(state);
self.dropped.load(Relaxed).hash(state);
}
}
impl Default for PopProcessQueue {
fn default() -> Self {
PopProcessQueue {
last_pop_timestamp: Arc::new(AtomicU64::new(current_millis())),
wait_ack_counter: Arc::new(AtomicI32::new(0)),
dropped: Arc::new(AtomicBool::new(false)),
}
}
}
impl PopProcessQueue {
#[inline]
pub fn new() -> Self {
PopProcessQueue::default()
}
#[inline]
pub fn get_last_pop_timestamp(&self) -> u64 {
self.last_pop_timestamp.load(Ordering::Relaxed)
}
#[inline]
pub fn set_last_pop_timestamp(&self, last_pop_timestamp: u64) {
self.last_pop_timestamp.store(last_pop_timestamp, Ordering::Release);
}
#[inline]
pub fn inc_found_msg(&self, count: usize) {
self.wait_ack_counter
.fetch_add(count.min(i32::MAX as usize) as i32, Ordering::AcqRel);
}
#[inline]
pub fn ack(&self) -> i32 {
self.wait_ack_counter.fetch_sub(1, Ordering::AcqRel)
}
#[inline]
pub fn dec_found_msg(&self, count: usize) {
self.wait_ack_counter
.fetch_sub(count.min(i32::MAX as usize) as i32, Ordering::AcqRel);
}
#[inline]
pub fn get_wai_ack_msg_count(&self) -> i32 {
self.wait_ack_counter.load(Ordering::Acquire)
}
#[inline]
pub fn is_dropped(&self) -> bool {
self.dropped.load(Ordering::Acquire)
}
#[inline]
pub fn set_dropped(&self, dropped: bool) {
self.dropped.store(dropped, Ordering::Release);
}
#[inline]
pub fn fill_pop_process_queue_info(&self, info: &mut PopProcessQueueInfo) {
info.set_wait_ack_count(self.get_wai_ack_msg_count());
info.set_droped(self.is_dropped());
info.set_last_pop_timestamp(self.get_last_pop_timestamp());
}
#[inline]
pub fn is_pull_expired(&self) -> bool {
let current_time = current_millis();
current_time.saturating_sub(self.last_pop_timestamp.load(Acquire)) > *PULL_MAX_IDLE_TIME
}
}
impl Display for PopProcessQueue {
#[inline]
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"PopProcessQueue[waitAckCounter:{}, lastPopTimestamp:{}, drop:{}]",
self.wait_ack_counter.load(Ordering::Relaxed),
self.last_pop_timestamp.load(Ordering::Relaxed),
self.dropped.load(Ordering::Relaxed)
)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn pop_process_queue_initializes_correctly() {
let queue = PopProcessQueue::new();
assert!(queue.get_last_pop_timestamp() > 0);
assert_eq!(queue.get_wai_ack_msg_count(), 0);
assert!(!queue.is_dropped());
}
#[test]
fn pop_process_queue_updates_last_pop_timestamp() {
let queue = PopProcessQueue::new();
let new_timestamp = queue.get_last_pop_timestamp() + 1000;
queue.set_last_pop_timestamp(new_timestamp);
assert_eq!(queue.get_last_pop_timestamp(), new_timestamp);
}
#[test]
fn pop_process_queue_increments_found_msg() {
let queue = PopProcessQueue::new();
queue.inc_found_msg(5);
assert_eq!(queue.get_wai_ack_msg_count(), 5);
}
#[test]
fn pop_process_queue_decrements_found_msg() {
let queue = PopProcessQueue::new();
queue.inc_found_msg(5);
queue.dec_found_msg(3);
assert_eq!(queue.get_wai_ack_msg_count(), 2);
}
#[test]
fn pop_process_queue_acknowledges_msg() {
let queue = PopProcessQueue::new();
queue.inc_found_msg(5);
assert_eq!(queue.ack(), 5);
assert_eq!(queue.get_wai_ack_msg_count(), 4);
}
#[test]
fn pop_process_queue_ack_at_zero_matches_java_signed_counter() {
let queue = PopProcessQueue::new();
assert_eq!(queue.ack(), 0);
assert_eq!(queue.get_wai_ack_msg_count(), -1);
}
#[test]
fn pop_process_queue_sets_dropped_flag() {
let queue = PopProcessQueue::new();
queue.set_dropped(true);
assert!(queue.is_dropped());
}
#[test]
fn pop_process_queue_detects_pull_expired() {
let queue = PopProcessQueue::new();
queue.set_last_pop_timestamp(queue.get_last_pop_timestamp() - *PULL_MAX_IDLE_TIME - 1);
assert!(queue.is_pull_expired());
}
#[test]
fn pop_process_queue_info_preserves_negative_wait_ack_count_like_java() {
let queue = PopProcessQueue::new();
queue.ack();
let mut info = PopProcessQueueInfo::new(0, false, 0);
queue.fill_pop_process_queue_info(&mut info);
assert_eq!(info.wait_ack_count(), -1);
}
}