Improve object pool safety

Handlebars-Net · Mar 4, 2022 · 300f2f6 · 300f2f6
1 parent d193100
commit 300f2f6
Show file tree

Hide file tree

Showing 13 changed files with 42 additions and 145 deletions.
diff --git a/source/Handlebars/Arguments.cs b/source/Handlebars/Arguments.cs
@@ -322,7 +322,7 @@ public void Dispose()
             private struct Policy : IInternalObjectPoolPolicy<Enumerator>
             {
                 public Enumerator Create() => new Enumerator();
-
+                public bool TryClaim(Enumerator item) => true;
                 public bool Return(Enumerator item)
                 {
                     item.Reset();

diff --git a/source/Handlebars/BindingContext.cs b/source/Handlebars/BindingContext.cs
@@ -17,8 +17,10 @@ public sealed partial class BindingContext : IDisposable, IHelpersRegistry
 
         internal readonly DeferredValue<BindingContext, ObjectDescriptor> Descriptor;
 
-        private BindingContext()
+        private BindingContext(bool claim)
         {
+            _claimed = claim ? 1 : 0;
+
             InlinePartialTemplates = new CascadeIndex<string, Action<EncodedTextWriter, BindingContext>, StringEqualityComparer>(new StringEqualityComparer(StringComparison.OrdinalIgnoreCase));
             Helpers = new CascadeIndex<string, IHelperDescriptor<HelperOptions>, StringEqualityComparer>(new StringEqualityComparer());
             BlockHelpers = new CascadeIndex<string, IHelperDescriptor<BlockHelperOptions>, StringEqualityComparer>(new StringEqualityComparer());

diff --git a/source/Handlebars/Collections/ImmutableStack.cs b/source/Handlebars/Collections/ImmutableStack.cs
diff --git a/source/Handlebars/IO/EncodedTextWriterWrapper.cs b/source/Handlebars/IO/EncodedTextWriterWrapper.cs
@@ -66,6 +66,7 @@ public override void Write(object value)
         private readonly struct Policy : IInternalObjectPoolPolicy<EncodedTextWriterWrapper>
         {
 	        public EncodedTextWriterWrapper Create() => new EncodedTextWriterWrapper();
+	        public bool TryClaim(EncodedTextWriterWrapper item) => true;
 
 	        public bool Return(EncodedTextWriterWrapper obj)
 	        {

diff --git a/source/Handlebars/IO/PolledStringWriter.cs b/source/Handlebars/IO/PolledStringWriter.cs
@@ -40,7 +40,7 @@ public Policy(int initialCapacity, int maximumRetainedCapacity = 4096)
             }
 
             public ReusableStringWriter Create() => new ReusableStringWriter();
-
+            public bool TryClaim(ReusableStringWriter item) => true;
             public bool Return(ReusableStringWriter item) => _policy.Return(item.GetStringBuilder());
         }
     }

diff --git a/source/Handlebars/IO/TextEncoderWrapper.cs b/source/Handlebars/IO/TextEncoderWrapper.cs
@@ -66,7 +66,7 @@ public void Dispose()
         private struct Policy : IInternalObjectPoolPolicy<TextEncoderWrapper>
         {
             public TextEncoderWrapper Create() => new TextEncoderWrapper();
-
+            public bool TryClaim(TextEncoderWrapper item) => true;
             public bool Return(TextEncoderWrapper item)
             {
                 item._enabled = true;

diff --git a/source/Handlebars/Polyfills/AsyncLocal.cs b/source/Handlebars/Polyfills/AsyncLocal.cs
diff --git a/source/Handlebars/Pools/BindingContext.Pool.cs b/source/Handlebars/Pools/BindingContext.Pool.cs
@@ -1,3 +1,4 @@
+using System.Threading;
 using HandlebarsDotNet.Compiler;
 using HandlebarsDotNet.Pools;
 
@@ -17,6 +18,8 @@ internal static BindingContext Create(ICompiledHandlebarsConfiguration configura
         {
             return Pool.CreateContext(configuration, value, parent, partialBlockTemplate);
         }
+
+        private int _claimed;
 
         public void Dispose() => Pool.Return(this);
 
@@ -39,12 +42,23 @@ public BindingContext CreateContext(ICompiledHandlebarsConfiguration configurati
                 return context;
             }
 
-            internal struct BindingContextPolicy : IInternalObjectPoolPolicy<BindingContext>
+            internal readonly struct BindingContextPolicy : IInternalObjectPoolPolicy<BindingContext>
             {
-                public BindingContext Create() => new BindingContext();
+                public BindingContext Create() => new(claim: true);
+
+                public bool TryClaim(BindingContext item)
+                {
+                    return Interlocked
+                        .CompareExchange(ref item._claimed, 1, 0) == 0;
+                }
 
                 public bool Return(BindingContext item)
                 {
+                    var shouldReturn = Interlocked
+                        .CompareExchange(ref item._claimed, 0, 1) == 1;
+
+                    if (!shouldReturn) return false;
+
                     item.Root = null;
                     item.Value = null;
                     item.ParentContext = null;

diff --git a/source/Handlebars/Pools/ClosureBuilder.Pool.cs b/source/Handlebars/Pools/ClosureBuilder.Pool.cs
@@ -36,6 +36,7 @@ public ClosureBuilderPool(Policy policy) : base(policy)
         private readonly struct Policy : IInternalObjectPoolPolicy<ClosureBuilder>
         {
             public ClosureBuilder Create() => new ();
+            public bool TryClaim(ClosureBuilder item) => true;
             public bool Return(ClosureBuilder item) => true;
         }
     }

diff --git a/source/Handlebars/Pools/GenericObjectPool.cs b/source/Handlebars/Pools/GenericObjectPool.cs
@@ -10,6 +10,7 @@ private GenericObjectPool() : base(new Policy()) { }
         public readonly struct Policy : IInternalObjectPoolPolicy<T>
         {
             public T Create() => new ();
+            public bool TryClaim(T item) => true;
             public bool Return(T item) => true;
         }
     }

diff --git a/source/Handlebars/Pools/ObjectPool.cs b/source/Handlebars/Pools/ObjectPool.cs
@@ -18,11 +18,11 @@ public InternalObjectPool(TPolicy policy)
 
         public T Get()
         {
-            if (_queue.TryDequeue(out var item))
+            if (_queue.TryDequeue(out var item) && _policy.TryClaim(item))
             {
                 return item;
             }
-    
+
             return _policy.Create();
         }
 
@@ -36,6 +36,7 @@ public void Return(T obj)
     internal interface IInternalObjectPoolPolicy<T>
     {
         T Create();
+        bool TryClaim(T item);
         bool Return(T item);
     }
 }
diff --git a/source/Handlebars/Pools/StringBuilderPool.cs b/source/Handlebars/Pools/StringBuilderPool.cs
@@ -31,6 +31,8 @@ public StringBuilder Create()
                 return new StringBuilder(InitialCapacity);
             }
 
+            public bool TryClaim(StringBuilder item) => true;
+
             public bool Return(StringBuilder item)
             {
                 if (item.Capacity > MaximumRetainedCapacity)

diff --git a/source/Handlebars/Runtime/AmbientContext.cs b/source/Handlebars/Runtime/AmbientContext.cs
@@ -1,12 +1,6 @@
 using System;
 using System.Collections.Generic;
 using System.Runtime.CompilerServices;
-#if !NET451 && !NET452
-using System.Threading;
-#else
-using HandlebarsDotNet.Polyfills;
-#endif
-using HandlebarsDotNet.Collections;
 using HandlebarsDotNet.IO;
 using HandlebarsDotNet.ObjectDescriptors;
 using HandlebarsDotNet.PathStructure;
@@ -16,14 +10,18 @@ namespace HandlebarsDotNet.Runtime
 {
     public sealed class AmbientContext : IDisposable
     {
-        private static readonly InternalObjectPool<AmbientContext, Policy> Pool = new InternalObjectPool<AmbientContext, Policy>(new Policy()); 
+        private static readonly InternalObjectPool<AmbientContext, Policy> Pool = new(new Policy()); 
 
-        private static readonly AsyncLocal<ImmutableStack<AmbientContext>> Local = new AsyncLocal<ImmutableStack<AmbientContext>>();
+        [ThreadStatic]
+        private static Stack<AmbientContext> _local;
+
+        private static Stack<AmbientContext> Local => 
+            _local ??= new Stack<AmbientContext>();
 
         public static AmbientContext Current
         {
             [MethodImpl(MethodImplOptions.AggressiveInlining)]
-            get => Local.Value.Peek();
+            get => Local.Count > 0 ? Local.Peek() : null;
         }
 
         public static AmbientContext Create(
@@ -67,12 +65,9 @@ public static AmbientContext Create(
 
         public static DisposableContainer Use(AmbientContext ambientContext)
         {
-            Local.Value = Local.Value.Push(ambientContext);
+            Local.Push(ambientContext);
 
-            return new DisposableContainer(() =>
-            {
-                Local.Value = Local.Value.Pop(out _);
-            });
+            return new DisposableContainer(() => Local.Pop());
         }
 
         private AmbientContext()
@@ -93,10 +88,9 @@ private AmbientContext()
 
         private struct Policy : IInternalObjectPoolPolicy<AmbientContext>
         {
-            public AmbientContext Create()
-            {
-                return new AmbientContext();
-            }
+            public AmbientContext Create() => new();
+
+            public bool TryClaim(AmbientContext item) => true;
 
             public bool Return(AmbientContext item)
             {
-Original file line number
+Diff line change
@@ Expand Up @@
                 }
                 public ReusableStringWriter Create() => new ReusableStringWriter();
+                public bool TryClaim(ReusableStringWriter item) => true;
                 public bool Return(ReusableStringWriter item) => _policy.Return(item.GetStringBuilder());
             }
         }
@@ Expand Down @@