From c8c9f5504fcfd0d6e1266f8f3e29027be9d71dca Mon Sep 17 00:00:00 2001
From: Saleem Abdulrasool <abdulras@google.com>
Date: Thu, 16 Jul 2020 10:12:10 -0700
Subject: [PATCH] parallel: add support to explicitly destroy TLS keys

In most modern pthread implementations, the TLS key does not create heap
allocations.  However, this is not true for all implementations of TLS.
Extend the interface to permit the destruction of the key to release
resources.

Also add an implementation which conforms to the POSIX threading
interfaces.

Extend the change a bit further to allow for the resetting of the key
itself to a TLS implementation specific invalid value.  Although this
unconstifies the value, it will make it easier to identify invalid uses
of the key (e.g. use-after-free).
---
 Sources/PenguinParallel/ConcurrencyPlatform.swift | 9 ++++++++-
 Sources/PenguinParallel/ThreadLocalStorage.swift  | 9 +++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/Sources/PenguinParallel/ConcurrencyPlatform.swift b/Sources/PenguinParallel/ConcurrencyPlatform.swift
index 1209c4b3..28e18e83 100644
--- a/Sources/PenguinParallel/ConcurrencyPlatform.swift
+++ b/Sources/PenguinParallel/ConcurrencyPlatform.swift
@@ -146,6 +146,8 @@ public protocol RawThreadLocalStorage {
 
   /// Makes a new key; the returned key should be used for the entire process lifetime.
   static func makeKey() -> Key
+  /// Invalidates a previously constructed key, freeing resources
+  static func destroyKey(_ key: inout Key)
   /// Retrieves the raw pointer associated with the given key.
   static func get(for key: Key) -> UnsafeMutableRawPointer?
   /// Sets the raw pointer associated with the given key.
@@ -158,7 +160,7 @@ public struct TypedThreadLocalStorage<Underlying: RawThreadLocalStorage> {
 
   /// Token used to index into the thread local storage.
   public struct Key<Value: AnyObject> {
-    fileprivate let key: Underlying.Key
+    fileprivate var key: Underlying.Key
 
     /// The thread-local value associated with `self`.
     public var localValue: Value? {
@@ -176,6 +178,11 @@ public struct TypedThreadLocalStorage<Underlying: RawThreadLocalStorage> {
     Key(key: Underlying.makeKey())
   }
 
+  /// Deallocates a key for type `T`.
+  public static func destroyKey<T: AnyObject>(_ key: inout Key<T>) {
+    Underlying.destroyKey(&key.key)
+  }
+
   /// Retrieves the thread-local value for `key`, if it exists.
   public static func get<T: AnyObject>(_ key: Key<T>) -> T? {
     guard let ptr = Underlying.get(for: key.key) else { return nil }
diff --git a/Sources/PenguinParallel/ThreadLocalStorage.swift b/Sources/PenguinParallel/ThreadLocalStorage.swift
index a2f6c117..cd768918 100644
--- a/Sources/PenguinParallel/ThreadLocalStorage.swift
+++ b/Sources/PenguinParallel/ThreadLocalStorage.swift
@@ -59,6 +59,15 @@ public struct PosixThreadLocalStorage: RawThreadLocalStorage {
     Key()
   }
 
+  public static func destroyKey(_ key: inout Key) {
+    #if os(Windows)
+      fatalError("Unimplemented!")
+    #else
+      pthread_key_delete(key.value)
+      key.value = 0
+    #endif
+  }
+
   public static func get(for key: Key) -> UnsafeMutableRawPointer? {
     pthread_getspecific(key.value)
   }