diff --git a/include/crow/websocket.h b/include/crow/websocket.h index 4369dfb01..79dc3386f 100644 --- a/include/crow/websocket.h +++ b/include/crow/websocket.h @@ -118,18 +118,51 @@ namespace crow start(crow::utility::base64encode((unsigned char*)digest, 20)); } + ~Connection() noexcept override + { + // Do not modify anchor_ here since writing shared_ptr is not atomic. + auto watch = std::weak_ptr{anchor_}; + + // Wait until all unhandled asynchronous operations to join. + // As the deletion occurs inside 'check_destroy()', which already locks + // anchor, use count can be 1 on valid deletion context. + while (watch.use_count() > 2) // 1 for 'check_destroy() routine', 1 for 'this->anchor_' + { + std::this_thread::yield(); + } + } + + template + struct WeakWrappedMessage + { + Callable callable; + std::weak_ptr watch; + + void operator()() + { + if (auto anchor = watch.lock()) + { + std::move(callable)(); + } + } + }; + /// Send data through the socket. template void dispatch(CompletionHandler&& handler) { - asio::dispatch(adaptor_.get_io_service(), std::forward(handler)); + asio::dispatch(adaptor_.get_io_service(), + WeakWrappedMessage::type>{ + std::forward(handler), anchor_}); } /// Send data through the socket and return immediately. template void post(CompletionHandler&& handler) { - asio::post(adaptor_.get_io_service(), std::forward(handler)); + asio::post(adaptor_.get_io_service(), + WeakWrappedMessage::type>{ + std::forward(handler), anchor_}); } /// Send a "Ping" message. @@ -605,12 +638,13 @@ namespace crow { buffers.emplace_back(asio::buffer(s)); } + auto watch = std::weak_ptr{anchor_}; asio::async_write( adaptor_.socket(), buffers, - [&](const asio::error_code& ec, std::size_t /*bytes_transferred*/) { - sending_buffers_.clear(); + [&, watch](const asio::error_code& ec, std::size_t /*bytes_transferred*/) { if (!ec && !close_connection_) { + sending_buffers_.clear(); if (!write_buffers_.empty()) do_write(); if (has_sent_close_) @@ -618,6 +652,10 @@ namespace crow } else { + auto anchor = watch.lock(); + if (anchor == nullptr) { return; } + + sending_buffers_.clear(); close_connection_ = true; check_destroy(); } @@ -694,6 +732,8 @@ namespace crow bool pong_received_{false}; bool is_close_handler_called_{false}; + std::shared_ptr anchor_ = std::make_shared(); // Value is just for placeholding + std::function open_handler_; std::function message_handler_; std::function close_handler_;