Use ExecuteDelete in token cleanup

Duende.IdentityServer.sln

@@ -25,7 +25,7 @@ Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Duende.IdentityServer.Entit
 EndProject
 Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Duende.IdentityServer.EntityFramework", "src\EntityFramework\Duende.IdentityServer.EntityFramework.csproj", "{376FD801-0E35-4145-9322-28FFB219E668}"
 EndProject
-Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "EntityFramework.Tests", "test\EntityFramework.Tests\EntityFramework.Tests.csproj", "{0772AE76-46E3-42A2-AA2F-B8EB56B4EB0D}"
+Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "EntityFramework.IntegrationTests", "test\EntityFramework.IntegrationTests\EntityFramework.IntegrationTests.csproj", "{0772AE76-46E3-42A2-AA2F-B8EB56B4EB0D}"
 EndProject
 Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "EntityFramework.Storage.UnitTests", "test\EntityFramework.Storage.UnitTests\EntityFramework.Storage.UnitTests.csproj", "{3A72EDFA-1E19-46E6-B983-ECF3EFBF192E}"
 EndProject

migrations/IdentityServerDb/Migrations/ConfigurationDb.sql

@@ -367,7 +367,7 @@ CREATE UNIQUE INDEX [IX_IdentityResources_Name] ON [IdentityResources] ([Name]);
 GO
 
 INSERT INTO [__EFMigrationsHistory] ([MigrationId], [ProductVersion])
-VALUES (N'20231110071401_Configuration', N'8.0.0-rc.2.23480.1');
+VALUES (N'20240104192404_Configuration', N'8.0.0');
 GO
 
 COMMIT;

migrations/IdentityServerDb/Migrations/ConfigurationDb/ConfigurationDbContextModelSnapshot.cs

@@ -17,7 +17,7 @@ protected override void BuildModel(ModelBuilder modelBuilder)
         {
 #pragma warning disable 612, 618
             modelBuilder
-                .HasAnnotation("ProductVersion", "8.0.0-rc.2.23480.1")
+                .HasAnnotation("ProductVersion", "8.0.0")
                 .HasAnnotation("Relational:MaxIdentifierLength", 128);
 
             SqlServerModelBuilderExtensions.UseIdentityColumns(modelBuilder);

migrations/IdentityServerDb/Migrations/PersistedGrantDb.sql

@@ -102,6 +102,9 @@ GO
 CREATE INDEX [IX_PersistedGrants_SubjectId_SessionId_Type] ON [PersistedGrants] ([SubjectId], [SessionId], [Type]);
 GO
 
+CREATE INDEX [IX_PushedAuthorizationRequests_ExpiresAtUtc] ON [PushedAuthorizationRequests] ([ExpiresAtUtc]);
+GO
+
 CREATE UNIQUE INDEX [IX_PushedAuthorizationRequests_ReferenceValueHash] ON [PushedAuthorizationRequests] ([ReferenceValueHash]);
 GO
 
@@ -121,7 +124,7 @@ CREATE INDEX [IX_ServerSideSessions_SubjectId] ON [ServerSideSessions] ([Subject
 GO
 
 INSERT INTO [__EFMigrationsHistory] ([MigrationId], [ProductVersion])
-VALUES (N'20231110071347_Grants', N'8.0.0-rc.2.23480.1');
+VALUES (N'20240104192353_Grants', N'8.0.0');
 GO
 
 COMMIT;

migrations/IdentityServerDb/Migrations/PersistedGrantDb/20231110071347_Grants.cs → migrations/IdentityServerDb/Migrations/PersistedGrantDb/20240104192353_Grants.cs

@@ -149,6 +149,11 @@ protected override void Up(MigrationBuilder migrationBuilder)
                 table: "PersistedGrants",
                 columns: new[] { "SubjectId", "SessionId", "Type" });
 
+            migrationBuilder.CreateIndex(
+                name: "IX_PushedAuthorizationRequests_ExpiresAtUtc",
+                table: "PushedAuthorizationRequests",
+                column: "ExpiresAtUtc");
+
             migrationBuilder.CreateIndex(
                 name: "IX_PushedAuthorizationRequests_ReferenceValueHash",
                 table: "PushedAuthorizationRequests",

migrations/IdentityServerDb/Migrations/PersistedGrantDb/PersistedGrantDbContextModelSnapshot.cs

@@ -17,7 +17,7 @@ protected override void BuildModel(ModelBuilder modelBuilder)
         {
 #pragma warning disable 612, 618
             modelBuilder
-                .HasAnnotation("ProductVersion", "8.0.0-rc.2.23480.1")
+                .HasAnnotation("ProductVersion", "8.0.0")
                 .HasAnnotation("Relational:MaxIdentifierLength", 128);
 
             SqlServerModelBuilderExtensions.UseIdentityColumns(modelBuilder);
@@ -195,6 +195,8 @@ protected override void BuildModel(ModelBuilder modelBuilder)
 
                     b.HasKey("Id");
 
+                    b.HasIndex("ExpiresAtUtc");
+
                     b.HasIndex("ReferenceValueHash")
                         .IsUnique();
 

src/EntityFramework.Storage/DbContexts/PersistedGrantDbContext.cs

@@ -82,7 +82,6 @@ protected override void OnModelCreating(ModelBuilder modelBuilder)
         }
 
         modelBuilder.ConfigurePersistedGrantContext(StoreOptions);
-        modelBuilder.ConfigurePushedAuthorizationRequestContext(StoreOptions);
 
         base.OnModelCreating(modelBuilder);
     }

src/EntityFramework.Storage/Extensions/ModelBuilderExtensions.cs

@@ -221,6 +221,19 @@ public static void ConfigurePersistedGrantContext(this ModelBuilder modelBuilder
             entity.HasIndex(x => x.SessionId);
             entity.HasIndex(x => x.DisplayName);
         });
+
+        modelBuilder.Entity<PushedAuthorizationRequest>(entity =>
+        {
+            entity.ToTable(storeOptions.PushedAuthorizationRequests);
+
+            entity.HasKey(x => x.Id);
+            entity.Property(x => x.ReferenceValueHash).HasMaxLength(64).IsRequired();
+            entity.Property(x => x.ExpiresAtUtc).IsRequired();
+            entity.Property(x => x.Parameters).IsRequired();
+
+            entity.HasIndex(x => x.ReferenceValueHash).IsUnique();
+            entity.HasIndex(x => x.ExpiresAtUtc);
+        });
     }
 
     /// <summary>
@@ -370,25 +383,4 @@ public static void ConfigureIdentityProviderContext(this ModelBuilder modelBuild
             entity.HasIndex(x => x.Scheme).IsUnique();
         });
     }
-
-    /// <summary>
-    /// Configures the pushed authorization requests.
-    /// </summary>
-    /// <param name="modelBuilder">The model builder.</param>
-    /// <param name="storeOptions">The store options.</param>
-    public static void ConfigurePushedAuthorizationRequestContext(this ModelBuilder modelBuilder, OperationalStoreOptions storeOptions)
-    {
-        if (!string.IsNullOrWhiteSpace(storeOptions.DefaultSchema)) modelBuilder.HasDefaultSchema(storeOptions.DefaultSchema);
-
-        modelBuilder.Entity<PushedAuthorizationRequest>(entity =>
-        {
-            entity.ToTable(storeOptions.PushedAuthorizationRequests).HasKey(x => x.Id);
-
-            entity.Property(x => x.ReferenceValueHash).HasMaxLength(64).IsRequired();
-            entity.Property(x => x.ExpiresAtUtc).IsRequired();
-            entity.Property(x => x.Parameters).IsRequired();
-
-            entity.HasIndex(x => x.ReferenceValueHash).IsUnique();
-        });
-    }
 }

src/EntityFramework.Storage/Options/OperationalStoreOptions.cs

@@ -109,6 +109,16 @@ public class OperationalStoreOptions
     /// </value>
     public int TokenCleanupInterval { get; set; } = 3600;
 
+    /// <summary>
+    /// If multiple nodes are running the token cleanup at the same time, there will be
+    /// concurrency issues in the database updates. To reduce the risk, the startup time
+    /// of the first run can be fuzzed (randomized). The default is <c>true</c>.
+    /// </summary>
+    /// <value>
+    /// <c>true</c> if startup time should be fuzzed, otherwise false.
+    /// </value>
+    public bool FuzzTokenCleanupStart { get; set; } = true;
+
     /// <summary>
     /// Gets or sets the number of records to remove at a time. Defaults to 100.
     /// </summary>

src/EntityFramework.Storage/TokenCleanup/TokenCleanupService.cs

@@ -4,8 +4,10 @@
 
 using System;
 using System.Linq;
+using System.Runtime.InteropServices;
 using System.Threading;
 using System.Threading.Tasks;
+using Duende.IdentityServer.EntityFramework.Entities;
 using Duende.IdentityServer.EntityFramework.Extensions;
 using Duende.IdentityServer.EntityFramework.Interfaces;
 using Duende.IdentityServer.EntityFramework.Options;
@@ -84,10 +86,16 @@ protected virtual async Task RemoveExpiredPersistedGrantsAsync(CancellationToken
 
         while (found >= _options.TokenCleanupBatchSize)
         {
-            var expiredGrants = await _persistedGrantDbContext.PersistedGrants
+            // Filter and order on expiration which is indexed, this allows the
+            // DB engine to just take the first N items from the index
+            var query = _persistedGrantDbContext.PersistedGrants
                 .Where(x => x.Expiration < DateTime.UtcNow)
-                .OrderBy(x => x.Expiration)
+                .OrderBy(x => x.Expiration);
+
+            // Get the batch to delete.
+            var expiredGrants = await query
                 .Take(_options.TokenCleanupBatchSize)
+                .AsNoTracking()
                 .ToArrayAsync(cancellationToken);
 
             found = expiredGrants.Length;
@@ -96,10 +104,38 @@ protected virtual async Task RemoveExpiredPersistedGrantsAsync(CancellationToken
             {
                 _logger.LogInformation("Removing {grantCount} expired grants", found);
 
-                _persistedGrantDbContext.PersistedGrants.RemoveRange(expiredGrants);
-
-                var list = await _persistedGrantDbContext.SaveChangesWithConcurrencyCheckAsync<Entities.PersistedGrant>(_logger, cancellationToken);
-                expiredGrants = expiredGrants.Except(list).ToArray();
+                var foundIds = expiredGrants.Select(pg => pg.Id).ToArray();
+
+                // Using two where clauses should be more DB engine friendly as the
+                // first clause can be resolved using the expiration index.
+                var deleteCount = await query
+                    // Run the same query, but now use an interval instead of Take(). This is to
+                    // ensure we get all the elements, even if a new element was added in the middle
+                    // of the set.
+                    .Where(pg => 
+                        pg.Expiration >= expiredGrants.First().Expiration 
+                        && pg.Expiration <= expiredGrants.Last().Expiration)
+                    // To be on the safe side, filter out any possibly newly added item within the interval
+                    .Where(pg => foundIds.Contains(pg.Id))
+                    // And delete them.
+                    .ExecuteDeleteAsync(cancellationToken);
+
+                if (deleteCount != found)
+                {
+                    if (_operationalStoreNotification != null)
+                    {
+                        _logger.LogWarning("Tried to remove {grantCount} expired grants, but only {deleteCount} " +
+                            "was deleted. This indicates that another process has already removed the items. Duplicate " +
+                            "notifications may be sent to the registered IOperationalStoreNotification.", 
+                            found, deleteCount);
+                    }
+                    else
+                    {
+                        _logger.LogDebug("Tried to remove {grantCount} expired grants, but only {deleteCount} " +
+                            "was deleted. This indicates that another process has already removed the items.",
+                            found, deleteCount);
+                    }
+                }
 
                 if (_operationalStoreNotification != null)
                 {
@@ -122,26 +158,50 @@ protected virtual async Task RemoveConsumedPersistedGrantsAsync(CancellationToke
 
         while (found >= _options.TokenCleanupBatchSize)
         {
-            var expiredGrants = await _persistedGrantDbContext.PersistedGrants
+            var query = _persistedGrantDbContext.PersistedGrants
                 .Where(x => x.ConsumedTime < consumedTimeThreshold)
-                .OrderBy(x => x.ConsumedTime)
+                .OrderBy(pg => pg.ConsumedTime);
+
+            var consumedGrants = await query
                 .Take(_options.TokenCleanupBatchSize)
+                .AsNoTracking()
                 .ToArrayAsync(cancellationToken);
 
-            found = expiredGrants.Length;
+            found = consumedGrants.Length;
 
             if (found > 0)
             {
                 _logger.LogInformation("Removing {grantCount} consumed grants", found);
 
-                _persistedGrantDbContext.PersistedGrants.RemoveRange(expiredGrants);
+                var foundIds = consumedGrants.Select(pg => pg.Id).ToArray();
 
-                var list = await _persistedGrantDbContext.SaveChangesWithConcurrencyCheckAsync<Entities.PersistedGrant>(_logger, cancellationToken);
-                expiredGrants = expiredGrants.Except(list).ToArray();
+                var deleteCount = await query
+                    .Where(pg => 
+                        pg.ConsumedTime >= consumedGrants.First().ConsumedTime 
+                        && pg.ConsumedTime <= consumedGrants.Last().ConsumedTime)
+                    .Where(pg => foundIds.Contains(pg.Id))
+                    .ExecuteDeleteAsync(cancellationToken);
+
+                if (deleteCount != found)
+                {
+                    if (_operationalStoreNotification != null)
+                    {
+                        _logger.LogWarning("Tried to remove {grantCount} consumed grants, but only {deleteCount} " +
+                            "was deleted. This indicates that another process has already removed the items. Duplicate " +
+                            "notifications may be sent to the registered IOperationalStoreNotification.",
+                            found, deleteCount);
+                    }
+                    else
+                    {
+                        _logger.LogDebug("Tried to remove {grantCount} consumed grants, but only {deleteCount} " +
+                            "was deleted. This indicates that another process has already removed the items.",
+                            found, deleteCount);
+                    }
+                }
 
                 if (_operationalStoreNotification != null)
                 {
-                    await _operationalStoreNotification.PersistedGrantsRemovedAsync(expiredGrants);
+                    await _operationalStoreNotification.PersistedGrantsRemovedAsync(consumedGrants);
                 }
             }
         }
@@ -158,10 +218,13 @@ protected virtual async Task RemoveDeviceCodesAsync(CancellationToken cancellati
 
         while (found >= _options.TokenCleanupBatchSize)
         {
-            var expiredCodes = await _persistedGrantDbContext.DeviceFlowCodes
+            var query = _persistedGrantDbContext.DeviceFlowCodes
                 .Where(x => x.Expiration < DateTime.UtcNow)
-                .OrderBy(x => x.DeviceCode)
+                .OrderBy(x => x.Expiration);
+
+            var expiredCodes = await query
                 .Take(_options.TokenCleanupBatchSize)
+                .AsNoTracking()
                 .ToArrayAsync(cancellationToken);
 
             found = expiredCodes.Length;
@@ -170,10 +233,29 @@ protected virtual async Task RemoveDeviceCodesAsync(CancellationToken cancellati
             {
                 _logger.LogInformation("Removing {deviceCodeCount} device flow codes", found);
 
-                _persistedGrantDbContext.DeviceFlowCodes.RemoveRange(expiredCodes);
+                var foundCodes = expiredCodes.Select(c => c.DeviceCode).ToArray();
 
-                var list = await _persistedGrantDbContext.SaveChangesWithConcurrencyCheckAsync<Entities.DeviceFlowCodes>(_logger, cancellationToken);
-                expiredCodes = expiredCodes.Except(list).ToArray();
+                var deleteCount = await query
+                    .Where(c => c.Expiration >= expiredCodes.First().Expiration && c.Expiration <= expiredCodes.Last().Expiration)
+                    .Where(c => foundCodes.Contains(c.DeviceCode))
+                    .ExecuteDeleteAsync();
+
+                if (deleteCount != found)
+                {
+                    if (_operationalStoreNotification != null)
+                    {
+                        _logger.LogWarning("Tried to remove {grantCount} expired device codes, but only {deleteCount} " +
+                            "was deleted. This indicates that another process has already removed the items. Duplicate " +
+                            "notifications may be sent to the registered IOperationalStoreNotification.",
+                            found, deleteCount);
+                    }
+                    else
+                    {
+                        _logger.LogDebug("Tried to remove {grantCount} expired device codes, but only {deleteCount} " +
+                            "was deleted. This indicates that another process has already removed the items.",
+                            found, deleteCount);
+                    }
+                }
 
                 if (_operationalStoreNotification != null)
                 {
@@ -188,9 +270,22 @@ protected virtual async Task RemoveDeviceCodesAsync(CancellationToken cancellati
     /// </summary>
     protected virtual async Task RemovePushedAuthorizationRequestsAsync(CancellationToken cancellationToken = default)
     {
-        var x = await _persistedGrantDbContext.PushedAuthorizationRequests
-            .Where(p => p.ExpiresAtUtc < DateTime.UtcNow)
-            .ExecuteDeleteAsync(cancellationToken);
-        _logger.LogInformation("Removed {parCount} stale pushed authorization requests", x);
+        var found = Int32.MaxValue;
+
+        var query = _persistedGrantDbContext.PushedAuthorizationRequests
+            .Where(par => par.ExpiresAtUtc < DateTime.UtcNow)
+            .OrderBy(par => par.ExpiresAtUtc);
+
+        while (found >= _options.TokenCleanupBatchSize)
+        {
+            found = await query
+                .Take(_options.TokenCleanupBatchSize)
+                .ExecuteDeleteAsync(cancellationToken);
+
+            if (found > 0)
+            {
+                _logger.LogInformation("Removed {parCount} stale pushed authorization requests", found);
+            }
+        }
     }
 }