tlbmc/hft/core/manager_impl.h

#ifndef THIRD_PARTY_MILOTIC_EXTERNAL_CC_TLBMC_HFT_CORE_MANAGER_IMPL_H_
#define THIRD_PARTY_MILOTIC_EXTERNAL_CC_TLBMC_HFT_CORE_MANAGER_IMPL_H_

#include <atomic>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <limits>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "absl/base/thread_annotations.h"
#include "absl/container/btree_set.h"
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/functional/any_invocable.h"
#include "absl/hash/hash.h"
#include "absl/log/log.h"
#include "absl/memory/memory.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/mutex.h"
#include "absl/time/time.h"
#include "tlbmc/adapter/data_source.h"
#include "identifier.pb.h"
#include "payload.pb.h"
#include "sensor_identifier.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"

namespace milotic_hft {

class SubscriptionManagerImpl;

struct IdentifierHash {
  size_t operator()(const Identifier& v) const {
    return v.identifier_case() == Identifier::kSensorIdentifier
               ? absl::HashOf(1, v.sensor_identifier().name())
               : absl::HashOf(2, v.fru_identifier().barepath());
  }
};

struct IdentifierEqual {
  bool operator()(const Identifier& lhs, const Identifier& rhs) const {
    if (lhs.identifier_case() != rhs.identifier_case()) {
      return false;
    }
    return lhs.identifier_case() == Identifier::kSensorIdentifier
               ? lhs.sensor_identifier().name() ==
                     rhs.sensor_identifier().name()
               : lhs.fru_identifier().barepath() ==
                     rhs.fru_identifier().barepath();
  }
};

class ActiveSubscription
    : public SubscriptionManager::Subscription,
      public std::enable_shared_from_this<ActiveSubscription> {
 public:
  using SubscriptionId = std::string;

  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);

  ~ActiveSubscription() override {
    if (task_scheduler_ != nullptr) {
      task_scheduler_->Cancel(task_id_);
    }
    LOG(WARNING) << "ActiveSubscription " << id_ << " destroyed.";
  }

  void Begin();

  SubscriptionId GetId() const { return id_; }

  const milotic_hft::SubscriptionParams& GetParams() const { return params_; }

  void DeliverData(Payload&& data);

  bool IsActive() const { return batches_remaining_ != 0; }

  // Debug accessors, consumed by the tlbmc-side JSON rendering (the core
  // deliberately carries no serialization dependency).
  int BatchesRemaining() const { return batches_remaining_; }
  std::string GetDataSourceName() const;
  // Last sampled time (Unix nanos) per identifier, parallel to
  // GetParams().identifiers().
  std::vector<int64_t> GetLastSampledTimesNs() const;

 private:
  // Collects sensor data for the given identifier and updates the
  // `identifier_payload`, returns the last sampled time in this collection.
  absl::Time ProcessSensorIdentifier(const milotic_hft::Identifier& identifier,
                                     int identifier_index,
                                     Payload& identifier_payload);

  // Collects fru data for the given identifier and updates the
  // `identifier_payload`, returns the last sampled time in this collection.
  absl::Time ProcessFruIdentifier(const milotic_hft::Identifier& identifier,
                                  int identifier_index,
                                  Payload& consolidated_data);

  // Appends a status-only sensor entry (no readings) for `identifier` to
  // `payload`. Edge-triggered: a marker is emitted when `status` differs from
  // the status last signaled for this identifier, so a persistently unhealthy
  // sensor does not repeat its marker every export cycle. If a heartbeat is
  // configured for the identifier (OnChangeFilter.heartbeat_interval_ms), an
  // unchanged non-OK status is additionally re-emitted once every heartbeat
  // interval (in whole export windows) to assert liveness. Records the new
  // status. Used for the no-fresh-sample paths where the sensor reports
  // STALE/MISSING.
  void SignalStatusIfChanged(int identifier_index,
                             const milotic_hft::Identifier& identifier,
                             milotic_hft::Status status,
                             absl::string_view status_message,
                             Payload& payload);

  int task_id_ = -1;  // -1 means no task is scheduled.
  const SubscriptionId id_;
  const milotic_hft::SubscriptionParams params_;
  absl::AnyInvocable<void(Payload&&)> on_data_callback_;
  // Back-pointer used to call Unsubscribe when the subscription drains its
  // batch budget. Non-owning: the manager outlives all subscriptions it
  // produced.
  SubscriptionManagerImpl* manager_;
  int batches_remaining_ = -1;  // -1 means unlimited.
  // Last sampled time (Unix nanos) per identifier, parallel to
  // params_.identifiers(). Written by the scheduled task (single writer);
  // read by the debug accessors from arbitrary threads. Relaxed atomics
  // suffice — readers tolerate a stale snapshot.
  std::vector<std::atomic<int64_t>> last_sampled_time_ns_;
  // Per-identifier stream state, parallel to params_.identifiers(). Touched
  // ONLY by the scheduled task (single-threaded; no atomics needed). Bundling
  // these task-private fields keeps them distinct from last_sampled_time_ns_
  // above, which is shared with debug readers and therefore atomic.
  struct IdentifierStreamState {
    // Highest absolute bin index emitted so far. The resampler emits a reading
    // only when its bin is strictly greater, which dedups across collect cycles
    // (without parking the watermark ahead of real data) and makes decimation
    // tolerant of non-monotonic input. INT64_MIN so the first reading emits.
    int64_t last_emitted_bin = std::numeric_limits<int64_t>::min();
    // Status last signaled to the subscriber. A non-OK status (STALE/MISSING)
    // is emitted only on a transition (edge-triggered), so a persistently
    // unhealthy sensor does not repeat its marker every cycle. STATUS_UNKNOWN
    // initial so the first observed non-OK status is treated as a transition
    // and emitted.
    milotic_hft::Status last_signaled_status = milotic_hft::STATUS_UNKNOWN;
    // Heartbeat cadence for re-emitting an unchanged non-OK status marker,
    // expressed in whole export windows; 0 disables it, leaving status purely
    // edge-triggered. Taken from the sample filter's HeartbeatWindows() at
    // construction, so the heartbeat window is defined in exactly one place.
    int status_heartbeat_windows = 0;
    // Export windows elapsed since the current status was last conveyed to the
    // subscriber. Drives the status heartbeat above; reset whenever the status
    // is (re-)signaled or rides along on delivered readings.
    int windows_since_status_signal = 0;
    // Edge filter (e.g. on-change) for this identifier, or null if none/unknown
    // was configured. Built once at construction; carries state across cycles
    // and is reset on recovery from a non-OK status.
    std::unique_ptr<SampleFilter> sample_filter;

    // Records that `status` was just conveyed to the subscriber (either as a
    // status-only marker or riding along on delivered readings) and restarts
    // the status heartbeat clock.
    void MarkStatusSignaled(milotic_hft::Status status) {
      last_signaled_status = status;
      windows_since_status_signal = 0;
    }
  };
  std::vector<IdentifierStreamState> stream_state_;
  Scheduler* task_scheduler_;
  DataSource* data_source_;
};

class SubscriptionManagerImpl : public SubscriptionManager {
 public:
  using SubscriptionId = ActiveSubscription::SubscriptionId;

  ~SubscriptionManagerImpl() override;

  // Creates a subscription manager. `scheduler` runs the per-subscription
  // export tasks; it is injected so the core stays independent of any
  // concrete scheduler implementation (and its dependencies).
  static std::unique_ptr<SubscriptionManagerImpl> Create(
      std::unique_ptr<DataSource> data_source,
      std::unique_ptr<Scheduler> scheduler) {
    return absl::WrapUnique(new SubscriptionManagerImpl(std::move(data_source),
                                                        std::move(scheduler)));
  }

  absl::StatusOr<std::shared_ptr<Subscription>> AddSubscription(
      const SubscriptionParams& params,
      absl::AnyInvocable<void(Payload&&)> on_data_callback) override
      ABSL_LOCKS_EXCLUDED(mutex_);

  absl::Status Unsubscribe(const std::shared_ptr<Subscription>& subscription)
      override ABSL_LOCKS_EXCLUDED(mutex_);

  // Manages the sampling interval of a resource.
  // Tracks the effective sampling interval of a resource as the minimum
  // over all subscribers and updates it as subscribers are added or removed.
  //
  // All accesses are serialized by the enclosing SubscriptionManagerImpl's
  // mutex_.
  class ResourceMonitor {
   public:
    explicit ResourceMonitor(const milotic_hft::Identifier& identifier)
        : identifier_(identifier),
          current_sampling_interval_ms_(0),
          current_max_batch_size_(0) {}

    const milotic_hft::Identifier& GetIdentifier() const { return identifier_; }

    absl::Status AddSubscriber(ActiveSubscription* subscription,
                               SensorMutationBatch* mutation);

    absl::Status RemoveSubscriber(ActiveSubscription* subscription,
                                  SensorMutationBatch* mutation);

    absl::Status ConfigureDataSource(SensorMutationBatch* mutation);

    int GetCurrentSamplingIntervalMs() const {
      return current_sampling_interval_ms_;
    }

    bool HasSubscribers() const { return !subscribers_.empty(); }

    // Point-in-time copy of the monitor's state for debug rendering.
    struct DebugSnapshot {
      Identifier identifier;
      int current_sampling_interval_ms = 0;
      int current_max_batch_size = 0;
      std::vector<SubscriptionId> subscriber_ids;
      std::vector<int> subscribed_sampling_intervals_ms;
      std::vector<int> subscribed_max_batch_sizes;
    };
    DebugSnapshot GetDebugSnapshot() const;

   private:
    void RemoveSubscriberInternal(ActiveSubscription* subscription);

    const milotic_hft::Identifier identifier_;
    int current_sampling_interval_ms_;
    int current_max_batch_size_;
    // Non-owning. Subscriptions are owned by their gRPC reactor (or whoever
    // called AddSubscription). Erasure happens under the enclosing
    // SubscriptionManagerImpl's mutex_, atomic with the caller's drop of the
    // owning shared_ptr in Unsubscribe paths.
    absl::flat_hash_set<ActiveSubscription*> subscribers_;
    absl::btree_multiset<int> lowest_sampling_interval_subscriptions_;
    absl::btree_multiset<int, std::greater<>> max_batch_sizes_;
  };

  // Point-in-time debug snapshot of the manager's state, consumed by the
  // tlbmc-side JSON rendering (the core deliberately carries no serialization
  // dependency).
  struct DebugSnapshot {
    size_t next_subscription_id = 0;
    std::string data_source_name;
    std::vector<ResourceMonitor::DebugSnapshot> resource_monitors;
  };
  DebugSnapshot GetDebugSnapshot() const ABSL_LOCKS_EXCLUDED(mutex_);

  // Debug-only access to the injected scheduler, so a consumer that knows the
  // concrete type can render its stats alongside the manager's.
  const Scheduler& GetScheduler() const { return *task_scheduler_; }

 private:
  SubscriptionManagerImpl(std::unique_ptr<DataSource> data_source,
                          std::unique_ptr<Scheduler> scheduler);

  SubscriptionId GenerateSubscriptionId();

  // Walks subscription's identifiers and registers it with each
  // resource_monitors_ entry, creating new entries as needed. Threads
  // `mutation` through to the data source. On failure, rolls back any
  // partial registration via RemoveFromAllResources before returning.
  absl::Status AddToAllResources(ActiveSubscription* subscription,
                                 SensorMutationBatch* mutation)
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);

  // Walks subscription's identifiers and unregisters it from each
  // resource_monitors_ entry that contains it; erases entries that
  // become empty. Idempotent: missing entries and "subscriber not
  // present" both no-op. `mutation` may be nullptr (in-memory only) for
  // post-failure rollback, where the data source must not be touched.
  void RemoveFromAllResources(ActiveSubscription* subscription,
                              SensorMutationBatch* mutation)
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);

  std::atomic<size_t> next_subscription_id_;
  std::unique_ptr<Scheduler> task_scheduler_;
  std::unique_ptr<DataSource> data_source_;

  mutable absl::Mutex mutex_;
  absl::flat_hash_map<milotic_hft::Identifier, std::unique_ptr<ResourceMonitor>,
                      IdentifierHash, IdentifierEqual>
      resource_monitors_ ABSL_GUARDED_BY(mutex_);
};

}  // namespace milotic_hft

#endif  // THIRD_PARTY_MILOTIC_EXTERNAL_CC_TLBMC_HFT_CORE_MANAGER_IMPL_H_