线程安全设计：互斥锁、读写锁、无锁队列在显存池中的应用

说实话，多线程环境下的显存池设计，是我这些年踩坑最多的领域之一。你想想看，GPU显存本身就是一个稀缺资源，多个线程同时申请释放，稍不注意就是各种诡异崩溃。今天我就把这块硬骨头掰开了讲。

为什么显存池必须考虑线程安全？

深度学习框架的推理场景，往往是多线程并发的。比如TensorRT部署时，多个请求线程同时向显存池要内存。如果没有保护机制，两个线程同时修改空闲链表，轻则内存泄漏，重则直接Segmentation Fault。

我遇到过最离谱的一次，线上服务跑了三天突然挂掉。查到最后，就是显存池的free操作没有加锁，两个线程同时释放相邻的block，链表指针被写坏了。嗯，从那以后我再也不敢轻视线程安全了。

方案一：互斥锁（std::mutex）

最朴素的做法，就是给整个显存池加一把大锁。所有申请和释放操作，都得先拿到这把锁。

class MemoryPool {
    std::mutex mtx_;
    std::vector<void*> free_blocks_;
public:
    void* allocate(size_t size) {
        std::lock_guard<std::mutex> lock(mtx_);
        if (free_blocks_.empty()) {
            // 向GPU申请新内存
            return cudaMalloc(size);
        }
        auto ptr = free_blocks_.back();
        free_blocks_.pop_back();
        return ptr;
    }
    
    void deallocate(void* ptr) {
        std::lock_guard<std::mutex> lock(mtx_);
        free_blocks_.push_back(ptr);
    }
};

这种方案的好处是简单，不容易出错。但问题也很明显——锁的粒度太粗了。如果线程A在申请大块内存（比如512MB），线程B只能干等着。在高并发场景下，这会导致严重的性能瓶颈。

适用场景：并发度不高（比如4-8个线程），或者显存操作不频繁的场景。我个人习惯在原型验证阶段先用互斥锁，等性能瓶颈出现了再优化。

方案二：读写锁（std::shared_mutex）

读写锁的思路是：读操作可以共享，写操作必须独占。在显存池里，什么算读？什么算写？

读操作：查询空闲列表大小、检查某个地址是否属于本池
写操作：申请内存（修改空闲列表）、释放内存（修改空闲列表）

说白了，大部分场景下读写锁并不比互斥锁好多少。因为显存池的核心操作——申请和释放——都是写操作。只有在「先检查再操作」这种场景下，读写锁才有优势。

class MemoryPool {
    mutable std::shared_mutex rw_mtx_;
    std::vector<void*> free_blocks_;
public:
    bool has_free_block() const {
        std::shared_lock lock(rw_mtx_);
        return !free_blocks_.empty();
    }
    
    void* allocate() {
        std::unique_lock lock(rw_mtx_);
        // 实际分配逻辑
    }
};

我的经验：读写锁在显存池中实际收益有限。除非你的场景是「大量线程频繁检查空闲状态，但很少实际申请释放」，否则直接用互斥锁反而更简单高效。

方案三：无锁队列（Lock-Free Queue）

这才是真正的高性能方案。无锁队列利用CAS（Compare-And-Swap）原子操作，避免了线程上下文切换的开销。

我曾在PyTorch的CachingAllocator里看到类似的设计。核心思路是：每个线程有自己的本地缓存，优先从本地拿；本地不够了，再从全局无锁队列里取。

class LockFreeMemoryPool {
    std::atomic<BlockHeader*> head_;
    
    void* allocate() {
        BlockHeader* old_head;
        BlockHeader* new_head;
        do {
            old_head = head_.load();
            if (!old_head) return nullptr;
            new_head = old_head->next;
        } while (!head_.compare_exchange_weak(old_head, new_head));
        return old_head;
    }
    
    void deallocate(void* ptr) {
        BlockHeader* block = reinterpret_cast<BlockHeader*>(ptr);
        BlockHeader* old_head;
        do {
            old_head = head_.load();
            block->next = old_head;
        } while (!head_.compare_exchange_weak(old_head, block));
    }
};

注意：无锁队列不是银弹。ABA问题、内存回收问题（怎么安全地释放节点？）、以及多生产者多消费者场景下的正确性验证，都非常复杂。我曾经花了两周调试一个无锁队列的bug，最后发现是内存序（memory order）用错了。

三种方案的对比

方案	性能	实现复杂度	适用场景
互斥锁	中等	低	低并发、原型验证
读写锁	中等	中	读多写少的特殊场景
无锁队列	高	高	高并发、延迟敏感

实战中的选择建议

我个人建议这样选：

第一步：先用互斥锁把功能跑通。别一上来就搞无锁，bug查死你。
第二步：用perf工具分析锁竞争情况。如果锁等待时间占比超过5%，再考虑优化。
第三步：如果确实需要优化，优先考虑「线程本地缓存+全局互斥锁」的混合方案。这比纯无锁队列容易实现得多。

你想想看，很多框架的显存池其实都是这么设计的。PyTorch的CachingAllocator就是每个线程一个缓存，全局用互斥锁保护。效果已经很好了。

核心逻辑流程图

下面这张图展示了三种方案在显存池中的工作流程对比：

从流程图可以清楚看到：互斥锁是串行化的，读写锁在读操作时允许并发，而无锁队列则让所有线程都能同时操作。但别忘了，无锁队列的每个CAS操作都可能失败重试，这在高冲突场景下反而会降低性能。

总结一下：线程安全设计没有银弹。互斥锁简单可靠，读写锁适合特定场景，无锁队列性能上限高但实现门槛也高。我的建议是——从简单的开始，用数据说话，哪里瓶颈优化哪里。

想要更多资料，联系可可

📚 课程目录

线程安全设计：互斥锁、读写锁、无锁队列在显存池中的应用

为什么显存池必须考虑线程安全？

方案一：互斥锁（std::mutex）

方案二：读写锁（std::shared_mutex）

方案三：无锁队列（Lock-Free Queue）

三种方案的对比

实战中的选择建议

核心逻辑流程图