Prevent leaks in ThreadLocalPtr initialization

Summary: While making an unrelated change, (D3271563, which was needed from an unrelated change, (D3237530)), I noticed a lack of exception safety here. If `std::bad_alloc` were thrown, then we don't want to leak memory.

Reviewed By: ericniebler

Differential Revision: D3271911

fbshipit-source-id: 0d316c0fa865a7d64622c1d62160bb0c2b061d78

diff --git a/folly/ThreadLocal.h b/folly/ThreadLocal.h
index abc9b64af3ac4af2d0a4666035948622599b3aa3..d1bc430c92e820c0e6262b155655c3fa2f3c3627 100644 (file)
--- a/folly/ThreadLocal.h
+++ b/folly/ThreadLocal.h
@@ -36,11 +36,12 @@
 
 #pragma once
 
+#include <folly/Likely.h>
 #include <folly/Portability.h>
+#include <folly/ScopeGuard.h>
 #include <boost/iterator/iterator_facade.hpp>
-#include <folly/Likely.h>
 #include <type_traits>
-
+#include <utility>
 
 namespace folly {
 enum class TLPDestructionMode {
@@ -180,12 +181,12 @@ class ThreadLocalPtr {
   }
 
   void reset(T* newPtr = nullptr) {
+    auto guard = makeGuard([&] { delete newPtr; });
     threadlocal_detail::ElementWrapper& w = StaticMeta::instance().get(&id_);
 
-    if (w.ptr != newPtr) {
-      w.dispose(TLPDestructionMode::THIS_THREAD);
-      w.set(newPtr);
-    }
+    w.dispose(TLPDestructionMode::THIS_THREAD);
+    guard.dismiss();
+    w.set(newPtr);
   }
 
   explicit operator bool() const {
@@ -197,15 +198,20 @@ class ThreadLocalPtr {
    * deleter(T* ptr, TLPDestructionMode mode)
    * "mode" is ALL_THREADS if we're destructing this ThreadLocalPtr (and thus
    * deleting pointers for all threads), and THIS_THREAD if we're only deleting
-   * the member for one thread (because of thread exit or reset())
+   * the member for one thread (because of thread exit or reset()).
+   * Invoking the deleter must not throw.
    */
   template <class Deleter>
-  void reset(T* newPtr, Deleter deleter) {
+  void reset(T* newPtr, const Deleter& deleter) {
+    auto guard = makeGuard([&] {
+      if (newPtr) {
+        deleter(newPtr, TLPDestructionMode::THIS_THREAD);
+      }
+    });
     threadlocal_detail::ElementWrapper& w = StaticMeta::instance().get(&id_);
-    if (w.ptr != newPtr) {
-      w.dispose(TLPDestructionMode::THIS_THREAD);
-      w.set(newPtr, deleter);
-    }
+    w.dispose(TLPDestructionMode::THIS_THREAD);
+    guard.dismiss();
+    w.set(newPtr, deleter);
   }
 
   // Holds a global lock for iteration through all thread local child objects.

diff --git a/folly/detail/ThreadLocalDetail.h b/folly/detail/ThreadLocalDetail.h
index 29e06b1f3bf1cc73a339656ae1608f69f26ed127..ea9ae136b617a4f74899bf19514af3df36805978 100644 (file)
--- a/folly/detail/ThreadLocalDetail.h
+++ b/folly/detail/ThreadLocalDetail.h
@@ -33,6 +33,7 @@
 #include <folly/Malloc.h>
 #include <folly/MicroSpinLock.h>
 #include <folly/Portability.h>
+#include <folly/ScopeGuard.h>
 
 #include <folly/detail/StaticSingletonManager.h>
 
@@ -81,6 +82,7 @@ struct ElementWrapper {
 
   template <class Ptr>
   void set(Ptr p) {
+    auto guard = makeGuard([&] { delete p; });
     DCHECK(ptr == nullptr);
     DCHECK(deleter1 == nullptr);
 
@@ -90,20 +92,27 @@ struct ElementWrapper {
         delete static_cast<Ptr>(pt);
       };
       ownsDeleter = false;
+      guard.dismiss();
     }
   }
 
   template <class Ptr, class Deleter>
-  void set(Ptr p, Deleter d) {
+  void set(Ptr p, const Deleter& d) {
+    auto guard = makeGuard([&] {
+      if (p) {
+        d(p, TLPDestructionMode::THIS_THREAD);
+      }
+    });
     DCHECK(ptr == nullptr);
     DCHECK(deleter2 == nullptr);
     if (p) {
       ptr = p;
-      deleter2 = new std::function<DeleterFunType>(
-          [d](void* pt, TLPDestructionMode mode) {
-            d(static_cast<Ptr>(pt), mode);
-          });
+      deleter2 = new std::function<DeleterFunType>([d = d](
+          void* pt, TLPDestructionMode mode) {
+        d(static_cast<Ptr>(pt), mode);
+      });
       ownsDeleter = true;
+      guard.dismiss();
     }
   }