From 46669cfda91ae40f1527b53cd2844aacba68a2af Mon Sep 17 00:00:00 2001
From: Robin Davies <rerdavies@gmail.com>
Date: Sat, 5 Oct 2024 02:14:30 -0400
Subject: [PATCH] Lifecycle management for sockets.

---
 src/PiPedalSocket.cpp |  1 -
 src/WebServer.cpp     | 11 +++++++++--
 src/WebServer.hpp     | 14 +++++++++++---
 3 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/src/PiPedalSocket.cpp b/src/PiPedalSocket.cpp
index aabe239..4c00411 100644
--- a/src/PiPedalSocket.cpp
+++ b/src/PiPedalSocket.cpp
@@ -552,7 +552,6 @@ public:
         {
             FinalCleanup();
         }
-        Lv2Log::error("PiPedalSocketHandler deleted");
 
     }
 
diff --git a/src/WebServer.cpp b/src/WebServer.cpp
index 66e498f..806f260 100644
--- a/src/WebServer.cpp
+++ b/src/WebServer.cpp
@@ -754,7 +754,11 @@ namespace pipedal
         public:
             ~WebSocketSession()
             {
-                this->socketHandler = nullptr;
+                if (this->socketHandler)
+                {
+                    this->socketHandler->onSocketClosed();
+                    this->socketHandler = nullptr;
+                }
                 webSocket = nullptr;
                 pServer = nullptr;
                 Lv2Log::info(SS("WebSocketSession closed. " << fromAddress));
@@ -819,7 +823,10 @@ namespace pipedal
             std::lock_guard<std::recursive_mutex> lock{m_sessionsMutex};
             m_sessions.erase(session);
             m_connections.erase(hConnection);
-            session = nullptr; // probably delete here.
+            if (session != nullptr)
+            {
+                session = nullptr; // probably delete here.
+            }
         }
 
         void NotFound(server::connection_type &connection, const std::string &filename)
diff --git a/src/WebServer.hpp b/src/WebServer.hpp
index c7fa19e..d7d402d 100644
--- a/src/WebServer.hpp
+++ b/src/WebServer.hpp
@@ -81,7 +81,6 @@ class WebServerImpl;
 class SocketHandler {
     friend class WebServerImpl;
 public:
-
     class IWriteCallback {
     public:
         virtual void close() = 0;
@@ -100,6 +99,7 @@ public:
     }
 
 public:
+    virtual void onSocketClosed() = 0;
     virtual void onReceive(const std::string_view&text) = 0;
 public:
     std::string getFromAddress() const { return writeCallback_->getFromAddress(); }
@@ -112,12 +112,20 @@ public:
             writeCallback_->writeCallback(text);
         }
     }
+    virtual void OnSocketClosed()
+    {
+        writeCallback_ = nullptr;
+    }
     virtual void Close()
     {
-        writeCallback_->close();
+        if (writeCallback_ != nullptr)
+        {
+            writeCallback_->close();
+            writeCallback_ =  nullptr;
+        }
     }
 
-    virtual ~SocketHandler() = default;
+    virtual ~SocketHandler() {}
 
     virtual void onAttach() { }