From 8cd6be8ccca47ede7c9a5e302b47256ebe28a2f6 Mon Sep 17 00:00:00 2001
From: stevenaw <sweerdenburg@gmail.com>
Date: Thu, 11 Apr 2024 10:23:06 -0400
Subject: [PATCH] pool the pinned array and avoid double-allocations

---
 src/Confluent.Kafka/Internal/Util.cs | 40 +++++++++++++++++++++++++---
 1 file changed, 37 insertions(+), 3 deletions(-)

diff --git a/src/Confluent.Kafka/Internal/Util.cs b/src/Confluent.Kafka/Internal/Util.cs
index 9647bbdaa..05819e154 100644
--- a/src/Confluent.Kafka/Internal/Util.cs
+++ b/src/Confluent.Kafka/Internal/Util.cs
@@ -19,6 +19,8 @@
 using SystemMarshal = System.Runtime.InteropServices.Marshal;
 using SystemGCHandle = System.Runtime.InteropServices.GCHandle;
 using SystemGCHandleType = System.Runtime.InteropServices.GCHandleType;
+using System.Buffers;
+using System.Runtime.InteropServices;
 
 
 namespace Confluent.Kafka.Internal
@@ -31,16 +33,45 @@ internal static class Marshal
             ///     Convenience class for generating and pinning the UTF8
             ///     representation of a string.
             /// </summary>
-            public class StringAsPinnedUTF8 : IDisposable
+            public sealed class StringAsPinnedUTF8 : IDisposable
             {
-                private SystemGCHandle gch;
+                private readonly GCHandle gch;
+
+#if NET6_0_OR_GREATER
+                private readonly byte[] strBytesNulTerminated;
+#endif
 
                 public StringAsPinnedUTF8(string str)
                 {
+#if NET6_0_OR_GREATER
+                    var size = Encoding.UTF8.GetMaxByteCount(str.Length);
+
+                    // Ask for one extra for the null byte
+                    strBytesNulTerminated = ArrayPool<byte>.Shared.Rent(size + 1);
+
+                    try
+                    {
+                        Span<byte> slice = strBytesNulTerminated.AsSpan().Slice(0, size);
+                        int bytesWritten = Encoding.UTF8.GetBytes(str, slice);
+
+                        // 0-init the remainder
+                        Array.Clear(strBytesNulTerminated, bytesWritten, strBytesNulTerminated.Length - bytesWritten);
+                    }
+                    catch
+                    {
+                        // An exception above means the object never creates, potentially leaving a
+                        // memory leak as there would be no object to Dispose() of later
+                        ArrayPool<byte>.Shared.Return(strBytesNulTerminated);
+                        throw;
+                    }
+
+                    this.gch = GCHandle.Alloc(strBytesNulTerminated, GCHandleType.Pinned);
+#else
                     byte[] strBytes = System.Text.UTF8Encoding.UTF8.GetBytes(str);
                     byte[] strBytesNulTerminated = new byte[strBytes.Length + 1]; // initialized to all 0's.
                     Array.Copy(strBytes, strBytesNulTerminated, strBytes.Length);
-                    this.gch = SystemGCHandle.Alloc(strBytesNulTerminated, SystemGCHandleType.Pinned);
+                    this.gch = GCHandle.Alloc(strBytesNulTerminated, GCHandleType.Pinned);
+#endif
                 }
 
                 public IntPtr Ptr { get => this.gch.AddrOfPinnedObject(); }
@@ -48,6 +79,9 @@ public StringAsPinnedUTF8(string str)
                 public void Dispose()
                 {
                     gch.Free();
+#if NET6_0_OR_GREATER
+                    ArrayPool<byte>.Shared.Return(strBytesNulTerminated);
+#endif
                 }
             }