From cf888e7cb75394862af0ff0cf1e066b7c328788d Mon Sep 17 00:00:00 2001
From: Tom Bevan <tom.bevan@r3labs.io>
Date: Tue, 5 Feb 2019 18:58:02 +0000
Subject: [PATCH] cleanup of subscriber

---
 handler.go | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/handler.go b/handler.go
index 2c153ae..5bce4b8 100644
--- a/handler.go
+++ b/handler.go
@@ -21,6 +21,8 @@ var upgrader = websocket.Upgrader{
 }
 
 func handler(w http.ResponseWriter, r *http.Request) {
+	var session *Session
+
 	reqid := uuid.New().String()
 
 	log.Printf("[%s] client connected\n", reqid)
@@ -34,9 +36,15 @@ func handler(w http.ResponseWriter, r *http.Request) {
 	defer func() {
 		log.Printf("[%s] client disconnected\n", reqid)
 		_ = c.Close()
-	}()
 
-	var session *Session
+		if session == nil {
+			return
+		}
+
+		if session.subscriber != nil {
+			session.subscriber.Disconnect(session.channel)
+		}
+	}()
 
 	for {
 		if session == nil {