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_async.h"
16
17 #include <mutex>
18
19 #include "pw_assert/check.h" // IWYU pragma: keep
20 #include "pw_async2/dispatcher.h" // IWYU pragma: keep
21 #include "pw_multibuf/allocator.h"
22
23 namespace pw::multibuf {
24
25 using ::pw::async2::Context;
26 using ::pw::async2::Poll;
27 using ::pw::async2::Waker;
28
29 // ########## MultiBufAllocatorAsync
30
AllocateAsync(size_t size)31 MultiBufAllocationFuture MultiBufAllocatorAsync::AllocateAsync(size_t size) {
32 return MultiBufAllocationFuture(
33 mbuf_allocator_, size, size, kAllowDiscontiguous);
34 }
35
AllocateAsync(size_t min_size,size_t desired_size)36 MultiBufAllocationFuture MultiBufAllocatorAsync::AllocateAsync(
37 size_t min_size, size_t desired_size) {
38 return MultiBufAllocationFuture(
39 mbuf_allocator_, min_size, desired_size, kAllowDiscontiguous);
40 }
41
AllocateContiguousAsync(size_t size)42 MultiBufAllocationFuture MultiBufAllocatorAsync::AllocateContiguousAsync(
43 size_t size) {
44 return MultiBufAllocationFuture(
45 mbuf_allocator_, size, size, kNeedsContiguous);
46 }
47
AllocateContiguousAsync(size_t min_size,size_t desired_size)48 MultiBufAllocationFuture MultiBufAllocatorAsync::AllocateContiguousAsync(
49 size_t min_size, size_t desired_size) {
50 return MultiBufAllocationFuture(
51 mbuf_allocator_, min_size, desired_size, kNeedsContiguous);
52 }
53
54 // ########## MultiBufAllocationFuture
55
HandleMemoryAvailable(MultiBufAllocator & alloc,size_t size_available,size_t contiguous_size_available) const56 bool MultiBufAllocationFuture::HandleMemoryAvailable(
57 MultiBufAllocator& alloc,
58 size_t size_available,
59 size_t contiguous_size_available) const {
60 PW_CHECK_PTR_EQ(allocator_, &alloc);
61 bool should_wake_and_remove =
62 ((min_size_ <= contiguous_size_available) ||
63 (contiguity_requirement_ == kAllowDiscontiguous &&
64 min_size_ <= size_available));
65 if (should_wake_and_remove) {
66 std::move(const_cast<Waker&>(waker_)).Wake();
67 }
68 return should_wake_and_remove;
69 }
70
MultiBufAllocationFuture(MultiBufAllocationFuture && other)71 MultiBufAllocationFuture::MultiBufAllocationFuture(
72 MultiBufAllocationFuture&& other)
73 : allocator_(other.allocator_),
74 waker_(),
75 min_size_(other.min_size_),
76 desired_size_(other.desired_size_),
77 contiguity_requirement_(other.contiguity_requirement_) {
78 std::lock_guard lock(allocator_->lock_);
79 if (!other.unlisted()) {
80 allocator_->RemoveMemoryAvailableDelegate(other);
81 allocator_->AddMemoryAvailableDelegate(*this);
82 // We must move the waker under the lock in order to ensure that there is no
83 // race between swapping ``MultiBufAllocationFuture``s and the waker being
84 // awoken by the allocator.
85 waker_ = std::move(other.waker_);
86 }
87 }
88
operator =(MultiBufAllocationFuture && other)89 MultiBufAllocationFuture& MultiBufAllocationFuture::operator=(
90 MultiBufAllocationFuture&& other) {
91 {
92 std::lock_guard lock(allocator_->lock_);
93 if (!this->unlisted()) {
94 allocator_->RemoveMemoryAvailableDelegate(*this);
95 }
96 }
97
98 allocator_ = other.allocator_;
99 min_size_ = other.min_size_;
100 desired_size_ = other.desired_size_;
101 contiguity_requirement_ = other.contiguity_requirement_;
102
103 std::lock_guard lock(allocator_->lock_);
104 if (!other.unlisted()) {
105 allocator_->RemoveMemoryAvailableDelegate(other);
106 allocator_->AddMemoryAvailableDelegate(*this);
107 // We must move the waker under the lock in order to ensure that there is no
108 // race between swapping ``MultiBufAllocationFuture``s and the waker being
109 // awoken by the allocator.
110 waker_ = std::move(other.waker_);
111 }
112
113 return *this;
114 }
115
~MultiBufAllocationFuture()116 MultiBufAllocationFuture::~MultiBufAllocationFuture() {
117 std::lock_guard lock(allocator_->lock_);
118 if (!this->unlisted()) {
119 allocator_->RemoveMemoryAvailableDelegate(*this);
120 }
121 }
122
SetDesiredSizes(size_t new_min_size,size_t new_desired_size,ContiguityRequirement new_contiguity_requirement)123 void MultiBufAllocationFuture::SetDesiredSizes(
124 size_t new_min_size,
125 size_t new_desired_size,
126 ContiguityRequirement new_contiguity_requirement) {
127 // No-op if the sizes are unchanged.
128 if (new_min_size == min_size_ && new_desired_size == desired_size_ &&
129 new_contiguity_requirement == contiguity_requirement_) {
130 return;
131 }
132 // Acquire the lock so the allocator doesn't touch the sizes while we're
133 // modifying them.
134 std::lock_guard lock(allocator_->lock_);
135
136 // If our needs decreased, try allocating again next time rather than
137 // waiting for a wake.
138 if (new_min_size < min_size_ ||
139 ((new_contiguity_requirement == kAllowDiscontiguous) &&
140 (contiguity_requirement_ == kNeedsContiguous))) {
141 if (!this->unlisted()) {
142 allocator_->RemoveMemoryAvailableDelegate(*this);
143 }
144 }
145 min_size_ = new_min_size;
146 desired_size_ = new_desired_size;
147 contiguity_requirement_ = new_contiguity_requirement;
148 }
149
Pend(Context & cx)150 Poll<std::optional<MultiBuf>> MultiBufAllocationFuture::Pend(Context& cx) {
151 std::lock_guard lock(allocator_->lock_);
152 // If we're still listed waiting for a wakeup, don't bother to try again.
153 if (this->unlisted()) {
154 auto result = TryAllocate();
155 if (result.IsReady()) {
156 return result;
157 }
158 allocator_->AddMemoryAvailableDelegate(*this);
159 }
160 // We set the waker while still holding the lock to ensure there is no gap
161 // between us checking TryAllocate above and the waker being reset here.
162 PW_ASYNC_STORE_WAKER(
163 cx,
164 waker_,
165 "MultiBufAllocationFuture is waiting for memory to become available");
166 return async2::Pending();
167 }
168
TryAllocate()169 Poll<std::optional<MultiBuf>> MultiBufAllocationFuture::TryAllocate() {
170 Result<MultiBuf> buf_opt =
171 allocator_->DoAllocate(min_size_, desired_size_, contiguity_requirement_);
172 if (buf_opt.ok()) {
173 return async2::Ready<std::optional<MultiBuf>>(std::move(*buf_opt));
174 }
175 if (buf_opt.status().IsOutOfRange()) {
176 return async2::Ready<std::optional<MultiBuf>>(std::nullopt);
177 }
178 return async2::Pending();
179 }
180
181 } // namespace pw::multibuf
182