bug 1366287 - Part 2.1: Track GMP memory allocation from XPCOM. r=njn

Based on similar functionality for ICU. Define a GMPReporter class and
use its methods for libgmp allocation.

js/public/Initialization.h

@@ -53,6 +53,26 @@ JS_SetICUMemoryFunctions(JS_ICUAllocFn allocFn,
                          JS_ICUReallocFn reallocFn,
                          JS_ICUFreeFn freeFn);
 
+#ifdef ENABLE_BIGINT
+namespace JS {
+
+// These types are documented as allocate_function, reallocate_function,
+// and free_function in the Info node `(gmp)Custom Allocation`.
+using GMPAllocFn = void* (*)(size_t allocSize);
+using GMPReallocFn = void* (*)(void* p, size_t oldSize, size_t newSize);
+using GMPFreeFn = void (*)(void* p, size_t size);
+
+// This function can be used to track memory used by GMP. If it is
+// called, it *must* be called before JS_Init so that same functions are
+// used for all allocations.
+extern JS_PUBLIC_API(void)
+SetGMPMemoryFunctions(GMPAllocFn allocFn,
+                      GMPReallocFn reallocFn,
+                      GMPFreeFn freeFn);
+
+}; // namespace JS
+#endif
+
 /**
  * Initialize SpiderMonkey, returning true only if initialization succeeded.
  * Once this method has succeeded, it is safe to call JS_NewContext and other

js/src/vm/BigIntType.cpp

@@ -16,6 +16,7 @@
 
 #include "gc/Allocator.h"
 #include "gc/Tracer.h"
+#include "js/Initialization.h"
 #include "js/Utility.h"
 #include "vm/JSContext.h"
 #include "vm/SelfHosting.h"
@@ -40,10 +41,27 @@ js_mp_free(void* ptr, size_t size)
     return js_free(ptr);
 }
 
+static bool memoryFunctionsInitialized = false;
+
+JS_PUBLIC_API(void)
+JS::SetGMPMemoryFunctions(JS::GMPAllocFn allocFn,
+                          JS::GMPReallocFn reallocFn,
+                          JS::GMPFreeFn freeFn)
+{
+    MOZ_ASSERT(JS::detail::libraryInitState == JS::detail::InitState::Uninitialized);
+    memoryFunctionsInitialized = true;
+    mp_set_memory_functions(allocFn, reallocFn, freeFn);
+}
+
 void
 BigInt::init()
 {
-    mp_set_memory_functions(js_malloc, js_mp_realloc, js_mp_free);
+    // Don't override custom allocation functions if
+    // JS::SetGMPMemoryFunctions was called.
+    if (!memoryFunctionsInitialized) {
+        memoryFunctionsInitialized = true;
+        mp_set_memory_functions(js_malloc, js_mp_realloc, js_mp_free);
+    }
 }
 
 BigInt*

js/src/vm/Initialization.cpp

@@ -80,7 +80,7 @@ JS::detail::InitWithFailureDiagnostic(bool isDebugBuild)
 
     MOZ_ASSERT(libraryInitState == InitState::Uninitialized,
                "must call JS_Init once before any JSAPI operation except "
-               "JS_SetICUMemoryFunctions");
+               "JS_SetICUMemoryFunctions or JS::SetGMPMemoryFunctions");
     MOZ_ASSERT(!JSRuntime::hasLiveRuntimes(),
                "how do we have live runtimes before JS_Init?");
 

xpcom/build/XPCOMInit.cpp

@@ -428,6 +428,50 @@ NS_IMPL_ISUPPORTS(VPXReporter, nsIMemoryReporter)
 CountingAllocatorBase<VPXReporter>::sAmount(0);
 #endif /* MOZ_VPX */
 
+#ifdef ENABLE_BIGINT
+class GMPReporter final
+  : public nsIMemoryReporter
+  , public CountingAllocatorBase<GMPReporter>
+{
+public:
+  NS_DECL_ISUPPORTS
+
+  static void* Alloc(size_t size)
+  {
+    return CountingMalloc(size);
+  }
+
+  static void* Realloc(void* ptr, size_t oldSize, size_t newSize)
+  {
+    return CountingRealloc(ptr, newSize);
+  }
+
+  static void Free(void* ptr, size_t size)
+  {
+    return CountingFree(ptr);
+  }
+
+private:
+  NS_IMETHOD
+  CollectReports(nsIHandleReportCallback* aHandleReport, nsISupports* aData,
+                 bool aAnonymize) override
+  {
+    MOZ_COLLECT_REPORT(
+      "explicit/gmp", KIND_HEAP, UNITS_BYTES, MemoryAllocated(),
+      "Memory allocated through libgmp for BigInt arithmetic.");
+
+    return NS_OK;
+  }
+
+  ~GMPReporter() {}
+};
+
+NS_IMPL_ISUPPORTS(GMPReporter, nsIMemoryReporter)
+
+/* static */ template<> Atomic<size_t>
+CountingAllocatorBase<GMPReporter>::sAmount(0);
+#endif // ENABLE_BIGINT
+
 static bool sInitializedJS = false;
 
 // Note that on OSX, aBinDirectory will point to .app/Contents/Resources/browser
@@ -637,6 +681,11 @@ NS_InitXPCOM2(nsIServiceManager** aResult,
                         memmove);
 #endif
 
+#ifdef ENABLE_BIGINT
+  // And for libgmp.
+  mozilla::SetGMPMemoryFunctions();
+#endif
+
   // Initialize the JS engine.
   const char* jsInitFailureReason = JS_InitWithFailureDiagnostic();
   if (jsInitFailureReason) {
@@ -797,6 +846,20 @@ SetICUMemoryFunctions()
   }
 }
 
+#ifdef ENABLE_BIGINT
+void
+SetGMPMemoryFunctions()
+{
+  static bool sGMPReporterInitialized = false;
+  if (!sGMPReporterInitialized) {
+    JS::SetGMPMemoryFunctions(GMPReporter::Alloc,
+                              GMPReporter::Realloc,
+                              GMPReporter::Free);
+    sGMPReporterInitialized = true;
+  }
+}
+#endif
+
 nsresult
 ShutdownXPCOM(nsIServiceManager* aServMgr)
 {

xpcom/build/nsXPCOMPrivate.h

@@ -41,6 +41,10 @@ ShutdownXPCOM(nsIServiceManager* aServMgr);
 
 void SetICUMemoryFunctions();
 
+#ifdef ENABLE_BIGINT
+void SetGMPMemoryFunctions();
+#endif
+
 /**
  * C++ namespaced version of NS_LogTerm.
  */