tlbmc/hft/core/manager_impl.cc

#include "tlbmc/hft/core/manager_impl.h"

#include <sys/param.h>
#include <unistd.h>

#include <algorithm>
#include <atomic>
#include <cstdint>
#include <iterator>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/functional/any_invocable.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/strings/substitute.h"
#include "absl/synchronization/mutex.h"
#include "absl/time/clock.h"
#include "absl/time/time.h"
#include "g3/macros.h"
#include "tlbmc/adapter/data_source.h"
#include "fru_identifier.pb.h"
#include "fru_payload.pb.h"
#include "identifier.pb.h"
#include "payload.pb.h"
#include "sensor_payload.pb.h"
#include "subscription_params.pb.h"
#include "tlbmc/hft/core/edge_filter.h"
#include "tlbmc/hft/core/manager.h"
#include "tlbmc/hft/core/scheduler.h"
#include "resource.pb.h"
#include "status.pb.h"

namespace milotic_hft {

namespace {

bool IsDataPopulated(const Payload& data) {
  // For all sensors in the batch, check if there are any readings.
  if (data.has_high_frequency_sensors_readings_batch()) {
    for (const auto& sensor : data.high_frequency_sensors_readings_batch()
                                  .high_frequency_sensors()) {
      if (!sensor.timestamped_readings().empty()) {
        return true;
      }
    }
  }
  return false;
}

absl::StatusOr<absl::Time> GetLastSampledTime(const Payload& data) {
  if (data.has_high_frequency_sensors_readings_batch()) {
    const auto& sensors =
        data.high_frequency_sensors_readings_batch().high_frequency_sensors();
    if (sensors.empty()) {
      return absl::NotFoundError("No sensors found in the batch.");
    }
    const auto& readings = sensors.rbegin()->timestamped_readings();
    if (readings.empty()) {
      return absl::NotFoundError("No readings found.");
    }

    // Get the last sampled time from the last sensor.
    return absl::FromUnixNanos(readings.rbegin()->timestamp_ns());
  }
  return absl::InvalidArgumentError("Unsupported payload type.");
}

// Returns the absolute, power-of-two-aligned bin index that `timestamp_ns`
// falls into for a `sampling_interval_ms`-wide bin. The resampler and EndOfBin
// share this so their notion of a bin stays identical.
int64_t BinIndex(int64_t timestamp_ns, int64_t sampling_interval_ms) {
  return absl::ToInt64Milliseconds(absl::Nanoseconds(timestamp_ns)) /
         sampling_interval_ms;
}

// Decimates `readings` to one reading per absolute, power-of-two-aligned bin,
// keeping the first reading seen in each bin whose index is strictly greater
// than `last_emitted_bin`. `last_emitted_bin` is updated to the highest bin
// emitted and is carried across collect cycles by the caller.
//
// The subscription interval is always a power of two (clamped in
// AddSubscription). Decimating by the absolute power-of-two bin is correct
// regardless of the source poll rate: the data source is not guaranteed to poll
// at a power-of-two interval -- it may run at a sensor's
// static_refresh_interval, an override, or kDefaultSensorSamplingInterval
// (1000ms), faster or slower than this subscription -- but at most one reading
// per bin is ever emitted. Because bins are absolute and the interval is a
// power of two, a coarser subscription's bins are unions of a finer one's, so
// its emitted samples are a strict subset of the finer subscription's: multiple
// subscribers to the same sensor do not perturb each other.
//
// Requiring a strictly greater bin (rather than merely "different from the
// last") does double duty: it dedups across collect cycles (a bin straddling a
// batch boundary is not re-emitted, since last_emitted_bin persists), and it
// makes decimation tolerant of non-monotonic input -- if a wall-clock step
// makes readings arrive out of order, readings in an already-emitted bin are
// dropped rather than double-emitted. (A large backward step can still hide
// readings from the source-buffer slice upstream; fully closing that needs a
// monotonic poll clock at the sensor layer.)
std::vector<HighFrequencySensorReading> ResampleHighFrequencySensorsReadings(
    const HighFrequencySensorsReadings& readings, int sampling_interval_ms,
    int64_t& last_emitted_bin) {
  std::vector<HighFrequencySensorReading> resampled_readings;
  const int64_t interval_ms = sampling_interval_ms;
  for (const auto& reading : readings.timestamped_readings()) {
    const int64_t bin = BinIndex(reading.timestamp_ns(), interval_ms);
    if (bin > last_emitted_bin) {
      resampled_readings.push_back(reading);
      last_emitted_bin = bin;
    }
  }
  return resampled_readings;
}

void ResampleReadingsInplace(Payload& raw_data, int sampling_interval_ms,
                             int64_t& last_emitted_bin) {
  // Resample high frequency sensors readings. A sensor identifier yields a
  // single sensor per collect, so one last_emitted_bin tracks that sensor.
  if (raw_data.has_high_frequency_sensors_readings_batch()) {
    for (auto& sensor :
         *raw_data.mutable_high_frequency_sensors_readings_batch()
              ->mutable_high_frequency_sensors()) {
      if (sensor.timestamped_readings().empty()) {
        continue;
      }
      std::vector<HighFrequencySensorReading> resampled_readings =
          ResampleHighFrequencySensorsReadings(sensor, sampling_interval_ms,
                                               last_emitted_bin);
      auto* readings = sensor.mutable_timestamped_readings();
      readings->Clear();
      readings->Reserve(static_cast<int>(resampled_readings.size()));
      for (auto& resampled_reading : resampled_readings) {
        (*readings->Add()) = std::move(resampled_reading);
      }
    }
  }
}

// Whether a per-identifier payload has anything worth delivering: at least one
// reading, a STATUS_STALE/STATUS_MISSING marker, or FRU data. A sensor with no
// fresh sample this window and a healthy status leaves only an empty OK sensor
// entry, which we suppress rather than ship as an empty batch.
bool HasDeliverableContent(const Payload& payload) {
  if (payload.has_fru_batch() && !payload.fru_batch().frus().empty()) {
    return true;
  }
  if (payload.has_high_frequency_sensors_readings_batch()) {
    for (const auto& sensor : payload.high_frequency_sensors_readings_batch()
                                  .high_frequency_sensors()) {
      if (!sensor.timestamped_readings().empty()) {
        return true;
      }
      const milotic_hft::Status status = sensor.state().status();
      if (status == milotic_hft::STATUS_STALE ||
          status == milotic_hft::STATUS_MISSING) {
        return true;
      }
    }
  }
  return false;
}

// Returns the State the data source stamped on the (single) sensor entry of a
// collected payload, or a default (STATUS_UNKNOWN) State if absent. The store
// sets this on every payload regardless of whether readings are present, so it
// is the source of truth for the sensor's health.
milotic_hft::State GetCarriedState(const Payload& data) {
  if (data.has_high_frequency_sensors_readings_batch()) {
    const auto& sensors =
        data.high_frequency_sensors_readings_batch().high_frequency_sensors();
    if (!sensors.empty()) {
      return sensors.begin()->state();
    }
  }
  return milotic_hft::State();
}

}  // namespace

// ResourceMonitor implementation.
// A nullptr mutation will not update the data source.
absl::Status SubscriptionManagerImpl::ResourceMonitor::AddSubscriber(
    ActiveSubscription* subscription, SensorMutationBatch* mutation) {
  LOG(INFO) << "Adding subscriber " << subscription->GetId() << " to "
            << identifier_;
  subscribers_.insert(subscription);
  lowest_sampling_interval_subscriptions_.insert(
      subscription->GetParams().sampling_interval_ms());
  // Calculate the batch size for the data source.
  int current_batch_size =
      subscription->GetParams().export_interval_ms() /
      std::max(1, subscription->GetParams().sampling_interval_ms());
  max_batch_sizes_.insert(current_batch_size);
  LOG(INFO) << "Max batch size for " << identifier_ << " is "
            << current_batch_size;

  absl::Status status = ConfigureDataSource(mutation);
  if (!status.ok()) {
    // If configuration fails, we need to remove the subscriber we just added
    // so that we don't leave the resource monitor in an inconsistent state.
    // We do this manually instead of calling RemoveSubscriber because
    // RemoveSubscriber itself calls ConfigureDataSource, which might fail again
    // or reset state undesirably.
    RemoveSubscriberInternal(subscription);

    return status;
  }
  return absl::OkStatus();
}

void SubscriptionManagerImpl::ResourceMonitor::RemoveSubscriberInternal(
    ActiveSubscription* subscription) {
  // Get the sampling interval for the subscription.
  int sampling_interval_ms_to_remove =
      subscription->GetParams().sampling_interval_ms();

  // Get the batch size for the subscription.
  int current_batch_size =
      subscription->GetParams().export_interval_ms() /
      std::max(1, subscription->GetParams().sampling_interval_ms());

  // Remove the subscription from the set of subscribers.
  subscribers_.erase(subscription);

  // Remove the sampling interval from the set of lowest sampling intervals.
  auto it_interval = lowest_sampling_interval_subscriptions_.find(
      sampling_interval_ms_to_remove);
  if (it_interval != lowest_sampling_interval_subscriptions_.end()) {
    lowest_sampling_interval_subscriptions_.erase(it_interval);
  }

  // Remove the batch size from the set of max batch sizes for the data
  // source.
  auto it_batch = max_batch_sizes_.find(current_batch_size);
  if (it_batch != max_batch_sizes_.end()) {
    max_batch_sizes_.erase(it_batch);
  }
}

absl::Status SubscriptionManagerImpl::ResourceMonitor::RemoveSubscriber(
    ActiveSubscription* subscription, SensorMutationBatch* mutation) {
  LOG(INFO) << "Removing subscriber " << subscription->GetId() << " from "
            << identifier_ << " with sampling interval "
            << current_sampling_interval_ms_;
  if (!subscribers_.contains(subscription)) {
    return absl::InvalidArgumentError(
        absl::StrCat("Subscriber ", subscription->GetId(), " not found."));
  }
  RemoveSubscriberInternal(subscription);
  ECCLESIA_RETURN_IF_ERROR(ConfigureDataSource(mutation));
  return absl::OkStatus();
}

SubscriptionManagerImpl::ResourceMonitor::DebugSnapshot
SubscriptionManagerImpl::ResourceMonitor::GetDebugSnapshot() const {
  DebugSnapshot snapshot;
  snapshot.identifier = identifier_;
  snapshot.current_sampling_interval_ms = current_sampling_interval_ms_;
  snapshot.current_max_batch_size = current_max_batch_size_;
  snapshot.subscriber_ids.reserve(subscribers_.size());
  for (const ActiveSubscription* sub : subscribers_) {
    snapshot.subscriber_ids.push_back(sub->GetId());
  }
  snapshot.subscribed_sampling_intervals_ms.assign(
      lowest_sampling_interval_subscriptions_.begin(),
      lowest_sampling_interval_subscriptions_.end());
  snapshot.subscribed_max_batch_sizes.assign(max_batch_sizes_.begin(),
                                             max_batch_sizes_.end());
  return snapshot;
}

absl::Status SubscriptionManagerImpl::ResourceMonitor::ConfigureDataSource(
    SensorMutationBatch* mutation) {
  int lowest_sampling_interval_ms =
      lowest_sampling_interval_subscriptions_.empty()
          ? -1
          : *lowest_sampling_interval_subscriptions_.begin();
  // Update the max batch size for the data source.
  int max_batch_size =
      max_batch_sizes_.empty() ? -1 : *max_batch_sizes_.begin();

  if (mutation != nullptr) {
    // Reset is always applied in unison for both sampling interval and max
    // batch size, as lowest_sampling_interval_subscriptions_ and
    // max_batch_sizes_ are guaranteed to be empty at the same time.
    if (max_batch_size == -1) {
      LOG(INFO) << "Resetting sensor " << identifier_;
      ECCLESIA_RETURN_IF_ERROR(mutation->ResetSensor(identifier_));

    } else {
      if (max_batch_size != current_max_batch_size_) {
        LOG(INFO) << "Setting max batch size for " << identifier_ << " to "
                  << max_batch_size;
        ECCLESIA_RETURN_IF_ERROR(
            mutation->SetBatchSize(identifier_, max_batch_size));
      }

      if (lowest_sampling_interval_ms != current_sampling_interval_ms_) {
        LOG(INFO) << "Setting effective sampling interval for " << identifier_
                  << " to " << lowest_sampling_interval_ms;
        ECCLESIA_RETURN_IF_ERROR(mutation->SetSamplingInterval(
            identifier_, lowest_sampling_interval_ms));
      }
    }
  }
  current_sampling_interval_ms_ = lowest_sampling_interval_ms;
  current_max_batch_size_ = max_batch_size;
  return absl::OkStatus();
}

SubscriptionManagerImpl::SubscriptionManagerImpl(
    std::unique_ptr<DataSource> data_source,
    std::unique_ptr<Scheduler> scheduler)
    : next_subscription_id_(0),
      task_scheduler_(std::move(scheduler)),
      data_source_(std::move(data_source)) {}

SubscriptionManagerImpl::~SubscriptionManagerImpl() { task_scheduler_->Stop(); }

// ActiveSubscription implementation.
void ActiveSubscription::DeliverData(Payload&& raw_data) {
  if (batches_remaining_ == 0) {
    return;
  }

  // Deliver the data to the client.
  on_data_callback_(std::move(raw_data));

  // If batches_remaining_ is -1, the subscription is indefinite and will be
  // ended when the client unsubscribes.
  if (batches_remaining_ < 0) {
    return;
  }

  // Subscription is limited to a fixed number of batches.
  // If the number of batches is reached, the subscription will be ended.
  --batches_remaining_;
  DLOG(INFO) << "Subscription " << id_ << " delivered batch. "
             << batches_remaining_ << " batches remaining.";
  if (batches_remaining_ == 0) {
    (void)manager_->Unsubscribe(shared_from_this());
  }
}

// ActiveSubscription implementation.
ActiveSubscription::ActiveSubscription(
    const SubscriptionId& id, const milotic_hft::SubscriptionParams& params,
    absl::AnyInvocable<void(Payload&&)> on_data_callback,
    SubscriptionManagerImpl* manager, Scheduler* task_scheduler,
    DataSource* data_source, std::vector<std::unique_ptr<SampleFilter>> filters)
    : id_(id),
      params_(params),
      on_data_callback_(std::move(on_data_callback)),
      manager_(manager),
      batches_remaining_(params.num_batches() == 0 ? -1 : params.num_batches()),
      last_sampled_time_ns_(params.identifiers_size()),
      stream_state_(params.identifiers_size()),
      task_scheduler_(task_scheduler),
      data_source_(data_source) {
  LOG(INFO) << "ActiveSubscription " << id_
            << " created. num_batches: " << params_.num_batches();
  // Install pre-built edge filters and derive each identifier's
  // status-heartbeat cadence from the filter's heartbeat window.
  for (int i = 0; i < params_.identifiers_size(); ++i) {
    // NO_CDC: stream_state_ allocated parallel to params_.identifiers().
    stream_state_[i].sample_filter = std::move(filters[i]);
    // Reuse the filter's heartbeat cadence for the out-of-band status signal so
    // the heartbeat window is defined in exactly one place (the filter).
    if (stream_state_[i].sample_filter != nullptr) {
      stream_state_[i].status_heartbeat_windows =
          stream_state_[i].sample_filter->HeartbeatWindows();
    }
  }
}

// ActiveSubscription implementation.
void ActiveSubscription::Begin() {
  // Schedule all the tasks for this subscription.
  // The task calls the data source to collect the data and then goes over
  // the data and downsample to the desired sampling interval in the
  // subscription params.

  auto task = [weak_self = weak_from_this()](
                  absl::AnyInvocable<void()> on_done) mutable {
    std::shared_ptr<ActiveSubscription> self = weak_self.lock();
    if (self == nullptr) {
      on_done();
      return;
    }
    bool should_deliver = false;
    // Collect the data for all the identifiers in the subscription params.
    Payload consolidated_data;
    const auto& identifiers = self->params_.identifiers();
    for (int i = 0; i < identifiers.size(); ++i) {
      const auto& identifier = identifiers[i];
      Payload identifier_payload;
      absl::Time last_sampled_time;
      switch (identifier.identifier_case()) {
        case milotic_hft::Identifier::kSensorIdentifier:
          last_sampled_time =
              self->ProcessSensorIdentifier(identifier, i, identifier_payload);
          break;
        case milotic_hft::Identifier::kFruIdentifier:
          last_sampled_time =
              self->ProcessFruIdentifier(identifier, i, identifier_payload);
          break;
        default:
          continue;
      }
      if (HasDeliverableContent(identifier_payload)) {
        // Only deliver when this identifier actually contributed: new readings,
        // a STATUS_MISSING marker, or a FRU snapshot. A steady or fully
        // edge-filtered sensor adds nothing, so we avoid an empty batch (which
        // would also wrongly consume the num_batches budget).
        consolidated_data.MergeFrom(identifier_payload);
        should_deliver = true;
      }

      // NO_CDC: last_sampled_time_ns_ Allocated parallel to
      // params_.identifiers().
      self->last_sampled_time_ns_[i].store(absl::ToUnixNanos(last_sampled_time),
                                           std::memory_order_relaxed);
    }
    if (should_deliver) {
      self->DeliverData(std::move(consolidated_data));
    }
    on_done();
  };

  task_id_ = task_scheduler_->ScheduleAsync(
      task, absl::Milliseconds(params_.export_interval_ms()));
}

void ActiveSubscription::SignalStatusIfChanged(int identifier_index,
                                               const Identifier& identifier,
                                               milotic_hft::Status status,
                                               absl::string_view status_message,
                                               Payload& payload) {
  // NO_CDC: stream_state_ allocated parallel to params_.identifiers().
  IdentifierStreamState& state = stream_state_[identifier_index];
  if (status == state.last_signaled_status) {
    // Same status as last conveyed. Stay edge-triggered unless a heartbeat is
    // configured and a full heartbeat interval of export windows has elapsed,
    // in which case re-emit the marker to assert liveness.
    if (state.status_heartbeat_windows <= 0 ||
        ++state.windows_since_status_signal < state.status_heartbeat_windows) {
      return;
    }
  }
  state.MarkStatusSignaled(status);
  HighFrequencySensorsReadings* sensor_readings =
      payload.mutable_high_frequency_sensors_readings_batch()
          ->add_high_frequency_sensors();
  *sensor_readings->mutable_sensor_identifier() =
      identifier.sensor_identifier();
  sensor_readings->mutable_state()->set_status(status);
  if (!status_message.empty()) {
    sensor_readings->mutable_state()->set_status_message(status_message);
  }
}

absl::Time ActiveSubscription::ProcessSensorIdentifier(
    const Identifier& identifier, int identifier_index,
    Payload& identifier_payload) {
  IdentifierStreamState& state = stream_state_[identifier_index];
  const absl::Time last_sampled_time = absl::FromUnixNanos(
      // NO_CDC: last_sampled_time_ns_ allocated parallel to identifiers().
      last_sampled_time_ns_[identifier_index].load(std::memory_order_relaxed));
  absl::StatusOr<Payload> data =
      data_source_->Collect(identifier, last_sampled_time);

  // Collection failed outright. NotFound means the sensor is gone from the
  // store (MISSING); any other error is a transient/adapter problem with a
  // sensor that is still present (STALE). Edge-triggered so a persistently
  // failing sensor does not repeat the marker every cycle; the watermark is
  // left where it was.
  if (!data.ok()) {
    LOG_EVERY_N_SEC(ERROR, 10) << "Failed to collect data for subscription "
                               << id_ << ": " << data.status();
    const milotic_hft::Status status = absl::IsNotFound(data.status())
                                           ? milotic_hft::STATUS_MISSING
                                           : milotic_hft::STATUS_STALE;
    SignalStatusIfChanged(
        identifier_index, identifier, status,
        absl::StrCat("Failed to collect data for identifier ", identifier,
                     " with status: ", data.status()),
        identifier_payload);
    return last_sampled_time;
  }

  // Collection succeeded; the payload carries the sensor's own status (OK while
  // healthy, STALE/MISSING when the sensor reports trouble), set by the store
  // whether or not readings are present.
  const milotic_hft::State carried_state = GetCarriedState(*data);

  if (IsDataPopulated(*data)) {
    // Always decimate to the subscription's (power-of-two) interval, regardless
    // of the source poll rate. The data source's configured interval is not
    // under our control -- it may be a non-power-of-two default and may shift
    // as other subscribers come and go -- so resampling unconditionally keeps a
    // subscriber's output deterministic: exactly one sample per absolute bin.
    // stream_state_.last_emitted_bin carries the highest bin emitted so far
    // across cycles, so a bin straddling a batch boundary is not re-emitted.
    ResampleReadingsInplace(*data, params_.sampling_interval_ms(),
                            state.last_emitted_bin);

    // The watermark is the last emitted reading's own timestamp -- an observed
    // value, never parked ahead of real data. Cross-cycle bin dedup is handled
    // by last_emitted_bin above, not by rounding the watermark up to the bin
    // end.
    absl::StatusOr<absl::Time> updated_sample_time = GetLastSampledTime(*data);
    if (updated_sample_time.ok()) {
      // Recovery: if the sensor was last signaled non-OK and is now OK, reset
      // the edge filter so the first recovered reading is emitted even if its
      // value equals the one last emitted before the sensor went unhealthy --
      // otherwise the on-change filter would suppress it, leaving the consumer
      // with a status marker followed by silence.
      const bool recovered =
          carried_state.status() == milotic_hft::STATUS_OK &&
          state.last_signaled_status != milotic_hft::STATUS_OK &&
          state.last_signaled_status != milotic_hft::STATUS_UNKNOWN;
      if (recovered && state.sample_filter != nullptr) {
        state.sample_filter->Reset();
      }

      // Apply this identifier's edge filter (e.g. on-change) to the decimated
      // stream. Readings dropped here have already advanced the watermark
      // above, so they are not reconsidered next cycle. A null filter means
      // no/unknown edge filter was configured -- the stream passes through
      // unfiltered.
      if (state.sample_filter != nullptr &&
          data->has_high_frequency_sensors_readings_batch()) {
        for (auto& sensor :
             *data->mutable_high_frequency_sensors_readings_batch()
                  ->mutable_high_frequency_sensors()) {
          state.sample_filter->Filter(sensor);
        }
      }

      // Fresh, newly-binned readings were available this cycle. The status
      // rides along on the payload, so record what the subscriber is about to
      // see and restart the status heartbeat clock.
      state.MarkStatusSignaled(carried_state.status());
      identifier_payload = std::move(*data);
      return *updated_sample_time;
    }
    // else: had raw readings but all fell in already-emitted bins (source
    // polled faster than the sampling interval, or non-monotonic input).
    // Nothing survived decimation -- fall through to the stall handling below.
  }

  // Nothing survived decimation this export window (no readings at all, or
  // every reading fell in an already-emitted bin). Because the export interval
  // is at least twice the sampling interval (enforced in AddSubscription), a
  // sensor that is actually producing always advances at least one bin per
  // window, so an empty window means the source is not advancing -- a stall --
  // and is signaled STALE. If the sensor already reports a more specific non-OK
  // status (e.g. MISSING because it is gone from the store), surface that
  // instead. Edge-triggered: the marker is emitted once on the status
  // transition, so a persistent stall does not repeat it every window, and a
  // spurious one-window gap self-heals to OK on the next surviving sample. The
  // (edge-filtered-to- empty) case never reaches here: it returns above with
  // the sensor healthy.
  const milotic_hft::Status carried = carried_state.status();
  const bool reported_unhealthy = carried != milotic_hft::STATUS_OK &&
                                  carried != milotic_hft::STATUS_UNKNOWN;
  const milotic_hft::Status status =
      reported_unhealthy ? carried : milotic_hft::STATUS_STALE;
  const std::string status_message =
      reported_unhealthy ? carried_state.status_message()
                         : "No sample produced this export window.";
  SignalStatusIfChanged(identifier_index, identifier, status, status_message,
                        identifier_payload);
  return last_sampled_time;
}

absl::Time ActiveSubscription::ProcessFruIdentifier(
    const Identifier& identifier, int identifier_index,
    Payload& consolidated_data) {
  // NO_CDC: last_sampled_time_ns_ allocated parallel to params_.identifiers().
  const absl::Time last_sampled_time = absl::FromUnixNanos(
      last_sampled_time_ns_[identifier_index].load(std::memory_order_relaxed));
  absl::StatusOr<Payload> data =
      data_source_->Collect(identifier, last_sampled_time);
  if (!data.ok() || !data->has_fru_batch() ||
      data->fru_batch().frus().empty()) {
    LOG_EVERY_N_SEC(ERROR, 10) << "Failed to collect data for subscription "
                               << id_ << ": " << data.status();
    return last_sampled_time;
  }

  consolidated_data.MergeFrom(*data);
  return absl::Now();
}

std::string ActiveSubscription::GetDataSourceName() const {
  return data_source_->GetName();
}

std::vector<int64_t> ActiveSubscription::GetLastSampledTimesNs() const {
  std::vector<int64_t> times;
  times.reserve(last_sampled_time_ns_.size());
  for (const std::atomic<int64_t>& ns : last_sampled_time_ns_) {
    times.push_back(ns.load(std::memory_order_relaxed));
  }
  return times;
}

absl::Status SubscriptionManagerImpl::AddToAllResources(
    ActiveSubscription* subscription, SensorMutationBatch* mutation) {
  for (const auto& identifier : subscription->GetParams().identifiers()) {
    auto& monitor = resource_monitors_[identifier];
    if (monitor == nullptr) {
      monitor = std::make_unique<ResourceMonitor>(identifier);
    }
    if (absl::Status status = monitor->AddSubscriber(subscription, mutation);
        !status.ok()) {
      LOG(ERROR) << "Failed to add subscriber " << subscription->GetId()
                 << " to " << identifier << ": " << status;
      // Rolls back any partial registration; also cleans the brand-new
      // empty monitor we may have just inserted for this identifier.
      RemoveFromAllResources(subscription, nullptr);
      return status;
    }
  }
  return absl::OkStatus();
}

void SubscriptionManagerImpl::RemoveFromAllResources(
    ActiveSubscription* subscription, SensorMutationBatch* mutation) {
  for (const auto& identifier : subscription->GetParams().identifiers()) {
    auto it = resource_monitors_.find(identifier);
    if (it == resource_monitors_.end()) {
      // Idempotent: already removed (e.g. concurrent unsubscribe +
      // completion, or a brand-new empty monitor erased on a prior pass).
      continue;
    }
    absl::Status remove_status =
        it->second->RemoveSubscriber(subscription, mutation);
    if (!remove_status.ok() && !absl::IsInvalidArgument(remove_status)) {
      LOG(ERROR) << "Failed to remove subscriber " << subscription->GetId()
                 << " from resource " << identifier << ": " << remove_status;
    }
    if (!it->second->HasSubscribers()) {
      resource_monitors_.erase(it);
    }
  }
}

absl::StatusOr<std::shared_ptr<SubscriptionManager::Subscription>>
SubscriptionManagerImpl::AddSubscription(
    const milotic_hft::SubscriptionParams& params,
    absl::AnyInvocable<void(Payload&&)> on_data_callback) {
  LOG(INFO) << "Add Subscription with params: " << params;
  if (on_data_callback == nullptr) {
    return absl::InvalidArgumentError("on_data_callback cannot be null.");
  }

  if (!params.subscription_policy().has_resource_type()) {
    return absl::InvalidArgumentError(
        "resource_type must be provided in the subscription policy.");
  }

  ECCLESIA_ASSIGN_OR_RETURN(
      std::vector<Identifier> identifiers,
      data_source_->GetIdentifiersForResourceType(
          {params.subscription_policy().resource_type()}));

  // Copy the subscription params to append the identifiers to subscribe to.
  milotic_hft::SubscriptionParams internal_params = params;
  if (params.subscription_policy().configuration_type() ==
      milotic_hft::SubscriptionPolicy::
          CONFIGURATION_TYPE_CONFIGURE_ALL_RESOURCES) {
    internal_params.mutable_identifiers()->Add(identifiers.begin(),
                                               identifiers.end());
  }

  if (internal_params.identifiers().empty()) {
    return absl::InvalidArgumentError(
        "Subscription parameters must include at least one identifier or "
        "specify CONFIGURATION_TYPE_ALL_RESOURCES.");
  }

  // Reject duplicate identifiers. The same identifier listed twice would
  // register this subscription against one ResourceMonitor more than once. The
  // monitor counts the subscriber once (a set) but tracks the sampling interval
  // and batch size in multisets, so each duplicate inserts a surplus entry that
  // is never reclaimed: removal reconciles a single entry per subscriber, so
  // the extras leak and pin the resource at a phantom (faster) rate after the
  // owner unsubscribes. The duplicate would also be collected and delivered
  // twice per batch.
  {
    absl::flat_hash_set<Identifier, IdentifierHash, IdentifierEqual> seen;
    for (const Identifier& identifier : internal_params.identifiers()) {
      if (!seen.insert(identifier).second) {
        return absl::InvalidArgumentError(
            "Subscription contains duplicate identifiers.");
      }
    }
  }

  // Validate sensor subscription intervals (FRU subscriptions are full
  // snapshots re-sampled on export and are exempt; no subscription mixes
  // sensors and FRUs). Sensor sampling must be a power of two so the lowest
  // requested period divides every other -- resampling reduces to power-of-two
  // decimation and a coarser subscriber's samples are a strict subset of a
  // finer one's. The export interval must be at least twice the sampling
  // interval: the source polls and the subscription exports on independent
  // schedulers with jitter, so 2x ensures a producing sensor clears decimation
  // with at least one sample every export window (a window with none then
  // reliably signals a stall; see ProcessSensorIdentifier). These are rejected
  // rather than silently adjusted so a caller's request is never quietly
  // changed underneath it.
  if (internal_params.subscription_policy().resource_type() ==
      milotic_hft::SubscriptionPolicy::RESOURCE_TYPE_SENSOR) {
    const int32_t sampling_ms = internal_params.sampling_interval_ms();
    if (!IsValidSensorSamplingIntervalMs(sampling_ms)) {
      return absl::InvalidArgumentError(absl::Substitute(
          "Sensor sampling_interval_ms must be a power of two in [1, $0] ms; "
          "got $1.",
          kMaxSensorSamplingIntervalMs, sampling_ms));
    }
    if (internal_params.export_interval_ms() < 2 * sampling_ms) {
      return absl::InvalidArgumentError(absl::Substitute(
          "Sensor export_interval_ms ($0) must be at least twice "
          "sampling_interval_ms ($1).",
          internal_params.export_interval_ms(), sampling_ms));
    }
  }

  // Build per-identifier edge filters before constructing the subscription so
  // malformed configs (e.g. negative float_epsilon) surface as errors here.
  const int export_interval_ms = internal_params.export_interval_ms();
  std::vector<std::unique_ptr<SampleFilter>> filters;
  filters.reserve(internal_params.identifiers_size());
  for (const auto& identifier : internal_params.identifiers()) {
    ECCLESIA_ASSIGN_OR_RETURN(
        std::unique_ptr<SampleFilter> filter,
        MakeSampleFilter(identifier.edge_filter(), export_interval_ms));
    filters.push_back(std::move(filter));
  }

  SubscriptionId sub_id = GenerateSubscriptionId();
  auto active_sub = std::make_shared<ActiveSubscription>(
      sub_id, internal_params, std::move(on_data_callback), this,
      task_scheduler_.get(), data_source_.get(), std::move(filters));

  {
    absl::MutexLock lock(mutex_);
    std::unique_ptr<SensorMutationBatch> mutation =
        data_source_->CreateSensorMutationBatch(
            params.subscription_policy().resource_type(),
            std::move(identifiers));
    ECCLESIA_RETURN_IF_ERROR(
        AddToAllResources(active_sub.get(), mutation.get()));
    if (auto status = std::move(*mutation).Apply(); !status.ok()) {
      LOG(ERROR) << "Failed to apply batch sensor mutation: " << status;
      RemoveFromAllResources(active_sub.get(), nullptr);
      return status;
    }
  }
  active_sub->Begin();
  return active_sub;
}
absl::Status SubscriptionManagerImpl::Unsubscribe(
    const std::shared_ptr<Subscription>& subscription) {
  auto active_sub = std::dynamic_pointer_cast<ActiveSubscription>(subscription);
  if (active_sub == nullptr) {
    return absl::InvalidArgumentError(
        "subscription handle was not produced by this manager");
  }
  absl::MutexLock lock(mutex_);
  LOG(INFO) << "Removing subscription " << active_sub->GetId();
  absl::StatusOr<std::unique_ptr<SensorMutationBatch>> mutation =
      data_source_->CreateSensorMutationBatch(
          active_sub->GetParams().subscription_policy().resource_type());
  SensorMutationBatch* mutation_ptr = nullptr;
  if (!mutation.ok()) {
    LOG(ERROR) << "While removing subscription, failed to create batch sensor "
                  "mutation: "
               << mutation.status();
  } else {
    mutation_ptr = mutation->get();
  }
  RemoveFromAllResources(active_sub.get(), mutation_ptr);
  if (mutation.ok()) {
    if (absl::Status apply_status = std::move(*mutation.value()).Apply();
        !apply_status.ok()) {
      LOG(ERROR) << "While removing subscription, failed to apply batch sensor "
                    "mutation: "
                 << apply_status;
    }
  }
  return absl::OkStatus();
}

SubscriptionManagerImpl::SubscriptionId
SubscriptionManagerImpl::GenerateSubscriptionId() {
  return "sub_" + std::to_string(next_subscription_id_++);
}

SubscriptionManagerImpl::DebugSnapshot
SubscriptionManagerImpl::GetDebugSnapshot() const {
  DebugSnapshot snapshot;
  snapshot.next_subscription_id = next_subscription_id_.load();
  snapshot.data_source_name = data_source_->GetName();
  absl::MutexLock lock(mutex_);
  snapshot.resource_monitors.reserve(resource_monitors_.size());
  for (const auto& [identifier, resource_monitor] : resource_monitors_) {
    snapshot.resource_monitors.push_back(resource_monitor->GetDebugSnapshot());
  }
  return snapshot;
}

}  // namespace milotic_hft