package lifecycle import ( "context" "fmt" "os" "os/signal" "sync" "syscall" "time" ) // ShutdownManager handles graceful shutdown of the application type ShutdownManager struct { registry *ModuleRegistry shutdownTasks []ShutdownTask shutdownHooks []ShutdownHook config ShutdownConfig signalChannel chan os.Signal shutdownChannel chan struct{} state ShutdownState startTime time.Time shutdownStarted time.Time mu sync.RWMutex wg sync.WaitGroup ctx context.Context cancel context.CancelFunc } // ShutdownTask represents a task to be executed during shutdown type ShutdownTask struct { Name string Priority int Timeout time.Duration Task func(ctx context.Context) error OnError func(error) Critical bool Enabled bool } // ShutdownHook is called at different stages of shutdown type ShutdownHook interface { OnShutdownStarted(ctx context.Context) error OnModulesStopped(ctx context.Context) error OnCleanupStarted(ctx context.Context) error OnShutdownCompleted(ctx context.Context) error OnShutdownFailed(ctx context.Context, err error) error } // ShutdownConfig configures shutdown behavior type ShutdownConfig struct { GracefulTimeout time.Duration `json:"graceful_timeout"` ForceTimeout time.Duration `json:"force_timeout"` SignalBufferSize int `json:"signal_buffer_size"` MaxRetries int `json:"max_retries"` RetryDelay time.Duration `json:"retry_delay"` ParallelShutdown bool `json:"parallel_shutdown"` SaveState bool `json:"save_state"` CleanupTempFiles bool `json:"cleanup_temp_files"` NotifyExternal bool `json:"notify_external"` WaitForConnections bool `json:"wait_for_connections"` EnableMetrics bool `json:"enable_metrics"` } // ShutdownState represents the current shutdown state type ShutdownState string const ( ShutdownStateRunning ShutdownState = "running" ShutdownStateInitiated ShutdownState = "initiated" ShutdownStateModuleStop ShutdownState = "stopping_modules" ShutdownStateCleanup ShutdownState = "cleanup" ShutdownStateCompleted ShutdownState = "completed" ShutdownStateFailed ShutdownState = "failed" ShutdownStateForced ShutdownState = "forced" ) // ShutdownMetrics tracks shutdown performance type ShutdownMetrics struct { ShutdownInitiated time.Time `json:"shutdown_initiated"` ModuleStopTime time.Duration `json:"module_stop_time"` CleanupTime time.Duration `json:"cleanup_time"` TotalShutdownTime time.Duration `json:"total_shutdown_time"` TasksExecuted int `json:"tasks_executed"` TasksSuccessful int `json:"tasks_successful"` TasksFailed int `json:"tasks_failed"` RetryAttempts int `json:"retry_attempts"` ForceShutdown bool `json:"force_shutdown"` Signal string `json:"signal"` } // ShutdownProgress tracks shutdown progress type ShutdownProgress struct { State ShutdownState `json:"state"` Progress float64 `json:"progress"` CurrentTask string `json:"current_task"` CompletedTasks int `json:"completed_tasks"` TotalTasks int `json:"total_tasks"` ElapsedTime time.Duration `json:"elapsed_time"` EstimatedRemaining time.Duration `json:"estimated_remaining"` Message string `json:"message"` } // NewShutdownManager creates a new shutdown manager func NewShutdownManager(registry *ModuleRegistry, config ShutdownConfig) *ShutdownManager { ctx, cancel := context.WithCancel(context.Background()) sm := &ShutdownManager{ registry: registry, shutdownTasks: make([]ShutdownTask, 0), shutdownHooks: make([]ShutdownHook, 0), config: config, signalChannel: make(chan os.Signal, config.SignalBufferSize), shutdownChannel: make(chan struct{}), state: ShutdownStateRunning, startTime: time.Now(), ctx: ctx, cancel: cancel, } // Set default configuration if sm.config.GracefulTimeout == 0 { sm.config.GracefulTimeout = 30 * time.Second } if sm.config.ForceTimeout == 0 { sm.config.ForceTimeout = 60 * time.Second } if sm.config.SignalBufferSize == 0 { sm.config.SignalBufferSize = 10 } if sm.config.MaxRetries == 0 { sm.config.MaxRetries = 3 } if sm.config.RetryDelay == 0 { sm.config.RetryDelay = time.Second } // Setup default shutdown tasks sm.setupDefaultTasks() // Setup signal handling sm.setupSignalHandling() return sm } // Start starts the shutdown manager func (sm *ShutdownManager) Start() error { sm.mu.Lock() defer sm.mu.Unlock() if sm.state != ShutdownStateRunning { return fmt.Errorf("shutdown manager not in running state: %s", sm.state) } // Start signal monitoring go sm.signalHandler() return nil } // Shutdown initiates graceful shutdown func (sm *ShutdownManager) Shutdown(ctx context.Context) error { sm.mu.Lock() defer sm.mu.Unlock() if sm.state != ShutdownStateRunning { return fmt.Errorf("shutdown already initiated: %s", sm.state) } sm.state = ShutdownStateInitiated sm.shutdownStarted = time.Now() // Close shutdown channel to signal shutdown close(sm.shutdownChannel) return sm.performShutdown(ctx) } // ForceShutdown forces immediate shutdown func (sm *ShutdownManager) ForceShutdown(ctx context.Context) error { sm.mu.Lock() defer sm.mu.Unlock() sm.state = ShutdownStateForced sm.cancel() // Cancel all operations // Force stop all modules immediately if sm.registry != nil { forceCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() sm.registry.StopAll(forceCtx) } os.Exit(1) return nil } // AddShutdownTask adds a task to be executed during shutdown func (sm *ShutdownManager) AddShutdownTask(task ShutdownTask) { sm.mu.Lock() defer sm.mu.Unlock() sm.shutdownTasks = append(sm.shutdownTasks, task) sm.sortTasksByPriority() } // AddShutdownHook adds a hook to be called during shutdown phases func (sm *ShutdownManager) AddShutdownHook(hook ShutdownHook) { sm.mu.Lock() defer sm.mu.Unlock() sm.shutdownHooks = append(sm.shutdownHooks, hook) } // GetState returns the current shutdown state func (sm *ShutdownManager) GetState() ShutdownState { sm.mu.RLock() defer sm.mu.RUnlock() return sm.state } // GetProgress returns the current shutdown progress func (sm *ShutdownManager) GetProgress() ShutdownProgress { sm.mu.RLock() defer sm.mu.RUnlock() totalTasks := len(sm.shutdownTasks) if sm.registry != nil { totalTasks += len(sm.registry.List()) } var progress float64 var completedTasks int var currentTask string switch sm.state { case ShutdownStateRunning: progress = 0 currentTask = "Running" case ShutdownStateInitiated: progress = 0.1 currentTask = "Shutdown initiated" case ShutdownStateModuleStop: progress = 0.3 currentTask = "Stopping modules" completedTasks = totalTasks / 3 case ShutdownStateCleanup: progress = 0.7 currentTask = "Cleanup" completedTasks = (totalTasks * 2) / 3 case ShutdownStateCompleted: progress = 1.0 currentTask = "Completed" completedTasks = totalTasks case ShutdownStateFailed: progress = 0.8 currentTask = "Failed" case ShutdownStateForced: progress = 1.0 currentTask = "Forced shutdown" completedTasks = totalTasks } elapsedTime := time.Since(sm.shutdownStarted) var estimatedRemaining time.Duration if progress > 0 && progress < 1.0 { totalEstimated := time.Duration(float64(elapsedTime) / progress) estimatedRemaining = totalEstimated - elapsedTime } return ShutdownProgress{ State: sm.state, Progress: progress, CurrentTask: currentTask, CompletedTasks: completedTasks, TotalTasks: totalTasks, ElapsedTime: elapsedTime, EstimatedRemaining: estimatedRemaining, Message: fmt.Sprintf("Shutdown %s", sm.state), } } // Wait waits for shutdown to complete func (sm *ShutdownManager) Wait() { <-sm.shutdownChannel sm.wg.Wait() } // WaitWithTimeout waits for shutdown with timeout func (sm *ShutdownManager) WaitWithTimeout(timeout time.Duration) error { done := make(chan struct{}) go func() { sm.Wait() close(done) }() select { case <-done: return nil case <-time.After(timeout): return fmt.Errorf("shutdown timeout after %v", timeout) } } // Private methods func (sm *ShutdownManager) setupSignalHandling() { signal.Notify(sm.signalChannel, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGHUP, ) } func (sm *ShutdownManager) setupDefaultTasks() { // Task: Save application state sm.shutdownTasks = append(sm.shutdownTasks, ShutdownTask{ Name: "save_state", Priority: 100, Timeout: 10 * time.Second, Task: func(ctx context.Context) error { if sm.config.SaveState { return sm.saveApplicationState(ctx) } return nil }, Critical: false, Enabled: sm.config.SaveState, }) // Task: Close external connections sm.shutdownTasks = append(sm.shutdownTasks, ShutdownTask{ Name: "close_connections", Priority: 90, Timeout: 5 * time.Second, Task: func(ctx context.Context) error { return sm.closeExternalConnections(ctx) }, Critical: false, Enabled: sm.config.WaitForConnections, }) // Task: Cleanup temporary files sm.shutdownTasks = append(sm.shutdownTasks, ShutdownTask{ Name: "cleanup_temp_files", Priority: 10, Timeout: 5 * time.Second, Task: func(ctx context.Context) error { if sm.config.CleanupTempFiles { return sm.cleanupTempFiles(ctx) } return nil }, Critical: false, Enabled: sm.config.CleanupTempFiles, }) // Task: Notify external systems sm.shutdownTasks = append(sm.shutdownTasks, ShutdownTask{ Name: "notify_external", Priority: 80, Timeout: 3 * time.Second, Task: func(ctx context.Context) error { if sm.config.NotifyExternal { return sm.notifyExternalSystems(ctx) } return nil }, Critical: false, Enabled: sm.config.NotifyExternal, }) } func (sm *ShutdownManager) signalHandler() { for { select { case sig := <-sm.signalChannel: switch sig { case syscall.SIGINT, syscall.SIGTERM: // Graceful shutdown ctx, cancel := context.WithTimeout(context.Background(), sm.config.GracefulTimeout) if err := sm.Shutdown(ctx); err != nil { cancel() // Force shutdown if graceful fails forceCtx, forceCancel := context.WithTimeout(context.Background(), sm.config.ForceTimeout) sm.ForceShutdown(forceCtx) forceCancel() } cancel() return case syscall.SIGQUIT: // Force shutdown ctx, cancel := context.WithTimeout(context.Background(), sm.config.ForceTimeout) sm.ForceShutdown(ctx) cancel() return case syscall.SIGHUP: // Reload signal - could be used for configuration reload // For now, just log it continue } case <-sm.ctx.Done(): return } } } func (sm *ShutdownManager) performShutdown(ctx context.Context) error { sm.wg.Add(1) defer sm.wg.Done() // Create timeout context for entire shutdown shutdownCtx, cancel := context.WithTimeout(ctx, sm.config.GracefulTimeout) defer cancel() var shutdownErr error // Phase 1: Call shutdown started hooks if err := sm.callHooks(shutdownCtx, "OnShutdownStarted"); err != nil { shutdownErr = fmt.Errorf("shutdown hooks failed: %w", err) } // Phase 2: Stop modules sm.state = ShutdownStateModuleStop if sm.registry != nil { if err := sm.registry.StopAll(shutdownCtx); err != nil { shutdownErr = fmt.Errorf("failed to stop modules: %w", err) } } // Call modules stopped hooks if err := sm.callHooks(shutdownCtx, "OnModulesStopped"); err != nil { if shutdownErr == nil { shutdownErr = fmt.Errorf("modules stopped hooks failed: %w", err) } } // Phase 3: Execute shutdown tasks sm.state = ShutdownStateCleanup if err := sm.callHooks(shutdownCtx, "OnCleanupStarted"); err != nil { if shutdownErr == nil { shutdownErr = fmt.Errorf("cleanup hooks failed: %w", err) } } if err := sm.executeShutdownTasks(shutdownCtx); err != nil { if shutdownErr == nil { shutdownErr = fmt.Errorf("shutdown tasks failed: %w", err) } } // Phase 4: Final cleanup if shutdownErr != nil { sm.state = ShutdownStateFailed sm.callHooks(shutdownCtx, "OnShutdownFailed") } else { sm.state = ShutdownStateCompleted sm.callHooks(shutdownCtx, "OnShutdownCompleted") } return shutdownErr } func (sm *ShutdownManager) executeShutdownTasks(ctx context.Context) error { if sm.config.ParallelShutdown { return sm.executeTasksParallel(ctx) } else { return sm.executeTasksSequential(ctx) } } func (sm *ShutdownManager) executeTasksSequential(ctx context.Context) error { var lastErr error for _, task := range sm.shutdownTasks { if !task.Enabled { continue } if err := sm.executeTask(ctx, task); err != nil { lastErr = err if task.Critical { return fmt.Errorf("critical task %s failed: %w", task.Name, err) } } } return lastErr } func (sm *ShutdownManager) executeTasksParallel(ctx context.Context) error { var wg sync.WaitGroup errors := make(chan error, len(sm.shutdownTasks)) // Group tasks by priority priorityGroups := sm.groupTasksByPriority() // Execute each priority group sequentially, but tasks within group in parallel for _, tasks := range priorityGroups { for _, task := range tasks { if !task.Enabled { continue } wg.Add(1) go func(t ShutdownTask) { defer wg.Done() if err := sm.executeTask(ctx, t); err != nil { errors <- fmt.Errorf("task %s failed: %w", t.Name, err) } }(task) } wg.Wait() } close(errors) // Collect errors var criticalErr error var lastErr error for err := range errors { lastErr = err // Check if this was from a critical task for _, task := range sm.shutdownTasks { if task.Critical && fmt.Sprintf("task %s failed:", task.Name) == err.Error()[:len(fmt.Sprintf("task %s failed:", task.Name))] { criticalErr = err break } } } if criticalErr != nil { return criticalErr } return lastErr } func (sm *ShutdownManager) executeTask(ctx context.Context, task ShutdownTask) error { // Create timeout context for the task taskCtx, cancel := context.WithTimeout(ctx, task.Timeout) defer cancel() // Execute task with retry var lastErr error for attempt := 0; attempt <= sm.config.MaxRetries; attempt++ { if attempt > 0 { select { case <-time.After(sm.config.RetryDelay): case <-taskCtx.Done(): return taskCtx.Err() } } err := task.Task(taskCtx) if err == nil { return nil } lastErr = err // Call error handler if provided if task.OnError != nil { task.OnError(err) } } return fmt.Errorf("task failed after %d attempts: %w", sm.config.MaxRetries, lastErr) } func (sm *ShutdownManager) callHooks(ctx context.Context, hookMethod string) error { var lastErr error for _, hook := range sm.shutdownHooks { var err error switch hookMethod { case "OnShutdownStarted": err = hook.OnShutdownStarted(ctx) case "OnModulesStopped": err = hook.OnModulesStopped(ctx) case "OnCleanupStarted": err = hook.OnCleanupStarted(ctx) case "OnShutdownCompleted": err = hook.OnShutdownCompleted(ctx) case "OnShutdownFailed": err = hook.OnShutdownFailed(ctx, lastErr) } if err != nil { lastErr = err } } return lastErr } func (sm *ShutdownManager) sortTasksByPriority() { // Simple bubble sort by priority (descending) for i := 0; i < len(sm.shutdownTasks); i++ { for j := i + 1; j < len(sm.shutdownTasks); j++ { if sm.shutdownTasks[j].Priority > sm.shutdownTasks[i].Priority { sm.shutdownTasks[i], sm.shutdownTasks[j] = sm.shutdownTasks[j], sm.shutdownTasks[i] } } } } func (sm *ShutdownManager) groupTasksByPriority() [][]ShutdownTask { groups := make(map[int][]ShutdownTask) for _, task := range sm.shutdownTasks { groups[task.Priority] = append(groups[task.Priority], task) } // Convert to sorted slice var priorities []int for priority := range groups { priorities = append(priorities, priority) } // Sort priorities descending for i := 0; i < len(priorities); i++ { for j := i + 1; j < len(priorities); j++ { if priorities[j] > priorities[i] { priorities[i], priorities[j] = priorities[j], priorities[i] } } } var result [][]ShutdownTask for _, priority := range priorities { result = append(result, groups[priority]) } return result } // Default task implementations func (sm *ShutdownManager) saveApplicationState(ctx context.Context) error { // Save application state to disk // This would save things like current configuration, runtime state, etc. return nil } func (sm *ShutdownManager) closeExternalConnections(ctx context.Context) error { // Close database connections, external API connections, etc. return nil } func (sm *ShutdownManager) cleanupTempFiles(ctx context.Context) error { // Remove temporary files, logs, caches, etc. return nil } func (sm *ShutdownManager) notifyExternalSystems(ctx context.Context) error { // Notify external systems that this instance is shutting down return nil } // DefaultShutdownHook provides a basic implementation of ShutdownHook type DefaultShutdownHook struct { name string } func NewDefaultShutdownHook(name string) *DefaultShutdownHook { return &DefaultShutdownHook{name: name} } func (dsh *DefaultShutdownHook) OnShutdownStarted(ctx context.Context) error { return nil } func (dsh *DefaultShutdownHook) OnModulesStopped(ctx context.Context) error { return nil } func (dsh *DefaultShutdownHook) OnCleanupStarted(ctx context.Context) error { return nil } func (dsh *DefaultShutdownHook) OnShutdownCompleted(ctx context.Context) error { return nil } func (dsh *DefaultShutdownHook) OnShutdownFailed(ctx context.Context, err error) error { return nil }