From 4b830be1bbbd0918b13ba00417f4698777949570 Mon Sep 17 00:00:00 2001
From: lihuiba <huiba.lhb@alibaba-inc.com>
Date: Mon, 2 Dec 2024 12:20:29 +0800
Subject: [PATCH] thread_pause_work_stealing()

---
 thread/test/test.cpp |  6 ++++++
 thread/thread.cpp    | 23 +++++++++++++----------
 thread/thread.h      | 16 ++++++++++++++--
 3 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/thread/test/test.cpp b/thread/test/test.cpp
index 4534be9d..543a1085 100644
--- 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
index 407c7ba9..a79a2a46 100644
--- 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
index 7183a422..d0f56235 100644
--- 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;