shmem_exchange.h

Go to the documentation of this file.

// Copyright (c) Space-ng, inc. All rights reserved.
 
#pragma once
#include "sol3/core/message_port.h"
#include "sol3/core/msg/exchange_config.h"
#include "sol3/core/shmem_buffer.h"
#include "sol3/cpp/thread_safe_value.h"
 
#include <boost/asio.hpp>
#include <boost/asio/local/datagram_protocol.hpp>
 
#include <atomic>
#include <cstddef>
#include <cstdint>
#include <map>
#include <memory>
#include <set>
#include <string>
#include <unordered_set>
#include <utility>
#include <vector>
namespace sol3::core {
 
enum ExchangeCode : uint32_t {
  buffers_request = 0,
  buffers_response = 1,
  buffer_removed = 2,
  peer_disconnect = 3,
};
 
msg::ExchangeConfigT loadExchangeConfigFromEnv();
 
msg::ExchangeConfigT loadExchangeConfigFromPath(std::string const& config_path);
 
class ShmemExchange : public IBufferExchange {
 public:
  using Endpoint = IBufferExchange::Endpoint;
  using EndpointPort = sol3::core::EndpointPort;
  using PortIdx = sol3::core::PortIdx;
 
  ShmemExchange(Endpoint const& local_endpoint);
 
  ShmemExchange(
      Endpoint const& local_endpoint, msg::ExchangeConfigT exchange_config);
  ~ShmemExchange() override;
 
  // Non-copyable and non-movable: ShmemExchange owns IO and socket resources.
  ShmemExchange(ShmemExchange const&) = delete;
  ShmemExchange& operator=(ShmemExchange const&) = delete;
  ShmemExchange(ShmemExchange&&) = delete;
  ShmemExchange& operator=(ShmemExchange&&) = delete;
 
  void start() override;
  void stop() override;
  bool stopped() const override;
  void addPeer(Endpoint const& peer_endpoint) override;
  void mapLocalToPeer(
      uint32_t local_port,
      Endpoint peer_endpoint,
      uint32_t peer_port = 0) override;
  void configure(msg::ExchangeConfigT const& config) override;
  void registerBuffer(IBufferMutable const& buffer) override;
  void unregisterBuffer(IBufferMutable const& buffer) override;
  Endpoint const& localEndpoint() const override { return local_endpoint_; }
  std::shared_ptr<IBufferConst> get(uint32_t port, uint32_t idx) const override;
  std::shared_ptr<IBufferConst> get(
      Endpoint const& ep, uint32_t port, uint32_t idx) const override;
  void dispose(std::shared_ptr<IBufferConst>&& buffer) override;
  void listPorts(std::vector<EndpointPort>& ports) const override;
  void listMappings(
      std::map<uint32_t, std::vector<std::pair<Endpoint, uint32_t>>>& mappings)
      const override;
  void listResolvedBuffers(
      std::map<PortIdx, std::shared_ptr<IBufferConst>>& buffers) const override;
 
 private:
  struct ExchangePacket {
    ExchangeCode what;
    uint64_t peer_session_id;
  } __attribute__((packed));
 
  static size_t constexpr kMaxRecvBuffer = 8192;
  static size_t constexpr kMaxPayload = sizeof(uint64_t);
  static size_t constexpr kMaxMissedPollsBeforeEvict = 4;
  static size_t constexpr kMaxBuffers =
      (kMaxRecvBuffer - sizeof(ExchangePacket)) / sizeof(BufferInfo);
  // Process-lifetime peer identity for protocol sessions.
  // Not a UUID: generated randomly at startup and changes on restart.
  using PeerSessionId = uint64_t;
 
  using Socket = boost::asio::local::datagram_protocol::socket;
  using ShmemBufferMap = std::map<PortIdx, std::shared_ptr<IBufferConst>>;
 
  // State information about downstream peers that are requesting
  // buffers from this exchange.
  struct SentState {
    // buffers that have been registered to share.
    ShmemBufferMap buffers_to_share;
 
    // A list of all peers that we're interacting with.
    std::vector<Endpoint> peer_endpoints;
 
    // Track the last seen peer session id for each polling endpoint.
    std::map<Endpoint, PeerSessionId> peer_session_ids;
 
    // Track which buffers we've sent to each requesting peer endpoint.
    std::map<Endpoint, std::unordered_set<PortIdx>> sent_buffers_per_peer;
  };
 
  // State information and buffers that are received from other peer endpoints
  // in response to this exchange's explicit polling.
  struct ReceivedState {
    Endpoint local_endpoint;
    // Track the last seen peer session id for each polled endpoint.
    // Includes this exchange's endpoint.
    std::map<Endpoint, PeerSessionId> peer_session_ids;
 
    // A list of all peers that we're interacting with.
    // Includes this exchange's endpoint.
    std::set<Endpoint> peer_endpoints;
 
    // Buffers that have been shared by each peer that we're polling
    // Note: this will also include buffers that are registered with the
    // exchange.
    std::map<Endpoint, ShmemBufferMap> buffers_per_peer;
 
    // Track consecutive failed polls for each peer endpoint. Used to evict
    // stale buffers when a preferred peer dies so fallback peers can take over.
    std::map<Endpoint, uint32_t> peer_missed_polls;
 
    // Resolution mappings:
    // - local_port != 0: explicit local-port to peer-port mappings.
    // - local_port == 0: wildcard fallback mappings for any local port.
    std::map<uint32_t, std::vector<std::pair<Endpoint, uint32_t>>>
        local_port_to_peer_ports;
 
    struct ResolvedCacheEntry {
      std::shared_ptr<IBufferConst> buffer;
      uint64_t generation = 0;
    };
 
    struct GetResolvedResult {
      PortIdx port_idx;
      std::shared_ptr<IBufferConst> buffer;
      bool cache_needs_update = false;
      uint64_t generation = 0;
    };
 
    // Cache of resolved buffers, invalidated by resolve_generation.
    std::map<PortIdx, ResolvedCacheEntry> resolved_buffers;
    // Monotonically increasing counter; bumped whenever the authoritative
    // buffer state changes, which obsoletes the resolved_buffers cache.
    uint64_t resolve_generation = 0;
 
    // Bump the generation counter, invalidating all resolved_buffers entries.
    void bumpResolvedGeneration();
 
    // Resolve from a mapping bucket in local_port_to_peer_ports. When
    // mapping_local_port == 0 and mapped peer_port == 0, peer port falls back
    // to requested_local_port.
    std::shared_ptr<IBufferConst> resolveFromPortMappings(
        uint32_t mapping_local_port,
        uint32_t requested_local_port,
        uint32_t idx) const;
 
    // Walk explicit port mappings, then wildcard fallback mappings
    // (local_port_to_peer_ports[0]), then finally local endpoint, returning
    // the first matching buffer for (port, idx).
    std::shared_ptr<IBufferConst> resolveEntry(PortIdx local_port_idx) const;
 
    // Shared-lock lookup for (port, idx). Returns resolved buffer and whether
    // the caller should schedule a deferred cache update.
    GetResolvedResult getResolved(uint32_t port, uint32_t idx) const;
 
    // Exclusive-lock cache write for deferred updates from getResolved().
    void updateResolvedCache(GetResolvedResult const& resolved);
 
    std::shared_ptr<IBufferConst> get(
        Endpoint const& ep, uint32_t port, uint32_t idx) const;
  };
 
  template <typename TFunction>
  auto accessReceivedState(TFunction&& cb) const {
    return received_state_.access(std::forward<TFunction>(cb));
  }
  void run();
  void registerBuffer(std::shared_ptr<IBufferConst> const& buffer);
  void unregisterBuffer(std::shared_ptr<IBufferConst> const& buffer);
  void unregisterBuffer(BufferInfo const& info);
  void startReceive();
  void startPeerPoll();
  void handleReceive();
  void sendSomeBuffers(Endpoint const& sender_ep);
  void requestPeerBuffers(Endpoint const& ep);
  bool sendPacketOnStrand(
      Endpoint const& ep,
      ExchangeCode what,
      uint8_t const* payload = nullptr,
      size_t payload_size = 0);
  void sendExchangePacket(
      Endpoint const& ep,
      ExchangeCode what,
      void const* payload = nullptr,
      size_t payload_size = 0);
 
  void handleBufferRequest(
      Endpoint const& sender_ep,
      std::array<char, ShmemExchange::kMaxRecvBuffer> const& msgbuf);
  void handleIncomingBuffers(
      Endpoint const& sender_ep,
      std::array<char, ShmemExchange::kMaxRecvBuffer> const& msgbuf,
      const struct msghdr& msgh);
  void handleBufferRemoved(
      Endpoint const& sender_ep,
      std::array<char, ShmemExchange::kMaxRecvBuffer> const& msgbuf);
  void handlePeerDisconnect(
      Endpoint const& sender_ep,
      std::array<char, ShmemExchange::kMaxRecvBuffer> const& msgbuf);
  void sendBufferRemoved(Endpoint const& ep, PortIdx port_idx);
  void sendDisconnect();
 
  PeerSessionId peer_session_id_;
  boost::asio::io_context ctx_;
  boost::asio::io_context::strand strand_;
  boost::asio::steady_timer peer_poll_timer_;
  Endpoint local_endpoint_;
  Socket socket_;
  cpp::ThreadSafeValue<ReceivedState> received_state_;
  cpp::ThreadSafeValue<SentState> sent_state_;
  std::atomic<bool> stop_requested_{false};
 
  std::thread run_thread_;
};
 
}  // namespace sol3::core