#pragma once

#include "shm_helper.h"

class ShmManager {
 public:
  ShmManager();
  ~ShmManager();

  void set_gpu_info(int gpu_id, uint64_t vmem_size, uint64_t shared_handle);

  // request
  void process_requests();
  // heart beat
  void check_heart_beats();

  void run_busy_loop();

  void stop_busy_loop() {
    stop_loop_flag.store(true, std::memory_order_release);
  }

  void register_callback_on_worker_change(
      std::function<void(const std::vector<int32_t> &)> cb) {
    cb_on_worker_change = std::move(cb);
  }

 private:
  ShmHelper *shm_helper;
  std::vector<int32_t> local_worker_tgids;
  std::vector<int32_t> alive_worker_tgids;
  std::vector<int> valid_gpu_ids;
  std::atomic<bool> stop_loop_flag;
  std::function<void(const std::vector<int32_t> &)> cb_on_worker_change;
};