Towards/modules/message_protocol/src/MessageEndpoint.cpp

#include "message_protocol/MessageEndpoint.hpp"

#include "quicr/QuicrAddress.hpp"
#include "quicr/QuicrConnection.hpp"
#include "quicr/QuicrConnectionListener.hpp"
#include "quicr/QuicrEndpoint.hpp"

#include <spdlog/spdlog.h>

namespace tw::msg {

namespace {
constexpr size_t RECEIVE_BUFFER_SIZE = 64 * 1024;
}

MessageEndpoint::MessageEndpoint(std::unique_ptr<net::quicr::QuicrEndpoint> endpoint) :
    m_endpoint(std::move(endpoint)),
    m_receive_buffer(RECEIVE_BUFFER_SIZE) {
}

MessageEndpoint::~MessageEndpoint() = default;

tl::expected<std::unique_ptr<MessageEndpoint>, MessageError> MessageEndpoint::create() {
    auto endpoint_r = net::quicr::QuicrEndpoint::create();
    if(!endpoint_r) {
        return tl::make_unexpected(
            MessageError(MessageErrorType::BindFailed, endpoint_r.error().message()));
    }

    return std::unique_ptr<MessageEndpoint>(new MessageEndpoint(std::move(endpoint_r.value())));
}

tl::expected<std::unique_ptr<MessageEndpoint>, MessageError> MessageEndpoint::bind(int port) {
    auto endpoint_r = create();
    if(!endpoint_r) {
        return endpoint_r;
    }

    auto& endpoint = endpoint_r.value();

    auto bind_r = endpoint->m_endpoint->bind(port);
    if(!bind_r) {
        return tl::make_unexpected(MessageError(MessageErrorType::BindFailed, bind_r.error().message()));
    }

    auto listener_r = net::quicr::QuicrConnectionListener::listen(endpoint->m_endpoint.get());
    if(!listener_r) {
        return tl::make_unexpected(
            MessageError(MessageErrorType::BindFailed, listener_r.error().message()));
    }

    endpoint->m_listener = std::move(listener_r.value());

    return endpoint_r;
}

MessageConnection* MessageEndpoint::add_peer(net::quicr::QuicrConnection* connection) {
    const PeerId id = m_next_peer_id++;

    auto peer = std::make_unique<MessageConnection>(id, connection, &m_dispatcher);
    auto* raw = peer.get();

    m_peers.emplace(id, std::move(peer));

    if(m_on_peer_connected) {
        m_on_peer_connected(id);
    }

    return raw;
}

tl::expected<MessageConnection*, MessageError> MessageEndpoint::connect(const std::string& host, int port) {
    auto connection_r = m_endpoint->connect(net::quicr::QuicrAddress(host, port));
    if(!connection_r) {
        return tl::make_unexpected(
            MessageError(MessageErrorType::ConnectFailed, connection_r.error().message()));
    }

    return add_peer(connection_r.value());
}

void MessageEndpoint::accept_peers() {
    if(!m_listener) {
        return;
    }

    while(net::quicr::QuicrConnection* connection = m_listener->listen()) {
        add_peer(connection);
    }
}

void MessageEndpoint::receive() {
    for(auto& [id, peer] : m_peers) {
        m_bytes_received += peer->receive(m_receive_buffer);
    }
}

void MessageEndpoint::update() {
    m_endpoint->poll();

    accept_peers();
    receive();

    const auto now = std::chrono::steady_clock::now();
    for(auto& [id, peer] : m_peers) {
        peer->expire_requests(now);
    }
}

MessageConnection* MessageEndpoint::peer(PeerId id) {
    auto peer = m_peers.find(id);
    return peer != m_peers.end() ? peer->second.get() : nullptr;
}

std::vector<MessageConnection*> MessageEndpoint::peers() const {
    std::vector<MessageConnection*> result;
    result.reserve(m_peers.size());

    for(const auto& [id, peer] : m_peers) {
        result.push_back(peer.get());
    }

    return result;
}

void MessageEndpoint::broadcast(MessageType type, std::span<const std::byte> body, bool reliable) {
    for(auto& [id, peer] : m_peers) {
        auto send_r = peer->send(type, body, reliable);
        if(!send_r) {
            spdlog::error("Failed to send to peer {}: {}", id, send_r.error().message());
        }
    }
}

uint64_t MessageEndpoint::bytes_sent() const {
    uint64_t total = 0;
    for(const auto& [id, peer] : m_peers) {
        total += peer->bytes_sent();
    }

    return total;
}

uint64_t MessageEndpoint::messages_sent() const {
    uint64_t total = 0;
    for(const auto& [id, peer] : m_peers) {
        total += peer->messages_sent();
    }

    return total;
}

uint64_t MessageEndpoint::messages_received() const {
    uint64_t total = 0;
    for(const auto& [id, peer] : m_peers) {
        total += peer->messages_received();
    }

    return total;
}

}