tlbmc/pacemaker/pacemaker.h

#ifndef THIRD_PARTY_MILOTIC_EXTERNAL_CC_TLBMC_PACEMAKER_PACEMAKER_H_
#define THIRD_PARTY_MILOTIC_EXTERNAL_CC_TLBMC_PACEMAKER_PACEMAKER_H_

#include <cstdint>
#include <memory>
#include <string>
#include <utility>

#include "absl/base/thread_annotations.h"
#include "absl/memory/memory.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/synchronization/mutex.h"
#include "absl/time/time.h"
#include "boost/circular_buffer.hpp"
#include "nlohmann/json.hpp"
#include "tlbmc/scheduler/scheduler.h"

namespace milotic_tlbmc {

class ShellCommandExecutor {
 public:
  virtual ~ShellCommandExecutor() = default;

  virtual absl::StatusOr<std::string> Execute(const std::string& command);
};

enum class ErrorType : uint8_t {
  kUnknown = 0,
  kServiceInactive,
  kPortNotListening,
  kMemoryUsageAboveThreshold,
};

class Pacemaker {
 public:
  struct ErrorInfo {
    nlohmann::json ToJson() const;

    ErrorType type;
    absl::Time timestamp;
  };

  struct MonitoredData {
    nlohmann::json ToJson() const;

    boost::circular_buffer<int> memory_usage_bytes =
        boost::circular_buffer<int>(50);
    boost::circular_buffer<ErrorInfo> restart_log =
        boost::circular_buffer<ErrorInfo>(50);
    int pid = -1;
    int cpu_usage = -1;
    // ActiveEnterTimestampMonotonic value of the process as captured from
    // systemctl show --property=ActiveEnterTimestampMonotonic.
    // This is only used to detect if the process was restarted outside of
    // pacemaker.
    int64_t last_active_timestamp = -1;
    absl::Time last_reset_time = absl::InfinitePast();
    bool restart_triggered = false;
    int consecutive_restart_attempts = 0;
  };

  static std::unique_ptr<Pacemaker> Create(
      absl::Duration interval,
      std::unique_ptr<ShellCommandExecutor> shell_command_executor =
          std::make_unique<ShellCommandExecutor>()) {
    return absl::WrapUnique(
        new Pacemaker(interval, std::move(shell_command_executor)));
  }

  nlohmann::json GetMonitoringData() const ABSL_LOCKS_EXCLUDED(mutex_);

  absl::Status PerformChecks();
  absl::StatusOr<int> GetPid(const std::string& process_name) const;
  absl::StatusOr<bool> IsPortListening(int port) const;
  absl::StatusOr<bool> IsServiceActive(const std::string& service_name) const;
  absl::StatusOr<int> GetMemoryUsage(int pid) const;
  absl::StatusOr<int> GetCpuUsage(const std::string& process_name) const;
  absl::StatusOr<int64_t> GetLastActiveTimestamp(
      const std::string& process_name) const;
  absl::Status RestartService(const std::string& service_name)
      ABSL_LOCKS_EXCLUDED(mutex_);
  void RecordError(ErrorType type) ABSL_LOCKS_EXCLUDED(mutex_);

  ~Pacemaker() { scheduler_.Stop(); }

 private:
  explicit Pacemaker(
      absl::Duration interval,
      std::unique_ptr<ShellCommandExecutor> shell_command_executor);

  mutable absl::Mutex mutex_;
  std::unique_ptr<ShellCommandExecutor> shell_command_executor_;
  TaskScheduler scheduler_;
  MonitoredData monitored_data_ ABSL_GUARDED_BY(mutex_);
};

}  // namespace milotic_tlbmc

#endif  // THIRD_PARTY_MILOTIC_EXTERNAL_CC_TLBMC_PACEMAKER_PACEMAKER_H_