tlbmc/hal/shared_mem/demo/shared_mem_cli.cc

#include <csignal>  // NOLINT needed for SIGKILL
#include <cstdint>
#include <filesystem>  // NOLINT
#include <iostream>
#include <string>
#include <string_view>
#include <system_error>  // NOLINT
#include <utility>
#include <vector>

#include "absl/flags/flag.h"
#include "absl/flags/parse.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_split.h"
#include "absl/types/span.h"
#include "boost/interprocess/exceptions.hpp"  // NOLINT
#include "boost/interprocess/managed_shared_memory.hpp"  //NOLINT
#include "tlbmc/hal/shared_mem/metrics.h"
#include "tlbmc/hal/shared_mem/static_client_impl.h"
#include "tlbmc/hal/shared_mem/static_server_impl.h"
#include "tlbmc/hal/shared_mem/static_shm_common.h"

ABSL_FLAG(std::string, command, "",
          "Command to execute: 'get_metrics', 'check_segment', "
          "'remove_segment', 'get_metrics_debug', 'get_sensor_value', "
          "'set_sensor_value', 'create_segment', 'list_sensors', 'crash_test'");
ABSL_FLAG(bool, unit_test, false, "Use shared memory for unit test.");
ABSL_FLAG(std::string, sensor_name, "",
          "The name of the sensor. For create_segment, it can be a "
          "comma-separated list of sensors.");
ABSL_FLAG(float, sensor_value, 0.0f, "The value to set for the sensor.");

int main(int argc, char** argv) {
  absl::ParseCommandLine(argc, argv);
  using ::milotic_tlbmc::StaticSharedMemoryClient;
  using ::milotic_tlbmc::StaticSharedMemoryServer;

  using ::milotic_tlbmc::TlbmcMetrics;
  std::string command = absl::GetFlag(FLAGS_command);

  bool unit_test = absl::GetFlag(FLAGS_unit_test);

  if (unit_test) {
    StaticSharedMemoryClient::SetupInstanceForUnitTest();
  }

  if (command == "create_segment") {
    if (!unit_test) {
      std::cerr << "create_segment is only supported in unit_test mode\n";
      return 1;
    }
    std::string sensors_str = absl::GetFlag(FLAGS_sensor_name);
    std::vector<std::string> sensors =
        absl::StrSplit(sensors_str, ',', absl::SkipEmpty());
    if (sensors.empty()) {
      std::cerr << "At least one sensor must be provided for create_segment.\n";
      return 1;
    }
    auto status =
        StaticSharedMemoryServer::InitializeSharedMemoryForUnitTest(sensors);
    if (!status.ok()) {
      std::cerr << "Failed to initialize segment: " << status << "\n";
      return 1;
    }
    return 0;
  }

  if (command == "get_sensor_value") {
    std::string sensor_name = absl::GetFlag(FLAGS_sensor_name);
    if (sensor_name.empty()) {
      std::cerr << "--sensor_name is required for get_sensor_value\n";
      return 1;
    }

    auto& client = static_cast<StaticSharedMemoryClient&>(
        StaticSharedMemoryClient::GetInstance());

    absl::StatusOr<std::pair<float, uint64_t>> result =
        client.ReadSensorValue(sensor_name);
    if (!result.ok()) {
      std::cerr << "Failed to read sensor value for " << sensor_name << ": "
                << result.status() << "\n";
      return 1;
    }
    std::cout << "Sensor: " << sensor_name << ", Value: " << result->first
              << ", Timestamp: " << result->second << "\n";
    return 0;
  }

  if (command == "set_sensor_value") {
    std::string sensor_name = absl::GetFlag(FLAGS_sensor_name);
    if (sensor_name.empty()) {
      std::cerr << "--sensor_name is required for set_sensor_value\n";
      return 1;
    }
    float val = absl::GetFlag(FLAGS_sensor_value);
    absl::Status status =
        StaticSharedMemoryClient::GetInstance().UpdateSensorValue(sensor_name,
                                                                  val);
    if (!status.ok()) {
      std::cerr << "Failed to set sensor value for " << sensor_name << ": "
                << status << "\n";
      return 1;
    }
    return 0;
  }

  if (command == "list_sensors") {
    try {
      milotic_tlbmc::TlbmcSharedMemoryType segment(
          boost::interprocess::open_only, milotic_tlbmc::kStaticShmName);

      for (auto it = segment.named_begin(); it != segment.named_end(); ++it) {
        const std::string_view name(it->name(), it->name_length());
        if (name.empty()) {
          continue;
        }
        if (name.starts_with("__")) {
          continue;  // Skip prefix metadata
        }
        std::cout << name << "\n";
      }
    } catch (const boost::interprocess::interprocess_exception& e) {
      std::cerr << "Failed to open shared memory to list sensors: " << e.what()
                << "\n";
      return 1;
    }
    return 0;
  }

  if (command == "crash_test") {
    milotic_tlbmc::TlbmcSharedMemoryType sensors_memory(
        boost::interprocess::open_only, milotic_tlbmc::kStaticShmName);

    auto crash_fn = [] {
      std::cerr << "Crashing the process under lock\n";
      pid_t current_pid = getpid();
      kill(current_pid, SIGKILL);
    };
    sensors_memory.atomic_func(crash_fn);
  }

  if (command == "check_segment") {
    auto& client = static_cast<StaticSharedMemoryClient&>(
        StaticSharedMemoryClient::GetInstance());
    bool segment_exists = client.IsSharedMemoryReady();
    std::cerr << "segment_exists: " << segment_exists << '\n';
    return 0;
  }

  if (command == "get_metrics") {
    const TlbmcMetrics* metrics =
        StaticSharedMemoryClient::GetInstance().GetMetrics();
    std::cerr << "Metrics pointer is " << metrics << '\n';
    if (metrics != nullptr) {
      std::cerr << "Metrics is " << metrics->ToJson().dump(2) << '\n';
    } else {
      return 1;
    }
    return 0;
  }

  if (command == "get_metrics_debug") {
    const TlbmcMetrics* metrics =
        StaticSharedMemoryClient::GetInstance().GetMetrics();
    std::cerr << "Metrics pointer is " << metrics << '\n';
    if (metrics != nullptr) {
      std::cerr << "Metrics is " << metrics->ToDebugJson().dump(2) << '\n';
    }
    return 0;
  }

  if (command == "remove_segment") {
    if (!unit_test) {
      std::cerr << "Removing segment is only allowed in unit test mode.\n";
      return 1;
    }
    boost::interprocess::shared_memory_object::remove(
        milotic_tlbmc::kStaticShmName);
    std::error_code ec;
    std::filesystem::remove("/tmp/tlbmc/static_shm_initialized", ec);
    if (ec) {
      std::cerr << "Error removing initialization file: " << ec.message()
                << '\n';
    }
    return 0;
  }

  return 0;
}