ApplicationShutdownCleanupService.cs 7.45 KB
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;
using Rcs.Domain.Entities;
using Rcs.Infrastructure.DB.MsSql;
using StackExchange.Redis;

namespace Rcs.Infrastructure.Services
{
    /// <summary>
    /// Performs shutdown cleanup work when the host is stopping:
    /// 1) clear all Redis keys
    /// 2) mark all robots as offline in database
    /// </summary>
    public class ApplicationShutdownCleanupService : IHostedService
    {
        private const int RedisDeleteBatchSize = 500;

        private readonly ILogger<ApplicationShutdownCleanupService> _logger;
        private readonly IHostApplicationLifetime _hostApplicationLifetime;
        private readonly IConnectionMultiplexer _redis;
        private readonly IServiceScopeFactory _scopeFactory;
        private readonly object _sync = new();

        private Task? _cleanupTask;

        public ApplicationShutdownCleanupService(
            ILogger<ApplicationShutdownCleanupService> logger,
            IHostApplicationLifetime hostApplicationLifetime,
            IConnectionMultiplexer redis,
            IServiceScopeFactory scopeFactory)
        {
            _logger = logger;
            _hostApplicationLifetime = hostApplicationLifetime;
            _redis = redis;
            _scopeFactory = scopeFactory;
        }

        public Task StartAsync(CancellationToken cancellationToken)
        {
            _hostApplicationLifetime.ApplicationStopping.Register(() =>
            {
                _ = EnsureCleanupStarted(CancellationToken.None);
            });

            return Task.CompletedTask;
        }

        public async Task StopAsync(CancellationToken cancellationToken)
        {
            var cleanupTask = EnsureCleanupStarted(CancellationToken.None);
            await cleanupTask;
        }

        private Task EnsureCleanupStarted(CancellationToken cancellationToken)
        {
            lock (_sync)
            {
                _cleanupTask ??= ExecuteCleanupAsync(cancellationToken);
                return _cleanupTask;
            }
        }

        private async Task ExecuteCleanupAsync(CancellationToken cancellationToken)
        {
            _logger.LogInformation("[ShutdownCleanup] Shutdown detected, starting cleanup workflow.");

            try
            {
                await ClearRedisAsync(cancellationToken);
                _logger.LogInformation("[ShutdownCleanup] Redis cleanup completed.");
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "[ShutdownCleanup] Redis cleanup failed.");
            }

            try
            {
                var affectedRows = await MarkRobotsOfflineAsync(cancellationToken);
                _logger.LogInformation("[ShutdownCleanup] Robot offline update completed. Affected rows: {Count}", affectedRows);
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "[ShutdownCleanup] Robot offline update failed.");
            }

            _logger.LogInformation("[ShutdownCleanup] Cleanup workflow finished.");
        }

        private async Task ClearRedisAsync(CancellationToken cancellationToken)
        {
            var endPoints = _redis.GetEndPoints();
            if (endPoints.Length == 0)
            {
                _logger.LogWarning("[ShutdownCleanup] No Redis endpoints found. Skip Redis cleanup.");
                return;
            }

            foreach (var endPoint in endPoints)
            {
                cancellationToken.ThrowIfCancellationRequested();

                IServer server;
                try
                {
                    server = _redis.GetServer(endPoint);
                }
                catch (Exception ex)
                {
                    _logger.LogWarning(ex, "[ShutdownCleanup] Failed to get Redis server for endpoint {Endpoint}.", endPoint);
                    continue;
                }

                if (!server.IsConnected || server.IsReplica)
                {
                    continue;
                }

                var flushed = await TryFlushAllDatabasesAsync(server);
                if (flushed)
                {
                    continue;
                }

                var deletedCount = await DeleteAllKeysByScanAsync(server, cancellationToken);
                _logger.LogInformation(
                    "[ShutdownCleanup] Redis scan-delete fallback completed for {Endpoint}. Deleted keys: {Count}",
                    endPoint,
                    deletedCount);
            }
        }

        private async Task<bool> TryFlushAllDatabasesAsync(IServer server)
        {
            try
            {
                await server.FlushAllDatabasesAsync();
                _logger.LogInformation("[ShutdownCleanup] Redis flush-all succeeded on server {Endpoint}.", server.EndPoint);
                return true;
            }
            catch (Exception ex)
            {
                _logger.LogWarning(ex,
                    "[ShutdownCleanup] Redis flush-all failed on server {Endpoint}, falling back to key scan delete.",
                    server.EndPoint);
                return false;
            }
        }

        private async Task<long> DeleteAllKeysByScanAsync(IServer server, CancellationToken cancellationToken)
        {
            var totalDeleted = 0L;
            var databaseCount = server.DatabaseCount > 0 ? server.DatabaseCount : 1;

            for (var dbIndex = 0; dbIndex < databaseCount; dbIndex++)
            {
                cancellationToken.ThrowIfCancellationRequested();

                try
                {
                    var db = _redis.GetDatabase(dbIndex);
                    var batch = new List<RedisKey>(RedisDeleteBatchSize);

                    foreach (var key in server.Keys(database: dbIndex, pageSize: RedisDeleteBatchSize))
                    {
                        batch.Add(key);
                        if (batch.Count < RedisDeleteBatchSize)
                        {
                            continue;
                        }

                        totalDeleted += await db.KeyDeleteAsync(batch.ToArray());
                        batch.Clear();
                    }

                    if (batch.Count > 0)
                    {
                        totalDeleted += await db.KeyDeleteAsync(batch.ToArray());
                    }
                }
                catch (Exception ex)
                {
                    _logger.LogWarning(
                        ex,
                        "[ShutdownCleanup] Failed to delete keys by scan on Redis endpoint {Endpoint}, db index {DbIndex}.",
                        server.EndPoint,
                        dbIndex);
                }
            }

            return totalDeleted;
        }

        private async Task<int> MarkRobotsOfflineAsync(CancellationToken cancellationToken)
        {
            using var scope = _scopeFactory.CreateScope();
            var dbContext = scope.ServiceProvider.GetRequiredService<AppDbContext>();

            return await dbContext.Robots
                .Where(r => r.Online != OnlineStatus.Offline || r.Driving)
                .ExecuteUpdateAsync(
                    setters => setters
                        .SetProperty(r => r.Online, OnlineStatus.Offline)
                        .SetProperty(r => r.Driving, false)
                        .SetProperty(r => r.UpdatedAt, _ => DateTime.Now),
                    cancellationToken);
        }
    }
}