Bug 1917844 - Expose ForkServer-specific mozjemalloc reinit r=glandium

Differential Revision: https://phabricator.services.mozilla.com/D220578

ipc/glue/ForkServer.cpp

@@ -318,6 +318,17 @@ bool ForkServer::RunForkServer(int* aArgc, char*** aArgv) {
   MOZ_ASSERT(!XRE_IsForkServerProcess(),
              "fork server created another fork server?");
 
+  // This is now a child process, and it may even be a Content process.
+  // It is required that the PRNG at least is re-initialized so the same state
+  // is not shared accross all child processes, and in case of a Content process
+  // it is also required that the small allocation are not being randomized ;
+  // failing to do so will lead to performance regressions, e.g. as in
+  // bug 1912262.
+#if defined(MOZ_MEMORY)
+  jemalloc_reset_small_alloc_randomization(
+      /* aRandomizeSmall */ !XRE_IsContentProcess());
+#endif
+
   // Open log files again with right names and the new PID.
   nsTraceRefcnt::ReopenLogFilesAfterFork(XRE_GetProcessTypeString());
 

memory/build/malloc_decls.h

@@ -114,6 +114,13 @@ MALLOC_DECL(jemalloc_free_dirty_pages, void)
 // after lowering the max dirty pages threshold to get RSS back to normal.
 MALLOC_DECL(jemalloc_free_excess_dirty_pages, void)
 
+// Only used by ForkServer after forking new child processes.
+// Change the value of opt_randomize_small to control small allocation
+// randomization and maybe perform a reinitialization of the arena's PRNG.
+#    if defined(MOZ_ENABLE_FORKSERVER)
+MALLOC_DECL(jemalloc_reset_small_alloc_randomization, void, bool)
+#    endif
+
 // Opt in or out of a thread local arena (bool argument is whether to opt-in
 // (true) or out (false)).
 MALLOC_DECL(jemalloc_thread_local_arena, void, bool)

memory/build/mozjemalloc.cpp

@@ -1204,6 +1204,10 @@ struct arena_t {
   explicit arena_t(arena_params_t* aParams, bool aIsPrivate);
   ~arena_t();
 
+  void ResetSmallAllocRandomization();
+
+  void InitPRNG() MOZ_REQUIRES(mLock);
+
  private:
   void InitChunk(arena_chunk_t* aChunk, size_t aMinCommittedPages);
 
@@ -3391,6 +3395,40 @@ void arena_bin_t::Init(SizeClass aSizeClass) {
   mSizeDivisor = FastDivisor<uint16_t>(aSizeClass.Size(), try_run_size);
 }
 
+void arena_t::ResetSmallAllocRandomization() {
+  if (MOZ_UNLIKELY(opt_randomize_small)) {
+    MaybeMutexAutoLock lock(mLock);
+    InitPRNG();
+  }
+  mRandomizeSmallAllocations = opt_randomize_small;
+}
+
+void arena_t::InitPRNG() {
+  // Both another thread could race and the code backing RandomUint64
+  // (arc4random for example) may allocate memory while here, so we must
+  // ensure to start the mPRNG initialization only once and to not hold
+  // the lock while initializing.
+  mIsPRNGInitializing = true;
+  {
+    mLock.Unlock();
+    mozilla::Maybe<uint64_t> prngState1 = mozilla::RandomUint64();
+    mozilla::Maybe<uint64_t> prngState2 = mozilla::RandomUint64();
+    mLock.Lock();
+
+    mozilla::non_crypto::XorShift128PlusRNG prng(prngState1.valueOr(0),
+                                                 prngState2.valueOr(0));
+    if (mPRNG) {
+      *mPRNG = prng;
+    } else {
+      void* backing =
+          base_alloc(sizeof(mozilla::non_crypto::XorShift128PlusRNG));
+      mPRNG = new (backing)
+          mozilla::non_crypto::XorShift128PlusRNG(std::move(prng));
+    }
+  }
+  mIsPRNGInitializing = false;
+}
+
 void* arena_t::MallocSmall(size_t aSize, bool aZero) {
   void* ret;
   arena_bin_t* bin;
@@ -3429,25 +3467,7 @@ void* arena_t::MallocSmall(size_t aSize, bool aZero) {
 
     if (MOZ_UNLIKELY(mRandomizeSmallAllocations && mPRNG == nullptr &&
                      !mIsPRNGInitializing)) {
-      // Both another thread could race and the code backing RandomUint64
-      // (arc4random for example) may allocate memory while here, so we must
-      // ensure to start the mPRNG initialization only once and to not hold
-      // the lock while initializing.
-      mIsPRNGInitializing = true;
-      mozilla::non_crypto::XorShift128PlusRNG* prng;
-      {
-        // TODO: I think no MaybeMutexAutoUnlock or similar exists, should it?
-        mLock.Unlock();
-        mozilla::Maybe<uint64_t> prngState1 = mozilla::RandomUint64();
-        mozilla::Maybe<uint64_t> prngState2 = mozilla::RandomUint64();
-        void* backing =
-            base_alloc(sizeof(mozilla::non_crypto::XorShift128PlusRNG));
-        prng = new (backing) mozilla::non_crypto::XorShift128PlusRNG(
-            prngState1.valueOr(0), prngState2.valueOr(0));
-        mLock.Lock();
-      }
-      mPRNG = prng;
-      mIsPRNGInitializing = false;
+      InitPRNG();
     }
     MOZ_ASSERT(!mRandomizeSmallAllocations || mPRNG);
 
@@ -5168,6 +5188,39 @@ inline void MozJemalloc::moz_set_max_dirty_page_modifier(int32_t aModifier) {
   gArenas.SetDefaultMaxDirtyPageModifier(aModifier);
 }
 
+#if defined(MOZ_ENABLE_FORKSERVER)
+inline void MozJemalloc::jemalloc_reset_small_alloc_randomization(
+    bool aRandomizeSmall) {
+  // When this process got forked by ForkServer then it inherited the existing
+  // state of mozjemalloc. Specifically, parsing of MALLOC_OPTIONS has already
+  // been done but it may not reflect anymore the current set of options after
+  // the fork().
+  //
+  // Content process will have randomization on small malloc disabled via the
+  // MALLOC_OPTIONS environment variable set by parent process, missing this
+  // will lead to serious performance regressions because CPU prefetch will
+  // break, cf bug 1912262.  However on forkserver-forked Content processes, the
+  // environment is not yet reset when the postfork child handler is being
+  // called.
+  //
+  // This API is here to allow those forkserver-forked Content processes to
+  // notify jemalloc to turn off the randomization on small allocations and
+  // perform the required reinitialization of already existing arena's PRNG.
+  // It is important to make sure that the PRNG state is properly re-initialized
+  // otherwise child processes would share all the same state.
+
+  {
+    AutoLock<StaticMutex> lock(gInitLock);
+    opt_randomize_small = aRandomizeSmall;
+  }
+
+  MutexAutoLock lock(gArenas.mLock);
+  for (auto* arena : gArenas.iter()) {
+    arena->ResetSmallAllocRandomization();
+  }
+}
+#endif
+
 #define MALLOC_DECL(name, return_type, ...)                          \
   inline return_type MozJemalloc::moz_arena_##name(                  \
       arena_id_t aArenaId, ARGS_HELPER(TYPED_ARGS, ##__VA_ARGS__)) { \