From 30a49217e4b231f666f01b40004f09fce38e65ac Mon Sep 17 00:00:00 2001
From: Erik Brakkee <erik@brakkee.org>
Date: Tue, 13 Aug 2024 10:57:05 +0200
Subject: [PATCH] Much imporoved websocket connection closure will now detect
 closing of websockets immediately.

---
 cmd/converge/converge.go             | 4 ++--
 pkg/server/matchmaker/websessions.go | 9 +++++++--
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/cmd/converge/converge.go b/cmd/converge/converge.go
index 0f91a4f..a5daa1c 100644
--- a/cmd/converge/converge.go
+++ b/cmd/converge/converge.go
@@ -136,13 +136,13 @@ func main() {
 	sessionService := websocketutil.WebSocketService{
 		Handler: func(w http.ResponseWriter, r *http.Request, conn net.Conn) {
 			ctx, cancel := context.WithCancel(context.Background())
-			websession := websessions.NewSession(conn, ctx)
+			websession := websessions.NewSession(conn, ctx, cancel)
 			defer websessions.SessionClosed(websession)
 			location, err := matchmaker.GetUserLocation(r)
 			if err != nil {
 				panic(err)
 			}
-			websession.WriteNotifications(location, cancel)
+			websession.WriteNotifications(location, ctx, cancel)
 		},
 		Text: true,
 	}
diff --git a/pkg/server/matchmaker/websessions.go b/pkg/server/matchmaker/websessions.go
index c12efd2..5186765 100644
--- a/pkg/server/matchmaker/websessions.go
+++ b/pkg/server/matchmaker/websessions.go
@@ -51,7 +51,8 @@ func (sessions *WebSessions) notifyWebSessions(notification *models.State) {
 	}
 }
 
-func (sessions *WebSessions) NewSession(wsConnection net.Conn, ctx context.Context) *WebSession {
+func (sessions *WebSessions) NewSession(wsConnection net.Conn, ctx context.Context,
+	cancel context.CancelFunc) *WebSession {
 	sessions.mutex.Lock()
 	defer sessions.mutex.Unlock()
 	session := &WebSession{
@@ -70,6 +71,7 @@ func (sessions *WebSessions) NewSession(wsConnection net.Conn, ctx context.Conte
 				session.notifications <- sessions.lastNotification
 			} else {
 				log.Printf("Got error reading %v", err)
+				cancel()
 				return
 			}
 		}
@@ -84,19 +86,22 @@ func GetUserLocation(r *http.Request) (*time.Location, error) {
 	if tzName == "" {
 		tzName = r.Header.Get("X-Timezone")
 	}
+	//log.Printf("Got timezone from request %v %v", tzName, r.URL.Path)
 	if tzName == "" {
 		tzName = "UTC"
 	}
 	return time.LoadLocation(tzName)
 }
 
-func (session *WebSession) WriteNotifications(location *time.Location, cancel context.CancelFunc) {
+func (session *WebSession) WriteNotifications(location *time.Location, ctx context.Context, cancel context.CancelFunc) {
 	timer := time.NewTicker(10 * time.Second)
 	defer timer.Stop()
 	// if for some reason we cannot send notifications to the web client then the context is canceled.
 	defer cancel()
 	for {
 		select {
+		case <-ctx.Done():
+			return
 		case notification, ok := <-session.notifications:
 			if !ok {
 				log.Println("channel closed")