Prevent use-after free if memory pool is disposed

csharp.test/TestProxyMemoryPool.cs

@@ -171,6 +171,82 @@ public static void TestDisposedProxy()
             Assert.Throws<NullReferenceException>(() => new ResizableBuffer(128, proxy));
         }
 
+        [TestCase(0)]
+        [TestCase(2)]
+        public static void TestProxyDisposedWhileWriterInUse(int rowGroupsBeforeDispose)
+        {
+            const int numRows = 10_000;
+            using var testDir = new TempWorkingDirectory();
+            var path = Path.Combine(testDir.DirectoryPath, "test.parquet");
+
+            var proxy = new ProxyMemoryPool(MemoryPool.GetDefaultMemoryPool());
+            using (var writer = CreateWriter(path, proxy))
+            {
+                for (var rowGroup = 0; rowGroup < rowGroupsBeforeDispose; ++rowGroup)
+                {
+                    using var rowGroupWriter = writer.AppendRowGroup();
+                    WriteRows(rowGroupWriter, numRows);
+                    rowGroupWriter.Close();
+                }
+
+                proxy.Dispose();
+
+                using (var rowGroupWriter = writer.AppendRowGroup())
+                {
+                    WriteRows(rowGroupWriter, numRows);
+                    rowGroupWriter.Close();
+                }
+
+                writer.Close();
+            }
+
+            Assert.Throws<NullReferenceException>(() => CreateWriter(Path.Combine(testDir.DirectoryPath, "other.parquet"), proxy));
+
+            using var readPool = new ProxyMemoryPool(MemoryPool.GetDefaultMemoryPool());
+            CheckRows(path, numRows, readPool);
+        }
+
+        [Test]
+        public static void TestProxyDisposedWhileBuilderInUse()
+        {
+            const int numRows = 10_000;
+            using var testDir = new TempWorkingDirectory();
+            var path = Path.Combine(testDir.DirectoryPath, "test.parquet");
+
+            var proxy = new ProxyMemoryPool(MemoryPool.GetDefaultMemoryPool());
+            using var builder = new WriterPropertiesBuilder();
+            builder.MemoryPool(proxy);
+            proxy.Dispose();
+
+            using (var properties = builder.Build())
+            using (var writer = new ParquetFileWriter(path, Columns, properties))
+            {
+                using var rowGroupWriter = writer.AppendRowGroup();
+                WriteRows(rowGroupWriter, numRows);
+                rowGroupWriter.Close();
+                writer.Close();
+            }
+
+            using var readPool = new ProxyMemoryPool(MemoryPool.GetDefaultMemoryPool());
+            CheckRows(path, numRows, readPool);
+        }
+
+        [Test]
+        public static void TestProxyDisposedWhileReaderInUse()
+        {
+            const int numRows = 10_000;
+            using var testDir = new TempWorkingDirectory();
+            var path = Path.Combine(testDir.DirectoryPath, "test.parquet");
+            WriteFile(path, numRows);
+
+            var proxy = new ProxyMemoryPool(MemoryPool.GetDefaultMemoryPool());
+            using var readerProperties = ReaderProperties.WithMemoryPool(proxy);
+            using var reader = new ParquetFileReader(path, readerProperties);
+            proxy.Dispose();
+
+            CheckRows(reader, numRows);
+        }
+
         [Test]
         public static void TestNullWrappedPool()
         {

csharp/MemoryPool.cs

@@ -140,7 +140,22 @@ internal MemoryPool(IntPtr handle)
         internal virtual IntPtr Handle => _handle;
 
         /// <summary>
-        /// Wrap a native free function so that it keeps a memory pool alive until the native object has been freed.
+        /// Take a lease on the native memory pool, which keeps it valid until the lease is released,
+        /// even if this pool is disposed in the meantime.
+        /// </summary>
+        internal virtual void AddLease()
+        {
+        }
+
+        /// <summary>
+        /// Release a lease taken with <see cref="AddLease"/>.
+        /// </summary>
+        internal virtual void ReleaseLease()
+        {
+        }
+
+        /// <summary>
+        /// Wrap a native free function so that it keeps a memory pool alive and valid until the native object has been freed.
         /// </summary>
         internal static Action<IntPtr> KeepAliveUntilFreed(Action<IntPtr> free, MemoryPool? memoryPool)
         {
@@ -149,10 +164,11 @@ internal static Action<IntPtr> KeepAliveUntilFreed(Action<IntPtr> free, MemoryPo
                 return free;
             }
 
+            memoryPool.AddLease();
             return ptr =>
             {
                 free(ptr);
-                GC.KeepAlive(memoryPool);
+                memoryPool.ReleaseLease();
             };
         }
 

csharp/ProxyMemoryPool.cs

@@ -13,8 +13,11 @@ namespace ParquetSharp
     /// even when they share the same underlying memory pool.
     ///
     /// Unlike the global memory pools, a proxy memory pool owns a native resource.
-    /// It must not be disposed until every object that was configured to use it has been disposed,
+    /// It should not be disposed until every object that was configured to use it has been disposed,
     /// for example any writer, reader, properties or output stream it was passed to.
+    /// After disposal, statistics are no longer available and the pool cannot be used to configure new objects.
+    /// Objects already using the pool keep the native pool alive so that disposing it early does not cause
+    /// memory errors, but continuing to use them after the pool is disposed is not recommended.
     ///
     /// Memory allocated from the pool may safely outlive it, for example buffers or Arrow arrays
     /// read using it. The native pool is only freed once all allocations made through it have been freed,
@@ -39,14 +42,17 @@ private ProxyMemoryPool(ProxyMemoryPoolHandle handle)
 
         public void Dispose()
         {
+            _disposed = true;
             _handle.Dispose();
         }
 
         internal override IntPtr Handle
         {
             get
             {
-                if (_handle.IsClosed)
+                // Objects already using the pool hold a lease that keeps the native pool alive after disposal,
+                // so check the disposed flag rather than only whether the native handle has been released.
+                if (_disposed || _handle.IsClosed)
                 {
                     throw new NullReferenceException("null native handle");
                 }
@@ -55,6 +61,18 @@ internal override IntPtr Handle
             }
         }
 
+        internal override void AddLease()
+        {
+            // Throws ObjectDisposedException if the native pool has already been freed
+            var success = false;
+            _handle.DangerousAddRef(ref success);
+        }
+
+        internal override void ReleaseLease()
+        {
+            _handle.DangerousRelease();
+        }
+
         private static ProxyMemoryPoolHandle CreateHandle(MemoryPool wrappedPool)
         {
             if (wrappedPool == null) throw new ArgumentNullException(nameof(wrappedPool));
@@ -94,5 +112,6 @@ protected override bool ReleaseHandle()
         // If the wrapped pool is also a proxy, the native proxy holds a reference to it,
         // so we don't need to keep a reference to the managed wrapped pool.
         private readonly ProxyMemoryPoolHandle _handle;
+        private volatile bool _disposed;
     }
 }

csharp/WriterPropertiesBuilder.cs

@@ -17,7 +17,11 @@ public sealed class WriterPropertiesBuilder : IDisposable
         public WriterPropertiesBuilder()
         {
             ExceptionInfo.Check(WriterPropertiesBuilder_Create(out var handle));
-            _handle = new ParquetHandle(handle, WriterPropertiesBuilder_Free);
+            _handle = new ParquetHandle(handle, ptr =>
+            {
+                WriterPropertiesBuilder_Free(ptr);
+                _memoryPool?.ReleaseLease();
+            });
             ApplyDefaults();
         }
 
@@ -544,6 +548,10 @@ public WriterPropertiesBuilder SortingColumns(WriterProperties.SortingColumn[] s
         public WriterPropertiesBuilder MemoryPool(MemoryPool memoryPool)
         {
             ExceptionInfo.Check(WriterPropertiesBuilder_Memory_Pool(_handle.IntPtr, memoryPool.Handle));
+
+            // The native builder holds a pointer to the pool, so keep it valid until the builder is freed
+            memoryPool.AddLease();
+            _memoryPool?.ReleaseLease();
             _memoryPool = memoryPool;
 
             GC.KeepAlive(_handle);