[MERGE #5099 @VSadov] Adding a configuration flag for setting the allocation policy limit.

Merge pull request #5099 from VSadov:allocLimit

The flag is useful for testing purposes and for setting the policy limit in non-JSRT scenarios.
- Fixed some missing `ReportExternalMemoryFree` calls.
Not having those was tripping consistency asserts in memory policy manager in debug and could potentially result in false OOMs in release.

Fixes:#5031

Author: Cellule

lib/Common/ConfigFlagsList.h

@@ -535,6 +535,7 @@ PHASE(All)
 #define DEFAULT_CONFIG_LowMemoryCap         (0xB900000) // 185 MB - based on memory cap for process on low-capacity device
 #define DEFAULT_CONFIG_NewPagesCapDuringBGSweeping    (15000 * 4)
 #define DEFAULT_CONFIG_MaxSingleAllocSizeInMB  (2048)
+#define DEFAULT_CONFIG_AllocationPolicyLimit    (-1)
 
 #define DEFAULT_CONFIG_MaxCodeFill          (500)
 #define DEFAULT_CONFIG_MaxLoopsPerFunction  (10)
@@ -1448,6 +1449,7 @@ FLAGNR(Boolean, RecyclerVerifyMark    , "verify concurrent gc", false)
 #endif
 FLAGR (Number,  LowMemoryCap          , "Memory cap indicating a low-memory process", DEFAULT_CONFIG_LowMemoryCap)
 FLAGNR(Number,  NewPagesCapDuringBGSweeping, "New pages count allowed to be allocated during background sweeping", DEFAULT_CONFIG_NewPagesCapDuringBGSweeping)
+FLAGR (Number,  AllocPolicyLimit      , "Memory allocation policy limit in MB (default: -1, which means no allocation policy limit).", DEFAULT_CONFIG_AllocationPolicyLimit)
 #ifdef RUNTIME_DATA_COLLECTION
 FLAGNR(String,  RuntimeDataOutputFile, "Filename to write the dynamic profile info", nullptr)
 #endif

lib/Common/Memory/AllocationPolicyManager.h

@@ -41,6 +41,11 @@ typedef bool (__stdcall * PageAllocatorMemoryAllocationCallback)(__in LPVOID con
         context(NULL),
         memoryAllocationCallback(NULL)
     {
+        Js::Number limitMB = Js::Configuration::Global.flags.AllocPolicyLimit;
+        if (limitMB > 0)
+        {
+            memoryLimit = (size_t)limitMB * 1024 * 1024;
+        }
     }
 
     ~AllocationPolicyManager()

lib/Common/Memory/Recycler.h

@@ -2648,7 +2648,6 @@ bool Recycler::DoExternalAllocation(size_t size, ExternalAllocFunc externalAlloc
     AutoExternalAllocation externalAllocation(this, size);
     if (externalAllocFunc())
     {
-        this->AddExternalMemoryUsage(size);
         externalAllocation.allocationSucceeded = true;
         return true;
     }

lib/Runtime/Base/ThreadBoundThreadContextManager.cpp

@@ -33,11 +33,15 @@ ThreadContext * ThreadBoundThreadContextManager::EnsureContextForCurrentThread()
     // Just reinitialize the thread context.
     if (threadContext == nullptr)
     {
-        threadContext = HeapNew(ThreadContext);
+        bool requireConcurrencySupport = true;
+        AllocationPolicyManager * policyManager = HeapNew(AllocationPolicyManager, requireConcurrencySupport);
+        threadContext = HeapNew(ThreadContext, policyManager);
+
         threadContext->SetIsThreadBound();
         if (!ThreadContextTLSEntry::TrySetThreadContext(threadContext))
         {
             HeapDelete(threadContext);
+            HeapDelete(policyManager);
             return NULL;
         }
     }
@@ -181,17 +185,7 @@ void ThreadBoundThreadContextManager::DestroyAllContextsAndEntries(bool shouldDe
 
         if (threadContext != nullptr)
         {
-#if DBG
-            PageAllocator* pageAllocator = threadContext->GetPageAllocator();
-            if (pageAllocator)
-            {
-                pageAllocator->SetConcurrentThreadId(::GetCurrentThreadId());
-            }
-#endif
-
-            threadContext->ShutdownThreads();
-
-            HeapDelete(threadContext);
+            ShutdownThreadContext(threadContext);
         }
 
         if (currentThreadEntry != entry)
@@ -264,6 +258,12 @@ void ThreadContextManagerBase::ShutdownThreadContext(
 
     if (deleteThreadContext)
     {
+        AllocationPolicyManager * policyManager = threadContext->IsThreadBound() ? threadContext->GetAllocationPolicyManager() : nullptr;
         HeapDelete(threadContext);
+
+        if (policyManager)
+        {
+            HeapDelete(policyManager);
+        }
     }
 }

lib/Runtime/Library/ArrayBuffer.cpp

@@ -218,6 +218,11 @@ namespace Js
     {
         Assert(!this->isDetached);
 
+        // we are about to lose track of the buffer to another owner
+        // report that we no longer own the memory
+        Recycler* recycler = GetType()->GetLibrary()->GetRecycler();
+        recycler->ReportExternalMemoryFree(bufferLength);
+
         this->buffer = nullptr;
         this->bufferLength = 0;
         this->isDetached = true;
@@ -696,6 +701,13 @@ namespace Js
         Recycler* recycler = type->GetScriptContext()->GetRecycler();
         JavascriptArrayBuffer* result = RecyclerNewFinalized(recycler, JavascriptArrayBuffer, buffer, length, type);
         Assert(result);
+
+        // we take the ownership of the buffer and will have to free it so charge it to our quota.
+        if (!recycler->RequestExternalMemoryAllocation(length))
+        {
+            JavascriptError::ThrowOutOfMemoryError(result->GetScriptContext());
+        }
+
         recycler->AddExternalMemoryUsage(length);
         return result;
     }
@@ -707,7 +719,6 @@ namespace Js
         {
             freeFunction(this->buffer);
             this->buffer = nullptr;
-            this->recycler->ReportExternalMemoryFree(this->bufferLength);
         }
         this->bufferLength = 0;
     }
@@ -770,12 +781,6 @@ namespace Js
 
     void JavascriptArrayBuffer::Finalize(bool isShutdown)
     {
-        // In debugger scenario, ScriptAuthor can create scriptContext and delete scriptContext
-        // explicitly. So for the builtin, while javascriptLibrary is still alive fine, the
-        // matching scriptContext might have been deleted and the javascriptLibrary->scriptContext
-        // field reset (but javascriptLibrary is still alive).
-        // Use the recycler field off library instead of scriptcontext to avoid av.
-
         // Recycler may not be available at Dispose. We need to
         // free the memory and report that it has been freed at the same
         // time. Otherwise, AllocationPolicyManager is unable to provide correct feedback
@@ -885,6 +890,12 @@ namespace Js
         if (buffer)
         {
             result = RecyclerNewFinalized(recycler, WebAssemblyArrayBuffer, buffer, length, type);
+
+            // we take the ownership of the buffer and will have to free it so charge it to our quota.
+            if (!recycler->RequestExternalMemoryAllocation(length))
+            {
+                JavascriptError::ThrowOutOfMemoryError(result->GetScriptContext());
+            }
         }
         else
         {
@@ -898,9 +909,9 @@ namespace Js
             {
                 result = RecyclerNewFinalized(recycler, WebAssemblyArrayBuffer, length, type);
             }
-            // Only add external memory when we create a new internal buffer
-            recycler->AddExternalMemoryUsage(length);
         }
+
+        recycler->AddExternalMemoryUsage(length);
         Assert(result);
         return result;
     }
@@ -955,6 +966,11 @@ namespace Js
             {
                 return nullptr;
             }
+
+            // We are transferring the buffer to the new owner. 
+            // To avoid double-charge to the allocation quota we will free the "diff" amount here.
+            this->GetRecycler()->ReportExternalMemoryFree(growSize);
+
             return finalizeGrowMemory(this->GetLibrary()->CreateWebAssemblyArrayBuffer(this->buffer, newBufferLength));
         }
 #endif
@@ -990,6 +1006,10 @@ namespace Js
                 return nullptr;
             }
 
+            // We are transferring the buffer to the new owner. 
+            // To avoid double-charge to the allocation quota we will free the "diff" amount here.
+            this->GetRecycler()->ReportExternalMemoryFree(growSize);
+
             WebAssemblyArrayBuffer* newArrayBuffer = finalizeGrowMemory(this->GetLibrary()->CreateWebAssemblyArrayBuffer(newBuffer, newBufferLength));
             // We've successfully Detached this buffer and created a new WebAssemblyArrayBuffer
             autoDisableInterrupt.Completed();
@@ -1018,15 +1038,37 @@ namespace Js
     ProjectionArrayBuffer* ProjectionArrayBuffer::Create(byte* buffer, uint32 length, DynamicType * type)
     {
         Recycler* recycler = type->GetScriptContext()->GetRecycler();
+        ProjectionArrayBuffer* result = RecyclerNewFinalized(recycler, ProjectionArrayBuffer, buffer, length, type);
+
+        // we take the ownership of the buffer and will have to free it so charge it to our quota.
+        if (!recycler->RequestExternalMemoryAllocation(length))
+        {
+            JavascriptError::ThrowOutOfMemoryError(result->GetScriptContext());
+        }
+
         // This is user passed [in] buffer, user should AddExternalMemoryUsage before calling jscript, but
         // I don't see we ask everyone to do this. Let's add the memory pressure here as well.
         recycler->AddExternalMemoryUsage(length);
-        return RecyclerNewFinalized(recycler, ProjectionArrayBuffer, buffer, length, type);
+        return result;
+
     }
 
-    void ProjectionArrayBuffer::Dispose(bool isShutdown)
+    void ProjectionArrayBuffer::Finalize(bool isShutdown)
     {
         CoTaskMemFree(buffer);
+        // Recycler may not be available at Dispose. We need to
+        // free the memory and report that it has been freed at the same
+        // time. Otherwise, AllocationPolicyManager is unable to provide correct feedback
+        Recycler* recycler = GetType()->GetLibrary()->GetRecycler();
+        recycler->ReportExternalMemoryFree(bufferLength);
+
+        buffer = nullptr;
+        bufferLength = 0;
+    }
+
+    void ProjectionArrayBuffer::Dispose(bool isShutdown)
+    {
+        /* See ProjectionArrayBuffer::Finalize */
     }
 
     ExternalArrayBuffer::ExternalArrayBuffer(byte *buffer, uint32 length, DynamicType *type)

lib/Runtime/Library/ArrayBuffer.h

@@ -325,7 +325,7 @@ namespace Js
         // take over ownership. a CoTaskMemAlloc'ed buffer passed in via projection.
         static ProjectionArrayBuffer* Create(byte* buffer, DECLSPEC_GUARD_OVERFLOW uint32 length, DynamicType * type);
         virtual void Dispose(bool isShutdown) override;
-        virtual void Finalize(bool isShutdown) override {};
+        virtual void Finalize(bool isShutdown) override;
     private:
         ProjectionArrayBuffer(uint32 length, DynamicType * type);
         ProjectionArrayBuffer(byte* buffer, uint32 length, DynamicType * type);

lib/Runtime/Library/SharedArrayBuffer.cpp

@@ -617,6 +617,8 @@ namespace Js
             {
                 return false;
             }
+
+            this->GetRecycler()->AddExternalMemoryUsage(growSize);
         }
         else
         {