Files
2026-05-07 10:16:46 +08:00

151 lines
3.5 KiB
Go

package control
import (
"sync"
"time"
)
// ModuleState represents the lifecycle state of a module
type ModuleState string
const (
ModuleStateActive ModuleState = "active"
ModuleStateClosing ModuleState = "closing"
ModuleStateClosed ModuleState = "closed"
)
// ModuleGate controls the enable/disable/close lifecycle of a module
type ModuleGate struct {
mu sync.RWMutex
enabled bool
state ModuleState
closedAt *time.Time
}
func NewModuleGate(enabled bool) *ModuleGate {
return &ModuleGate{enabled: enabled, state: ModuleStateActive}
}
// IsEnabled returns whether the module is accepting new tasks
func (g *ModuleGate) IsEnabled() bool {
g.mu.RLock()
defer g.mu.RUnlock()
return g.enabled && g.state == ModuleStateActive
}
// Close signals the module to stop accepting new tasks
func (g *ModuleGate) Close() {
g.mu.Lock()
defer g.mu.Unlock()
if g.state == ModuleStateActive {
g.state = ModuleStateClosing
now := time.Now().UTC()
g.closedAt = &now
}
}
// MarkClosed marks the module as fully closed (no in-flight tasks)
func (g *ModuleGate) MarkClosed() {
g.mu.Lock()
defer g.mu.Unlock()
g.state = ModuleStateClosed
g.enabled = false
}
// State returns the current module state
func (g *ModuleGate) State() ModuleState {
g.mu.RLock()
defer g.mu.RUnlock()
return g.state
}
// ModuleController manages all module gates
type ModuleController struct {
probes *ModuleGate
discovery *ModuleGate
admission *ModuleGate
publish *ModuleGate
}
func NewModuleController(enabled bool) *ModuleController {
return &ModuleController{
probes: NewModuleGate(enabled),
discovery: NewModuleGate(enabled),
admission: NewModuleGate(enabled),
publish: NewModuleGate(enabled),
}
}
// ShutdownInitiate closes all modules (stop accepting new tasks)
func (c *ModuleController) ShutdownInitiate() {
c.probes.Close()
c.discovery.Close()
c.admission.Close()
c.publish.Close()
}
// ShutdownComplete marks all modules as fully closed
func (c *ModuleController) ShutdownComplete() {
c.probes.MarkClosed()
c.discovery.MarkClosed()
c.admission.MarkClosed()
c.publish.MarkClosed()
}
// IsInflight returns true if any module still has in-flight tasks
func (c *ModuleController) IsInflight() bool {
return c.probes.State() == ModuleStateClosing ||
c.discovery.State() == ModuleStateClosing ||
c.admission.State() == ModuleStateClosing ||
c.publish.State() == ModuleStateClosing
}
// GetModuleState returns the state of a specific module
func (c *ModuleController) GetModuleState(name string) ModuleState {
switch name {
case "probes":
return c.probes.State()
case "discovery":
return c.discovery.State()
case "admission":
return c.admission.State()
case "publish":
return c.publish.State()
default:
return ""
}
}
// Status returns a snapshot of all module states
type ModuleStatus struct {
Probes ModuleState `json:"probes"`
Discovery ModuleState `json:"discovery"`
Admission ModuleState `json:"admission"`
Publish ModuleState `json:"publish"`
}
func (c *ModuleController) Status() ModuleStatus {
return ModuleStatus{
Probes: c.probes.State(),
Discovery: c.discovery.State(),
Admission: c.admission.State(),
Publish: c.publish.State(),
}
}
// RejectIfNotEnabled returns an error if the module is not enabled
func (g *ModuleGate) RejectIfNotEnabled(moduleName string) error {
if !g.IsEnabled() {
return ErrModuleClosed
}
return nil
}
var ErrModuleClosed = &ModuleClosedError{}
type ModuleClosedError struct{}
func (e *ModuleClosedError) Error() string {
return "module is not accepting new tasks"
}