1 // Copyright 2024 The Pigweed Authors
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License"); you may not
4 // use this file except in compliance with the License. You may obtain a copy of
5 // 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, WITHOUT
11 // WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
12 // License for the specific language governing permissions and limitations under
13 // the License.
14
15 #include "pw_multibuf/allocator.h"
16
17 #include "gtest/gtest.h"
18 #include "pw_async2/dispatcher.h"
19 #include "pw_async2/poll.h"
20 #include "pw_multibuf/allocator_async.h"
21
22 namespace pw::multibuf {
23 namespace {
24
25 using ::pw::async2::Context;
26 using ::pw::async2::Dispatcher;
27 using ::pw::async2::Pending;
28 using ::pw::async2::Poll;
29 using ::pw::async2::Ready;
30 using ::pw::async2::Task;
31 using ::pw::multibuf::ContiguityRequirement;
32 using ::pw::multibuf::kAllowDiscontiguous;
33
34 struct AllocateExpectation {
35 size_t min_size;
36 size_t desired_size;
37 ContiguityRequirement contiguous;
38 pw::Result<MultiBuf> result;
39 };
40
41 class MockMultiBufAllocator : public MultiBufAllocator {
42 public:
~MockMultiBufAllocator()43 ~MockMultiBufAllocator() override {
44 // All expectations should have been met and removed.
45 EXPECT_EQ(expected_allocate_, std::nullopt);
46 }
47
ExpectAllocateAndReturn(size_t min_size,size_t desired_size,ContiguityRequirement contiguous,pw::Result<MultiBuf> result)48 void ExpectAllocateAndReturn(size_t min_size,
49 size_t desired_size,
50 ContiguityRequirement contiguous,
51 pw::Result<MultiBuf> result) {
52 // Multiple simultaneous expectations are not supported.
53 ASSERT_FALSE(expected_allocate_.has_value());
54 expected_allocate_ = {
55 min_size, desired_size, contiguous, std::move(result)};
56 }
57
58 using MultiBufAllocator::MoreMemoryAvailable;
59
60 private:
DoAllocate(size_t min_size,size_t desired_size,ContiguityRequirement contiguous)61 pw::Result<MultiBuf> DoAllocate(size_t min_size,
62 size_t desired_size,
63 ContiguityRequirement contiguous) final {
64 EXPECT_NE(expected_allocate_, std::nullopt);
65 if (!expected_allocate_.has_value()) {
66 return Status::FailedPrecondition();
67 }
68 AllocateExpectation expected = std::move(*expected_allocate_);
69 expected_allocate_ = std::nullopt;
70 EXPECT_EQ(min_size, expected.min_size);
71 EXPECT_EQ(desired_size, expected.desired_size);
72 EXPECT_EQ(contiguous, expected.contiguous);
73 return std::move(expected.result);
74 }
DoGetBackingCapacity()75 std::optional<size_t> DoGetBackingCapacity() final { return {}; }
76
77 std::optional<AllocateExpectation> expected_allocate_;
78 };
79
80 // ########## MultiBufAllocatorAsync
81
82 class AllocateTask : public Task {
83 public:
AllocateTask(MultiBufAllocationFuture && future)84 AllocateTask(MultiBufAllocationFuture&& future)
85 : future_(std::move(future)), last_result_(Pending()) {}
86
87 MultiBufAllocationFuture future_;
88 Poll<std::optional<MultiBuf>> last_result_;
89
90 private:
DoPend(Context & cx)91 Poll<> DoPend(Context& cx) override {
92 last_result_ = future_.Pend(cx);
93 if (last_result_.IsReady()) {
94 return Ready();
95 }
96 return Pending();
97 }
98 };
99
TEST(MultiBufAllocatorAsync,MultiBufAllocationFutureCtor)100 TEST(MultiBufAllocatorAsync, MultiBufAllocationFutureCtor) {
101 MockMultiBufAllocator mbuf_alloc;
102 MultiBufAllocatorAsync async_alloc{mbuf_alloc};
103 {
104 MultiBufAllocationFuture fut = async_alloc.AllocateAsync(44u, 33u);
105 EXPECT_EQ(44u, fut.min_size());
106 EXPECT_EQ(33u, fut.desired_size());
107 EXPECT_FALSE(fut.needs_contiguous());
108 EXPECT_EQ(&mbuf_alloc, &fut.allocator());
109 }
110 {
111 MultiBufAllocationFuture fut =
112 async_alloc.AllocateContiguousAsync(66u, 55u);
113 EXPECT_EQ(66u, fut.min_size());
114 EXPECT_EQ(55u, fut.desired_size());
115 EXPECT_TRUE(fut.needs_contiguous());
116 EXPECT_EQ(&mbuf_alloc, &fut.allocator());
117 }
118 }
119
TEST(MultiBufAllocatorAsync,AllocateAsyncReturnsImmediatelyAvailableAllocation)120 TEST(MultiBufAllocatorAsync,
121 AllocateAsyncReturnsImmediatelyAvailableAllocation) {
122 MockMultiBufAllocator mbuf_alloc;
123 MultiBufAllocatorAsync async_alloc{mbuf_alloc};
124
125 AllocateTask task(async_alloc.AllocateAsync(44, 33));
126 mbuf_alloc.ExpectAllocateAndReturn(44, 33, kAllowDiscontiguous, MultiBuf());
127
128 Dispatcher dispatcher;
129 dispatcher.Post(task);
130 EXPECT_EQ(dispatcher.RunUntilStalled(), Ready());
131
132 ASSERT_TRUE(task.last_result_.IsReady());
133 ASSERT_TRUE(task.last_result_->has_value());
134 }
135
TEST(MultiBufAllocatorAsync,AllocateAsyncWillNotPollUntilMoreMemoryAvailable)136 TEST(MultiBufAllocatorAsync, AllocateAsyncWillNotPollUntilMoreMemoryAvailable) {
137 MockMultiBufAllocator mbuf_alloc;
138 MultiBufAllocatorAsync async_alloc{mbuf_alloc};
139
140 AllocateTask task(async_alloc.AllocateAsync(44, 33));
141 Dispatcher dispatcher;
142 dispatcher.Post(task);
143
144 // First attempt will return `ResourceExhausted` to signal temporary OOM.
145 mbuf_alloc.ExpectAllocateAndReturn(
146 44, 33, kAllowDiscontiguous, Status::ResourceExhausted());
147 EXPECT_TRUE(dispatcher.RunUntilStalled().IsPending());
148 EXPECT_TRUE(task.last_result_.IsPending());
149
150 // Re-running the dispatcher should not poll the pending task since it has
151 // not been awoken. `AllocateAndReturn` should *not* be called.
152 EXPECT_TRUE(dispatcher.RunUntilStalled().IsPending());
153
154 // Insufficient memory should not awaken the task.
155 mbuf_alloc.MoreMemoryAvailable(30, 30);
156 EXPECT_TRUE(dispatcher.RunUntilStalled().IsPending());
157
158 // Sufficient memory will awaken and return the memory
159 mbuf_alloc.MoreMemoryAvailable(50, 50);
160 mbuf_alloc.ExpectAllocateAndReturn(44, 33, kAllowDiscontiguous, MultiBuf());
161 EXPECT_TRUE(dispatcher.RunUntilStalled().IsReady());
162 }
163
TEST(MultiBufAllocatorAsync,MoveMultiBufAllocationFuture)164 TEST(MultiBufAllocatorAsync, MoveMultiBufAllocationFuture) {
165 MockMultiBufAllocator mbuf_alloc;
166 MultiBufAllocatorAsync async_alloc{mbuf_alloc};
167 MultiBufAllocationFuture fut1 = async_alloc.AllocateAsync(44u, 33u);
168 EXPECT_EQ(44u, fut1.min_size());
169 EXPECT_EQ(33u, fut1.desired_size());
170
171 // Test move ctor
172 MultiBufAllocationFuture fut2{std::move(fut1)};
173 EXPECT_EQ(44u, fut2.min_size());
174 EXPECT_EQ(33u, fut2.desired_size());
175
176 // Test move assign
177 MultiBufAllocationFuture fut3{std::move(fut2)};
178 EXPECT_EQ(44u, fut3.min_size());
179 EXPECT_EQ(33u, fut3.desired_size());
180
181 // Test task behavior works after two moves
182 AllocateTask task(std::move(fut3));
183 mbuf_alloc.ExpectAllocateAndReturn(44, 33, kAllowDiscontiguous, MultiBuf());
184
185 Dispatcher dispatcher;
186 dispatcher.Post(task);
187 EXPECT_EQ(dispatcher.RunUntilStalled(), Ready());
188
189 ASSERT_TRUE(task.last_result_.IsReady());
190 ASSERT_TRUE(task.last_result_->has_value());
191 }
192
193 } // namespace
194 } // namespace pw::multibuf
195