From 1e59e63e064552de0c5f9f2625efbf430c252b2b Mon Sep 17 00:00:00 2001 From: Shanicky Chen Date: Wed, 17 Jul 2024 21:32:34 +0800 Subject: [PATCH] Refactor ActorMapping iteration and worker lookup in mapping.rs --- src/common/src/hash/consistent_hash/mapping.rs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/common/src/hash/consistent_hash/mapping.rs b/src/common/src/hash/consistent_hash/mapping.rs index f95d89bad0cbd..6029c080afddb 100644 --- a/src/common/src/hash/consistent_hash/mapping.rs +++ b/src/common/src/hash/consistent_hash/mapping.rs @@ -332,7 +332,12 @@ impl ActorMapping { /// Transform the actor mapping to the worker slot mapping. Note that the parameter is a mapping from actor to worker. pub fn to_worker_slot(&self, actor_to_worker: &HashMap) -> WorkerSlotMapping { let mut worker_actors = HashMap::new(); - for (actor_id, worker_id) in actor_to_worker { + for actor_id in self.iter_unique() { + let worker_id = actor_to_worker + .get(&actor_id) + .cloned() + .expect(&format!("location for actor {} not found", actor_id)); + worker_actors .entry(worker_id) .or_insert(BTreeSet::new()) @@ -342,7 +347,7 @@ impl ActorMapping { let mut actor_location = HashMap::new(); for (worker, actors) in worker_actors { for (idx, &actor) in actors.iter().enumerate() { - actor_location.insert(*actor, WorkerSlotId::new(*worker, idx)); + actor_location.insert(actor, WorkerSlotId::new(worker, idx)); } }