rocketmq-observability 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::sync::atomic::AtomicI64;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::Duration;

use parking_lot::Mutex;
use rocketmq_runtime::OperationContext;
use rocketmq_runtime::ScheduledExecutionPolicy;
use rocketmq_runtime::ScheduledTaskConfig;
use rocketmq_runtime::ScheduledTaskGroup;
use rocketmq_runtime::TaskGroup;
use rocketmq_runtime::TaskKind;
use tracing::info;
use tracing::warn;

use rocketmq_runtime::common::time_utils::compute_next_minutes_time_millis;
use rocketmq_runtime::common::time_utils::current_millis;

const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5);

struct StatsTaskOwner {
    task_group: TaskGroup,
    operation: OperationContext,
}

#[derive(Clone)]
pub struct MomentStatsItem {
    value: Arc<AtomicI64>,
    stats_name: String,
    stats_key: String,
    task_owner: Arc<Mutex<Option<StatsTaskOwner>>>,
    parent_task_group: TaskGroup,
}

impl MomentStatsItem {
    pub fn new(stats_name: String, stats_key: String, parent_task_group: TaskGroup) -> Self {
        MomentStatsItem {
            value: Arc::new(AtomicI64::new(0)),
            stats_name,
            stats_key,
            task_owner: Arc::new(Mutex::new(None)),
            parent_task_group,
        }
    }

    pub fn new_with_task_group(stats_name: String, stats_key: String, parent_task_group: TaskGroup) -> Self {
        Self::new(stats_name, stats_key, parent_task_group)
    }

    pub fn init(self: Arc<Self>) {
        if self.task_owner.lock().is_some() {
            warn!(
                "[{}] [{}] MomentStatsItem already initialized",
                self.stats_name, self.stats_key
            );
            return;
        }

        let task_group = self.parent_task_group.clone();
        let operation = OperationContext::without_deadline(TaskKind::ScheduledDriver);
        let scheduled_tasks = ScheduledTaskGroup::new(task_group.clone());
        let self_clone = self.clone();
        let mut config =
            ScheduledTaskConfig::fixed_rate_no_overlap("common.moment-stats.print", Duration::from_secs(300));
        config.initial_delay =
            Duration::from_millis((compute_next_minutes_time_millis() as i64 - current_millis() as i64).unsigned_abs());

        if let Err(error) =
            scheduled_tasks.schedule_operation(&operation, config, ScheduledExecutionPolicy::default(), move || {
                let self_clone = self_clone.clone();
                async move {
                    self_clone.print_at_minutes();
                    self_clone.value.store(0, Ordering::Relaxed);
                }
            })
        {
            warn!(
                "[{}] [{}] failed to spawn MomentStatsItem task: {}",
                self.stats_name, self.stats_key, error
            );
            return;
        }
        *self.task_owner.lock() = Some(StatsTaskOwner { task_group, operation });
    }

    pub async fn shutdown(&self) {
        let task_owner = { self.task_owner.lock().take() };
        if let Some(task_owner) = task_owner {
            match task_owner
                .operation
                .cancel_and_wait(&task_owner.task_group, SHUTDOWN_TIMEOUT)
                .await
            {
                Ok(true) => {}
                Ok(false) => warn!(
                    "[{}] [{}] MomentStatsItem shutdown exceeded its deadline",
                    self.stats_name, self.stats_key
                ),
                Err(error) => warn!(
                    "[{}] [{}] MomentStatsItem used an invalid task owner: {}",
                    self.stats_name, self.stats_key, error
                ),
            }
        }
    }

    pub fn print_at_minutes(&self) {
        info!(
            "[{}] [{}] Stats Every 5 Minutes, Value: {}",
            self.stats_name,
            self.stats_key,
            self.value.load(Ordering::Relaxed)
        );
    }

    pub fn get_value(&self) -> Arc<AtomicI64> {
        self.value.clone()
    }

    pub fn get_stats_key(&self) -> &str {
        &self.stats_key
    }

    pub fn get_stats_name(&self) -> &str {
        &self.stats_name
    }
}

#[cfg(test)]
mod tests {
    use std::sync::atomic::Ordering;

    use super::*;

    fn test_parent(name: &'static str) -> TaskGroup {
        rocketmq_runtime::RuntimeContext::from_current(name)
            .service_context("moment-stats-service")
            .task_group()
            .clone()
    }

    #[tokio::test]
    async fn moment_stats_item_initializes_with_zero_value() {
        let stats_item = MomentStatsItem::new(
            "TestName".to_string(),
            "TestKey".to_string(),
            test_parent("moment-stats-value-test"),
        );
        assert_eq!(stats_item.get_value().load(Ordering::Relaxed), 0);
    }

    #[tokio::test]
    async fn moment_stats_item_returns_correct_stats_name() {
        let stats_item = MomentStatsItem::new(
            "TestName".to_string(),
            "TestKey".to_string(),
            test_parent("moment-stats-name-test"),
        );
        assert_eq!(stats_item.get_stats_name(), "TestName");
    }

    #[tokio::test]
    async fn moment_stats_item_returns_correct_stats_key() {
        let stats_item = MomentStatsItem::new(
            "TestName".to_string(),
            "TestKey".to_string(),
            test_parent("moment-stats-key-test"),
        );
        assert_eq!(stats_item.get_stats_key(), "TestKey");
    }
}