Repository: ray-project/ray
File: src/ray/raylet/scheduling/local_resource_manager.h (C/C++)

Code before:
                                 bool record_idle_resource = true);

  void UpdateAvailableObjectStoreMemResource();

  void SetResourceIdle(const scheduling::ResourceID &resource_id);

  void SetResourceNonIdle(const scheduling::ResourceID &resource_id);

  std::optional<absl::Time> GetResourceIdleTime() const;

  /// Get the draining deadline if node is in draining state.
  ///
  /// \return The draining deadline if node is in draining state, otherwise -1.
  int64_t GetDrainingDeadline() const {
    return drain_request_.has_value() ? drain_request_->deadline_timestamp_ms() : -1;
  }
  /// Identifier of local node.
  scheduling::NodeID local_node_id_;
  /// Resources of local node.
  NodeResourceInstances local_resources_;

  /// A map storing when the resource was last idle.
  absl::flat_hash_map<WorkArtifact, std::optional<absl::Time>> last_idle_times_;
  /// Saved idle times for footprints marked busy via MaybeMarkFootprintAsBusy().
  /// Used to restore idle time when the speculative busy state is cleared.
  absl::flat_hash_map<WorkFootprint, absl::Time> saved_footprint_idle_times_;
  /// Function to get current time. Defaults to absl::Now() if not provided.
  std::function<absl::Time()> now_fn_;
  /// Function to get used object store memory.
  std::function<int64_t(void)> get_used_object_store_memory_;
  /// Function to get whether the pull manager is at capacity.
  std::function<bool(void)> get_pull_manager_at_capacity_;
  /// Function to shutdown the raylet gracefully.
  std::function<void(const rpc::NodeDeathInfo &)> shutdown_raylet_gracefully_;

  /// Subscribes to resource changes.
  std::function<void(const NodeResources &)> resource_change_subscriber_;

  // Version of this resource. It will incr by one whenever the state changed.
  int64_t version_ = 0;

  /// The draining request this node received.
  std::optional<rpc::DrainRayletRequest> drain_request_;

  ray::observability::MetricInterface &resource_usage_gauge_;

  FRIEND_TEST(ClusterResourceSchedulerTest, SchedulingUpdateTotalResourcesTest);
  FRIEND_TEST(ClusterResourceSchedulerTest, AvailableResourceInstancesOpsTest);
  FRIEND_TEST(ClusterResourceSchedulerTest, TaskResourceInstancesTest);
  FRIEND_TEST(ClusterResourceSchedulerTest, TaskResourceInstancesAllocationFailureTest);
  FRIEND_TEST(ClusterResourceSchedulerTest, TaskResourceInstancesTest2);

Review comment (suggestion):
Is it possible that we merge the 2 maps? I think this should save us some lookup costs.

```c
struct FootprintIdleTimes {
  std::optional<absl::Time> idle_ts;
  std::optional<absl::Time> saved_ts;
};

absl::flat_hash_map<WorkArtifact, IdleTimes> last_idle_times_;
```

Code after:
  /// added back to the node's local available resources.
  ///
  /// \param task_allocation: Task's resources to be freed.
  /// \param record_idle_resource: Whether to record the idle resource. This is false
  ///   when the resource was allocated partially so its idle state is actually not
  ///   affected.
  void FreeTaskResourceInstances(std::shared_ptr<TaskResourceInstances> task_allocation,
                                 bool record_idle_resource = true);

  void UpdateAvailableObjectStoreMemResource();

  void SetResourceIdle(const scheduling::ResourceID &resource_id);

  void SetResourceNonIdle(const scheduling::ResourceID &resource_id);

  std::optional<absl::Time> GetResourceIdleTime() const;

  /// Get the draining deadline if node is in draining state.
  ///
  /// \return The draining deadline if node is in draining state, otherwise -1.
  int64_t GetDrainingDeadline() const {
    return drain_request_.has_value() ? drain_request_->deadline_timestamp_ms() : -1;
  }
  /// Identifier of local node.
  scheduling::NodeID local_node_id_;
  /// Resources of local node.
  NodeResourceInstances local_resources_;

  /// A map storing idle state for resources and work footprints.
  /// Each entry tracks the current idle time (or nullopt if busy) and optionally
  /// a saved time for speculative busy marking (used only for WorkFootprints).
  absl::flat_hash_map<WorkArtifact, IdleTimeState> idle_time_states_;
  /// Function to get current time. Defaults to absl::Now() if not provided.
  std::function<absl::Time()> now_fn_;
  /// Function to get used object store memory.
  std::function<int64_t(void)> get_used_object_store_memory_;
  /// Function to get whether the pull manager is at capacity.
  std::function<bool(void)> get_pull_manager_at_capacity_;
  /// Function to shutdown the raylet gracefully.
  std::function<void(const rpc::NodeDeathInfo &)> shutdown_raylet_gracefully_;

  /// Subscribes to resource changes.
  std::function<void(const NodeResources &)> resource_change_subscriber_;

  // Version of this resource. It will incr by one whenever the state changed.
  int64_t version_ = 0;

  /// The draining request this node received.
  std::optional<rpc::DrainRayletRequest> drain_request_;

  ray::observability::MetricInterface &resource_usage_gauge_;