Search code examples
c#.netgarbage-collection

Is there a way to store a MemoryHandle exclusively in unmanaged memory?


I have a method which takes in a Memory<T> object. I want to pin it, and store a pointer to it exclusively in unmanaged memory. I know I can get this pinned pointer by using memory.Pin() to create a MemoryHandle and then grabbing memoryHandle.Pointer. However, because I only want to store this in unmanaged memory I need to ensure that the managed object the memory handle points somewhere into is not collected until I am done with it. Is there a way I can store the MemoryHandle in unmanaged memory without the underlying object being collected, than free it at a later time?

I know that in theory this can be done by storing a GCHandle alongside my pointer than disposing it when I'm done. The problem is I don't see any way to retrieve a GCHandle from the Memory or MemoryHandle objects, even though I'm pretty sure at least the memory handle contains a GC handle internally. So, if there where a way to get a GCHandle from the MemoryHandle that would be a good solution to this problem.


Solution

  • Here is the solution I came up with. It is a simple wrapper around a GCHandle. There is a method which creates the handle from a MemoryHandle in a safe way which avoids object or GCHandle allocations in the vast majority of cases. It is important that you do not dispose the MemoryHandle you created the UnmanagedMemoryHandle from.

    public unsafe readonly struct UnmanagedMemoryHandle : IDisposable {
    
        // https://github.com/dotnet/runtime/blob/main/src/libraries/System.Private.CoreLib/src/System/Buffers/MemoryHandle.cs
        private struct ExposedMemoryHandle {
            public void* Pointer;
            public GCHandle Handle;
            public IPinnable? Pinnable;
        }
    
        public static UnmanagedMemoryHandle FromMemoryHandle(MemoryHandle memoryHandle) {
            return FromMemoryHandle(ref memoryHandle);
        }
    
        public static UnmanagedMemoryHandle FromMemoryHandle(ref MemoryHandle memoryHandle) {
            ref ExposedMemoryHandle exposedMemoryHandle = ref Unsafe.As<MemoryHandle, ExposedMemoryHandle>(ref memoryHandle);
    
            bool hasHandle = exposedMemoryHandle.Handle.IsAllocated;
            bool hasPinnable = exposedMemoryHandle.Pinnable != null;
    
            GCHandle handle;
    
            if (hasHandle && !hasPinnable) {
                if (exposedMemoryHandle.Handle.Target is not IPinnable) {
                    // If we only have a handle, we just store that handle
                    handle = exposedMemoryHandle.Handle;
                } else {
                    // Later, we check if the handle is an instance to IPinnable to see if we need to unpin it.
                    //  That's a problem in the case the the Handle just happens to be an IPinnable. In this case
                    //  we wrap our handle in a type that is not IPinnable
                    handle = GCHandle.Alloc(new DoubleMemoryHandle(exposedMemoryHandle.Handle, null));
                }
            } else if (hasPinnable && !hasHandle) {
                // If we only have a pinnable, we store a handle to that pinnable
                handle = GCHandle.Alloc(exposedMemoryHandle.Pinnable);
            } else if (hasHandle && hasPinnable) {
                // If we have both, store a handle to an object containing the handle and pinnable
                handle = GCHandle.Alloc(new DoubleMemoryHandle(exposedMemoryHandle.Handle, exposedMemoryHandle.Pinnable));
            } else {
                // If we have neither, we don't need to store anything
                handle = default;
            }
    
            return new UnmanagedMemoryHandle(handle);
        }
    
        private record DoubleMemoryHandle(GCHandle Handle, IPinnable? Pinnable);
    
        public readonly GCHandle Handle;
    
        private UnmanagedMemoryHandle(GCHandle handle) {
            Handle = handle;
        }
    
        public void Dispose() {
            if (!Handle.IsAllocated) return;
    
            object? target = Handle.Target;
    
            if (target is IPinnable pinnable) {
                pinnable.Unpin();
            } else if (target is DoubleMemoryHandle doubleMemoryHandle) {
                doubleMemoryHandle.Handle.Free();
                doubleMemoryHandle.Pinnable?.Unpin();
            }
    
            Handle.Free();
        }
    }