Response to comments

DataDog · Sep 15, 2023 · a6c44b7 · a6c44b7
1 parent 51b4a45
commit a6c44b7
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 25 deletions.
diff --git a/tracer/src/Datadog.Trace/AppSec/Waf/Context.cs b/tracer/src/Datadog.Trace/AppSec/Waf/Context.cs
@@ -83,10 +83,16 @@ private Context(IntPtr contextHandle, Waf waf, WafLibraryInvoker wafLibraryInvok
             lock (_stopwatch)
             {
                 var pool = Encoder.Pool;
-                var pwArgs = Encoder.Encode(addresses, applySafetyLimits: true, argToFree: _argCache, pool: pool);
-                code = _waf.Run(_contextHandle, ref pwArgs, ref retNative, timeoutMicroSeconds);
-                pool.Return(_argCache);
-                _argCache.Clear();
+                try
+                {
+                    var pwArgs = Encoder.Encode(addresses, applySafetyLimits: true, argToFree: _argCache, pool: pool);
+                    code = _waf.Run(_contextHandle, ref pwArgs, ref retNative, timeoutMicroSeconds);
+                }
+                finally
+                {
+                    pool.Return(_argCache);
+                    _argCache.Clear();
+                }
             }
 
             _stopwatch.Stop();

diff --git a/tracer/src/Datadog.Trace/Util/UnmanagedMemoryPool.cs b/tracer/src/Datadog.Trace/Util/UnmanagedMemoryPool.cs
@@ -12,6 +12,9 @@
 
 namespace Datadog.Trace.Util;
 
+/// <summary>
+/// Beware that this type is not thread safe and should be used with [ThreadStatic]
+/// </summary>
 internal unsafe class UnmanagedMemoryPool : IDisposable
 {
     private readonly IntPtr* _items;
@@ -40,22 +43,24 @@ public UnmanagedMemoryPool(int blockSize, int poolSize)
 
     public bool IsDisposed => _isDisposed;
 
+    /// <summary>
+    /// Beware that this method is not thread safe, and needs to be used with [ThreadStatic] in case of multiple thread scenarios
+    /// </summary>
+    /// <returns>Pointer to a memory block of size specified in the constructor</returns>
     public IntPtr Rent()
     {
         if (IsDisposed)
         {
             ThrowObjectDisposedException();
         }
 
-        var items = _items;
-        var length = _length;
-        for (var i = _initialSearchIndex; i < length; i++)
+        for (var i = _initialSearchIndex; i < _length; i++)
         {
-            var inst = items[i];
+            var inst = _items[i];
             if (inst != IntPtr.Zero)
             {
                 _initialSearchIndex = i + 1;
-                items[i] = IntPtr.Zero;
+                _items[i] = IntPtr.Zero;
                 return inst;
             }
         }
@@ -75,13 +80,11 @@ public void Return(IntPtr block)
             ThrowObjectDisposedException();
         }
 
-        var items = _items;
-        var length = _length;
-        for (var i = 0; i < length; i++)
+        for (var i = 0; i < _length; i++)
         {
-            if (items[i] == IntPtr.Zero)
+            if (_items[i] == IntPtr.Zero)
             {
-                items[i] = block;
+                _items[i] = block;
                 _initialSearchIndex = 0;
                 return;
             }
@@ -102,14 +105,12 @@ public void Return(IList<IntPtr> blocks)
             return;
         }
 
-        var items = _items;
-        var length = _length;
         var blockIndex = 0;
-        for (var i = 0; i < length; i++)
+        for (var i = 0; i < _length; i++)
         {
-            if (items[i] == IntPtr.Zero)
+            if (_items[i] == IntPtr.Zero)
             {
-                items[i] = blocks[blockIndex++];
+                _items[i] = blocks[blockIndex++];
                 if (blockIndex == blocks.Count)
                 {
                     _initialSearchIndex = 0;
@@ -136,14 +137,12 @@ public void Dispose()
             return;
         }
 
-        var items = _items;
-        var length = _length;
-        for (var i = 0; i < length; i++)
+        for (var i = 0; i < _length; i++)
         {
-            if (items[i] != IntPtr.Zero)
+            if (_items[i] != IntPtr.Zero)
             {
-                Marshal.FreeCoTaskMem(items[i]);
-                items[i] = IntPtr.Zero;
+                Marshal.FreeCoTaskMem(_items[i]);
+                _items[i] = IntPtr.Zero;
             }
         }