DynamicCounter.h 1.5 KB

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