Optimize MongoDB SCIM attribute queries for large groups using explicit batch sizes

src/Scim/SimpleIdServer.Scim.Persistence.MongoDB/Models/SCIMRepresentationModel.cs

@@ -6,6 +6,7 @@
 using SimpleIdServer.Scim.Persistence.MongoDB.Infrastructures;
 using System.Collections.Generic;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace SimpleIdServer.Scim.Persistence.MongoDB.Models
@@ -33,22 +34,32 @@ public SCIMRepresentationModel(SCIMRepresentation representation, string schemaC
         public ICollection<CustomMongoDBRef> SchemaRefs { get; set; }
         public ICollection<CustomMongoDBRef> AttributeRefs { get; set; }
 
-        public async Task IncludeAll(SCIMDbContext dbContext)
+        public async Task IncludeAll(SCIMDbContext dbContext, CancellationToken cancellationToken = default)
         {
             IncludeSchemas(dbContext.Database);
-            await IncludeAttributes(dbContext);
+            await IncludeAttributes(dbContext, cancellationToken);
         }
 
         public void IncludeSchemas(IMongoDatabase database)
         {
             Schemas = MongoDBEntity.GetReferences<SCIMSchema>(SchemaRefs, database);
         }
 
-        public async Task IncludeAttributes(SCIMDbContext dbContext)
+        public async Task IncludeAttributes(SCIMDbContext dbContext, CancellationToken cancellationToken = default)
         {
-            FlatAttributes = await dbContext.SCIMRepresentationAttributeLst.AsQueryable()
-                .Where(a => a.RepresentationId == Id)
-                .ToMongoListAsync();
+            // Optimize MongoDB query for large result sets by using Find with explicit options
+            // This provides better performance than LINQ for representations with many attributes (e.g., a group with 20k+ member attributes)
+            var filter = Builders<SCIMRepresentationAttribute>.Filter.Eq(a => a.RepresentationId, Id);
+            var findOptions = new FindOptions<SCIMRepresentationAttribute>
+            {
+                // Use configured batch size to reduce round trips to MongoDB
+                BatchSize = dbContext.Options.BatchSize
+            };
+
+            using var cursor = await dbContext.SCIMRepresentationAttributeLst
+                .FindAsync(filter, findOptions, cancellationToken);
+
+            FlatAttributes = await cursor.ToListAsync(cancellationToken);
         }
     }
 }

src/Scim/SimpleIdServer.Scim.Persistence.MongoDB/MongoDbOptions.cs

@@ -1,5 +1,6 @@
 // Copyright (c) SimpleIdServer. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See LICENSE in the project root for license information.
+using System;
 
 namespace SimpleIdServer.Scim.Persistence.MongoDB
 {
@@ -16,6 +17,7 @@ public MongoDbOptions()
 			CollectionProvisioningLst = "provisioningLst";
 			CollectionRealms = "realms";
             SupportTransaction = true;
+			BatchSize = 10000;
 		}
 
 		public string ConnectionString { get; set; }
@@ -28,5 +30,23 @@ public MongoDbOptions()
 		public string CollectionRealms { get; set; }
 
         public bool SupportTransaction { get; set; }
+
+		private int _batchSize;
+		/// <summary>
+		/// MongoDB cursor batch size for large result sets. 
+		/// Default is 10000 to optimize performance when loading groups with many members.
+		/// Higher values reduce network round trips but increase memory usage.
+		/// Minimum value is 1.
+		/// </summary>
+		public int BatchSize 
+		{ 
+			get => _batchSize;
+			set
+			{
+				if (value <= 0)
+					throw new ArgumentOutOfRangeException(nameof(BatchSize), "BatchSize must be greater than 0");
+				_batchSize = value;
+			}
+		}
 	}
 }

src/Scim/SimpleIdServer.Scim.Persistence.MongoDB/SCIMRepresentationQueryRepository.cs

@@ -97,9 +97,15 @@ join b in _scimDbContext.SCIMRepresentationAttributeLst.AsQueryable() on a.Paren
             var representationsList = await filteredRepresentations.ToListAsync(cancellationToken);
             var representationIds = representationsList.Select(r => r.Id).ToList();
 
-            var attributes = await _scimDbContext.SCIMRepresentationAttributeLst.AsQueryable()
-                .Where(a => representationIds.Contains(a.RepresentationId))
-                .ToListAsync(cancellationToken);
+            // Optimize MongoDB query for potentially large result sets (e.g., when querying groups with many members)
+            // Use Find API with explicit batch size instead of LINQ for better performance
+            var filter = Builders<SCIMRepresentationAttribute>.Filter.In(a => a.RepresentationId, representationIds);
+            var findOptions = new FindOptions<SCIMRepresentationAttribute>
+            {
+                BatchSize = _scimDbContext.Options.BatchSize
+            };
+            using var cursor = await _scimDbContext.SCIMRepresentationAttributeLst.FindAsync(filter, findOptions, cancellationToken);
+            var attributes = await cursor.ToListAsync(cancellationToken);
 
             var attributesByRepId = attributes.GroupBy(a => a.RepresentationId)
                 .ToDictionary(g => g.Key, g => g.ToList());
@@ -147,7 +153,7 @@ public async Task<SCIMRepresentation> FindSCIMRepresentationById(string realm, s
                 return null;
             }
 
-            await result.IncludeAll(_scimDbContext);
+            await result.IncludeAll(_scimDbContext, cancellationToken);
             return result;
         }