From 9418bc40d25c721a8939a3f202a7ac452ab83c25 Mon Sep 17 00:00:00 2001
From: Ben Adams <thundercat@illyriad.co.uk>
Date: Fri, 12 Feb 2016 09:25:33 +0000
Subject: [PATCH] Slimmer QueueUserWorkItemCallback

---
 src/mscorlib/model.xml                          |  3 +
 src/mscorlib/src/System/Threading/ThreadPool.cs | 93 ++++++++++++++++++++-----
 2 files changed, 78 insertions(+), 18 deletions(-)

diff --git a/src/mscorlib/model.xml b/src/mscorlib/model.xml
index 7d51ae8..72ba47f 100644
--- a/src/mscorlib/model.xml
+++ b/src/mscorlib/model.xml
@@ -7837,6 +7837,9 @@
     <Type Status="ImplRoot" Name="System.Threading.QueueUserWorkItemCallback">
       <Member Status="ImplRoot" Name="#ctor(System.Threading.WaitCallback,System.Object,System.Threading.ExecutionContext)"/>
     </Type>
+    <Type Status="ImplRoot" Name="System.Threading.QueueUserWorkItemCallbackDefaultContext">
+      <Member Status="ImplRoot" Name="#ctor(System.Threading.WaitCallback,System.Object)"/>
+    </Type>
     <Type Name="System.Threading.SynchronizationContext">
       <Member Name="#ctor"/>
       <Member Name="Send(System.Threading.SendOrPostCallback,System.Object)"/>
diff --git a/src/mscorlib/src/System/Threading/ThreadPool.cs b/src/mscorlib/src/System/Threading/ThreadPool.cs
index b1ffc77..f9d8e7d 100644
--- a/src/mscorlib/src/System/Threading/ThreadPool.cs
+++ b/src/mscorlib/src/System/Threading/ThreadPool.cs
@@ -1195,22 +1195,6 @@ namespace System.Threading
 #endif
 
         [SecurityCritical]
-        internal QueueUserWorkItemCallback(WaitCallback waitCallback, Object stateObj, bool compressStack, ref StackCrawlMark stackMark)
-        {
-            callback = waitCallback;
-            state = stateObj;
-            if (compressStack && !ExecutionContext.IsFlowSuppressed())
-            {
-                // clone the exection context
-                context = ExecutionContext.Capture(
-                    ref stackMark,
-                    ExecutionContext.CaptureOptions.IgnoreSyncCtx | ExecutionContext.CaptureOptions.OptimizeDefaultCase);
-            }
-        }
-
-        //
-        // internal test hook - used by tests to exercise work-stealing, etc.
-        //
         internal QueueUserWorkItemCallback(WaitCallback waitCallback, Object stateObj, ExecutionContext ec)
         {
             callback = waitCallback;
@@ -1224,7 +1208,6 @@ namespace System.Threading
 #if DEBUG
             MarkExecuted(false);
 #endif
-
             // call directly if it is an unsafe call OR EC flow is suppressed
             if (context == null)
             {
@@ -1261,6 +1244,73 @@ namespace System.Threading
         }
     }
 
+    internal sealed class QueueUserWorkItemCallbackDefaultContext : IThreadPoolWorkItem
+    {
+        [System.Security.SecuritySafeCritical]
+        static QueueUserWorkItemCallbackDefaultContext() { }
+
+        private WaitCallback callback;
+        private Object state;
+
+#if DEBUG
+        private volatile int executed;
+
+        ~QueueUserWorkItemCallbackDefaultContext()
+        {
+            Contract.Assert(
+                executed != 0 || Environment.HasShutdownStarted || AppDomain.CurrentDomain.IsFinalizingForUnload(),
+                "A QueueUserWorkItemCallbackDefaultContext was never called!");
+        }
+
+        void MarkExecuted(bool aborted)
+        {
+            GC.SuppressFinalize(this);
+            Contract.Assert(
+                0 == Interlocked.Exchange(ref executed, 1) || aborted,
+                "A QueueUserWorkItemCallbackDefaultContext was called twice!");
+        }
+#endif
+
+        [SecurityCritical]
+        internal QueueUserWorkItemCallbackDefaultContext(WaitCallback waitCallback, Object stateObj)
+        {
+            callback = waitCallback;
+            state = stateObj;
+        }
+
+        [SecurityCritical]
+        void IThreadPoolWorkItem.ExecuteWorkItem()
+        {
+#if DEBUG
+            MarkExecuted(false);
+#endif
+            ExecutionContext.Run(ExecutionContext.PreAllocatedDefault, ccb, this, true);
+        }
+
+        [SecurityCritical]
+        void IThreadPoolWorkItem.MarkAborted(ThreadAbortException tae)
+        {
+#if DEBUG
+            // this workitem didn't execute because we got a ThreadAbortException prior to the call to ExecuteWorkItem.  
+            // This counts as being executed for our purposes.
+            MarkExecuted(true);
+#endif
+        }
+
+        [System.Security.SecurityCritical]
+        static internal ContextCallback ccb = new ContextCallback(WaitCallback_Context);
+
+        [System.Security.SecurityCritical]
+        static private void WaitCallback_Context(Object state)
+        {
+            QueueUserWorkItemCallbackDefaultContext obj = (QueueUserWorkItemCallbackDefaultContext)state;
+            WaitCallback wc = obj.callback as WaitCallback;
+            Contract.Assert(null != wc);
+            obj.callback = null;
+            wc(obj.state);
+        }
+    }
+
     internal class _ThreadPoolWaitOrTimerCallback
     {
         [System.Security.SecuritySafeCritical]
@@ -1625,7 +1675,14 @@ namespace System.Threading
                 try { }
                 finally
                 {
-                    QueueUserWorkItemCallback tpcallBack = new QueueUserWorkItemCallback(callBack, state, compressStack, ref stackMark);
+                    ExecutionContext context = compressStack && !ExecutionContext.IsFlowSuppressed() ?
+                        ExecutionContext.Capture(ref stackMark, ExecutionContext.CaptureOptions.IgnoreSyncCtx | ExecutionContext.CaptureOptions.OptimizeDefaultCase) :
+                        null;
+
+                    IThreadPoolWorkItem tpcallBack = context == ExecutionContext.PreAllocatedDefault ?
+                                 new QueueUserWorkItemCallbackDefaultContext(callBack, state) :
+                                 (IThreadPoolWorkItem)new QueueUserWorkItemCallback(callBack, state, context);
+
                     ThreadPoolGlobals.workQueue.Enqueue(tpcallBack, true);
                     success = true;
                 }
-- 
2.7.4