1 // Copyright 2025 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_async2/pend_func_awaitable.h"
16
17 #include <optional>
18 #include <utility>
19
20 #include "pw_allocator/testing.h"
21 #include "pw_async2/coro.h"
22 #include "pw_async2/coro_or_else_task.h"
23 #include "pw_async2/dispatcher.h"
24 #include "pw_async2/poll.h"
25 #include "pw_function/function.h"
26 #include "pw_status/status.h"
27 #include "pw_unit_test/framework.h"
28
29 namespace {
30
31 using ::pw::Function;
32 using ::pw::OkStatus;
33 using ::pw::Result;
34 using ::pw::Status;
35 using ::pw::allocator::test::AllocatorForTest;
36 using ::pw::async2::Context;
37 using ::pw::async2::Coro;
38 using ::pw::async2::CoroContext;
39 using ::pw::async2::CoroOrElseTask;
40 using ::pw::async2::Dispatcher;
41 using ::pw::async2::PendFuncAwaitable;
42 using ::pw::async2::Pending;
43 using ::pw::async2::Poll;
44 using ::pw::async2::Ready;
45 using ::pw::async2::Waker;
46
47 template <typename T>
48 class Mailbox {
49 public:
PendGetValue(Context & cx)50 Poll<Result<T>> PendGetValue(Context& cx) {
51 ++poll_count_;
52 if (value_) {
53 auto v = *value_;
54 value_.reset();
55 return Ready(v);
56 }
57 PW_ASYNC_STORE_WAKER(cx, waker_, "Mailbox is waiting for a value");
58 return Pending();
59 }
60
SetValue(T v)61 void SetValue(T v) {
62 value_ = v;
63 std::move(waker_).Wake();
64 }
65
PollCount()66 int PollCount() { return poll_count_; }
67
68 private:
69 std::optional<T> value_;
70 int poll_count_ = 0;
71 Waker waker_;
72 };
73
74 template <typename T>
ReadMailbox(CoroContext &,Mailbox<T> & mailbox,int & out)75 Coro<Status> ReadMailbox(CoroContext&, Mailbox<T>& mailbox, int& out) {
76 PW_CO_TRY_ASSIGN(out, co_await PendFuncAwaitable([&](Context& cx) {
77 return mailbox.PendGetValue(cx);
78 }));
79 co_return OkStatus();
80 }
81
TEST(PendFuncAwaitable,TestMailbox)82 TEST(PendFuncAwaitable, TestMailbox) {
83 Mailbox<int> mailbox;
84
85 AllocatorForTest<256> alloc;
86 CoroContext coro_cx(alloc);
87 int output = 0;
88 bool error_handler_did_run = false;
89 CoroOrElseTask task(
90 ReadMailbox(coro_cx, mailbox, output),
91 [&error_handler_did_run](Status) { error_handler_did_run = true; });
92
93 Dispatcher dispatcher;
94 dispatcher.Post(task);
95
96 EXPECT_EQ(mailbox.PollCount(), 0);
97 EXPECT_EQ(dispatcher.RunUntilStalled(), Pending());
98 EXPECT_EQ(mailbox.PollCount(), 1);
99
100 // Unwoken mailbox is not polled.
101 EXPECT_EQ(dispatcher.RunUntilStalled(), Pending());
102 EXPECT_EQ(mailbox.PollCount(), 1);
103
104 mailbox.SetValue(5);
105 EXPECT_EQ(dispatcher.RunUntilStalled(), Ready());
106 EXPECT_EQ(mailbox.PollCount(), 2);
107 EXPECT_EQ(output, 5);
108 EXPECT_FALSE(error_handler_did_run);
109 }
110
ReturnsReady8(Context &)111 Poll<int> ReturnsReady8(Context&) { return Ready(8); }
112
TEST(PendFuncAwaitable,TestTemplateDeductionAndSize)113 TEST(PendFuncAwaitable, TestTemplateDeductionAndSize) {
114 // A PendFuncAwaitable with an unspecified Func template parameter will
115 // default to pw::Function. This allows the same container to hold a variety
116 // of different callables, but it may either reserve extra inline storage or
117 // dynamically allocate momeory, depending on how pw::Function is configured.
118 std::optional<PendFuncAwaitable<int>> a;
119 a.emplace([](Context&) -> Poll<int> { return Ready(4); });
120 a.emplace(&ReturnsReady8);
121 static_assert(sizeof(decltype(a)::value_type::CallableType) ==
122 sizeof(Function<Poll<int>(Context&)>));
123
124 // When constructing a PendFuncAwaitable directly from a callable, CTAD will
125 // match the Func template parameter to that of the callable. This has the
126 // benefit of reducing the amount of storage needed vs that of a pw::Function.
127 //
128 // A lambda without any captures doesn't require any storage.
129 auto b = PendFuncAwaitable([](Context&) -> Poll<int> { return Ready(4); });
130 static_assert(sizeof(decltype(b)::CallableType) <= 1);
131
132 // A lambda with captures requires storage to hold the captures.
133 int scratch = 6;
134 auto c = PendFuncAwaitable(
135 [&scratch](Context&) -> Poll<int> { return Ready(scratch); });
136 static_assert(sizeof(decltype(c)::CallableType) == sizeof(&scratch));
137
138 // A raw function pointer just needs storage for the pointer value.
139 auto d = PendFuncAwaitable(&ReturnsReady8);
140 static_assert(sizeof(decltype(d)::CallableType) == sizeof(&ReturnsReady8));
141 }
142
143 } // namespace
144