From 39902e32544e834e312cf28d61172e2970ad366e Mon Sep 17 00:00:00 2001
From: Son Roy Almerol <github@snry.me>
Date: Sat, 11 Jan 2025 00:06:44 -0500
Subject: [PATCH] use client connections for agent ping instead

---
 internal/store/agent_ping.go  | 32 ++++++++------------------------
 internal/store/db.go          | 29 +++--------------------------
 internal/websockets/server.go | 22 ----------------------
 3 files changed, 11 insertions(+), 72 deletions(-)

diff --git a/internal/store/agent_ping.go b/internal/store/agent_ping.go
index 99f3be3..e8c3806 100644
--- a/internal/store/agent_ping.go
+++ b/internal/store/agent_ping.go
@@ -4,38 +4,22 @@ package store
 
 import (
 	"strings"
-	"time"
-
-	"github.com/sonroyaalmerol/pbs-plus/internal/websockets"
 )
 
-func (storeInstance *Store) AgentPing(agentTarget *Target) (bool, error) {
+func (storeInstance *Store) AgentPing(agentTarget *Target) bool {
 	splittedName := strings.Split(agentTarget.Name, " - ")
 	agentHostname := splittedName[0]
 
 	if storeInstance.WSHub == nil {
-		return false, nil
+		return false
 	}
 
-	err := storeInstance.WSHub.SendCommand(agentHostname, websockets.Message{
-		Type:    "ping",
-		Content: "ping",
-	})
-	if err != nil {
-		return false, err
+	storeInstance.WSHub.ClientsMux.RLock()
+	if client, ok := storeInstance.WSHub.Clients[agentHostname]; ok && client != nil {
+		storeInstance.WSHub.ClientsMux.RUnlock()
+		return true
 	}
+	storeInstance.WSHub.ClientsMux.RUnlock()
 
-	listener := storeInstance.WSHub.Broadcast.Subscribe()
-	defer storeInstance.WSHub.Broadcast.CancelSubscription(listener)
-
-	for {
-		select {
-		case resp := <-listener:
-			if resp.Type == "ping" && resp.Content == "pong" {
-				return true, nil
-			}
-		case <-time.After(time.Second * 2):
-			return false, nil
-		}
-	}
+	return false
 }
diff --git a/internal/store/db.go b/internal/store/db.go
index 7299ee7..18c61df 100644
--- a/internal/store/db.go
+++ b/internal/store/db.go
@@ -15,7 +15,6 @@ import (
 	"time"
 
 	_ "github.com/mattn/go-sqlite3"
-	"github.com/sonroyaalmerol/pbs-plus/internal/syslog"
 	"github.com/sonroyaalmerol/pbs-plus/internal/utils"
 	"github.com/sonroyaalmerol/pbs-plus/internal/websockets"
 )
@@ -469,16 +468,8 @@ func (store *Store) GetTarget(name string) (*Target, error) {
 		return nil, fmt.Errorf("GetTarget: error scanning row from targets table -> %w", err)
 	}
 
-	syslogger, err := syslog.InitializeLogger()
-	if err != nil {
-		return nil, fmt.Errorf("GetTarget: failed to initialize logger -> %w", err)
-	}
-
 	if strings.HasPrefix(target.Path, "agent://") {
-		target.ConnectionStatus, err = store.AgentPing(&target)
-		if err != nil {
-			syslogger.Errorf("GetTarget: error agent ping -> %v", err)
-		}
+		target.ConnectionStatus = store.AgentPing(&target)
 		target.IsAgent = true
 	} else {
 		target.ConnectionStatus = utils.IsValid(target.Path)
@@ -505,15 +496,8 @@ func (store *Store) GetAllTargetsByIP(ip string) ([]Target, error) {
 			return nil, fmt.Errorf("GetAllTargetsByIP: error scanning row from targets -> %w", err)
 		}
 
-		syslogger, err := syslog.InitializeLogger()
-		if err != nil {
-			return nil, fmt.Errorf("GetAllTargetsByIP: failed to initialize logger -> %w", err)
-		}
 		if strings.HasPrefix(target.Path, "agent://") {
-			target.ConnectionStatus, err = store.AgentPing(&target)
-			if err != nil {
-				syslogger.Errorf("GetAllTargetsByIP: error agent ping -> %v", err)
-			}
+			target.ConnectionStatus = store.AgentPing(&target)
 			target.IsAgent = true
 		} else {
 			target.ConnectionStatus = utils.IsValid(target.Path)
@@ -564,15 +548,8 @@ func (store *Store) GetAllTargets() ([]Target, error) {
 			return nil, fmt.Errorf("GetAllTargets: error scanning row from targets -> %w", err)
 		}
 
-		syslogger, err := syslog.InitializeLogger()
-		if err != nil {
-			return nil, fmt.Errorf("GetTarget: failed to initialize logger -> %w", err)
-		}
 		if strings.HasPrefix(target.Path, "agent://") {
-			target.ConnectionStatus, err = store.AgentPing(&target)
-			if err != nil {
-				syslogger.Errorf("GetTarget: error agent ping -> %v", err)
-			}
+			target.ConnectionStatus = store.AgentPing(&target)
 			target.IsAgent = true
 		} else {
 			target.ConnectionStatus = utils.IsValid(target.Path)
diff --git a/internal/websockets/server.go b/internal/websockets/server.go
index 86627bf..d2a4881 100644
--- a/internal/websockets/server.go
+++ b/internal/websockets/server.go
@@ -80,32 +80,10 @@ func (s *Server) HandleClientConnection(w http.ResponseWriter, r *http.Request)
 	clientID := initMessage.Content
 	client := NewClient(clientID, conn, s.Broadcast)
 
-	// Subscribe to broadcasts
-	subscription := s.Broadcast.Subscribe()
-
 	s.ClientsMux.Lock()
 	s.Clients[clientID] = client
 	s.ClientsMux.Unlock()
 
-	// Start broadcast handler
-	go func() {
-		defer s.Broadcast.CancelSubscription(subscription)
-		for {
-			select {
-			case msg := <-subscription:
-				select {
-				case client.send <- msg:
-					// Message queued successfully
-				default:
-					// Client's buffer is full, consider disconnecting
-					log.Printf("Client %s message buffer full, dropping message", clientID)
-				}
-			case <-client.done:
-				return
-			}
-		}
-	}()
-
 	// Start the read/write pumps
 	go client.readPump(s)
 	go client.writePump()