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_async2/coro.h"
16
17 #include "pw_allocator/null_allocator.h"
18 #include "pw_allocator/testing.h"
19 #include "pw_async2/dispatcher_base.h"
20 #include "pw_status/status.h"
21
22 namespace {
23
24 using ::pw::OkStatus;
25 using ::pw::Result;
26 using ::pw::Status;
27 using ::pw::allocator::Allocator;
28 using ::pw::allocator::GetNullAllocator;
29 using ::pw::allocator::test::AllocatorForTest;
30 using ::pw::async2::Context;
31 using ::pw::async2::Coro;
32 using ::pw::async2::CoroContext;
33 using ::pw::async2::Dispatcher;
34 using ::pw::async2::Pending;
35 using ::pw::async2::Poll;
36 using ::pw::async2::Ready;
37 using ::pw::async2::Task;
38 using ::pw::async2::Waker;
39
40 class ExpectCoroTask final : public Task {
41 public:
ExpectCoroTask(Coro<pw::Status> && coro)42 ExpectCoroTask(Coro<pw::Status>&& coro) : coro_(std::move(coro)) {}
43
44 private:
DoPend(Context & cx)45 Poll<> DoPend(Context& cx) final {
46 Poll<Status> result = coro_.Pend(cx);
47 if (result.IsPending()) {
48 return Pending();
49 }
50 EXPECT_EQ(*result, OkStatus());
51 return Ready();
52 }
53 Coro<pw::Status> coro_;
54 };
55
ImmediatelyReturnsFive(CoroContext &)56 Coro<Result<int>> ImmediatelyReturnsFive(CoroContext&) { co_return 5; }
57
StoresFiveThenReturns(CoroContext & coro_cx,int & out)58 Coro<Status> StoresFiveThenReturns(CoroContext& coro_cx, int& out) {
59 PW_CO_TRY_ASSIGN(out, co_await ImmediatelyReturnsFive(coro_cx));
60 co_return OkStatus();
61 }
62
TEST(CoroTest,BasicFunctionsWithoutYieldingRun)63 TEST(CoroTest, BasicFunctionsWithoutYieldingRun) {
64 AllocatorForTest<256> alloc;
65 CoroContext coro_cx(alloc);
66 int output = 0;
67 ExpectCoroTask task = StoresFiveThenReturns(coro_cx, output);
68 Dispatcher dispatcher;
69 dispatcher.Post(task);
70 EXPECT_TRUE(dispatcher.RunUntilStalled().IsReady());
71 EXPECT_EQ(output, 5);
72 }
73
TEST(CoroTest,AllocationFailureProducesInvalidCoro)74 TEST(CoroTest, AllocationFailureProducesInvalidCoro) {
75 CoroContext coro_cx(GetNullAllocator());
76 EXPECT_FALSE(ImmediatelyReturnsFive(coro_cx).IsValid());
77 int x = 0;
78 EXPECT_FALSE(StoresFiveThenReturns(coro_cx, x).IsValid());
79 }
80
81 struct MockPendable {
MockPendable__anond25c46c40111::MockPendable82 MockPendable() : poll_count(0), return_value(Pending()), last_waker() {}
83 MockPendable(const MockPendable&) = delete;
84 MockPendable& operator=(const MockPendable&) = delete;
85 MockPendable(MockPendable&&) = delete;
86 MockPendable& operator=(MockPendable&&) = delete;
87
Pend__anond25c46c40111::MockPendable88 Poll<int> Pend(Context& cx) {
89 ++poll_count;
90 last_waker = cx.GetWaker(pw::async2::WaitReason::Unspecified());
91 return return_value;
92 }
93
94 int poll_count;
95 Poll<int> return_value;
96 Waker last_waker;
97 };
98
AddTwo(CoroContext &,MockPendable & a,MockPendable & b)99 Coro<Result<int>> AddTwo(CoroContext&, MockPendable& a, MockPendable& b) {
100 co_return co_await a + co_await b;
101 }
102
AddTwoThenStore(CoroContext & alloc,MockPendable & a,MockPendable & b,int & out)103 Coro<Status> AddTwoThenStore(CoroContext& alloc,
104 MockPendable& a,
105 MockPendable& b,
106 int& out) {
107 PW_CO_TRY_ASSIGN(out, co_await AddTwo(alloc, a, b));
108 co_return OkStatus();
109 }
110
TEST(CoroTest,AwaitMultipleAndAwakenRuns)111 TEST(CoroTest, AwaitMultipleAndAwakenRuns) {
112 AllocatorForTest<512> alloc;
113 CoroContext coro_cx(alloc);
114 MockPendable a;
115 MockPendable b;
116 int output = 0;
117 ExpectCoroTask task = AddTwoThenStore(coro_cx, a, b, output);
118 Dispatcher dispatcher;
119 dispatcher.Post(task);
120
121 EXPECT_TRUE(dispatcher.RunUntilStalled().IsPending());
122 EXPECT_EQ(a.poll_count, 1);
123 EXPECT_EQ(b.poll_count, 0);
124
125 EXPECT_TRUE(dispatcher.RunUntilStalled().IsPending());
126 EXPECT_EQ(a.poll_count, 1);
127 EXPECT_EQ(b.poll_count, 0);
128
129 int a_value = 4;
130 a.return_value = a_value;
131 std::move(a.last_waker).Wake();
132 EXPECT_TRUE(dispatcher.RunUntilStalled().IsPending());
133 EXPECT_EQ(a.poll_count, 2);
134 EXPECT_EQ(b.poll_count, 1);
135
136 int b_value = 5;
137 b.return_value = b_value;
138 std::move(b.last_waker).Wake();
139 EXPECT_TRUE(dispatcher.RunUntilStalled().IsReady());
140 EXPECT_EQ(a.poll_count, 2);
141 EXPECT_EQ(b.poll_count, 2);
142 EXPECT_EQ(output, a_value + b_value);
143 }
144
145 } // namespace
146