[core] SetTaskStatus should only be called within the same lock scope where task_entry is retrieved

[kevin85421]

Why are these changes needed?

This PR reverts #52695 and adds comments to explain where should SetTaskStatus be called.

#52695 updates the value in submissble_tasks_ without acquiring the mutex lock. If multiple threads or coroutines write to the map, a rehash or deletion may occur, causing the pointer to the value to become invalid.

Outdated PR statement

Question

#52695 (comment)

Pointers to values in a flat_hash_map become invalid after a rehash. Additionally, we dereference those pointers in RetryTask, which doesn’t hold a mutex lock. Hence, it’s possible for the pointers to become invalid when other coroutines or threads insert or delete elements from the map, triggering a rehash.

"Iterators, references, and pointers to elements are invalidated on rehash." (reference)

Solution

Changing submissible_tasks_ from absl::flat_hash_map<TaskID, TaskEntry> to absl::flat_hash_map<TaskID, std::unique_ptr<TaskEntry>> requires a lot of changes.

Hence, this PR implements a short-term solution by copying the value (i.e., TaskEntry) while holding the mutex lock where rehash will not be triggered by other threads / coroutine.

Related issue number

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[kevin85421]

@dentiny does this solution make sense to you? Thanks!

[dentiny]

Or even simplier, you could use std::unordered_map, which guarantees pointer stability.

Internally, flat hash map / swiss table uses open addressing to resolve collision,
while unordered map leverages hash chaining, so each ValueEntry is heap allocated.

[kevin85421]

this is the only change. Others are for reverting.

[kevin85421]

@dentiny I found that the original solution (copy it) doesn't work because we need to update the task_entry. Hence, I decide to revert the old solution for now until I have time to implement a better solution.