Graceful Shutdown

/**
 * Interface for beans that support graceful shutdown
 * Enables coordinated shutdown with resource cleanup
 */
public interface GracefulShutdownCapable {
    
    /**
     * Perform graceful shutdown of this component
     * Should complete ongoing operations and release resources
     * @return CompletionStage that completes when shutdown is finished
     */
    CompletionStage<?> shutdownGracefully();
    
    /**
     * Report the number of currently active tasks/operations
     * Used for monitoring shutdown progress
     * @return Optional number of active tasks, empty if not trackable
     */
    OptionalLong reportActiveTasks();
    
    /**
     * Combine multiple completion stages into a single stage
     * Utility method for coordinating multiple async operations
     * @param stages Stream of completion stages to combine
     * @return CompletionStage that completes when all input stages complete
     */
    @NonNull
    static CompletionStage<?> allOf(@NonNull Stream<CompletionStage<?>> stages);
    
    /**
     * Shutdown all provided GracefulShutdownCapable instances
     * Coordinates shutdown across multiple components with exception handling
     * @param stages Stream of components to shutdown
     * @return CompletionStage that completes when all components are shutdown
     */
    @NonNull
    static CompletionStage<?> shutdownAll(@NonNull Stream<? extends GracefulShutdownCapable> stages);
    
    /**
     * Combine active task counts from multiple GracefulShutdownCapable components
     * Aggregates active task reports across multiple components
     * @param delegates Iterable of components to query for active tasks
     * @return Combined active task count, or empty if no tasks are trackable
     */
    @NonNull
    static OptionalLong combineActiveTasks(@NonNull Iterable<? extends GracefulShutdownCapable> delegates);
}

import io.micronaut.runtime.graceful.GracefulShutdownCapable;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.CompletableFuture;

// Database connection pool with graceful shutdown
@Singleton
public class DatabaseConnectionPool implements GracefulShutdownCapable {
    
    private final HikariDataSource dataSource;
    private final AtomicLong activeConnections = new AtomicLong(0);
    
    public DatabaseConnectionPool(HikariDataSource dataSource) {
        this.dataSource = dataSource;
    }
    
    @Override
    public CompletionStage<?> shutdownGracefully() {
        logger.info("Starting graceful shutdown of database connection pool");
        
        return CompletableFuture.runAsync(() -> {
            try {
                // Wait for active connections to complete
                long maxWaitSeconds = 30;
                long waitStart = System.currentTimeMillis();
                
                while (activeConnections.get() > 0 && 
                       (System.currentTimeMillis() - waitStart) < maxWaitSeconds * 1000) {
                    Thread.sleep(100);
                }
                
                // Close the data source
                dataSource.close();
                logger.info("Database connection pool shutdown completed");
                
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                logger.warn("Database shutdown interrupted", e);
            }
        });
    }
    
    @Override
    public OptionalLong reportActiveTasks() {
        return OptionalLong.of(activeConnections.get());
    }
    
    public Connection getConnection() throws SQLException {
        activeConnections.incrementAndGet();
        try {
            return new ConnectionWrapper(dataSource.getConnection()) {
                @Override
                public void close() throws SQLException {
                    super.close();
                    activeConnections.decrementAndGet();
                }
            };
        } catch (SQLException e) {
            activeConnections.decrementAndGet();
            throw e;
        }
    }
}

// Message processing service with graceful shutdown
@Singleton
public class MessageProcessor implements GracefulShutdownCapable {
    
    private final ExecutorService processingExecutor;
    private final AtomicLong activeMessages = new AtomicLong(0);
    private volatile boolean shutdownRequested = false;
    
    public MessageProcessor() {
        this.processingExecutor = Executors.newFixedThreadPool(10);
    }
    
    @Override
    public CompletionStage<?> shutdownGracefully() {
        logger.info("Starting graceful shutdown of message processor");
        shutdownRequested = true;
        
        return CompletableFuture.runAsync(() -> {
            // Stop accepting new messages
            processingExecutor.shutdown();
            
            try {
                // Wait for active messages to complete (max 60 seconds)
                boolean terminated = processingExecutor.awaitTermination(60, TimeUnit.SECONDS);
                
                if (!terminated) {
                    logger.warn("Message processor shutdown timed out, forcing shutdown");
                    processingExecutor.shutdownNow();
                }
                
                logger.info("Message processor shutdown completed. Active messages: {}", 
                           activeMessages.get());
                           
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                processingExecutor.shutdownNow();
            }
        });
    }
    
    @Override
    public OptionalLong reportActiveTasks() {
        return OptionalLong.of(activeMessages.get());
    }
    
    public CompletableFuture<Void> processMessage(Message message) {
        if (shutdownRequested) {
            return CompletableFuture.failedFuture(
                new IllegalStateException("Shutdown in progress"));
        }
        
        activeMessages.incrementAndGet();
        
        return CompletableFuture.runAsync(() -> {
            try {
                // Process the message
                handleMessage(message);
            } finally {
                activeMessages.decrementAndGet();
            }
        }, processingExecutor);
    }
}

/**
 * Singleton service that manages graceful shutdown across all components
 * Automatically discovers and coordinates GracefulShutdownCapable beans
 */
@Singleton
@Requires(classes = GracefulShutdownCapable.class)
@Experimental
public final class GracefulShutdownManager {
    
    /**
     * Initiate graceful shutdown of all GracefulShutdownCapable beans
     * Automatically discovers all beans implementing GracefulShutdownCapable
     * @return CompletionStage that completes when all components are shutdown
     */
    @NonNull
    public CompletionStage<?> shutdownGracefully();
    
    /**
     * Report combined active tasks from all GracefulShutdownCapable beans
     * Aggregates active task counts from all discovered components
     * @return Optional total count of active tasks across all components
     */
    @NonNull
    public OptionalLong reportActiveTasks();
}

import io.micronaut.runtime.graceful.GracefulShutdownManager;

@Singleton
public class ApplicationShutdownCoordinator {
    
    private final GracefulShutdownManager shutdownManager;
    
    public ApplicationShutdownCoordinator(GracefulShutdownManager shutdownManager) {
        this.shutdownManager = shutdownManager;
    }
    
    @EventListener
    public void onShutdownEvent(ShutdownEvent event) {
        logger.info("Application shutdown initiated");
        
        // Get current active task count
        OptionalLong activeTasks = shutdownManager.reportActiveTasks();
        if (activeTasks.isPresent()) {
            logger.info("Active tasks at shutdown: {}", activeTasks.getAsLong());
        }
        
        // Initiate graceful shutdown
        CompletionStage<?> shutdownStage = shutdownManager.shutdownGracefully();
        
        // Wait for completion (with timeout)
        try {
            shutdownStage.toCompletableFuture().get(90, TimeUnit.SECONDS);
            logger.info("✅ Graceful shutdown completed successfully");
        } catch (TimeoutException e) {
            logger.warn("⚠️ Graceful shutdown timed out after 90 seconds");
        } catch (Exception e) {
            logger.error("❌ Error during graceful shutdown", e);
        }
    }
    
    // Manual shutdown trigger (e.g., for admin endpoint)
    public CompletableFuture<ShutdownReport> initiateShutdown() {
        return CompletableFuture.supplyAsync(() -> {
            long startTime = System.currentTimeMillis();
            OptionalLong initialActiveTasks = shutdownManager.reportActiveTasks();
            
            try {
                shutdownManager.shutdownGracefully()
                              .toCompletableFuture()
                              .get(60, TimeUnit.SECONDS);
                              
                long duration = System.currentTimeMillis() - startTime;
                return new ShutdownReport(true, duration, initialActiveTasks);
                
            } catch (Exception e) {
                long duration = System.currentTimeMillis() - startTime;
                return new ShutdownReport(false, duration, initialActiveTasks, e);
            }
        });
    }
}

// Shutdown report data class
public class ShutdownReport {
    private final boolean successful;
    private final long durationMs;
    private final OptionalLong initialActiveTasks;
    private final Exception error;
    
    // Constructors and getters...
}

/**
 * Listener interface for graceful shutdown events
 * Allows components to respond to shutdown phases
 */
public interface GracefulShutdownListener {
    
    /**
     * Called when graceful shutdown begins
     * Opportunity to prepare for shutdown
     */
    default void onShutdownStart() {
        // Default implementation does nothing
    }
    
    /**
     * Called during shutdown progress
     * @param remainingTasks Number of tasks still active
     */
    default void onShutdownProgress(long remainingTasks) {
        // Default implementation does nothing
    }
    
    /**
     * Called when graceful shutdown completes successfully
     */
    default void onShutdownComplete() {
        // Default implementation does nothing
    }
    
    /**
     * Called if graceful shutdown fails or times out
     * @param cause The exception that caused the failure, if any
     */
    default void onShutdownFailed(Throwable cause) {
        // Default implementation does nothing
    }
}

import io.micronaut.runtime.graceful.GracefulShutdownListener;

@Singleton
public class ShutdownProgressLogger implements GracefulShutdownListener {
    
    private long shutdownStartTime;
    
    @Override
    public void onShutdownStart() {
        shutdownStartTime = System.currentTimeMillis();
        logger.info("🔄 Graceful shutdown initiated");
    }
    
    @Override
    public void onShutdownProgress(long remainingTasks) {
        long elapsed = System.currentTimeMillis() - shutdownStartTime;
        logger.info("⏳ Shutdown progress: {} tasks remaining after {}ms", 
                   remainingTasks, elapsed);
    }
    
    @Override
    public void onShutdownComplete() {
        long totalTime = System.currentTimeMillis() - shutdownStartTime;
        logger.info("✅ Graceful shutdown completed in {}ms", totalTime);
    }
    
    @Override
    public void onShutdownFailed(Throwable cause) {
        long totalTime = System.currentTimeMillis() - shutdownStartTime;
        logger.error("❌ Graceful shutdown failed after {}ms", totalTime, cause);
    }
}

// Metrics collection during shutdown
@Singleton
public class ShutdownMetricsCollector implements GracefulShutdownListener {
    
    private final MeterRegistry meterRegistry;
    private Timer.Sample shutdownTimer;
    
    public ShutdownMetricsCollector(MeterRegistry meterRegistry) {
        this.meterRegistry = meterRegistry;
    }
    
    @Override
    public void onShutdownStart() {
        shutdownTimer = Timer.start(meterRegistry);
        meterRegistry.counter("shutdown.started").increment();
    }
    
    @Override
    public void onShutdownProgress(long remainingTasks) {
        meterRegistry.gauge("shutdown.remaining.tasks", remainingTasks);
    }
    
    @Override
    public void onShutdownComplete() {
        if (shutdownTimer != null) {
            shutdownTimer.stop(Timer.builder("shutdown.duration")
                                   .description("Time taken for graceful shutdown")
                                   .register(meterRegistry));
        }
        meterRegistry.counter("shutdown.completed").increment();
    }
    
    @Override
    public void onShutdownFailed(Throwable cause) {
        if (shutdownTimer != null) {
            shutdownTimer.stop(Timer.builder("shutdown.duration.failed")
                                   .description("Time taken for failed shutdown")
                                   .register(meterRegistry));
        }
        meterRegistry.counter("shutdown.failed").increment();
    }
}

/**
 * Configuration for graceful shutdown behavior
 * Allows customization of shutdown timeouts and policies
 */
@ConfigurationProperties("micronaut.application.graceful-shutdown")
public class GracefulShutdownConfiguration {
    
    /**
     * Enable or disable graceful shutdown
     * @return true if graceful shutdown is enabled
     */
    public boolean isEnabled();
    
    /**
     * Maximum time to wait for graceful shutdown to complete
     * @return Shutdown timeout duration
     */
    public Duration getTimeout();
    
    /**
     * Time to wait between shutdown progress checks
     * @return Progress check interval
     */
    public Duration getProgressCheckInterval();
    
    /**
     * Whether to force shutdown if graceful shutdown times out
     * @return true to force shutdown on timeout
     */
    public boolean isForceShutdownOnTimeout();
    
    /**
     * Log level for shutdown progress messages
     * @return Log level for shutdown logging
     */
    public LogLevel getLogLevel();
}

# application.yml
micronaut:
  application:
    graceful-shutdown:
      enabled: true
      timeout: 60s
      progress-check-interval: 5s
      force-shutdown-on-timeout: true
      log-level: INFO

// Using configuration in custom shutdown logic
@Singleton
public class ConfigurableShutdownManager {
    
    private final GracefulShutdownConfiguration config;
    private final GracefulShutdownManager shutdownManager;
    
    public ConfigurableShutdownManager(GracefulShutdownConfiguration config,
                                     GracefulShutdownManager shutdownManager) {
        this.config = config;
        this.shutdownManager = shutdownManager;
    }
    
    public CompletableFuture<Void> shutdown() {
        if (!config.isEnabled()) {
            logger.info("Graceful shutdown disabled, performing immediate shutdown");
            return CompletableFuture.completedFuture(null);
        }
        
        return CompletableFuture.runAsync(() -> {
            try {
                Duration timeout = config.getTimeout();
                logger.log(config.getLogLevel(), "Starting graceful shutdown with {}s timeout", 
                          timeout.getSeconds());
                
                CompletableFuture<?> shutdownFuture = shutdownManager
                    .shutdownGracefully()
                    .toCompletableFuture();
                
                // Wait with configured timeout
                shutdownFuture.get(timeout.toMillis(), TimeUnit.MILLISECONDS);
                
                logger.log(config.getLogLevel(), "Graceful shutdown completed successfully");
                
            } catch (TimeoutException e) {
                if (config.isForceShutdownOnTimeout()) {
                    logger.warn("Graceful shutdown timed out, forcing shutdown");
                    // Force shutdown logic here
                } else {
                    logger.error("Graceful shutdown timed out and force shutdown disabled");
                    throw new RuntimeException("Shutdown timeout", e);
                }
            } catch (Exception e) {
                logger.error("Error during graceful shutdown", e);
                throw new RuntimeException("Shutdown error", e);
            }
        });
    }
}

@Singleton
public class ConnectionManager implements GracefulShutdownCapable {
    
    private final Map<String, ConnectionPool> connectionPools = new ConcurrentHashMap<>();
    private volatile boolean shutdownInitiated = false;
    
    @Override
    public CompletionStage<?> shutdownGracefully() {
        shutdownInitiated = true;
        logger.info("Shutting down {} connection pools", connectionPools.size());
        
        List<CompletableFuture<Void>> shutdownFutures = connectionPools.values()
            .stream()
            .map(pool -> CompletableFuture.runAsync(() -> {
                try {
                    pool.close();
                } catch (Exception e) {
                    logger.error("Error closing connection pool", e);
                }
            }))
            .collect(Collectors.toList());
        
        return CompletableFuture.allOf(shutdownFutures.toArray(new CompletableFuture[0]));
    }
    
    @Override
    public OptionalLong reportActiveTasks() {
        return OptionalLong.of(
            connectionPools.values().stream()
                          .mapToLong(ConnectionPool::getActiveConnectionCount)
                          .sum()
        );
    }
}

@Singleton
public class BackgroundTaskProcessor implements GracefulShutdownCapable {
    
    private final ScheduledExecutorService scheduler;
    private final Map<String, CompletableFuture<Void>> activeTasks = new ConcurrentHashMap<>();
    
    public BackgroundTaskProcessor() {
        this.scheduler = Executors.newScheduledThreadPool(5);
    }
    
    @Override
    public CompletionStage<?> shutdownGracefully() {
        logger.info("Shutting down background task processor");
        
        // Stop accepting new tasks
        scheduler.shutdown();
        
        return CompletableFuture.runAsync(() -> {
            try {
                // Wait for scheduled tasks to complete
                boolean terminated = scheduler.awaitTermination(30, TimeUnit.SECONDS);
                
                if (!terminated) {
                    logger.warn("Scheduled tasks did not complete within timeout");
                    scheduler.shutdownNow();
                }
                
                // Wait for active background tasks
                CompletableFuture<Void> allTasks = CompletableFuture.allOf(
                    activeTasks.values().toArray(new CompletableFuture[0])
                );
                
                allTasks.get(60, TimeUnit.SECONDS);
                logger.info("All background tasks completed");
                
            } catch (Exception e) {
                logger.error("Error during background task shutdown", e);
                // Cancel remaining tasks
                activeTasks.values().forEach(task -> task.cancel(true));
            }
        });
    }
    
    @Override
    public OptionalLong reportActiveTasks() {
        return OptionalLong.of(activeTasks.size());
    }
}