Optimize query helper memory

1 files changed, 135 insertions, 108 deletions

diff --git a/src/Jellyfin.Database/Jellyfin.Database.Implementations/DescendantQueryHelper.cs b/src/Jellyfin.Database/Jellyfin.Database.Implementations/DescendantQueryHelper.cs
index bfd0fac34a..9a42c86f7d 100644
--- a/src/Jellyfin.Database/Jellyfin.Database.Implementations/DescendantQueryHelper.cs
+++ b/src/Jellyfin.Database/Jellyfin.Database.Implementations/DescendantQueryHelper.cs
@@ -8,7 +8,7 @@ using Jellyfin.Database.Implementations.MatchCriteria;
 namespace Jellyfin.Database.Implementations;
 
 /// <summary>
-/// Provides methods for querying item hierarchies using iterative traversal.
+/// Provides methods for querying item hierarchies.
 /// Uses AncestorIds and LinkedChildren tables for parent-child traversal.
 /// </summary>
 public static class DescendantQueryHelper
@@ -32,11 +32,18 @@ public static class DescendantQueryHelper
     {
         ArgumentNullException.ThrowIfNull(context);
 
-        var descendants = TraverseHierarchyDown(context, [parentId]);
+        var (closureRoots, linkRoots) = ResolveLinkedRoots(context, parentId);
 
-        descendants.Remove(parentId);
+        var hierarchyDescendants = ClosureDescendants(context, closureRoots);
 
-        return descendants.AsQueryable();
+        var linkedDescendants = context.LinkedChildren
+            .WhereOneOrMany(linkRoots, e => e.ParentId)
+            .Select(e => e.ChildId);
+
+        return hierarchyDescendants
+            .Concat(linkedDescendants)
+            .Where(e => !e.Equals(parentId))
+            .Distinct();
     }
 
     /// <summary>
@@ -51,11 +58,9 @@ public static class DescendantQueryHelper
     {
         ArgumentNullException.ThrowIfNull(context);
 
-        var descendants = TraverseHierarchyDownOwned(context, [parentId]);
-
-        descendants.Remove(parentId);
-
-        return descendants.AsQueryable();
+        return ClosureDescendants(context, [parentId])
+            .Where(e => !e.Equals(parentId))
+            .Distinct();
     }
 
     /// <summary>
@@ -76,11 +81,12 @@ public static class DescendantQueryHelper
             return [];
         }
 
-        var seedSet = new HashSet<Guid>(parentIds);
-        var descendants = TraverseHierarchyDownOwned(context, seedSet);
+        var descendants = ClosureDescendants(context, parentIds)
+            .Distinct()
+            .ToHashSet();
 
-        // Remove the seed IDs — callers want only descendants
-        descendants.ExceptWith(seedSet);
+        // The callers want only descendants, and an item is never its own descendant.
+        descendants.ExceptWith(parentIds);
 
         return descendants;
     }
@@ -96,28 +102,48 @@ public static class DescendantQueryHelper
     {
         ArgumentNullException.ThrowIfNull(context);
         ArgumentNullException.ThrowIfNull(criteria);
+
         var matchingItemIds = criteria switch
         {
             HasSubtitles => context.MediaStreamInfos
                 .Where(ms => ms.StreamType == MediaStreamTypeEntity.Subtitle)
-                .Select(ms => ms.ItemId)
-                .Distinct()
-                .ToHashSet(),
+                .Select(ms => ms.ItemId),
             HasChapterImages => context.Chapters
                 .Where(c => c.ImagePath != null)
-                .Select(c => c.ItemId)
-                .Distinct()
-                .ToHashSet(),
+                .Select(c => c.ItemId),
             HasMediaStreamType m => GetMatchingMediaStreamItemIds(context, m),
             _ => throw new ArgumentOutOfRangeException(nameof(criteria), $"Unknown criteria type: {criteria.GetType().Name}")
         };
 
-        var ancestors = TraverseHierarchyUp(context, matchingItemIds);
+        // One hop up the closure covers every ancestor level.
+        var hierarchyAncestors = context.AncestorIds
+            .Where(e => matchingItemIds.Contains(e.ItemId))
+            .Select(e => e.ParentItemId);
 
-        return ancestors.AsQueryable();
+        var linkParents = ResolveLinkParents(context, matchingItemIds, hierarchyAncestors);
+
+        // The link parents are resolved ids, so they are read back as a sub-select to keep the result
+        // composable. An id without a BaseItem row could never match a caller's row anyway.
+        var linkedParents = context.BaseItems
+            .WhereOneOrMany(linkParents, e => e.Id)
+            .Select(e => e.Id);
+
+        var linkedParentAncestors = context.AncestorIds
+            .WhereOneOrMany(linkParents, e => e.ItemId)
+            .Select(e => e.ParentItemId);
+
+        var seamAncestors = context.AncestorIds
+            .Where(e => hierarchyAncestors.Contains(e.ItemId) || linkedParentAncestors.Contains(e.ItemId))
+            .Select(e => e.ParentItemId);
+
+        return hierarchyAncestors
+            .Concat(linkedParents)
+            .Concat(linkedParentAncestors)
+            .Concat(seamAncestors)
+            .Distinct();
     }
 
-    private static HashSet<Guid> GetMatchingMediaStreamItemIds(JellyfinDbContext context, HasMediaStreamType criteria)
+    private static IQueryable<Guid> GetMatchingMediaStreamItemIds(JellyfinDbContext context, HasMediaStreamType criteria)
     {
         var query = context.MediaStreamInfos
             .Where(ms => ms.StreamType == criteria.StreamType
@@ -130,130 +156,131 @@ public static class DescendantQueryHelper
             query = query.Where(ms => ms.IsExternal == isExternal);
         }
 
-        return query.Select(ms => ms.ItemId).Distinct().ToHashSet();
+        return query.Select(ms => ms.ItemId);
     }
 
-    /// <summary>
-    /// Traverses DOWN the hierarchy from parent folders to find all descendants.
-    /// </summary>
-    private static HashSet<Guid> TraverseHierarchyDown(JellyfinDbContext context, ICollection<Guid> startIds)
+    private static IQueryable<Guid> ClosureDescendants(JellyfinDbContext context, IReadOnlyList<Guid> roots)
     {
-        var visited = new HashSet<Guid>(startIds);
-        var folderStack = new HashSet<Guid>(startIds);
+        var direct = context.AncestorIds
+            .WhereOneOrMany(roots, e => e.ParentItemId)
+            .Select(e => e.ItemId);
 
-        while (folderStack.Count != 0)
-        {
-            var currentFolders = folderStack.ToArray();
-            folderStack.Clear();
+        // An item carries its own chain plus its collection folders, never the UserRootFolder.
+        var indirect = context.AncestorIds
+            .Where(e => direct.Contains(e.ParentItemId))
+            .Select(e => e.ItemId);
 
-            var directChildren = context.AncestorIds
-                .WhereOneOrMany(currentFolders, e => e.ParentItemId)
-                .Select(e => e.ItemId)
-                .ToArray();
+        return direct.Concat(indirect);
+    }
 
-            var linkedChildren = context.LinkedChildren
-                .WhereOneOrMany(currentFolders, e => e.ParentId)
-                .Select(e => e.ChildId)
-                .ToArray();
+    /// <summary>
+    /// Resolves every folder that reaches one of the matching items through a linked edge.
+    /// </summary>
+    /// <returns>The ids of the folders whose linked children lead, at any depth, to a matching item.</returns>
+    private static List<Guid> ResolveLinkParents(JellyfinDbContext context, IQueryable<Guid> matchingItemIds, IQueryable<Guid> ancestorsOfMatches)
+    {
+        // A link sits above the closure as well as above another link: a BoxSet holds a Series whose
+        // episode matches, and another BoxSet holds that BoxSet. So the hop repeats until it stops
+        // finding anything new, and each hop takes the links landing on the set itself or on a folder
+        // that contains it. Only folders owning linked children are ever collected, which bounds this
+        // by the number of BoxSets and Playlists rather than by the item count.
+        var resolved = context.LinkedChildren
+            .Where(e => matchingItemIds.Contains(e.ChildId) || ancestorsOfMatches.Contains(e.ChildId))
+            .Select(e => e.ParentId)
+            .Distinct()
+            .ToHashSet();
+
+        var frontier = resolved.ToList();
+
+        while (frontier.Count != 0)
+        {
+            var containingFolders = context.AncestorIds
+                .WhereOneOrMany(frontier, e => e.ItemId)
+                .Select(e => e.ParentItemId);
 
-            var allChildren = directChildren.Concat(linkedChildren).Distinct().ToArray();
+            var directLinkParents = context.LinkedChildren
+                .WhereOneOrMany(frontier, e => e.ChildId)
+                .Select(e => e.ParentId);
 
-            if (allChildren.Length == 0)
-            {
-                break;
-            }
+            var indirectLinkParents = context.LinkedChildren
+                .Where(e => containingFolders.Contains(e.ChildId))
+                .Select(e => e.ParentId);
 
-            var childFolders = context.BaseItems
-                .WhereOneOrMany(allChildren, e => e.Id)
-                .Where(e => e.IsFolder)
-                .Select(e => e.Id)
-                .ToHashSet();
+            var next = directLinkParents
+                .Concat(indirectLinkParents)
+                .Distinct()
+                .ToArray();
 
-            foreach (var childId in allChildren)
+            frontier = [];
+            foreach (var id in next)
             {
-                if (visited.Add(childId) && childFolders.Contains(childId))
+                // Cyclic links (a BoxSet holding itself, directly or not) terminate on the resolved set.
+                if (resolved.Add(id))
                 {
-                    folderStack.Add(childId);
+                    frontier.Add(id);
                 }
             }
         }
 
-        return visited;
+        return [.. resolved];
     }
 
     /// <summary>
-    /// Traverses DOWN the hierarchy using only AncestorIds (ownership), not LinkedChildren.
+    /// Resolves the roots the descendant sub-selects have to be anchored on.
     /// </summary>
-    private static HashSet<Guid> TraverseHierarchyDownOwned(JellyfinDbContext context, ICollection<Guid> startIds)
+    /// <returns>
+    /// The roots whose AncestorIds closure belongs to the result, and the roots whose LinkedChildren
+    /// belong to the result.
+    /// </returns>
+    private static (List<Guid> ClosureRoots, List<Guid> LinkRoots) ResolveLinkedRoots(JellyfinDbContext context, Guid parentId)
     {
-        var visited = new HashSet<Guid>(startIds);
-        var folderStack = new HashSet<Guid>(startIds);
+        // A folder found through the closure needs no closure hop of its own.
+        var closureRoots = new List<Guid> { parentId };
+        var linkRoots = new List<Guid> { parentId };
+        var visited = new HashSet<Guid> { parentId };
+        var frontier = new List<Guid> { parentId };
 
-        while (folderStack.Count != 0)
+        while (frontier.Count != 0)
         {
-            var currentFolders = folderStack.ToArray();
-            folderStack.Clear();
+            var closureIds = ClosureDescendants(context, frontier);
 
-            var directChildren = context.AncestorIds
-                .WhereOneOrMany(currentFolders, e => e.ParentItemId)
-                .Select(e => e.ItemId)
-                .ToArray();
+            var linkedIds = context.LinkedChildren
+                .WhereOneOrMany(frontier, e => e.ParentId)
+                .Select(e => e.ChildId);
 
-            if (directChildren.Length == 0)
-            {
-                break;
-            }
+            // Folders that own linked children, i.e. the only items whose links are worth following.
+            var linkOwners = context.BaseItems
+                .Where(e => e.IsFolder
+                    && (closureIds.Contains(e.Id) || linkedIds.Contains(e.Id))
+                    && context.LinkedChildren.Any(l => l.ParentId.Equals(e.Id)))
+                .Select(e => e.Id)
+                .ToArray();
 
-            var childFolders = context.BaseItems
-                .WhereOneOrMany(directChildren, e => e.Id)
-                .Where(e => e.IsFolder)
+            var linkedFolders = context.BaseItems
+                .Where(e => e.IsFolder && linkedIds.Contains(e.Id))
                 .Select(e => e.Id)
                 .ToHashSet();
 
-            foreach (var childId in directChildren)
+            frontier = [];
+            foreach (var id in linkOwners.Concat(linkedFolders))
             {
-                if (visited.Add(childId) && childFolders.Contains(childId))
+                if (!visited.Add(id))
                 {
-                    folderStack.Add(childId);
+                    continue;
                 }
-            }
-        }
-
-        return visited;
-    }
-
-    /// <summary>
-    /// Traverses UP the hierarchy from items to find all ancestor folders.
-    /// </summary>
-    private static HashSet<Guid> TraverseHierarchyUp(JellyfinDbContext context, ICollection<Guid> startIds)
-    {
-        var ancestors = new HashSet<Guid>();
-        var itemStack = new HashSet<Guid>(startIds);
 
-        while (itemStack.Count != 0)
-        {
-            var currentItems = itemStack.ToArray();
-            itemStack.Clear();
+                frontier.Add(id);
+                linkRoots.Add(id);
 
-            var ancestorParents = context.AncestorIds
-                .WhereOneOrMany(currentItems, e => e.ItemId)
-                .Select(e => e.ParentItemId)
-                .ToArray();
-
-            var linkedParents = context.LinkedChildren
-                .WhereOneOrMany(currentItems, e => e.ChildId)
-                .Select(e => e.ParentId)
-                .ToArray();
-
-            foreach (var parentId in ancestorParents.Concat(linkedParents))
-            {
-                if (ancestors.Add(parentId))
+                // Only a folder reached through a link contributes a closure that is not covered by
+                // the roots already collected.
+                if (linkedFolders.Contains(id))
                 {
-                    itemStack.Add(parentId);
+                    closureRoots.Add(id);
                 }
             }
         }
 
-        return ancestors;
+        return (closureRoots, linkRoots);
     }
 }