Lift RequestContext root-id into task stats

Summary: We want to use observer API to correlate tasks across same Thrift request. `RequestContext::getRootId` is the thing for it.

Reviewed By: yfeldblum

Differential Revision: D20600485

fbshipit-source-id: 1a79bde0f6c5b3692b4fb3805c98f377b1b317c9

folly/executors/EDFThreadPoolExecutor.cpp

@@ -78,7 +78,6 @@ class EDFThreadPoolExecutor::Task {
   std::atomic<int> iter_{0};
   int total_;
   uint64_t deadline_;
-  TaskStats stats_;
   std::shared_ptr<RequestContext> context_ = RequestContext::saveContext();
   std::chrono::steady_clock::time_point enqueueTime_ =
       std::chrono::steady_clock::now();
@@ -369,7 +368,13 @@ void EDFThreadPoolExecutor::threadRun(ThreadPtr thread) {
 
     thread->idle = false;
     auto startTime = std::chrono::steady_clock::now();
-    task->stats_.waitTime = startTime - task->enqueueTime_;
+    TaskStats stats;
+    stats.enqueueTime = task->enqueueTime_;
+    if (task->context_) {
+      stats.requestId = task->context_->getRootId();
+    }
+
+    stats.waitTime = startTime - stats.enqueueTime;
     try {
       task->run(iter);
     } catch (const std::exception& e) {
@@ -379,7 +384,7 @@ void EDFThreadPoolExecutor::threadRun(ThreadPtr thread) {
       LOG(ERROR)
           << "EDFThreadPoolExecutor: func threw unhandled non-exception object";
     }
-    task->stats_.runTime = std::chrono::steady_clock::now() - startTime;
+    stats.runTime = std::chrono::steady_clock::now() - startTime;
     thread->idle = true;
     thread->lastActiveTime = std::chrono::steady_clock::now();
     thread->taskStatsCallbacks->callbackList.withRLock([&](auto& callbacks) {
@@ -389,7 +394,7 @@ void EDFThreadPoolExecutor::threadRun(ThreadPtr thread) {
       };
       try {
         for (auto& callback : callbacks) {
-          callback(task->stats_);
+          callback(stats);
         }
       } catch (const std::exception& e) {
         LOG(ERROR) << "EDFThreadPoolExecutor: task stats callback threw "

folly/executors/ThreadPoolExecutor.cpp

@@ -76,10 +76,16 @@ ThreadPoolExecutor::Task::Task(
 void ThreadPoolExecutor::runTask(const ThreadPtr& thread, Task&& task) {
   thread->idle = false;
   auto startTime = std::chrono::steady_clock::now();
-  task.stats_.waitTime = startTime - task.enqueueTime_;
+  TaskStats stats;
+  stats.enqueueTime = task.enqueueTime_;
+  if (task.context_) {
+    stats.requestId = task.context_->getRootId();
+  }
+  stats.waitTime = startTime - task.enqueueTime_;
+
   if (task.expiration_ > std::chrono::milliseconds(0) &&
-      task.stats_.waitTime >= task.expiration_) {
-    task.stats_.expired = true;
+      stats.waitTime >= task.expiration_) {
+    stats.expired = true;
     if (task.expireCallback_ != nullptr) {
       task.expireCallback_();
     }
@@ -94,7 +100,7 @@ void ThreadPoolExecutor::runTask(const ThreadPtr& thread, Task&& task) {
       LOG(ERROR) << "ThreadPoolExecutor: func threw unhandled non-exception "
                     "object";
     }
-    task.stats_.runTime = std::chrono::steady_clock::now() - startTime;
+    stats.runTime = std::chrono::steady_clock::now() - startTime;
   }
   thread->idle = true;
   thread->lastActiveTime = std::chrono::steady_clock::now();
@@ -105,7 +111,7 @@ void ThreadPoolExecutor::runTask(const ThreadPtr& thread, Task&& task) {
     };
     try {
       for (auto& callback : callbacks) {
-        callback(task.stats_);
+        callback(stats);
       }
     } catch (const std::exception& e) {
       LOG(ERROR) << "ThreadPoolExecutor: task stats callback threw "

folly/executors/ThreadPoolExecutor.h

@@ -125,10 +125,12 @@ class ThreadPoolExecutor : public DefaultKeepAliveExecutor {
   std::string getName() const;
 
   struct TaskStats {
-    TaskStats() : expired(false), waitTime(0), runTime(0) {}
+    TaskStats() : expired(false), waitTime(0), runTime(0), requestId(0) {}
     bool expired;
     std::chrono::nanoseconds waitTime;
     std::chrono::nanoseconds runTime;
+    std::chrono::steady_clock::time_point enqueueTime;
+    uint64_t requestId;
   };
 
   using TaskStatsCallback = std::function<void(TaskStats)>;
@@ -204,7 +206,6 @@ class ThreadPoolExecutor : public DefaultKeepAliveExecutor {
         std::chrono::milliseconds expiration,
         Func&& expireCallback);
     Func func_;
-    TaskStats stats_;
     std::chrono::steady_clock::time_point enqueueTime_;
     std::chrono::milliseconds expiration_;
     Func expireCallback_;

folly/executors/test/ThreadPoolExecutorTest.cpp

@@ -266,14 +266,20 @@ template <class TPE>
 static void taskStats() {
   TPE tpe(1);
   std::atomic<int> c(0);
+  auto now = std::chrono::steady_clock::now();
   tpe.subscribeToTaskStats([&](ThreadPoolExecutor::TaskStats stats) {
     int i = c++;
+    EXPECT_LT(now, stats.enqueueTime);
     EXPECT_LT(milliseconds(0), stats.runTime);
     if (i == 1) {
       EXPECT_LT(milliseconds(0), stats.waitTime);
+      EXPECT_NE(0, stats.requestId);
+    } else {
+      EXPECT_EQ(0, stats.requestId);
     }
   });
   tpe.add(burnMs(10));
+  RequestContextScopeGuard rctx;
   tpe.add(burnMs(10));
   tpe.join();
   EXPECT_EQ(2, c);