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