Implement stream pooling and more fine grained lock

src/Lumina/Data/SqPack.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.IO;
 using System.Reflection;
@@ -10,13 +11,14 @@ namespace Lumina.Data
     public class SqPackInflateException : Exception
     {
         public SqPackInflateException( string message ) : base( message )
-        {
-        }
+        { }
     }
 
+    /// <summary>Represents a .dat file of a sqpack file group.</summary>
     public class SqPack
     {
         private readonly GameData _gameData;
+        private readonly ConcurrentDictionary< long, WeakReference< FileResource > > _cache;
 
         /// <summary>
         /// Where the actual file is located on disk
@@ -33,22 +35,26 @@ public class SqPack
         /// </summary>
         public string FullName => File.FullName;
 
+        /// <summary>
+        /// Gets the file header of this SqPack file.
+        /// </summary>
         public SqPackHeader SqPackHeader { get; private set; }
 
         /// <summary>
         /// PC and PS4 are LE, PS3 is BE
         /// </summary>
         [Obsolete( "Use property \"ConvertEndianness\" from \"LuminaBinaryReader\" instead" )]
-        public bool ShouldConvertEndianness
-        {
+        public bool ShouldConvertEndianness {
             // todo: what about reading LE files on BE device? do we even care?
             get => BitConverter.IsLittleEndian && SqPackHeader.platformId == PlatformId.PS3 ||
-                   !BitConverter.IsLittleEndian && SqPackHeader.platformId != PlatformId.PS3;
+                !BitConverter.IsLittleEndian && SqPackHeader.platformId != PlatformId.PS3;
         }
 
-        protected readonly Dictionary< long, WeakReference< FileResource > > FileCache;
-
-        protected readonly object CacheLock = new object();
+        /// <summary>Unused field.</summary>
+        [Obsolete( "Not used." )] protected readonly Dictionary< long, WeakReference< FileResource > > FileCache = new();
+
+        /// <summary>Unused field.</summary>
+        [Obsolete( "Not used." )] protected readonly object CacheLock = new();
 
         internal SqPack( FileInfo file, GameData gameData )
         {
@@ -60,77 +66,70 @@ internal SqPack( FileInfo file, GameData gameData )
             _gameData = gameData;
 
             // always init the cache just in case the should cache setting is changed later
-            FileCache = new Dictionary< long, WeakReference< FileResource > >();
+            _cache = new();
 
             File = file;
 
-            using var ss = new SqPackStream( File );
-
-            SqPackHeader = ss.GetSqPackHeader();
-        }
-
-        protected T? GetCachedFile< T >( long offset ) where T : FileResource
-        {
-            if( !FileCache.TryGetValue( offset, out var weakRef ) )
-            {
-                return null;
-            }
-
-            if( !weakRef.TryGetTarget( out var cachedFile ) )
-            {
-                return null;
-            }
-
-            // only return from cache if target type matches
-            // otherwise we'll force a cache miss and parse it as per usual
-            if( cachedFile is T obj )
-            {
-                return obj;
-            }
-
-            return null;
+            using var sh = GetStreamHolder();
+            SqPackHeader = sh.Stream.GetSqPackHeader();
         }
 
+        /// <summary>Gets the file metadata of a file at the given offset.</summary>
+        /// <param name="offset">The byte offset of a file in this SqPack .dat file.</param>
+        /// <returns>The corresponding file metadata.</returns>
+        /// <remarks>No verification is performed on whether <paramref name="offset"/> points to a valid file metadata.</remarks>
         public SqPackFileInfo GetFileMetadata( long offset )
         {
-            using var ss = new SqPackStream( File, SqPackHeader.platformId );
-
-            return ss.GetFileMetadata( offset );
+            using var sh = GetStreamHolder();
+            return sh.Stream.GetFileMetadata( offset );
         }
 
+        /// <summary>Reads the file at the given offset.</summary>
+        /// <param name="offset">The byte offset of a file in this SqPack .dat file.</param>
+        /// <typeparam name="T">The file resource type.</typeparam>
+        /// <returns>The corresponding file resource.</returns>
         public T ReadFile< T >( long offset ) where T : FileResource
         {
             var cacheBehaviour = FileOptionsAttribute.FileCacheBehaviour.None;
-
-            var fileOpts = typeof( T ).GetCustomAttribute< FileOptionsAttribute >();
-            if( fileOpts != null )
-            {
+
+            if( typeof( T ).GetCustomAttribute< FileOptionsAttribute >() is { } fileOpts )
                 cacheBehaviour = fileOpts.CacheBehaviour;
-            }
-
+
             if( !_gameData.Options.CacheFileResources || cacheBehaviour == FileOptionsAttribute.FileCacheBehaviour.Never )
             {
-                using var ss = new SqPackStream( File, SqPackHeader.platformId );
-                return ss.ReadFile< T >( offset );
+                using var rsh = GetStreamHolder();
+                return rsh.Stream.ReadFile< T >( offset );
             }
-
-
-            lock( CacheLock )
-            {
-                var obj = GetCachedFile< T >( offset );
 
-                if( obj != null )
-                {
-                    return obj;
-                }
+            // WeakReference ctor accepts null.
+            var weakRef = _cache.GetOrAdd( offset, static _ => new( null! ) );
+            if( weakRef.TryGetTarget( out var valueUntyped ) && valueUntyped is T value )
+                return value;
 
-                using var ss = new SqPackStream( File, SqPackHeader.platformId );
-                var file = ss.ReadFile< T >( offset );
-
-                FileCache[ offset ] = new WeakReference< FileResource >( file );
+            lock( weakRef )
+            {
+                if( weakRef.TryGetTarget( out valueUntyped ) && valueUntyped is T value2 )
+                    return value2;
 
-                return file;
+                using var sh = GetStreamHolder();
+                var f = sh.Stream.ReadFile< T >( offset );
+                weakRef.SetTarget( f );
+                return f;
             }
         }
+
+        /// <summary>Gets the file, if the cached instance of <see cref="FileResource"/> can be cast as <typeparamref name="T"/>.</summary>
+        /// <param name="offset">The byte offset of a file in this SqPack .dat file.</param>
+        /// <typeparam name="T">The file resource type.</typeparam>
+        /// <returns>The corresponding file resource, or <c>null</c> if the file was not cached.</returns>
+        protected T? GetCachedFile< T >( long offset ) where T : FileResource =>
+            _cache.TryGetValue( offset, out var cached ) && cached.TryGetTarget( out var valueUntyped ) ? valueUntyped as T : null;
+
+        /// <summary>Rents a stream from the pool if stream pooling is enabled, or creates a new stream.</summary>
+        /// <returns>A stream holder that will either return the stream to the pool or close the stream.</returns>
+        private SqPackStreamPool.StreamHolder GetStreamHolder() =>
+            _gameData.StreamPool is { } streamPool
+                ? streamPool.RentScoped( this )
+                : new( new SqPackStream( File, SqPackHeader.platformId ) );
     }
 }

src/Lumina/Data/SqPackStreamPool.cs

@@ -0,0 +1,110 @@
+using System;
+using System.Collections.Concurrent;
+using System.Runtime.CompilerServices;
+using System.Threading;
+using Microsoft.Extensions.ObjectPool;
+
+namespace Lumina.Data;
+
+/// <summary>A stream pool for use with <see cref="GameData.StreamPool"/>.</summary>
+/// <remarks>
+/// As creating a handle for a file on the filesystem takes much longer than seeking an already opened handle, reusing handles will help with the overall
+/// performance, both single-threaded and multi-threaded.
+/// </remarks>
+public sealed class SqPackStreamPool : IDisposable
+{
+    private readonly ObjectPoolProvider _objectPoolProvider;
+    private ConcurrentDictionary< SqPack, ObjectPool< SqPackStream > >? _pools = new();
+
+    /// <summary>Initializes a new instance of <see cref="SqPackStreamPool"/> class.</summary>
+    /// <param name="maximumRetained">The maximum number of streams to retain in the pool per file. If a non-positive number is given, the default value of
+    /// <see cref="DefaultObjectPoolProvider"/> will be used.</param>
+    public SqPackStreamPool( int maximumRetained = 0 )
+    {
+        var dopp = new DefaultObjectPoolProvider();
+        if( maximumRetained > 0 )
+            dopp.MaximumRetained = maximumRetained;
+        _objectPoolProvider = dopp;
+    }
+
+    /// <summary>Initializes a new instance of <see cref="SqPackStreamPool"/> class.</summary>
+    /// <param name="objectPoolProvider">The object pool provider to use.</param>
+    public SqPackStreamPool( ObjectPoolProvider objectPoolProvider ) => _objectPoolProvider = objectPoolProvider;
+
+    /// <inheritdoc/>
+    public void Dispose()
+    {
+        if( Interlocked.Exchange( ref _pools, null ) is not { } pools )
+            return;
+
+        foreach( var pool in pools.Values )
+            ( pool as IDisposable )?.Dispose();
+        pools.Clear();
+    }
+
+    /// <summary>Rents a stream from the pool.</summary>
+    /// <param name="dat">An instance of <see cref="SqPack"/> to get the stream of.</param>
+    /// <returns>The rented stream.</returns>
+    /// <remarks>If this pool has been disposed, a new instance of <see cref="SqPackStream"/> will be created instead.</remarks>
+    internal StreamHolder RentScoped( SqPack dat )
+    {
+        if( _pools is not { } pools )
+            return new( new SqPackStream( dat.File, dat.SqPackHeader.platformId ) );
+
+        var pool = pools.GetOrAdd(
+            dat,
+            static ( key, owner ) => owner._objectPoolProvider.Create( new SqPackStreamObjectPoolProvider( owner, key ) ),
+            this );
+        return new( pool );
+    }
+
+    /// <summary>A struct facilitating the scoped borrowing from a pool.</summary>
+    /// <remarks>Multithreaded usage is not supported.</remarks>
+    internal struct StreamHolder : IDisposable
+    {
+        private readonly ObjectPool< SqPackStream >? _pool;
+        private SqPackStream? _stream;
+
+        [MethodImpl( MethodImplOptions.AggressiveInlining )]
+        public StreamHolder( ObjectPool< SqPackStream > pool )
+        {
+            _pool = pool;
+            _stream = pool.Get();
+        }
+
+        [MethodImpl( MethodImplOptions.AggressiveInlining )]
+        public StreamHolder( SqPackStream stream ) => _stream = stream;
+
+        public SqPackStream Stream {
+            [MethodImpl( MethodImplOptions.AggressiveInlining )]
+            get => _stream ?? throw new ObjectDisposedException( nameof( StreamHolder ) );
+        }
+
+        public void Dispose()
+        {
+            if( _stream is null )
+                return;
+            if( _pool is null )
+                _stream.Dispose();
+            else
+                _pool.Return( _stream );
+            _stream = null;
+        }
+    }
+
+    private class SqPackStreamObjectPoolProvider : IPooledObjectPolicy< SqPackStream >
+    {
+        private readonly SqPackStreamPool _owner;
+        private readonly SqPack _sqpack;
+
+        public SqPackStreamObjectPoolProvider( SqPackStreamPool owner, SqPack sqpack )
+        {
+            _owner = owner;
+            _sqpack = sqpack;
+        }
+
+        public SqPackStream Create() => new( _sqpack.File, _sqpack.SqPackHeader.platformId );
+
+        public bool Return( SqPackStream obj ) => _owner._pools is not null && obj.BaseStream.CanRead;
+    }
+}

src/Lumina/GameData.cs

@@ -13,8 +13,10 @@
 
 namespace Lumina
 {
-    public class GameData
+    public class GameData : IDisposable
     {
+        ~GameData() => Dispose( false );
+
         /// <summary>
         /// The current data path that Lumina is using to load files.
         /// </summary>
@@ -50,6 +52,12 @@ public class GameData
         /// easily defer file loading onto another thread.
         /// </summary>
         public FileHandleManager FileHandleManager { get; private set; }
+
+        /// <summary>
+        /// Provides a pool for file streams for .dat files.
+        /// </summary>
+        /// <remarks>The pool will be disposed when <see cref="Dispose"/> is called.</remarks>
+        public SqPackStreamPool? StreamPool { get; set; }
 
         internal ILogger? Logger { get; private set; }
 
@@ -152,6 +160,13 @@ public GameData( string dataPath, ILogger logger, LuminaOptions? options = null!
             };
         }
 
+        /// <inheritdoc/>
+        public void Dispose()
+        {
+            Dispose( true );
+            GC.SuppressFinalize( this );
+        }
+
         /// <summary>
         /// Load a raw file given a game file path
         /// </summary>
@@ -315,6 +330,17 @@ public void ProcessFileHandleQueue()
             FileHandleManager.ProcessQueue();
         }
 
+        /// <summary>Disposes this object.</summary>
+        /// <param name="disposing">Whether this function is being called from <see cref="Dispose"/>.</param>
+        protected virtual void Dispose( bool disposing )
+        {
+            if( disposing )
+            {
+                StreamPool?.Dispose();
+                StreamPool = null;
+            }
+        }
+
         internal void SetCurrentContext()
         {
             SetCurrentContext( this );

src/Lumina/Lumina.csproj

@@ -27,6 +27,7 @@
     </PropertyGroup>
 
     <ItemGroup>
+        <PackageReference Include="Microsoft.Extensions.ObjectPool" Version="9.0.0-preview.1.24081.5" />
         <PackageReference Include="Microsoft.SourceLink.GitHub" Version="1.0.0" PrivateAssets="All" />
         <PackageReference Include="MinVer" Version="2.3.1">
           <PrivateAssets>all</PrivateAssets>