"Compiled" static header maps instead of big trie

source/common/common/BUILD

@@ -473,6 +473,11 @@ envoy_cc_library(
     ],
 )
 
+envoy_cc_library(
+    name = "compiled_string_map_lib",
+    hdrs = ["compiled_string_map.h"],
+)
+
 envoy_cc_library(
     name = "packed_struct_lib",
     hdrs = ["packed_struct.h"],

source/common/common/compiled_string_map.h

@@ -0,0 +1,157 @@
+#pragma once
+
+#include <algorithm>
+#include <string>
+#include <vector>
+
+#include "absl/strings/string_view.h"
+
+namespace Envoy {
+
+/**
+ * This is a specialized structure intended for static header maps, but
+ * there may be other use cases.
+ * The structure is:
+ * 1. a length-based lookup table so only keys the same length as the
+ * target key are considered.
+ * 2. a trie that branches on the "most divisions" position of the key.
+ *
+ * For example, if we consider the case where the set of headers is
+ * `x-prefix-banana`
+ * `x-prefix-babana`
+ * `x-prefix-apple`
+ * `x-prefix-pineapple`
+ * `x-prefix-barana`
+ * `x-prefix-banaka`
+ *
+ * A standard front-first trie looking for `x-prefix-banana` would walk
+ * 7 nodes through the tree, first for `x`, then for `-`, etc.
+ *
+ * This structure first jumps to matching length, eliminating in this
+ * example case apple and pineapple.
+ * Then the "best split" node is on
+ *   `x-prefix-banana`
+ *               ^
+ * so the first node has 3 non-miss branches, n, b and r for that position.
+ * Down that n branch, the "best split" is on
+ *   `x-prefix-banana`
+ *                 ^
+ * which has two branches, n or k.
+ * Down the n branch is the leaf node (only `x-prefix-banana` remains) - at
+ * this point a regular string-compare checks if the key is an exact match
+ * for the string node.
+ */
+template <class Value> class CompiledStringMap {
+  using FindFn = std::function<Value(const absl::string_view&)>;
+
+public:
+  using KV = std::pair<std::string, Value>;
+  /**
+   * Returns the value with a matching key, or the default value
+   * (typically nullptr) if the key was not present.
+   * @param key the key to look up.
+   */
+  Value find(const absl::string_view& key) const {
+    if (key.size() >= table_.size() || table_[key.size()] == nullptr) {
+      return {};
+    }
+    return table_[key.size()](key);
+  };
+  /**
+   * Construct the lookup table. This is a somewhat slow multi-pass
+   * operation - using this structure is not recommended unless the
+   * table is initialize-once, use-many.
+   * @param initial a vector of key->value pairs.
+   */
+  void compile(std::vector<KV> initial) {
+    if (initial.empty()) {
+      return;
+    }
+    size_t longest = 0;
+    for (const KV& pair : initial) {
+      longest = std::max(pair.first.size(), longest);
+    }
+    table_.resize(longest + 1);
+    std::sort(initial.begin(), initial.end(),
+              [](const KV& a, const KV& b) { return a.first.size() < b.first.size(); });
+    auto it = initial.begin();
+    for (size_t i = 0; i <= longest; i++) {
+      auto start = it;
+      while (it != initial.end() && it->first.size() == i) {
+        it++;
+      }
+      if (it != start) {
+        std::vector<KV> node_contents;
+        node_contents.reserve(it - start);
+        std::copy(start, it, std::back_inserter(node_contents));
+        table_[i] = createEqualLengthNode(node_contents);
+      }
+    }
+  }
+
+private:
+  static FindFn createEqualLengthNode(std::vector<KV> node_contents) {
+    if (node_contents.size() == 1) {
+      return [pair = node_contents[0]](const absl::string_view& key) -> Value {
+        if (key != pair.first) {
+          return {};
+        }
+        return pair.second;
+      };
+    }
+    struct IndexSplitInfo {
+      uint8_t index, min, max, count;
+    } best{0, 0, 0, 0};
+    for (size_t i = 0; i < node_contents[0].first.size(); i++) {
+      std::array<bool, 256> hits{};
+      IndexSplitInfo info{static_cast<uint8_t>(i), 255, 0, 0};
+      for (size_t j = 0; j < node_contents.size(); j++) {
+        uint8_t v = node_contents[j].first[i];
+        if (!hits[v]) {
+          hits[v] = true;
+          info.count++;
+          info.min = std::min(v, info.min);
+          info.max = std::max(v, info.max);
+        }
+      }
+      if (info.count > best.count) {
+        best = info;
+      }
+    }
+    std::vector<FindFn> nodes;
+    nodes.resize(best.max - best.min + 1);
+    std::sort(node_contents.begin(), node_contents.end(), [&best](const KV& a, const KV& b) {
+      return a.first[best.index] < b.first[best.index];
+    });
+    auto it = node_contents.begin();
+    for (int i = best.min; i <= best.max; i++) {
+      auto start = it;
+      while (it != node_contents.end() && it->first[best.index] == i) {
+        it++;
+      }
+      if (it != start) {
+        // Optimization was tried here, std::array<KV, 256> rather than
+        // a smaller-range vector with bounds, to keep locality and reduce
+        // comparisons. It didn't help.
+        std::vector<KV> next_contents;
+        next_contents.reserve(it - start);
+        std::copy(start, it, std::back_inserter(next_contents));
+        nodes[i - best.min] = createEqualLengthNode(next_contents);
+      }
+    }
+    return [nodes = std::move(nodes), min = best.min,
+            index = best.index](const absl::string_view& key) -> Value {
+      uint8_t k = static_cast<uint8_t>(key[index]);
+      // Possible optimization was tried here, populating empty nodes with
+      // a function that returns {} to reduce branching vs checking for null
+      // nodes. Checking for null nodes benchmarked faster.
+      if (k < min || k >= min + nodes.size() || nodes[k - min] == nullptr) {
+        return {};
+      }
+      return nodes[k - min](key);
+    };
+  }
+  std::vector<FindFn> table_;
+};
+
+} // namespace Envoy

source/common/http/BUILD

@@ -449,6 +449,7 @@ envoy_cc_library(
         ":headers_lib",
         "//envoy/http:header_map_interface",
         "//source/common/common:assert_lib",
+        "//source/common/common:compiled_string_map_lib",
         "//source/common/common:dump_state_utils",
         "//source/common/common:empty_string",
         "//source/common/common:non_copyable",

source/common/http/header_map_impl.cc

@@ -119,15 +119,15 @@ template <> HeaderMapImpl::StaticLookupTable<RequestHeaderMap>::StaticLookupTabl
   INLINE_REQ_HEADERS(REGISTER_DEFAULT_REQUEST_HEADER)
   INLINE_REQ_RESP_HEADERS(REGISTER_DEFAULT_REQUEST_HEADER)
 
-  finalizeTable();
-
   // Special case where we map a legacy host header to :authority.
   const auto handle =
       CustomInlineHeaderRegistry::getInlineHeader<RequestHeaderMap::header_map_type>(
           Headers::get().Host);
-  add(Headers::get().HostLegacy.get().c_str(), [handle](HeaderMapImpl& h) -> StaticLookupResponse {
-    return {&h.inlineHeaders()[handle.value().it_->second], &handle.value().it_->first};
-  });
+  finalizeTable(
+      {{std::string{Headers::get().HostLegacy.get()},
+        [handle](HeaderMapImpl& h) -> StaticLookupResponse {
+          return {&h.inlineHeaders()[handle.value().it_->second], &handle.value().it_->first};
+        }}});
 }
 
 template <> HeaderMapImpl::StaticLookupTable<RequestTrailerMap>::StaticLookupTable() {

source/common/http/header_map_impl.h

@@ -11,6 +11,7 @@
 #include "envoy/config/core/v3/base.pb.h"
 #include "envoy/http/header_map.h"
 
+#include "source/common/common/compiled_string_map.h"
 #include "source/common/common/non_copyable.h"
 #include "source/common/common/utility.h"
 #include "source/common/http/headers.h"
@@ -146,18 +147,23 @@ class HeaderMapImpl : NonCopyable {
    */
   template <class Interface>
   struct StaticLookupTable
-      : public TrieLookupTable<std::function<StaticLookupResponse(HeaderMapImpl&)>> {
+      : public CompiledStringMap<std::function<StaticLookupResponse(HeaderMapImpl&)>> {
     StaticLookupTable();
 
-    void finalizeTable() {
+    void finalizeTable(std::vector<KV> extra = {}) {
       CustomInlineHeaderRegistry::finalize<Interface::header_map_type>();
       auto& headers = CustomInlineHeaderRegistry::headers<Interface::header_map_type>();
-      size_ = headers.size();
+      size_ = headers.size() + extra.size();
+      std::vector<KV> input;
+      input.reserve(size_);
       for (const auto& header : headers) {
-        this->add(header.first.get().c_str(), [&header](HeaderMapImpl& h) -> StaticLookupResponse {
-          return {&h.inlineHeaders()[header.second], &header.first};
-        });
+        input.emplace_back(std::make_pair(
+            std::string{header.first.get()}, [&header](HeaderMapImpl& h) -> StaticLookupResponse {
+              return {&h.inlineHeaders()[header.second], &header.first};
+            }));
       }
+      std::copy(extra.begin(), extra.end(), std::back_inserter(input));
+      compile(input);
     }
 
     static size_t size() {
@@ -345,6 +351,8 @@ class HeaderMapImpl : NonCopyable {
   const uint32_t max_headers_kb_ = UINT32_MAX;
   // This holds the max count of the headers in the HeaderMap.
   const uint32_t max_headers_count_ = UINT32_MAX;
+
+  template <class T> friend class StaticLookupBenchmarker;
 };
 
 /**

test/common/common/BUILD

@@ -281,6 +281,11 @@ envoy_cc_test(
     ],
 )
 
+envoy_cc_test(
+    name = "compiled_string_map_test",
+    srcs = ["compiled_string_map_test.cc"],
+)
+
 envoy_cc_test(
     name = "packed_struct_test",
     srcs = ["packed_struct_test.cc"],

test/common/common/compiled_string_map_test.cc

@@ -0,0 +1,38 @@
+#include "source/common/common/compiled_string_map.h"
+
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+namespace Envoy {
+
+using testing::IsNull;
+
+TEST(CompiledStringMapTest, FindsEntriesCorrectly) {
+  CompiledStringMap<const char*> map;
+  map.compile({
+      {"key-1", "value-1"},
+      {"key-2", "value-2"},
+      {"longer-key", "value-3"},
+      {"bonger-key", "value-4"},
+      {"bonger-bey", "value-5"},
+      {"only-key-of-this-length", "value-6"},
+  });
+  EXPECT_EQ(map.find("key-1"), "value-1");
+  EXPECT_EQ(map.find("key-2"), "value-2");
+  EXPECT_THAT(map.find("key-0"), IsNull());
+  EXPECT_THAT(map.find("key-3"), IsNull());
+  EXPECT_EQ(map.find("longer-key"), "value-3");
+  EXPECT_EQ(map.find("bonger-key"), "value-4");
+  EXPECT_EQ(map.find("bonger-bey"), "value-5");
+  EXPECT_EQ(map.find("only-key-of-this-length"), "value-6");
+  EXPECT_THAT(map.find("songer-key"), IsNull());
+  EXPECT_THAT(map.find("absent-length-key"), IsNull());
+}
+
+TEST(CompiledStringMapTest, EmptyMapReturnsNull) {
+  CompiledStringMap<const char*> map;
+  map.compile({});
+  EXPECT_THAT(map.find("key-1"), IsNull());
+}
+
+} // namespace Envoy

test/common/http/header_map_impl_speed_test.cc

@@ -295,5 +295,65 @@ static void headerMapImplRemovePrefix(benchmark::State& state) {
 }
 BENCHMARK(headerMapImplRemovePrefix)->Arg(0)->Arg(1)->Arg(5)->Arg(10)->Arg(50);
 
+template <class HeaderMapType> class StaticLookupBenchmarker {
+public:
+  absl::optional<HeaderMapImpl::StaticLookupResponse> lookup(absl::string_view key) {
+    return table_.lookup(*ignored_, key);
+  }
+
+private:
+  std::unique_ptr<RequestHeaderMapImpl> ignored_ = RequestHeaderMapImpl::create();
+  HeaderMapImpl::StaticLookupTable<HeaderMapType> table_;
+};
+
+template <class HeaderMapType>
+static void headerMapImplStaticLookups(benchmark::State& state,
+                                       const std::vector<std::string>& keys) {
+  int i = keys.size();
+  StaticLookupBenchmarker<HeaderMapType> table;
+  for (auto _ : state) {
+    UNREFERENCED_PARAMETER(_);
+    auto result = table.lookup(keys[--i]);
+    if (i == 0) {
+      i = keys.size();
+    }
+    benchmark::DoNotOptimize(result);
+  }
+}
+
+static std::vector<std::string> makeMismatchedHeaders() {
+  return {
+      "x-envoy-banana",
+      "some-unknown-header",
+      "what-is-this-header",
+      "nobody-expects-this-header",
+      "another-unexpected-header",
+      "x-is-a-letter",
+      "x-y-problems-are-the-worst",
+  };
+}
+
+#define ADD_HEADER_TO_KEYS(name) keys.emplace_back(Http::Headers::get().name);
+static void bmHeaderMapImplRequestStaticLookupHits(benchmark::State& state) {
+  std::vector<std::string> keys;
+  INLINE_REQ_HEADERS(ADD_HEADER_TO_KEYS);
+  headerMapImplStaticLookups<RequestHeaderMap>(state, keys);
+}
+static void bmHeaderMapImplResponseStaticLookupHits(benchmark::State& state) {
+  std::vector<std::string> keys;
+  INLINE_RESP_HEADERS(ADD_HEADER_TO_KEYS);
+  headerMapImplStaticLookups<ResponseHeaderMap>(state, keys);
+}
+static void bmHeaderMapImplRequestStaticLookupMisses(benchmark::State& state) {
+  headerMapImplStaticLookups<RequestHeaderMap>(state, makeMismatchedHeaders());
+}
+static void bmHeaderMapImplResponseStaticLookupMisses(benchmark::State& state) {
+  headerMapImplStaticLookups<ResponseHeaderMap>(state, makeMismatchedHeaders());
+}
+BENCHMARK(bmHeaderMapImplRequestStaticLookupHits);
+BENCHMARK(bmHeaderMapImplResponseStaticLookupHits);
+BENCHMARK(bmHeaderMapImplRequestStaticLookupMisses);
+BENCHMARK(bmHeaderMapImplResponseStaticLookupMisses);
+
 } // namespace Http
 } // namespace Envoy