Rust中线程安全高效字符串驻留Map实现优化问询

当前实现了一个基于Arc和DashMap的全局字符串驻留模块Text，核心特性包括：

• 全局一致性：STRINGS集合确保每个活跃Text对应唯一条目
• 字符串去重：内存中仅保留一份字符串实例
• 快速比较与哈希：基于Arc指针地址实现O(1)的相等性判断和哈希计算

多线程场景下的核心挑战：

• 避免字符串去重时的竞态条件
• Drop过程中，其他线程可能克隆Arc导致引用计数判断错误
• 当前采用key-specific locks（KEY_LOCKS）保证字符串创建、访问、删除的原子性，需验证：
  1. 实现的多线程正确性
  2. 方案的效率优化空间，尤其是key锁的开销

当前代码实现：

use dashmap::{DashMap, DashSet};
use std::{
    sync::{Arc, LazyLock, Mutex},
    hash::{Hash, Hasher},
};

static STRINGS: LazyLock<DashSet<Arc<String>>> = LazyLock::new(|| DashSet::new());

static KEY_LOCKS: LazyLock<DashMap<Arc<String>, Arc<Mutex<()>>>> = LazyLock::new(|| DashMap::new());

fn get_key_lock(key: Arc<String>) -> Arc<Mutex<()>> {
    KEY_LOCKS
        .entry(key)
        .or_insert_with(|| Arc::new(Mutex::new(())))
        .clone()
}

#[derive(Debug, Clone)]
pub struct Text(Arc<String>);

impl Text {
    fn create_or_get_text(arc: Arc<String>) -> Self {
        let key_lock = get_key_lock(arc.clone());
        let _lock = key_lock.lock().unwrap();

        if let Some(existing) = STRINGS.get(&arc) {
            existing.clone()
        } else {
            STRINGS.insert(arc.clone());
            Text(arc)
        }
    }
}

impl Drop for Text {
    fn drop(&mut self) {
        let key_lock = get_key_lock(self.0.clone());
        let _lock = key_lock.lock().unwrap();
        if Arc::strong_count(&self.0) == 2 {
            STRINGS.remove(&self.0);
        }
    }
}

impl Hash for Text {
    fn hash<H: Hasher>(&self, state: &mut H) {
        let addr = Arc::as_ptr(&self.0) as usize as u64;
        state.write_u64(addr);
    }
}

impl PartialEq for Text {
    fn eq(&self, other: &Self) -> bool {
        Arc::ptr_eq(&self.0, &other.0)
    }
}

impl Eq for Text {}

备注：

• STRINGS用于快速查找与去重
• KEY_LOCKS用于保证线程安全，与STRINGS分离是为了避免给每个驻留字符串绑定Mutex的开销
• Drop语义：当Text实例被销毁时，仅在无其他外部引用时从STRINGS中移除条目；当前代码中引用计数判断存在误差，实际应为3（STRINGS持有1个，当前Text实例持有1个，get_key_lock调用时克隆又生成1个）

一、正确性分析 #

当前实现存在3个关键问题，导致多线程场景下不可靠：

1. Drop中的引用计数判断错误 #

调用get_key_lock(self.0.clone())会额外克隆一次Arc，使self.0的强引用计数临时+1。假设原本只有STRINGS和即将被Drop的Text持有引用（计数为2），克隆后计数变为3，此时Arc::strong_count(&self.0) == 2的判断永远不会成立，导致条目永远无法从STRINGS中移除，最终造成内存泄漏。

正确的判断逻辑应基于锁持有期间的实际引用数：

impl Drop for Text {
    fn drop(&mut self) {
        let arc = self.0.clone();
        let key_lock = get_key_lock(arc.clone());
        let _lock = key_lock.lock().unwrap();
        // 此时arc的引用包括：STRINGS（1）、当前Text（1）、get_key_lock中的克隆（1），总共3
        if Arc::strong_count(&arc) == 3 {
            STRINGS.remove(&arc);
        }
    }
}

2. KEY_LOCKS的内存泄漏 #

KEY_LOCKS中的锁条目永远不会被移除，即使对应的字符串已经从STRINGS中删除。随着不同字符串的创建，KEY_LOCKS会无限增长，最终耗尽内存。

3. 已解决的竞态问题（需配合修复） #

key-specific locks本身能保证同一字符串的创建和删除操作互斥，避免了多线程同时创建同一字符串导致的重复插入，也避免了删除时其他线程正在克隆的问题，但需先解决前两个问题才能保证整体正确性。

二、效率优化方案 #

1. 替换轻量锁降低开销 #

用parking_lot::Mutex替代标准库Mutex，它在无竞争场景下开销更低，竞争场景下的调度效率也更高。只需修改依赖和导入即可：

use parking_lot::Mutex;

2. 清理KEY_LOCKS避免内存泄漏 #

在删除STRINGS条目时，尝试清理KEY_LOCKS中对应的锁条目，仅当锁的引用计数为1（即只有KEY_LOCKS持有该锁）时才安全移除：

impl Drop for Text {
    fn drop(&mut self) {
        let arc = self.0.clone();
        let key_lock = get_key_lock(arc.clone());
        let _lock = key_lock.lock().unwrap();

        if Arc::strong_count(&arc) == 3 {
            STRINGS.remove(&arc);
            // 尝试清理锁条目
            if let Some((_, lock)) = KEY_LOCKS.remove(&arc) {
                if Arc::strong_count(&lock) == 1 {
                    // 无其他线程使用，锁会自动销毁
                } else {
                    // 其他线程仍在使用，重新插入锁条目
                    KEY_LOCKS.insert(arc, lock);
                }
            }
        }
    }
}

3. 改用弱引用简化架构 #

放弃key-specific locks，用DashMap<Arc<String>, Weak<String>>存储驻留字符串，利用弱引用自动处理引用计数：

use dashmap::DashMap;
use std::{
    sync::{Arc, LazyLock, Weak},
    hash::{Hash, Hasher},
};

static STRINGS: LazyLock<DashMap<Arc<String>, Weak<String>>> = LazyLock::new(DashMap::new);

#[derive(Debug, Clone)]
pub struct Text(Arc<String>);

impl Text {
    pub fn create_or_get(s: &str) -> Self {
        let key = Arc::new(s.to_string());
        let mut entry = STRINGS.entry(key.clone()).or_insert_with(|| Arc::downgrade(&key));
        
        // 检查弱引用是否存活
        if let Some(existing) = entry.get().upgrade() {
            Text(existing)
        } else {
            // 弱引用失效，替换为新的弱引用
            *entry = Arc::downgrade(&key);
            Text(key)
        }
    }
}

// 无需Drop实现，弱引用自动失效；可定期清理无效条目

这种方案无需锁，完全依赖DashMap的原子操作和弱引用特性，避免了竞态条件，缺点是需要定期清理STRINGS中的无效弱引用条目（可在插入时随机触发清理）。

4. 减少不必要的Arc克隆 #

将create_or_get_text的参数改为&str，内部创建Arc，避免外部传入Arc带来的额外引用计数开销：

impl Text {
    pub fn create_or_get(s: &str) -> Self {
        let key = Arc::new(s.to_string());
        let key_lock = get_key_lock(key.clone());
        let _lock = key_lock.lock().unwrap();

        if let Some(existing) = STRINGS.get(&key) {
            existing.clone()
        } else {
            STRINGS.insert(key.clone());
            Text(key)
        }
    }
}

三、总结 #

• 正确性修复：优先修正Drop中的引用计数判断，再解决KEY_LOCKS的内存泄漏问题；
• 效率优化：替换轻量锁、改用弱引用架构、减少不必要的Arc克隆，三者任选或组合使用；
• 若不想自行实现复杂逻辑，可直接使用成熟的字符串驻留库（如string-cache或internment），已内置线程安全和高效去重机制。

内容的提问来源于stack exchange，提问作者optevo