[PROTON] Reduce msgpack serialization overhead-3

third_party/proton/csrc/include/Utility/MsgPackWriter.h

@@ -13,8 +13,16 @@ namespace proton {
 class MsgPackWriter {
 public:
   void reserve(size_t bytes);
+  size_t size() const { return out.size(); }
+  const uint8_t *data() const { return out.data(); }
 
   std::vector<uint8_t> take() &&;
+  void appendBytes(const std::vector<uint8_t> &bytes);
+  template <size_t N> void appendBytes(const uint8_t (&bytes)[N]) {
+    const auto offset = out.size();
+    out.resize(offset + N);
+    std::memcpy(out.data() + offset, bytes, N);
+  }
 
   void packNil();
   void packBool(bool value);

third_party/proton/csrc/lib/Data/TreeData.cpp

@@ -10,6 +10,7 @@
 #include <array>
 #include <cstddef>
 #include <cstdint>
+#include <deque>
 #include <functional>
 #include <limits>
 #include <map>
@@ -124,10 +125,7 @@ class TreeData::Tree {
     friend class Tree;
   };
 
-  Tree() {
-    treeNodeMap.try_emplace(TreeNode::RootId, TreeNode::RootId,
-                            TreeNode::RootId, "ROOT");
-  }
+  Tree() { treeNodes.emplace_back(TreeNode::RootId, TreeNode::RootId, "ROOT"); }
 
   size_t addNode(const std::vector<Context> &contexts, size_t parentId) {
     for (const auto &context : contexts) {
@@ -137,14 +135,14 @@ class TreeData::Tree {
   }
 
   size_t addNode(const Context &context, size_t parentId) {
-    auto &parent = treeNodeMap.at(parentId);
+    auto &parent = getNode(parentId);
     std::string_view contextName = context.name;
     auto existingChildId = parent.findChild(contextName);
     if (existingChildId != TreeNode::DummyId)
       return existingChildId;
     auto id = nextContextId++;
-    auto [it, inserted] = treeNodeMap.try_emplace(id, id, parentId, context);
-    parent.addChild(it->second.name, id);
+    treeNodes.emplace_back(id, parentId, context);
+    parent.addChild(treeNodes.back().name, id);
     return id;
   }
 
@@ -156,11 +154,11 @@ class TreeData::Tree {
     return parentId;
   }
 
-  TreeNode &getNode(size_t id) { return treeNodeMap.at(id); }
+  TreeNode &getNode(size_t id) { return treeNodes.at(id); }
 
   void upsertFlexibleMetric(size_t contextId,
                             const FlexibleMetric &flexibleMetric) {
-    auto &node = treeNodeMap.at(contextId);
+    auto &node = getNode(contextId);
     auto &flexibleMetrics = node.metricSet.flexibleMetrics;
     auto valueName = std::string(flexibleMetric.getValueName(0));
     auto it = flexibleMetrics.find(valueName);
@@ -188,31 +186,12 @@ class TreeData::Tree {
     }
   }
 
-  size_t size() const { return nextContextId; }
-
-  Tree structure() const {
-    Tree cloned;
-    cloned.nextContextId = nextContextId;
-
-    for (const auto &[id, node] : treeNodeMap) {
-      cloned.treeNodeMap.try_emplace(id, id, node.parentId, node);
-    }
-
-    for (const auto &[id, node] : treeNodeMap) {
-      auto &clonedNode = cloned.treeNodeMap.at(id);
-      clonedNode.children.reserve(node.children.size());
-      for (const auto &child : node.children) {
-        clonedNode.addChild(cloned.treeNodeMap[child.id].name, child.id);
-      }
-    }
-
-    return cloned;
-  }
+  size_t size() const { return treeNodes.size(); }
 
 private:
   size_t nextContextId = TreeNode::RootId + 1;
-  // tree node id -> tree node
-  std::unordered_map<size_t, TreeNode> treeNodeMap;
+  // Node ids are dense and assigned sequentially, so index lookup is enough.
+  std::deque<TreeNode> treeNodes;
 };
 
 json TreeData::buildHatchetJson(TreeData::Tree *tree,
@@ -461,15 +440,44 @@ TreeData::buildHatchetMsgPack(TreeData::Tree *tree,
         }
       });
   const auto &virtualRootNode = virtualTree->getNode(Tree::TreeNode::RootId);
+  auto packUncachedHatchetFrameHeader = [](MsgPackWriter &out,
+                                           std::string_view name) {
+    static constexpr uint8_t kHatchetFrameHeaderPrefix[] = {
+        0x83, // map(3)
+        0xa5, 'f', 'r', 'a', 'm', 'e',
+        0x82, // map(2)
+        0xa4, 'n', 'a', 'm', 'e'};
+    static constexpr uint8_t kHatchetFrameHeaderSuffix[] = {
+        0xa4, 't', 'y', 'p',  'e', 0xa8, 'f', 'u', 'n', 'c', 't',
+        'i',  'o', 'n', 0xa7, 'm', 'e',  't', 'r', 'i', 'c', 's'};
+    out.appendBytes(kHatchetFrameHeaderPrefix);
+    out.packStr(name);
+    out.appendBytes(kHatchetFrameHeaderSuffix);
+  };
+  // Names that fit in MsgPack fixstr are cheap enough to encode directly. Cache
+  // only longer headers so repeated linked virtual frames can skip the larger
+  // string copy without adding hash-table overhead to every small frame name.
+  constexpr size_t kCachedFrameHeaderMinNameBytes = 64;
+  std::unordered_map<std::string_view, std::vector<uint8_t>> frameHeaderCache;
   auto packHatchetFrameHeader = [&](std::string_view name) {
-    writer.packMap(3);
-    writer.packFixStrLiteral("frame");
-    writer.packMap(2);
-    writer.packFixStrLiteral("name");
-    writer.packStr(name);
-    writer.packFixStrLiteral("type");
-    writer.packFixStrLiteral("function");
-    writer.packFixStrLiteral("metrics");
+    if (name.size() < kCachedFrameHeaderMinNameBytes) {
+      packUncachedHatchetFrameHeader(writer, name);
+      return;
+    }
+
+    auto it = frameHeaderCache.find(name);
+    if (it != frameHeaderCache.end()) {
+      writer.appendBytes(it->second);
+      return;
+    }
+
+    const auto offset = writer.size();
+    packUncachedHatchetFrameHeader(writer, name);
+    std::vector<uint8_t> headerBytes;
+    headerBytes.reserve(kCachedFrameHeaderMinNameBytes + name.size());
+    headerBytes.insert(headerBytes.end(), writer.data() + offset,
+                       writer.data() + writer.size());
+    frameHeaderCache.emplace(name, std::move(headerBytes));
   };
 
   // Root metrics only carry inclusive aggregate fields. Non-root metrics also
@@ -667,6 +675,49 @@ TreeData::buildHatchetMsgPack(TreeData::Tree *tree,
               flexibleMetric.getValues()[0]);
         }
       };
+  auto packLinkedVirtualNode = [&](auto &&packLinkedVirtualNode,
+                                   const TreeData::Tree::TreeNode &treeNode,
+                                   size_t virtualNodeId) -> void {
+    const auto &virtualNode = virtualTree->getNode(virtualNodeId);
+    const auto &linkedMetrics = treeNode.metricSet.linkedMetrics;
+    const auto &linkedFlexibleMetrics =
+        treeNode.metricSet.linkedFlexibleMetrics;
+    // Write the header
+    packHatchetFrameHeader(virtualNode.name);
+    // Count linked metrics
+    auto metricEntries = 0u;
+    const auto metricsIt = linkedMetrics.find(virtualNodeId);
+    if (metricsIt != linkedMetrics.end()) {
+      metricEntries += countMetricEntries(metricsIt->second, /*isRoot=*/false);
+    }
+    // Count linked flexible metrics exist in the child <metric> helpers
+    if (!linkedFlexibleMetrics.empty()) {
+      for (const auto &child : virtualNode.children) {
+        auto it = linkedFlexibleMetrics.find(child.id);
+        if (it != linkedFlexibleMetrics.end()) {
+          metricEntries += static_cast<uint32_t>(it->second.size());
+        }
+      }
+    }
+    // Pack
+    writer.packMap(metricEntries);
+    if (metricsIt != linkedMetrics.end()) {
+      packMetrics(metricsIt->second, /*isRoot=*/false);
+    }
+    if (!linkedFlexibleMetrics.empty()) {
+      for (const auto &child : virtualNode.children) {
+        auto it = linkedFlexibleMetrics.find(child.id);
+        if (it != linkedFlexibleMetrics.end()) {
+          packFlexibleMetrics(it->second);
+        }
+      }
+    }
+    writer.packFixStrLiteral("children");
+    writer.packArray(static_cast<uint32_t>(virtualNode.children.size()));
+    for (const auto &child : virtualNode.children) {
+      packLinkedVirtualNode(packLinkedVirtualNode, treeNode, child.id);
+    }
+  };
   auto packNode = [&](auto &&packNode,
                       TreeData::Tree::TreeNode &treeNode) -> void {
     // Write the header
@@ -679,49 +730,6 @@ TreeData::buildHatchetMsgPack(TreeData::Tree *tree,
     packMetrics(treeNode.metricSet.metrics, isRoot);
     packFlexibleMetrics(treeNode.metricSet.flexibleMetrics);
 
-    auto packLinkedVirtualNode = [&](auto &&packLinkedVirtualNode,
-                                     size_t virtualNodeId) -> void {
-      const auto &virtualNode = virtualTree->getNode(virtualNodeId);
-      auto &linkedMetrics = treeNode.metricSet.linkedMetrics;
-      auto &linkedFlexibleMetrics = treeNode.metricSet.linkedFlexibleMetrics;
-      // Write the header
-      packHatchetFrameHeader(virtualNode.name);
-      // Count linked metrics
-      auto metricEntries = 0u;
-      const auto metricsIt = linkedMetrics.find(virtualNodeId);
-      if (metricsIt != linkedMetrics.end()) {
-        metricEntries +=
-            countMetricEntries(metricsIt->second, /*isRoot=*/false);
-      }
-      // Count linked flexible metrics exist in the child <metric> helpers
-      if (!linkedFlexibleMetrics.empty()) {
-        for (const auto &child : virtualNode.children) {
-          auto it = linkedFlexibleMetrics.find(child.id);
-          if (it != linkedFlexibleMetrics.end()) {
-            metricEntries += static_cast<uint32_t>(it->second.size());
-          }
-        }
-      }
-      // Pack
-      writer.packMap(metricEntries);
-      if (metricsIt != treeNode.metricSet.linkedMetrics.end()) {
-        packMetrics(metricsIt->second, /*isRoot=*/false);
-      }
-      if (!linkedFlexibleMetrics.empty()) {
-        for (const auto &child : virtualNode.children) {
-          auto it = linkedFlexibleMetrics.find(child.id);
-          if (it != linkedFlexibleMetrics.end()) {
-            packFlexibleMetrics(it->second);
-          }
-        }
-      }
-      writer.packFixStrLiteral("children");
-      writer.packArray(static_cast<uint32_t>(virtualNode.children.size()));
-      for (const auto &child : virtualNode.children) {
-        packLinkedVirtualNode(packLinkedVirtualNode, child.id);
-      }
-    };
-
     const bool hasLinkedTargets =
         !treeNode.metricSet.linkedMetrics.empty() ||
         !treeNode.metricSet.linkedFlexibleMetrics.empty();
@@ -737,7 +745,7 @@ TreeData::buildHatchetMsgPack(TreeData::Tree *tree,
     }
     if (hasLinkedTargets) {
       for (const auto &virtualChild : virtualRootNode.children) {
-        packLinkedVirtualNode(packLinkedVirtualNode, virtualChild.id);
+        packLinkedVirtualNode(packLinkedVirtualNode, treeNode, virtualChild.id);
       }
     }
   };

third_party/proton/csrc/lib/Utility/MsgPackWriter.cpp

@@ -22,6 +22,12 @@ void MsgPackWriter::reserve(size_t bytes) { out.reserve(bytes); }
 
 std::vector<uint8_t> MsgPackWriter::take() && { return std::move(out); }
 
+void MsgPackWriter::appendBytes(const std::vector<uint8_t> &bytes) {
+  const auto offset = out.size();
+  out.resize(offset + bytes.size());
+  std::memcpy(out.data() + offset, bytes.data(), bytes.size());
+}
+
 void MsgPackWriter::packNil() { out.push_back(0xc0); }
 
 void MsgPackWriter::packBool(bool value) { out.push_back(value ? 0xc3 : 0xc2); }