用 process_vm_readv() 实现多进程内存队列

#include <assert.h>
#include <fcntl.h>
#include <mqueue.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/uio.h>
#include <sys/wait.h>
#include <unistd.h>

#include <iostream>
#include <sstream>
#include <vector>
using namespace std;

struct MemMessage {
    pid_t pid = 0;
    vector<iovec> mem_vec;

public:
    MemMessage() {}

    string DebugString() const {
        ostringstream oss;
        oss << "pid=" << pid;
        for (const auto& mem : mem_vec) {
            oss << " iov={" << mem.iov_base << "," << mem.iov_len << "}";
        }
        return oss.str();
    }

    string ToString() const {
        pid_t from_pid = getpid();
        string buff;
        buff.append((const char*)&from_pid, sizeof(from_pid));
        for (const auto& mem : mem_vec) {
            buff.append((const char*)&mem, sizeof(mem));
        }
        return buff;
    }

    bool FromString(const char* buff, size_t buff_len) {
        size_t pos = 0;
        if (buff_len < sizeof(pid)) {
            return false;
        }
        pid = *((pid_t*)buff);
        pos += sizeof(pid);
        for (; pos + sizeof(iovec) <= buff_len; pos += sizeof(iovec)) {
            mem_vec.resize(mem_vec.size() + 1);
            auto& iov = mem_vec.back();
            memcpy(reinterpret_cast<char*>(&iov), buff + pos, sizeof(iov));
        }
        return true;
    }

    void CopyToMem(string& buff) {
        size_t sum = 0;
        for (const auto& iov : mem_vec) {
            sum += iov.iov_len;
        }

        buff.resize(sum, '\0');
        iovec local;
        local.iov_base = &buff[0];
        local.iov_len = buff.size();

        auto nread = process_vm_readv(pid, &local, 1, &mem_vec[0], mem_vec.size(), 0);
        // assert(nread >= 0);
        if (nread <= 0) {
            cout << getpid() << " process_vm_readv from=" << DebugString() << " sum=" << sum << " len=" << nread
                 << endl;
            cout << "pid=" << pid << " " << strerror(errno) << endl;
        }
    }
};

class MemQueue {
    mqd_t m_mq;
    struct mq_attr m_attr;

public:
    int Init(const string& queue_name, bool read_write = false) {
        memset(&m_attr, 0, sizeof(m_attr));
        m_attr.mq_flags = 0;
        m_attr.mq_maxmsg = 10;
        m_attr.mq_msgsize = 4096;
        const int oflag = read_write ? O_WRONLY | O_CREAT : O_RDONLY;
        m_mq = mq_open(queue_name.c_str(), oflag, S_IRWXU, &m_attr);
        // m_mq = mq_open(queue_name.c_str(), oflag);
        cout << getpid() << " mq_open=" << queue_name << " oflag=" << oflag << " ret=" << m_mq << endl;
    }

    int PushMemMsg(const MemMessage& msg) {
        string buff = msg.ToString();
        return mq_send(m_mq, buff.c_str(), buff.size(), 0);
    }

    bool PopMemMsg(MemMessage& msg) {
        string buff(m_attr.mq_msgsize, '\0');
        unsigned prio = 0;
        const auto ret = mq_receive(m_mq, &buff[0], buff.size(), &prio);
        if (ret < 0) {
            return false;
        }
        buff.resize(ret);
        msg.FromString(buff.data(), buff.size());
        return true;
    }
};

const static string queue_name = "/test_queue";

enum STATUS { STATUS_WAITING = 0, STATUS_RUNNING = 1, STATUS_EXIT = 2 };

struct ControlCenter {
    int proc_status;
};

ControlCenter* control = nullptr;

void child_proc() {
    const auto ret = fork();
    assert(ret >= 0);
    if (ret != 0) {
        return;
    }

    while (STATUS_WAITING == control->proc_status) {
        usleep(1000);
    }

    MemQueue q;
    q.Init(queue_name);
    size_t bytes = 0;
    while (STATUS_RUNNING == control->proc_status) {
        MemMessage msg;
        if (q.PopMemMsg(msg)) {
            string buff;
            msg.CopyToMem(buff);
            bytes += buff.size();
        } else {
            usleep(1000);
        }
    }
    cout << getpid() << " bytes=" << bytes << " exit" << endl;

    exit(0);
}

int main(int argc, char* argv[]) {
    control = (ControlCenter*)mmap(NULL, sizeof(*control), PROT_READ | PROT_WRITE, MAP_SHARED | MAP_ANONYMOUS, -1, 0);
    control->proc_status = STATUS_WAITING;

    const static int proc_num = 3;
    for (int i = 0; i < proc_num; ++i) {
        child_proc();
    }

    MemQueue q;
    q.Init(queue_name, true);

    control->proc_status = STATUS_RUNNING;
    string str(1024 * 1024 * 10, 'a');
    for (int i = 0; i < 10000; ++i) {
        MemMessage msg;
        for (int t = 0; t < 3; ++t) {
            size_t base = rand() % str.size();
            size_t size = rand() % (str.size() - base);
            msg.mem_vec.emplace_back(iovec{&str[base], size});
        }
        const auto ret = q.PushMemMsg(msg);
        // cout << getpid() << " PushMemMsg=" << msg.DebugString() << " ret=" << ret << endl;

        // sleep(1);
    }

    control->proc_status = STATUS_EXIT;

    for (int i = 0; i < proc_num; ++i) {
        int status = 0;
        wait(&status);
    }

    return 0;
}