Auto alloc TLS for file system instead of requiring explicit Setup/Cleanup

include/llvm/Support/FileSystem.h

@@ -53,7 +53,10 @@ namespace fs {
 class MSFileSystem;
 typedef _Inout_ MSFileSystem* MSFileSystemRef;
 
-// Nesting [Setup/Cleanup]PerThreadFileSystem() calls is supported.
+std::error_code GetFileSystemTlsError() throw();
+
+// Setup/Cleanup are no longer necessary, but SetupPerThreadFileSystem will
+//  return GetFileSystemTlsError();
 std::error_code SetupPerThreadFileSystem() throw();
 void CleanupPerThreadFileSystem() throw();
 
@@ -79,9 +82,7 @@ private:
   std::error_code ec;
 public:
   AutoPerThreadSystem(_In_ ::llvm::sys::fs::MSFileSystem *value)
-      : m_pOrigValue(nullptr) {
-    ::llvm::sys::fs::SetupPerThreadFileSystem();
-    m_pOrigValue = ::llvm::sys::fs::GetCurrentThreadFileSystem();
+      : m_pOrigValue(::llvm::sys::fs::GetCurrentThreadFileSystem()) {
     SetCurrentThreadFileSystem(nullptr);
     ec = ::llvm::sys::fs::SetCurrentThreadFileSystem(value);
   }
@@ -94,7 +95,6 @@ public:
     } else if (!ec) {
       ::llvm::sys::fs::SetCurrentThreadFileSystem(nullptr);
     }
-    ::llvm::sys::fs::CleanupPerThreadFileSystem();
   }
 
   const std::error_code& error_code() const { return ec; }

lib/Support/Windows/MSFileSystem.inc.cpp

@@ -51,40 +51,57 @@ namespace fs {
 ///////////////////////////////////////////////////////////////////////////////////////////////////
 // Per-thread MSFileSystem support.
 
-static DWORD g_FileSystemTls = 0;
+namespace {
 
-// skip, but keep track of setup/cleanup nesting
-static volatile unsigned g_FileSystemSetupNested = 0;
+template <typename _T>
+class ThreadLocalStorage {
+  DWORD m_Tls;
+  DWORD m_dwError;
+public:
+  ThreadLocalStorage() : m_Tls(TlsAlloc()), m_dwError(0) {
+    if (m_Tls == TLS_OUT_OF_INDEXES)
+      m_dwError = ::GetLastError();
+  }
+  ~ThreadLocalStorage() { TlsFree(m_Tls); }
+  _T GetValue() throw() {
+    if (m_Tls != TLS_OUT_OF_INDEXES)
+      return (_T)TlsGetValue(m_Tls);
+    else
+      return nullptr;
+  }
+  bool SetValue(_T value) throw() {
+    if (m_Tls != TLS_OUT_OF_INDEXES) {
+      return TlsSetValue(m_Tls, (void*)value);
+    } else {
+      ::SetLastError(m_dwError);
+      return false;
+    }
+  }
+  // Retrieve error code if TlsAlloc() failed
+  DWORD GetError() const {
+    return m_dwError;
+  }
+};
+static ThreadLocalStorage<MSFileSystemRef> g_PerThreadSystem;
+}
 
-error_code SetupPerThreadFileSystem() throw()
-{
-  if (g_FileSystemSetupNested++)
+error_code GetFileSystemTlsError() throw() {
+  DWORD dwError = g_PerThreadSystem.GetError();
+  if (dwError)
+    return error_code(dwError, system_category());
+  else
     return error_code();
-
-  assert(g_FileSystemTls == 0 && "otherwise this has already been initialized, and nesting guard failed.");
-  g_FileSystemTls = TlsAlloc();
-  if (g_FileSystemTls == TLS_OUT_OF_INDEXES)
-  {
-    g_FileSystemTls = 0;
-    return mapWindowsError(::GetLastError());
-  }
-  return error_code();
 }
 
-void CleanupPerThreadFileSystem() throw()
-{
-  assert(g_FileSystemSetupNested && "otherwise, Cleanup called without matching Setup");
-  if (--g_FileSystemSetupNested)
-    return;
-  assert(g_FileSystemTls != 0 && "otherwise this has not been initialized");
-  TlsFree(g_FileSystemTls);
-  g_FileSystemTls = 0;
+// No longer necessary to call, but returns error code if TlsAlloc() failed.
+error_code SetupPerThreadFileSystem() throw() {
+  return GetFileSystemTlsError();
 }
+void CleanupPerThreadFileSystem() throw() {}
 
 MSFileSystemRef GetCurrentThreadFileSystem() throw()
 {
-  assert(g_FileSystemTls != 0 && "otherwise this has not been initialized");
-  return (MSFileSystemRef)TlsGetValue(g_FileSystemTls);
+  return g_PerThreadSystem.GetValue();
 }
 
 error_code SetCurrentThreadFileSystem(MSFileSystemRef value) throw()
@@ -99,7 +116,7 @@ error_code SetCurrentThreadFileSystem(MSFileSystemRef value) throw()
     }
   }
 
-  if (!TlsSetValue(g_FileSystemTls, value))
+  if (!g_PerThreadSystem.SetValue(value))
   {
     return mapWindowsError(::GetLastError());
   }

tools/clang/unittests/HLSL/CompilerTest.cpp

@@ -389,7 +389,6 @@ public:
   END_TEST_CLASS()
 
   TEST_CLASS_SETUP(InitSupport);
-  TEST_CLASS_CLEANUP(Cleanup);
 
   TEST_METHOD(CompileWhenDebugThenDIPresent)
   TEST_METHOD(CompileDebugLines)
@@ -1751,14 +1750,8 @@ bool CompilerTest::InitSupport() {
   if (!m_dllSupport.IsEnabled()) {
     VERIFY_SUCCEEDED(m_dllSupport.Initialize());
     m_ver.Initialize(m_dllSupport);
-    llvm::sys::fs::SetupPerThreadFileSystem();
-  }
-  return true;
-}
-bool CompilerTest::Cleanup() {
-  if (m_dllSupport.IsEnabled()) {
-    m_dllSupport.Cleanup();
-    llvm::sys::fs::CleanupPerThreadFileSystem();
+    if (llvm::sys::fs::GetFileSystemTlsResult())
+      return false;
   }
   return true;
 }

tools/clang/unittests/HLSL/DxilModuleTest.cpp

@@ -42,7 +42,6 @@ public:
   END_TEST_CLASS()
 
   TEST_CLASS_SETUP(InitSupport);
-  TEST_CLASS_CLEANUP(Cleanup);
 
   dxc::DxcDllSupport m_dllSupport;
 
@@ -65,14 +64,8 @@ bool DxilModuleTest::InitSupport() {
   if (!m_dllSupport.IsEnabled()) {
     VERIFY_SUCCEEDED(m_dllSupport.Initialize());
   }
-  llvm::sys::fs::SetupPerThreadFileSystem();
-  return true;
-}
-bool DxilModuleTest::Cleanup() {
-  if (m_dllSupport.IsEnabled()) {
-    m_dllSupport.Cleanup();
-  }
-  llvm::sys::fs::CleanupPerThreadFileSystem();
+  if (llvm::sys::fs::GetFileSystemTlsResult())
+    return false;
   return true;
 }