bthread copies tls_bls to its local_storage before switching to next bthread

src/bthread/key.cpp

@@ -472,10 +472,6 @@ void* bthread_getspecific(bthread_key_t key) {
 
 void bthread_assign_data(void* data) {
     bthread::tls_bls.assigned_data = data;
-    bthread::TaskGroup* const g = bthread::tls_task_group;
-    if (g) {
-        g->current_task()->local_storage.assigned_data = data;
-    }
 }
 
 void* bthread_get_assigned_data() {

src/bthread/task_group.cpp

@@ -55,6 +55,9 @@ const bool ALLOW_UNUSED dummy_show_per_worker_usage_in_vars =
                                     pass_bool);
 
 __thread TaskGroup* tls_task_group = NULL;
+// Sync with TaskMeta::local_storage when a bthread is created or destroyed.
+// During running, the two fields may be inconsistent, use tls_bls as the
+// groundtruth.
 __thread LocalStorage tls_bls = BTHREAD_LOCAL_STORAGE_INITIALIZER;
 
 // defined in bthread/key.cpp
@@ -229,14 +232,7 @@ int TaskGroup::init(size_t runqueue_capacity) {
     m->about_to_quit = false;
     m->fn = NULL;
     m->arg = NULL;
-    // In current implementation, even if we set m->local_storage to empty,
-    // everything should be fine because a non-worker pthread never context
-    // switches to a bthread, inconsistency between m->local_storage and tls_bls
-    // does not result in bug. However to avoid potential future bug,
-    // TaskMeta.local_storage is better to be sync with tls_bls otherwise
-    // context switching back to this main bthread will restore tls_bls
-    // with NULL values which is incorrect.
-    m->local_storage = tls_bls;
+    m->local_storage = LOCAL_STORAGE_INIT;
     m->cpuwide_start_ns = butil::cpuwide_time_ns();
     m->stat = EMPTY_STAT;
     m->attr = BTHREAD_ATTR_TASKGROUP;
@@ -316,12 +312,12 @@ void TaskGroup::task_runner(intptr_t skip_remained) {
         // Clean tls variables, must be done before changing version_butex
         // otherwise another thread just joined this thread may not see side
         // effects of destructing tls variables.
-        KeyTable* kt = m->local_storage.keytable;
+        KeyTable* kt = tls_bls.keytable;
         if (kt != NULL) {
             return_keytable(m->attr.keytable_pool, kt);
             // After deletion: tls may be set during deletion.
-            m->local_storage.keytable = NULL;
             tls_bls.keytable = NULL;
+            m->local_storage.keytable = NULL; // optional
         }
         
         // Increase the version and wake up all joiners, if resulting version
@@ -588,6 +584,8 @@ void TaskGroup::sched_to(TaskGroup** pg, TaskMeta* next_meta) {
     // Switch to the task
     if (__builtin_expect(next_meta != cur_meta, 1)) {
         g->_cur_meta = next_meta;
+        // Switch tls_bls
+        cur_meta->local_storage = tls_bls;
         tls_bls = next_meta->local_storage;
 
         // Logging must be done after switching the local storage, since the logging lib 

src/bthread/task_meta.h

@@ -83,7 +83,9 @@ struct TaskMeta {
     int64_t cpuwide_start_ns;
     TaskStatistics stat;
 
-    // bthread local storage.
+    // bthread local storage, sync with tls_bls (defined in task_group.cpp)
+    // when the bthread is created or destroyed.
+    // DO NOT use this field directly, use tls_bls instead.
     LocalStorage local_storage;
 
 public: