diff --git a/tests/test_server.cpp b/tests/test_server.cpp index 02230c3..f7bd247 100644 --- a/tests/test_server.cpp +++ b/tests/test_server.cpp @@ -3,22 +3,36 @@ #include "connection_handler.hpp" #include "server.hpp" +#include #include #include #include #include +struct Event { + void wait() { done_.wait(); } + void set() { + if (!set_.exchange(true, std::memory_order_relaxed)) { + done_.count_down(); + } + } + +private: + std::latch done_{1}; + std::atomic set_; +}; + struct EchoHandler : ConnectionHandler { Arena arena; std::span reply; WeakRef wconn; - std::latch done{1}; + Event done; void on_data_arrived(std::string_view data, Connection &conn) override { reply = arena.allocate_span(1); reply[0] = arena.copy_string(data); wconn = conn.get_weak_ref(); CHECK(wconn.lock()); - done.count_down(); + done.set(); } }; @@ -60,8 +74,8 @@ struct ShutdownTestHandler : ConnectionHandler { Arena arena; std::span reply; WeakRef wconn; - std::latch received_data{1}; - std::latch connection_closed_latch{1}; + Event received_data; + Event connection_closed_latch; ConnectionShutdown shutdown_mode = ConnectionShutdown::None; std::atomic connection_closed{false}; @@ -69,12 +83,12 @@ struct ShutdownTestHandler : ConnectionHandler { reply = arena.allocate_span(1); reply[0] = arena.copy_string(data); wconn = conn.get_weak_ref(); - received_data.count_down(); + received_data.set(); } void on_connection_closed(Connection &) override { connection_closed = true; - connection_closed_latch.count_down(); + connection_closed_latch.set(); } };