Skip to content

feat: Add ProxyMemoryPool - #736

Open
adamreeve wants to merge 4 commits into
G-Research:masterfrom
adamreeve:proxy-mem-pool
Open

adamreeve wants to merge 4 commits into
G-Research:masterfrom
adamreeve:proxy-mem-pool

Conversation

@adamreeve

@adamreeve adamreeve commented Nov 2, 2026 •

Copy link
Copy Markdown
Contributor

Adds a wrapper for the Arrow C++ ProxyMemoryPool class, and also exposes the native total_bytes_allocated and num_allocations methods on memory pools.

This implements part 1 of #729.

Unlike the existing memory pool types, a ProxyMemoryPool is non-static so must be kept alive while objects that use memory allocated through it remain alive. This required keeping references to the memory pool in things that depend on it.

In order to prevent use-after free when native memory allocations are freed, a native RefCountedProxyMemoryPool class was added that deletes itself only after all allocations made through it have been freed.

Also add total_bytes_allocated and num_allocations methods
on memory pools.

Unlike the existing memory pool types, a ProxyMemoryPool
is non-static so must be kept alive while objects that use
memory allocated through it remain alive.
This required adding some extra book-keeping.
@adamreeve
adamreeve marked this pull request as ready for review November 2, 2026 22:38
@adamreeve
adamreeve requested a balanced review from Copilot November 2, 2026 22:38

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

Critical pool-ownership gaps in properties getters and derived readers, streams, and batches allow live native consumers to outlast the owning proxy, causing use-after-free.

6 open issues
What changed in this PR

Adds a managed ProxyMemoryPool wrapper for tracking native allocations independently per writer, implementing part 1 of #729. The PR also exposes allocation counters and adds pool-lifetime retention to dependent objects.

Changes:

  • Adds disposable proxy pools and native bindings for cumulative allocation statistics.
  • Retains pool references in properties, readers, writers, and buffers.
  • Tests proxy statistics, nested pools, independent writer accounting, and lifetime behavior.
File Description
csharp/​WriterPropertiesBuilder.cs Passes the configured pool to built properties.
csharp/​WriterProperties.cs Retains the pool alongside the native properties handle.
csharp/​ReaderProperties.cs Retains the configured reader pool.
csharp/​PublicAPI.Unshipped.txt Declares the proxy pool and allocation counters.
csharp/​ProxyMemoryPool.cs Implements disposable proxy ownership using a safe handle.
csharp/​ParquetFileWriter.cs Retains the configured pool through native writer release.
csharp/​ParquetFileReader.cs Retains the configured pool through native reader release.
csharp/​MemoryPool.cs Adds allocation counters and lifetime-retention helpers.
csharp/​IO/​ResizableBuffer.cs Passes pool ownership to the buffer wrapper.
csharp/​IO/​BufferOutputStream.cs Passes its stored pool to finished buffers.
csharp/​IO/​Buffer.cs Retains the pool through native buffer release.
csharp/​Arrow/​FileWriter.cs Retains the configured pool for Arrow writers.
csharp/​Arrow/​FileReader.cs Retains the configured pool for Arrow readers.
csharp.test/​TestProxyMemoryPool.cs Tests accounting, ownership, and disposal behavior.
cpp/​MemoryPool.cpp Exports proxy creation, deletion, and allocation counters.

Review effort: Balanced

Comment thread csharp/Arrow/FileReader.cs Resolved
Comment thread csharp/IO/Buffer.cs Resolved
Comment thread csharp/IO/BufferOutputStream.cs Resolved
Comment thread csharp/ParquetFileWriter.cs Resolved
Comment thread csharp/ReaderProperties.cs Resolved
Comment thread csharp/WriterProperties.cs Resolved
This allow allocations to outlive the managed pool without risk of
use-after-free.
This ensures that the pool isn't freed before use.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

A critical use-after-free risk remains because both buffer constructors pass the pool’s raw pointer to native code without keeping memoryPool alive through construction; both need GC.KeepAlive(memoryPool).

1 open issue
6 resolved since last review

Review effort: Balanced

Comment thread csharp/ProxyMemoryPool.cs Resolved

@marcin-krystianc marcin-krystianc left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants