tlbmc/collector/sel_collector.h

#ifndef THIRD_PARTY_MILOTIC_EXTERNAL_CC_TLBMC_COLLECTOR_SEL_COLLECTOR_H_
#define THIRD_PARTY_MILOTIC_EXTERNAL_CC_TLBMC_COLLECTOR_SEL_COLLECTOR_H_

#include <cstddef>
#include <cstdint>
#include <fstream>
#include <memory>
#include <ostream>
#include <string>
#include <vector>

#include "absl/base/thread_annotations.h"
#include "absl/container/btree_map.h"
#include "absl/container/flat_hash_map.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/synchronization/mutex.h"
#include "time/clock.h"
#include "gbmc_sel_sub.h"
#include "nlohmann/json.hpp"
#include "tlbmc/collector/collector.h"
#include "sel_config.pb.h"
#include "tlbmc/scheduler/scheduler.h"

namespace milotic_tlbmc {

struct SelThreadManager {
  explicit SelThreadManager(ecclesia::Clock* clock)
      : task_scheduler(std::make_unique<TaskScheduler>(clock)) {}
  std::unique_ptr<TaskScheduler> task_scheduler;
  std::vector<int> task_ids;
};

// SelCollector is a collector that collects SEL events from the gbmc SEL
// framework and writes them to a file. The file is a jsonl file that contains a
// list of json objects. The file is rotated every chunk_size events or when the
// file is full. The file is named sel_<start_id>.<suffix> where start_id is
// the id of the first event in the file and suffix is the suffix of the file.
// The file is stored in the sel_save_path directory. The collector will recover
// existing SEL files from the directory on creation.
class SelCollector : public Collector {
 public:
  struct Params {
    ecclesia::Clock* clock = ecclesia::Clock::RealClock();
    SelConfig config;
  };
  // Creates a SelCollector with the default subscriber.
  static absl::StatusOr<std::unique_ptr<SelCollector>> Create(
      const Params& params);

  // Creates a SelCollector with a provided subscriber.
  static absl::StatusOr<std::unique_ptr<SelCollector>> Create(
      const Params& params,
      std::unique_ptr<gbmc_sel_framework::GbmcSelSubscriber> subscriber);
  ~SelCollector() override;

  // Starts the default subscriber to collect SEL events. Returns error if the
  // subscriber fails to start or recover existing SEL files fails.
  virtual absl::Status StartDefaultSubscriber();
  void Stop();

  // Returns a list of SEL entries in the forward direction starting from the
  // given ID. The list is sorted by ID in ascending order.
  virtual std::vector<nlohmann::json> GetSelEntriesForward(
      uint64_t start_id, size_t max_count) const;
  // Returns a list of SEL entries in the backward direction starting from the
  // given ID. The list is sorted by ID in descending order.
  virtual std::vector<nlohmann::json> GetSelEntriesBackward(
      uint64_t start_id, size_t max_count) const;
  // Returns a list of up to max_count SEL entries in the backward direction
  // starting from the latest ID. The list is sorted by ID in descending order.
  virtual std::vector<nlohmann::json> GetLatestSelEntries(
      size_t max_count) const;

  // Dumps SEL entries starting from the given ID to the output stream.
  // The output is in the same context as the SEL files.
  virtual void DumpSelEntriesRaw(uint64_t start_id, std::ostream& os) const;

  nlohmann::json ToJson() const override;
  nlohmann::json GetSchedulerStats() const override;

 protected:
  SelCollector() = default;

 private:
  SelCollector(
      const SelConfig& config, std::unique_ptr<SelThreadManager> thread_manager,
      std::unique_ptr<gbmc_sel_framework::GbmcSelSubscriber> subscriber);
  void DefaultSubscriberTask();
  absl::Status RecoverExistingSelFiles();
  void RotateSelFile();
  void WriteSelToFile(
      const absl::flat_hash_map<std::string, std::string>& event);

  SelConfig config_;
  std::unique_ptr<SelThreadManager> thread_manager_;
  // The default subscriber is used to collect SEL events from the gbmc SEL
  // framework, with subscribing all SEL events in this boot without history.
  // It is created by default when the SelCollector is created.
  std::unique_ptr<gbmc_sel_framework::GbmcSelSubscriber> default_subscriber_;
  std::ofstream current_sel_file_;
  // The number of SEL events written to the current SEL file.
  int current_sel_file_count_ = 0;
  // The ID of the first SEL event in the current SEL file.
  uint64_t current_sel_file_start_id_ = 0;

  mutable absl::Mutex mutex_;
  // Map of SEL file start ID to file path.
  // Use btree_map to keep file IDs in order. This allows quickly finding
  // the file containing a specific SEL ID via binary search (e.g., given an ID,
  // find the file with the greatest start_id less than or equal to it), and
  // deleting the oldest files when rotating.
  absl::btree_map<uint64_t, std::string> existing_sel_files_
      ABSL_GUARDED_BY(mutex_);
};

class EmptySelCollector final : public SelCollector {
 public:
  static std::unique_ptr<EmptySelCollector> Create();

  absl::Status StartDefaultSubscriber() override;

  std::vector<nlohmann::json> GetSelEntriesForward(
      uint64_t start_id, size_t max_count) const override;
  std::vector<nlohmann::json> GetSelEntriesBackward(
      uint64_t start_id, size_t max_count) const override;
  std::vector<nlohmann::json> GetLatestSelEntries(
      size_t max_count) const override;

  void DumpSelEntriesRaw(uint64_t start_id, std::ostream& os) const override;

  nlohmann::json ToJson() const override;

  nlohmann::json GetSchedulerStats() const override;
};

}  // namespace milotic_tlbmc

#endif  // THIRD_PARTY_MILOTIC_EXTERNAL_CC_TLBMC_COLLECTOR_SEL_COLLECTOR_H_