[1.6>1.7] [MERGE #3468 @Cellule] Strengthen WebAssemblyArrayBuffer.Gr…

…owMemory

Merge pull request #3468 from Cellule:users/micfer/wasm/grow_memory

Make reporting external memory allocation cleaner with a wrapper and ensure we don't end up in a bad state after ReAlloc

lib/Common/Memory/Recycler.cpp

@@ -1104,7 +1104,7 @@ Recycler::AddExternalMemoryUsage(size_t size)
     CollectNow<CollectOnAllocation>();
 }
 
-BOOL Recycler::ReportExternalMemoryAllocation(size_t size)
+bool Recycler::RequestExternalMemoryAllocation(size_t size)
 {
     return recyclerPageAllocator.RequestAlloc(size);
 }

lib/Common/Memory/Recycler.h

@@ -1202,9 +1202,12 @@ class Recycler
 
     template <CollectionFlags flags>
     bool FinishDisposeObjectsNow();
-    BOOL ReportExternalMemoryAllocation(size_t size);
+    bool RequestExternalMemoryAllocation(size_t size);
     void ReportExternalMemoryFailure(size_t size);
     void ReportExternalMemoryFree(size_t size);
+    // ExternalAllocFunc returns true when allocation succeeds
+    template <typename ExternalAllocFunc>
+    bool DoExternalAllocation(size_t size, ExternalAllocFunc externalAllocFunc);
 
 #ifdef TRACE_OBJECT_LIFETIME
 #define DEFINE_RECYCLER_ALLOC_TRACE(AllocFunc, AllocWithAttributesFunc, attributes) \
@@ -2532,3 +2535,35 @@ extern bool IsLikelyRuntimeFalseReference(
 #else
 #define DECLARE_RECYCLER_VERIFY_MARK_FRIEND()
 #endif
+
+template <typename ExternalAllocFunc>
+bool Recycler::DoExternalAllocation(size_t size, ExternalAllocFunc externalAllocFunc)
+{
+    // Request external memory allocation
+    if (!RequestExternalMemoryAllocation(size))
+    {
+        // Attempt to free some memory then try again
+        CollectNow<CollectOnTypedArrayAllocation>();
+        if (!RequestExternalMemoryAllocation(size))
+        {
+            return false;
+        }
+    }
+    struct AutoExternalAllocation
+    {
+        bool allocationSucceeded = false;
+        Recycler* recycler;
+        size_t size;
+        AutoExternalAllocation(Recycler* recycler, size_t size): recycler(recycler), size(size) {}
+        // In case the externalAllocFunc throws or fails, the destructor will report the failure
+        ~AutoExternalAllocation() { if (!allocationSucceeded) recycler->ReportExternalMemoryFailure(size); }
+    };
+    AutoExternalAllocation externalAllocation(this, size);
+    if (externalAllocFunc())
+    {
+        this->AddExternalMemoryUsage(size);
+        externalAllocation.allocationSucceeded = true;
+        return true;
+    }
+    return false;
+}

lib/Runtime/Base/ThreadContext.h

@@ -1802,7 +1802,7 @@ class AutoDisableInterrupt
             AssertOrFailFast(false);
         }
     }
-
+    void RequireExplicitCompletion() { m_explicitCompletion = true; }
     void Completed() { m_operationCompleted = true; }
 
 private:

lib/Runtime/Library/ArrayBuffer.cpp

@@ -208,12 +208,10 @@ namespace Js
         }
     }
 
-    ArrayBufferDetachedStateBase* ArrayBuffer::DetachAndGetState()
+    void ArrayBuffer::Detach()
     {
         Assert(!this->isDetached);
 
-        AutoPtr<ArrayBufferDetachedStateBase> arrayBufferState(this->CreateDetachedState(this->buffer, this->bufferLength));
-
         this->buffer = nullptr;
         this->bufferLength = 0;
         this->isDetached = true;
@@ -235,7 +233,13 @@ namespace Js
                 this->DetachBufferFromParent(item->Get());
             });
         }
+    }
 
+    ArrayBufferDetachedStateBase* ArrayBuffer::DetachAndGetState()
+    {
+        // Save the state before detaching
+        AutoPtr<ArrayBufferDetachedStateBase> arrayBufferState(this->CreateDetachedState(this->buffer, this->bufferLength));
+        Detach();
         return arrayBufferState.Detach();
     }
 
@@ -594,7 +598,7 @@ namespace Js
         else if (length > 0)
         {
             Recycler* recycler = GetType()->GetLibrary()->GetRecycler();
-            if (recycler->ReportExternalMemoryAllocation(length))
+            if (recycler->RequestExternalMemoryAllocation(length))
             {
                 buffer = (BYTE*)allocator(length);
                 if (buffer == nullptr)
@@ -607,7 +611,7 @@ namespace Js
             {
                 recycler->CollectNow<CollectOnTypedArrayAllocation>();
 
-                if (recycler->ReportExternalMemoryAllocation(length))
+                if (recycler->RequestExternalMemoryAllocation(length))
                 {
                     buffer = (BYTE*)allocator(length);
                     if (buffer == nullptr)
@@ -909,10 +913,10 @@ namespace Js
             // Expanding buffer
             if (newBufferLength > this->bufferLength)
             {
-                if (!recycler->ReportExternalMemoryAllocation(newBufferLength - this->bufferLength))
+                if (!recycler->RequestExternalMemoryAllocation(newBufferLength - this->bufferLength))
                 {
                     recycler->CollectNow<CollectOnTypedArrayAllocation>();
-                    if (!recycler->ReportExternalMemoryAllocation(newBufferLength - this->bufferLength))
+                    if (!recycler->RequestExternalMemoryAllocation(newBufferLength - this->bufferLength))
                     {
                         reportFailureFn();
                     }
@@ -1001,7 +1005,7 @@ namespace Js
     WebAssemblyArrayBuffer* WebAssemblyArrayBuffer::Create(byte* buffer, uint32 length, DynamicType * type)
     {
         Recycler* recycler = type->GetScriptContext()->GetRecycler();
-        WebAssemblyArrayBuffer* result;
+        WebAssemblyArrayBuffer* result = nullptr;
         if (buffer)
         {
             result = RecyclerNewFinalized(recycler, WebAssemblyArrayBuffer, buffer, length, type);
@@ -1018,9 +1022,10 @@ namespace Js
             {
                 result = RecyclerNewFinalized(recycler, WebAssemblyArrayBuffer, length, type);
             }
+            // Only add external memory when we create a new internal buffer
+            recycler->AddExternalMemoryUsage(length);
         }
         Assert(result);
-        recycler->AddExternalMemoryUsage(length);
         return result;
     }
 
@@ -1040,71 +1045,75 @@ namespace Js
             Assert(UNREACHED);
             JavascriptError::ThrowTypeError(GetScriptContext(), WASMERR_BufferGrowOnly);
         }
-        uint32 growSize = newBufferLength - this->bufferLength;
 
-        bool failedReport = false;
-        const auto reportFailedFn = [&failedReport] { failedReport = true; };
+        uint32 growSize = newBufferLength - this->bufferLength;
+        const auto finalizeGrowMemory = [&](WebAssemblyArrayBuffer* newArrayBuffer)
+        {
+            AssertOrFailFast(newArrayBuffer && newArrayBuffer->GetByteLength() == newBufferLength);
+            // Detach the buffer from this ArrayBuffer
+            this->Detach();
+            return newArrayBuffer;
+        };
+
+        // We're not growing the buffer, just create a new WebAssemblyArrayBuffer and detach this
+        if (growSize == 0)
+        {
+            return finalizeGrowMemory(this->GetLibrary()->CreateWebAssemblyArrayBuffer(this->buffer, this->bufferLength));
+        }
 
-        WebAssemblyArrayBuffer* newArrayBuffer = nullptr;
 #if ENABLE_FAST_ARRAYBUFFER
+        // 8Gb Array case
         if (CONFIG_FLAG(WasmFastArray))
         {
             AssertOrFailFast(this->buffer);
-            ReportDifferentialAllocation(newBufferLength, reportFailedFn);
-            if (failedReport)
+            const auto virtualAllocFunc = [&]
             {
-                return nullptr;
-            }
-
-            if (growSize > 0)
+                return !!VirtualAlloc(this->buffer + this->bufferLength, growSize, MEM_COMMIT, PAGE_READWRITE);
+            };
+            if (!this->GetRecycler()->DoExternalAllocation(growSize, virtualAllocFunc))
             {
-                LPVOID newMem = VirtualAlloc(this->buffer + this->bufferLength, growSize, MEM_COMMIT, PAGE_READWRITE);
-                if (!newMem)
-                {
-                    Recycler* recycler = this->GetRecycler();
-                    recycler->ReportExternalMemoryFailure(newBufferLength);
-                    return nullptr;
-                }
+                return nullptr;
             }
-            newArrayBuffer = GetLibrary()->CreateWebAssemblyArrayBuffer(this->buffer, newBufferLength);
+            return finalizeGrowMemory(this->GetLibrary()->CreateWebAssemblyArrayBuffer(this->buffer, newBufferLength));
         }
-        else
 #endif
+
+        // No previous buffer case
         if (this->GetByteLength() == 0)
         {
-            if (growSize > 0)
-            {
-                newArrayBuffer = GetLibrary()->CreateWebAssemblyArrayBuffer(newBufferLength);
-            }
-            else
-            {
-                newArrayBuffer = GetLibrary()->CreateWebAssemblyArrayBuffer(this->buffer, 0);
-            }
+            Assert(newBufferLength == growSize);
+            // Creating a new buffer will do the external memory allocation report
+            return finalizeGrowMemory(this->GetLibrary()->CreateWebAssemblyArrayBuffer(newBufferLength));
         }
-        else
+
+        // Regular growing case
         {
-            ReportDifferentialAllocation(newBufferLength, reportFailedFn);
-            if (failedReport)
+            // Disable Interrupts while doing a ReAlloc to minimize chances to end up in a bad state
+            AutoDisableInterrupt autoDisableInterrupt(this->GetScriptContext()->GetThreadContext(), false);
+
+            byte* newBuffer = nullptr;
+            const auto reallocFunc = [&]
             {
-                return nullptr;
-            }
-            byte* newBuffer = ReallocZero(this->buffer, this->bufferLength, newBufferLength);
-            if (!newBuffer)
+                newBuffer = ReallocZero(this->buffer, this->bufferLength, newBufferLength);
+                if (newBuffer != nullptr)
+                {
+                    // Realloc freed this->buffer
+                    // if anything goes wrong before we detach, we can't recover the state and should failfast
+                    autoDisableInterrupt.RequireExplicitCompletion();
+                }
+                return !!newBuffer;
+            };
+
+            if (!this->GetRecycler()->DoExternalAllocation(growSize, reallocFunc))
             {
-                this->GetRecycler()->ReportExternalMemoryFailure(newBufferLength - this->bufferLength);
                 return nullptr;
             }
-            newArrayBuffer = GetLibrary()->CreateWebAssemblyArrayBuffer(newBuffer, newBufferLength);
-        }
 
-        if (!newArrayBuffer || newArrayBuffer->GetByteLength() != newBufferLength)
-        {
-            return nullptr;
+            WebAssemblyArrayBuffer* newArrayBuffer = finalizeGrowMemory(this->GetLibrary()->CreateWebAssemblyArrayBuffer(newBuffer, newBufferLength));
+            // We've successfully Detached this buffer and created a new WebAssemblyArrayBuffer
+            autoDisableInterrupt.Completed();
+            return newArrayBuffer;
         }
-
-        AutoDiscardPTR<Js::ArrayBufferDetachedStateBase> state(DetachAndGetState());
-        state->MarkAsClaimed();
-        return newArrayBuffer;
     }
 
     ArrayBuffer * WebAssemblyArrayBuffer::TransferInternal(uint32 newBufferLength)

lib/Runtime/Library/ArrayBuffer.h

@@ -159,7 +159,7 @@ namespace Js
         virtual BOOL GetDiagTypeString(StringBuilder<ArenaAllocator>* stringBuilder, ScriptContext* requestContext) override;
         virtual BOOL GetDiagValueString(StringBuilder<ArenaAllocator>* stringBuilder, ScriptContext* requestContext) override;
 
-        virtual ArrayBufferDetachedStateBase* DetachAndGetState();
+        ArrayBufferDetachedStateBase* DetachAndGetState();
         virtual uint32 GetByteLength() const override { return bufferLength; }
         virtual BYTE* GetBuffer() const override { return buffer; }
 
@@ -185,6 +185,7 @@ namespace Js
 
         virtual ArrayBuffer * TransferInternal(DECLSPEC_GUARD_OVERFLOW uint32 newBufferLength) = 0;
     protected:
+        void Detach();
 
         typedef void __cdecl FreeFn(void* ptr);
         virtual ArrayBufferDetachedStateBase* CreateDetachedState(BYTE* buffer, DECLSPEC_GUARD_OVERFLOW uint32 bufferLength) = 0;

lib/Runtime/Library/SharedArrayBuffer.cpp

@@ -270,7 +270,7 @@ namespace Js
             }
         };
 
-        if (recycler->ReportExternalMemoryAllocation(length))
+        if (recycler->RequestExternalMemoryAllocation(length))
         {
             buffer = alloc(length);
             if (buffer == nullptr)