From ea9f583dc0bd5b87d7235c34406ccf5926c2f0da Mon Sep 17 00:00:00 2001
From: ZOt6 <57828206+YHM404@users.noreply.github.com>
Date: Fri, 3 Jan 2025 22:33:32 +0800
Subject: [PATCH] Replace unsafe cache init with OnceLock (#1331)

---
 server/src/streaming/cache/memory_tracker.rs | 21 +++++++++-----------
 1 file changed, 9 insertions(+), 12 deletions(-)

diff --git a/server/src/streaming/cache/memory_tracker.rs b/server/src/streaming/cache/memory_tracker.rs
index 082a999e0..835c44ffe 100644
--- a/server/src/streaming/cache/memory_tracker.rs
+++ b/server/src/streaming/cache/memory_tracker.rs
@@ -4,12 +4,11 @@ use crate::configs::resource_quota::MemoryResourceQuota;
 use crate::configs::system::CacheConfig;
 use iggy::utils::byte_size::IggyByteSize;
 use std::sync::atomic::{AtomicU64, Ordering};
-use std::sync::{Arc, Once};
+use std::sync::{Arc, OnceLock};
 use sysinfo::System;
 use tracing::info;
 
-static ONCE: Once = Once::new();
-static mut INSTANCE: Option<Arc<CacheMemoryTracker>> = None;
+static INSTANCE: OnceLock<Option<Arc<CacheMemoryTracker>>> = OnceLock::new();
 
 #[derive(Debug)]
 pub struct CacheMemoryTracker {
@@ -21,22 +20,20 @@ type MessageSize = u64;
 
 impl CacheMemoryTracker {
     pub fn initialize(config: &CacheConfig) -> Option<Arc<CacheMemoryTracker>> {
-        unsafe {
-            ONCE.call_once(|| {
+        INSTANCE
+            .get_or_init(|| {
                 if config.enabled {
-                    INSTANCE = Some(Arc::new(CacheMemoryTracker::new(config.size.clone())));
-                    info!("Cache memory tracker initialized");
+                    Some(Arc::new(CacheMemoryTracker::new(config.size.clone())))
                 } else {
-                    INSTANCE = None;
                     info!("Cache memory tracker disabled");
+                    None
                 }
-            });
-            INSTANCE.clone()
-        }
+            })
+            .clone()
     }
 
     pub fn get_instance() -> Option<Arc<CacheMemoryTracker>> {
-        unsafe { INSTANCE.clone() }
+        INSTANCE.get().cloned().flatten()
     }
 
     fn new(limit: MemoryResourceQuota) -> Self {