Add C++ test for demo

Author: phoebusm

cpp/arcticdb/async/test/test_async.cpp

@@ -141,6 +141,90 @@ TEST(Async, CollectWithThrow) {
    ARCTICDB_DEBUG(log::version(), "Collect returned");
 }
 
+TEST(Async, StatsQueryDemo) {
+    using namespace arcticdb::util::query_stats;
+    class EnableQueryStatsRAII {
+    public:
+        EnableQueryStatsRAII() {
+            QueryStats::instance().is_enabled_ = true;
+        }
+        ~EnableQueryStatsRAII() {
+            QueryStats::instance().is_enabled_ = false;
+        }
+    };
+    EnableQueryStatsRAII enable_query_stats;
+    async::TaskScheduler sched{20, 20};
+    auto& instance = QueryStats::instance();
+    auto work = [&]() {
+        std::vector<folly::Future<folly::Unit>> stuff;
+        {
+            /*
+            "AddFuture": {
+                "count": 2,
+            }
+            */
+            QUERY_STATS_ADD_GROUP(storage_ops, "AddFuture") // Always add group at the bottom of the C++ call stack
+            QUERY_STATS_ADD(count, 1)
+            QUERY_STATS_ADD(count, 1)
+            stuff.push_back(sched.submit_cpu_task(MaybeThrowTask(false))
+                .thenValue([](auto)
+                {
+                    /*
+                    "key_type": {
+                        "l": {
+                            "storage_ops": {
+                                "ListObjectsV2": {
+                                    "result_count": 123,
+                                }
+                            }
+                        }
+                    }
+                    */
+                    QUERY_STATS_ADD_GROUP(key_type, "l")
+                    QUERY_STATS_ADD_GROUP(storage_ops, "ListObjectsV2")
+                    QUERY_STATS_ADD(result_count, 123)
+                    return folly::Unit{};
+                })
+                .via(&async::io_executor()) // switching executor in the chain won't create a new child layer
+            );
+            stuff.push_back(sched.submit_io_task(MaybeThrowTask(false))
+                .thenValue([](auto)
+                {
+                    /*
+                    "key_type": {
+                        "l": {
+                            "storage_ops": {
+                                "ListObjectsV2": {
+                                    "result_count": 456,
+                                }
+                            }
+                        }
+                    }
+                    */
+                    QUERY_STATS_ADD_GROUP(key_type, "l")
+                    QUERY_STATS_ADD_GROUP(storage_ops, "ListObjectsV2")
+                    QUERY_STATS_ADD(result_count, 456)
+                    return folly::Unit{};
+                })
+            );
+            folly::collect(stuff).get();
+            ASSERT_EQ(instance.thread_local_var_.child_layers_.size(), 2); // One child_layers_ for each chain of tasks
+            const auto& add_future_layer = instance.root_layer()->next_layer_maps_[static_cast<size_t>(StatsGroupName::storage_ops)]["AddFuture"];
+            ASSERT_EQ(add_future_layer->stats_[static_cast<size_t>(StatsName::count)], 2);
+        }
+        ASSERT_EQ(instance.thread_local_var_.child_layers_.size(), 0); // child maps should be folded into root map when the root add group stat deconstructs
+        const auto& list_objects_layer = instance.root_layer()->next_layer_maps_[static_cast<size_t>(StatsGroupName::storage_ops)]["AddFuture"]
+            ->next_layer_maps_[static_cast<size_t>(StatsGroupName::key_type)]["l"]
+            ->next_layer_maps_[static_cast<size_t>(StatsGroupName::storage_ops)]["ListObjectsV2"]; 
+        ASSERT_EQ(list_objects_layer->stats_[static_cast<size_t>(StatsName::result_count)], 579); // child map stats should be summed
+    };
+    std::thread t1(work), t2(work); // mimic multithreading at python level
+    t1.join();
+    t2.join();
+    auto root_layers = instance.root_layers();
+    ASSERT_EQ(root_layers.size(), 2); // Each pseudo-python thread will create one root layer
+}
+
 using IndexSegmentReader = int;
 
 int get_index_segment_reader_impl(arcticdb::StreamId id) {

cpp/arcticdb/util/query_stats.cpp

@@ -136,6 +136,5 @@ StatsGroup::~StatsGroup() {
     if (!async::is_folly_thread && query_stats_instance.current_layer() == query_stats_instance.root_layer()) {
         query_stats_instance.merge_layers();
     }
-    auto& root_layers = QueryStats::instance().root_layers();
 }
 }

cpp/arcticdb/util/query_stats.hpp

@@ -85,8 +85,7 @@ enum class StatsName : size_t {
     count = 2
 };
 
-class StatsGroupLayer {
-public:
+struct StatsGroupLayer {
     std::array<int64_t, 3> stats_ = {0}; // sizeof(StatsName)
     std::array<std::map<std::string, std::shared_ptr<StatsGroupLayer>>, 3> next_layer_maps_; // sizeof(StatsGroupName)
     void reset_stats();
@@ -98,7 +97,7 @@ struct ChildLayer{
     std::shared_ptr<StatsGroupLayer> root_layer_;
 };
 
-// Collection of thread-local variables that reside in class QueryStats simply for the sake of convenience
+// Collection of thread-local variables. Resides in class QueryStats simply for the sake of convenience
 struct ThreadLocalQueryStatsVar {
     std::mutex child_layer_creation_mutex_;
     std::vector<ChildLayer> child_layers_;
@@ -124,7 +123,8 @@ class QueryStats {
     thread_local inline static ThreadLocalQueryStatsVar thread_local_var_;
 private:
     std::mutex root_layer_mutex_; 
-    std::vector<std::shared_ptr<StatsGroupLayer>> root_layers_; 
+    // As GIL could be released, each python thread will have their own root layer.
+    std::vector<std::shared_ptr<StatsGroupLayer>> root_layers_;
 };