state_updater.h

#ifndef PRODUCTION_BORG_MGMT_NODE_PROXY_SAFEPOWER_SAFEPOWER_AGENT_STATE_UPDATER_H_
#define PRODUCTION_BORG_MGMT_NODE_PROXY_SAFEPOWER_SAFEPOWER_AGENT_STATE_UPDATER_H_

#include <algorithm>
#include <array>
#include <cstddef>
#include <list>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "callback_manager.h"
#include "state_merge.h"
#include "absl/base/nullability.h"
#include "absl/base/thread_annotations.h"
#include "absl/functional/any_invocable.h"
#include "absl/functional/function_ref.h"
#include "absl/log/log.h"
#include "absl/synchronization/mutex.h"

namespace safepower_agent {
template <typename StateProto>
class StateUpdater {
 public:
  class Listener;
  using ListenerIterator = typename std::list<Listener* >::iterator;

  class Listener {
    struct StateUpdaterRef {
      std::weak_ptr<StateUpdater> ptr;
      StateUpdater* raw_ptr;  // For compares in destructor.
      ListenerIterator itr;
    };

   public:
    explicit Listener(bool is_passive_listener = false)
        : is_passive_listener_(is_passive_listener) {}

    virtual ~Listener() {
      if (!state_updater_.empty()) {
        LOG(DFATAL) << "Listener destroyed without calling StopListening";
      }
      StopListening();
    }

    // Derived classes must explicitly call this at the start of their
    // destructor if they access any member variables or resources during
    // callbacks, since those resources will be destroyed before the
    // base ~Listener() runs.
    void StopListening() {
      for (const StateUpdaterRef& ref : state_updater_) {
        std::shared_ptr<StateUpdater> state_updater = ref.ptr.lock();
        if (state_updater != nullptr) {
          state_updater->RemoveListener(ref.itr);
        }
      }
      state_updater_.clear();
    }

    template <typename... Args>
    void Listen(const Args&... args) {
      static_assert(sizeof...(args) > 0);
      if (!state_updater_.empty()) {
        LOG(DFATAL) << "Listener::Listen called multiple times";
        return;
      }
      state_updater_.reserve(sizeof...(args));
      std::array done = {Insert(args)...};
      for (size_t i = 0; i < done.size(); ++i) {
        if (done[i]) {
          MaybeDone(*state_updater_[i].raw_ptr);
        }
      }
    }

    bool IsPassiveListener() const { return is_passive_listener_; }
    virtual void UpdateState([[maybe_unused]] const StateProto& previous_state,
                             [[maybe_unused]] const StateProto& update) {
      LOG(DFATAL) << "UpdateState called on base Listener";
    }

    void MaybeDone(const StateUpdater& updater) {
      for (StateUpdaterRef& ref : state_updater_) {
        if (ref.raw_ptr == &updater) {
          ref.raw_ptr = nullptr;
          ++done_count_;
          if (done_count_ >= state_updater_.size()) {
            Done();
          }
          return;
        }
      }
      LOG(DFATAL) << "Listener not listening to this updater";
    }

   protected:
    virtual void Done() { LOG(DFATAL) << "Done called on base Listener"; }

   private:
    bool Insert(const  std::shared_ptr<StateUpdater>& updater) {
      for (const StateUpdaterRef& ref : state_updater_) {
        if (ref.ptr.lock() == updater) {
          LOG(DFATAL) << "Listener already listening to this updater";
          return true;
        }
      }
      bool final = updater->AddListener(
          *this, [&](ListenerIterator itr) {
            state_updater_.push_back({updater, updater.get(), itr});
          });
      return final;
    }

    std::vector<StateUpdaterRef> state_updater_;
    size_t done_count_ = 0;
    bool is_passive_listener_ = false;
  };
  friend class Listener;

  explicit StateUpdater(StateProto initial_state = {}, bool final_state = false)
      : state_(std::move(initial_state)), final_(final_state) {}

  ~StateUpdater() {
    if (!listeners_mutex_.try_lock()) {
      LOG(DFATAL) << "StateUpdater destroyed while mutex was held";
      return;
    }
    for (Listener* listener : listeners_) {
      listener->MaybeDone(*this);
    }
    listeners_.clear();
    listeners_mutex_.unlock();
  }

  // Returns a copy of the state. If you want to access the state without
  // copying (e.g. for performance reasons with large states), prefer using
  // WithState.
  StateProto state() const {
    absl::MutexLock lock(state_mutex_);
    return state_;
  }

  template <typename Func>
  auto WithState(Func&& f) const ABSL_LOCKS_EXCLUDED(state_mutex_) {
    absl::MutexLock lock(state_mutex_);
    return std::forward<Func>(f)(state_);
  }

  std::string DebugString() const ABSL_LOCKS_EXCLUDED(state_mutex_) {
    return WithState(
        [](const StateProto& state) { return state.DebugString(); });
  }

  void UpdateState(const StateProto& update, bool final_state = false)
      ABSL_LOCKS_EXCLUDED(state_mutex_, listeners_mutex_) {
    absl::MutexLock lock_state(state_mutex_);
    final_ = final_state;
    absl::MutexLock lock_listeners(listeners_mutex_);
    for (auto listener : listeners_) {
      listener->UpdateState(state_, update);
      if (final_state) {
        listener->MaybeDone(*this);
      }
    }
    ApplyStateUpdate(state_, update);
  }

  // Runs a callback when the updater has no more listeners. The callback will
  // be called immediately if the updater is already idle. The move-only
  // returned handle can be destroyed to cancel the callback.
  // Idle callbacks are run in LIFO order, while active callbacks are run in
  // FIFO order. This reduces the likelihood of deadlocks: For example, if a
  // resource is acquired when active and released when idle, another pair of
  // callbacks that are added later will always see that resource acquired.
  CallbackManager::Handle OnIdle(absl::AnyInvocable<void() &&> callback)
      ABSL_LOCKS_EXCLUDED(listeners_mutex_) {
    if (IsIdle()) {
      std::move(callback)();
      return {};
    }
    return activity_callbacks_.RunFirst(std::move(callback));
  }

  // Runs a callback when the updater has at least one listener. The callback
  // will be called immediately if the updater already has listeners. See OnIdle
  // for more details.
  CallbackManager::Handle OnActive(absl::AnyInvocable<void() &&> callback)
      ABSL_LOCKS_EXCLUDED(listeners_mutex_) {
    if (!IsIdle()) {
      std::move(callback)();
      return {};
    }
    return activity_callbacks_.RunLast(std::move(callback));
  }

 private:
  bool IsIdleInternal() const ABSL_EXCLUSIVE_LOCKS_REQUIRED(listeners_mutex_) {
    for (const auto* listener : listeners_) {
      if (!listener->IsPassiveListener()) {
        return false;
      }
    }
    return true;
  }

  bool IsIdle() const ABSL_LOCKS_EXCLUDED(listeners_mutex_) {
    absl::MutexLock lock(listeners_mutex_);
    return IsIdleInternal();
  }

  bool AddListener(Listener& listener,
                   absl::FunctionRef<void(ListenerIterator)> register_func)
      ABSL_LOCKS_EXCLUDED(state_mutex_, listeners_mutex_) {
    bool run_callbacks;
    bool result;
    {
      absl::MutexLock lock_state(state_mutex_);
      listener.UpdateState({}, state_);

      absl::MutexLock lock_listeners(listeners_mutex_);
      // IsIdleInternal() must be checked before push_front modifies the list
      // to correctly evaluate the pre-transition idle state.
      run_callbacks = IsIdleInternal() && !listener.IsPassiveListener();
      // push_front must happen before register_func so the iterator is valid.
      listeners_.push_front(&listener);
      register_func(listeners_.begin());
      result = final_;
    }
    if (run_callbacks) {
      activity_callbacks_.RunCallbacks();
    }
    return result;
  }

  void RemoveListener(ListenerIterator handle)
      ABSL_LOCKS_EXCLUDED(listeners_mutex_) {
    bool becomes_idle = false;
    {
      absl::MutexLock lock(listeners_mutex_);
      if (handle == listeners_.end()) return;
      bool was_idle = IsIdleInternal();
      listeners_.erase(handle);
      becomes_idle = !was_idle && IsIdleInternal();
    }
    if (becomes_idle) {
      activity_callbacks_.RunCallbacks();
    }
  }

  mutable absl::Mutex state_mutex_ ABSL_ACQUIRED_BEFORE(listeners_mutex_);
  mutable absl::Mutex listeners_mutex_;
  // This is a std::list so that we have stable iterators
  std::list<Listener* > listeners_
      ABSL_GUARDED_BY(listeners_mutex_);
  StateProto state_ ABSL_GUARDED_BY(state_mutex_);
  bool final_ ABSL_GUARDED_BY(state_mutex_);
  // Callbacks that are run when the updater becomes idle (no listeners) or
  // active (at least one listener).
  CallbackManager activity_callbacks_;
};

}  // namespace safepower_agent

#endif  // PRODUCTION_BORG_MGMT_NODE_PROXY_SAFEPOWER_SAFEPOWER_AGENT_STATE_UPDATER_H_