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