From 10c8dbcb840c744002a824950e8293ef0a1fd61a Mon Sep 17 00:00:00 2001
From: Kadir Cetinkaya <kadircet@google.com>
Date: Wed, 23 Oct 2019 10:18:09 +0200
Subject: [PATCH] [clangd] Propogate context in TUScheduler::run

Reviewers: sammccall

Subscribers: ilya-biryukov, javed.absar, MaskRay, jkorous, arphaman, usaxena95, cfe-commits

Tags: #clang

Differential Revision: https://reviews.llvm.org/D69328
---
 clang-tools-extra/clangd/TUScheduler.cpp                | 7 ++++++-
 clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp | 9 +++++++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/clang-tools-extra/clangd/TUScheduler.cpp b/clang-tools-extra/clangd/TUScheduler.cpp
index c9cd7cf..6436e7a 100644
--- a/clang-tools-extra/clangd/TUScheduler.cpp
+++ b/clang-tools-extra/clangd/TUScheduler.cpp
@@ -44,6 +44,7 @@
 #include "TUScheduler.h"
 #include "Cancellation.h"
 #include "Compiler.h"
+#include "Context.h"
 #include "Diagnostics.h"
 #include "GlobalCompilationDatabase.h"
 #include "Logger.h"
@@ -919,7 +920,11 @@ void TUScheduler::run(llvm::StringRef Name,
                       llvm::unique_function<void()> Action) {
   if (!PreambleTasks)
     return Action();
-  PreambleTasks->runAsync(Name, std::move(Action));
+  PreambleTasks->runAsync(Name, [Ctx = Context::current().clone(),
+                                 Action = std::move(Action)]() mutable {
+    WithContext WC(std::move(Ctx));
+    Action();
+  });
 }
 
 void TUScheduler::runWithAST(
diff --git a/clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp b/clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp
index db05c70..e34ce0b 100644
--- a/clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp
+++ b/clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp
@@ -688,6 +688,15 @@ TEST_F(TUSchedulerTests, Run) {
   S.run("add 2", [&] { Counter += 2; });
   ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
   EXPECT_EQ(Counter.load(), 3);
+
+  Notification TaskRun;
+  Key<int> TestKey;
+  WithContextValue CtxWithKey(TestKey, 10);
+  S.run("props context", [&] {
+    EXPECT_EQ(Context::current().getExisting(TestKey), 10);
+    TaskRun.notify();
+  });
+  TaskRun.wait();
 }
 
 TEST_F(TUSchedulerTests, TUStatus) {
-- 
2.7.4