Merge pull request #58 from sonroyaalmerol/remove-passthrough-proxy

Merge websockets with the golang broadcast channel implementation
sonroyaalmerol · Jan 12, 2025 · 1992baa · 1992baa
2 parents c3ffad7 + b58545d
commit 1992baa
Show file tree

Hide file tree

Showing 6 changed files with 216 additions and 217 deletions.
diff --git a/cmd/pbs_plus/main.go b/cmd/pbs_plus/main.go
@@ -3,6 +3,7 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"log"
 	"net/http"
@@ -32,7 +33,9 @@ func main() {
 	jobRun := flag.String("job", "", "Job ID to execute")
 	flag.Parse()
 
-	wsHub := websockets.NewServer()
+	wsHub := websockets.NewServer(context.Background())
+	go wsHub.Run()
+
 	storeInstance, err := store.Initialize(wsHub)
 	if err != nil {
 		s.Errorf("Failed to initialize store: %v", err)

diff --git a/internal/proxy/controllers/plus/plus.go b/internal/proxy/controllers/plus/plus.go
@@ -47,7 +47,7 @@ func MountHandler(storeInstance *store.Store) http.HandlerFunc {
 		agentDrive := string(agentDriveBytes)
 
 		if r.Method == http.MethodPost {
-			err := storeInstance.WSHub.SendCommand(targetHostname, websockets.Message{
+			err := storeInstance.WSHub.SendToClient(targetHostname, websockets.Message{
 				Type:    "backup_start",
 				Content: agentDrive,
 			})
@@ -56,9 +56,8 @@ func MountHandler(storeInstance *store.Store) http.HandlerFunc {
 				return
 			}
 
-			listener := storeInstance.WSHub.Broadcast.Subscribe()
-			defer storeInstance.WSHub.Broadcast.CancelSubscription(listener)
-
+			listener, closeListener := storeInstance.WSHub.Subscribe()
+			defer closeListener()
 		respWait:
 			for {
 				select {
@@ -79,7 +78,7 @@ func MountHandler(storeInstance *store.Store) http.HandlerFunc {
 		}
 
 		if r.Method == http.MethodDelete {
-			_ = storeInstance.WSHub.SendCommand(targetHostname, websockets.Message{
+			_ = storeInstance.WSHub.SendToClient(targetHostname, websockets.Message{
 				Type:    "backup_close",
 				Content: agentDrive,
 			})

diff --git a/internal/proxy/controllers/plus/ws.go b/internal/proxy/controllers/plus/ws.go
@@ -20,6 +20,6 @@ func WSHandler(storeInstance *store.Store) http.HandlerFunc {
 			return
 		}
 
-		storeInstance.WSHub.HandleClientConnection(w, r)
+		storeInstance.WSHub.ServeWS(w, r)
 	}
 }
diff --git a/internal/store/agent_ping.go b/internal/store/agent_ping.go
@@ -14,12 +14,5 @@ func (storeInstance *Store) AgentPing(agentTarget *Target) bool {
 		return false
 	}
 
-	storeInstance.WSHub.ClientsMux.RLock()
-	if client, ok := storeInstance.WSHub.Clients[agentHostname]; ok && client != nil {
-		storeInstance.WSHub.ClientsMux.RUnlock()
-		return true
-	}
-	storeInstance.WSHub.ClientsMux.RUnlock()
-
-	return false
+	return storeInstance.WSHub.IsClientConnected(agentHostname)
 }
diff --git a/internal/websockets/message_broadcast.go b/internal/websockets/message_broadcast.go
-Original file line number
+Diff line change
@@ Expand Up / @@ -20,6 +20,6 @@ func WSHandler(storeInstance *store.Store) http.HandlerFunc { @@
     			return
     		}
-    		storeInstance.WSHub.HandleClientConnection(w, r)
+    		storeInstance.WSHub.ServeWS(w, r)
     	}
     }