feat(federated): allow to pickup a specific worker, improve loadbalancing (#3243)

* feat(explorer): allow to specify a worker target

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* feat(explorer): correctly load balance requests

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* feat(explorer): mark load balanced by default

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* fix: make sure to delete tunnels that might not exist anymore

If a worker goes off and on might change tunnel address, and we want to
load balance only on the active tunnels.

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

---------

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

core/p2p/federated.go

@@ -1,6 +1,12 @@
 package p2p
 
-import "fmt"
+import (
+	"fmt"
+	"math/rand/v2"
+	"sync"
+
+	"github.com/rs/zerolog/log"
+)
 
 const FederatedID = "federated"
 
@@ -12,22 +18,70 @@ func NetworkID(networkID, serviceID string) string {
 }
 
 type FederatedServer struct {
+	sync.Mutex
 	listenAddr, service, p2ptoken string
 	requestTable                  map[string]int
 	loadBalanced                  bool
+	workerTarget                  string
 }
 
-func NewFederatedServer(listenAddr, service, p2pToken string, loadBalanced bool) *FederatedServer {
+func NewFederatedServer(listenAddr, service, p2pToken string, loadBalanced bool, workerTarget string) *FederatedServer {
 	return &FederatedServer{
 		listenAddr:   listenAddr,
 		service:      service,
 		p2ptoken:     p2pToken,
 		requestTable: map[string]int{},
 		loadBalanced: loadBalanced,
+		workerTarget: workerTarget,
+	}
+}
+
+func (fs *FederatedServer) RandomServer() string {
+	var tunnelAddresses []string
+	for _, v := range GetAvailableNodes(fs.service) {
+		if v.IsOnline() {
+			tunnelAddresses = append(tunnelAddresses, v.TunnelAddress)
+		} else {
+			delete(fs.requestTable, v.TunnelAddress) // make sure it's not tracked
+			log.Info().Msgf("Node %s is offline", v.ID)
+		}
+	}
+
+	if len(tunnelAddresses) == 0 {
+		return ""
+	}
+
+	return tunnelAddresses[rand.IntN(len(tunnelAddresses))]
+}
+
+func (fs *FederatedServer) syncTableStatus() {
+	fs.Lock()
+	defer fs.Unlock()
+	currentTunnels := make(map[string]struct{})
+
+	for _, v := range GetAvailableNodes(fs.service) {
+		if v.IsOnline() {
+			fs.ensureRecordExist(v.TunnelAddress)
+			currentTunnels[v.TunnelAddress] = struct{}{}
+		}
+	}
+
+	// delete tunnels that don't exist anymore
+	for t := range fs.requestTable {
+		if _, ok := currentTunnels[t]; !ok {
+			delete(fs.requestTable, t)
+		}
 	}
 }
 
 func (fs *FederatedServer) SelectLeastUsedServer() string {
+	fs.syncTableStatus()
+
+	fs.Lock()
+	defer fs.Unlock()
+
+	log.Debug().Any("request_table", fs.requestTable).Msgf("Current request table")
+
 	// cycle over requestTable and find the entry with the lower number
 	// if there are multiple entries with the same number, select one randomly
 	// if there are no entries, return an empty string
@@ -39,18 +93,26 @@ func (fs *FederatedServer) SelectLeastUsedServer() string {
 			minKey = k
 		}
 	}
+	log.Debug().Any("requests_served", min).Msgf("Selected tunnel %s", minKey)
+
 	return minKey
 }
 
 func (fs *FederatedServer) RecordRequest(nodeID string) {
+	fs.Lock()
+	defer fs.Unlock()
 	// increment the counter for the nodeID in the requestTable
 	fs.requestTable[nodeID]++
+
+	log.Debug().Any("request_table", fs.requestTable).Msgf("Current request table")
 }
 
-func (fs *FederatedServer) EnsureRecordExist(nodeID string) {
+func (fs *FederatedServer) ensureRecordExist(nodeID string) {
 	// if the nodeID is not in the requestTable, add it with a counter of 0
 	_, ok := fs.requestTable[nodeID]
 	if !ok {
 		fs.requestTable[nodeID] = 0
 	}
+
+	log.Debug().Any("request_table", fs.requestTable).Msgf("Current request table")
 }