From c2ba7e49b2095019f9a7df2d71bb4b42b31a357e Mon Sep 17 00:00:00 2001
From: Florian Fischer <florian.fischer@muhq.space>
Date: Sun, 16 Jul 2023 14:43:29 +0200
Subject: do not remove active connections when removing stale ones

If the connection is renewed by the client removing the old connection
and adding the new one is racy.
Fix this by remembering a connection generation number and only to
prevent removing the new connection.
---
 wrms.go | 24 +++++++++++++++++++-----
 1 file changed, 19 insertions(+), 5 deletions(-)

(limited to 'wrms.go')

diff --git a/wrms.go b/wrms.go
index 056848c..cd6cb70 100644
--- a/wrms.go
+++ b/wrms.go
@@ -72,7 +72,9 @@ func newNotification(notification string) Event {
 }
 
 type Connection struct {
+	wrms      *Wrms
 	Id        uuid.UUID
+	nr        uint64
 	closing   atomic.Bool
 	refs      atomic.Int64
 	Events    chan Event
@@ -85,21 +87,29 @@ type Connection struct {
 
 const EVENT_BUFFER_SIZE = 3
 
-func newConnection(id uuid.UUID, w http.ResponseWriter, ctx context.Context, cncl func()) *Connection {
+func (wrms *Wrms) newConnection(id uuid.UUID,
+	w http.ResponseWriter,
+	ctx context.Context,
+	cncl func()) *Connection {
 	// prepare the flusher
 	flusher, ok := w.(http.Flusher)
 	if !ok {
 		llog.Fatal("Flusher type assertion failed")
 	}
 
-	return &Connection{
+	conn := &Connection{
+		wrms:    wrms,
 		Id:      id,
+		nr:      wrms.nextConnNr.Add(1),
 		Events:  make(chan Event, EVENT_BUFFER_SIZE),
 		w:       w,
 		flusher: flusher,
 		ctx:     ctx,
 		cancel:  cncl,
 	}
+
+	wrms.addConn(conn)
+	return conn
 }
 
 func (c *Connection) Send(ev Event) {
@@ -175,6 +185,7 @@ func (c *Connection) Close() {
 
 type Wrms struct {
 	Connections sync.Map
+	nextConnNr  atomic.Uint64
 	// The rwlock must be held when using most internal state
 	rwlock      sync.RWMutex
 	Songs       []Song
@@ -200,12 +211,15 @@ func (wrms *Wrms) addConn(conn *Connection) {
 
 func (wrms *Wrms) delConn(conn *Connection) {
 	llog.DDebug("Deleting Connection %s", conn.Id)
-	wrms.Connections.Delete(conn.Id)
+	_oldConn, _ := wrms.Connections.Load(conn.Id)
+	oldConn := _oldConn.(*Connection)
+	// Only delete the connection if no newer one is currently active
+	if oldConn.nr == conn.nr {
+		wrms.Connections.Delete(conn.Id)
+	}
 }
 
 func (wrms *Wrms) initConn(conn *Connection) error {
-	wrms.addConn(conn)
-
 	wrms.rwlock.RLock()
 	defer wrms.rwlock.RUnlock()
 
-- 
cgit v1.2.3