action_context.h

#ifndef PRODUCTION_BORG_MGMT_NODE_PROXY_SAFEPOWER_SAFEPOWER_AGENT_ACTION_CONTEXT_H_
#define PRODUCTION_BORG_MGMT_NODE_PROXY_SAFEPOWER_SAFEPOWER_AGENT_ACTION_CONTEXT_H_

#include <cstdint>
#include <memory>
#include <optional>
#include <ostream>
#include <string>
#include <utility>
#include <vector>

#include "action_hash.h"
#include "condition.h"
#include "flight_record.h"
#include "action.pb.h"
#include "safepower_agent.pb.h"
#include "state_persistence.pb.h"
#include "state_updater.h"
#include "utils.pb.h"
#include "absl/base/nullability.h"
#include "absl/base/thread_annotations.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/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/mutex.h"
#include "absl/time/time.h"
#include "google/protobuf/repeated_field.h"
#include "google/protobuf/repeated_ptr_field.h"

namespace safepower_agent {

class ActionContext;

// An active listener that does nothing on state updates, but serves purely
// to keep the local StateUpdater active (non-idle) during action execution.
// This prevents background state monitors from going idle
// and pausing while an action is in flight.
class KeepAliveListener
    : public StateUpdater<safepower_agent_proto::SystemState>::Listener {
 public:
  KeepAliveListener()
      : StateUpdater<safepower_agent_proto::SystemState>::Listener(
            /*is_passive_listener=*/false) {}
  void UpdateState(const safepower_agent_proto::SystemState&,
                   const safepower_agent_proto::SystemState&) override {}

 protected:
  void Done() override {}
};

// Used to start and manage actions. Handles retrieving persistent state and
// scheduling new actions.
class ActionContextManager {
 public:
  using FlightRecordRequest =
      ::borg_mgmt::node_proxy::safepower::utils::FlightRecordRequest;

  // Action is a function that is called when the action is ready to be
  // executed.
  // The first argument is the proto representation of the action.
  // The second argument is a callback function that is called when the action
  // is finished. The callback takes a single argument: the status of the
  // action.
  // It is worth noting that actions are non-blocking. To allow for this, the
  // callback function is passed by value. This allows the caller to invoke the
  // callback at a later time.
  using Action =
      absl::AnyInvocable<void(const safepower_agent_proto::Action&,
                              absl::AnyInvocable<void(absl::Status) &&>)>;

  explicit ActionContextManager(
      std::shared_ptr<StateUpdater<safepower_agent_proto::SystemState>>
          state_updater,
      std::shared_ptr<StateUpdater<safepower_agent_proto::SystemState>>
          global_state_updater,
      bool read_only_mode = false)
      : state_updater_(std::move(state_updater)),
        global_state_updater_(std::move(global_state_updater)),
        read_only_mode_(read_only_mode),
        forwarder_(global_state_updater_) {
    forwarder_.Listen(state_updater_);
  }
  ~ActionContextManager() { forwarder_.StopListening(); }

  ActionContextManager(ActionContextManager&& other) = delete;
  ActionContextManager& operator=(ActionContextManager&& other) = delete;

  absl::Status RegisterAction(const safepower_agent_proto::Action& action,
                              Action action_impl)
      ABSL_LOCKS_EXCLUDED(actions_mutex_);

  absl::Status LoadSavedActions();

  absl::StatusOr<ActionContext* > StartAction(
      safepower_agent_proto::StartActionRequest request)
      ABSL_LOCKS_EXCLUDED(actions_mutex_);
  void FinishAction(const safepower_agent_proto::Action& action,
                    Action action_impl) ABSL_LOCKS_EXCLUDED(actions_mutex_);

  std::vector<FlightRecordWrapper> ListRunningActions() const
      ABSL_LOCKS_EXCLUDED(actions_mutex_);

  // Dumps debug information about the running actions to the provided stream.
  void DumpDebugInfo(std::ostream& os) const
      ABSL_LOCKS_EXCLUDED(actions_mutex_);

  ActionContext*  GetActionContext(
      const FlightRecordRequest& flight_record)
      ABSL_LOCKS_EXCLUDED(actions_mutex_);

  bool IsActionReloading(const FlightRecordRequest& flight_record) const
      ABSL_LOCKS_EXCLUDED(actions_mutex_);

  void GetSupportedActions(
      google::protobuf::RepeatedPtrField<safepower_agent_proto::Action>& actions) const
      ABSL_LOCKS_EXCLUDED(actions_mutex_);

  const 
  std::shared_ptr<StateUpdater<safepower_agent_proto::SystemState>>&
  system_state_updater() const {
    return state_updater_;
  }
  bool read_only_mode() const { return read_only_mode_; }

  const 
  std::shared_ptr<StateUpdater<safepower_agent_proto::SystemState>>&
  global_system_state_updater() const {
    return global_state_updater_;
  }

  // This is just for use by mutex annotations. Since it return const, it
  // can't be used to acquire the mutex.
  const absl::Mutex& actions_mutex() const { return actions_mutex_; }

  absl::Status validateFlightRecord(const FlightRecordRequest& record);

 private:
  absl::Status ValidateRequest(
      const safepower_agent_proto::StartActionRequest& request,
      absl::Time start_time);

  absl::StatusOr<Action> ReserveAction(
      const safepower_agent_proto::Action& action)
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(actions_mutex_);

  absl::StatusOr< std::unique_ptr<ActionContext>>
  ReinitActionContext(
      safepower_agent_persistence_proto::SavedAction saved_action)
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(actions_mutex_);

  absl::Status ReloadAndInsertAction(
      safepower_agent_persistence_proto::SavedActionRecord saved_action_record,
      std::unique_ptr<ActionContext>& context_to_cleanup)
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(actions_mutex_);

  std::shared_ptr<StateUpdater<safepower_agent_proto::SystemState>>
      state_updater_;
  std::shared_ptr<StateUpdater<safepower_agent_proto::SystemState>>
      global_state_updater_;
  bool read_only_mode_;

  // Forwards updates from `state_updater_` to `global_state_updater_`.
  class Forwarder
      : public StateUpdater<safepower_agent_proto::SystemState>::Listener {
   public:
    explicit Forwarder(
        std::shared_ptr<StateUpdater<safepower_agent_proto::SystemState>>
            global_state_updater)
        : StateUpdater<safepower_agent_proto::SystemState>::Listener(
              /*is_passive_listener=*/true),
          global_state_updater_(std::move(global_state_updater)) {}

    void UpdateState(
        const safepower_agent_proto::SystemState& /*previous_state*/,
        const safepower_agent_proto::SystemState& update) override {
      global_state_updater_->UpdateState(update);
    }

   protected:
    void Done() override {}

   private:
    std::shared_ptr<StateUpdater<safepower_agent_proto::SystemState>>
        global_state_updater_;
  };
  Forwarder forwarder_;

  mutable absl::Mutex actions_mutex_;
  absl::flat_hash_map<safepower_agent_proto::Action, Action, ActionHash,
                      ActionEq>
      actions_ ABSL_GUARDED_BY(actions_mutex_);
  // ActionContext owns the key.

  absl::flat_hash_map<FlightRecordWrapper, std::unique_ptr<ActionContext>>
      running_actions_ ABSL_GUARDED_BY(actions_mutex_);

  absl::flat_hash_set<FlightRecordWrapper> reloading_actions_
      ABSL_GUARDED_BY(actions_mutex_);
};

// Represents a single action execution.
class ActionContext {
  constexpr static int kMaxStateChangeRetryCount = 4;

 public:
  // Only allow ActionContext to be created by ActionContextManager.
  struct CreationToken {
   private:
    CreationToken() = default;
    friend class ActionContextManager;
  };

  using Action = ActionContextManager::Action;

  explicit ActionContext(
      CreationToken token, ActionContextManager& manager,
      safepower_agent_proto::StartActionRequest request,
      Action action_impl = {},
      safepower_agent_proto::ActionStateLog initial_state = NewInitialState());
  ~ActionContext();

  ActionContext(const ActionContext& other) = delete;
  ActionContext& operator=(const ActionContext& other) = delete;

  // Dumps debug information about this action to the provided stream.
  void DumpDebugInfo(std::ostream& os) const ABSL_LOCKS_EXCLUDED(mutex_);

  const FlightRecordWrapper& flight_record() const { return flight_record_; }

  const safepower_agent_proto::StartActionRequest& request() const {
    return request_;
  }

  
  std::shared_ptr<StateUpdater<safepower_agent_proto::ActionStateLog>>
  action_state_updater() const {
    return action_state_updater_;
  }

  absl::Status Activate();

  static bool IsFinalState(safepower_agent_proto::ActionState state) {
    return state == safepower_agent_proto::ACTION_STATE_SUCCESS ||
           state == safepower_agent_proto::ACTION_STATE_ERROR ||
           state == safepower_agent_proto::ACTION_STATE_ABORTED;
  }

 private:
  std::optional<absl::Status> CheckValidation(
      const safepower_agent_proto::SystemState& state);

  absl::Status EnterStateInit() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
  absl::Status EnterStateRunningAction() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
  absl::Status EnterStateCheckingPrecondition()
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
  absl::Status EnterStateValidatingFinalState()
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);

  void NextStateInit() ABSL_LOCKS_EXCLUDED(mutex_);
  void NextStatePreconditionMatched(absl::Status status,
                                    Condition::MatchList matches)
      ABSL_LOCKS_EXCLUDED(mutex_);
  void NextStateActionRan(absl::Status status) ABSL_LOCKS_EXCLUDED(mutex_);
  void NextStateValidationCompleted(absl::Status status,
                                    Condition::MatchList matches)
      ABSL_LOCKS_EXCLUDED(mutex_);

  void RunAction() ABSL_LOCKS_EXCLUDED(mutex_);

  absl::Status StartCheckingCondition(
      Condition& condition, absl::string_view condition_id,
      absl::AnyInvocable<void(absl::Status, Condition::MatchList)> callback)
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);

  std::string execution_task_name() const {
    // The execution task name must be unique for each action.
    // As the flight name and step can be any characters, we add a hash to
    // ensure that the name is unique.
    uint64_t hash = absl::HashOf(flight_record());
    return absl::StrCat(flight_record().flight_name(), "-",
                        flight_record().step_id(), "-", hash, ".execution");
  }

  std::string init_task_name() const {
    return absl::StrCat(execution_task_name(), ".init");
  }

  std::string response_task_name() const {
    return absl::StrCat(execution_task_name(), ".response");
  }

  std::string final_state_task_name() const {
    return absl::StrCat(execution_task_name(), ".final");
  }

  void SetState(safepower_agent_proto::ActionState new_state,
                safepower_agent_proto::ActionStateChange change_info = {})
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
  void Finish(Action action_impl);

  mutable absl::Mutex mutex_ ABSL_ACQUIRED_AFTER(manager_.actions_mutex());
  ActionContextManager& manager_;
  const safepower_agent_proto::StartActionRequest request_;
  const FlightRecordWrapper flight_record_;
  Action action_impl_ ABSL_GUARDED_BY(mutex_);
  std::optional<Condition> precondition_;
  std::optional<Condition> validation_;
  std::unique_ptr<KeepAliveListener> keep_alive_listener_
      ABSL_GUARDED_BY(mutex_);
  
  std::shared_ptr<StateUpdater<safepower_agent_proto::ActionStateLog>>
      action_state_updater_;
  static safepower_agent_proto::ActionStateLog NewInitialState();
};

}  // namespace safepower_agent

#endif  // PRODUCTION_BORG_MGMT_NODE_PROXY_SAFEPOWER_SAFEPOWER_AGENT_ACTION_CONTEXT_H_