[management] Store connected proxies in DB (#5472)

Co-authored-by: mlsmaycon <mlsmaycon@gmail.com>

management/internals/modules/reverseproxy/proxy/manager/controller.go

@@ -0,0 +1,88 @@
+package manager
+
+import (
+	"context"
+	"sync"
+
+	log "github.com/sirupsen/logrus"
+	"go.opentelemetry.io/otel/metric"
+
+	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
+	nbgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
+	"github.com/netbirdio/netbird/shared/management/proto"
+)
+
+// GRPCController is a concrete implementation that manages proxy clusters and sends updates directly via gRPC.
+type GRPCController struct {
+	proxyGRPCServer *nbgrpc.ProxyServiceServer
+	// Map of cluster address -> set of proxy IDs
+	clusterProxies sync.Map
+	metrics        *metrics
+}
+
+// NewGRPCController creates a new GRPCController.
+func NewGRPCController(proxyGRPCServer *nbgrpc.ProxyServiceServer, meter metric.Meter) (*GRPCController, error) {
+	m, err := newMetrics(meter)
+	if err != nil {
+		return nil, err
+	}
+
+	return &GRPCController{
+		proxyGRPCServer: proxyGRPCServer,
+		metrics:         m,
+	}, nil
+}
+
+// SendServiceUpdateToCluster sends a service update to a specific proxy cluster.
+func (c *GRPCController) SendServiceUpdateToCluster(ctx context.Context, accountID string, update *proto.ProxyMapping, clusterAddr string) {
+	c.proxyGRPCServer.SendServiceUpdateToCluster(ctx, update, clusterAddr)
+	c.metrics.IncrementServiceUpdateSendCount(clusterAddr)
+}
+
+// GetOIDCValidationConfig returns the OIDC validation configuration from the gRPC server.
+func (c *GRPCController) GetOIDCValidationConfig() proxy.OIDCValidationConfig {
+	return c.proxyGRPCServer.GetOIDCValidationConfig()
+}
+
+// RegisterProxyToCluster registers a proxy to a specific cluster for routing.
+func (c *GRPCController) RegisterProxyToCluster(ctx context.Context, clusterAddr, proxyID string) error {
+	if clusterAddr == "" {
+		return nil
+	}
+	proxySet, _ := c.clusterProxies.LoadOrStore(clusterAddr, &sync.Map{})
+	proxySet.(*sync.Map).Store(proxyID, struct{}{})
+	log.WithContext(ctx).Debugf("Registered proxy %s to cluster %s", proxyID, clusterAddr)
+
+	c.metrics.IncrementProxyConnectionCount(clusterAddr)
+
+	return nil
+}
+
+// UnregisterProxyFromCluster removes a proxy from a cluster.
+func (c *GRPCController) UnregisterProxyFromCluster(ctx context.Context, clusterAddr, proxyID string) error {
+	if clusterAddr == "" {
+		return nil
+	}
+	if proxySet, ok := c.clusterProxies.Load(clusterAddr); ok {
+		proxySet.(*sync.Map).Delete(proxyID)
+		log.WithContext(ctx).Debugf("Unregistered proxy %s from cluster %s", proxyID, clusterAddr)
+
+		c.metrics.DecrementProxyConnectionCount(clusterAddr)
+	}
+	return nil
+}
+
+// GetProxiesForCluster returns all proxy IDs registered for a specific cluster.
+func (c *GRPCController) GetProxiesForCluster(clusterAddr string) []string {
+	proxySet, ok := c.clusterProxies.Load(clusterAddr)
+	if !ok {
+		return nil
+	}
+
+	var proxies []string
+	proxySet.(*sync.Map).Range(func(key, _ interface{}) bool {
+		proxies = append(proxies, key.(string))
+		return true
+	})
+	return proxies
+}

management/internals/modules/reverseproxy/proxy/manager/manager.go

@@ -0,0 +1,115 @@
+package manager
+
+import (
+	"context"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+	"go.opentelemetry.io/otel/metric"
+
+	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
+)
+
+// store defines the interface for proxy persistence operations
+type store interface {
+	SaveProxy(ctx context.Context, p *proxy.Proxy) error
+	UpdateProxyHeartbeat(ctx context.Context, proxyID string) error
+	GetActiveProxyClusterAddresses(ctx context.Context) ([]string, error)
+	CleanupStaleProxies(ctx context.Context, inactivityDuration time.Duration) error
+}
+
+// Manager handles all proxy operations
+type Manager struct {
+	store   store
+	metrics *metrics
+}
+
+// NewManager creates a new proxy Manager
+func NewManager(store store, meter metric.Meter) (*Manager, error) {
+	m, err := newMetrics(meter)
+	if err != nil {
+		return nil, err
+	}
+
+	return &Manager{
+		store:   store,
+		metrics: m,
+	}, nil
+}
+
+// Connect registers a new proxy connection in the database
+func (m Manager) Connect(ctx context.Context, proxyID, clusterAddress, ipAddress string) error {
+	now := time.Now()
+	p := &proxy.Proxy{
+		ID:             proxyID,
+		ClusterAddress: clusterAddress,
+		IPAddress:      ipAddress,
+		LastSeen:       now,
+		ConnectedAt:    &now,
+		Status:         "connected",
+	}
+
+	if err := m.store.SaveProxy(ctx, p); err != nil {
+		log.WithContext(ctx).Errorf("failed to register proxy %s: %v", proxyID, err)
+		return err
+	}
+
+	log.WithContext(ctx).WithFields(log.Fields{
+		"proxyID":        proxyID,
+		"clusterAddress": clusterAddress,
+		"ipAddress":      ipAddress,
+	}).Info("proxy connected")
+
+	return nil
+}
+
+// Disconnect marks a proxy as disconnected in the database
+func (m Manager) Disconnect(ctx context.Context, proxyID string) error {
+	now := time.Now()
+	p := &proxy.Proxy{
+		ID:             proxyID,
+		Status:         "disconnected",
+		DisconnectedAt: &now,
+		LastSeen:       now,
+	}
+
+	if err := m.store.SaveProxy(ctx, p); err != nil {
+		log.WithContext(ctx).Errorf("failed to disconnect proxy %s: %v", proxyID, err)
+		return err
+	}
+
+	log.WithContext(ctx).WithFields(log.Fields{
+		"proxyID": proxyID,
+	}).Info("proxy disconnected")
+
+	return nil
+}
+
+// Heartbeat updates the proxy's last seen timestamp
+func (m Manager) Heartbeat(ctx context.Context, proxyID string) error {
+	if err := m.store.UpdateProxyHeartbeat(ctx, proxyID); err != nil {
+		log.WithContext(ctx).Debugf("failed to update proxy %s heartbeat: %v", proxyID, err)
+		return err
+	}
+	m.metrics.IncrementProxyHeartbeatCount()
+	return nil
+}
+
+// GetActiveClusterAddresses returns all unique cluster addresses for active proxies
+func (m Manager) GetActiveClusterAddresses(ctx context.Context) ([]string, error) {
+	addresses, err := m.store.GetActiveProxyClusterAddresses(ctx)
+	if err != nil {
+		log.WithContext(ctx).Errorf("failed to get active proxy cluster addresses: %v", err)
+		return nil, err
+	}
+	return addresses, nil
+}
+
+// CleanupStale removes proxies that haven't sent heartbeat in the specified duration
+func (m Manager) CleanupStale(ctx context.Context, inactivityDuration time.Duration) error {
+	if err := m.store.CleanupStaleProxies(ctx, inactivityDuration); err != nil {
+		log.WithContext(ctx).Errorf("failed to cleanup stale proxies: %v", err)
+		return err
+	}
+	return nil
+}

management/internals/modules/reverseproxy/proxy/manager/metrics.go

@@ -0,0 +1,74 @@
+package manager
+
+import (
+	"context"
+
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/metric"
+)
+
+type metrics struct {
+	proxyConnectionCount   metric.Int64UpDownCounter
+	serviceUpdateSendCount metric.Int64Counter
+	proxyHeartbeatCount    metric.Int64Counter
+}
+
+func newMetrics(meter metric.Meter) (*metrics, error) {
+	proxyConnectionCount, err := meter.Int64UpDownCounter(
+		"management_proxy_connection_count",
+		metric.WithDescription("Number of active proxy connections"),
+		metric.WithUnit("{connection}"),
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	serviceUpdateSendCount, err := meter.Int64Counter(
+		"management_proxy_service_update_send_count",
+		metric.WithDescription("Total number of service updates sent to proxies"),
+		metric.WithUnit("{update}"),
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	proxyHeartbeatCount, err := meter.Int64Counter(
+		"management_proxy_heartbeat_count",
+		metric.WithDescription("Total number of proxy heartbeats received"),
+		metric.WithUnit("{heartbeat}"),
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	return &metrics{
+		proxyConnectionCount:   proxyConnectionCount,
+		serviceUpdateSendCount: serviceUpdateSendCount,
+		proxyHeartbeatCount:    proxyHeartbeatCount,
+	}, nil
+}
+
+func (m *metrics) IncrementProxyConnectionCount(clusterAddr string) {
+	m.proxyConnectionCount.Add(context.Background(), 1,
+		metric.WithAttributes(
+			attribute.String("cluster", clusterAddr),
+		))
+}
+
+func (m *metrics) DecrementProxyConnectionCount(clusterAddr string) {
+	m.proxyConnectionCount.Add(context.Background(), -1,
+		metric.WithAttributes(
+			attribute.String("cluster", clusterAddr),
+		))
+}
+
+func (m *metrics) IncrementServiceUpdateSendCount(clusterAddr string) {
+	m.serviceUpdateSendCount.Add(context.Background(), 1,
+		metric.WithAttributes(
+			attribute.String("cluster", clusterAddr),
+		))
+}
+
+func (m *metrics) IncrementProxyHeartbeatCount() {
+	m.proxyHeartbeatCount.Add(context.Background(), 1)
+}