Support for custom memory allocation (Sdk.SetCustomMemoryAllocator me…

…thod)

K4AdotNet.Tests.Unit/Sensor/ImageTests.cs

@@ -426,5 +426,95 @@ public void TestImageSizeCalculationCustom()
         }
 
         #endregion
+
+        #region Test custom memory management
+
+        private sealed class TestMemoryAllocator : ICustomMemoryAllocator
+        {
+            public int AllocateCount { get; private set; }
+            public nint AllocContextValue { get; set; }
+            public int LastAllocSizeValue { get; private set; }
+            public nint LastAllocReturnValue { get; private set; }
+
+            public int FreeCount { get; private set; }
+            public nint LastFreeContextValue { get; private set; }
+
+            nint ICustomMemoryAllocator.Allocate(int size, out nint context)
+            {
+                AllocateCount++;
+                context = AllocContextValue;
+                LastAllocSizeValue = size;
+                return LastAllocReturnValue = Marshal.AllocHGlobal(size);
+            }
+
+            void ICustomMemoryAllocator.Free(nint buffer, nint context)
+            {
+                FreeCount++;
+                LastFreeContextValue = context;
+                Marshal.FreeHGlobal(buffer);
+            }
+        }
+
+        [TestMethod]
+        public void TestCustomMemoryManagement()
+        {
+            // Set our test custom allocator
+            var testAllocator = new TestMemoryAllocator();
+            Sdk.SetCustomMemoryAllocator(testAllocator);
+            // Check initial state
+            Assert.AreEqual(0, testAllocator.AllocateCount);
+            Assert.AreEqual(0, testAllocator.FreeCount);
+
+            // The first test image - should result in one memory allocation
+            var testContextA = testAllocator.AllocContextValue = 12345;
+            var testImageA = new Image(ImageFormat.Depth16, 1, 1);
+            // One allocation but no calls of Free
+            Assert.AreEqual(1, testAllocator.AllocateCount);    // 0 -> 1 !
+            Assert.AreEqual(0, testAllocator.FreeCount);        // unchanged
+            // SDK can request bigger buffer and place actual image buffer somewhere inside allocated one
+            Assert.IsTrue(testAllocator.LastAllocSizeValue >= 2);
+            Assert.IsTrue(testAllocator.LastAllocReturnValue <= testImageA.Buffer);
+            Assert.IsTrue(testAllocator.LastAllocReturnValue + testAllocator.LastAllocSizeValue >= testImageA.Buffer + 2);
+
+            // The second test image - should result in one more memory allocation
+            var testContextB = testAllocator.AllocContextValue = 98765;
+            var testImageB = new Image(ImageFormat.ColorBgra32, 1, 1);
+            // Two allocations but still no calls of Free
+            Assert.AreEqual(2, testAllocator.AllocateCount);    // 1 -> 2 !
+            Assert.AreEqual(0, testAllocator.FreeCount);        // unchanged
+            // SDK can request bigger buffer and place actual image buffer somewhere inside allocated one
+            Assert.IsTrue(testAllocator.LastAllocSizeValue >= 4);
+            Assert.IsTrue(testAllocator.LastAllocReturnValue <= testImageB.Buffer);
+            Assert.IsTrue(testAllocator.LastAllocReturnValue + testAllocator.LastAllocSizeValue >= testImageB.Buffer + 4);
+
+            // Disposing of the first test image - should result in appropriate call of Free method
+            testImageA.Dispose();
+            // Now, one call to Free
+            Assert.AreEqual(2, testAllocator.AllocateCount);    // unchanged
+            Assert.AreEqual(1, testAllocator.FreeCount);        // 0 -> 1 !
+            Assert.AreEqual(testContextA, testAllocator.LastFreeContextValue);
+
+            // Clear custom allocator
+            Sdk.SetCustomMemoryAllocator(null);
+
+            // Now creation of test image does not result in calls to our testAllocator instance
+            var testImageC = new Image(ImageFormat.ColorYUY2, testWidth, testHeight);
+            Assert.AreEqual(2, testAllocator.AllocateCount);
+            Assert.AreEqual(1, testAllocator.FreeCount);
+
+            // The second test image was created with the aid of out test allocator,
+            // this why it should be releasing using the same allocator in spite of the fact that this allocator is not active anymore
+            testImageB.Dispose();
+            Assert.AreEqual(2, testAllocator.AllocateCount);    // unchanged
+            Assert.AreEqual(2, testAllocator.FreeCount);        // 1 -> 2 !
+            Assert.AreEqual(testContextB, testAllocator.LastFreeContextValue);
+
+            // But for the testImageC our testAllocator shouldn't be called on releasing either
+            testImageC.Dispose();
+            Assert.AreEqual(2, testAllocator.AllocateCount);    // unchanged
+            Assert.AreEqual(2, testAllocator.FreeCount);        // unchanged value!
+        }
+
+        #endregion
     }
 }

K4AdotNet/ICustomMemoryAllocator.cs

@@ -0,0 +1,22 @@
+using System;
+
+namespace K4AdotNet
+{
+    /// <summary>
+    /// Base interface for custom memory manager that can be used to allocate and free memory used by images.
+    /// </summary>
+    /// <seealso cref="Sdk.SetCustomMemoryAllocator(ICustomMemoryAllocator?)"/>
+    public interface ICustomMemoryAllocator
+    {
+        /// <summary>Function for a memory allocation. Will be called by internals of Azure Kinect SDK.</summary>
+        /// <param name="size">Minimum size in bytes needed for the buffer.</param>
+        /// <param name="context">Output parameter for a context that will be provided in the subsequent call to the <see cref="Free(IntPtr, IntPtr)"/> callback.</param>
+        /// <returns>A pointer to the newly allocated memory.</returns>
+        IntPtr Allocate(int size, out IntPtr context);
+
+        /// <summary>Function for a memory object being destroyed. Will be called by internals of Azure Kinect SDK.</summary>
+        /// <param name="buffer">The buffer pointer that was supplied by the <see cref="Allocate(int, out IntPtr)"/> method and that should be free.</param>
+        /// <param name="context">The context for    the memory object that needs to be destroyed that was supplied by the <see cref="Allocate(int, out IntPtr)"/> method.</param>
+        void Free(IntPtr buffer, IntPtr context);
+    }
+}

K4AdotNet/Sdk.cs

@@ -1,10 +1,10 @@
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.IO;
 using System.Linq;
-using System.Reflection;
 
 [assembly: CLSCompliant(isCompliant: true)]
 
@@ -144,6 +144,48 @@ public static void ConfigureBodyTrackingLogging(TraceLevel level, bool logToStdo
 
         #endregion
 
+        #region Custom memory allocation
+
+        // to keep callbacks alive
+        private static readonly ConcurrentBag<Sensor.NativeApi.MemoryAllocateCallback> allocateCallbacks = new();
+        private static readonly ConcurrentBag<Sensor.NativeApi.MemoryDestroyCallback> destroyCallbacks = new();
+
+        /// <summary>
+        /// Sets or clears custom memory allocator to be use by internals of Azure Kinect SDK.
+        /// </summary>
+        /// <param name="allocator">
+        /// Instance of <see cref="ICustomMemoryAllocator"/> to allocate and free memory in internals of Azure Kinect SDK
+        /// or <see langword="null"/> to "clear" custom memory allocator (that is, to switch to standard built-in SDK's allocator).
+        /// </param>
+        /// <exception cref="InvalidOperationException">Setting or clearing of custom memory allocator was failed. See log for details.</exception>
+        /// <remarks>
+        /// All instances of <paramref name="allocator"/> will be keeping alive forever because they can be used
+        /// to free memory even after setting custom allocator to <see langword="null"/>.
+        /// </remarks>
+        public static void SetCustomMemoryAllocator(ICustomMemoryAllocator? allocator)
+        {
+            if (allocator is null)
+            {
+                var res = Sensor.NativeApi.SetAllocator(null, null);
+                if (res != NativeCallResults.Result.Succeeded)
+                    throw new InvalidOperationException("Cannot clear custom memory allocator");
+                return;
+            }
+
+            var allocateCallback = new Sensor.NativeApi.MemoryAllocateCallback(allocator.Allocate);
+            var destroyCallback = new Sensor.NativeApi.MemoryDestroyCallback(allocator.Free);
+
+            // to keep callbacks alive
+            allocateCallbacks.Add(allocateCallback);
+            destroyCallbacks.Add(destroyCallback);
+
+            var rs = Sensor.NativeApi.SetAllocator(allocateCallback, destroyCallback);
+            if (rs != NativeCallResults.Result.Succeeded)
+                throw new InvalidOperationException("Cannot set custom memory allocator");
+        }
+
+        #endregion
+
         #region Body tracking SDK availability and initialization
 
         /// <summary>URL to step-by-step instruction "How to set up Body Tracking SDK". Helpful for UI and user messages.</summary>
@@ -278,7 +320,7 @@ private static bool TryGetBodyTrackingRuntimePath(
 
 #if NET461 || NETSTANDARD2_0
             // Try location of this assembly
-            var asm = Assembly.GetExecutingAssembly();
+            var asm = System.Reflection.Assembly.GetExecutingAssembly();
             var asmDir = Path.GetFullPath(Path.GetDirectoryName(new Uri(asm.GetName().CodeBase!).LocalPath)!);
             if (!asmDir.Equals(currentDir, StringComparison.InvariantCultureIgnoreCase)
                 && !asmDir.Equals(baseDir, StringComparison.InvariantCultureIgnoreCase))

K4AdotNet/Sensor/NativeApi.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Runtime.ConstrainedExecution;
 using System.Runtime.InteropServices;
 
 namespace K4AdotNet.Sensor
@@ -187,13 +188,51 @@ public static extern NativeCallResults.Result ImageCreate(
             int strideBytes,
             out NativeHandles.ImageHandle imageHandle);
 
+        // typedef uint8_t *(k4a_memory_allocate_cb_t)(int size, void **context);
+        /// <summary>Callback function for a memory allocation.</summary>
+        /// <param name="size">Minimum size in bytes needed for the buffer.</param>
+        /// <param name="context">Output parameter for a context that will be provided in the subsequent call to the <see cref="MemoryDestroyCallback"/> callback.</param>
+        /// <returns>A pointer to the newly allocated memory.</returns>
+        [UnmanagedFunctionPointer(CallingConvention.Cdecl)]
+        public delegate IntPtr MemoryAllocateCallback(int size, out IntPtr context);
+
+
         // typedef void(k4a_memory_destroy_cb_t)(void *buffer, void *context);
         /// <summary>Callback function for a memory object being destroyed.</summary>
         /// <param name="buffer">The buffer pointer that was supplied by the caller.</param>
         /// <param name="context">The context for the memory object that needs to be destroyed that was supplied by the caller.</param>
         [UnmanagedFunctionPointer(CallingConvention.Cdecl)]
         public delegate void MemoryDestroyCallback(IntPtr buffer, IntPtr context);
 
+        // K4A_EXPORT k4a_result_t k4a_set_allocator(k4a_memory_allocate_cb_t allocate, k4a_memory_destroy_cb_t free);
+        /// <summary>Sets the callback functions for the SDK allocator</summary>
+        /// <param name="allocate">
+        /// The callback function to allocate memory. When the SDK requires memory allocation this callback will be
+        /// called and the application can provide a buffer and a context.
+        /// </param>
+        /// <param name="free">
+        /// The callback function to free memory.
+        /// The SDK will call this function when memory allocated by <paramref name="allocate"/> is no longer needed.</param>
+        /// <returns>
+        /// <see cref="NativeCallResults.Result.Succeeded"/> if the callback function was set or cleared successfully.
+        /// <see cref="NativeCallResults.Result.Failed"/> if an error is encountered or the callback function has already been set.
+        /// </returns>
+        /// <remarks>
+        /// Call this function to hook memory allocation by the SDK. Calling with both <paramref name="allocate"/> and <paramref name="free"/>
+        /// as <see langword="null"/> will clear the hook and reset to the default allocator.
+        /// 
+        /// If this function is called after memory has been allocated, the previous version of <paramref name="free"/> function may still be
+        /// called in the future. The SDK will always call the <paramref name="free"/> function that was set at the time that the memory
+        /// was allocated.
+        /// 
+        /// Not all memory allocation by the SDK is performed by this allocate function.
+        /// Small allocations or allocations from special pools may come from other sources.
+        /// </remarks>
+        [DllImport(Sdk.SENSOR_DLL_NAME, EntryPoint = "k4a_set_allocator", CallingConvention = CallingConvention.Cdecl)]
+        public static extern NativeCallResults.Result SetAllocator(
+            MemoryAllocateCallback? allocate,
+            MemoryDestroyCallback? free);
+
         // K4A_EXPORT k4a_result_t k4a_image_create_from_buffer(k4a_image_format_t format,
         //                                                      int width_pixels,
         //                                                      int height_pixels,