Context helper functions (open-telemetry#225)

ziqizh · Jul 31, 2020 · ec0b87e · ec0b87e
1 parent ce7d76b
commit ec0b87e
Show file tree

Hide file tree

Showing 2 changed files with 81 additions and 0 deletions.
diff --git a/api/include/opentelemetry/context/runtime_context.h b/api/include/opentelemetry/context/runtime_context.h
@@ -47,6 +47,46 @@ class RuntimeContext
 
   static RuntimeContext *context_handler_;
 
+  // Sets the Key and Value into the passed in context or if a context is not
+  // passed in, the RuntimeContext.
+  // Should be used to SetValues to the current RuntimeContext, is essentially
+  // equivalent to RuntimeContext::GetCurrent().SetValue(key,value). Keep in
+  // mind that the current RuntimeContext will not be changed, and the new
+  // context will be returned.
+  static Context SetValue(nostd::string_view key,
+                          ContextValue value,
+                          Context *context = nullptr) noexcept
+  {
+    Context temp_context;
+    if (context == nullptr)
+    {
+      temp_context = GetCurrent();
+    }
+    else
+    {
+      temp_context = *context;
+    }
+    return temp_context.SetValue(key, value);
+  }
+
+  // Returns the value associated with the passed in key and either the
+  // passed in context* or the runtime context if a context is not passed in.
+  // Should be used to get values from the current RuntimeContext, is
+  // essentially equivalent to RuntimeContext::GetCurrent().GetValue(key).
+  static ContextValue GetValue(nostd::string_view key, Context *context = nullptr) noexcept
+  {
+    Context temp_context;
+    if (context == nullptr)
+    {
+      temp_context = GetCurrent();
+    }
+    else
+    {
+      temp_context = *context;
+    }
+    return temp_context.GetValue(key);
+  }
+
 protected:
   // Provides a token with the passed in context
   Token CreateToken(Context context) noexcept { return Token(context); }

diff --git a/api/test/context/runtime_context_test.cc b/api/test/context/runtime_context_test.cc
@@ -59,3 +59,44 @@ TEST(RuntimeContextTest, ThreeAttachDetach)
   EXPECT_TRUE(context::RuntimeContext::Detach(foo_context_token));
   EXPECT_TRUE(context::RuntimeContext::Detach(test_context_token));
 }
+
+// Tests that SetValue returns a context with the passed in data and the
+// RuntimeContext data when a context is not passed into the
+// RuntimeContext::SetValue method.
+TEST(RuntimeContextTest, SetValueRuntimeContext)
+{
+  context::Context foo_context                     = context::Context("foo_key", (int64_t)596);
+  context::RuntimeContext::Token old_context_token = context::RuntimeContext::Attach(foo_context);
+  context::Context test_context = context::RuntimeContext::SetValue("test_key", (int64_t)123);
+  EXPECT_EQ(nostd::get<int64_t>(test_context.GetValue("test_key")), 123);
+  EXPECT_EQ(nostd::get<int64_t>(test_context.GetValue("foo_key")), 596);
+}
+
+// Tests that SetValue returns a context with the passed in data and the
+// passed in context data when a context* is passed into the
+// RuntimeContext::SetValue method.
+TEST(RuntimeContextTest, SetValueOtherContext)
+{
+  context::Context foo_context = context::Context("foo_key", (int64_t)596);
+  context::Context test_context =
+      context::RuntimeContext::SetValue("test_key", (int64_t)123, &foo_context);
+  EXPECT_EQ(nostd::get<int64_t>(test_context.GetValue("test_key")), 123);
+  EXPECT_EQ(nostd::get<int64_t>(test_context.GetValue("foo_key")), 596);
+}
+
+// Tests that SetValue returns the ContextValue associated with the
+// passed in string and the current Runtime Context
+TEST(RuntimeContextTest, GetValueRuntimeContext)
+{
+  context::Context foo_context                     = context::Context("foo_key", (int64_t)596);
+  context::RuntimeContext::Token old_context_token = context::RuntimeContext::Attach(foo_context);
+  EXPECT_EQ(nostd::get<int64_t>(context::RuntimeContext::GetValue("foo_key")), 596);
+}
+
+// Tests that SetValue returns the ContextValue associated with the
+// passed in string and the passed in context
+TEST(RuntimeContextTest, GetValueOtherContext)
+{
+  context::Context foo_context = context::Context("foo_key", (int64_t)596);
+  EXPECT_EQ(nostd::get<int64_t>(context::RuntimeContext::GetValue("foo_key", &foo_context)), 596);
+}