client.hpp

/* -*- mode: c++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*-
 *
 * Copyright (C) 2025 Marco Craveiro <marco.craveiro@gmail.com>
 *
 * This program is free software; you can redistribute it and/or modify it under
 * the terms of the GNU General Public License as published by the Free Software
 * Foundation; either version 3 of the License, or (at your option) any later
 * version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
 * details.
 *
 * You should have received a copy of the GNU General Public License along with
 * this program; if not, write to the Free Software Foundation, Inc., 51
 * Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 *
 */
#ifndef ORES_COMMS_NET_CLIENT_HPP
#define ORES_COMMS_NET_CLIENT_HPP
 
#include <mutex>
#include <atomic>
#include <thread>
#include <memory>
#include <iosfwd>
#include <cstdint>
#include <concepts>
#include <expected>
#include <functional>
#include <string_view>
#include <condition_variable>
#include <boost/asio/awaitable.hpp>
#include <boost/asio/ssl.hpp>
#include <boost/asio/strand.hpp>
#include <boost/asio/ip/tcp.hpp>
#include <boost/asio/io_context.hpp>
#include "ores.logging/make_logger.hpp"
#include "ores.eventing/service/event_bus.hpp"
#include "ores.comms/net/client_options.hpp"
#include "ores.comms/net/connection.hpp"
#include "ores.comms/net/pending_request_map.hpp"
#include "ores.comms/messaging/frame.hpp"
#include "ores.comms/messaging/message_traits.hpp"
#include "ores.comms/recording/session_recorder.hpp"
 
namespace ores::comms::net {
 
enum class connection_state {
    disconnected,   
    connecting,     
    connected,      
    reconnecting    
};
 
std::ostream& operator<<(std::ostream& s, connection_state v);
 
using disconnect_callback_t = std::function<void()>;
 
using reconnecting_callback_t = std::function<void()>;
 
using reconnected_callback_t = std::function<void()>;
 
using notification_callback_t = std::function<void(
    const std::string& event_type,
    std::chrono::system_clock::time_point timestamp,
    const std::vector<std::string>& entity_ids)>;
 
class client final {
private:
    inline static std::string_view logger_name = "ores.comms.net.client";
 
    static auto& lg() {
        using namespace ores::logging;
        static auto instance = make_logger(logger_name);
        return instance;
    }
 
    void setup_ssl_context();
 
    boost::asio::awaitable<void> perform_handshake();
 
    boost::asio::awaitable<void> run_heartbeat();
 
    boost::asio::awaitable<void> run_message_loop();
 
    boost::asio::awaitable<void> run_reconnect_loop();
 
    boost::asio::awaitable<void> perform_connection();
 
    std::chrono::milliseconds calculate_backoff(std::uint32_t attempt) const;
 
    void notify_shutdown_if_complete();
 
    boost::asio::awaitable<void> write_frame(const messaging::frame& f);
 
    std::uint32_t next_correlation_id();
 
public:
    explicit client(client_options config,
        std::shared_ptr<eventing::service::event_bus> event_bus = nullptr);
 
    explicit client(client_options config, boost::asio::any_io_executor executor,
        std::shared_ptr<eventing::service::event_bus> event_bus = nullptr);
 
    ~client();
 
    boost::asio::awaitable<void> connect();
 
    void connect_sync();
 
    boost::asio::awaitable<void> connect_with_retry();
 
    void connect_with_retry_sync();
 
    void disconnect();
 
    bool await_shutdown(std::chrono::milliseconds timeout = std::chrono::milliseconds{5000});
 
    bool is_connected() const;
 
    connection_state get_state() const;
 
    void set_disconnect_callback(disconnect_callback_t callback);
 
    void set_reconnecting_callback(reconnecting_callback_t callback);
 
    void set_reconnected_callback(reconnected_callback_t callback);
 
    void set_notification_callback(notification_callback_t callback);
 
    // =========================================================================
    // Session Recording
    // =========================================================================
 
    std::expected<std::filesystem::path, recording::session_file_error>
    enable_recording(const std::filesystem::path& output_directory);
 
    void disable_recording();
 
    bool is_recording() const;
 
    boost::asio::awaitable<std::expected<messaging::frame, ores::utility::serialization::error_code>>
    send_request(messaging::frame request_frame);
 
    std::expected<messaging::frame, ores::utility::serialization::error_code>
    send_request_sync(messaging::frame request_frame);
 
    template <typename RequestType, typename ResponseType,
              messaging::message_type RequestMsgType>
    std::expected<ResponseType, ores::utility::serialization::error_code>
    process_request(RequestType request) {
        using namespace ores::logging;
 
        auto payload = request.serialize();
        messaging::frame request_frame(RequestMsgType, 0, std::move(payload));
 
        auto result = send_request_sync(std::move(request_frame));
        if (!result) {
            return std::unexpected(result.error());
        }
 
        auto response_payload = result->decompressed_payload();
        if (!response_payload) {
            BOOST_LOG_SEV(lg(), error) << "Failed to decompress response";
            return std::unexpected(response_payload.error());
        }
 
        auto response = ResponseType::deserialize(*response_payload);
        if (!response) {
            BOOST_LOG_SEV(lg(), error) << "Failed to deserialize response";
            return std::unexpected(response.error());
        }
 
        return *response;
    }
 
    template <typename RequestType, typename ResponseType,
              messaging::message_type RequestMsgType>
    boost::asio::awaitable<std::expected<ResponseType, ores::utility::serialization::error_code>>
    process_request_async(RequestType request) {
        using namespace ores::logging;
 
        auto payload = request.serialize();
        messaging::frame request_frame(RequestMsgType, 0, std::move(payload));
 
        auto result = co_await send_request(std::move(request_frame));
        if (!result) {
            co_return std::unexpected(result.error());
        }
 
        auto response_payload = result->decompressed_payload();
        if (!response_payload) {
            BOOST_LOG_SEV(lg(), error) << "Failed to decompress response";
            co_return std::unexpected(response_payload.error());
        }
 
        auto response = ResponseType::deserialize(*response_payload);
        if (!response) {
            BOOST_LOG_SEV(lg(), error) << "Failed to deserialize response";
            co_return std::unexpected(response.error());
        }
 
        co_return *response;
    }
 
    // =========================================================================
    // Traits-based process_request overloads
    // =========================================================================
    // These overloads use message_traits to infer the response type and
    // message_type enum from the request type, simplifying the API.
 
    template <typename RequestType>
        requires messaging::has_message_traits<RequestType>
    std::expected<typename messaging::message_traits<RequestType>::response_type,
                  ores::utility::serialization::error_code>
    process_request(RequestType request) {
        using traits = messaging::message_traits<RequestType>;
        return process_request<
            RequestType,
            typename traits::response_type,
            traits::request_message_type>(std::move(request));
    }
 
    template <typename RequestType>
        requires messaging::has_message_traits<RequestType>
    boost::asio::awaitable<
        std::expected<typename messaging::message_traits<RequestType>::response_type,
                      ores::utility::serialization::error_code>>
    process_request_async(RequestType request) {
        using traits = messaging::message_traits<RequestType>;
        return process_request_async<
            RequestType,
            typename traits::response_type,
            traits::request_message_type>(std::move(request));
    }
 
private:
    client_options config_;
    std::unique_ptr<boost::asio::io_context> io_ctx_; // Owned io_context for sync operations
    boost::asio::any_io_executor executor_;
    boost::asio::ssl::context ssl_ctx_;
    std::unique_ptr<connection> conn_;
    std::uint32_t sequence_number_;
    std::atomic<connection_state> state_;
    mutable std::mutex state_mutex_; // Protects conn_, sequence_number_, callbacks, and recorder_
    disconnect_callback_t disconnect_callback_;
    reconnecting_callback_t reconnecting_callback_;
    reconnected_callback_t reconnected_callback_;
    notification_callback_t notification_callback_;
    std::shared_ptr<eventing::service::event_bus> event_bus_; // Optional event bus for connection events
 
    // Infrastructure for unified message loop
    std::unique_ptr<boost::asio::strand<boost::asio::any_io_executor>> write_strand_;
    std::atomic<bool> write_in_progress_{false}; // Async spinlock for serializing SSL writes
    std::unique_ptr<pending_request_map> pending_requests_;
    std::atomic<std::uint32_t> correlation_id_counter_{1};
    std::atomic<bool> message_loop_running_{false};
    std::atomic<bool> reconnect_loop_running_{false};
    std::atomic<bool> heartbeat_loop_running_{false};
    std::atomic<bool> shutdown_requested_{false};
    std::condition_variable shutdown_cv_;
    std::mutex shutdown_mutex_;
    std::unique_ptr<std::thread> io_thread_; // Background thread for io_context
    std::unique_ptr<boost::asio::executor_work_guard<boost::asio::io_context::executor_type>> work_guard_;
 
    // Session compression type negotiated during handshake
    messaging::compression_type session_compression_{messaging::compression_type::none};
 
    // Session recording (protected by state_mutex_ for thread-safe access)
    std::shared_ptr<recording::session_recorder> recorder_;
};
 
}
 
#endif