mirror of
https://github.com/netbirdio/netbird.git
synced 2026-09-25 08:09:07 +02:00
endpoint model discovery and proxy integration
This commit is contained in:
@@ -4,11 +4,16 @@ package proxy
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/netbirdio/netbird/shared/management/proto"
|
||||
)
|
||||
|
||||
// ErrModelDiscoveryUnavailable is returned when no connected proxy in the
|
||||
// requested cluster can execute a model-discovery request.
|
||||
var ErrModelDiscoveryUnavailable = errors.New("model discovery proxy unavailable")
|
||||
|
||||
// Manager defines the interface for proxy operations
|
||||
type Manager interface {
|
||||
Connect(ctx context.Context, proxyID, sessionID, clusterAddress, ipAddress string, accountID *string, capabilities *Capabilities) (*Proxy, error)
|
||||
@@ -38,6 +43,7 @@ type OIDCValidationConfig struct {
|
||||
// Controller is responsible for managing proxy clusters and routing service updates.
|
||||
type Controller interface {
|
||||
SendServiceUpdateToCluster(ctx context.Context, accountID string, update *proto.ProxyMapping, clusterAddr string)
|
||||
DiscoverModels(ctx context.Context, accountID, clusterAddr string, req *proto.ModelDiscoveryRequest) (*proto.ModelDiscoveryResult, error)
|
||||
GetOIDCValidationConfig() OIDCValidationConfig
|
||||
RegisterProxyToCluster(ctx context.Context, clusterAddr, proxyID string) error
|
||||
UnregisterProxyFromCluster(ctx context.Context, clusterAddr, proxyID string) error
|
||||
|
||||
@@ -39,6 +39,12 @@ func (c *GRPCController) SendServiceUpdateToCluster(ctx context.Context, account
|
||||
c.metrics.IncrementServiceUpdateSendCount(clusterAddr)
|
||||
}
|
||||
|
||||
// DiscoverModels executes a correlated model-discovery request on one capable
|
||||
// proxy connected to the requested cluster.
|
||||
func (c *GRPCController) DiscoverModels(ctx context.Context, accountID, clusterAddr string, req *proto.ModelDiscoveryRequest) (*proto.ModelDiscoveryResult, error) {
|
||||
return c.proxyGRPCServer.DiscoverModels(ctx, accountID, clusterAddr, req)
|
||||
}
|
||||
|
||||
// GetOIDCValidationConfig returns the OIDC validation configuration from the gRPC server.
|
||||
func (c *GRPCController) GetOIDCValidationConfig() proxy.OIDCValidationConfig {
|
||||
return c.proxyGRPCServer.GetOIDCValidationConfig()
|
||||
@@ -50,10 +56,12 @@ func (c *GRPCController) RegisterProxyToCluster(ctx context.Context, clusterAddr
|
||||
return nil
|
||||
}
|
||||
proxySet, _ := c.clusterProxies.LoadOrStore(clusterAddr, &sync.Map{})
|
||||
proxySet.(*sync.Map).Store(proxyID, struct{}{})
|
||||
_, alreadyRegistered := proxySet.(*sync.Map).LoadOrStore(proxyID, struct{}{})
|
||||
log.WithContext(ctx).Debugf("Registered proxy %s to cluster %s", proxyID, clusterAddr)
|
||||
|
||||
c.metrics.IncrementProxyConnectionCount(clusterAddr)
|
||||
if !alreadyRegistered {
|
||||
c.metrics.IncrementProxyConnectionCount(clusterAddr)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -64,10 +72,10 @@ func (c *GRPCController) UnregisterProxyFromCluster(ctx context.Context, cluster
|
||||
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)
|
||||
if _, registered := proxySet.(*sync.Map).LoadAndDelete(proxyID); registered {
|
||||
log.WithContext(ctx).Debugf("Unregistered proxy %s from cluster %s", proxyID, clusterAddr)
|
||||
c.metrics.DecrementProxyConnectionCount(clusterAddr)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -259,6 +259,21 @@ func (m *MockController) EXPECT() *MockControllerMockRecorder {
|
||||
return m.recorder
|
||||
}
|
||||
|
||||
// DiscoverModels mocks base method.
|
||||
func (m *MockController) DiscoverModels(ctx context.Context, accountID, clusterAddr string, req *proto.ModelDiscoveryRequest) (*proto.ModelDiscoveryResult, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "DiscoverModels", ctx, accountID, clusterAddr, req)
|
||||
ret0, _ := ret[0].(*proto.ModelDiscoveryResult)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// DiscoverModels indicates an expected call of DiscoverModels.
|
||||
func (mr *MockControllerMockRecorder) DiscoverModels(ctx, accountID, clusterAddr, req interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DiscoverModels", reflect.TypeOf((*MockController)(nil).DiscoverModels), ctx, accountID, clusterAddr, req)
|
||||
}
|
||||
|
||||
// GetOIDCValidationConfig mocks base method.
|
||||
func (m *MockController) GetOIDCValidationConfig() OIDCValidationConfig {
|
||||
m.ctrl.T.Helper()
|
||||
|
||||
Reference in New Issue
Block a user