sample_recorder.h 7.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230
  1. // Copyright 2018 The Abseil Authors.
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // https://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. //
  15. // -----------------------------------------------------------------------------
  16. // File: sample_recorder.h
  17. // -----------------------------------------------------------------------------
  18. //
  19. // This header file defines a lock-free linked list for recording samples
  20. // collected from a random/stochastic process.
  21. //
  22. // This utility is internal-only. Use at your own risk.
  23. #ifndef ABSL_PROFILING_INTERNAL_SAMPLE_RECORDER_H_
  24. #define ABSL_PROFILING_INTERNAL_SAMPLE_RECORDER_H_
  25. #include <atomic>
  26. #include <cstddef>
  27. #include <functional>
  28. #include "absl/base/config.h"
  29. #include "absl/base/thread_annotations.h"
  30. #include "absl/synchronization/mutex.h"
  31. #include "absl/time/time.h"
  32. namespace absl {
  33. ABSL_NAMESPACE_BEGIN
  34. namespace profiling_internal {
  35. // Sample<T> that has members required for linking samples in the linked list of
  36. // samples maintained by the SampleRecorder. Type T defines the sampled data.
  37. template <typename T>
  38. struct Sample {
  39. // Guards the ability to restore the sample to a pristine state. This
  40. // prevents races with sampling and resurrecting an object.
  41. absl::Mutex init_mu;
  42. T* next = nullptr;
  43. T* dead ABSL_GUARDED_BY(init_mu) = nullptr;
  44. };
  45. // Holds samples and their associated stack traces with a soft limit of
  46. // `SetHashtablezMaxSamples()`.
  47. //
  48. // Thread safe.
  49. template <typename T>
  50. class SampleRecorder {
  51. public:
  52. SampleRecorder();
  53. ~SampleRecorder();
  54. // Registers for sampling. Returns an opaque registration info.
  55. T* Register();
  56. // Unregisters the sample.
  57. void Unregister(T* sample);
  58. // The dispose callback will be called on all samples the moment they are
  59. // being unregistered. Only affects samples that are unregistered after the
  60. // callback has been set.
  61. // Returns the previous callback.
  62. using DisposeCallback = void (*)(const T&);
  63. DisposeCallback SetDisposeCallback(DisposeCallback f);
  64. // Iterates over all the registered `StackInfo`s. Returning the number of
  65. // samples that have been dropped.
  66. int64_t Iterate(const std::function<void(const T& stack)>& f);
  67. void SetMaxSamples(int32_t max);
  68. private:
  69. void PushNew(T* sample);
  70. void PushDead(T* sample);
  71. T* PopDead();
  72. std::atomic<size_t> dropped_samples_;
  73. std::atomic<size_t> size_estimate_;
  74. std::atomic<int32_t> max_samples_{1 << 20};
  75. // Intrusive lock free linked lists for tracking samples.
  76. //
  77. // `all_` records all samples (they are never removed from this list) and is
  78. // terminated with a `nullptr`.
  79. //
  80. // `graveyard_.dead` is a circular linked list. When it is empty,
  81. // `graveyard_.dead == &graveyard`. The list is circular so that
  82. // every item on it (even the last) has a non-null dead pointer. This allows
  83. // `Iterate` to determine if a given sample is live or dead using only
  84. // information on the sample itself.
  85. //
  86. // For example, nodes [A, B, C, D, E] with [A, C, E] alive and [B, D] dead
  87. // looks like this (G is the Graveyard):
  88. //
  89. // +---+ +---+ +---+ +---+ +---+
  90. // all -->| A |--->| B |--->| C |--->| D |--->| E |
  91. // | | | | | | | | | |
  92. // +---+ | | +->| |-+ | | +->| |-+ | |
  93. // | G | +---+ | +---+ | +---+ | +---+ | +---+
  94. // | | | | | |
  95. // | | --------+ +--------+ |
  96. // +---+ |
  97. // ^ |
  98. // +--------------------------------------+
  99. //
  100. std::atomic<T*> all_;
  101. T graveyard_;
  102. std::atomic<DisposeCallback> dispose_;
  103. };
  104. template <typename T>
  105. typename SampleRecorder<T>::DisposeCallback
  106. SampleRecorder<T>::SetDisposeCallback(DisposeCallback f) {
  107. return dispose_.exchange(f, std::memory_order_relaxed);
  108. }
  109. template <typename T>
  110. SampleRecorder<T>::SampleRecorder()
  111. : dropped_samples_(0), size_estimate_(0), all_(nullptr), dispose_(nullptr) {
  112. absl::MutexLock l(&graveyard_.init_mu);
  113. graveyard_.dead = &graveyard_;
  114. }
  115. template <typename T>
  116. SampleRecorder<T>::~SampleRecorder() {
  117. T* s = all_.load(std::memory_order_acquire);
  118. while (s != nullptr) {
  119. T* next = s->next;
  120. delete s;
  121. s = next;
  122. }
  123. }
  124. template <typename T>
  125. void SampleRecorder<T>::PushNew(T* sample) {
  126. sample->next = all_.load(std::memory_order_relaxed);
  127. while (!all_.compare_exchange_weak(sample->next, sample,
  128. std::memory_order_release,
  129. std::memory_order_relaxed)) {
  130. }
  131. }
  132. template <typename T>
  133. void SampleRecorder<T>::PushDead(T* sample) {
  134. if (auto* dispose = dispose_.load(std::memory_order_relaxed)) {
  135. dispose(*sample);
  136. }
  137. absl::MutexLock graveyard_lock(&graveyard_.init_mu);
  138. absl::MutexLock sample_lock(&sample->init_mu);
  139. sample->dead = graveyard_.dead;
  140. graveyard_.dead = sample;
  141. }
  142. template <typename T>
  143. T* SampleRecorder<T>::PopDead() {
  144. absl::MutexLock graveyard_lock(&graveyard_.init_mu);
  145. // The list is circular, so eventually it collapses down to
  146. // graveyard_.dead == &graveyard_
  147. // when it is empty.
  148. T* sample = graveyard_.dead;
  149. if (sample == &graveyard_) return nullptr;
  150. absl::MutexLock sample_lock(&sample->init_mu);
  151. graveyard_.dead = sample->dead;
  152. sample->dead = nullptr;
  153. sample->PrepareForSampling();
  154. return sample;
  155. }
  156. template <typename T>
  157. T* SampleRecorder<T>::Register() {
  158. int64_t size = size_estimate_.fetch_add(1, std::memory_order_relaxed);
  159. if (size > max_samples_.load(std::memory_order_relaxed)) {
  160. size_estimate_.fetch_sub(1, std::memory_order_relaxed);
  161. dropped_samples_.fetch_add(1, std::memory_order_relaxed);
  162. return nullptr;
  163. }
  164. T* sample = PopDead();
  165. if (sample == nullptr) {
  166. // Resurrection failed. Hire a new warlock.
  167. sample = new T();
  168. PushNew(sample);
  169. }
  170. return sample;
  171. }
  172. template <typename T>
  173. void SampleRecorder<T>::Unregister(T* sample) {
  174. PushDead(sample);
  175. size_estimate_.fetch_sub(1, std::memory_order_relaxed);
  176. }
  177. template <typename T>
  178. int64_t SampleRecorder<T>::Iterate(
  179. const std::function<void(const T& stack)>& f) {
  180. T* s = all_.load(std::memory_order_acquire);
  181. while (s != nullptr) {
  182. absl::MutexLock l(&s->init_mu);
  183. if (s->dead == nullptr) {
  184. f(*s);
  185. }
  186. s = s->next;
  187. }
  188. return dropped_samples_.load(std::memory_order_relaxed);
  189. }
  190. template <typename T>
  191. void SampleRecorder<T>::SetMaxSamples(int32_t max) {
  192. max_samples_.store(max, std::memory_order_release);
  193. }
  194. } // namespace profiling_internal
  195. ABSL_NAMESPACE_END
  196. } // namespace absl
  197. #endif // ABSL_PROFILING_INTERNAL_SAMPLE_RECORDER_H_