Make sure we clear the logger observer before starting a new session

libkineto/src/AsyncActivityProfilerHandler.cpp

@@ -266,6 +266,7 @@ void AsyncActivityProfilerHandler::activateConfig(
     std::chrono::time_point<std::chrono::system_clock> now) {
   logger_ = ActivityProfilerController::makeLogger(*asyncRequestConfig_);
   profiler_.setLogger(logger_.get());
+  Logger::resetLoggerObservers();
   LOGGER_OBSERVER_SET_TRIGGER_ON_DEMAND();
   profiler_.configure(*asyncRequestConfig_, now);
   asyncRequestConfig_ = nullptr;

libkineto/src/GenericActivityProfiler.cpp

@@ -143,6 +143,7 @@ const std::unordered_set<std::string>& getLoggerMedataAllowList() {
 } // namespace
 
 void GenericActivityProfiler::processTraceInternal(ActivityLogger& logger) {
+  UST_LOGGER_STAGE_SCOPE(kPostProcessingStage);
   LOG(INFO) << "Processing " << traceBuffers_->cpu.size() << " CPU buffers";
   VLOG(0) << "Profile time range: " << captureWindowStartTime_ << " - "
           << captureWindowEndTime_;
@@ -974,7 +975,6 @@ time_point<system_clock> GenericActivityProfiler::performRunLoopStep(
       // for quickly handling trace request via synchronous API
       std::lock_guard<std::recursive_mutex> guard(mutex_);
       processTraceInternal(*logger_);
-      UST_LOGGER_MARK_COMPLETED(kPostProcessingStage);
       resetInternal();
       VLOG(0) << "ProcessTrace -> WaitForRequest";
       break;

libkineto/src/Logger.cpp

@@ -157,6 +157,13 @@ void Logger::setLoggerObserverOnDemand() {
   }
 }
 
+void Logger::resetLoggerObservers() {
+  std::lock_guard<std::mutex> guard(loggerObserversMutex());
+  for (auto observer : loggerObservers()) {
+    observer->reset();
+  }
+}
+
 void Logger::addLoggerObserverAddMetadata(
     const std::string& key,
     const std::string& value) {
@@ -166,6 +173,13 @@ void Logger::addLoggerObserverAddMetadata(
   }
 }
 
+USTLoggerStageGuard::~USTLoggerStageGuard() {
+  try {
+    UST_LOGGER_MARK_COMPLETED(stage_);
+  } catch (...) {
+  }
+}
+
 } // namespace KINETO_NAMESPACE
 
 #endif // USE_GOOGLE_LOG

libkineto/src/Logger.h

@@ -27,6 +27,7 @@
 #define LOGGER_OBSERVER_SET_TRIGGER_ON_DEMAND()
 #define LOGGER_OBSERVER_ADD_METADATA(key, value)
 #define UST_LOGGER_MARK_COMPLETED(stage)
+#define UST_LOGGER_STAGE_SCOPE(stage)
 #define USDT_LOGGER_EMIT_MESSAGE(usdt_type)
 #define USDT_EMIT_START_TRACE()
 #define USDT_EMIT_STOP_TRACE()
@@ -129,6 +130,8 @@ class Logger {
 
   static void setLoggerObserverOnDemand();
 
+  static void resetLoggerObservers();
+
   static void addLoggerObserverAddMetadata(const std::string& key, const std::string& value);
 
  private:
@@ -155,6 +158,23 @@ class VoidLogger {
   void operator&(std::ostream&) {}
 };
 
+// RAII helper used to ensure a UST stage row is emitted on every exit from a scope.
+// Bucketed LOG(ERROR) / LOG(WARNING) calls within the scope are picked up by the
+// emitted row as usual.
+class USTLoggerStageGuard {
+ public:
+  explicit USTLoggerStageGuard(const std::string& stage) : stage_(stage) {}
+  ~USTLoggerStageGuard();
+
+  USTLoggerStageGuard(const USTLoggerStageGuard&) = delete;
+  USTLoggerStageGuard& operator=(const USTLoggerStageGuard&) = delete;
+  USTLoggerStageGuard(USTLoggerStageGuard&&) = delete;
+  USTLoggerStageGuard& operator=(USTLoggerStageGuard&&) = delete;
+
+ private:
+  std::string stage_;
+};
+
 } // namespace KINETO_NAMESPACE
 
 #ifdef LOG // Undefine in case these are already defined (quite likely)
@@ -254,6 +274,9 @@ struct __to_constant__ {
 // UST Logger Semantics to describe when a stage is complete.
 #define UST_LOGGER_MARK_COMPLETED(stage) LOG(libkineto::LoggerOutputType::STAGE) << "Completed Stage: " << stage
 
+// RAII helper that fires UST_LOGGER_MARK_COMPLETED(stage) on scope exit.
+#define UST_LOGGER_STAGE_SCOPE(stage) libkineto::USTLoggerStageGuard LOCAL_VARNAME(ust_stage_guard)(stage)
+
 #define USDT_LOGGER_EMIT_MESSAGE(usdt_type) LOG(libkineto::LoggerOutputType::USDT) << usdt_type
 #define USDT_EMIT_START_TRACE() USDT_LOGGER_EMIT_MESSAGE("profiler_start")
 #define USDT_EMIT_STOP_TRACE() USDT_LOGGER_EMIT_MESSAGE("profiler_stop")

libkineto/src/LoggerCollector.h

@@ -41,6 +41,8 @@ class LoggerCollector : public ILoggerObserver {
   }
 
   void reset() override {
+    buckets_.clear();
+    devices.clear();
     trace_duration_ms = 0;
     event_count = 0;
     destinations.clear();

libkineto/src/SyncActivityProfilerHandler.cpp

@@ -44,6 +44,7 @@ void SyncActivityProfilerHandler::prepareTrace(const Config& config) {
     profiler_.reset();
   }
 
+  Logger::resetLoggerObservers();
   profiler_.configure(config, now);
 }
 
@@ -62,7 +63,6 @@ std::unique_ptr<ActivityTraceInterface> SyncActivityProfilerHandler::
   profiler_.processTrace(*logger);
   // Will follow up with another patch for logging URLs when ActivityTrace
   // is moved.
-  UST_LOGGER_MARK_COMPLETED(kPostProcessingStage);
 
   // Logger Metadata contains a map of LOGs collected in Kineto
   //   logger_level -> List of log lines