Persist query performance stats in RocksDB.

osquery · zwass · Feb 21, 2024 · Jan 19, 2024 · Jan 19, 2024 · Jan 19, 2024
commit cc8ea639bc0bb3ef046320cd7dea7b9f8e236836
@@ -782,26 +782,25 @@ Status Config::updateSource(const std::string& source,
     if (newQueries.find(oldPack.first) == newQueries.end()) {
       // This pack was removed. Also remove performance stats.
       for (const auto& oldQuery : oldPack.second) {
-        performance_.erase(getQueryName(oldPack.first, oldQuery.first));
+        deleteDatabaseValue(kQueryPerformance, getQueryName(oldPack.first,
+                                                            oldQuery.first));
       }
       continue;
     }
     for (const auto& oldQuery : oldPack.second) {
       if (newQueries[oldPack.first].find(oldQuery.first) ==
           newQueries[oldPack.first].end()) {
         // This query was removed. Also remove performance stats.
-        performance_.erase(getQueryName(oldPack.first, oldQuery.first));
+        deleteDatabaseValue(kQueryPerformance, getQueryName(oldPack.first, oldQuery.first));
         continue;
       }
       if (queries[oldPack.first][oldQuery.first] !=
           newQueries[oldPack.first][oldQuery.first]) {
         // This query was updated. Clear the performance stats.
         auto fullName = getQueryName(oldPack.first, oldQuery.first);
         RecursiveLock lock(config_performance_mutex_);
-        if (performance_.count(fullName) != 0) {
-          LOG(INFO) << "Clearing performance stats for query: " << fullName;
-          performance_[fullName] = QueryPerformance();
-        }
+        LOG(INFO) << "Clearing performance stats for query: " << fullName;
+        setDatabaseValue(kQueryPerformance, fullName, QueryPerformance().toCSV());
       }
     }
   }
@@ -1043,7 +1042,6 @@ void Config::reset() {
   setStartTime(getUnixTime());
 
   schedule_ = std::make_unique<Schedule>();
-  std::map<std::string, QueryPerformance>().swap(performance_);
   std::map<std::string, FileCategories>().swap(files_);
   std::map<std::string, std::string>().swap(hash_);
   valid_ = false;
@@ -1087,12 +1085,13 @@ void Config::recordQueryPerformance(const std::string& name,
                                     const Row& r0,
                                     const Row& r1) {
   RecursiveLock lock(config_performance_mutex_);
-  if (performance_.count(name) == 0) {
-    performance_[name] = QueryPerformance();
+  std::string csv;
+  QueryPerformance query;
+  auto status = getDatabaseValue(kQueryPerformance, name, csv);
+  if (status.ok()) {
+    query = QueryPerformance(csv);
   }
 
-  // Grab access to the non-const schedule item.
-  auto& query = performance_.at(name);
   if (!r1.at("user_time").empty() && !r0.at("user_time").empty()) {
     auto ut1 = tryTo<long long>(r1.at("user_time"));
     auto ut0 = tryTo<long long>(r0.at("user_time"));
@@ -1132,6 +1131,11 @@ void Config::recordQueryPerformance(const std::string& name,
   query.executions += 1;
   query.last_executed = getUnixTime();
 
+  status = setDatabaseValue(kQueryPerformance, name, query.toCSV());
+  if (!status.ok()) {
+    LOG(WARNING) << "Could not write performance stats for query " << name << " to the database: " << status.getMessage();
+  }
+
   /* Clear the executing query only if a resource limit has not been hit.
      This is used by the next worker execution to denylist a query
      that triggered a watchdog resource limit. */
@@ -1153,10 +1157,11 @@ void Config::recordQueryStart(const std::string& name) {
 
 void Config::getPerformanceStats(
     const std::string& name,
-    std::function<void(const QueryPerformance& query)> predicate) const {
-  if (performance_.count(name) > 0) {
-    RecursiveLock lock(config_performance_mutex_);
-    predicate(performance_.at(name));
+    std::function<void(const QueryPerformance& query)> predicate) {
+  std::string csv;
+  auto status = getDatabaseValue(kQueryPerformance, name, csv);
+  if (status.ok()) {
+    predicate(QueryPerformance(csv));
   }
 }
 

@@ -215,16 +215,16 @@ class Config : private boost::noncopyable {
    * QueryPerformance struct, if it exists.
    *
    * @code{.cpp}
-   *   Config::get().getPerformanceStats(
+   *   Config::getPerformanceStats(
    *     "my_awesome_query",
    *     [](const QueryPerformance& query) {
    *       // use "query" here
    *     });
    * @endcode
    */
-  void getPerformanceStats(
+  static void getPerformanceStats(
       const std::string& name,
-      std::function<void(const QueryPerformance& query)> predicate) const;
+      std::function<void(const QueryPerformance& query)> predicate);
 
   /**
    * @brief Helper to access config parsers via the registry
@@ -331,9 +331,6 @@ class Config : private boost::noncopyable {
   /// Schedule of packs and their queries.
   std::unique_ptr<Schedule> schedule_;
 
-  /// A set of performance stats for each query in the schedule.
-  std::map<std::string, QueryPerformance> performance_;
-
   /// A set of named categories filled with filesystem globbing paths.
   using FileCategories = std::map<std::string, std::vector<std::string>>;
   std::map<std::string, FileCategories> files_;

@@ -8,3 +8,80 @@
  */
 
 #include "query_performance.h"
+#include "boost/lexical_cast.hpp"
+#include <boost/algorithm/string/classification.hpp>
+#include <boost/algorithm/string/split.hpp>
+#include <string>
+
+namespace osquery {
+
+// Helper function to convert a string to a given type.
+template <typename Result>
+Result convert(const std::string& source) {
+  Result result;
+  if (!boost::conversion::try_lexical_convert<Result>(source, result)) {
+    return Result();
+  }
+  return result;
+}
+
+QueryPerformance::QueryPerformance(const std::string& csv) {
+  std::vector<std::string> parts;
+  boost::split(parts, csv, boost::is_any_of(","));
+  // future proofing the size, in case additional fields are added
+  if (parts.size() < 12) {
+    return;
+  }
+
+  executions = convert<std::size_t>(parts[0]);
+  last_executed = convert<std::uint64_t>(parts[1]);
+  wall_time = convert<std::uint64_t>(parts[2]);
+  wall_time_ms = convert<std::uint64_t>(parts[3]);
+  last_wall_time_ms = convert<std::uint64_t>(parts[4]);
+  user_time = convert<std::uint64_t>(parts[5]);
+  last_user_time = convert<std::uint64_t>(parts[6]);
+  system_time = convert<std::uint64_t>(parts[7]);
+  last_system_time = convert<std::uint64_t>(parts[8]);
+  average_memory = convert<std::uint64_t>(parts[9]);
+  last_memory = convert<std::uint64_t>(parts[10]);
+  output_size = convert<std::uint64_t>(parts[11]);
+}
+
+std::string QueryPerformance::toCSV() const {
+  return std::to_string(executions) + "," + std::to_string(last_executed) +
+         "," + std::to_string(wall_time) + "," + std::to_string(wall_time_ms) +
+         "," + std::to_string(last_wall_time_ms) + "," +
+         std::to_string(user_time) + "," + std::to_string(last_user_time) +
+         "," + std::to_string(system_time) + "," +
+         std::to_string(last_system_time) + "," +
+         std::to_string(average_memory) + "," + std::to_string(last_memory) +
+         "," + std::to_string(output_size);
+}
+
+bool operator==(const QueryPerformance& l, const QueryPerformance& r) {
+  return std::tie(l.executions,
+                  l.last_executed,
+                  l.wall_time,
+                  l.wall_time_ms,
+                  l.last_wall_time_ms,
+                  l.user_time,
+                  l.last_user_time,
+                  l.system_time,
+                  l.last_system_time,
+                  l.average_memory,
+                  l.last_memory,
+                  l.output_size) == std::tie(r.executions,
+                                             r.last_executed,
+                                             r.wall_time,
+                                             r.wall_time_ms,
+                                             r.last_wall_time_ms,
+                                             r.user_time,
+                                             r.last_user_time,
+                                             r.system_time,
+                                             r.last_system_time,
+                                             r.average_memory,
+                                             r.last_memory,
+                                             r.output_size);
+}
+
+} // namespace osquery
@@ -11,6 +11,7 @@
 
 #include <cstddef>
 #include <cstdint>
+#include <string>
 
 namespace osquery {
 
@@ -55,6 +56,17 @@ struct QueryPerformance {
 
   /// Total bytes for the query
   std::uint64_t output_size{0};
+
+  // Default constructor
+  QueryPerformance() = default;
+
+  // Constructor from a CSV string
+  explicit QueryPerformance(const std::string& csv);
+
+  // Convert this struct to a CSV string
+  [[nodiscard]] std::string toCSV() const;
+
+  friend bool operator==(const QueryPerformance& l, const QueryPerformance& r);
 };
 
 } // namespace osquery
@@ -23,6 +23,7 @@ endfunction()
 function(generateOsqueryCoreTestsMergedtestsTest)
   set(source_files
     flags_tests.cpp
+    query_performance_tests.cpp
     system_test.cpp
     tables_tests.cpp
     watcher_tests.cpp

@@ -0,0 +1,56 @@
+/**
+ * Copyright (c) 2014-present, The osquery authors
+ *
+ * This source code is licensed as defined by the LICENSE file found in the
+ * root directory of this source tree.
+ *
+ * SPDX-License-Identifier: (Apache-2.0 OR GPL-2.0-only)
+ */
+
+#include "osquery/core/sql/query_performance.h"
+#include <gtest/gtest.h>
+
+namespace osquery {
+
+class QueryPerformanceTests : public testing::Test {};
+
+TEST_F(QueryPerformanceTests, test_query_performance) {
+  // Default case
+  QueryPerformance defaultStats;
+  auto emptyStats = QueryPerformance("");
+  ASSERT_EQ(defaultStats, emptyStats);
+  ASSERT_EQ("0,0,0,0,0,0,0,0,0,0,0,0", defaultStats.toCSV());
+
+  // Normal case
+  {
+    QueryPerformance expected;
+    expected.executions = 1;
+    expected.last_executed = 2;
+    expected.wall_time = 3;
+    expected.wall_time_ms = 4;
+    expected.last_wall_time_ms = 5;
+    expected.user_time = 6;
+    expected.last_user_time = 7;
+    expected.system_time = 8;
+    expected.last_system_time = 9;
+    expected.average_memory = 10;
+    expected.last_memory = 11;
+    expected.output_size = 12;
+    std::string csv = "1,2,3,4,5,6,7,8,9,10,11,12";
+    auto filledStats = QueryPerformance(csv);
+    ASSERT_EQ(expected, filledStats);
+    ASSERT_EQ(csv, expected.toCSV());
+    ASSERT_EQ(csv, filledStats.toCSV());
+  }
+
+  // Invalid case
+  {
+    std::string csv = "1,,bozo,4,5,6,7,8,9,10,11,12";
+    auto filledStats = QueryPerformance(csv);
+    ASSERT_EQ(0, filledStats.last_executed);
+    ASSERT_EQ(0, filledStats.wall_time);
+    ASSERT_EQ("1,0,0,4,5,6,7,8,9,10,11,12", filledStats.toCSV());
+  }
+}
+
+} // namespace osquery
@@ -47,6 +47,7 @@ const std::string kCarves = "carves";
 const std::string kLogs = "logs";
 const std::string kDistributedQueries = "distributed";
 const std::string kDistributedRunningQueries = "distributed_running";
+const std::string kQueryPerformance = "query_performance";
 
 const std::string kDbEpochSuffix = "epoch";
 const std::string kDbCounterSuffix = "counter";
@@ -59,7 +60,8 @@ const std::vector<std::string> kDomains = {kPersistentSettings,
                                            kLogs,
                                            kCarves,
                                            kDistributedQueries,
-                                           kDistributedRunningQueries};
+                                           kDistributedRunningQueries,
+                                           kQueryPerformance};
 
 std::atomic<bool> kDBAllowOpen(false);
 std::atomic<bool> kDBInitialized(false);

@@ -58,6 +58,9 @@ extern const std::string kDistributedQueries;
 /// The "domain" where currently running distributed queries are stored.
 extern const std::string kDistributedRunningQueries;
 
+/// The "domain" where query performance stats are stored.
+extern const std::string kQueryPerformance;
+
 /// The running version of our database schema
 const int kDbCurrentVersion = 2;