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
21 namespace pw::multibuf {
22 namespace {
23
24 using ::pw::async2::Context;
25 using ::pw::async2::Dispatcher;
26 using ::pw::async2::Pending;
27 using ::pw::async2::Poll;
28 using ::pw::async2::Ready;
29 using ::pw::async2::Task;
30
31 struct AllocateExpectation {
32 size_t min_size;
33 size_t desired_size;
34 bool contiguous;
35 pw::Result<MultiBuf> result;
36 };
37
38 class MockMultiBufAllocator : public MultiBufAllocator {
39 public:
~MockMultiBufAllocator()40 ~MockMultiBufAllocator() override {
41 // All expectations should have been met and removed.
42 EXPECT_EQ(expected_allocate_, std::nullopt);
43 }
44
ExpectAllocateAndReturn(size_t min_size,size_t desired_size,bool contiguous,pw::Result<MultiBuf> result)45 void ExpectAllocateAndReturn(size_t min_size,
46 size_t desired_size,
47 bool contiguous,
48 pw::Result<MultiBuf> result) {
49 // Multiple simultaneous expectations are not supported.
50 ASSERT_FALSE(expected_allocate_.has_value());
51 expected_allocate_ = {
52 min_size, desired_size, contiguous, std::move(result)};
53 }
54
55 using MultiBufAllocator::MoreMemoryAvailable;
56
57 private:
DoAllocate(size_t min_size,size_t desired_size,bool contiguous)58 pw::Result<MultiBuf> DoAllocate(size_t min_size,
59 size_t desired_size,
60 bool contiguous) final {
61 EXPECT_NE(expected_allocate_, std::nullopt);
62 if (!expected_allocate_.has_value()) {
63 return Status::FailedPrecondition();
64 }
65 AllocateExpectation expected = std::move(*expected_allocate_);
66 expected_allocate_ = std::nullopt;
67 EXPECT_EQ(min_size, expected.min_size);
68 EXPECT_EQ(desired_size, expected.desired_size);
69 EXPECT_EQ(contiguous, expected.contiguous);
70 return std::move(expected.result);
71 }
72
73 std::optional<AllocateExpectation> expected_allocate_;
74 };
75
76 class AllocateTask : public Task {
77 public:
AllocateTask(MultiBufAllocationFuture && future)78 AllocateTask(MultiBufAllocationFuture&& future)
79 : future_(std::move(future)), last_result_(Pending()) {}
80
81 MultiBufAllocationFuture future_;
82 Poll<std::optional<MultiBuf>> last_result_;
83
84 private:
DoPend(Context & cx)85 Poll<> DoPend(Context& cx) override {
86 last_result_ = future_.Pend(cx);
87 if (last_result_.IsReady()) {
88 return Ready();
89 }
90 return Pending();
91 }
92 };
93
TEST(MultiBufAllocator,AllocateAsyncReturnsImmediatelyAvailableAllocation)94 TEST(MultiBufAllocator, AllocateAsyncReturnsImmediatelyAvailableAllocation) {
95 MockMultiBufAllocator alloc;
96 AllocateTask task(alloc.AllocateAsync(44, 33));
97 alloc.ExpectAllocateAndReturn(44, 33, false, MultiBuf());
98
99 Dispatcher dispatcher;
100 dispatcher.Post(task);
101 EXPECT_EQ(dispatcher.RunUntilStalled(), Ready());
102
103 ASSERT_TRUE(task.last_result_.IsReady());
104 ASSERT_TRUE(task.last_result_->has_value());
105 }
106
TEST(MultiBufAllocator,AllocateAsyncWillNotPollUntilMoreMemoryAvailable)107 TEST(MultiBufAllocator, AllocateAsyncWillNotPollUntilMoreMemoryAvailable) {
108 MockMultiBufAllocator alloc;
109 AllocateTask task(alloc.AllocateAsync(44, 33));
110 Dispatcher dispatcher;
111 dispatcher.Post(task);
112
113 // First attempt will return `ResourceExhausted` to signal temporary OOM.
114 alloc.ExpectAllocateAndReturn(44, 33, false, Status::ResourceExhausted());
115 EXPECT_TRUE(dispatcher.RunUntilStalled().IsPending());
116 EXPECT_TRUE(task.last_result_.IsPending());
117
118 // Re-running the dispatcher should not poll the pending task since it has
119 // not been awoken. `AllocateAndReturn` should *not* be called.
120 EXPECT_TRUE(dispatcher.RunUntilStalled().IsPending());
121
122 // Insufficient memory should not awaken the task.
123 alloc.MoreMemoryAvailable(30, 30);
124 EXPECT_TRUE(dispatcher.RunUntilStalled().IsPending());
125
126 // Sufficient memory will awaken and return the memory
127 alloc.MoreMemoryAvailable(50, 50);
128 alloc.ExpectAllocateAndReturn(44, 33, false, MultiBuf());
129 EXPECT_TRUE(dispatcher.RunUntilStalled().IsReady());
130 }
131
132 } // namespace
133 } // namespace pw::multibuf
134