1 // Copyright 2023 The Chromium Authors. All rights reserved.
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4
5 #include "quiche/common/btree_scheduler.h"
6
7 #include <optional>
8
9 #include "absl/status/status.h"
10 #include "absl/status/statusor.h"
11 #include "absl/types/span.h"
12 #include "quiche/common/platform/api/quiche_test.h"
13 #include "quiche/common/test_tools/quiche_test_utils.h"
14
15 namespace quiche::test {
16 namespace {
17
18 using ::testing::ElementsAre;
19 using ::testing::Optional;
20
21 template <typename Id, typename Priority>
ScheduleIds(BTreeScheduler<Id,Priority> & scheduler,absl::Span<const Id> ids)22 void ScheduleIds(BTreeScheduler<Id, Priority>& scheduler,
23 absl::Span<const Id> ids) {
24 for (Id id : ids) {
25 QUICHE_EXPECT_OK(scheduler.Schedule(id));
26 }
27 }
28
29 template <typename Id, typename Priority>
PopAll(BTreeScheduler<Id,Priority> & scheduler)30 std::vector<Id> PopAll(BTreeScheduler<Id, Priority>& scheduler) {
31 std::vector<Id> result;
32 result.reserve(scheduler.NumScheduled());
33 for (;;) {
34 absl::StatusOr<Id> id = scheduler.PopFront();
35 if (id.ok()) {
36 result.push_back(*id);
37 } else {
38 EXPECT_THAT(id, StatusIs(absl::StatusCode::kNotFound));
39 break;
40 }
41 }
42 return result;
43 }
44
TEST(BTreeSchedulerTest,SimplePop)45 TEST(BTreeSchedulerTest, SimplePop) {
46 BTreeScheduler<int, int> scheduler;
47 QUICHE_EXPECT_OK(scheduler.Register(1, 100));
48 QUICHE_EXPECT_OK(scheduler.Register(2, 101));
49 QUICHE_EXPECT_OK(scheduler.Register(3, 102));
50
51 EXPECT_THAT(scheduler.GetPriorityFor(1), Optional(100));
52 EXPECT_THAT(scheduler.GetPriorityFor(3), Optional(102));
53 EXPECT_EQ(scheduler.GetPriorityFor(5), std::nullopt);
54
55 EXPECT_EQ(scheduler.NumScheduled(), 0u);
56 EXPECT_FALSE(scheduler.HasScheduled());
57 QUICHE_EXPECT_OK(scheduler.Schedule(1));
58 QUICHE_EXPECT_OK(scheduler.Schedule(2));
59 QUICHE_EXPECT_OK(scheduler.Schedule(3));
60 EXPECT_EQ(scheduler.NumScheduled(), 3u);
61 EXPECT_TRUE(scheduler.HasScheduled());
62
63 EXPECT_THAT(scheduler.PopFront(), IsOkAndHolds(3));
64 EXPECT_THAT(scheduler.PopFront(), IsOkAndHolds(2));
65 EXPECT_THAT(scheduler.PopFront(), IsOkAndHolds(1));
66
67 QUICHE_EXPECT_OK(scheduler.Schedule(2));
68 QUICHE_EXPECT_OK(scheduler.Schedule(1));
69 QUICHE_EXPECT_OK(scheduler.Schedule(3));
70
71 EXPECT_THAT(scheduler.PopFront(), IsOkAndHolds(3));
72 EXPECT_THAT(scheduler.PopFront(), IsOkAndHolds(2));
73 EXPECT_THAT(scheduler.PopFront(), IsOkAndHolds(1));
74
75 QUICHE_EXPECT_OK(scheduler.Schedule(3));
76 QUICHE_EXPECT_OK(scheduler.Schedule(1));
77
78 EXPECT_THAT(scheduler.PopFront(), IsOkAndHolds(3));
79 EXPECT_THAT(scheduler.PopFront(), IsOkAndHolds(1));
80 }
81
TEST(BTreeSchedulerTest,FIFO)82 TEST(BTreeSchedulerTest, FIFO) {
83 BTreeScheduler<int, int> scheduler;
84 QUICHE_EXPECT_OK(scheduler.Register(1, 100));
85 QUICHE_EXPECT_OK(scheduler.Register(2, 100));
86 QUICHE_EXPECT_OK(scheduler.Register(3, 100));
87
88 ScheduleIds(scheduler, {2, 1, 3});
89 EXPECT_THAT(PopAll(scheduler), ElementsAre(2, 1, 3));
90
91 QUICHE_EXPECT_OK(scheduler.Register(4, 101));
92 QUICHE_EXPECT_OK(scheduler.Register(5, 99));
93
94 ScheduleIds(scheduler, {5, 1, 2, 3, 4});
95 EXPECT_THAT(PopAll(scheduler), ElementsAre(4, 1, 2, 3, 5));
96 ScheduleIds(scheduler, {1, 5, 2, 4, 3});
97 EXPECT_THAT(PopAll(scheduler), ElementsAre(4, 1, 2, 3, 5));
98 ScheduleIds(scheduler, {3, 5, 2, 4, 1});
99 EXPECT_THAT(PopAll(scheduler), ElementsAre(4, 3, 2, 1, 5));
100 ScheduleIds(scheduler, {3, 2, 1, 2, 3});
101 EXPECT_THAT(PopAll(scheduler), ElementsAre(3, 2, 1));
102 }
103
TEST(BTreeSchedulerTest,NumEntriesInRange)104 TEST(BTreeSchedulerTest, NumEntriesInRange) {
105 BTreeScheduler<int, int> scheduler;
106 QUICHE_EXPECT_OK(scheduler.Register(1, 0));
107 QUICHE_EXPECT_OK(scheduler.Register(2, 0));
108 QUICHE_EXPECT_OK(scheduler.Register(3, 0));
109 QUICHE_EXPECT_OK(scheduler.Register(4, -2));
110 QUICHE_EXPECT_OK(scheduler.Register(5, -5));
111 QUICHE_EXPECT_OK(scheduler.Register(6, 10));
112 QUICHE_EXPECT_OK(scheduler.Register(7, 16));
113 QUICHE_EXPECT_OK(scheduler.Register(8, 32));
114 QUICHE_EXPECT_OK(scheduler.Register(9, 64));
115
116 EXPECT_EQ(scheduler.NumScheduled(), 0u);
117 EXPECT_EQ(scheduler.NumScheduledInPriorityRange(std::nullopt, std::nullopt),
118 0u);
119 EXPECT_EQ(scheduler.NumScheduledInPriorityRange(-1, 1), 0u);
120
121 for (int stream = 1; stream <= 9; ++stream) {
122 QUICHE_ASSERT_OK(scheduler.Schedule(stream));
123 }
124
125 EXPECT_EQ(scheduler.NumScheduled(), 9u);
126 EXPECT_EQ(scheduler.NumScheduledInPriorityRange(std::nullopt, std::nullopt),
127 9u);
128 EXPECT_EQ(scheduler.NumScheduledInPriorityRange(0, 0), 3u);
129 EXPECT_EQ(scheduler.NumScheduledInPriorityRange(std::nullopt, -1), 2u);
130 EXPECT_EQ(scheduler.NumScheduledInPriorityRange(1, std::nullopt), 4u);
131 }
132
TEST(BTreeSchedulerTest,Registration)133 TEST(BTreeSchedulerTest, Registration) {
134 BTreeScheduler<int, int> scheduler;
135 QUICHE_EXPECT_OK(scheduler.Register(1, 0));
136 QUICHE_EXPECT_OK(scheduler.Register(2, 0));
137
138 QUICHE_EXPECT_OK(scheduler.Schedule(1));
139 QUICHE_EXPECT_OK(scheduler.Schedule(2));
140 EXPECT_EQ(scheduler.NumScheduled(), 2u);
141 EXPECT_TRUE(scheduler.IsScheduled(2));
142
143 EXPECT_THAT(scheduler.Register(2, 0),
144 StatusIs(absl::StatusCode::kAlreadyExists));
145 QUICHE_EXPECT_OK(scheduler.Unregister(2));
146 EXPECT_EQ(scheduler.NumScheduled(), 1u);
147 EXPECT_FALSE(scheduler.IsScheduled(2));
148
149 EXPECT_THAT(scheduler.UpdatePriority(2, 1234),
150 StatusIs(absl::StatusCode::kNotFound));
151 EXPECT_THAT(scheduler.Unregister(2), StatusIs(absl::StatusCode::kNotFound));
152 EXPECT_THAT(scheduler.Schedule(2), StatusIs(absl::StatusCode::kNotFound));
153 QUICHE_EXPECT_OK(scheduler.Register(2, 0));
154 EXPECT_EQ(scheduler.NumScheduled(), 1u);
155 EXPECT_TRUE(scheduler.IsScheduled(1));
156 EXPECT_FALSE(scheduler.IsScheduled(2));
157 }
158
TEST(BTreeSchedulerTest,UpdatePriorityUp)159 TEST(BTreeSchedulerTest, UpdatePriorityUp) {
160 BTreeScheduler<int, int> scheduler;
161 QUICHE_EXPECT_OK(scheduler.Register(1, 0));
162 QUICHE_EXPECT_OK(scheduler.Register(2, 0));
163 QUICHE_EXPECT_OK(scheduler.Register(3, 0));
164
165 ScheduleIds(scheduler, {1, 2, 3});
166 QUICHE_EXPECT_OK(scheduler.UpdatePriority(2, 1000));
167 EXPECT_THAT(PopAll(scheduler), ElementsAre(2, 1, 3));
168 }
169
TEST(BTreeSchedulerTest,UpdatePriorityDown)170 TEST(BTreeSchedulerTest, UpdatePriorityDown) {
171 BTreeScheduler<int, int> scheduler;
172 QUICHE_EXPECT_OK(scheduler.Register(1, 0));
173 QUICHE_EXPECT_OK(scheduler.Register(2, 0));
174 QUICHE_EXPECT_OK(scheduler.Register(3, 0));
175
176 ScheduleIds(scheduler, {1, 2, 3});
177 QUICHE_EXPECT_OK(scheduler.UpdatePriority(2, -1000));
178 EXPECT_THAT(PopAll(scheduler), ElementsAre(1, 3, 2));
179 }
180
TEST(BTreeSchedulerTest,UpdatePriorityEqual)181 TEST(BTreeSchedulerTest, UpdatePriorityEqual) {
182 BTreeScheduler<int, int> scheduler;
183 QUICHE_EXPECT_OK(scheduler.Register(1, 0));
184 QUICHE_EXPECT_OK(scheduler.Register(2, 0));
185 QUICHE_EXPECT_OK(scheduler.Register(3, 0));
186
187 ScheduleIds(scheduler, {1, 2, 3});
188 QUICHE_EXPECT_OK(scheduler.UpdatePriority(2, 0));
189 EXPECT_THAT(PopAll(scheduler), ElementsAre(1, 2, 3));
190 }
191
TEST(BTreeSchedulerTest,UpdatePriorityIntoSameBucket)192 TEST(BTreeSchedulerTest, UpdatePriorityIntoSameBucket) {
193 BTreeScheduler<int, int> scheduler;
194 QUICHE_EXPECT_OK(scheduler.Register(1, 0));
195 QUICHE_EXPECT_OK(scheduler.Register(2, -100));
196 QUICHE_EXPECT_OK(scheduler.Register(3, 0));
197
198 ScheduleIds(scheduler, {1, 2, 3});
199 QUICHE_EXPECT_OK(scheduler.UpdatePriority(2, 0));
200 EXPECT_THAT(PopAll(scheduler), ElementsAre(1, 2, 3));
201 }
202
TEST(BTreeSchedulerTest,ShouldYield)203 TEST(BTreeSchedulerTest, ShouldYield) {
204 BTreeScheduler<int, int> scheduler;
205 QUICHE_EXPECT_OK(scheduler.Register(10, 100));
206 QUICHE_EXPECT_OK(scheduler.Register(20, 101));
207 QUICHE_EXPECT_OK(scheduler.Register(21, 101));
208 QUICHE_EXPECT_OK(scheduler.Register(30, 102));
209
210 EXPECT_THAT(scheduler.ShouldYield(10), IsOkAndHolds(false));
211 EXPECT_THAT(scheduler.ShouldYield(20), IsOkAndHolds(false));
212 EXPECT_THAT(scheduler.ShouldYield(21), IsOkAndHolds(false));
213 EXPECT_THAT(scheduler.ShouldYield(30), IsOkAndHolds(false));
214 EXPECT_THAT(scheduler.ShouldYield(40), StatusIs(absl::StatusCode::kNotFound));
215
216 QUICHE_EXPECT_OK(scheduler.Schedule(20));
217
218 EXPECT_THAT(scheduler.ShouldYield(10), IsOkAndHolds(true));
219 EXPECT_THAT(scheduler.ShouldYield(20), IsOkAndHolds(false));
220 EXPECT_THAT(scheduler.ShouldYield(21), IsOkAndHolds(true));
221 EXPECT_THAT(scheduler.ShouldYield(30), IsOkAndHolds(false));
222 }
223
224 struct CustomPriority {
225 int a;
226 int b;
227
operator <quiche::test::__anonb13aa43e0111::CustomPriority228 bool operator<(const CustomPriority& other) const {
229 return std::make_tuple(a, b) < std::make_tuple(other.a, other.b);
230 }
231 };
232
TEST(BTreeSchedulerTest,CustomPriority)233 TEST(BTreeSchedulerTest, CustomPriority) {
234 BTreeScheduler<int, CustomPriority> scheduler;
235 QUICHE_EXPECT_OK(scheduler.Register(10, CustomPriority{0, 1}));
236 QUICHE_EXPECT_OK(scheduler.Register(11, CustomPriority{0, 0}));
237 QUICHE_EXPECT_OK(scheduler.Register(12, CustomPriority{0, 0}));
238 QUICHE_EXPECT_OK(scheduler.Register(13, CustomPriority{10, 0}));
239 QUICHE_EXPECT_OK(scheduler.Register(14, CustomPriority{-10, 0}));
240
241 ScheduleIds(scheduler, {10, 11, 12, 13, 14});
242 EXPECT_THAT(PopAll(scheduler), ElementsAre(13, 10, 11, 12, 14));
243 }
244
245 struct CustomId {
246 int a;
247 std::string b;
248
operator ==quiche::test::__anonb13aa43e0111::CustomId249 bool operator==(const CustomId& other) const {
250 return a == other.a && b == other.b;
251 }
252
253 template <typename H>
AbslHashValue(H h,const CustomId & c)254 friend H AbslHashValue(H h, const CustomId& c) {
255 return H::combine(std::move(h), c.a, c.b);
256 }
257 };
258
operator <<(std::ostream & os,const CustomId & id)259 std::ostream& operator<<(std::ostream& os, const CustomId& id) {
260 os << id.a << ":" << id.b;
261 return os;
262 }
263
TEST(BTreeSchedulerTest,CustomIds)264 TEST(BTreeSchedulerTest, CustomIds) {
265 BTreeScheduler<CustomId, int> scheduler;
266 QUICHE_EXPECT_OK(scheduler.Register(CustomId{1, "foo"}, 10));
267 QUICHE_EXPECT_OK(scheduler.Register(CustomId{1, "bar"}, 12));
268 QUICHE_EXPECT_OK(scheduler.Register(CustomId{2, "foo"}, 11));
269 EXPECT_THAT(scheduler.Register(CustomId{1, "foo"}, 10),
270 StatusIs(absl::StatusCode::kAlreadyExists));
271
272 ScheduleIds(scheduler,
273 {CustomId{1, "foo"}, CustomId{1, "bar"}, CustomId{2, "foo"}});
274 EXPECT_THAT(scheduler.ShouldYield(CustomId{1, "foo"}), IsOkAndHolds(true));
275 EXPECT_THAT(scheduler.ShouldYield(CustomId{1, "bar"}), IsOkAndHolds(false));
276 EXPECT_THAT(
277 PopAll(scheduler),
278 ElementsAre(CustomId{1, "bar"}, CustomId{2, "foo"}, CustomId{1, "foo"}));
279 }
280
281 } // namespace
282 } // namespace quiche::test
283