thread_pause_work_stealing()

alibaba · Dec 2, 2024 · 4b830be · 4b830be
1 parent 586c167
commit 4b830be
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 12 deletions.
diff --git a/thread/test/test.cpp b/thread/test/test.cpp
@@ -1751,8 +1751,14 @@ TEST(WorkStealing, basic) {
     });
     auto th = thread_create(&ws_basic, &stolen, 0, 0,
             THREAD_ENABLE_WORK_STEALING | THREAD_JOINABLE);
+    thread_pause_work_stealing(true, th);
+    ::usleep(1000 * 10);    // emulate a busy work of 10ms
+    EXPECT_FALSE(stolen);
+
+    thread_pause_work_stealing(false, th);
     ::usleep(1000 * 10);    // emulate a busy work of 10ms
     EXPECT_TRUE(stolen);
+
     thread_join((join_handle*)th);
     running = false;
     vcpu.join();

diff --git a/thread/thread.cpp b/thread/thread.cpp
@@ -177,6 +177,7 @@ namespace photon
         int idx = -1;                       /* index in the sleep queue array */
         int error_number = 0;
         thread_list* waitq = nullptr;       /* the q if WAITING in a queue */
+        uint32_t flags = 0;
         uint16_t state = states::READY;
         spinlock lock, _;                   // Current usage of thread.lock:
                                             //  interrupt()
@@ -186,7 +187,6 @@ namespace photon
                                             //  thread_join()
                                             //  ScopedLockHead
                                             //  try_work_stealing()?
-        int flags = 0;
         uint64_t ts_wakeup = 0;             /* Wakeup time when thread is sleeping */
 // offset 64B
         union {
@@ -208,7 +208,8 @@ namespace photon
         enum shift {
             joinable = 0,
             enable_work_stealing = 1,
-            shutting_down = 2,              // the thread should cancel what is doing, and quit
+            pause_work_stealing = 2,
+            shutting_down = 3,              // the thread should cancel what is doing, and quit
         };                                  // current job ASAP; not allowed to sleep or block more
                                             // than 10ms, otherwise -1 will be returned and errno == EPERM
         bool is_bit(int i) { return flags & (1<<i); }
@@ -219,8 +220,8 @@ namespace photon
         void set_joinable(bool flag = true) { set_bit(shift::joinable, flag); }
         bool is_shutting_down() { return is_bit(shift::shutting_down); }
         void set_shutting_down(bool flag = true) { set_bit(shift::shutting_down, flag); }
-        bool allow_work_stealing() { return is_bit(shift::enable_work_stealing); }
-
+        bool allow_work_stealing() { return is_bit(shift::enable_work_stealing) &&
+                                          !(flags & THREAD_PAUSE_WORK_STEALING); }
         int set_error_number() {
             if (likely(error_number)) {
                 errno = error_number;
@@ -296,6 +297,7 @@ namespace photon
 #pragma GCC diagnostic ignored "-Winvalid-offsetof"
     static_assert(offsetof(thread, vcpu) == offsetof(partial_thread, vcpu), "...");
     static_assert(offsetof(thread,  tls) == offsetof(partial_thread,  tls), "...");
+    static_assert(offsetof(thread,flags) == offsetof(partial_thread,flags), "...");
 #pragma GCC diagnostic pop
 
     struct thread_list : public intrusive_list<thread>
@@ -657,12 +659,6 @@ namespace photon
         }
     };
 
-#pragma GCC diagnostic push
-#pragma GCC diagnostic ignored "-Winvalid-offsetof"
-    static_assert(offsetof(thread, arg)   == 0x40, "...");
-    static_assert(offsetof(thread, start) == 0x48, "...");
-#pragma GCC diagnostic pop
-
     inline void thread::dequeue_ready_atomic(states newstat)
     {
         assert("this is not in runq, and this->lock is locked");
@@ -687,6 +683,13 @@ namespace photon
         to->get_vcpu()->switch_count++;
     }
 
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Winvalid-offsetof"
+    // the offsets are used in _photon_thread_stub() assembly code
+    static_assert(offsetof(thread, arg)   == 0x40, "...");
+    static_assert(offsetof(thread, start) == 0x48, "...");
+#pragma GCC diagnostic pop
+
 #if defined(__x86_64__)
 #if !defined(_WIN64)
     asm(

diff --git a/thread/thread.h b/thread/thread.h
@@ -33,6 +33,7 @@ namespace photon
     constexpr uint8_t  VCPU_ENABLE_PASSIVE_WORK_STEALING    = 2;    // allow this vCPU to be stolen by other vCPUs
     constexpr uint32_t THREAD_JOINABLE                      = 1;    // allow this thread to be joined
     constexpr uint32_t THREAD_ENABLE_WORK_STEALING          = 2;    // allow this thread to be stolen by other vCPUs
+    constexpr uint32_t THREAD_PAUSE_WORK_STEALING           = 4;    // temporarily pause work-stealing for a thread
 
     int vcpu_init(uint64_t flags = 0);
     int vcpu_fini();
@@ -149,12 +150,23 @@ namespace photon
     // A helper struct in order to make some function calls inline.
     // The memory layout of its first 4 fields is the same as the one of thread.
     struct partial_thread {
-        uint64_t _, __;
+        uint64_t _[2];
         volatile vcpu_base* vcpu;
-        uint64_t ___[5];
+        uint64_t __[3];
+        uint32_t flags, ___[3];
         void* tls;
     };
 
+    // this function doesn't affect whether WS is enabled or not for the thread
+    inline void thread_pause_work_stealing(bool flag, thread* th = CURRENT) {
+        auto& flags = ((partial_thread*)th)->flags;
+        if (flag) {
+            flags |= THREAD_PAUSE_WORK_STEALING;
+        } else {
+            flags &= ~THREAD_PAUSE_WORK_STEALING;
+        }
+    }
+
     inline vcpu_base* get_vcpu(thread* th = CURRENT) {
         auto vcpu = ((partial_thread*)th) -> vcpu;
         return (vcpu_base*)vcpu;