pubsub.h

Go to the documentation of this file.

#pragma once
 
#include "util/rats_export.h"
#include "node/peer_network.h"
#include "peer/peer.h"
#include "core/bytes.h"
#include "peer/peer_id.h"
 
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstdint>
#include <deque>
#include <functional>
#include <mutex>
#include <random>
#include <string>
#include <thread>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
 
namespace librats {
 
enum class ValidationResult { Accept, Reject, Ignore };
 
class RATS_API PubSub final : public Subsystem {
public:
    using Handler   = std::function<void(const PeerId& from, const std::string& topic, ByteView data)>;
    using Validator = std::function<ValidationResult(const PeerId& from, const std::string& topic, ByteView data)>;
 
    struct Config {
        int mesh_target   = 6;   
        int mesh_low      = 4;   
        int mesh_high     = 12;  
        int fanout_size   = 6;   
        int gossip_factor = 6;   
        std::chrono::milliseconds fanout_ttl{60000};         
        std::chrono::milliseconds heartbeat_interval{1000};  
        int    history_length = 5;     
        int    history_gossip = 3;     
        size_t seen_limit     = 8192;  
    };
 
    PubSub();
    explicit PubSub(Config config);
    ~PubSub() override;
 
    void subscribe(const std::string& topic, Handler handler);
    void unsubscribe(const std::string& topic);
 
    void publish(const std::string& topic, ByteView data);
 
    bool                     is_subscribed(const std::string& topic) const;
    std::vector<std::string> subscribed_topics() const;
    std::vector<PeerId>      peers_for_topic(const std::string& topic) const;  
    std::vector<PeerId>      mesh_peers(const std::string& topic) const;       
 
    void set_validator(const std::string& topic, Validator validator);
 
    // Subsystem.
    void attach(NodeContext& ctx) override;
    void start() override;  
    void stop() override;   
 
private:
    struct Topic {
        std::unordered_set<PeerId, PeerId::Hash> subscribers;  
        std::unordered_set<PeerId, PeerId::Hash> mesh;         
        std::unordered_set<PeerId, PeerId::Hash> fanout;       
        std::chrono::steady_clock::time_point    last_fanout{};
    };
 
    struct CachedMessage {
        std::string topic;
        Bytes       frame;  
    };
 
    // Inbound dispatch (all run on a reactor thread).
    void on_new_peer(const Peer& peer);
    void on_peer_gone(const PeerId& id);
    void on_gossip(const Peer& peer, ByteView payload);
 
    void recv_subscription(const PeerId& from, const std::string& topic, bool subscribe);
    void recv_graft(const PeerId& from, const std::string& topic);
    void recv_prune(const PeerId& from, const std::string& topic);
    void recv_publish(const PeerId& from, ByteView frame, const PeerId& origin, uint64_t seqno,
                      const std::string& topic, ByteView data);
    void recv_ihave(const PeerId& from, const std::string& topic, const std::vector<std::string>& ids);
    void recv_iwant(const PeerId& from, const std::vector<std::string>& ids);
 
    // Heartbeat (background thread).
    void heartbeat_loop();
    void do_heartbeat();
 
    using CtrlList = std::vector<std::pair<PeerId, std::string>>;
    void maintain_mesh_locked(const std::string& topic, CtrlList& grafts, CtrlList& prunes);
 
    // Helpers.
    void send_ctrl(const PeerId& to, uint8_t op, const std::string& topic);
    void broadcast_ctrl(uint8_t op, const std::string& topic);
    void deliver_local(const PeerId& from, const std::string& topic, ByteView data);
    Handler          handler_for(const std::string& topic) const;
    ValidationResult validate(const PeerId& from, const std::string& topic, ByteView data);
    bool             mark_seen(const std::string& id);             
    void             cache_message(const std::string& id, const std::string& topic, const Bytes& frame);
    std::vector<PeerId> random_sample(std::vector<PeerId> in, int n);
 
    PeerNetwork* network_ = nullptr;
    Config       config_;
 
    mutable std::mutex mutex_;  
    uint64_t           seqno_ = 0;
    std::unordered_map<std::string, Handler>   subscriptions_;  
    std::unordered_map<std::string, Topic>     topics_;         
    std::unordered_map<std::string, Validator> validators_;     
    Validator                                  global_validator_;
 
    mutable std::mutex mcache_mutex_;  
    std::unordered_map<std::string, CachedMessage> mcache_;     
    std::deque<std::vector<std::string>>           history_;    
    std::unordered_set<std::string>                seen_;       
    std::deque<std::string>                        seen_order_; 
 
    std::mutex   rng_mutex_;
    std::mt19937 rng_;
 
    std::atomic<bool>       running_{false};
    std::thread             heartbeat_thread_;
    std::mutex              hb_mutex_;
    std::condition_variable hb_cv_;
};
 
} // namespace librats