基于mmap的高性能可持久化内存队列实现方案

7 阅读6分钟

背景

近期在某个系统的一部分实现细节中遇到了一个常见的生产者-消费者模型,即同一个服务(进程)内有一个生产者线程和一个消费者线程,二者通过一个内存队列进行解耦,生产者往队列写数据,消费者从队列取数据再进行异步处理。由于是内存队列,服务一但重启,如果队列中原来还有数据则会丢失,我们想尽可能避免这种情况,提升系统可靠性。因此,我们提出为队列增加持久化能力,进程退出的时候可将队列数据暂存到临时文件,重启后再从文件恢复。(为什么不选择kafka等成熟的组件? 因为我们的应用场景很简单,就是个单机部署的服务内部一写一读场景,希望实现尽可能轻量,引入第三方组件增加系统复杂性和维护成本。)

方案

整体实现思路很简单,就是基于mmap创建文件内存映射,将一个本地文件映射至进程内存空间,并将这块内存空间按照一个循环队列(数组)进行管理,读写数据直接操作这块内存,操作系统保证将内存中的脏页数据写回文件(除非系统宕机的情况,这时可能造成数据丢失)

整块内存分为两块:开头存放元数据(包含队列容量、元素大小、队列头尾指针下标等),之后是数据区(数据区看作是一个特定类型的循环数组)。具体实现如下:

template <typename T>
class PersistentQueue {
private:
    // 元数据结构,放在文件开头
    struct QueueHeader {
        std::atomic<uint64_t> head; // 读位置
        std::atomic<uint64_t> tail; // 写位置
        uint64_t capacity;          // 队列容量(实际容量+1)
        uint64_t element_size;      // 元素大小
        bool initialized;           // 是否已初始化
    };

    int fd_;              // 文件描述符
    char* data_;          // 映射的内存地址
    size_t mapped_size_;  // 映射大小
    QueueHeader* header_; // 元数据指针
    T* buffer_;           // 数据缓冲区指针
    size_t capacity_;     // 队列容量

public:
    // 构造函数
    PersistentQueue(const std::string& file_path, size_t capacity)
        : fd_(-1), data_(nullptr), header_(nullptr), buffer_(nullptr) {
        capacity_          = capacity;
        size_t header_size = sizeof(QueueHeader);
        size_t data_size   = capacity * sizeof(T);
        mapped_size_       = header_size + data_size;

        // 打开或创建文件
        fd_ = open(file_path.c_str(), O_RDWR | O_CREAT, 0644);
        if (fd_ < 0) {
            throw std::runtime_error("Failed to open file");
        }

        // 设置文件大小
        if (ftruncate(fd_, mapped_size_) < 0) {
            close(fd_);
            throw std::runtime_error("Failed to truncate file");
        }

        // 映射文件到内存
        data_ = static_cast<char*>(
            mmap(nullptr, mapped_size_, PROT_READ | PROT_WRITE, MAP_SHARED, fd_, 0));
        if (data_ == MAP_FAILED) {
            close(fd_);
            throw std::runtime_error("Failed to mmap file");
        }

        // 设置指针
        header_ = reinterpret_cast<QueueHeader*>(data_);
        buffer_ = reinterpret_cast<T*>(data_ + header_size);

        // 如果是新文件,初始化元数据
        initialize();
    }

    ~PersistentQueue() {
        if (data_) {
            // 确保数据写回磁盘
            msync(data_, mapped_size_, MS_SYNC);
            munmap(data_, mapped_size_);
        }
        if (fd_ >= 0) {
            close(fd_);
        }
    }

    // 禁止拷贝
    PersistentQueue(const PersistentQueue&) = delete;
    PersistentQueue& operator=(const PersistentQueue&) = delete;

    bool push(const T& item) {
        uint64_t current_tail = header_->tail.load(std::memory_order_relaxed);
        uint64_t next_tail    = (current_tail + 1) % capacity_;
        uint64_t current_head = header_->head.load(std::memory_order_acquire);

        // 检查队列是否已满
        if (next_tail == current_head) {
            return false;
        }

        // 写入数据
        new (&buffer_[current_tail]) T(item);

        // 更新尾指针
        header_->tail.store(next_tail, std::memory_order_release);

        // 可选:立即同步到磁盘
        // msync(data_, mapped_size_, MS_ASYNC);

        return true;
    }

    bool pop(T& item) {
        uint64_t current_head = header_->head.load(std::memory_order_relaxed);
        uint64_t current_tail = header_->tail.load(std::memory_order_acquire);

        // 检查队列是否为空
        if (current_head == current_tail) {
            return false;
        }

        // 读取数据
        item = buffer_[current_head];
        buffer_[current_head].~T(); // 调用析构函数

        // 更新头指针
        uint64_t next_head = (current_head + 1) % capacity_;
        header_->head.store(next_head, std::memory_order_release);

        return true;
    }

    bool empty() const {
        return header_->head.load(std::memory_order_acquire) ==
            header_->tail.load(std::memory_order_acquire);
    }

    bool full() const {
        uint64_t next_tail = (header_->tail.load(std::memory_order_acquire) + 1) % capacity_;
        return next_tail == header_->head.load(std::memory_order_acquire);
    }

    size_t size() const {
        uint64_t head = header_->head.load(std::memory_order_acquire);
        uint64_t tail = header_->tail.load(std::memory_order_acquire);

        if (tail >= head) {
            return tail - head;
        } else {
            return capacity_ - (head - tail);
        }
    }

    // 获取剩余容量
    size_t available() const { return capacity_ - 1 - size(); }

    // 强制同步到磁盘
    void sync() {
        if (data_) {
            msync(data_, mapped_size_, MS_SYNC);
        }
    }

private:
    void initialize() {
        // 检查是否需要初始化
        if (!header_->initialized) {
            header_->head.store(0, std::memory_order_relaxed);
            header_->tail.store(0, std::memory_order_relaxed);
            header_->capacity     = capacity_;
            header_->element_size = sizeof(T);
            header_->initialized  = true;

            // 使用内存屏障确保初始化完成
            std::atomic_thread_fence(std::memory_order_release);

            // 立即写回磁盘
            msync(data_, sizeof(QueueHeader), MS_SYNC);
        } else {
            // 验证文件格式
            if (header_->capacity != capacity_ || header_->element_size != sizeof(T)) {
                throw std::runtime_error("File format mismatch");
            }
        }
    }
};

创建队列时,会打开指定文件并创建一个文件内存映射,映射区域大小为:存储元素大小*容量 + header大小,之后初始化head指针、数据区指针、容量、元素大小、队列头尾指针下标等。如果是首次打开文件,初始化完成后会将元数据强制刷盘一次,非首次打开则验证打开文件数据格式是否和要创建的队列所需格式一致。

该队列用于一写一读场景,所有操作基于内存且lock free,具有非常高读写性能。

测试

#include <iostream>
#include <sstream>
#include <string>
#include <thread>
#include "PersistentQueue.h"

struct Data {
    int id;
    char name[8];
    operator std::string() {
        std::ostringstream os;
        os << "{\"id\":" << id << ", \"name\":\"" << name << "\"}";
        return os.str();
    }
};

PersistentQueue<Data> queue("./queue.dat", 10000);

int Write(int count) {
    for (int i = 0; i < count; ++i) {
        Data data{i, "abc"};
        if (!queue.push(data)) {
            std::cout << "Queue is full!" << std::endl;
        } else {
            std::cout << "Pushed: " << (std::string)(data) << std::endl;
        }
    }
    return 0;
}

int Read(int count) {
    Data data;
    while (count) {
        if(queue.pop(data)) {
            std::cout << "Popped: " << (std::string)(data) << std::endl;
            --count;
        }
    }
    return 0;
}

int main(int argc, char* argv[]) {
    std::thread producer_thread(Write, 10000);
    std::thread consumer_thread(Read, 10000);
    producer_thread.join();
    consumer_thread.join();
    return 0;
}

注意事项

内存映射文件只能存储连续内存,无法存储指针指向的堆,内存队列的push方法中,我们使用了placement new在映射区域内存上直接创建对象并将元素拷贝过去,这要求所存储对象类型必须是支持平凡拷贝的(也即可以通过简单复制内存实现对象拷贝),因此元素中不能包含指针、std::string、std::容器类型这段,因为这些类型底层都包含指向堆内存的指针,指针的值虽然能持久化到文件,但程序重启后指针所指向的堆内存已失效。

如果要支持带非平凡拷贝类型字段的元素,可以考虑以下两种方式:

  1. 修改结构体为平凡类型,如使用固定长度的char[]数组代替std::string等。
  2. 使用序列化存储,将元素序列化为字符串后,按char[]数组方式存储。

在我们的系统中使用了方式1支持非平凡拷贝类型,这样每个元素都固定大小,方便对堆数据区进行管理,虽然会在一定程度上造成内存空间浪费。 如对于以下非平凡拷贝类型:

struct OrderInfo {
    uint64_t uid;
    std::string order_id;
    std::map<std::string, std::string> params;
};

可以通过该转换为以下支持平凡拷贝类型,进而使用PersistentQueue:

struct PersistentOrderInfo {
    uint64_t uid;
    char order_id[64];     // 固定长度(确保长度够用就行)
    
    // 固定长度的参数存储
    struct Param {
        char key[32];
        char value[128];
    };
    
    uint32_t param_count;  // 参数个数
    Param params[10];      // 最多10个参数
    
    // 构造函数
    PersistentOrderInfo() : uid(0), order_id{0}, param_count(0) {
        memset(params, 0, sizeof(params));
    }
    
    / 从原始OrderInfo转换
    void fromOrderInfo(const OrderInfo& order) {
        uid = order.uid;
        strncpy(order_id, order.order_id.c_str(), sizeof(order_id) - 1);
      
        param_count = 0;
        for (const auto& [key, value] : order.params) {
            if (param_count >= 10) {
                break;
            }
            strncpy(params[param_count].key, key.c_str(), sizeof(Param::key) - 1);
            strncpy(params[param_count].value, value.c_str(), sizeof(Param::value) - 1);
            param_count++;
        }
    }
    
    // 转换为原始OrderInfo
    OrderInfo toOrderInfo() const {
        OrderInfo order;
        order.uid = uid;
        order.order_id = order_id;

        for (uint32_t i = 0; i < param_count; ++i) {
            order.params[params[i].key] = params[i].value;
        }
        
        return order;
    }