Allocation profiling fix - pthread TLS storage (#290)

- Manage TLS storage with pthread APIs.
A user provided a case where __tls_get_addr was calling back into Go, which was calling back into mmap.
This was causing an infinite recursive loop
By managing the lifetime of TLS variables, we can fix this.

- Add a test to check if TLS variables are being used

- Symbol overrides: Minor fix on the usage of wrong pointer for realloc array

- Add some settings to the allocation tracking benchmark

include/lib/allocation_tracker.hpp

@@ -5,39 +5,24 @@
 
 #pragma once
 
+#include "allocation_tracker_tls.hpp"
 #include "ddprof_base.hpp"
 #include "ddres_def.hpp"
 #include "pevent.hpp"
-#include "span.hpp"
+#include "reentry_guard.hpp"
 #include "unlikely.hpp"
 
 #include <atomic>
 #include <cstddef>
 #include <mutex>
-#include <random>
+#include <pthread.h>
 #include <unordered_set>
 
 namespace ddprof {
 
 class MPSCRingBufferWriter;
 struct RingBufferInfo;
 
-struct TrackerThreadLocalState {
-  int64_t remaining_bytes; // remaining allocation bytes until next sample
-  bool remaining_bytes_initialized; // false if remaining_bytes is not
-                                    // initialized
-  ddprof::span<const byte> stack_bounds;
-
-  // In the choice of random generators, this one is smaller
-  // - smaller than mt19937 (8 vs 5K)
-  std::minstd_rand _gen{std::random_device{}()};
-
-  pid_t tid; // cache of tid
-
-  bool reentry_guard; // prevent reentry in AllocationTracker (eg. when
-                      // allocation are done inside AllocationTracker)
-};
-
 class AllocationTracker {
 public:
   friend class AllocationTrackerDisablerForCurrentThread;
@@ -63,11 +48,17 @@ class AllocationTracker {
   static void allocation_tracking_free();
 
   static inline DDPROF_NO_SANITIZER_ADDRESS void
-  track_allocation(uintptr_t addr, size_t size);
-  static inline void track_deallocation(uintptr_t addr);
+  track_allocation_s(uintptr_t addr, size_t size,
+                     TrackerThreadLocalState &tl_state);
+
+  static inline void track_deallocation_s(uintptr_t addr);
 
   static inline bool is_active();
 
+  static TrackerThreadLocalState *init_tl_state();
+  // can return null (does not init)
+  static TrackerThreadLocalState *get_tl_state();
+
 private:
   using AdressSet = std::unordered_set<uintptr_t>;
 
@@ -97,6 +88,10 @@ class AllocationTracker {
 
   static AllocationTracker *create_instance();
 
+  static void delete_tl_state(void *tl_state);
+
+  static void make_key();
+
   void track_allocation(uintptr_t addr, size_t size,
                         TrackerThreadLocalState &tl_state);
   void track_deallocation(uintptr_t addr, TrackerThreadLocalState &tl_state);
@@ -120,11 +115,18 @@ class AllocationTracker {
   bool _deterministic_sampling;
   AdressSet _address_set;
 
-  static thread_local TrackerThreadLocalState _tl_state;
+  // These can not be tied to the internal state of the instance.
+  // The creation of the instance depends on this
+  static pthread_once_t _key_once; // ensures we call key creation a single time
+  static pthread_key_t _tl_state_key;
+  // For Thread reentry guard of init_tl_state
+  static ThreadEntries _thread_entries;
+
   static AllocationTracker *_instance;
 };
 
-void AllocationTracker::track_allocation(uintptr_t addr, size_t size) {
+void AllocationTracker::track_allocation_s(uintptr_t addr, size_t size,
+                                           TrackerThreadLocalState &tl_state) {
   AllocationTracker *instance = _instance;
 
   // Be safe, if allocation tracker has not been initialized, just bail out
@@ -136,10 +138,6 @@ void AllocationTracker::track_allocation(uintptr_t addr, size_t size) {
     return;
   }
 
-  // In shared libraries, TLS access requires a call to tls_get_addr,
-  // therefore obtain a pointer on TLS state once and pass it around
-  TrackerThreadLocalState &tl_state = _tl_state;
-
   tl_state.remaining_bytes += size;
   if (likely(tl_state.remaining_bytes < 0)) {
     return;
@@ -155,19 +153,18 @@ void AllocationTracker::track_allocation(uintptr_t addr, size_t size) {
   }
 }
 
-void AllocationTracker::track_deallocation(uintptr_t addr) {
+void AllocationTracker::track_deallocation_s(uintptr_t addr) {
   // same pattern as track_allocation
   AllocationTracker *instance = _instance;
-
   if (!instance) {
     return;
   }
-  TrackerThreadLocalState &tl_state = _tl_state;
-
   if (instance->_state.track_deallocations.load(std::memory_order_relaxed)) {
-    // not cool as we are always calling this (high overhead). Can we do better
-    // ?
-    instance->track_deallocation(addr, tl_state);
+    TrackerThreadLocalState *tl_state = get_tl_state();
+    if (unlikely(!tl_state)) {
+      return;
+    }
+    instance->track_deallocation(addr, *tl_state);
   }
 }
 

include/lib/allocation_tracker_tls.hpp

@@ -0,0 +1,28 @@
+#pragma once
+
+#include "span.hpp"
+
+#include <cstdint>
+#include <random>
+#include <sys/types.h>
+
+namespace ddprof {
+
+struct TrackerThreadLocalState {
+  int64_t remaining_bytes; // remaining allocation bytes until next sample
+  bool remaining_bytes_initialized; // false if remaining_bytes is not
+                                    // initialized
+  ddprof::span<const byte> stack_bounds;
+
+  // In the choice of random generators, this one is smaller
+  // - smaller than mt19937 (8 vs 5K)
+  std::minstd_rand _gen{std::random_device{}()};
+
+  pid_t tid; // cache of tid
+
+  bool reentry_guard;         // prevent reentry in AllocationTracker (eg. when
+                              // allocation are done inside AllocationTracker)
+  bool double_tracking_guard; // prevent mmap tracking within a malloc
+};
+
+} // namespace ddprof

include/lib/lib_logger.hpp

@@ -0,0 +1,22 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0. This product includes software
+// developed at Datadog (https://www.datadoghq.com/). Copyright 2021-Present
+// Datadog, Inc.
+
+#pragma once
+
+#include <cstdarg>
+#include <cstdio>
+#include <mutex>
+
+namespace ddprof {
+template <typename... Args>
+void log_once(char const *const format, Args... args) {
+#ifndef DEBUG
+  static std::once_flag flag;
+  std::call_once(flag, [&, format]() { fprintf(stderr, format, args...); });
+#else
+  fprintf(stderr, format, args...);
+#endif
+}
+} // namespace ddprof

include/lib/reentry_guard.hpp

@@ -0,0 +1,106 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0. This product includes software
+// developed at Datadog (https://www.datadoghq.com/). Copyright 2021-Present
+// Datadog, Inc.
+
+#pragma once
+
+#include <array>
+#include <atomic>
+#include <thread>
+
+namespace ddprof {
+
+class ThreadEntries {
+public:
+  static constexpr size_t max_threads = 10;
+  std::array<std::atomic<pid_t>, max_threads> thread_entries;
+  ThreadEntries() { reset(); }
+  void reset() {
+    for (auto &entry : thread_entries) {
+      entry.store(-1, std::memory_order_relaxed);
+    }
+  }
+};
+
+class TLReentryGuard {
+public:
+  explicit TLReentryGuard(ThreadEntries &entries, pid_t tid)
+      : _entries(entries), _tid(tid), _ok(false), _index(-1) {
+    while (true) {
+      for (size_t i = 0; i < ThreadEntries::max_threads; ++i) {
+        pid_t expected = -1;
+        if (_entries.thread_entries[i].compare_exchange_weak(
+                expected, tid, std::memory_order_acq_rel)) {
+          _ok = true;
+          _index = i;
+          return;
+        } else if (expected == tid) {
+          // This thread is already in the entries.
+          return;
+        }
+      }
+      // If we've reached here, all slots are occupied and none of them belongs
+      // to this thread. Let's yield to other threads and then try again.
+      std::this_thread::yield();
+    }
+  }
+
+  ~TLReentryGuard() {
+    if (_ok) {
+      _entries.thread_entries[_index].store(-1, std::memory_order_release);
+    }
+  }
+
+  explicit operator bool() const { return _ok; }
+
+  TLReentryGuard(const TLReentryGuard &) = delete;
+  TLReentryGuard &operator=(const TLReentryGuard &) = delete;
+
+private:
+  ThreadEntries &_entries;
+  pid_t _tid;
+  bool _ok;
+  int _index;
+};
+
+class ReentryGuard {
+public:
+  explicit ReentryGuard(bool *reentry_guard)
+      : _reentry_guard(reentry_guard), _ok(false) {
+    if (_reentry_guard) {
+      _ok = (!*_reentry_guard);
+      *_reentry_guard = true;
+    }
+  }
+
+  ~ReentryGuard() {
+    if (_ok) {
+      *_reentry_guard = false;
+    }
+  }
+
+  bool register_guard(bool *reentry_guard) {
+    if (_reentry_guard) {
+      // not supported (already registered to other bool)
+      return false;
+    }
+    if (reentry_guard) {
+      _reentry_guard = reentry_guard;
+      _ok = (!*_reentry_guard);
+      *_reentry_guard = true;
+    }
+    return _ok;
+  }
+
+  explicit operator bool() const { return _ok; }
+
+  ReentryGuard(const ReentryGuard &) = delete;
+  ReentryGuard &operator=(const ReentryGuard &) = delete;
+
+private:
+  bool *_reentry_guard;
+  bool _ok;
+};
+
+} // namespace ddprof