xref: /aosp_15_r20/external/pytorch/c10/util/DynamicCounter.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 
3 #include <functional>
4 #include <memory>
5 #include <string_view>
6 
7 #include <c10/macros/Macros.h>
8 
9 namespace c10::monitor {
10 
11 class C10_API DynamicCounter {
12  public:
13   using Callback = std::function<int64_t()>;
14 
15   // Creates a dynamic counter that can be queried at any point in time by
16   // multiple backends. Only one counter with a given key can exist at any point
17   // in time.
18   //
19   // The callback is invoked every time the counter is queried.
20   // The callback must be thread-safe.
21   // The callback must not throw.
22   // The callback must not block.
23   DynamicCounter(std::string_view key, Callback getCounterCallback);
24 
25   // Unregisters the callback.
26   // Waits for all ongoing callback invocations to finish.
27   ~DynamicCounter();
28 
29  private:
30   struct Guard;
31   std::unique_ptr<Guard> guard_;
32 };
33 
34 namespace detail {
35 class DynamicCounterBackendIf {
36  public:
37   virtual ~DynamicCounterBackendIf() = default;
38 
39   virtual void registerCounter(
40       std::string_view key,
41       DynamicCounter::Callback getCounterCallback) = 0;
42   // MUST wait for all ongoing callback invocations to finish
43   virtual void unregisterCounter(std::string_view key) = 0;
44 };
45 
46 void C10_API
47     registerDynamicCounterBackend(std::unique_ptr<DynamicCounterBackendIf>);
48 } // namespace detail
49 } // namespace c10::monitor
50