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 #include "absl/profiling/internal/sample_recorder.h"
16
17 #include <atomic>
18 #include <random>
19 #include <vector>
20
21 #include "gmock/gmock.h"
22 #include "absl/base/thread_annotations.h"
23 #include "absl/synchronization/internal/thread_pool.h"
24 #include "absl/synchronization/mutex.h"
25 #include "absl/synchronization/notification.h"
26 #include "absl/time/time.h"
27
28 namespace absl {
29 ABSL_NAMESPACE_BEGIN
30 namespace profiling_internal {
31
32 namespace {
33 using ::absl::synchronization_internal::ThreadPool;
34 using ::testing::IsEmpty;
35 using ::testing::UnorderedElementsAre;
36
37 struct Info : public Sample<Info> {
38 public:
PrepareForSamplingabsl::profiling_internal::__anon952a8c050111::Info39 void PrepareForSampling(int64_t w) { weight = w; }
40 std::atomic<size_t> size;
41 absl::Time create_time;
42 };
43
GetSizes(SampleRecorder<Info> * s)44 std::vector<size_t> GetSizes(SampleRecorder<Info>* s) {
45 std::vector<size_t> res;
46 s->Iterate([&](const Info& info) {
47 res.push_back(info.size.load(std::memory_order_acquire));
48 });
49 return res;
50 }
51
GetWeights(SampleRecorder<Info> * s)52 std::vector<int64_t> GetWeights(SampleRecorder<Info>* s) {
53 std::vector<int64_t> res;
54 s->Iterate([&](const Info& info) { res.push_back(info.weight); });
55 return res;
56 }
57
Register(SampleRecorder<Info> * s,int64_t weight,size_t size)58 Info* Register(SampleRecorder<Info>* s, int64_t weight, size_t size) {
59 auto* info = s->Register(weight);
60 assert(info != nullptr);
61 info->size.store(size);
62 return info;
63 }
64
TEST(SampleRecorderTest,Registration)65 TEST(SampleRecorderTest, Registration) {
66 SampleRecorder<Info> sampler;
67 auto* info1 = Register(&sampler, 31, 1);
68 EXPECT_THAT(GetSizes(&sampler), UnorderedElementsAre(1));
69 EXPECT_THAT(GetWeights(&sampler), UnorderedElementsAre(31));
70
71 auto* info2 = Register(&sampler, 32, 2);
72 EXPECT_THAT(GetSizes(&sampler), UnorderedElementsAre(1, 2));
73 info1->size.store(3);
74 EXPECT_THAT(GetSizes(&sampler), UnorderedElementsAre(3, 2));
75 EXPECT_THAT(GetWeights(&sampler), UnorderedElementsAre(31, 32));
76
77 sampler.Unregister(info1);
78 sampler.Unregister(info2);
79 }
80
TEST(SampleRecorderTest,Unregistration)81 TEST(SampleRecorderTest, Unregistration) {
82 SampleRecorder<Info> sampler;
83 std::vector<Info*> infos;
84 for (size_t i = 0; i < 3; ++i) {
85 infos.push_back(Register(&sampler, 33 + i, i));
86 }
87 EXPECT_THAT(GetSizes(&sampler), UnorderedElementsAre(0, 1, 2));
88 EXPECT_THAT(GetWeights(&sampler), UnorderedElementsAre(33, 34, 35));
89
90 sampler.Unregister(infos[1]);
91 EXPECT_THAT(GetSizes(&sampler), UnorderedElementsAre(0, 2));
92 EXPECT_THAT(GetWeights(&sampler), UnorderedElementsAre(33, 35));
93
94 infos.push_back(Register(&sampler, 36, 3));
95 infos.push_back(Register(&sampler, 37, 4));
96 EXPECT_THAT(GetSizes(&sampler), UnorderedElementsAre(0, 2, 3, 4));
97 EXPECT_THAT(GetWeights(&sampler), UnorderedElementsAre(33, 35, 36, 37));
98 sampler.Unregister(infos[3]);
99 EXPECT_THAT(GetSizes(&sampler), UnorderedElementsAre(0, 2, 4));
100 EXPECT_THAT(GetWeights(&sampler), UnorderedElementsAre(33, 35, 37));
101
102 sampler.Unregister(infos[0]);
103 sampler.Unregister(infos[2]);
104 sampler.Unregister(infos[4]);
105 EXPECT_THAT(GetSizes(&sampler), IsEmpty());
106 }
107
TEST(SampleRecorderTest,MultiThreaded)108 TEST(SampleRecorderTest, MultiThreaded) {
109 SampleRecorder<Info> sampler;
110 Notification stop;
111 ThreadPool pool(10);
112
113 for (int i = 0; i < 10; ++i) {
114 pool.Schedule([&sampler, &stop, i]() {
115 std::random_device rd;
116 std::mt19937 gen(rd());
117
118 std::vector<Info*> infoz;
119 while (!stop.HasBeenNotified()) {
120 if (infoz.empty()) {
121 infoz.push_back(sampler.Register(i));
122 }
123 switch (std::uniform_int_distribution<>(0, 2)(gen)) {
124 case 0: {
125 infoz.push_back(sampler.Register(i));
126 break;
127 }
128 case 1: {
129 size_t p =
130 std::uniform_int_distribution<>(0, infoz.size() - 1)(gen);
131 Info* info = infoz[p];
132 infoz[p] = infoz.back();
133 infoz.pop_back();
134 EXPECT_EQ(info->weight, i);
135 sampler.Unregister(info);
136 break;
137 }
138 case 2: {
139 absl::Duration oldest = absl::ZeroDuration();
140 sampler.Iterate([&](const Info& info) {
141 oldest = std::max(oldest, absl::Now() - info.create_time);
142 });
143 ASSERT_GE(oldest, absl::ZeroDuration());
144 break;
145 }
146 }
147 }
148 });
149 }
150 // The threads will hammer away. Give it a little bit of time for tsan to
151 // spot errors.
152 absl::SleepFor(absl::Seconds(3));
153 stop.Notify();
154 }
155
TEST(SampleRecorderTest,Callback)156 TEST(SampleRecorderTest, Callback) {
157 SampleRecorder<Info> sampler;
158
159 auto* info1 = Register(&sampler, 39, 1);
160 auto* info2 = Register(&sampler, 40, 2);
161
162 static const Info* expected;
163
164 auto callback = [](const Info& info) {
165 // We can't use `info` outside of this callback because the object will be
166 // disposed as soon as we return from here.
167 EXPECT_EQ(&info, expected);
168 };
169
170 // Set the callback.
171 EXPECT_EQ(sampler.SetDisposeCallback(callback), nullptr);
172 expected = info1;
173 sampler.Unregister(info1);
174
175 // Unset the callback.
176 EXPECT_EQ(callback, sampler.SetDisposeCallback(nullptr));
177 expected = nullptr; // no more calls.
178 sampler.Unregister(info2);
179 }
180
181 } // namespace
182 } // namespace profiling_internal
183 ABSL_NAMESPACE_END
184 } // namespace absl
185