[PartitionAlloc] Add TLS support for Windows, enable the thread cache.

The thread cache is disabled on Windows, as TLS is not implemented
there. Add a Windows TLS implementation, and enable the thread cache.

Bug: 998048
Change-Id: Ia9d90d5d5e78611e00e6456a770bdf00f44bb773
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2426323
Commit-Queue: Benoit L <lizeb@chromium.org>
Reviewed-by: Kentaro Hara <haraken@chromium.org>
Cr-Commit-Position: refs/heads/master@{#816603}

base/BUILD.gn

@@ -1793,8 +1793,10 @@ component("base") {
         "allocator/partition_allocator/yield_processor.h",
       ]
       if (is_win) {
-        sources +=
-            [ "allocator/partition_allocator/page_allocator_internals_win.h" ]
+        sources += [
+          "allocator/partition_allocator/page_allocator_internals_win.h",
+          "allocator/partition_allocator/partition_tls_win.cc",
+        ]
       } else if (is_posix) {
         sources +=
             [ "allocator/partition_allocator/page_allocator_internals_posix.h" ]

base/allocator/partition_allocator/partition_alloc.cc

@@ -320,7 +320,7 @@ void PartitionRoot<thread_safe>::Init(PartitionOptions opts) {
   // this operation is idempotent, so there is no harm.
   InitBucketIndexLookup(this);
 
-#if !defined(OS_POSIX)
+#if !defined(PA_THREAD_CACHE_SUPPORTED)
   // TLS in ThreadCache not supported on other OSes.
   with_thread_cache = false;
 #else

base/allocator/partition_allocator/partition_tls.h

@@ -6,12 +6,17 @@
 #define BASE_ALLOCATOR_PARTITION_ALLOCATOR_PARTITION_TLS_H_
 
 #include "base/allocator/partition_allocator/partition_alloc_check.h"
+#include "base/compiler_specific.h"
 #include "build/build_config.h"
 
 #if defined(OS_POSIX)
 #include <pthread.h>
 #endif
 
+#if defined(OS_WIN)
+#include <windows.h>
+#endif
+
 // Barebones TLS implementation for use in PartitionAlloc. This doesn't use the
 // general chromium TLS handling to avoid dependencies, but more importantly
 // because it allocates memory.
@@ -21,33 +26,48 @@ namespace internal {
 #if defined(OS_POSIX)
 typedef pthread_key_t PartitionTlsKey;
 
-inline bool PartitionTlsCreate(PartitionTlsKey* key,
-                               void (*destructor)(void*)) {
+ALWAYS_INLINE bool PartitionTlsCreate(PartitionTlsKey* key,
+                                      void (*destructor)(void*)) {
   return !pthread_key_create(key, destructor);
 }
-inline void* PartitionTlsGet(PartitionTlsKey key) {
+ALWAYS_INLINE void* PartitionTlsGet(PartitionTlsKey key) {
   return pthread_getspecific(key);
 }
-inline void PartitionTlsSet(PartitionTlsKey key, void* value) {
+ALWAYS_INLINE void PartitionTlsSet(PartitionTlsKey key, void* value) {
   int ret = pthread_setspecific(key, value);
   PA_DCHECK(!ret);
 }
+#elif defined(OS_WIN)
+// Note: supports only a single TLS key on Windows. Not a hard constraint, may
+// be lifted.
+typedef unsigned long PartitionTlsKey;
+
+BASE_EXPORT bool PartitionTlsCreate(PartitionTlsKey* key,
+                                    void (*destructor)(void*));
+
+ALWAYS_INLINE void* PartitionTlsGet(PartitionTlsKey key) {
+  return TlsGetValue(key);
+}
+
+ALWAYS_INLINE void PartitionTlsSet(PartitionTlsKey key, void* value) {
+  BOOL ret = TlsSetValue(key, value);
+  PA_DCHECK(ret);
+}
 #else
-// Not implemented.
+// Not supported.
 typedef int PartitionTlsKey;
-
-inline bool PartitionTlsCreate(PartitionTlsKey* key,
-                               void (*destructor)(void*)) {
-  // Cannot use NOIMPLEMENTED() as it may allocate.
+ALWAYS_INLINE bool PartitionTlsCreate(PartitionTlsKey* key,
+                                      void (*destructor)(void*)) {
+  // NOTIMPLEMENTED() may allocate, crash instead.
   IMMEDIATE_CRASH();
 }
-inline void* PartitionTlsGet(PartitionTlsKey key) {
+ALWAYS_INLINE void* PartitionTlsGet(PartitionTlsKey key) {
   IMMEDIATE_CRASH();
 }
-inline void PartitionTlsSet(PartitionTlsKey key, void* value) {
+ALWAYS_INLINE void PartitionTlsSet(PartitionTlsKey key, void* value) {
   IMMEDIATE_CRASH();
 }
-#endif  // defined(OS_POSIX)
+#endif  // defined(OS_WIN)
 
 }  // namespace internal
 }  // namespace base

base/allocator/partition_allocator/partition_tls_win.cc

@@ -0,0 +1,101 @@
+// Copyright 2020 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "base/allocator/partition_allocator/partition_tls.h"
+
+namespace base {
+namespace internal {
+namespace {
+
+// Store the key as the thread destruction callback doesn't get it.
+PartitionTlsKey g_key;
+void (*g_destructor)(void*) = nullptr;
+
+// Static callback function to call with each thread termination.
+void NTAPI PartitionTlsOnThreadExit(PVOID module,
+                                    DWORD reason,
+                                    PVOID reserved) {
+  if (reason != DLL_THREAD_DETACH && reason != DLL_PROCESS_DETACH)
+    return;
+
+  if (g_destructor) {
+    void* per_thread_data = PartitionTlsGet(g_key);
+    if (per_thread_data)
+      g_destructor(per_thread_data);
+  }
+}
+
+}  // namespace
+
+bool PartitionTlsCreate(PartitionTlsKey* key, void (*destructor)(void*)) {
+  PA_CHECK(g_destructor == nullptr);  // Only one TLS key supported at a time.
+  PartitionTlsKey value = TlsAlloc();
+  if (value != TLS_OUT_OF_INDEXES) {
+    *key = value;
+
+    g_key = value;
+    g_destructor = destructor;
+    return true;
+  }
+  return false;
+}
+
+}  // namespace internal
+}  // namespace base
+
+// See thread_local_storage_win.cc for details and reference.
+//
+// The callback has to be in any section between .CRT$XLA and .CRT$XLZ, as these
+// are sentinels used by the TLS code to find the callback array bounds. As we
+// don't particularly care about where we are called but would prefer to be
+// deinitialized towards the end (in particular after Chromium's TLS), we locate
+// ourselves in .CRT$XLY.
+
+// Force a reference to _tls_used to make the linker create the TLS directory if
+// it's not already there.  (e.g. if __declspec(thread) is not used).  Force a
+// reference to partition_tls_thread_exit_callback to prevent whole program
+// optimization from discarding the variable.
+#ifdef _WIN64
+
+#pragma comment(linker, "/INCLUDE:_tls_used")
+#pragma comment(linker, "/INCLUDE:partition_tls_thread_exit_callback")
+
+#else  // _WIN64
+
+#pragma comment(linker, "/INCLUDE:__tls_used")
+#pragma comment(linker, "/INCLUDE:_partition_tls_thread_exit_callback")
+
+#endif  // _WIN64
+
+// extern "C" suppresses C++ name mangling so we know the symbol name for the
+// linker /INCLUDE:symbol pragma above.
+extern "C" {
+// The linker must not discard partition_tls_thread_exit_callback.  (We force a
+// reference to this variable with a linker /INCLUDE:symbol pragma to ensure
+// that.) If this variable is discarded, PartitionTlsOnThreadExit will never be
+// called.
+#ifdef _WIN64
+
+// .CRT section is merged with .rdata on x64 so it must be constant data.
+#pragma const_seg(".CRT$XLY")
+// When defining a const variable, it must have external linkage to be sure the
+// linker doesn't discard it.
+extern const PIMAGE_TLS_CALLBACK partition_tls_thread_exit_callback;
+const PIMAGE_TLS_CALLBACK partition_tls_thread_exit_callback =
+    base::internal::PartitionTlsOnThreadExit;
+
+// Reset the default section.
+#pragma const_seg()
+
+#else  // _WIN64
+
+#pragma data_seg(".CRT$XLY")
+PIMAGE_TLS_CALLBACK partition_tls_thread_exit_callback =
+    base::internal::PartitionTlsOnThreadExit;
+
+// Reset the default section.
+#pragma data_seg()
+
+#endif  // _WIN64
+}

base/allocator/partition_allocator/thread_cache.h

@@ -19,6 +19,11 @@
 #include "base/partition_alloc_buildflags.h"
 #include "base/synchronization/lock.h"
 
+// Need TLS support.
+#if defined(OS_POSIX) || defined(OS_WIN)
+#define PA_THREAD_CACHE_SUPPORTED
+#endif
+
 namespace base {
 
 namespace internal {

base/allocator/partition_allocator/thread_cache_unittest.cc

@@ -28,7 +28,7 @@
 // disable the cache (and tests)
 #if !BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC) && \
     !defined(MEMORY_TOOL_REPLACES_ALLOCATOR) &&  \
-    (defined(OS_LINUX) || defined(OS_CHROMEOS))
+    defined(PA_THREAD_CACHE_SUPPORTED)
 
 namespace base {
 namespace internal {
@@ -404,4 +404,4 @@ TEST_F(ThreadCacheTest, PurgeAll) NO_THREAD_SAFETY_ANALYSIS {
 
 #endif  // !BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC) &&
         // !defined(MEMORY_TOOL_REPLACES_ALLOCATOR) &&
-        // (defined(OS_LINUX) || defined(OS_CHROMEOS))
+        // defined(PA_THREAD_CACHE_SUPPORTED)