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
24 #ifndef ABSL_PROFILING_INTERNAL_SAMPLE_RECORDER_H_
25 #define ABSL_PROFILING_INTERNAL_SAMPLE_RECORDER_H_
26
27 #include <atomic>
28 #include <cstddef>
29 #include <functional>
30
31 #include "absl/base/config.h"
32 #include "absl/base/thread_annotations.h"
33 #include "absl/synchronization/mutex.h"
34 #include "absl/time/time.h"
35
36 namespace absl {
37 ABSL_NAMESPACE_BEGIN
38 namespace profiling_internal {
39
40 // Sample<T> that has members required for linking samples in the linked list of
41 // samples maintained by the SampleRecorder. Type T defines the sampled data.
42 template <typename T>
43 struct Sample {
44 // Guards the ability to restore the sample to a pristine state. This
45 // prevents races with sampling and resurrecting an object.
46 absl::Mutex init_mu;
47 T* next = nullptr;
48 T* dead ABSL_GUARDED_BY(init_mu) = nullptr;
49 int64_t weight; // How many sampling events were required to sample this one.
50 };
51
52 // Holds samples and their associated stack traces with a soft limit of
53 // `SetHashtablezMaxSamples()`.
54 //
55 // Thread safe.
56 template <typename T>
57 class SampleRecorder {
58 public:
59 SampleRecorder();
60 ~SampleRecorder();
61
62 // Registers for sampling. Returns an opaque registration info.
63 template <typename... Targs>
64 T* Register(Targs&&... args);
65
66 // Unregisters the sample.
67 void Unregister(T* sample);
68
69 // The dispose callback will be called on all samples the moment they are
70 // being unregistered. Only affects samples that are unregistered after the
71 // callback has been set.
72 // Returns the previous callback.
73 using DisposeCallback = void (*)(const T&);
74 DisposeCallback SetDisposeCallback(DisposeCallback f);
75
76 // Iterates over all the registered `StackInfo`s. Returning the number of
77 // samples that have been dropped.
78 int64_t Iterate(const std::function<void(const T& stack)>& f);
79
80 size_t GetMaxSamples() const;
81 void SetMaxSamples(size_t max);
82
83 private:
84 void PushNew(T* sample);
85 void PushDead(T* sample);
86 template <typename... Targs>
87 T* PopDead(Targs... args);
88
89 std::atomic<size_t> dropped_samples_;
90 std::atomic<size_t> size_estimate_;
91 std::atomic<size_t> max_samples_{1 << 20};
92
93 // Intrusive lock free linked lists for tracking samples.
94 //
95 // `all_` records all samples (they are never removed from this list) and is
96 // terminated with a `nullptr`.
97 //
98 // `graveyard_.dead` is a circular linked list. When it is empty,
99 // `graveyard_.dead == &graveyard`. The list is circular so that
100 // every item on it (even the last) has a non-null dead pointer. This allows
101 // `Iterate` to determine if a given sample is live or dead using only
102 // information on the sample itself.
103 //
104 // For example, nodes [A, B, C, D, E] with [A, C, E] alive and [B, D] dead
105 // looks like this (G is the Graveyard):
106 //
107 // +---+ +---+ +---+ +---+ +---+
108 // all -->| A |--->| B |--->| C |--->| D |--->| E |
109 // | | | | | | | | | |
110 // +---+ | | +->| |-+ | | +->| |-+ | |
111 // | G | +---+ | +---+ | +---+ | +---+ | +---+
112 // | | | | | |
113 // | | --------+ +--------+ |
114 // +---+ |
115 // ^ |
116 // +--------------------------------------+
117 //
118 std::atomic<T*> all_;
119 T graveyard_;
120
121 std::atomic<DisposeCallback> dispose_;
122 };
123
124 template <typename T>
125 typename SampleRecorder<T>::DisposeCallback
SetDisposeCallback(DisposeCallback f)126 SampleRecorder<T>::SetDisposeCallback(DisposeCallback f) {
127 return dispose_.exchange(f, std::memory_order_relaxed);
128 }
129
130 template <typename T>
SampleRecorder()131 SampleRecorder<T>::SampleRecorder()
132 : dropped_samples_(0), size_estimate_(0), all_(nullptr), dispose_(nullptr) {
133 absl::MutexLock l(&graveyard_.init_mu);
134 graveyard_.dead = &graveyard_;
135 }
136
137 template <typename T>
~SampleRecorder()138 SampleRecorder<T>::~SampleRecorder() {
139 T* s = all_.load(std::memory_order_acquire);
140 while (s != nullptr) {
141 T* next = s->next;
142 delete s;
143 s = next;
144 }
145 }
146
147 template <typename T>
PushNew(T * sample)148 void SampleRecorder<T>::PushNew(T* sample) {
149 sample->next = all_.load(std::memory_order_relaxed);
150 while (!all_.compare_exchange_weak(sample->next, sample,
151 std::memory_order_release,
152 std::memory_order_relaxed)) {
153 }
154 }
155
156 template <typename T>
PushDead(T * sample)157 void SampleRecorder<T>::PushDead(T* sample) {
158 if (auto* dispose = dispose_.load(std::memory_order_relaxed)) {
159 dispose(*sample);
160 }
161
162 absl::MutexLock graveyard_lock(&graveyard_.init_mu);
163 absl::MutexLock sample_lock(&sample->init_mu);
164 sample->dead = graveyard_.dead;
165 graveyard_.dead = sample;
166 }
167
168 template <typename T>
169 template <typename... Targs>
PopDead(Targs...args)170 T* SampleRecorder<T>::PopDead(Targs... args) {
171 absl::MutexLock graveyard_lock(&graveyard_.init_mu);
172
173 // The list is circular, so eventually it collapses down to
174 // graveyard_.dead == &graveyard_
175 // when it is empty.
176 T* sample = graveyard_.dead;
177 if (sample == &graveyard_) return nullptr;
178
179 absl::MutexLock sample_lock(&sample->init_mu);
180 graveyard_.dead = sample->dead;
181 sample->dead = nullptr;
182 sample->PrepareForSampling(std::forward<Targs>(args)...);
183 return sample;
184 }
185
186 template <typename T>
187 template <typename... Targs>
Register(Targs &&...args)188 T* SampleRecorder<T>::Register(Targs&&... args) {
189 size_t size = size_estimate_.fetch_add(1, std::memory_order_relaxed);
190 if (size > max_samples_.load(std::memory_order_relaxed)) {
191 size_estimate_.fetch_sub(1, std::memory_order_relaxed);
192 dropped_samples_.fetch_add(1, std::memory_order_relaxed);
193 return nullptr;
194 }
195
196 T* sample = PopDead(args...);
197 if (sample == nullptr) {
198 // Resurrection failed. Hire a new warlock.
199 sample = new T();
200 {
201 absl::MutexLock sample_lock(&sample->init_mu);
202 // If flag initialization happens to occur (perhaps in another thread)
203 // while in this block, it will lock `graveyard_` which is usually always
204 // locked before any sample. This will appear as a lock inversion.
205 // However, this code is run exactly once per sample, and this sample
206 // cannot be accessed until after it is returned from this method. This
207 // means that this lock state can never be recreated, so we can safely
208 // inform the deadlock detector to ignore it.
209 sample->init_mu.ForgetDeadlockInfo();
210 sample->PrepareForSampling(std::forward<Targs>(args)...);
211 }
212 PushNew(sample);
213 }
214
215 return sample;
216 }
217
218 template <typename T>
Unregister(T * sample)219 void SampleRecorder<T>::Unregister(T* sample) {
220 PushDead(sample);
221 size_estimate_.fetch_sub(1, std::memory_order_relaxed);
222 }
223
224 template <typename T>
Iterate(const std::function<void (const T & stack)> & f)225 int64_t SampleRecorder<T>::Iterate(
226 const std::function<void(const T& stack)>& f) {
227 T* s = all_.load(std::memory_order_acquire);
228 while (s != nullptr) {
229 absl::MutexLock l(&s->init_mu);
230 if (s->dead == nullptr) {
231 f(*s);
232 }
233 s = s->next;
234 }
235
236 return dropped_samples_.load(std::memory_order_relaxed);
237 }
238
239 template <typename T>
SetMaxSamples(size_t max)240 void SampleRecorder<T>::SetMaxSamples(size_t max) {
241 max_samples_.store(max, std::memory_order_release);
242 }
243
244 template <typename T>
GetMaxSamples()245 size_t SampleRecorder<T>::GetMaxSamples() const {
246 return max_samples_.load(std::memory_order_acquire);
247 }
248
249 } // namespace profiling_internal
250 ABSL_NAMESPACE_END
251 } // namespace absl
252
253 #endif // ABSL_PROFILING_INTERNAL_SAMPLE_RECORDER_H_
254