diff options
Diffstat (limited to '')
| -rw-r--r-- | internal/controlplane/manager.go | 20 |
1 files changed, 18 insertions, 2 deletions
diff --git a/internal/controlplane/manager.go b/internal/controlplane/manager.go index 42f5efe..cdafc36 100644 --- a/internal/controlplane/manager.go +++ b/internal/controlplane/manager.go @@ -10,6 +10,7 @@ import ( "aigw/internal/auth" "aigw/internal/catalog" + "aigw/internal/domain" ) const broadcastQueueSize = 128 @@ -22,6 +23,10 @@ type managerStore interface { RedisEnabled() bool } +type policyReplacer interface { + ReplacePolicies([]domain.LimitPolicy) +} + type Manager struct { store managerStore catalog *catalog.Catalog @@ -32,15 +37,20 @@ type Manager struct { redisConnected atomic.Bool reloadMu sync.Mutex broadcasts chan ChangeEvent + policyTarget policyReplacer } -func NewManager(store managerStore, modelCatalog *catalog.Catalog, authenticator *auth.StaticAuthenticator, logger *slog.Logger, pollInterval time.Duration) *Manager { +func NewManager(store managerStore, modelCatalog *catalog.Catalog, authenticator *auth.StaticAuthenticator, logger *slog.Logger, pollInterval time.Duration, policyTargets ...policyReplacer) *Manager { if pollInterval <= 0 { pollInterval = 30 * time.Second } + var policyTarget policyReplacer + if len(policyTargets) > 0 { + policyTarget = policyTargets[0] + } return &Manager{ store: store, catalog: modelCatalog, authenticator: authenticator, - logger: logger, pollInterval: pollInterval, broadcasts: make(chan ChangeEvent, broadcastQueueSize), + logger: logger, pollInterval: pollInterval, broadcasts: make(chan ChangeEvent, broadcastQueueSize), policyTarget: policyTarget, } } @@ -53,6 +63,9 @@ func (m *Manager) Reload(ctx context.Context) (int64, error) { } m.catalog.Replace(snapshot.Models) m.authenticator.ReplaceHashed(snapshot.APIKeys) + if m.policyTarget != nil { + m.policyTarget.ReplacePolicies(snapshot.Limits) + } m.generation.Store(snapshot.Generation) m.logger.Info("control_plane_reloaded", "generation", snapshot.Generation, "models", len(snapshot.Models), "api_keys", len(snapshot.APIKeys)) return snapshot.Generation, nil @@ -195,7 +208,10 @@ func (m *Manager) runBroadcasts(ctx context.Context) { err := m.store.PublishChange(publishContext, event) cancel() if err != nil { + m.redisConnected.Store(false) m.logger.Warn("control_plane_publish_failed", "generation", event.Generation, "resource", event.Resource, "id", event.ID, "error", err) + } else { + m.redisConnected.Store(true) } } } |
