From 1de3c90146d0ea8a6af6ec7446495dc56add0f40 Mon Sep 17 00:00:00 2001
From: Erik Brakkee <erik@brakkee.org>
Date: Thu, 15 Aug 2024 22:16:20 +0200
Subject: [PATCH] Now using a map of ClientId to clientConnection which is more
 efficient for deleting clients.

---
 pkg/server/admin/admin.go | 30 ++++++++++++------------------
 1 file changed, 12 insertions(+), 18 deletions(-)

diff --git a/pkg/server/admin/admin.go b/pkg/server/admin/admin.go
index bef5de8..9043cee 100644
--- a/pkg/server/admin/admin.go
+++ b/pkg/server/admin/admin.go
@@ -21,10 +21,9 @@ type agentConnection struct {
 	CommChannel comms.CommChannel
 }
 
-var agentIdGenerator = concurrency.NewAtomicCounter()
 var clientIdGenerator = concurrency.NewAtomicCounter()
 
-type ClientConnection struct {
+type clientConnection struct {
 	Info             *models.Client
 	agentConnection  net.Conn
 	clientConnection iowrappers2.ReadWriteAddrCloser
@@ -45,7 +44,7 @@ func newAgent(commChannel comms.CommChannel, publicId models.RendezVousId, agent
 }
 
 func newClient(publicId models.RendezVousId, clientConn iowrappers2.ReadWriteAddrCloser,
-	agentConn net.Conn, agentGuid models.AgentGuid) *ClientConnection {
+	agentConn net.Conn, agentGuid models.AgentGuid) *clientConnection {
 	client := models.Client{
 		Guid:       models.ClientGuid(strconv.Itoa(rand.Int())),
 		RemoteAddr: models.RemoteAddr(clientConn.RemoteAddr().String()),
@@ -54,14 +53,14 @@ func newClient(publicId models.RendezVousId, clientConn iowrappers2.ReadWriteAdd
 		ClientId:   models.ClientId(strconv.Itoa(clientIdGenerator.IncrementAndGet())),
 		StartTime:  time.Now(),
 	}
-	return &ClientConnection{
+	return &clientConnection{
 		Info:             &client,
 		agentConnection:  agentConn,
 		clientConnection: clientConn,
 	}
 }
 
-func (match *ClientConnection) Synchronize() {
+func (match *clientConnection) Synchronize() {
 	iowrappers2.SynchronizeStreams("client -- agent", match.clientConnection, match.agentConnection)
 }
 
@@ -75,7 +74,7 @@ type Admin struct {
 
 	// TODO: use linked map for both of these
 	agents  map[models.RendezVousId]*agentConnection
-	clients []*ClientConnection
+	clients map[models.ClientId]*clientConnection
 }
 
 func NewAdmin() *Admin {
@@ -83,7 +82,7 @@ func NewAdmin() *Admin {
 		mutex:   sync.Mutex{},
 		state:   models.NewState(),
 		agents:  make(map[models.RendezVousId]*agentConnection),
-		clients: make([]*ClientConnection, 0), // not strictly needed
+		clients: make(map[models.ClientId]*clientConnection),
 	}
 }
 
@@ -146,7 +145,7 @@ func (admin *Admin) AddAgent(publicId models.RendezVousId, agentInfo comms.Envir
 	return agent, nil
 }
 
-func (admin *Admin) AddClient(publicId models.RendezVousId, clientConn iowrappers2.ReadWriteAddrCloser) (*ClientConnection, error) {
+func (admin *Admin) AddClient(publicId models.RendezVousId, clientConn iowrappers2.ReadWriteAddrCloser) (*clientConnection, error) {
 	admin.mutex.Lock()
 	defer admin.mutex.Unlock()
 
@@ -177,7 +176,7 @@ func (admin *Admin) AddClient(publicId models.RendezVousId, clientConn iowrapper
 
 	admin.state = admin.state.Copy()
 	admin.state.Clients.Put(client.Info.Guid, client.Info)
-	admin.clients = append(admin.clients, client)
+	admin.clients[client.Info.ClientId] = client
 	return client, nil
 }
 
@@ -212,7 +211,7 @@ func (admin *Admin) RemoveAgent(publicId models.RendezVousId) error {
 	return nil
 }
 
-func (admin *Admin) RemoveClient(client *ClientConnection) error {
+func (admin *Admin) RemoveClient(client *clientConnection) error {
 	admin.mutex.Lock()
 	defer admin.mutex.Unlock()
 
@@ -222,14 +221,9 @@ func (admin *Admin) RemoveClient(client *ClientConnection) error {
 	_ = client.agentConnection.Close()
 	_ = client.clientConnection.Close()
 
-	for i, _client := range admin.clients {
-		if _client.Info.ClientId == client.Info.ClientId {
-			admin.state = admin.state.Copy()
-			admin.state.Clients.Delete(client.Info.Guid)
-			admin.clients = append(admin.clients[:i], admin.clients[i+1:]...)
-			break
-		}
-	}
+	admin.state = admin.state.Copy()
+	admin.state.Clients.Delete(client.Info.Guid)
+	delete(admin.clients, client.Info.ClientId)
 	return nil
 }