如何创建键值生命周期受闭包约束的Rust引用映射？

我想要创建一个仅在有限时间内持有引用的容器，其生命周期由传入的闭包约束，期望的API如下：

let mut refs = TempRefMap::<String, &usize>::new();
let (tx, rx) = channel::<(String, Sender<()>)>();

// 工作线程：收到值可用通知后读取键
thread::spawn(move || {
  while let Ok((key, tx_done)) = rx.recv() {
    // 假设内部用Arc<Mutex>封装了refs
    let value: &usize = refs.get_ref(&key).unwrap();
  
    // 执行任务
    println!("{}", value);

    // 通知任务完成
    tx_done.send(()).unwrap();
  }
});

// 插入引用到容器并通知工作线程值已可用
refs.insert_ref("key".to_string(), &42, || {
  // 用于通知工作线程完成的通道
  let (tx_done, rx_done) = channel::<()>();

  // 告知工作线程值已可用
  tx.send(("key".to_string(), tx_done)).unwrap();

  // 等待工作线程完成
  rx_done.recv().unwrap();

  // 闭包结束时键会被移除
});

我知道可以用Rc<T>替代&T来管理多引用，但希望尽量避免：一是这个场景下引用只会出现一次，改用Rc会污染大量代码；二是该操作每秒执行数百万次，Rc的引用计数可能带来性能损耗。

我可以接受用Unsafe Rust实现这个容器，但不确定正确的用法。我尝试实现了如下类型：

struct TempRefMap<K, V> {
  values: HashMap<K, &V>,
}

impl<K: Hash + Eq + PartialEq, V> TempRefMap<K, V> {
  fn new() -> Self {
    Self {
      values: HashMap::new()
    }
  }

  fn insert_ref(&mut self, key: K, value: &V, callback: impl FnOnce()) {
    self.values.insert(key.clone(), value);
    callback();
    self.values.remove(&key);
  }

  fn get_ref(&mut self, key: K) -> Option<&V> {
    self.values.get(&key)
  }
}

但存储&V时遇到很多生命周期问题。我曾尝试用unsafe解引用值来存储：

fn insert_ref(&mut self, key: T, value: &V, callback: impl FnOnce()) {
  let v = unsafe { *value };
  self.values.insert(key.clone(), v);
  callback();
  self.values.remove(&key);
}

但我知道这不是unsafe的正确用法，我对unsafe经验不足。

解决方案 #

核心思路 #

要解决生命周期问题，我们需要绕过Rust的静态生命周期检查，同时手动保证引用的安全性：

• 容器持有的引用必须在闭包执行期间有效，闭包结束后立即移除引用
• 工作线程的访问必须在闭包等待完成的时间段内，确保引用不会悬空

安全的Unsafe实现 #

我们可以用*const V（原始指针）存储引用，因为原始指针不受生命周期约束，然后在获取时转换回引用，同时手动保证安全：

use std::collections::HashMap;
use std::hash::Hash;
use std::sync::mpsc;
use std::thread;

struct TempRefMap<K, V> {
    values: HashMap<K, *const V>,
}

impl<K: Hash + Eq + Clone, V> TempRefMap<K, V> {
    fn new() -> Self {
        Self {
            values: HashMap::new(),
        }
    }

    /// 插入引用，闭包执行期间引用有效，闭包结束后自动移除
    fn insert_ref(&mut self, key: K, value: &V, callback: impl FnOnce()) {
        // 将共享引用转换为原始指针
        let ptr = value as *const V;
        self.values.insert(key.clone(), ptr);
        
        // 执行闭包：此时我们保证原始指针指向的内存仍然有效
        callback();
        
        // 闭包执行完毕，立即移除指针，避免后续访问悬空指针
        self.values.remove(&key);
    }

    /// 获取引用，必须保证此时该键对应的引用仍然有效
    fn get_ref(&self, key: &K) -> Option<&V> {
        self.values.get(key).map(|&ptr| {
            // 这里unsafe是安全的：
            // 1. 指针是从有效的&T转换而来
            // 2. 只有在insert_ref的闭包执行期间才会调用get_ref
            // 3. 闭包执行期间原始值仍然存活，且没有可变引用（因为我们传入的是&T）
            unsafe { &*ptr }
        })
    }
}

// 使用示例
fn main() {
    let mut refs = TempRefMap::<String, usize>::new();
    let (tx, rx) = mpsc::channel::<(String, mpsc::Sender<()>)>();

    // 工作线程
    thread::spawn(move || {
        while let Ok((key, tx_done)) = rx.recv() {
            let value = refs.get_ref(&key).unwrap();
            println!("{}", value);
            tx_done.send(()).unwrap();
        }
    });

    // 插入引用并等待工作线程完成
    let num = 42;
    refs.insert_ref("key".to_string(), &num, || {
        let (tx_done, rx_done) = mpsc::channel();
        tx.send(("key".to_string(), tx_done)).unwrap();
        rx_done.recv().unwrap();
    });
}

安全保证说明 #

我们需要手动维护以下不变量来确保unsafe代码的安全性：

• 指针有效性：insert_ref传入的&V必须在闭包执行的整个期间保持有效，也就是原始值不能在闭包执行时被销毁。
• 无数据竞争：在闭包执行期间，原始值不能被可变引用访问（因为我们传入的是共享引用&V，Rust的借用规则已经保证了这一点）。
• 及时清理：闭包执行完毕后立即移除指针，避免后续调用get_ref获取到悬空指针。

注意事项 #

• 如果你的场景中需要多线程访问TempRefMap，需要用Arc<Mutex<TempRefMap>>封装，确保线程安全。
• 必须严格保证get_ref只在insert_ref的闭包执行期间被调用，否则会导致未定义行为。

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