Merge pull request #49663 from CurtizJ/fix-race-context-copy

Fix race in `Context::createCopy`

src/Interpreters/Context.cpp

@@ -669,13 +669,15 @@ SharedContextHolder Context::createShared()
 
 ContextMutablePtr Context::createCopy(const ContextPtr & other)
 {
+    auto lock = other->getLock();
     return std::shared_ptr<Context>(new Context(*other));
 }
 
 ContextMutablePtr Context::createCopy(const ContextWeakPtr & other)
 {
     auto ptr = other.lock();
-    if (!ptr) throw Exception(ErrorCodes::LOGICAL_ERROR, "Can't copy an expired context");
+    if (!ptr)
+        throw Exception(ErrorCodes::LOGICAL_ERROR, "Can't copy an expired context");
     return createCopy(ptr);
 }
 

src/Interpreters/tests/gtest_context_race.cpp

@@ -0,0 +1,49 @@
+#include <Interpreters/Context.h>
+#include <Common/tests/gtest_global_context.h>
+#include <gtest/gtest.h>
+
+using namespace DB;
+
+template <typename Ptr>
+void run(Ptr context)
+{
+    for (size_t i = 0; i < 100; ++i)
+    {
+        std::thread t1([context]
+        {
+            if constexpr (std::is_same_v<ContextWeakPtr, Ptr>)
+                context.lock()->getAsyncReadCounters();
+            else
+                context->getAsyncReadCounters();
+        });
+
+        std::thread t2([context]
+        {
+            Context::createCopy(context);
+        });
+
+        t1.join();
+        t2.join();
+    }
+}
+
+TEST(Context, MutableRace)
+{
+    auto context = Context::createCopy(getContext().context);
+    context->makeQueryContext();
+    run<ContextMutablePtr>(context);
+}
+
+TEST(Context, ConstRace)
+{
+    auto context = Context::createCopy(getContext().context);
+    context->makeQueryContext();
+    run<ContextPtr>(context);
+}
+
+TEST(Context, WeakRace)
+{
+    auto context = Context::createCopy(getContext().context);
+    context->makeQueryContext();
+    run<ContextWeakPtr>(context);
+}