@samitouri / QOSAMI-WSL-2
Code Commits Issues Pulls Sessions CI/CD
master
master
cpp 175 lines 3.97 KB
Raw
1 // Copyright (C) Microsoft Corporation. All rights reserved.
2 #include "precomp.h"
3 #include "p9platform.h"
4 #include "p9scheduler.h"
5
6 namespace p9fs {
7
8 Scheduler g_Scheduler;
9 thread_local bool Scheduler::tls_Blocked{};
10 thread_local bool Scheduler::tls_SchedulerThread{};
11
12 Scheduler::Scheduler() : m_Work{CreateWorkItem(std::bind(&Scheduler::WorkerCallback, this))}
13 {
14 }
15
16 /// Schedules a coroutine to run. It will run sometime after this coroutine
17 /// yields or enters a blocking region.
18 void Scheduler::Schedule(Coroutine coroutine) noexcept
19 {
20 bool kick = false;
21
22 {
23 std::unique_lock<std::shared_mutex> lock(m_Lock);
24
25 // N.B. This could throw in very low memory situations, which would terminate the process.
26 m_Queue.push(coroutine);
27 if (!m_Running && !m_ThreadEnqueued)
28 {
29 m_ThreadEnqueued = true;
30 kick = true;
31 }
32 }
33
34 if (kick)
35 {
36 m_Work->Submit();
37 }
38 }
39
40 /// Donates the current thread to run coroutines and schedules the specified
41 /// coroutine to run.
42 void Scheduler::DonateThreadAndResume(Coroutine coroutine) noexcept
43 {
44 const bool run = Claim(false);
45 Schedule(coroutine);
46 if (run)
47 {
48 RunAndRelease();
49 }
50 }
51
52 /// Runs coroutines until there are no more in the queue or until this thread
53 /// gave up the queue in order to run blocking code.
54 ///
55 /// Must be called on the thread that called Claim().
56 void Scheduler::RunAndRelease() noexcept
57 {
58 WI_ASSERT(!tls_Blocked);
59
60 tls_SchedulerThread = true;
61 std::unique_lock<std::shared_mutex> lock(m_Lock);
62 while (!m_Queue.empty())
63 {
64 auto coroutine = m_Queue.front();
65 m_Queue.pop();
66 lock.unlock();
67 coroutine.resume();
68 if (tls_Blocked)
69 {
70 tls_Blocked = false;
71 tls_SchedulerThread = false;
72 return;
73 }
74
75 lock.lock();
76 }
77
78 WI_ASSERT(m_Queue.empty() || m_Running || m_ThreadEnqueued);
79
80 m_Running = false;
81 tls_SchedulerThread = false;
82 }
83
84 /// Called when the current thread may block for some time. Gives up queue
85 /// ownership, potentially scheduling another thread to resume running
86 /// non-blocking code.
87 bool Scheduler::Block() noexcept
88 {
89 if (!tls_SchedulerThread)
90 {
91 return false;
92 }
93
94 WI_ASSERT(!tls_Blocked);
95
96 tls_Blocked = true;
97
98 bool kick = false;
99
100 {
101 std::unique_lock<std::shared_mutex> lock(m_Lock);
102
103 WI_ASSERT(m_Running);
104
105 m_Running = false;
106 if (!m_Queue.empty() && !m_ThreadEnqueued)
107 {
108 m_ThreadEnqueued = true;
109 kick = true;
110 }
111 }
112
113 if (kick)
114 {
115 m_Work->Submit();
116 }
117
118 return true;
119 }
120
121 /// Awaitable function called when the current thread is done running blocking
122 /// code. Tries to reclaim ownership of the queue and resumes the current
123 /// coroutine.
124 Scheduler::Unblocker Scheduler::Unblock() noexcept
125 {
126 WI_ASSERT(tls_Blocked);
127
128 // Try to reuse this thread to run async tasks.
129 const bool run = Claim(false);
130 if (run)
131 {
132 tls_Blocked = false;
133 }
134
135 // Unblocker will either resume the current coroutine or schedule it to run
136 // on the new queue owner.
137 return Unblocker{*this, run};
138 }
139
140 /// Try to claim queue ownership for the current thread. If this function
141 /// returns true, then the caller must call RunAndRelease to process the queue.
142 ///
143 /// If fromKick, then the caller is the thread that was explicitly kicked to
144 /// process the queue. Otherwise, this is an IO completion or other
145 /// opportunistic thread.
146 bool Scheduler::Claim(bool fromKick) noexcept
147 {
148 std::unique_lock<std::shared_mutex> lock(m_Lock);
149
150 WI_ASSERT(!fromKick || m_ThreadEnqueued);
151
152 if (fromKick)
153 {
154 m_ThreadEnqueued = false;
155 }
156
157 if (m_Running)
158 {
159 return false;
160 }
161
162 m_Running = true;
163 return true;
164 }
165
166 /// Threadpool callback called to process the queue.
167 void Scheduler::WorkerCallback() noexcept
168 {
169 if (Claim(true))
170 {
171 RunAndRelease();
172 }
173 }
174
175 } // namespace p9fs