@samitouri / QOSAMI-WSL-2
Code Commits Issues Pulls Sessions CI/CD
master
master
h 288 lines 7.13 KB
Raw
1 // Copyright (C) Microsoft Corporation. All rights reserved.
2 #pragma once
3
4 #include "p9await.h"
5 #include "p9tracelogging.h"
6
7 namespace p9fs {
8
9 struct IoResult
10 {
11 int Error;
12 size_t BytesTransferred;
13 };
14
15 struct CoroutineIoOperation final : public ICancellable
16 {
17 aiocb ControlBlock;
18 IoResult Result;
19 std::coroutine_handle<> Coroutine{};
20 std::atomic<bool> DoneOrCoroutine{false};
21
22 void Cancel() override
23 {
24 aio_cancel(ControlBlock.aio_fildes, &ControlBlock);
25 }
26 };
27
28 struct CoroutineEpollOperation final : public ICancellable
29 {
30 std::atomic<int> Result{EWOULDBLOCK};
31 std::coroutine_handle<> Coroutine;
32
33 void Resume(int result)
34 {
35 if (SetResult(result))
36 {
37 g_Scheduler.Schedule(Coroutine);
38 }
39 }
40
41 bool SetResult(int result)
42 {
43 int expected = EWOULDBLOCK;
44 return Result.compare_exchange_strong(expected, result);
45 }
46
47 void Cancel() override
48 {
49 Resume(ECANCELED);
50 }
51 };
52
53 struct CoroutineIoIssuer
54 {
55 private:
56 struct Awaiter
57 {
58 CoroutineIoOperation& m_Operation;
59 CancelToken& m_Token;
60
61 bool await_ready() const
62 {
63 return m_Operation.DoneOrCoroutine;
64 }
65
66 bool await_suspend(std::coroutine_handle<> handle)
67 {
68 WI_ASSERT(m_Operation.Coroutine == nullptr);
69
70 m_Operation.Coroutine = handle;
71 if (m_Operation.DoneOrCoroutine.exchange(true))
72 {
73 return false;
74 }
75
76 return true;
77 }
78
79 IoResult await_resume() const
80 {
81 m_Token.Unregister();
82 return m_Operation.Result;
83 }
84 };
85
86 public:
87 CoroutineIoIssuer() = default;
88 CoroutineIoIssuer(int fd);
89
90 explicit operator bool() const
91 {
92 return m_FileDescriptor >= 0;
93 }
94
95 template <class T>
96 Awaiter Issue(CoroutineIoOperation& operation, CancelToken& token, T&& func)
97 {
98 if (PreIssue(operation, token))
99 {
100 IoResult result;
101 try
102 {
103 result = func(operation.ControlBlock);
104 }
105 catch (...)
106 {
107 IssueFailed(token);
108 throw;
109 }
110
111 PostIssue(operation, token, result);
112 }
113
114 return Awaiter{operation, token};
115 }
116
117 private:
118 static void Callback(sigval value);
119
120 bool PreIssue(CoroutineIoOperation& operation, CancelToken& token);
121 static void IssueFailed(CancelToken& token);
122 void PostIssue(CoroutineIoOperation& operation, CancelToken& token, IoResult result);
123
124 int m_FileDescriptor{-1};
125 };
126
127 // Class that handles suspending and resuming operations based on EPOLLIN and EPOLLOUT events.
128 class EpollDispatcher
129 {
130 public:
131 bool Register(int event, CoroutineEpollOperation& operation);
132 void Remove(int event);
133 void Notify(int events);
134
135 private:
136 std::mutex m_lock;
137 int m_currentEvents{};
138 CoroutineEpollOperation* m_outOperation{};
139 CoroutineEpollOperation* m_inOperation{};
140 };
141
142 class EpollWatcher
143 {
144 public:
145 void Run();
146 void Add(int fd, int events, EpollDispatcher& dispatcher);
147 void Remove(int fd);
148
149 explicit operator bool() const noexcept
150 {
151 return m_EpollFileDescriptor >= 0;
152 }
153
154 private:
155 static void WatchThread(EpollWatcher* watcher);
156
157 int m_EpollFileDescriptor{-1};
158 };
159
160 extern EpollWatcher g_Watcher;
161
162 class CoroutineEpollIssuer
163 {
164 private:
165 struct Awaiter
166 {
167 EpollDispatcher& Dispatcher;
168 int Fd;
169 CoroutineEpollOperation& Operation;
170 int Events;
171 CancelToken& Token;
172
173 static constexpr bool await_ready() noexcept
174 {
175 return false;
176 }
177
178 bool await_suspend(std::coroutine_handle<> handle)
179 {
180 Operation.Result = EWOULDBLOCK;
181 Operation.Coroutine = handle;
182
183 // Check if the operation needs to be suspended.
184 if (!Dispatcher.Register(Events, Operation))
185 {
186 return false;
187 }
188
189 // Register the operation for cancellation only if it actually needs to suspend.
190 if (!Token.Register(Operation))
191 {
192 // If the operation is already cancelled, attempt to mark it cancelled. The function
193 // must return true if that fails, because it means the dispatcher already scheduled
194 // it for resumption.
195 return !Operation.SetResult(ECANCELED);
196 }
197
198 return true;
199 }
200
201 int await_resume() const
202 {
203 // If the function was resumed by cancellation, make sure the dispatcher can't access
204 // it after it goes out of scope.
205 // N.B. If the operation was not registered with the dispatcher, calling remove is a
206 // no-op.
207 if (Operation.Result == ECANCELED)
208 {
209 Dispatcher.Remove(Events);
210 }
211
212 // Unregister from the cancel token.
213 // N.B. If the operation was not registered with the token, calling remove is a no-op.
214 Token.Unregister();
215 return Operation.Result;
216 }
217 };
218
219 public:
220 CoroutineEpollIssuer(EpollWatcher& watcher) : m_Watcher{watcher}
221 {
222 }
223
224 ~CoroutineEpollIssuer()
225 {
226 Reset();
227 }
228
229 explicit operator bool() const
230 {
231 return m_FileDescriptor >= 0;
232 }
233
234 template <typename TResult, typename T>
235 Task<TResult> Issue(CoroutineEpollOperation& operation, CancelToken& token, int events, T&& func)
236 {
237 for (;;)
238 {
239 TResult result = func(m_FileDescriptor);
240 if (result >= 0)
241 {
242 co_return result;
243 }
244
245 if (errno != EWOULDBLOCK)
246 {
247 co_return -errno;
248 }
249
250 // Wait for the epoll notification before retrying.
251 int waitResult = co_await Awaiter{m_Dispatcher, m_FileDescriptor, operation, events, token};
252 if (waitResult != 0)
253 {
254 co_return -waitResult;
255 }
256 }
257 }
258
259 void Reset(int fd = -1)
260 {
261 // If there is an existing file descriptor, remove it from the epoll.
262 if (m_FileDescriptor >= 0)
263 {
264 m_Watcher.Remove(m_FileDescriptor);
265 }
266
267 m_FileDescriptor = fd;
268
269 // Add the new file descriptor to epoll.
270 if (m_FileDescriptor >= 0)
271 {
272 m_Watcher.Add(m_FileDescriptor, EPOLLIN | EPOLLOUT | EPOLLET, m_Dispatcher);
273 }
274 }
275
276 private:
277 EpollWatcher& m_Watcher;
278 EpollDispatcher m_Dispatcher;
279 int m_FileDescriptor{-1};
280 };
281
282 Task<int> AcceptAsync(CoroutineEpollIssuer& listen, CancelToken& token);
283 Task<size_t> RecvAsync(CoroutineEpollIssuer& socket, gsl::span<gsl::byte> buffer, CancelToken& token);
284 Task<size_t> SendAsync(CoroutineEpollIssuer& socket, gsl::span<const gsl::byte> buffer, CancelToken& token);
285 Task<IoResult> ReadAsync(CoroutineIoIssuer& file, std::uint64_t offset, gsl::span<gsl::byte> buffer, CancelToken& token);
286 Task<IoResult> WriteAsync(CoroutineIoIssuer& file, std::uint64_t offset, gsl::span<const gsl::byte> buffer, CancelToken& token);
287
288 } // namespace p9fs