Sort output of maps when printing

For count() maps, sort by value
For others, sort by key

src/bpftrace.cpp

@@ -315,10 +315,10 @@ int BPFtrace::print_map(Map &map)
   }
   auto key(old_key);
 
+  std::vector<std::pair<std::vector<uint8_t>, std::vector<uint8_t>>> values_by_key;
+
   while (bpf_get_next_key(map.mapfd_, old_key.data(), key.data()) == 0)
   {
-    std::cout << map.name_ << map.key_.argument_value_list(*this, key) << ": ";
-
     int value_size = map.type_.size;
     if (map.type_.type == Type::count)
       value_size *= ncpus_;
@@ -329,6 +329,31 @@ int BPFtrace::print_map(Map &map)
       std::cerr << "Error looking up elem: " << err << std::endl;
       return -1;
     }
+
+    values_by_key.push_back({key, value});
+
+    old_key = key;
+  }
+
+  if (map.type_.type == Type::count)
+  {
+    std::sort(values_by_key.begin(), values_by_key.end(), [&](auto &a, auto &b)
+    {
+      return reduce_value(a.second, ncpus_) < reduce_value(b.second, ncpus_);
+    });
+  }
+  else
+  {
+    sort_by_key(map.key_.args_, values_by_key);
+  };
+
+  for (auto &pair : values_by_key)
+  {
+    auto key = pair.first;
+    auto value = pair.second;
+
+    std::cout << map.name_ << map.key_.argument_value_list(*this, key) << ": ";
+
     if (map.type_.type == Type::stack)
       std::cout << get_stack(*(uint32_t*)value.data(), false, 8);
     else if (map.type_.type == Type::ustack)
@@ -340,11 +365,9 @@ int BPFtrace::print_map(Map &map)
     else if (map.type_.type == Type::string)
       std::cout << value.data() << std::endl;
     else if (map.type_.type == Type::count)
-      std::cout << reduce_value(value) << std::endl;
+      std::cout << reduce_value(value, ncpus_) << std::endl;
     else
       std::cout << *(int64_t*)value.data() << std::endl;
-
-    old_key = key;
   }
 
   std::cout << std::endl;
@@ -396,7 +419,7 @@ int BPFtrace::print_map_quantize(Map &map)
       // New key - create a list of buckets for it
       values_by_key[key_prefix] = std::vector<uint64_t>(65);
     }
-    values_by_key[key_prefix].at(bucket) = reduce_value(value);
+    values_by_key[key_prefix].at(bucket) = reduce_value(value, ncpus_);
 
     old_key = key;
   }
@@ -456,7 +479,7 @@ int BPFtrace::print_quantize(const std::vector<uint64_t> &values) const
   return 0;
 }
 
-std::string BPFtrace::quantize_index_label(int power) const
+std::string BPFtrace::quantize_index_label(int power)
 {
   char suffix = '\0';
   if (power >= 40)
@@ -487,10 +510,10 @@ std::string BPFtrace::quantize_index_label(int power) const
   return label.str();
 }
 
-uint64_t BPFtrace::reduce_value(const std::vector<uint8_t> &value) const
+uint64_t BPFtrace::reduce_value(const std::vector<uint8_t> &value, int ncpus)
 {
   uint64_t sum = 0;
-  for (int i=0; i<ncpus_; i++)
+  for (int i=0; i<ncpus; i++)
   {
     sum += *(uint64_t*)(value.data() + i*sizeof(uint64_t*));
   }
@@ -574,4 +597,46 @@ std::string BPFtrace::resolve_usym(uint64_t addr) const
   return symbol.str();
 }
 
+void BPFtrace::sort_by_key(std::vector<SizedType> key_args,
+    std::vector<std::pair<std::vector<uint8_t>, std::vector<uint8_t>>> &values_by_key)
+{
+  int arg_offset = 0;
+  for (auto arg : key_args)
+  {
+    arg_offset += arg.size;
+  }
+
+  // Sort the key arguments in reverse order so the results are sorted by
+  // the first argument first, then the second, etc.
+  for (size_t i=key_args.size(); i-- > 0; )
+  {
+    auto arg = key_args.at(i);
+    arg_offset -= arg.size;
+
+    if (arg.type == Type::integer)
+    {
+      if (arg.size == 8)
+      {
+        std::stable_sort(values_by_key.begin(), values_by_key.end(), [&](auto &a, auto &b)
+        {
+          return *(uint64_t*)(a.first.data() + arg_offset) < *(uint64_t*)(b.first.data() + arg_offset);
+        });
+      }
+      else
+        abort();
+    }
+    else if (arg.type == Type::string)
+    {
+      std::stable_sort(values_by_key.begin(), values_by_key.end(), [&](auto &a, auto &b)
+      {
+        return strncmp((char*)(a.first.data() + arg_offset),
+                       (char*)(b.first.data() + arg_offset),
+                       STRING_SIZE) < 0;
+      });
+    }
+
+    // Other types don't get sorted
+  }
+}
+
 } // namespace bpftrace

src/bpftrace.h

@@ -34,6 +34,9 @@ public:
   std::unique_ptr<Map> stackid_map_;
   std::unique_ptr<Map> perf_event_map_;
 
+  static void sort_by_key(std::vector<SizedType> key_args,
+      std::vector<std::pair<std::vector<uint8_t>, std::vector<uint8_t>>> &values_by_key);
+
 protected:
   virtual std::set<std::string> find_wildcard_matches(std::string attach_point, std::string file_name);
   std::vector<Probe> probes_;
@@ -51,8 +54,8 @@ private:
   int print_map(Map &map);
   int print_map_quantize(Map &map);
   int print_quantize(const std::vector<uint64_t> &values) const;
-  uint64_t reduce_value(const std::vector<uint8_t> &value) const;
-  std::string quantize_index_label(int power) const;
+  static uint64_t reduce_value(const std::vector<uint8_t> &value, int ncpus);
+  static std::string quantize_index_label(int power);
   std::vector<uint8_t> find_empty_key(Map &map, size_t size) const;
 };
 

tests/bpftrace.cpp

@@ -20,6 +20,7 @@ public:
 };
 
 using ::testing::_;
+using ::testing::ContainerEq;
 using ::testing::Return;
 using ::testing::StrictMock;
 
@@ -196,6 +197,206 @@ TEST(bpftrace, add_probes_tracepoint_wildcard)
   EXPECT_EQ(bpftrace.get_special_probes().size(), 0);
 }
 
+std::pair<std::vector<uint8_t>, std::vector<uint8_t>> key_value_pair_int(std::vector<uint64_t> key, int val)
+{
+  std::pair<std::vector<uint8_t>, std::vector<uint8_t>> pair;
+  pair.first  = std::vector<uint8_t>(sizeof(uint64_t)*key.size());
+  pair.second = std::vector<uint8_t>(sizeof(uint64_t));
+
+  uint8_t *key_data = pair.first.data();
+  uint8_t *val_data = pair.second.data();
+
+  for (size_t i=0; i<key.size(); i++)
+  {
+    *(uint64_t*)(key_data + sizeof(uint64_t)*i) = key.at(i);
+  }
+  *(uint64_t*)val_data = val;
+
+  return pair;
+}
+
+std::pair<std::vector<uint8_t>, std::vector<uint8_t>> key_value_pair_str(std::vector<std::string> key, int val)
+{
+  std::pair<std::vector<uint8_t>, std::vector<uint8_t>> pair;
+  pair.first  = std::vector<uint8_t>(STRING_SIZE*key.size());
+  pair.second = std::vector<uint8_t>(sizeof(uint64_t));
+
+  uint8_t *key_data = pair.first.data();
+  uint8_t *val_data = pair.second.data();
+
+  for (size_t i=0; i<key.size(); i++)
+  {
+    strncpy((char*)key_data + STRING_SIZE*i, key.at(i).c_str(), STRING_SIZE);
+  }
+  *(uint64_t*)val_data = val;
+
+  return pair;
+}
+
+std::pair<std::vector<uint8_t>, std::vector<uint8_t>> key_value_pair_int_str(int myint, std::string mystr, int val)
+{
+  std::pair<std::vector<uint8_t>, std::vector<uint8_t>> pair;
+  pair.first  = std::vector<uint8_t>(sizeof(uint64_t) + STRING_SIZE);
+  pair.second = std::vector<uint8_t>(sizeof(uint64_t));
+
+  uint8_t *key_data = pair.first.data();
+  uint8_t *val_data = pair.second.data();
+
+  *(uint64_t*)key_data = myint;
+  strncpy((char*)key_data + sizeof(uint64_t), mystr.c_str(), STRING_SIZE);
+  *(uint64_t*)val_data = val;
+
+  return pair;
+}
+
+TEST(bpftrace, sort_by_key_int)
+{
+  StrictMock<MockBPFtrace> bpftrace;
+
+  std::vector<SizedType> key_args = { SizedType(Type::integer, 8) };
+  std::vector<std::pair<std::vector<uint8_t>, std::vector<uint8_t>>> values_by_key =
+  {
+    key_value_pair_int({2}, 12),
+    key_value_pair_int({3}, 11),
+    key_value_pair_int({1}, 10),
+  };
+  bpftrace.sort_by_key(key_args, values_by_key);
+
+  std::vector<std::pair<std::vector<uint8_t>, std::vector<uint8_t>>> expected_values =
+  {
+    key_value_pair_int({1}, 10),
+    key_value_pair_int({2}, 12),
+    key_value_pair_int({3}, 11),
+  };
+
+  EXPECT_THAT(values_by_key, ContainerEq(expected_values));
+}
+
+TEST(bpftrace, sort_by_key_int_int)
+{
+  StrictMock<MockBPFtrace> bpftrace;
+
+  std::vector<SizedType> key_args = {
+    SizedType(Type::integer, 8),
+    SizedType(Type::integer, 8),
+    SizedType(Type::integer, 8),
+  };
+  std::vector<std::pair<std::vector<uint8_t>, std::vector<uint8_t>>> values_by_key =
+  {
+    key_value_pair_int({5,2,1}, 1),
+    key_value_pair_int({5,3,1}, 2),
+    key_value_pair_int({5,1,1}, 3),
+    key_value_pair_int({2,2,2}, 4),
+    key_value_pair_int({2,3,2}, 5),
+    key_value_pair_int({2,1,2}, 6),
+  };
+  bpftrace.sort_by_key(key_args, values_by_key);
+
+  std::vector<std::pair<std::vector<uint8_t>, std::vector<uint8_t>>> expected_values =
+  {
+    key_value_pair_int({2,1,2}, 6),
+    key_value_pair_int({2,2,2}, 4),
+    key_value_pair_int({2,3,2}, 5),
+    key_value_pair_int({5,1,1}, 3),
+    key_value_pair_int({5,2,1}, 1),
+    key_value_pair_int({5,3,1}, 2),
+  };
+
+  EXPECT_THAT(values_by_key, ContainerEq(expected_values));
+}
+
+TEST(bpftrace, sort_by_key_str)
+{
+  StrictMock<MockBPFtrace> bpftrace;
+
+  std::vector<SizedType> key_args = {
+    SizedType(Type::string, STRING_SIZE),
+  };
+  std::vector<std::pair<std::vector<uint8_t>, std::vector<uint8_t>>> values_by_key =
+  {
+    key_value_pair_str({"z"}, 1),
+    key_value_pair_str({"a"}, 2),
+    key_value_pair_str({"x"}, 3),
+    key_value_pair_str({"d"}, 4),
+  };
+  bpftrace.sort_by_key(key_args, values_by_key);
+
+  std::vector<std::pair<std::vector<uint8_t>, std::vector<uint8_t>>> expected_values =
+  {
+    key_value_pair_str({"a"}, 2),
+    key_value_pair_str({"d"}, 4),
+    key_value_pair_str({"x"}, 3),
+    key_value_pair_str({"z"}, 1),
+  };
+
+  EXPECT_THAT(values_by_key, ContainerEq(expected_values));
+}
+
+TEST(bpftrace, sort_by_key_str_str)
+{
+  StrictMock<MockBPFtrace> bpftrace;
+
+  std::vector<SizedType> key_args = {
+    SizedType(Type::string, STRING_SIZE),
+    SizedType(Type::string, STRING_SIZE),
+    SizedType(Type::string, STRING_SIZE),
+  };
+  std::vector<std::pair<std::vector<uint8_t>, std::vector<uint8_t>>> values_by_key =
+  {
+    key_value_pair_str({"z", "a", "l"}, 1),
+    key_value_pair_str({"a", "a", "m"}, 2),
+    key_value_pair_str({"z", "c", "n"}, 3),
+    key_value_pair_str({"a", "c", "o"}, 4),
+    key_value_pair_str({"z", "b", "p"}, 5),
+    key_value_pair_str({"a", "b", "q"}, 6),
+  };
+  bpftrace.sort_by_key(key_args, values_by_key);
+
+  std::vector<std::pair<std::vector<uint8_t>, std::vector<uint8_t>>> expected_values =
+  {
+    key_value_pair_str({"a", "a", "m"}, 2),
+    key_value_pair_str({"a", "b", "q"}, 6),
+    key_value_pair_str({"a", "c", "o"}, 4),
+    key_value_pair_str({"z", "a", "l"}, 1),
+    key_value_pair_str({"z", "b", "p"}, 5),
+    key_value_pair_str({"z", "c", "n"}, 3),
+  };
+
+  EXPECT_THAT(values_by_key, ContainerEq(expected_values));
+}
+
+TEST(bpftrace, sort_by_key_int_str)
+{
+  StrictMock<MockBPFtrace> bpftrace;
+
+  std::vector<SizedType> key_args = {
+    SizedType(Type::integer, 8),
+    SizedType(Type::string, STRING_SIZE),
+  };
+  std::vector<std::pair<std::vector<uint8_t>, std::vector<uint8_t>>> values_by_key =
+  {
+    key_value_pair_int_str(1, "b", 1),
+    key_value_pair_int_str(2, "b", 2),
+    key_value_pair_int_str(3, "b", 3),
+    key_value_pair_int_str(1, "a", 4),
+    key_value_pair_int_str(2, "a", 5),
+    key_value_pair_int_str(3, "a", 6),
+  };
+  bpftrace.sort_by_key(key_args, values_by_key);
+
+  std::vector<std::pair<std::vector<uint8_t>, std::vector<uint8_t>>> expected_values =
+  {
+    key_value_pair_int_str(1, "a", 4),
+    key_value_pair_int_str(1, "b", 1),
+    key_value_pair_int_str(2, "a", 5),
+    key_value_pair_int_str(2, "b", 2),
+    key_value_pair_int_str(3, "a", 6),
+    key_value_pair_int_str(3, "b", 3),
+  };
+
+  EXPECT_THAT(values_by_key, ContainerEq(expected_values));
+}
+
 } // namespace bpftrace
 } // namespace test
 } // namespace bpftrace