IdentityServer4中自定义TokenCleanupService的实现及令牌清理停止控制问题

嘿，我明白你的困扰了！默认的IdentityServer4令牌清理服务是按设定的间隔持续运行的，只要有过期令牌就会一直周期性执行删除操作，这确实可能给数据库和应用带来额外压力，导致性能变慢。咱们可以通过自定义令牌清理服务，添加停止条件来解决这个问题——比如清理完指定批次（比如1批100条）就停止，或者到某个时间点自动停止。

第一步：禁用默认的令牌清理服务 #

首先要把你原来配置里的EnableTokenCleanup设为false，因为我们要替换成自己实现的服务：

opts.EnableTokenCleanup = false; // 禁用默认服务

第二步：自定义托管式令牌清理服务（推荐方案） #

在ASP.NET Core中，使用IHostedService来实现定时任务是更符合框架最佳实践的方式，它能更好地管理服务生命周期。下面是一个带停止条件的自定义实现：

public class CustomTokenCleanupHostedService : IHostedService, IDisposable
{
    private readonly IPersistedGrantStore _persistedGrantStore;
    private readonly ILogger<CustomTokenCleanupHostedService> _logger;
    private readonly int _batchSize;
    private readonly int _maxCleanupBatches; // 最多执行的批次数量
    private readonly DateTime? _stopTime; // 停止清理的时间点（UTC）
    private Timer _timer;
    private int _executedBatches;
    private bool _disposed;

    public CustomTokenCleanupHostedService(
        IPersistedGrantStore persistedGrantStore,
        ILogger<CustomTokenCleanupHostedService> logger,
        IConfiguration configuration)
    {
        _persistedGrantStore = persistedGrantStore ?? throw new ArgumentNullException(nameof(persistedGrantStore));
        _logger = logger ?? throw new ArgumentNullException(nameof(logger));
        
        // 从配置读取参数，也可以直接硬编码
        _batchSize = configuration.GetValue<int>("TokenCleanup:BatchSize", 100);
        _maxCleanupBatches = configuration.GetValue<int>("TokenCleanup:MaxBatches", 1);
        
        var stopTimeStr = configuration.GetValue<string>("TokenCleanup:StopTime");
        if (!string.IsNullOrEmpty(stopTimeStr) && DateTime.TryParse(stopTimeStr, out var stopTime))
        {
            _stopTime = stopTime.ToUniversalTime(); // 统一用UTC时间避免时区问题
        }
        _executedBatches = 0;
    }

    public Task StartAsync(CancellationToken cancellationToken)
    {
        _logger.LogInformation("启动自定义令牌清理服务");
        // 立即执行第一次清理
        _ = ExecuteCleanupAsync(cancellationToken);
        // 调度下一次执行（如果还没达到停止条件）
        ScheduleNextCleanup();
        return Task.CompletedTask;
    }

    private async Task ExecuteCleanupAsync(CancellationToken cancellationToken)
    {
        // 检查是否达到停止条件
        if (_executedBatches >= _maxCleanupBatches 
            || (_stopTime.HasValue && DateTime.UtcNow >= _stopTime.Value))
        {
            _logger.LogInformation("令牌清理即将停止：已达到最大批次或指定停止时间");
            Stop();
            return;
        }

        try
        {
            _logger.LogInformation("开始执行第 {BatchNumber} 批令牌清理", _executedBatches + 1);
            var expiredGrants = await _persistedGrantStore.GetExpiredAsync(_batchSize, cancellationToken);
            
            // 如果没有找到过期令牌，直接停止
            if (!expiredGrants.Any())
            {
                _logger.LogInformation("未发现过期令牌，停止清理服务");
                Stop();
                return;
            }

            // 批量删除过期令牌
            var removeTasks = expiredGrants.Select(grant => _persistedGrantStore.RemoveAsync(grant.Key, cancellationToken));
            await Task.WhenAll(removeTasks);

            _executedBatches++;
            _logger.LogInformation("第 {BatchNumber} 批清理完成，共删除 {Count} 条记录", _executedBatches, expiredGrants.Count());
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "令牌清理过程中发生错误");
        }
    }

    private void ScheduleNextCleanup()
    {
        // 已达到停止条件就不再调度
        if (_executedBatches >= _maxCleanupBatches 
            || (_stopTime.HasValue && DateTime.UtcNow >= _stopTime.Value))
        {
            return;
        }

        // 设置下次执行间隔（和你原来的3600秒一致）
        var interval = TimeSpan.FromSeconds(3600);
        _timer = new Timer(async _ =>
        {
            await ExecuteCleanupAsync(CancellationToken.None);
            ScheduleNextCleanup();
        }, null, interval, Timeout.InfiniteTimeSpan);
    }

    public Task StopAsync(CancellationToken cancellationToken)
    {
        _logger.LogInformation("停止自定义令牌清理服务");
        Stop();
        return Task.CompletedTask;
    }

    private void Stop()
    {
        _timer?.Change(Timeout.Infinite, Timeout.Infinite);
    }

    public void Dispose()
    {
        Dispose(true);
        GC.SuppressFinalize(this);
    }

    protected virtual void Dispose(bool disposing)
    {
        if (_disposed) return;

        if (disposing)
        {
            _timer?.Dispose();
        }

        _disposed = true;
    }
}

第三步：注册自定义服务 #

在Startup.cs的ConfigureServices方法中注册这个托管服务：

services.AddHostedService<CustomTokenCleanupHostedService>();

如果需要通过配置文件控制参数，可以在appsettings.json中添加：

"TokenCleanup": {
  "BatchSize": 100,
  "MaxBatches": 1,
  "StopTime": "2024-12-31 23:59:59"
}

方案说明 #

这个自定义服务会自动检查停止条件：

1. 当执行的批次数量达到你设定的MaxBatches（比如1批）时，停止清理
2. 当当前时间达到你设定的StopTime时，停止清理
3. 如果某次清理没有找到任何过期令牌，也会自动停止，避免空跑浪费资源

这样就能完美解决你遇到的“持续删除导致应用变慢”的问题啦！

备注：内容来源于stack exchange，提问作者Gouthami Shivakoti