diff --git a/api/include/opentelemetry/context/runtime_context.h b/api/include/opentelemetry/context/runtime_context.h index e60a8726ea..094e1a5bd9 100644 --- a/api/include/opentelemetry/context/runtime_context.h +++ b/api/include/opentelemetry/context/runtime_context.h @@ -175,15 +175,30 @@ class ThreadLocalContextStorage : public RuntimeContextStorage // Return the current context. Context GetCurrent() noexcept override { return GetStack().Top(); } - // Resets the context to a previous value stored in the - // passed in token. Returns true if successful, false otherwise + // Resets the context to the value previous to the passed in token. This will + // also detach all child contexts of the passed in token. + // Returns true if successful, false otherwise. bool Detach(Token &token) noexcept override { - if (!(token == GetStack().Top())) + // In most cases, the context to be detached is on the top of the stack. + if (token == GetStack().Top()) + { + GetStack().Pop(); + return true; + } + + if (!GetStack().Contains(token)) { return false; } + + while (!(token == GetStack().Top())) + { + GetStack().Pop(); + } + GetStack().Pop(); + return true; } @@ -214,6 +229,19 @@ class ThreadLocalContextStorage : public RuntimeContextStorage return base_[size_]; } + bool Contains(const Token &token) const noexcept + { + for (size_t pos = size_; pos > 0; --pos) + { + if (token == base_[pos - 1]) + { + return true; + } + } + + return false; + } + // Returns the Context at the top of the stack. Context Top() const noexcept { diff --git a/api/test/context/runtime_context_test.cc b/api/test/context/runtime_context_test.cc index 31c7d64a8e..3fe78ade44 100644 --- a/api/test/context/runtime_context_test.cc +++ b/api/test/context/runtime_context_test.cc @@ -11,7 +11,7 @@ TEST(RuntimeContextTest, GetCurrent) std::map map_test = {{"test_key", (int64_t)123}}; context::Context test_context = context::Context(map_test); auto old_context = context::RuntimeContext::Attach(test_context); - EXPECT_TRUE(context::RuntimeContext::GetCurrent() == test_context); + EXPECT_EQ(context::RuntimeContext::GetCurrent(), test_context); } // Tests that detach resets the context to the previous context @@ -25,7 +25,7 @@ TEST(RuntimeContextTest, Detach) auto foo_context_token = context::RuntimeContext::Attach(foo_context); foo_context_token.reset(); - EXPECT_TRUE(context::RuntimeContext::GetCurrent() == test_context); + EXPECT_EQ(context::RuntimeContext::GetCurrent(), test_context); test_context_token.reset(); } @@ -34,9 +34,8 @@ TEST(RuntimeContextTest, DetachWrongContext) { std::map map_test = {{"test_key", (int64_t)123}}; context::Context test_context = context::Context(map_test); - context::Context foo_context = context::Context(map_test); auto test_context_token = context::RuntimeContext::Attach(test_context); - auto foo_context_token = context::RuntimeContext::Attach(foo_context); + EXPECT_TRUE(context::RuntimeContext::Detach(*test_context_token)); EXPECT_FALSE(context::RuntimeContext::Detach(*test_context_token)); } @@ -96,3 +95,38 @@ TEST(RuntimeContextTest, GetValueOtherContext) context::Context foo_context = context::Context("foo_key", (int64_t)596); EXPECT_EQ(nostd::get(context::RuntimeContext::GetValue("foo_key", &foo_context)), 596); } + +// Test that any possible order of context detaching doesn't mess up the stack. +TEST(RuntimeContextTest, DetachOutOfOrder) +{ + std::vector indices; + indices.push_back(0); + indices.push_back(1); + indices.push_back(2); + indices.push_back(3); + + std::vector contexts; + for (auto i : indices) + { + contexts.push_back(context::Context("index", i)); + } + + do + { + std::vector> tokens; + + for (auto &c : contexts) + { + tokens.push_back(context::RuntimeContext::Attach(c)); + } + + for (size_t i : indices) + { + auto token = std::move(tokens.at(i)); + context::RuntimeContext::Detach(*token); + } + + EXPECT_EQ(context::RuntimeContext::GetCurrent(), context::Context()); + + } while (std::next_permutation(indices.begin(), indices.end())); +}