general tracing context interface

envoy/http/BUILD

@@ -115,6 +115,7 @@ envoy_cc_library(
     ],
     deps = [
         ":header_formatter_interface",
+        "//envoy/tracing:trace_context_interface",
         "//source/common/common:assert_lib",
         "//source/common/common:hash_lib",
     ],

envoy/http/header_map.h

@@ -11,6 +11,7 @@
 #include "envoy/common/optref.h"
 #include "envoy/common/pure.h"
 #include "envoy/http/header_formatter.h"
+#include "envoy/tracing/trace_context.h"
 
 #include "source/common/common/assert.h"
 #include "source/common/common/hash.h"
@@ -78,6 +79,9 @@ class LowerCaseString {
     return os << lower_case_string.string_;
   }
 
+  // Implicit conversion to absl::string_view.
+  operator absl::string_view() const { return string_; }
+
 private:
   void lower() {
     std::transform(string_.begin(), string_.end(), string_.begin(), absl::ascii_tolower);
@@ -817,7 +821,8 @@ class RequestOrResponseHeaderMap : public HeaderMap {
 // Request headers.
 class RequestHeaderMap
     : public RequestOrResponseHeaderMap,
-      public CustomInlineHeaderBase<CustomInlineHeaderRegistry::Type::RequestHeaders> {
+      public CustomInlineHeaderBase<CustomInlineHeaderRegistry::Type::RequestHeaders>,
+      public Tracing::TraceContext {
 public:
   INLINE_REQ_STRING_HEADERS(DEFINE_INLINE_STRING_HEADER)
   INLINE_REQ_NUMERIC_HEADERS(DEFINE_INLINE_NUMERIC_HEADER)

envoy/tracing/BUILD

@@ -40,3 +40,8 @@ envoy_cc_library(
         "//envoy/stream_info:stream_info_interface",
     ],
 )
+
+envoy_cc_library(
+    name = "trace_context_interface",
+    hdrs = ["trace_context.h"],
+)

envoy/tracing/trace_context.h

@@ -0,0 +1,68 @@
+#pragma once
+
+#include <string>
+
+#include "envoy/common/pure.h"
+
+#include "absl/strings/string_view.h"
+#include "absl/types/optional.h"
+
+namespace Envoy {
+namespace Tracing {
+
+/**
+ * Protocol-independent abstraction for traceable stream. It hides the differences between different
+ * protocol and provides tracer driver with common methods for obtaining and setting the tracing
+ * context.
+ */
+class TraceContext {
+public:
+  virtual ~TraceContext() = default;
+
+  /**
+   * Get tracing context value by key.
+   *
+   * @param key The context key of string view type.
+   * @return The optional context value of string_view type.
+   */
+  virtual absl::optional<absl::string_view> getTraceContext(absl::string_view key) const PURE;
+
+  /**
+   * Set new tracing context key/value pair.
+   *
+   * @param key The context key of string view type.
+   * @param val The context value of string view type.
+   */
+  virtual void setTraceContext(absl::string_view key, absl::string_view val) PURE;
+
+  /**
+   * Set new tracing context key/value pair. The key MUST point to data that will live beyond
+   * the lifetime of any traceable stream that using the string.
+   *
+   * @param key The context key of string view type.
+   * @param val The context value of string view type.
+   */
+  virtual void setTraceContextReferenceKey(absl::string_view key, absl::string_view val) {
+    // The reference semantics of key and value are ignored by default. Derived classes that wish to
+    // use reference semantics to improve performance or reduce memory overhead can override this
+    // method.
+    setTraceContext(key, val);
+  }
+
+  /**
+   * Set new tracing context key/value pair. Both key and val MUST point to data that will live
+   * beyond the lifetime of any traceable stream that using the string.
+   *
+   * @param key The context key of string view type.
+   * @param val The context value of string view type.
+   */
+  virtual void setTraceContextReference(absl::string_view key, absl::string_view val) {
+    // The reference semantics of key and value are ignored by default. Derived classes that wish to
+    // use reference semantics to improve performance or reduce memory overhead can override this
+    // method.
+    setTraceContext(key, val);
+  }
+};
+
+} // namespace Tracing
+} // namespace Envoy

source/common/http/header_map_impl.cc

@@ -40,6 +40,12 @@ InlineHeaderVector& getInVec(VariantHeader& buffer) {
 const InlineHeaderVector& getInVec(const VariantHeader& buffer) {
   return absl::get<InlineHeaderVector>(buffer);
 }
+
+bool validatedLowerCaseString(absl::string_view str) {
+  auto lower_case_str = LowerCaseString(std::string(str));
+  return lower_case_str == str;
+}
+
 } // namespace
 
 // Initialize as a Type::Inline
@@ -469,13 +475,13 @@ HeaderMap::GetResult HeaderMapImpl::get(const LowerCaseString& key) const {
   return HeaderMap::GetResult(const_cast<HeaderMapImpl*>(this)->getExisting(key));
 }
 
-HeaderMap::NonConstGetResult HeaderMapImpl::getExisting(const LowerCaseString& key) {
+HeaderMap::NonConstGetResult HeaderMapImpl::getExisting(absl::string_view key) {
   // Attempt a trie lookup first to see if the user is requesting an O(1) header. This may be
   // relatively common in certain header matching / routing patterns.
   // TODO(mattklein123): Add inline handle support directly to the header matcher code to support
   // this use case more directly.
   HeaderMap::NonConstGetResult ret;
-  auto lookup = staticLookup(key.get());
+  auto lookup = staticLookup(key);
   if (lookup.has_value()) {
     if (*lookup.value().entry_ != nullptr) {
       ret.push_back(*lookup.value().entry_);
@@ -486,7 +492,7 @@ HeaderMap::NonConstGetResult HeaderMapImpl::getExisting(const LowerCaseString& k
   // If the requested header is not an O(1) header try using the lazy map to
   // search for it instead of iterating the headers list.
   if (headers_.maybeMakeMap()) {
-    HeaderList::HeaderLazyMap::iterator iter = headers_.mapFind(key.get());
+    HeaderList::HeaderLazyMap::iterator iter = headers_.mapFind(key);
     if (iter != headers_.mapEnd()) {
       const HeaderList::HeaderNodeVector& v = iter->second;
       ASSERT(!v.empty()); // It's impossible to have a map entry with an empty vector as its value.
@@ -502,7 +508,7 @@ HeaderMap::NonConstGetResult HeaderMapImpl::getExisting(const LowerCaseString& k
   // scan. Doing the trie lookup is wasteful in the miss case, but is present for code consistency
   // with other functions that do similar things.
   for (HeaderEntryImpl& header : headers_) {
-    if (header.key() == key.get().c_str()) {
+    if (header.key() == key) {
       ret.push_back(&header);
     }
   }
@@ -556,16 +562,7 @@ size_t HeaderMapImpl::removeIf(const HeaderMap::HeaderMatchPredicate& predicate)
   return old_size - headers_.size();
 }
 
-size_t HeaderMapImpl::remove(const LowerCaseString& key) {
-  const size_t old_size = headers_.size();
-  auto lookup = staticLookup(key.get());
-  if (lookup.has_value()) {
-    removeInline(lookup.value().entry_);
-  } else {
-    subtractSize(headers_.remove(key.get()));
-  }
-  return old_size - headers_.size();
-}
+size_t HeaderMapImpl::remove(const LowerCaseString& key) { return removeExisting(key); }
 
 size_t HeaderMapImpl::removePrefix(const LowerCaseString& prefix) {
   return HeaderMapImpl::removeIf([&prefix](const HeaderEntry& entry) -> bool {
@@ -610,6 +607,17 @@ HeaderMapImpl::HeaderEntryImpl& HeaderMapImpl::maybeCreateInline(HeaderEntryImpl
   return **entry;
 }
 
+size_t HeaderMapImpl::removeExisting(absl::string_view key) {
+  const size_t old_size = headers_.size();
+  auto lookup = staticLookup(key);
+  if (lookup.has_value()) {
+    removeInline(lookup.value().entry_);
+  } else {
+    subtractSize(headers_.remove(key));
+  }
+  return old_size - headers_.size();
+}
+
 size_t HeaderMapImpl::removeInline(HeaderEntryImpl** ptr_to_entry) {
   if (!*ptr_to_entry) {
     return 0;
@@ -650,5 +658,49 @@ HeaderMapImplUtility::getAllHeaderMapImplInfo() {
   return ret;
 }
 
+absl::optional<absl::string_view>
+RequestHeaderMapImpl::getTraceContext(absl::string_view key) const {
+  ASSERT(validatedLowerCaseString(key));
+  auto result = const_cast<RequestHeaderMapImpl*>(this)->getExisting(key);
+
+  if (result.empty()) {
+    return absl::nullopt;
+  }
+  return result[0]->value().getStringView();
+}
+
+void RequestHeaderMapImpl::setTraceContext(absl::string_view key, absl::string_view val) {
+  ASSERT(validatedLowerCaseString(key));
+  HeaderMapImpl::removeExisting(key);
+
+  HeaderString new_key;
+  new_key.setCopy(key);
+  HeaderString new_val;
+  new_val.setCopy(val);
+
+  HeaderMapImpl::insertByKey(std::move(new_key), std::move(new_val));
+  ASSERT(new_key.empty()); // NOLINT(bugprone-use-after-move)
+  ASSERT(new_val.empty()); // NOLINT(bugprone-use-after-move)
+}
+
+void RequestHeaderMapImpl::setTraceContextReferenceKey(absl::string_view key,
+                                                       absl::string_view val) {
+  ASSERT(validatedLowerCaseString(key));
+  HeaderMapImpl::removeExisting(key);
+
+  HeaderString new_val;
+  new_val.setCopy(val);
+
+  HeaderMapImpl::insertByKey(HeaderString(key), std::move(new_val));
+  ASSERT(new_val.empty()); // NOLINT(bugprone-use-after-move)
+}
+
+void RequestHeaderMapImpl::setTraceContextReference(absl::string_view key, absl::string_view val) {
+  ASSERT(validatedLowerCaseString(key));
+  HeaderMapImpl::removeExisting(key);
+
+  HeaderMapImpl::insertByKey(HeaderString(key), HeaderString(val));
+}
+
 } // namespace Http
 } // namespace Envoy

source/common/http/header_map_impl.h

@@ -322,7 +322,10 @@ class HeaderMapImpl : NonCopyable {
   HeaderEntryImpl& maybeCreateInline(HeaderEntryImpl** entry, const LowerCaseString& key);
   HeaderEntryImpl& maybeCreateInline(HeaderEntryImpl** entry, const LowerCaseString& key,
                                      HeaderString&& value);
-  HeaderMap::NonConstGetResult getExisting(const LowerCaseString& key);
+
+  HeaderMap::NonConstGetResult getExisting(absl::string_view key);
+  size_t removeExisting(absl::string_view key);
+
   size_t removeInline(HeaderEntryImpl** entry);
   void updateSize(uint64_t from_size, uint64_t to_size);
   void addSize(uint64_t size);
@@ -480,6 +483,12 @@ class RequestHeaderMapImpl final : public TypedHeaderMapImpl<RequestHeaderMap>,
   INLINE_REQ_RESP_STRING_HEADERS(DEFINE_INLINE_HEADER_STRING_FUNCS)
   INLINE_REQ_RESP_NUMERIC_HEADERS(DEFINE_INLINE_HEADER_NUMERIC_FUNCS)
 
+  // Tracing::TraceContext
+  absl::optional<absl::string_view> getTraceContext(absl::string_view key) const override;
+  void setTraceContext(absl::string_view key, absl::string_view val) override;
+  void setTraceContextReferenceKey(absl::string_view key, absl::string_view val) override;
+  void setTraceContextReference(absl::string_view key, absl::string_view val) override;
+
 protected:
   // NOTE: Because inline_headers_ is a variable size member, it must be the last member in the
   // most derived class. This forces the definition of the following three functions to also be

test/common/http/header_map_impl_test.cc

@@ -1370,5 +1370,29 @@ TEST_P(HeaderMapImplTest, ValidHeaderString) {
   EXPECT_FALSE(validHeaderString("abc\n"));
 }
 
+TEST_P(HeaderMapImplTest, HttpTraceContextTest) {
+  {
+    TestRequestHeaderMapImpl request_headers{{"host", "foo"}};
+    EXPECT_EQ(request_headers.getTraceContext("host").value(), "foo");
+
+    request_headers.setTraceContext("trace_key", "trace_value");
+    EXPECT_EQ(request_headers.getTraceContext("trace_key").value(), "trace_value");
+
+    std::string trace_ref_key = "trace_ref_key";
+    request_headers.setTraceContextReferenceKey(trace_ref_key, "trace_value");
+    auto* header_entry = request_headers.get(Http::LowerCaseString(trace_ref_key))[0];
+    EXPECT_EQ(reinterpret_cast<intptr_t>(trace_ref_key.data()),
+              reinterpret_cast<intptr_t>(header_entry->key().getStringView().data()));
+
+    std::string trace_ref_value = "trace_ref_key";
+    request_headers.setTraceContextReference(trace_ref_key, trace_ref_value);
+    header_entry = request_headers.get(Http::LowerCaseString(trace_ref_key))[0];
+    EXPECT_EQ(reinterpret_cast<intptr_t>(trace_ref_key.data()),
+              reinterpret_cast<intptr_t>(header_entry->key().getStringView().data()));
+    EXPECT_EQ(reinterpret_cast<intptr_t>(trace_ref_value.data()),
+              reinterpret_cast<intptr_t>(header_entry->value().getStringView().data()));
+  }
+}
+
 } // namespace Http
 } // namespace Envoy

test/test_common/utility.h

@@ -843,6 +843,33 @@ class ConditionalInitializer {
   bool ready_ ABSL_GUARDED_BY(mutex_){false};
 };
 
+namespace Tracing {
+
+class TestTraceContextImpl : public Tracing::TraceContext {
+public:
+  TestTraceContextImpl(const std::initializer_list<std::pair<std::string, std::string>>& values) {
+    for (const auto& value : values) {
+      context_map_[value.first] = value.second;
+    }
+  }
+
+  absl::optional<absl::string_view> getTraceContext(absl::string_view key) const override {
+    auto iter = context_map_.find(std::string(key));
+    if (iter == context_map_.end()) {
+      return absl::nullopt;
+    }
+    return iter->second;
+  }
+
+  void setTraceContext(absl::string_view key, absl::string_view val) override {
+    context_map_.insert({std::string(key), std::string(val)});
+  }
+
+  absl::flat_hash_map<std::string, std::string> context_map_;
+};
+
+} // namespace Tracing
+
 namespace Http {
 
 /**
@@ -1056,6 +1083,25 @@ class TestRequestHeaderMapImpl
   INLINE_REQ_NUMERIC_HEADERS(DEFINE_TEST_INLINE_NUMERIC_HEADER_FUNCS)
   INLINE_REQ_RESP_STRING_HEADERS(DEFINE_TEST_INLINE_STRING_HEADER_FUNCS)
   INLINE_REQ_RESP_NUMERIC_HEADERS(DEFINE_TEST_INLINE_NUMERIC_HEADER_FUNCS)
+
+  absl::optional<absl::string_view> getTraceContext(absl::string_view key) const override {
+    ASSERT(header_map_);
+    return header_map_->getTraceContext(key);
+  }
+  void setTraceContext(absl::string_view key, absl::string_view value) override {
+    ASSERT(header_map_);
+    header_map_->setTraceContext(key, value);
+  }
+
+  void setTraceContextReferenceKey(absl::string_view key, absl::string_view val) override {
+    ASSERT(header_map_);
+    header_map_->setTraceContextReferenceKey(key, val);
+  }
+
+  void setTraceContextReference(absl::string_view key, absl::string_view val) override {
+    ASSERT(header_map_);
+    header_map_->setTraceContextReference(key, val);
+  }
 };
 
 using TestRequestTrailerMapImpl = TestHeaderMapImplBase<RequestTrailerMap, RequestTrailerMapImpl>;