@samitouri / QOSAMI-WSL-1
Code Commits Issues Pulls Sessions CI/CD
master
master
h 927 lines 21.4 KB
Raw
1 // Copyright (C) Microsoft Corporation. All rights reserved.
2 #pragma once
3
4 #include "p9scheduler.h"
5 #include "p9errors.h"
6
7 namespace p9fs {
8
9 // Provides a movable wrapper around std::coroutine_handle that
10 // destroys the coroutine when it goes out of scope.
11 template <class PromiseType>
12 class UniqueCoroutineHandle
13 {
14 public:
15 UniqueCoroutineHandle() = default;
16
17 ~UniqueCoroutineHandle()
18 {
19 Reset();
20 }
21
22 UniqueCoroutineHandle(std::coroutine_handle<PromiseType> handle) : m_Handle(handle)
23 {
24 }
25
26 UniqueCoroutineHandle(UniqueCoroutineHandle&& handle) : m_Handle(handle.m_Handle)
27 {
28 handle.m_Handle = nullptr;
29 }
30
31 UniqueCoroutineHandle& operator=(UniqueCoroutineHandle&& handle)
32 {
33 if (this != &handle)
34 {
35 Reset();
36 m_Handle = handle.m_Handle;
37 handle.m_Handle = nullptr;
38 }
39
40 return *this;
41 }
42
43 UniqueCoroutineHandle(const UniqueCoroutineHandle&) = delete;
44 UniqueCoroutineHandle& operator=(const UniqueCoroutineHandle&) = delete;
45
46 void Reset()
47 {
48 if (m_Handle)
49 {
50 FAIL_FAST_IF(!m_Handle.done());
51 m_Handle.destroy();
52 m_Handle = nullptr;
53 }
54 }
55
56 std::coroutine_handle<PromiseType> Get() const
57 {
58 return m_Handle;
59 }
60
61 explicit operator bool() const
62 {
63 return static_cast<bool>(m_Handle);
64 }
65
66 private:
67 std::coroutine_handle<PromiseType> m_Handle{};
68 };
69
70 // Async task is an awaitable task object whose resources are released
71 // asynchronously from being awaited on. This requires an extra heap allocation,
72 // so only use this when you need to have a task that you don't plan to
73 // immediately await.
74 class AsyncTask
75 {
76 private:
77 struct Storage
78 {
79 std::exception_ptr Exception;
80 std::mutex Mutex;
81 std::condition_variable Condition;
82 std::atomic<bool> Done;
83 std::atomic<void*> Waiter;
84 };
85
86 public:
87 AsyncTask() = default;
88
89 AsyncTask(const std::shared_ptr<Storage>& storage) : m_Storage(storage)
90 {
91 }
92
93 bool await_ready() const
94 {
95 return m_Storage->Done;
96 }
97
98 // Wait for the promise to resume us, unless the coroutine has
99 // already completed.
100 bool await_suspend(std::coroutine_handle<> handle)
101 {
102 m_Storage->Waiter = handle.address();
103 if ((m_Storage->Done || m_Storage->Exception) && m_Storage->Waiter.exchange(nullptr) != nullptr)
104 {
105 return false;
106 }
107
108 return true;
109 }
110
111 // Return the result from the promise object.
112 void await_resume() const
113 {
114 if (m_Storage->Exception)
115 {
116 std::rethrow_exception(m_Storage->Exception);
117 }
118 }
119
120 explicit operator bool()
121 {
122 return bool{m_Storage};
123 }
124
125 void Get()
126 {
127 std::unique_lock<std::mutex> lock{m_Storage->Mutex};
128 m_Storage->Condition.wait(lock, [this]() -> bool { return m_Storage->Done; });
129
130 await_resume();
131 }
132
133 class promise_type
134 {
135 public:
136 std::suspend_never initial_suspend()
137 {
138 return {};
139 }
140
141 std::suspend_never final_suspend() noexcept
142 {
143 return {};
144 }
145
146 void return_void()
147 {
148 Signal();
149 }
150
151 void unhandled_exception()
152 {
153 m_Storage->Exception = std::current_exception();
154 Signal();
155 }
156
157 AsyncTask get_return_object() const
158 {
159 return AsyncTask{m_Storage};
160 }
161
162 private:
163 void Signal()
164 {
165 {
166 std::lock_guard<std::mutex> lock{m_Storage->Mutex};
167 m_Storage->Done = true;
168 }
169
170 m_Storage->Condition.notify_all();
171 auto address = m_Storage->Waiter.exchange(nullptr);
172 if (address)
173 {
174 auto waiter = std::coroutine_handle<>::from_address(address);
175 g_Scheduler.Schedule(waiter);
176 }
177
178 m_Storage.reset();
179 }
180
181 std::shared_ptr<Storage> m_Storage{std::make_shared<Storage>()};
182 };
183
184 private:
185 std::shared_ptr<Storage> m_Storage;
186 };
187
188 template <class T>
189 class PromiseBase : public T
190 {
191 public:
192 auto initial_suspend()
193 {
194 return std::suspend_never{};
195 }
196
197 // When this coroutine function exits, signal the waiter and then suspend so
198 // that the promise stays alive long enough to get the result out of it.
199 auto final_suspend() noexcept
200 {
201 return FinalAwaiter{};
202 }
203
204 // Stores an exception object for rethrowing when awaiting on
205 // the associated task.
206 void unhandled_exception()
207 {
208 m_Exception = std::current_exception();
209 }
210
211 protected:
212 // Rethrows the coroutine's exception if there is one.
213 void CheckException()
214 {
215 if (m_Exception)
216 {
217 std::rethrow_exception(m_Exception);
218 }
219 }
220
221 private:
222 struct FinalAwaiter
223 {
224 static bool await_ready() noexcept
225 {
226 return false;
227 }
228
229 template <class U>
230 void await_suspend(std::coroutine_handle<U> handle) noexcept
231 {
232 handle.promise().Resume();
233 }
234
235 static void await_resume() noexcept
236 {
237 }
238 };
239
240 std::exception_ptr m_Exception;
241 };
242
243 // Promise type for coroutines that return values.
244 template <class T, class Task, class Base>
245 class Promise : public PromiseBase<Base>
246 {
247 public:
248 Task get_return_object()
249 {
250 return {*this};
251 }
252
253 template <class U>
254 void return_value(U&& value)
255 {
256 m_Value = std::forward<U>(value);
257 }
258
259 T GetResult()
260 {
261 PromiseBase<Base>::CheckException();
262 return std::move(m_Value);
263 }
264
265 private:
266 T m_Value;
267 };
268
269 // Promise type for coroutines that do not return values.
270 template <class Task, class Base>
271 class Promise<void, Task, Base> : public PromiseBase<Base>
272 {
273 public:
274 Task get_return_object()
275 {
276 return {*this};
277 }
278
279 void return_void()
280 {
281 }
282
283 void GetResult()
284 {
285 PromiseBase<Base>::CheckException();
286 }
287 };
288
289 // Promise type for Task objects. Represents a coroutine that will eventually
290 // return or throw an exception.
291 class TaskPromise
292 {
293 public:
294 // Sets the unique waiting task. Returns false if the promise already
295 // completed.
296 bool SetWaiter(std::coroutine_handle<> handle)
297 {
298 m_Waiter = handle.address();
299 return !m_WaiterOrDone.exchange(true);
300 }
301
302 // Resumes the waiting task, if there is one.
303 void Resume()
304 {
305 // Set the value to true to indicate we're done. If it was already true,
306 // it indicates a waiter was registered before completion so resume it.
307 if (m_WaiterOrDone.exchange(true))
308 {
309 auto waiter = std::coroutine_handle<>::from_address(m_Waiter);
310 m_Waiter = nullptr;
311 g_Scheduler.Schedule(waiter);
312 }
313 }
314
315 bool Done() const
316 {
317 return m_WaiterOrDone.load(std::memory_order_acquire);
318 }
319
320 private:
321 void* m_Waiter{};
322 std::atomic<bool> m_WaiterOrDone{false};
323 };
324
325 // Task is an awaitable task that is designed to be co_awaited immediately. It
326 // will not release the underlying coroutine's resources until it goes out of
327 // scope.
328 template <class T>
329 class Task
330 {
331 public:
332 using promise_type = Promise<T, Task<T>, TaskPromise>;
333
334 bool await_ready()
335 {
336 return Promise().Done();
337 }
338
339 // Start the coroutine and wait for the promise to resume us.
340 bool await_suspend(std::coroutine_handle<> handle)
341 {
342 // Try to set the waiter. If the coroutine has already completed, abort
343 // the suspend.
344 return Promise().SetWaiter(handle);
345 }
346
347 // Return the result from the promise object.
348 auto await_resume()
349 {
350 return Promise().GetResult();
351 }
352
353 Task() = default;
354
355 Task(promise_type& promise) : m_Coroutine(std::coroutine_handle<promise_type>::from_promise(promise))
356 {
357 }
358
359 private:
360 promise_type& Promise()
361 {
362 return m_Coroutine.Get().promise();
363 }
364
365 UniqueCoroutineHandle<promise_type> m_Coroutine;
366 };
367
368 class ScheduledTask
369 {
370 public:
371 struct promise_type
372 {
373 ScheduledTask get_return_object()
374 {
375 return {*this};
376 }
377
378 std::suspend_always initial_suspend()
379 {
380 return {};
381 }
382
383 std::suspend_never final_suspend() noexcept
384 {
385 return {};
386 }
387
388 void unhandled_exception()
389 {
390 try
391 {
392 std::rethrow_exception(std::current_exception());
393 }
394 catch (...)
395 {
396 FAIL_FAST_CAUGHT_EXCEPTION();
397 }
398 }
399
400 static void return_void()
401 {
402 }
403 };
404
405 ScheduledTask(promise_type& promise)
406 {
407 g_Scheduler.Schedule(std::coroutine_handle<promise_type>::from_promise(promise));
408 }
409 };
410
411 /// Non-awaitable wrapper to schedule a coroutine to run on another thread.
412 template <class T>
413 void RunScheduledTask(T&& awaitable)
414 {
415 [](T func) -> ScheduledTask { co_await std::move(func)(); }(std::forward<T>(awaitable));
416 }
417
418 template <typename T>
419 void RunAsyncTask(T&& awaitable)
420 {
421 [](T func) -> AsyncTask { co_await std::move(func)(); }(std::forward<T>(awaitable));
422 }
423
424 /// Awaitable wrapper to run synchronous blocking code without blocking
425 /// outstanding coroutines.
426 template <class T>
427 auto BlockingCode(T func) -> Task<decltype(func())>
428 {
429 const bool unblock = g_Scheduler.Block();
430 auto result = func();
431 if (unblock)
432 {
433 co_await g_Scheduler.Unblock();
434 }
435
436 co_return result;
437 }
438
439 // Awaitable semaphore.
440 class AsyncSemaphore
441 {
442 public:
443 class AsyncSemaphoreTask
444 {
445 public:
446 AsyncSemaphoreTask(AsyncSemaphore& evt, uint64_t count) : m_Semaphore{evt}, m_Count{count}
447 {
448 }
449
450 bool await_ready() const noexcept
451 {
452 return m_Count == 0;
453 }
454
455 bool await_suspend(std::coroutine_handle<> awaiter) noexcept
456 {
457 m_Awaiter = awaiter;
458 return m_Semaphore.Enqueue(this, m_Count);
459 }
460
461 static void await_resume() noexcept
462 {
463 }
464
465 private:
466 friend class AsyncSemaphore;
467 AsyncSemaphore& m_Semaphore;
468 AsyncSemaphoreTask* m_Next;
469 std::coroutine_handle<> m_Awaiter;
470 uint64_t m_Count;
471 };
472
473 AsyncSemaphore(uint64_t initialCount) : m_Count(initialCount)
474 {
475 }
476
477 AsyncSemaphore(AsyncSemaphore&) = delete;
478
479 AsyncSemaphoreTask Acquire(uint64_t count) noexcept
480 {
481 if (TryAcquire(count))
482 {
483 count = 0;
484 }
485
486 return AsyncSemaphoreTask(*this, count);
487 }
488
489 bool TryAcquire(uint64_t count) noexcept
490 {
491 std::lock_guard<std::mutex> lock{m_Lock};
492 if (m_Count >= count)
493 {
494 m_Count -= count;
495 return true;
496 }
497
498 return false;
499 }
500
501 void Release(uint64_t count) noexcept
502 {
503 std::lock_guard<std::mutex> lock{m_Lock};
504 m_Count += count;
505
506 // Wake all possible waiters.
507 AsyncSemaphoreTask** head = &m_Waiter;
508 if (*head != nullptr)
509 {
510 const auto waiter = *head;
511 if (m_Count >= waiter->m_Count)
512 {
513 m_Count -= waiter->m_Count;
514 *head = waiter->m_Next;
515 g_Scheduler.Schedule(waiter->m_Awaiter);
516 }
517 else
518 {
519 head = &waiter->m_Next;
520 }
521 }
522 }
523
524 private:
525 bool Enqueue(AsyncSemaphoreTask* task, uint64_t count)
526 {
527 std::lock_guard<std::mutex> lock{m_Lock};
528 if (m_Count >= count)
529 {
530 m_Count -= count;
531 return false;
532 }
533
534 task->m_Next = m_Waiter;
535 m_Waiter = task;
536 return true;
537 }
538
539 std::mutex m_Lock;
540 uint64_t m_Count;
541 AsyncSemaphoreTask* m_Waiter{};
542 };
543
544 class AsyncEvent
545 {
546 public:
547 class AsyncEventTask
548 {
549 public:
550 AsyncEventTask(AsyncEvent& evt) : m_Event{evt}
551 {
552 }
553
554 bool await_ready() const noexcept
555 {
556 return m_Event.IsSet();
557 }
558
559 bool await_suspend(std::coroutine_handle<> awaiter) noexcept
560 {
561 m_Awaiter = awaiter;
562 ULONG_PTR state = m_Event.m_State;
563 while (true)
564 {
565 if (state == m_SetState)
566 {
567 return false;
568 }
569
570 m_Next = reinterpret_cast<AsyncEventTask*>(state);
571 if (m_Event.m_State.compare_exchange_weak(state, reinterpret_cast<ULONG_PTR>(this)))
572 {
573 return true;
574 }
575 }
576 }
577
578 static void await_resume() noexcept
579 {
580 }
581
582 private:
583 friend class AsyncEvent;
584 AsyncEvent& m_Event;
585 AsyncEventTask* m_Next;
586 std::coroutine_handle<> m_Awaiter;
587 };
588
589 AsyncEvent() : m_State{m_UnsetState}
590 {
591 }
592
593 AsyncEvent(AsyncEvent&) = delete;
594
595 AsyncEventTask operator co_await() noexcept
596 {
597 return AsyncEventTask(*this);
598 }
599
600 bool IsSet() const noexcept
601 {
602 return m_State == m_SetState;
603 }
604
605 void Set() noexcept
606 {
607 const ULONG_PTR state = m_State.exchange(m_SetState);
608 if (state != m_SetState)
609 {
610 // Resume all waiters.
611 auto task = reinterpret_cast<AsyncEventTask*>(state);
612 while (task != nullptr)
613 {
614 const auto next = task->m_Next;
615 g_Scheduler.Schedule(task->m_Awaiter);
616 task = next;
617 }
618 }
619 }
620
621 void Reset()
622 {
623 ULONG_PTR state = m_SetState;
624 m_State.compare_exchange_strong(state, m_UnsetState);
625 }
626
627 private:
628 static constexpr ULONG_PTR m_SetState = 1;
629 static constexpr ULONG_PTR m_UnsetState = 0;
630
631 // State is NULL if the event is not set with no waiters, 1 if it's set,
632 // or a pointer to the first waiter.
633 std::atomic<ULONG_PTR> m_State;
634 };
635
636 // Mutex lock that can be waited on with co_await.
637 // N.B. Any waiters will be resumed on the thread that calls Unlock()
638 class AsyncLock
639 {
640 public:
641 class AsyncLockTask;
642
643 // RAII class that releases the lock on scope exit.
644 // N.B. Unlike std::lock_guard, it does not acquire the lock; it must be
645 // created when the lock is already owned.
646 class AsyncLockGuard
647 {
648 public:
649 AsyncLockGuard(AsyncLockGuard&& lock) : m_Lock(lock.m_Lock)
650 {
651 lock.m_Lock = nullptr;
652 }
653
654 AsyncLockGuard(const AsyncLockGuard&) = delete;
655
656 ~AsyncLockGuard()
657 {
658 if (m_Lock != nullptr)
659 {
660 m_Lock->Unlock();
661 }
662 }
663
664 private:
665 friend class AsyncLockTask;
666
667 AsyncLockGuard(AsyncLock& lock) : m_Lock{&lock}
668 {
669 }
670
671 AsyncLock* m_Lock;
672 };
673
674 class AsyncLockTask
675 {
676 public:
677 AsyncLockTask(AsyncLock& lock) : m_Lock{lock}, m_Next{nullptr}
678 {
679 }
680
681 static bool await_ready() noexcept
682 {
683 return false;
684 }
685
686 bool await_suspend(std::coroutine_handle<> awaiter)
687 {
688 m_Awaiter = awaiter;
689 ULONG_PTR state = m_Lock.m_State;
690 while (true)
691 {
692 if (state == m_UnlockedState)
693 {
694 // Acquire the lock and complete synchronously if it is
695 // not currently held.
696 if (m_Lock.m_State.compare_exchange_weak(state, m_LockedState))
697 {
698 return false;
699 }
700 }
701 else
702 {
703 // Add this instance to the list of new waiters.
704 if (state != m_LockedState)
705 {
706 m_Next = reinterpret_cast<AsyncLockTask*>(state);
707 }
708 else
709 {
710 // Must reset in case the loop ran more than once.
711 m_Next = nullptr;
712 }
713
714 if (m_Lock.m_State.compare_exchange_weak(state, reinterpret_cast<ULONG_PTR>(this)))
715 {
716 return true;
717 }
718 }
719 }
720 }
721
722 AsyncLockGuard await_resume() const noexcept
723 {
724 return AsyncLockGuard(m_Lock);
725 }
726
727 private:
728 friend class AsyncLock;
729 AsyncLock& m_Lock;
730 AsyncLockTask* m_Next;
731 std::coroutine_handle<> m_Awaiter;
732 };
733
734 AsyncLock() : m_State{m_UnlockedState}, m_WaitList{nullptr}
735 {
736 }
737
738 AsyncLock(const AsyncLock&) = delete;
739
740 AsyncLockTask Lock() noexcept
741 {
742 return AsyncLockTask(*this);
743 }
744
745 bool TryLock() noexcept
746 {
747 ULONG_PTR unlockedState = m_UnlockedState;
748 return m_State.compare_exchange_strong(unlockedState, m_LockedState);
749 }
750
751 void Unlock() noexcept
752 {
753 // If there are no existing waiters and no new waiters, unlock and
754 // return.
755 // N.B. Since m_WaitList is only accessed from the unlock method,
756 // which should only be called with the lock held, it needs no
757 // synchronization.
758 ULONG_PTR state = m_LockedState;
759 if ((m_WaitList == nullptr) && (m_State.compare_exchange_strong(state, m_UnlockedState)))
760 {
761 return;
762 }
763
764 // If there are no existing waiters, transfer the list of new waiters.
765 if (m_WaitList == nullptr)
766 {
767 // Take ownership of the list of new waiters, leaving the state
768 // locked.
769 state = m_State.exchange(m_LockedState);
770 FAIL_FAST_IF(state == m_LockedState || state == m_UnlockedState);
771
772 // Reverse the list and transfer it. This ensures waiters are
773 // awoken in FIFO order.
774 auto* current = reinterpret_cast<AsyncLockTask*>(state);
775 AsyncLockTask* previous = nullptr;
776 AsyncLockTask* next = nullptr;
777 while (current != nullptr)
778 {
779 next = current->m_Next;
780 current->m_Next = previous;
781 previous = current;
782 current = next;
783 }
784
785 m_WaitList = previous;
786 }
787
788 // Awake the first waiter; this transfers lock ownership to them.
789 AsyncLockTask* head = m_WaitList;
790 m_WaitList = head->m_Next;
791 g_Scheduler.Schedule(head->m_Awaiter);
792 }
793
794 private:
795 static constexpr ULONG_PTR m_LockedState = 1;
796 static constexpr ULONG_PTR m_UnlockedState = 0;
797
798 // State is NULL if the lock is not held, 1 if it's held and there are no
799 // new waiters, or a pointer to the first new waiter in LIFO order.
800 std::atomic<ULONG_PTR> m_State;
801 // List of existing waiters in FIFO order.
802 AsyncLockTask* m_WaitList;
803 };
804
805 class ICancellable
806 {
807 public:
808 virtual ~ICancellable() = default;
809
810 virtual void Cancel() = 0;
811 };
812
813 // Token used to cancel an outstanding IO operation.
814 class CancelToken
815 {
816 public:
817 CancelToken()
818 {
819 InitializeListHead(&m_Children);
820 }
821
822 CancelToken(CancelToken& parent) : CancelToken()
823 {
824 if (parent.AddChild(*this))
825 {
826 m_Parent = &parent;
827 }
828 else
829 {
830 m_Cancelled = true;
831 }
832 }
833
834 ~CancelToken()
835 {
836 if (m_Parent)
837 {
838 m_Parent->RemoveChild(*this);
839 }
840 }
841
842 // Register a running IO as cancellable.
843 bool Register(ICancellable& operation)
844 {
845 std::lock_guard<std::mutex> lock{m_Lock};
846 if (m_Cancelled)
847 {
848 return false;
849 }
850
851 m_Operation = &operation;
852 return true;
853 }
854
855 // Unregister the currently registered overlapped structure.
856 void Unregister()
857 {
858 std::lock_guard<std::mutex> lock{m_Lock};
859 m_Operation = nullptr;
860 }
861
862 // Cancels the token, cancelling any associated outstanding IO.
863 void Cancel()
864 {
865 std::lock_guard<std::mutex> lock{m_Lock};
866 const bool wasCancelled = m_Cancelled;
867 m_Cancelled = true;
868 if (!wasCancelled)
869 {
870 if (m_Operation != nullptr)
871 {
872 m_Operation->Cancel();
873 }
874
875 for (auto entry = m_Children.Flink; entry != &m_Children; entry = entry->Flink)
876 {
877 const auto child = CONTAINING_RECORD(entry, CancelToken, m_Link);
878 child->Cancel();
879 }
880 }
881
882 m_Cancelled = true;
883 }
884
885 // Returns whether the token has already been cancelled.
886 bool Cancelled()
887 {
888 return m_Cancelled;
889 }
890
891 // Resets the state of the token. This should only be used when
892 // the token is no longer in use by any IOs.
893 void Reset()
894 {
895 std::lock_guard<std::mutex> lock{m_Lock};
896 FAIL_FAST_IF(m_Operation != nullptr || !IsListEmpty(&m_Children));
897 m_Cancelled = false;
898 }
899
900 private:
901 bool AddChild(CancelToken& child)
902 {
903 std::lock_guard<std::mutex> lock{m_Lock};
904 if (m_Cancelled)
905 {
906 return false;
907 }
908
909 InsertTailList(&m_Children, &child.m_Link);
910 return true;
911 }
912
913 void RemoveChild(CancelToken& child)
914 {
915 std::lock_guard<std::mutex> lock{m_Lock};
916 RemoveEntryList(&child.m_Link);
917 }
918
919 std::mutex m_Lock{};
920 ICancellable* m_Operation{};
921 std::atomic<bool> m_Cancelled{};
922 CancelToken* m_Parent{};
923 LIST_ENTRY m_Children{};
924 LIST_ENTRY m_Link{};
925 };
926
927 } // namespace p9fs