57 lines
1.7 KiB
C++
57 lines
1.7 KiB
C++
#include <atomic>
|
|
#include <csignal>
|
|
#include <memory>
|
|
#include <thread>
|
|
|
|
#include "rdp_worker/common/logger.hpp"
|
|
#include "rdp_worker/config/config.hpp"
|
|
#include "rdp_worker/coordination/control_plane.hpp"
|
|
#include "rdp_worker/dataplane/direct_wss_server.hpp"
|
|
#include "rdp_worker/runtime/session_manager.hpp"
|
|
|
|
namespace {
|
|
|
|
std::atomic<bool> g_stop{false};
|
|
|
|
void HandleSignal(int) {
|
|
g_stop.store(true);
|
|
}
|
|
|
|
} // namespace
|
|
|
|
int main() {
|
|
std::signal(SIGINT, HandleSignal);
|
|
std::signal(SIGTERM, HandleSignal);
|
|
|
|
auto logger = std::make_shared<rdp_worker::common::Logger>("rdp-worker");
|
|
const auto config = rdp_worker::config::LoadFromEnv();
|
|
auto control_plane = std::make_shared<rdp_worker::coordination::ControlPlane>(config, logger);
|
|
control_plane->Connect();
|
|
control_plane->ReleaseOwnedLeasesOnStartup();
|
|
control_plane->RegisterWorker();
|
|
|
|
auto session_manager = std::make_shared<rdp_worker::runtime::SessionManager>(control_plane, logger);
|
|
rdp_worker::dataplane::DirectWssServer direct_wss_server(config, session_manager, logger);
|
|
direct_wss_server.Start();
|
|
|
|
std::thread heartbeat_thread([&]() {
|
|
while (!g_stop.load()) {
|
|
control_plane->SendHeartbeat();
|
|
std::this_thread::sleep_for(config.worker_heartbeat_interval);
|
|
}
|
|
});
|
|
|
|
while (!g_stop.load()) {
|
|
if (auto assignment = control_plane->PollAssignment(config.assignment_poll_interval); assignment.has_value()) {
|
|
session_manager->ApplyAssignment(*assignment);
|
|
}
|
|
}
|
|
|
|
direct_wss_server.Stop();
|
|
session_manager->StopAll();
|
|
if (heartbeat_thread.joinable()) {
|
|
heartbeat_thread.join();
|
|
}
|
|
return 0;
|
|
}
|