[release/3.1] Capture AVX context when redirecting threads for suspension

src/vm/threads.cpp

@@ -57,7 +57,10 @@
 uint64_t Thread::dead_threads_non_alloc_bytes = 0;
 
 SPTR_IMPL(ThreadStore, ThreadStore, s_pThreadStore);
-CONTEXT *ThreadStore::s_pOSContext = NULL;
+
+CONTEXT* ThreadStore::s_pOSContext = NULL;
+BYTE* ThreadStore::s_pOSContextBuffer = NULL;
+
 CLREvent *ThreadStore::s_pWaitForStackCrawlEvent;
 
 PTR_ThreadLocalModule ThreadLocalBlock::GetTLMIfExists(ModuleIndex index)
@@ -1452,7 +1455,6 @@ Thread::Thread()
     m_fHasDeadThreadBeenConsideredForGCTrigger = false;
     m_TraceCallCount = 0;
     m_ThrewControlForThread = 0;
-    m_OSContext = NULL;
     m_ThreadTasks = (ThreadTasks)0;
     m_pLoadLimiter= NULL;
     m_pLoadingFile = NULL;
@@ -1487,7 +1489,11 @@ Thread::Thread()
     NewHolder<CONTEXT> contextHolder(m_OSContext);
 
     m_pSavedRedirectContext = NULL;
-    NewHolder<CONTEXT> savedRedirectContextHolder(m_pSavedRedirectContext);
+    m_pOSContextBuffer = NULL;
+
+#ifdef _DEBUG
+    m_RedirectContextInUse = false;
+#endif
 
 #ifdef FEATURE_COMINTEROP
     m_pRCWStack = new RCWStackHeader();
@@ -1563,7 +1569,6 @@ Thread::Thread()
     trackSyncHolder.SuppressRelease();
 #endif
     contextHolder.SuppressRelease();
-    savedRedirectContextHolder.SuppressRelease();
 
 #ifdef FEATURE_COMINTEROP
     m_uliInitializeSpyCookie.QuadPart = 0ul;
@@ -2640,11 +2645,18 @@ Thread::~Thread()
     if (m_OSContext)
         delete m_OSContext;
 
-    if (GetSavedRedirectContext())
+    if (m_pOSContextBuffer)
     {
-        delete GetSavedRedirectContext();
-        SetSavedRedirectContext(NULL);
+        delete[] m_pOSContextBuffer;
+        m_pOSContextBuffer = NULL;
     }
+    else if (m_pSavedRedirectContext)
+    {
+        delete m_pSavedRedirectContext;
+    }
+
+    MarkRedirectContextInUse(m_pSavedRedirectContext);
+    m_pSavedRedirectContext = NULL;
 
 #ifdef FEATURE_COMINTEROP
     if (m_pRCWStack)

src/vm/threads.h

@@ -4339,8 +4339,21 @@ class Thread: public IUnknown
 #endif // _DEBUG
 
 private:
+    // context used during redirection of this thread
+    // NOTE: there is only one. Since redirection cannot be nested
+    //       if more than one are needed, something is wrong.
     PTR_CONTEXT m_pSavedRedirectContext;
 
+    // in a case when we need the redirection context to include CONTEXT_XSTATE
+    // this is the buffer that contains the context parts.
+    // we need the buffer so we could deallocate the whole deal.
+    BYTE* m_pOSContextBuffer;
+
+#ifdef _DEBUG
+    // validate that we use only one context per thread. 
+    bool m_RedirectContextInUse;
+#endif
+
     BOOL IsContextSafeToRedirect(T_CONTEXT* pContext);
 
 public:
@@ -4351,14 +4364,26 @@ class Thread: public IUnknown
     }
 
 #ifndef DACCESS_COMPILE
-    void     SetSavedRedirectContext(PT_CONTEXT pCtx)
+    void MarkRedirectContextInUse(PTR_CONTEXT pCtx)
     {
         LIMITED_METHOD_CONTRACT;
-        m_pSavedRedirectContext = pCtx;
-    }
+#ifdef _DEBUG
+        _ASSERTE(!m_RedirectContextInUse);
+        _ASSERTE(pCtx == m_pSavedRedirectContext);
+        m_RedirectContextInUse = true;
 #endif
+    }
 
-    void EnsurePreallocatedContext();
+    void UnmarkRedirectContextInUse(PTR_CONTEXT pCtx)
+    {
+        LIMITED_METHOD_CONTRACT;
+#ifdef _DEBUG
+        _ASSERTE(m_RedirectContextInUse);
+        _ASSERTE(pCtx == m_pSavedRedirectContext);
+        m_RedirectContextInUse = false;
+#endif
+    }
+#endif //DACCESS_COMPILE
 
     ThreadLocalBlock m_ThreadLocalBlock;
 
@@ -5290,12 +5315,21 @@ class ThreadStore
 
 #endif
 private:
+    static BYTE* s_pOSContextBuffer;
     static CONTEXT *s_pOSContext;
 public:
-    // We can not do any memory allocation after we suspend a thread in order ot
-    // avoid deadlock situation.
+    // Pre-allocate an OS context for possible use by a redirected thread and keep in a static variable.
+    // 
+    // There are two reasons for this pattern:
+    // - We can not do any memory allocation after we suspend a thread in order to avoid deadlock situation.
+    //   So, when anticipating a need, we must pre-allocate.
+    // 
+    // - Even though we know the thread we are suspending, we do not want to put the context directly on the
+    //   thread because the thread only _may_ need the context. Often it does not end up needing it,
+    //   then we will keep the context for the next time like this.
     static void AllocateOSContext();
-    static CONTEXT *GrabOSContext();
+    // Retrieves and detaches the pre-alocated context + optional containing buffer (when CONTEXT_XSTATE is used)
+    static CONTEXT* GrabOSContext(BYTE** contextBuffer);
 
 private:
     // Thread abort needs to walk stack to decide if thread abort can proceed.