Windows: release VirtualVector reservations only on FreeLibrary, not at process exit

CMakeLists.txt

@@ -379,6 +379,9 @@ if (WIN32)
   # VirtualAlloc2 is exposed by mincore.lib, not Kernel32.lib (as the
   # documentation says)
   target_link_libraries(snmalloc INTERFACE $<$<NOT:$<BOOL:${WIN8COMPAT}>>:mincore>)
+  # RtlDllShutdownInProgress is exported by ntdll.  MSVC links this via a
+  # pragma in pal_windows.h, but MinGW ignores that pragma.
+  target_link_libraries(snmalloc INTERFACE ntdll)
   message(STATUS "snmalloc: Avoiding Windows 10 APIs is ${WIN8COMPAT}")
 endif()
 
@@ -548,7 +551,8 @@ if(NOT SNMALLOC_HEADER_ONLY_LIBRARY)
   # These are mitigation-independent and can be compiled once, then linked
   # against both fast and check testlib variants.
   set(TESTLIB_ONLY_TESTS
-    bits first_operation memory memory_usage multi_atexit multi_threadatexit
+    bits dll_teardown first_operation fls_teardown memory memory_usage
+    multi_atexit multi_threadatexit
     redblack statistics teardown
     contention external_pointer large_alloc lotsofthreads post_teardown
     singlethread startup
@@ -634,6 +638,10 @@ if(NOT SNMALLOC_HEADER_ONLY_LIBRARY)
             "LLVM_PROFILE_FILE=${CMAKE_BINARY_DIR}/profiles/${TESTNAME}/%p-%m.profraw")
         endif()
         if(WIN32)
+          if (${TEST} STREQUAL "dll_teardown")
+            add_dependencies(${TESTNAME}
+              snmalloc-dll-teardown-alloc snmalloc-dll-teardown-touch)
+          endif()
           # On Windows these tests use a lot of memory as it doesn't support
           # lazy commit.
           if (${TEST} MATCHES "two_alloc_types")
@@ -803,6 +811,22 @@ if(NOT SNMALLOC_HEADER_ONLY_LIBRARY)
   endif()
 
   if (SNMALLOC_BUILD_TESTING)
+    if (WIN32)
+      # DLLs loaded by the func-dll_teardown tests.
+      add_library(snmalloc-dll-teardown-alloc SHARED
+        ${TESTDIR}/func/dll_teardown/dll/alloc.cc)
+      target_link_libraries(snmalloc-dll-teardown-alloc PRIVATE snmalloc)
+      target_compile_definitions(snmalloc-dll-teardown-alloc PRIVATE
+        "SNMALLOC_USE_${TEST_CLEANUP}")
+      add_library(snmalloc-dll-teardown-touch SHARED
+        ${TESTDIR}/func/dll_teardown/dll/touch.cc)
+      foreach(DLL snmalloc-dll-teardown-alloc snmalloc-dll-teardown-touch)
+        # The test loads the DLLs by name, so avoid MinGW's "lib" prefix.
+        set_target_properties(${DLL} PROPERTIES PREFIX "")
+        add_warning_flags(${DLL})
+      endforeach()
+    endif()
+
     set(FLAVOURS fast;check)
 
     foreach(FLAVOUR ${FLAVOURS})

snmalloc-rs/snmalloc-sys/build.rs

@@ -583,6 +583,7 @@ fn configure_linking(config: &BuildConfig) {
             println!("cargo:rustc-link-lib=ws2_32");
             println!("cargo:rustc-link-lib=userenv");
             println!("cargo:rustc-link-lib=bcrypt");
+            println!("cargo:rustc-link-lib=ntdll");
             if config.debug {
                 println!("cargo:rustc-link-lib=msvcrtd");
             } else {
@@ -592,6 +593,7 @@ fn configure_linking(config: &BuildConfig) {
         _ if config.is_windows() && config.is_gnu() => {
             println!("cargo:rustc-link-lib=kernel32");
             println!("cargo:rustc-link-lib=bcrypt");
+            println!("cargo:rustc-link-lib=ntdll");
             println!("cargo:rustc-link-lib=winpthread");
 
             if config.is_clang_msys() {

src/snmalloc/pal/pal_windows.h

@@ -18,6 +18,7 @@
 #  include <windows.h>
 #  pragma comment(lib, "bcrypt.lib")
 #  include <bcrypt.h>
+#  pragma comment(lib, "ntdll.lib")
 // VirtualAlloc2 is exposed in RS5 headers.
 #  ifdef NTDDI_WIN10_RS5
 #    if (NTDDI_VERSION >= NTDDI_WIN10_RS5) && \
@@ -35,17 +36,39 @@
  * allocations have been freed.
  *
  * One way to guarantee that the reservations get released
- * at the absolute end of the program is to force them to
+ * at the end of the CRT's teardown is to force them to
  * be initialized first. Statics and globals get destroyed
  * in FILO order of when they were initialized. The pragma
  * init_seg makes sure the statics and globals in this
  * file are handled first, and thus will be the last to
- * be destroyed when the program exits or the DLL is
- * unloaded.
+ * be destroyed by the CRT when the program exits or the
+ * DLL is unloaded.
+ *
+ * Fiber-local storage callbacks (used, for example, by
+ * Rust's thread-local destructors) and DLL detach
+ * notifications are run by the loader during process
+ * exit, which for an executable is after the CRT's
+ * teardown.  Hence, the reservations are only released
+ * when a DLL is unloaded while the process continues to
+ * run; see ~VirtualVector.
  */
 #  pragma warning(disable : 4075)
 #  pragma init_seg(".CRT$XCB")
 
+/**
+ * Linker-provided symbol at the base address of the current module.
+ */
+extern "C" IMAGE_DOS_HEADER __ImageBase;
+
+/**
+ * Returns TRUE if the loader has started shutting down the process, and FALSE
+ * otherwise, including when a DLL is unloaded by FreeLibrary while the process
+ * continues to run.  This is exported by ntdll.dll, but not declared in the
+ * SDK headers, so must be declared by the caller.  See
+ * https://learn.microsoft.com/en-us/windows/win32/devnotes/rtldllshutdowninprogress
+ */
+extern "C" __declspec(dllimport) BOOLEAN NTAPI RtlDllShutdownInProgress();
+
 namespace snmalloc
 {
   class PALWindows : public PalTimerDefaultImpl<PALWindows>
@@ -231,6 +254,18 @@ namespace snmalloc
       abort();
     }
 
+    /**
+     * Returns true if the module containing this code is a DLL, and false if
+     * it is the executable.
+     */
+    static bool current_module_is_dll()
+    {
+      const auto* base = reinterpret_cast<const char*>(&__ImageBase);
+      const auto* nt_headers =
+        reinterpret_cast<const IMAGE_NT_HEADERS*>(base + __ImageBase.e_lfanew);
+      return (nt_headers->FileHeader.Characteristics & IMAGE_FILE_DLL) != 0;
+    }
+
     /// Notify platform that we will not be using these pages
     static void notify_not_using(void* p, size_t size) noexcept
     {
@@ -414,6 +449,18 @@ namespace snmalloc
 
     ~VirtualVector()
     {
+      // Memory is only returned to the OS if this module is being unloaded
+      // while the process continues to run.  An executable's statics are
+      // only destroyed as part of process exit, and a DLL may be detached
+      // as part of process exit.  In both cases, code run later by the
+      // loader, such as fiber-local storage callbacks or other DLLs' detach
+      // notifications, may still access allocations, and the OS reclaims
+      // the memory when the process terminates.
+      if (
+        !snmalloc::PALWindows::current_module_is_dll() ||
+        RtlDllShutdownInProgress())
+        return;
+
       if (data)
       {
         for (size_t i = size; i > 0; i--)
@@ -588,7 +635,9 @@ namespace snmalloc
 
   /**
    * This will be destroyed last of all of the
-   * statics and globals due to init_seg
+   * statics and globals due to init_seg.  See
+   * ~VirtualVector for when the reservations are
+   * released.
    */
   static inline VirtualVector reservations;
 

src/test/func/dll_teardown/dll/alloc.cc

@@ -0,0 +1,15 @@
+/**
+ * DLL that uses snmalloc, for the dll_teardown test.
+ */
+
+#include <snmalloc/snmalloc.h>
+
+extern "C" __declspec(dllexport) void* dll_teardown_alloc(size_t size)
+{
+  return snmalloc::alloc(size);
+}
+
+extern "C" __declspec(dllexport) void dll_teardown_dealloc(void* p)
+{
+  snmalloc::dealloc(p);
+}

src/test/func/dll_teardown/dll/touch.cc

@@ -0,0 +1,50 @@
+/**
+ * DLL that accesses an object during its process-exit detach notification,
+ * for the dll_teardown test.  The object is allocated by a different DLL that
+ * uses snmalloc and was loaded after this one, so it is detached before this
+ * one.
+ */
+
+#ifndef WIN32_LEAN_AND_MEAN
+#  define WIN32_LEAN_AND_MEAN
+#endif
+#include <windows.h>
+
+namespace
+{
+  volatile size_t* object = nullptr;
+}
+
+extern "C" __declspec(dllexport) void dll_teardown_set_object(size_t* p)
+{
+  object = p;
+}
+
+extern "C" BOOL WINAPI DllMain(HINSTANCE, DWORD reason, LPVOID reserved)
+{
+  // A non-null reserved argument indicates process exit.
+  if (
+    (reason == DLL_PROCESS_DETACH) && (reserved != nullptr) &&
+    (object != nullptr))
+  {
+    // The loader may catch exceptions raised by DllMain, so check that the
+    // memory is still committed rather than relying on an access violation.
+    MEMORY_BASIC_INFORMATION info;
+    if (
+      (VirtualQuery(const_cast<size_t*>(object), &info, sizeof(info)) == 0) ||
+      (info.State != MEM_COMMIT))
+    {
+      const char msg[] = "Memory was released before process exit\n";
+      DWORD written;
+      WriteFile(
+        GetStdHandle(STD_ERROR_HANDLE),
+        msg,
+        sizeof(msg) - 1,
+        &written,
+        nullptr);
+      TerminateProcess(GetCurrentProcess(), 1);
+    }
+    *object = *object + 1;
+  }
+  return TRUE;
+}

src/test/func/dll_teardown/dll_teardown.cc

@@ -0,0 +1,110 @@
+/**
+ * Checks when a DLL that uses snmalloc returns its memory to the OS.
+ *
+ * When the DLL is unloaded with FreeLibrary while the process continues to
+ * run, the memory it reserved must be released.
+ *
+ * When the process exits, the memory must not be released, as code run later
+ * by the loader may still access allocations.  Here, a second DLL that was
+ * loaded first, so is detached last, accesses an allocation from the first
+ * DLL in its process-exit detach notification.
+ */
+
+#include <stdio.h>
+#include <stdlib.h>
+
+#ifdef _WIN32
+#  ifndef WIN32_LEAN_AND_MEAN
+#    define WIN32_LEAN_AND_MEAN
+#  endif
+#  include <windows.h>
+
+using AllocFn = void* (*)(size_t);
+using DeallocFn = void (*)(void*);
+using SetObjectFn = void (*)(size_t*);
+
+template<typename Fn>
+Fn get_function(HMODULE module, const char* name)
+{
+  FARPROC proc = GetProcAddress(module, name);
+  if (proc == nullptr)
+  {
+    printf("GetProcAddress(%s) failed\n", name);
+    exit(1);
+  }
+  return reinterpret_cast<Fn>(reinterpret_cast<void (*)()>(proc));
+}
+
+HMODULE load(const char* name)
+{
+  HMODULE module = LoadLibraryA(name);
+  if (module == nullptr)
+  {
+    printf("LoadLibrary(%s) failed: %lu\n", name, GetLastError());
+    exit(1);
+  }
+  return module;
+}
+
+DWORD state(void* p)
+{
+  MEMORY_BASIC_INFORMATION info;
+  if (VirtualQuery(p, &info, sizeof(info)) == 0)
+  {
+    printf("VirtualQuery failed: %lu\n", GetLastError());
+    exit(1);
+  }
+  return info.State;
+}
+
+constexpr const char* alloc_dll = "snmalloc-dll-teardown-alloc.dll";
+constexpr const char* touch_dll = "snmalloc-dll-teardown-touch.dll";
+
+int main()
+{
+  // The touch DLL is loaded first, so it is detached after the alloc DLL
+  // during process exit.
+  HMODULE touch = load(touch_dll);
+
+  // FreeLibrary while the process continues to run releases the memory.
+  {
+    HMODULE alloc = load(alloc_dll);
+    auto p = get_function<AllocFn>(alloc, "dll_teardown_alloc")(sizeof(size_t));
+    if (state(p) != MEM_COMMIT)
+    {
+      printf("Allocation is not committed\n");
+      return 1;
+    }
+    get_function<DeallocFn>(alloc, "dll_teardown_dealloc")(p);
+
+    FreeLibrary(alloc);
+    if (GetModuleHandleA(alloc_dll) != nullptr)
+    {
+      // Some runtimes, such as MinGW's, pin DLLs that use thread_local
+      // destructors, so FreeLibrary does not unload them.
+      printf("DLL was not unloaded; skipping release check\n");
+    }
+    else if (state(p) != MEM_FREE)
+    {
+      printf("Memory was not released by FreeLibrary\n");
+      return 1;
+    }
+  }
+
+  // Process exit must not release the memory, as the touch DLL accesses it
+  // after the alloc DLL has been detached.
+  HMODULE alloc = load(alloc_dll);
+  auto p = static_cast<size_t*>(
+    get_function<AllocFn>(alloc, "dll_teardown_alloc")(sizeof(size_t)));
+  *p = 0;
+  get_function<SetObjectFn>(touch, "dll_teardown_set_object")(p);
+
+  printf("Exiting\n");
+  return 0;
+}
+#else
+int main()
+{
+  return 0;
+}
+#endif