Libwebsockets: Fix message loss on reconnect

Signed-off-by: Cornelius Claussen <cc@pionix.de>

Author: corneliusclaussen

include/ocpp/common/websocket/websocket_base.hpp

@@ -47,7 +47,7 @@ enum class ConnectionFailedReason {
 ///
 class WebsocketBase {
 protected:
-    bool m_is_connected;
+    std::atomic_bool m_is_connected;
     WebsocketConnectionOptions connection_options;
     std::function<void(const int security_profile)> connected_callback;
     std::function<void()> disconnected_callback;
@@ -59,11 +59,11 @@ class WebsocketBase {
     websocketpp::connection_hdl handle;
     std::mutex reconnect_mutex;
     std::mutex connection_mutex;
-    long reconnect_backoff_ms;
+    std::atomic_int reconnect_backoff_ms;
     websocketpp::transport::timer_handler reconnect_callback;
-    int connection_attempts;
-    bool shutting_down;
-    bool reconnecting;
+    std::atomic_int connection_attempts;
+    std::atomic_bool shutting_down;
+    std::atomic_bool reconnecting;
 
     /// \brief Indicates if the required callbacks are registered
     /// \returns true if the websocket is properly initialized

include/ocpp/common/websocket/websocket_tls_tpm.hpp

@@ -67,7 +67,7 @@ class WebsocketTlsTPM final : public WebsocketBase {
 
     void request_write();
 
-    void poll_message(const std::shared_ptr<WebsocketMessage>& msg, bool wait_sendaf);
+    void poll_message(const std::shared_ptr<WebsocketMessage>& msg);
 
 private:
     std::shared_ptr<EvseSecurity> evse_security;
@@ -79,8 +79,10 @@ class WebsocketTlsTPM final : public WebsocketBase {
     std::condition_variable conn_cv;
 
     std::mutex queue_mutex;
+
     std::queue<std::shared_ptr<WebsocketMessage>> message_queue;
     std::condition_variable msg_send_cv;
+    std::mutex msg_send_cv_mutex;
 
     std::unique_ptr<std::thread> recv_message_thread;
     std::mutex recv_mutex;

lib/ocpp/common/websocket/websocket_tls_tpm.cpp

@@ -87,41 +87,47 @@ struct ConnectionData {
     }
 
     bool is_connecting() {
-        return (state == EConnectionState::CONNECTING);
+        return (state.load() == EConnectionState::CONNECTING);
     }
 
     bool is_close_requested() {
         return is_marked_close;
     }
 
     auto get_state() {
-        return state;
+        return state.load();
     }
 
     lws* get_conn() {
         return wsi;
     }
 
-    lws_context* get_ctx() {
-        return lws_ctx.get();
+    WebsocketTlsTPM* get_owner() {
+        return owner.load();
+    }
+
+    void set_owner(WebsocketTlsTPM* o) {
+        owner = o;
     }
 
 public:
+    // This public block will only be used from client loop thread, no locking needed
     // Openssl context, must be destroyed in this order
     std::unique_ptr<SSL_CTX> sec_context;
     std::unique_ptr<OSSL_LIB_CTX> sec_lib_context;
-
     // libwebsockets state
     std::unique_ptr<lws_context> lws_ctx;
-    lws* wsi;
 
-    WebsocketTlsTPM* owner;
+    lws* wsi;
 
 private:
+    std::atomic<WebsocketTlsTPM*> owner;
+
     std::thread::id lws_thread_id;
-    bool is_running;
-    bool is_marked_close;
-    EConnectionState state;
+
+    std::atomic_bool is_running;
+    std::atomic_bool is_marked_close;
+    std::atomic<EConnectionState> state;
 };
 
 struct WebsocketMessage {
@@ -140,7 +146,7 @@ struct WebsocketMessage {
     // just that these were sent to libwebsockets
     size_t sent_bytes;
     // If libwebsockets has sent all the bytes through the wire
-    volatile bool message_sent;
+    std::atomic_bool message_sent;
 };
 
 WebsocketTlsTPM::WebsocketTlsTPM(const WebsocketConnectionOptions& connection_options,
@@ -189,7 +195,12 @@ static int callback_minimal(struct lws* wsi, enum lws_callback_reasons reason, v
     // Get user safely, since on some callbacks (void *user) can be different than what we set
     if (wsi != nullptr) {
         if (ConnectionData* data = reinterpret_cast<ConnectionData*>(lws_wsi_user(wsi))) {
-            return data->owner->process_callback(wsi, static_cast<int>(reason), user, in, len);
+            auto owner = data->get_owner();
+            if (owner not_eq nullptr) {
+                return data->get_owner()->process_callback(wsi, static_cast<int>(reason), user, in, len);
+            } else {
+                EVLOG_error << "callback_minimal called, but data->owner is nullptr";
+            }
         }
     }
 
@@ -331,11 +342,14 @@ void WebsocketTlsTPM::recv_loop() {
 
     while (false == data->is_interupted()) {
         // Process all messages
-        while (false == recv_message_queue.empty()) {
+        while (true) {
             std::string message{};
 
             {
                 std::lock_guard lk(this->recv_mutex);
+                if (recv_message_queue.empty())
+                    break;
+
                 message = std::move(recv_message_queue.front());
                 recv_message_queue.pop();
             }
@@ -459,9 +473,14 @@ void WebsocketTlsTPM::client_loop() {
 
     while (n >= 0 && (false == data->is_interupted())) {
         // Set to -1 for continuous servicing, of required, not recommended
-        n = lws_service(data->get_ctx(), 0);
+        n = lws_service(data->lws_ctx.get(), 0);
 
-        if (false == message_queue.empty()) {
+        bool message_queue_empty;
+        {
+            std::lock_guard<std::mutex> lock(this->queue_mutex);
+            message_queue_empty = message_queue.empty();
+        }
+        if (false == message_queue_empty) {
             lws_callback_on_writable(data->get_conn());
         }
     }
@@ -470,6 +489,7 @@ void WebsocketTlsTPM::client_loop() {
     EVLOG_debug << "Exit client loop with ID: " << std::this_thread::get_id();
 }
 
+// Will be called from external threads as well
 bool WebsocketTlsTPM::connect() {
     if (!this->initialized()) {
         return false;
@@ -485,7 +505,7 @@ bool WebsocketTlsTPM::connect() {
     }
 
     auto conn_data = new ConnectionData();
-    conn_data->owner = this;
+    conn_data->set_owner(this);
 
     this->conn_data.reset(conn_data);
 
@@ -600,7 +620,8 @@ void WebsocketTlsTPM::on_conn_connected() {
     this->m_is_connected = true;
     this->reconnecting = false;
 
-    this->connected_callback(this->connection_options.security_profile);
+    std::thread connected([this]() { this->connected_callback(this->connection_options.security_profile); });
+    connected.detach();
 }
 
 void WebsocketTlsTPM::on_conn_close() {
@@ -717,11 +738,14 @@ void WebsocketTlsTPM::on_writable() {
         return;
     }
 
-    while (false == message_queue.empty()) {
+    while (true) {
         WebsocketMessage* message = nullptr;
 
         {
             std::lock_guard<std::mutex> lock(this->queue_mutex);
+            if (message_queue.empty()) {
+                break;
+            }
             message = message_queue.front().get();
         }
 
@@ -744,7 +768,7 @@ void WebsocketTlsTPM::on_writable() {
 
             EVLOG_debug << "Notifying waiting thread!";
             // Notify any waiting thread to check it's state
-            msg_send_cv.notify_one();
+            msg_send_cv.notify_all();
         } else {
             EVLOG_debug << "Client writable, sending message part!";
 
@@ -766,16 +790,18 @@ void WebsocketTlsTPM::request_write() {
     if (this->m_is_connected) {
         if (auto* data = conn_data.get()) {
             if (data->get_conn()) {
-                // Notify waiting processing thread to wake up
-                lws_cancel_service(data->get_ctx());
+                // Notify waiting processing thread to wake up. According to docs it is ok to call from another
+                // thread.
+                lws_cancel_service(data->lws_ctx.get());
             }
         }
     } else {
         EVLOG_warning << "Requested write with offline TLS websocket!";
     }
 }
 
-void WebsocketTlsTPM::poll_message(const std::shared_ptr<WebsocketMessage>& msg, bool wait_send) {
+void WebsocketTlsTPM::poll_message(const std::shared_ptr<WebsocketMessage>& msg) {
+
     if (std::this_thread::get_id() == conn_data->get_lws_thread_id()) {
         EVLOG_AND_THROW(std::runtime_error("Deadlock detected, polling send from client lws thread!"));
     }
@@ -790,17 +816,17 @@ void WebsocketTlsTPM::poll_message(const std::shared_ptr<WebsocketMessage>& msg,
     // Request a write callback
     request_write();
 
-    if (wait_send) {
-        std::unique_lock lock(this->queue_mutex);
-        msg_send_cv.wait_for(lock, std::chrono::seconds(10), [&] { return (true == msg->message_sent); });
+    {
+        std::unique_lock lock(this->msg_send_cv_mutex);
+        if (msg_send_cv.wait_for(lock, std::chrono::seconds(20), [&] { return (true == msg->message_sent); })) {
+            EVLOG_info << "Successfully sent last message over TLS websocket!";
+        } else {
+            EVLOG_warning << "Could not send last message over TLS websocket!";
+        }
     }
-
-    if (msg->message_sent)
-        EVLOG_info << "Successfully sent last message over TLS websocket!";
-    else
-        EVLOG_warning << "Could not send last message over TLS websocket!";
 }
 
+// Will be called from external threads
 bool WebsocketTlsTPM::send(const std::string& message) {
     if (!this->initialized()) {
         EVLOG_error << "Could not send message because websocket is not properly initialized.";
@@ -811,7 +837,7 @@ bool WebsocketTlsTPM::send(const std::string& message) {
     msg->payload = std::move(message);
     msg->protocol = LWS_WRITE_TEXT;
 
-    poll_message(msg, true);
+    poll_message(msg);
 
     return msg->message_sent;
 }
@@ -825,7 +851,7 @@ void WebsocketTlsTPM::ping() {
     msg->payload = this->connection_options.ping_payload;
     msg->protocol = LWS_WRITE_PING;
 
-    poll_message(msg, true);
+    poll_message(msg);
 }
 
 int WebsocketTlsTPM::process_callback(void* wsi_ptr, int callback_reason, void* user, void* in, size_t len) {
@@ -949,31 +975,53 @@ int WebsocketTlsTPM::process_callback(void* wsi_ptr, int callback_reason, void*
 
     case LWS_CALLBACK_CLIENT_WRITEABLE:
         on_writable();
-
-        if (false == message_queue.empty()) {
-            lws_callback_on_writable(wsi);
+        {
+            bool message_queue_empty;
+            {
+                std::lock_guard<std::mutex> lock(this->queue_mutex);
+                message_queue_empty = message_queue.empty();
+            }
+            if (false == message_queue_empty) {
+                lws_callback_on_writable(wsi);
+            }
         }
         break;
 
-    case LWS_CALLBACK_CLIENT_RECEIVE_PONG:
-        if (false == message_queue.empty()) {
+    case LWS_CALLBACK_CLIENT_RECEIVE_PONG: {
+        bool message_queue_empty;
+        {
+            std::lock_guard<std::mutex> lock(this->queue_mutex);
+            message_queue_empty = message_queue.empty();
+        }
+        if (false == message_queue_empty) {
             lws_callback_on_writable(data->get_conn());
         }
-        break;
+    } break;
 
     case LWS_CALLBACK_CLIENT_RECEIVE:
         on_message(in, len);
-
-        if (false == message_queue.empty()) {
-            lws_callback_on_writable(data->get_conn());
+        {
+            bool message_queue_empty;
+            {
+                std::lock_guard<std::mutex> lock(this->queue_mutex);
+                message_queue_empty = message_queue.empty();
+            }
+            if (false == message_queue_empty) {
+                lws_callback_on_writable(data->get_conn());
+            }
         }
         break;
 
-    case LWS_CALLBACK_EVENT_WAIT_CANCELLED:
-        if (false == message_queue.empty()) {
+    case LWS_CALLBACK_EVENT_WAIT_CANCELLED: {
+        bool message_queue_empty;
+        {
+            std::lock_guard<std::mutex> lock(this->queue_mutex);
+            message_queue_empty = message_queue.empty();
+        }
+        if (false == message_queue_empty) {
             lws_callback_on_writable(data->get_conn());
         }
-        break;
+    } break;
 
     default:
         EVLOG_info << "Callback with unhandled reason: " << reason;