Track Workers by uuid for Reconnection

server/worker.go

@@ -4,16 +4,24 @@ import (
 	"context"
 	"log/slog"
 	"net/http"
+	"sync"
+	"time"
 
 	"git.lastassault.de/speatzle/morffix/constants"
 
 	"nhooyr.io/websocket"
 )
 
-type Connection struct {
+type Worker struct {
+	Name              string
+	Address           string
+	Conn              *websocket.Conn
+	Connected         bool
+	ConnectionChanged time.Time
 }
 
-var Connections []Connection
+var Workers = map[string]*Worker{}
+var WorkersMutex sync.Mutex
 
 func handleWorkerWebsocket(w http.ResponseWriter, r *http.Request) {
 	if r.Header.Get(constants.SHARED_SECRET_HEADER) != conf.SharedSecret {
@@ -29,9 +37,48 @@ func handleWorkerWebsocket(w http.ResponseWriter, r *http.Request) {
 	}
 	defer c.CloseNow()
 
-	slog.InfoContext(r.Context(), "Worker Websocket Connection", "RemoteAddress", r.RemoteAddr)
+	// Connection ID
+	uuid := r.Header.Get(constants.UUID_HEADER)
 
-	// TODO Add Connection Tracking
+	// Track Connection
+	func() {
+		WorkersMutex.Lock()
+		defer WorkersMutex.Unlock()
+
+		w, ok := Workers[uuid]
+		if ok && w.Connected {
+			slog.WarnContext(r.Context(), "Worker Reconnected before the old connection died, killing...")
+			w.Conn.CloseNow()
+
+			// Since we Closed the Connection, we need to wait for the old connections deferred cleanup to finish before Writing down the new connection or we will be overwritten.
+			for {
+				WorkersMutex.Unlock()
+				time.Sleep(time.Millisecond * 10)
+				WorkersMutex.Lock()
+				if !Workers[uuid].Connected {
+					break
+				}
+			}
+		}
+
+		Workers[uuid] = &Worker{
+			Name:              r.Header.Get(constants.NAME_HEADER),
+			Address:           r.RemoteAddr,
+			Conn:              c,
+			Connected:         true,
+			ConnectionChanged: time.Now(),
+		}
+	}()
+
+	// Set Status on Disconnect
+	defer func() {
+		WorkersMutex.Lock()
+		defer WorkersMutex.Unlock()
+		Workers[uuid].Connected = false
+		Workers[uuid].ConnectionChanged = time.Now()
+	}()
+
+	slog.InfoContext(r.Context(), "Worker Websocket Connection", "RemoteAddress", r.RemoteAddr)
 
 	for {
 		err = readMessage(r.Context(), c)