<commit_before>#include "zcm.h"
#include <zmq.h>

#include <unistd.h>

#include <cstring>
#include <cassert>
#include <cctype>
#include <cstdio>

#include <vector>
using std::vector;

#include <unordered_map>
using std::unordered_map;

#include <string>
using std::string;

#include <mutex>
#include <thread>
#include <iostream>

#define METADATA_PUB_ADDR "ipc:///tmp/zcm-metadata-pub"
#define METADATA_SUB_ADDR "ipc:///tmp/zcm-metadata-sub"
#define ZMQ_IO_THREADS 1

struct Sub
{
    zcm_callback_t *cb;
    void *usr;
};

struct zcm_t
{
    std::thread recvThread;

    // The mutex protects all data below it
    std::mutex mut;
    bool recvThreadStarted = false;
    bool running = true;
    unordered_map<string, Sub> subs;

    void *ctx;
    unordered_map<string, void*> subsocks;
    unordered_map<string, void*> pubsocks;

    // TEMPORARY
    string subchannel;
    void *subsock;

    zcm_t()
    {
        ctx = zmq_init(ZMQ_IO_THREADS);
    }

    string _getAddress(const string& channel)
    {
        return "ipc:///tmp/zmq-channel-"+channel;
    }

    void *_getSubSock(const string& channel)
    {
        auto it = subsocks.find(channel);
        if (it != subsocks.end())
            return it->second;
        void *sock = zmq_socket(ctx, ZMQ_SUB);
        string address = _getAddress(channel);
        zmq_connect(sock, address.c_str());
        zmq_setsockopt(sock, ZMQ_SUBSCRIBE, "", 0);
        subsocks[channel] = sock;

        // TEMPORARY
        subchannel = channel;
        subsock = sock;

        return sock;
    }

    void *_getPubSock(const string& channel)
    {
        auto it = pubsocks.find(channel);
        if (it != pubsocks.end())
            return it->second;
        void *sock = zmq_socket(ctx, ZMQ_PUB);
        string address = _getAddress(channel);
        zmq_bind(sock, address.c_str());
        pubsocks[channel] = sock;
        return sock;
    }

    int publish(const string& channel, const char *data, size_t len)
    {
        std::unique_lock<std::mutex> lk(mut);

        void *sock = _getPubSock(channel);
        int rc = zmq_send(sock, data, len, 0);
        assert(rc == (int)len);
        return 0;
    }

    void recvMsgFromSock(const string& channel, void *sock)
    {
        const int BUFSZ = 1 << 20;
        char *buf = (char *) malloc(BUFSZ);

        int rc = zmq_recv(sock, buf, BUFSZ, 0);
        assert(0 < rc && rc < BUFSZ);

        // Dispatch
        {
            std::unique_lock<std::mutex> lk(mut);
            auto it = subs.find(channel);
            if (it != subs.end()) {
                auto& sub = it->second;
                zcm_recv_buf_t rbuf;
                rbuf.data = buf;
                rbuf.len = rc;
                rbuf.utime = 0;
                rbuf.zcm = this;
                sub.cb(&rbuf, channel.c_str(), sub.usr);
            }
        }
    }

    void recvThreadFunc()
    {
        while (1) {
            // Build up a list of poll items
            vector<zmq_pollitem_t> pitems;
            vector<string> pchannels;
            {
                std::unique_lock<std::mutex> lk(mut);
                pitems.resize(subsocks.size());
                int i = 0;
                for (auto& elt : subsocks) {
                    auto& channel = elt.first;
                    auto& sock = elt.second;
                    auto *p = &pitems[i];
                    memset(p, 0, sizeof(*p));
                    p->socket = sock;
                    p->events = ZMQ_POLLIN;
                    pchannels.emplace_back(channel);
                    i++;
                }
            }

            int rc = zmq_poll(pitems.data(), pitems.size(), 100);
            if (!running)
                break;
            if (rc >= 0) {
                for (size_t i = 0; i < pitems.size(); i++) {
                    auto& p = pitems[i];
                    if (p.revents != 0)
                        recvMsgFromSock(pchannels[i], p.socket);
                }
            }
        }
    }

    int subscribe(const string& channel, zcm_callback_t *cb, void *usr)
    {
        std::unique_lock<std::mutex> lk(mut);
        _getSubSock(channel);
        subs.emplace(channel, Sub{cb, usr});
        if (!recvThreadStarted) {
            recvThread = std::thread{&zcm_t::recvThreadFunc, this};
            recvThreadStarted = true;
        }
        return 0;
    }
};

static void printBytes(const char *buf, size_t len)
{
    printf("printBytes:");
    for (size_t i = 0; i < len; i++) {
        char c = buf[i];
        printf(" 0x%02x", c&0xff);
        if (isalpha(c))
            printf("(%c)", c);
    }
    printf("\n");
}

zcm_t *zcm_create(void)
{
    return new zcm_t();
}

void zcm_destroy(zcm_t *zcm)
{
    // TODO: What is the "correct" method to shut-down ZMQ?
}
<commit_msg>Implemented a subscribe all without using any crossbar! Still some known issues with the impl such as not supporting multiple subscriptions on the same channel, no unsubscribe, no real regex (just subscribe all via .*)<commit_after>#include "zcm.h"
#include <zmq.h>

#include <unistd.h>
#include <dirent.h>

#include <cstring>
#include <cassert>
#include <cctype>
#include <cstdio>

#include <vector>
using std::vector;

#include <unordered_map>
using std::unordered_map;

#include <string>
using std::string;

#include <mutex>
#include <thread>
#include <iostream>

#define METADATA_PUB_ADDR "ipc:///tmp/zcm-metadata-pub"
#define METADATA_SUB_ADDR "ipc:///tmp/zcm-metadata-sub"
#define ZMQ_IO_THREADS 1

struct Sub
{
    zcm_callback_t *cb;
    void *usr;
    void *sock;
};

struct SubRegex
{
    string channelRegex;
    zcm_callback_t *cb;
    void *usr;
};

struct Pub
{
    void *sock;
};

struct zcm_t
{
    std::thread recvThread;

    // The mutex protects all data below it
    std::mutex mut;
    bool recvThreadStarted = false;
    bool running = true;

    void *ctx;
    unordered_map<string, Sub> subs;
    unordered_map<string, Pub> pubs;

    vector<SubRegex> subRegex;

    zcm_t()
    {
        ctx = zmq_init(ZMQ_IO_THREADS);
    }

    string getIPCAddress(const string& channel)
    {
        return "ipc:///tmp/zmq-channel-"+channel;
    }

    Sub *findSub(const string& channel)
    {
        auto it = subs.find(channel);
        if (it != subs.end())
            return &it->second;
        else
            return nullptr;
    }

    Pub *findPub(const string& channel)
    {
        auto it = pubs.find(channel);
        if (it != pubs.end())
            return &it->second;
        else
            return nullptr;
    }

    int publish(const string& channel, const char *data, size_t len)
    {
        std::unique_lock<std::mutex> lk(mut);

        Pub *pub = findPub(channel);
        if (!pub) {
            void *sock = zmq_socket(ctx, ZMQ_PUB);
            string address = getIPCAddress(channel);
            zmq_bind(sock, address.c_str());
            pubs.emplace(channel, Pub{sock});
            pub = findPub(channel);
            assert(pub);
        }

        int rc = zmq_send(pub->sock, data, len, 0);
        assert(rc == (int)len);
        return 0;
    }

    void searchForRegexSubs()
    {
        std::unique_lock<std::mutex> lk(mut);

        // TODO: actually implement regex, or remove it
        if (subRegex.size() == 0)
            return;
        for (auto& r : subRegex) {
            if (r.channelRegex != ".*") {
                static bool warned = false;
                if (!warned) {
                    fprintf(stderr, "Er: only .* (aka subscribe-all) is implemented\n");
                    warned = true;
                }
            }
        }
        SubRegex& srex = subRegex[0];

        const char *prefix = "zmq-channel-";
        size_t prefixLen = strlen(prefix);

        DIR *d;
        dirent *ent;

        if (!(d=opendir("/tmp/")))
            return;

        while ((ent=readdir(d)) != nullptr) {
            if (strncmp(ent->d_name, prefix, prefixLen) == 0) {
                string channel(ent->d_name + prefixLen);
                if (!findSub(channel))
                    subOneInternal(channel, srex.cb, srex.usr);
            }
        }

        closedir(d);
    }

    void recvMsgFromSock(const string& channel, void *sock)
    {
        const int BUFSZ = 1 << 20;
        char *buf = (char *) malloc(BUFSZ);

        int rc = zmq_recv(sock, buf, BUFSZ, 0);
        assert(0 < rc && rc < BUFSZ);

        // Dispatch
        {
            std::unique_lock<std::mutex> lk(mut);
            if (Sub *sub = findSub(channel)) {
                zcm_recv_buf_t rbuf;
                rbuf.data = buf;
                rbuf.len = rc;
                rbuf.utime = 0;
                rbuf.zcm = this;
                sub->cb(&rbuf, channel.c_str(), sub->usr);
            }
        }
    }

    void recvThreadFunc()
    {
        while (1) {
            searchForRegexSubs();

            // Build up a list of poll items
            vector<zmq_pollitem_t> pitems;
            vector<string> pchannels;
            {
                std::unique_lock<std::mutex> lk(mut);
                pitems.resize(subs.size());
                int i = 0;
                for (auto& elt : subs) {
                    auto& channel = elt.first;
                    auto& sub = elt.second;
                    auto *p = &pitems[i];
                    memset(p, 0, sizeof(*p));
                    p->socket = sub.sock;
                    p->events = ZMQ_POLLIN;
                    pchannels.emplace_back(channel);
                    i++;
                }
            }

            int rc = zmq_poll(pitems.data(), pitems.size(), 100);
            if (!running)
                break;
            if (rc >= 0) {
                for (size_t i = 0; i < pitems.size(); i++) {
                    auto& p = pitems[i];
                    if (p.revents != 0)
                        recvMsgFromSock(pchannels[i], p.socket);
                }
            }
        }
    }

    static bool isRegexChannel(const string& channel)
    {
        // These chars are considered regex
        auto isRegexChar = [](char c) {
            return c == '(' || c == ')' || c == '|' ||
                   c == '.' || c == '*' || c == '+';
        };

        for (auto& c : channel)
            if (isRegexChar(c))
                return true;

        return false;
    }

    void subOneInternal(const string& channel, zcm_callback_t *cb, void *usr)
    {
        assert(!findSub(channel));

        void *sock = zmq_socket(ctx, ZMQ_SUB);
        string address = getIPCAddress(channel);
        zmq_connect(sock, address.c_str());
        zmq_setsockopt(sock, ZMQ_SUBSCRIBE, "", 0);

        subs.emplace(channel, Sub{cb, usr, sock});
    }

    int subscribe(const string& channel, zcm_callback_t *cb, void *usr)
    {
        std::unique_lock<std::mutex> lk(mut);

        if (isRegexChannel(channel)) {
            subRegex.emplace_back(SubRegex{channel, cb, usr});
        } else {
            subOneInternal(channel, cb, usr);
        }

        if (!recvThreadStarted) {
            recvThread = std::thread{&zcm_t::recvThreadFunc, this};
            recvThreadStarted = true;
        }
        return 0;
    }
};

static void printBytes(const char *buf, size_t len)
{
    printf("printBytes:");
    for (size_t i = 0; i < len; i++) {
        char c = buf[i];
        printf(" 0x%02x", c&0xff);
        if (isalpha(c))
            printf("(%c)", c);
    }
    printf("\n");
}

zcm_t *zcm_create(void)
{
    return new zcm_t();
}

void zcm_destroy(zcm_t *zcm)
{
    // TODO: What is the "correct" method to shut-down ZMQ?
}

{'fraction_non_alphanumeric': 0.12953054884363135, 'fraction_numerical': 0.004573696927856403, 'mean_word_length': 1.6907360111446483, 'pattern_counts': {'":': 0, '<': 59, '<?xml version=': 0, '>': 63, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3570595', 'n_tokens_mistral': 3910, 'n_tokens_neox': 3708, 'n_words': 875}