Close connection properly from transport closing (#17)

Author: BrennanConroy

samples/HubConnectionSample/HubConnectionSample.cpp

@@ -13,7 +13,7 @@ class logger : public signalr::log_writer
     // Inherited via log_writer
     virtual void __cdecl write(const std::string & entry) override
     {
-        std::cout << entry << std::endl;
+        std::cout << entry;
     }
 };
 
@@ -77,12 +77,12 @@ void chat()
         }
 
         std::cout << "Enter your message:";
-        for (;;)
+        while (connection.get_connection_state() == signalr::connection_state::connected)
         {
             std::string message;
             std::getline(std::cin, message);
 
-            if (message == ":q")
+            if (message == ":q" || connection.get_connection_state() != signalr::connection_state::connected)
             {
                 break;
             }

src/signalrclient/connection_impl.cpp

@@ -145,10 +145,7 @@ namespace signalr
             m_connection_id = "";
         }
 
-        start_negotiate(m_base_url, 0, [callback](std::exception_ptr exception)
-        {
-            callback(exception);
-        });
+        start_negotiate(m_base_url, 0, callback);
     }
 
     void connection_impl::start_negotiate(const std::string& url, int redirect_count, std::function<void(std::exception_ptr)> callback)
@@ -311,6 +308,21 @@ namespace signalr
         auto transport = connection->m_transport_factory->create_transport(
             transport_type::websockets, connection->m_logger, connection->m_signalr_client_config);
 
+        transport->on_close([weak_connection](std::exception_ptr exception)
+            {
+                auto connection = weak_connection.lock();
+                if (!connection)
+                {
+                    return;
+                }
+
+                // close callback will only be called if start on the transport has already returned
+                // wait for the event in order to avoid a race where the state hasn't changed from connecting
+                // yet and the transport errors out
+                connection->m_start_completed_event.wait();
+                connection->stop_connection(exception);
+            });
+
         transport->on_receive([disconnect_cts, connect_request_done, connect_request_lock, logger, weak_connection, callback](std::string&& message, std::exception_ptr exception)
             {
                 if (exception == nullptr)
@@ -532,50 +544,7 @@ namespace signalr
     {
         m_logger.log(trace_level::info, "stopping connection");
 
-        auto connection = shared_from_this();
-        shutdown([connection, callback](std::exception_ptr exception)
-            {
-                std::thread([connection, callback, exception]()
-                    {
-                        if (exception != nullptr)
-                        {
-                            callback(exception);
-                            return;
-                        }
-
-                        {
-                            // the lock prevents a race where the user calls `stop` on a disconnected connection and calls `start`
-                            // on a different thread at the same time. In this case we must not null out the transport if we are
-                            // not in the `disconnecting` state to not affect the 'start' invocation.
-                            std::lock_guard<std::mutex> lock(connection->m_stop_lock);
-                            if (connection->change_state(connection_state::disconnecting, connection_state::disconnected))
-                            {
-                                // we do let the exception through (especially the task_canceled exception)
-                                connection->m_transport = nullptr;
-                            }
-                        }
-
-                        try
-                        {
-                            connection->m_disconnected();
-                        }
-                        catch (const std::exception& e)
-                        {
-                            connection->m_logger.log(
-                                trace_level::errors,
-                                std::string("disconnected callback threw an exception: ")
-                                .append(e.what()));
-                        }
-                        catch (...)
-                        {
-                            connection->m_logger.log(
-                                trace_level::errors,
-                                std::string("disconnected callback threw an unknown exception"));
-                        }
-
-                        callback(nullptr);
-                    }).detach();
-            });
+        shutdown(callback);
     }
 
     // This function is called from the dtor so you must not use `shared_from_this` here (it will throw).
@@ -623,10 +592,61 @@ namespace signalr
             change_state(connection_state::disconnecting);
         }
 
-        m_transport->stop([callback](std::exception_ptr exception)
+        m_transport->stop(callback);
+    }
+
+    // do not use `shared_from_this` as it can be called via the destructor
+    void connection_impl::stop_connection(std::exception_ptr error)
+    {
+        {
+            // the lock prevents a race where the user calls `stop` on a disconnected connection and calls `start`
+            // on a different thread at the same time. In this case we must not null out the transport if we are
+            // not in the `disconnecting` state to not affect the 'start' invocation.
+            std::lock_guard<std::mutex> lock(m_stop_lock);
+
+            if (m_connection_state == connection_state::disconnected)
             {
-                callback(exception);
-            });
+                m_logger.log(trace_level::info, "Stopping was ignored because the connection is already in the disconnected state.");
+                return;
+            }
+
+            change_state(connection_state::disconnected);
+            m_transport = nullptr;
+        }
+
+        if (error)
+        {
+            try
+            {
+                std::rethrow_exception(error);
+            }
+            catch (const std::exception & ex)
+            {
+                m_logger.log(trace_level::errors, std::string("Connection closed with error: ").append(ex.what()));
+            }
+        }
+        else
+        {
+            m_logger.log(trace_level::info, "Connection closed.");
+        }
+
+        try
+        {
+            m_disconnected();
+        }
+        catch (const std::exception & e)
+        {
+            m_logger.log(
+                trace_level::errors,
+                std::string("disconnected callback threw an exception: ")
+                .append(e.what()));
+        }
+        catch (...)
+        {
+            m_logger.log(
+                trace_level::errors,
+                std::string("disconnected callback threw an unknown exception"));
+        }
     }
 
     connection_state connection_impl::get_connection_state() const noexcept

src/signalrclient/connection_impl.h

@@ -80,6 +80,7 @@ namespace signalr
         void process_response(std::string&& response);
 
         void shutdown(std::function<void(std::exception_ptr)> callback);
+        void stop_connection(std::exception_ptr);
 
         bool change_state(connection_state old_state, connection_state new_state);
         connection_state change_state(connection_state new_state);

src/signalrclient/hub_connection_impl.cpp

@@ -61,7 +61,11 @@ namespace signalr
             auto connection = weak_hub_connection.lock();
             if (connection)
             {
+                // start may be waiting on the handshake response so we complete it here, this no-ops if already set
                 connection->m_handshakeTask->set(std::make_exception_ptr(signalr_exception("connection closed while handshake was in progress.")));
+
+                connection->m_callback_manager.clear(signalr::value(std::map<std::string, signalr::value> { { std::string("error"), std::string("connection was stopped before invocation result was received") } }));
+
                 connection->m_disconnected();
             }
         });
@@ -125,7 +129,6 @@ namespace signalr
                                 callback(std::make_exception_ptr(signalr_exception("the hub connection has been deconstructed")));
                                 return;
                             }
-                            connection->m_handshakeTask->get();
                         }
                         catch (...) {}
 
@@ -171,8 +174,51 @@ namespace signalr
 
     void hub_connection_impl::stop(std::function<void(std::exception_ptr)> callback) noexcept
     {
-        m_callback_manager.clear(signalr::value(std::map<std::string, signalr::value> { { std::string("error"), std::string("connection was stopped before invocation result was received") } }));
-        m_connection->stop(callback);
+        if (get_connection_state() == connection_state::disconnected)
+        {
+            m_logger.log(trace_level::info, "Stop ignored because the connection is already disconnected.");
+            callback(nullptr);
+            return;
+        }
+        else
+        {
+            {
+                std::lock_guard<std::mutex> lock(m_stop_callback_lock);
+                m_stop_callbacks.push_back(callback);
+
+                if (m_stop_callbacks.size() > 1)
+                {
+                    m_logger.log(trace_level::info, "Stop is already in progress, waiting for it to finish.");
+                    // we already registered the callback
+                    // so we can just return now as the in-progress stop will trigger the callback when it completes
+                    return;
+                }
+            }
+            std::weak_ptr<hub_connection_impl> weak_connection = shared_from_this();
+            m_connection->stop([weak_connection](std::exception_ptr exception)
+                {
+                    auto connection = weak_connection.lock();
+                    if (!connection)
+                    {
+                        return;
+                    }
+
+                    std::vector<std::function<void(std::exception_ptr)>> callbacks;
+
+                    {
+                        std::lock_guard<std::mutex> lock(connection->m_stop_callback_lock);
+                        // copy the callbacks out and clear the list inside the lock
+                        // then run the callbacks outside of the lock
+                        callbacks = connection->m_stop_callbacks;
+                        connection->m_stop_callbacks.clear();
+                    }
+
+                    for (auto& callback : callbacks)
+                    {
+                        callback(exception);
+                    }
+                });
+        }
     }
 
     void hub_connection_impl::process_message(std::string&& response)
@@ -184,11 +230,11 @@ namespace signalr
                 signalr::value handshake;
                 std::tie(response, handshake) = handshake::parse_handshake(response);
 
-                auto obj = handshake.as_map();
+                auto& obj = handshake.as_map();
                 auto found = obj.find("error");
                 if (found != obj.end())
                 {
-                    auto error = found->second.as_string();
+                    auto& error = found->second.as_string();
                     m_logger.log(trace_level::errors, std::string("handshake error: ")
                         .append(error));
                     m_handshakeTask->set(std::make_exception_ptr(signalr_exception(std::string("Received an error during handshake: ").append(error))));
@@ -291,13 +337,13 @@ namespace signalr
             throw signalr_exception("expected object");
         }
 
-        auto invocationId = message.as_map().at("invocationId");
+        auto& invocationId = message.as_map().at("invocationId");
         if (!invocationId.is_string())
         {
             throw signalr_exception("invocationId is not a string");
         }
 
-        auto id = invocationId.as_string();
+        auto& id = invocationId.as_string();
         if (!m_callback_manager.invoke_callback(id, message, true))
         {
             m_logger.log(trace_level::info, std::string("no callback found for id: ").append(id));

src/signalrclient/hub_connection_impl.h

@@ -59,6 +59,9 @@ namespace signalr
         signalr_client_config m_signalr_client_config;
         std::shared_ptr<hub_protocol> m_protocol;
 
+        std::mutex m_stop_callback_lock;
+        std::vector<std::function<void(std::exception_ptr)>> m_stop_callbacks;
+
         void initialize();
 
         void process_message(std::string&& message);

src/signalrclient/hub_protocol.h

@@ -14,5 +14,6 @@ namespace signalr
         virtual std::vector<signalr::value> parse_messages(const std::string&) const = 0;
         virtual const std::string& name() const = 0;
         virtual int version() const = 0;
+        virtual ~hub_protocol() {}
     };
 }

src/signalrclient/json_hub_protocol.h

@@ -21,6 +21,8 @@ namespace signalr
         {
             return 1;
         }
+
+        ~json_hub_protocol() {}
     private:
         signalr::value parse_message(const std::string&) const;
 

src/signalrclient/websocket_transport.cpp

@@ -226,6 +226,7 @@ namespace signalr
             {
                 try
                 {
+                    close_callback(exception);
                     if (exception != nullptr)
                     {
                         std::rethrow_exception(exception);
@@ -241,8 +242,6 @@ namespace signalr
 
                     callback(exception);
                 }
-
-                close_callback(exception);
             });
     }
 

test/signalrclienttests/connection_impl_tests.cpp

@@ -1640,10 +1640,34 @@ TEST(connection_impl_stop, exception_for_disconnected_callback_caught_and_logged
     mre.get();
 
     auto log_entries = std::dynamic_pointer_cast<memory_log_writer>(writer)->get_log_entries();
-    ASSERT_EQ(1U, log_entries.size());
+    ASSERT_EQ(1U, log_entries.size()) << dump_vector(log_entries);
     ASSERT_EQ("[error       ] disconnected callback threw an unknown exception\n", remove_date_from_log_entry(log_entries[0]));
 }
 
+TEST(connection_impl_stop, transport_error_invokes_disconnected_callback)
+{
+    auto websocket_client = create_test_websocket_client();
+    auto connection = create_connection(websocket_client);
+
+    auto disconnect_mre = manual_reset_event<void>();
+    connection->set_disconnected([&disconnect_mre]() { disconnect_mre.set(); });
+
+    auto mre = manual_reset_event<void>();
+    connection->start([&mre](std::exception_ptr exception)
+        {
+            mre.set(exception);
+        });
+
+    mre.get();
+
+    ASSERT_FALSE(websocket_client->receive_loop_started.wait(1000));
+    websocket_client->receive_message(std::make_exception_ptr(std::runtime_error("error")));
+
+    disconnect_mre.get();
+
+    ASSERT_EQ(connection_state::disconnected, connection->get_connection_state());
+}
+
 TEST(connection_impl_config, custom_headers_set_in_requests)
 {
     auto writer = std::shared_ptr<log_writer>{ std::make_shared<memory_log_writer>() };