在 Rust 中实现线程安全且高效的字符串驻留映射

Question

我希望能够帮助确保线程安全并提高 Rust 中全局字符串驻留映射的效率，特别是在使用特定于键的锁来管理并发访问并防止字符串重复数据删除期间的竞争条件时。

下面显示的文本模块片段旨在通过重复数据删除以内存高效且高性能的方式管理字符串。它使用全局实习映射来存储唯一的字符串实例，利用 Arc 进行引用计数，并利用 DashMap 进行线程安全的并发访问。主要特点是：

• 全局一致性：实习映射（STRINGS）确保每个活动文本都对应于映射中的一个条目。
• 字符串去重：内存中只存在给定字符串的一个实例。
• 基于指针的相等和快速哈希：由于没有重复项，字符串相等和哈希基于 Arc 内存地址，提供极快的比较和哈希插入/查找。

挑战在于确保驻留映射（STRINGS）保持线程安全和高效，同时防止多线程上下文中出现重复字符串。当删除 Text 实例时，仅当不存在其他引用时，才应从驻留映射中删除其关联字符串（不一定立即）。但是，可能会出现竞争条件，因为其他线程可能会在 Drop 过程中克隆 Arc，从而造成引用计数不匹配。我考虑过使用队列来管理删除，但发现它不足以解决竞争条件（尽管可以帮助解决其他问题，例如删除大量字符串）。

我当前的方法是使用特定于键的锁（KEY_LOCKS）来确保文本创建、访问和删除每个字符串期间的原子访问。如果您有任何反馈或建议，我将不胜感激：

• 正确性：从多线程的角度来看，这个实现是否合理？
• 效率：这种方法是否可以改进，特别是在特定于键的锁的开销方面？

这是模块中的相关代码：

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 和 KEY_LOCKS 主要作为单独的结构维护，这样我就不会为每个内部字符串使用互斥锁。

删除语义：当删除文本实例时，它会检查是否应该从全局驻留映射中删除其相应的条目。 Arc 引用计数 (strong_count) 被测试为 3：当准备从实习映射中丢弃时，drop 将传递一个带有 2 个引用的文本：一个由 STRINGS 实习映射保存，一个引用由 Text 实例保存被丢弃。使用 KEY_LOCKS 时，在放置函数期间会创建一个附加引用。

Answer 1

如果您使用

remove_if

来处理驱逐，那么您似乎应该能够完全删除互斥体。在底层，这会取出包含存储桶的分片上的写锁，并将其保留到整个操作完成，这使得对内部

Text

的并发访问变得不可能。因此，你应该能够做这样的事情：

impl Drop for Text {
    fn drop(&mut self) {
        STRINGS.remove_if(&self.0, |s| Arc::strong_count(s) == 2);
    }
}

但是，进行此更改后，插入时存在竞争条件：如果同时对

create_or_get_text

进行两次调用，并且都没有找到存在的字符串，则它们都会尝试插入它。 我们可以通过检测插入失败的情况并重新启动操作来纠正此问题。

fn create_or_get_text(arc: Arc<String>) -> Self {
    loop {
        if let Some(existing) = STRINGS.get(&arc) {
            break Self(existing.clone());
        } else if STRINGS.insert(arc.clone()) {
            break Self(arc);
        }
    }
}

旁注：只需对指针本身进行哈希处理，即可使您的

Hash

实现对不同的指针大小更具弹性。

impl Hash for Text {
    fn hash<H: Hasher>(&self, state: &mut H) {
        Arc::as_ptr(&self.0).hash(state);
    }
}

上述修改后的完整代码：

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

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

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

impl Text {
    fn create_or_get_text(arc: Arc<String>) -> Self {
        loop {
            if let Some(existing) = STRINGS.get(&arc) {
                break Self(existing.clone());
            } else if STRINGS.insert(arc.clone()) {
                break Self(arc);
            }
        }
    }
}

impl Drop for Text {
    fn drop(&mut self) {
        STRINGS.remove_if(&self.0, |s| Arc::strong_count(s) == 2);
    }
}

impl Hash for Text {
    fn hash<H: Hasher>(&self, state: &mut H) {
        Arc::as_ptr(&self.0).hash(state);
    }
}

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

impl Eq for Text {}