MPSC消息队列设计

设计原理

初始化

• 建立2048个元素的数组，用于传输数据
• 为每个元素分配一个版本号：turn
• turn为奇数是，为可读
• turn为偶数时，为可写

生产者

• 抢夺head原子变量idx，增加1
• 轮询idx下turn是否为偶数
• 如果为偶数，写入数据，并turn = idx/2048*2 + 1
• 如果为奇数，则一直轮询，知道turn变为偶数

消费者

• 判断tail下的turn是否为奇数
• 如果为奇数，则消费
• 并写入turn = idx/2048*2 + 2
• 如果为偶数，则轮询，知道turn变为奇数

局部优化

• 原子变量head, tail, turn align 64，避免伪共享
• 使用位运算代替取余 idx & size => idx & (size -1)
• 使用左移和右移来代替乘除 idx / size * 2 => （(idx >> 11)<<1)

代码

#include <vector>
#include <atomic>
#include <thread>
#include <iostream>
#include <sys/time.h>
#include <sstream>
#include <cstring>

#define SLOT_DATA_BUF_SIZE 2048
#define MPSC_IDX(i, size_mask) i&size_mask
#define MPSC_TURN(i, size_bit) ((i>>size_bit)<<1)

class MpscQueueSlot {
public:
    std::atomic<uint32_t> turn_ = {0};
    uint32_t padding[15];
};
class MpscQueueSlotData {
public:
    char value[SLOT_DATA_BUF_SIZE];
};
class MpscQueue {
public:
    const uint32_t size = 2048;
    const uint32_t size_mask = 2047;
    const uint32_t size_bit = 11;
    const uint32_t slot_size = SLOT_DATA_BUF_SIZE;
    std::vector<MpscQueueSlotData> v;
    MpscQueueSlot* slot_;
    std::atomic<uint32_t> head_ = {0};
    uint32_t padding[15];
    std::atomic<uint32_t> tail_ = {0};
    MpscQueue() {
        v.resize(size);
        slot_ = new MpscQueueSlot[size];
    }
    ~MpscQueue() {
        delete[] slot_;
    }
    void peek(uint32_t& i, char*& buf ) {
        i = head_.fetch_add(1, std::memory_order_acquire);
        auto & turn = slot_[MPSC_IDX(i,size_mask)].turn_;
        auto new_turn = MPSC_TURN(i, size_bit);
        while (new_turn != turn.load(std::memory_order_acquire));
        buf = v[MPSC_IDX(i,size_mask)].value;
    }
    void send(uint32_t& i) {
        auto & turn = slot_[MPSC_IDX(i,size_mask)].turn_;
        auto new_turn = MPSC_TURN(i, size_bit);
        turn.store(new_turn + 1, std::memory_order_release);
    }
    bool try_pop(char*& buf) {
        auto i = tail_.load(std::memory_order_relaxed);
        auto & turn = slot_[MPSC_IDX(i,size_mask)].turn_;
        if (MPSC_TURN(i, size_bit) + 1 != turn.load(std::memory_order_acquire)) {
            return false;
        }
        buf = v[MPSC_IDX(i,size_mask)].value;
        return true;

    }
    void ack() {
        auto i = tail_.fetch_add(1, std::memory_order_release);
        auto new_turn = MPSC_TURN(i, size_bit);
        auto & turn = slot_[MPSC_IDX(i,size_mask)].turn_;
        turn.store(new_turn + 2, std::memory_order_release);
    }
};

void prodecer_thread(MpscQueue* q) {
    uint32_t i;
    char* buf;
    char data[2048];
    uint32_t count = 0;
    while(true) {
        q->peek(i, buf);
        memcpy(buf, data, 2048);
        *((uint32_t*)buf) = count;
        count++;
        q->send(i);
    }
}

int64_t getUs() {
    struct timeval cur_time;
    gettimeofday(&cur_time, nullptr);
    return (cur_time.tv_sec*1000000+cur_time.tv_usec);
}

void consumer_thread(MpscQueue* q) {
    char* buf;
    char tmp[2048];
    uint64_t count = 0;
    int64_t start = getUs();
    std::stringstream print_str;
    int print_count = 0;
    while(true) {
        if(!q->try_pop(buf)) {
            continue;
        }
        memcpy(tmp, buf, 2048);
        q->ack();
        count++;
        if(count == 0xFFFFF) {
            print_count++;
            int64_t cur = getUs();
            print_str<<" consume "<< count << " records take "<< cur - start << "us speed "
                << (count*1000000)/(cur - start) <<" qps " << (count*2048)/(cur - start) << " MB/s" << std::endl;
            count = 0;
            start = cur;

            if(print_count == 5) {
                std::cout<<print_str.str() << std::endl;
                exit(0);
            }
        }
    }
}


int main()  {
    MpscQueue q;
    auto producer_1 = std::thread(prodecer_thread, &q);
    auto producer_2 = std::thread(prodecer_thread, &q);
    auto producer_3 = std::thread(prodecer_thread, &q);
    auto producer_4 = std::thread(prodecer_thread, &q);
    auto consumer = std::thread(consumer_thread, &q);
    producer_1.join();
    producer_2.join();
    producer_3.join();
    producer_4.join();
    consumer.join();
    
    return 0;
}

性能数据

root@ubuntu:/home/marco/Documents/code/mpsc# ./a.out
 consume 1048575 records take 59722us speed 17557600 qps 35957 MB/s
 consume 1048575 records take 61869us speed 16948310 qps 34710 MB/s
 consume 1048575 records take 61062us speed 17172300 qps 35168 MB/s
 consume 1048575 records take 61801us speed 16966958 qps 34748 MB/s
 consume 1048575 records take 61486us speed 17053882 qps 34926 MB/s