Repository: tokio-rs/tokio
File: tokio-util/src/task/join_map.rs (Rust)

Code before:
    /// [`LocalSet`]: tokio::task::LocalSet
    /// [`join_next`]: Self::join_next
    #[track_caller]
    pub fn spawn_local_on<F>(&mut self, key: K, task: F, local_set: &LocalSet)
    where
        F: Future<Output = V>,
        F: 'static,
    {
        let task = self.tasks.spawn_local_on(task, local_set);
        self.insert(key, task)
    }

    fn insert(&mut self, key: K, abort: AbortHandle) {
        let hash_builder = self.hashes_by_task.hasher();
        let hash = hash_one(hash_builder, &key);
        let id = abort.id();

        // Insert the new key into the map of tasks by keys.
        let entry =
            self.tasks_by_key
                .entry(hash, |(k, _)| *k == key, |(k, _)| hash_one(hash_builder, k));
        match entry {
            Entry::Occupied(occ) => {
                // There was a previous task spawned with the same key! Cancel
                // that task, and remove its ID from the map of hashes by task IDs.
                let (_, abort) = std::mem::replace(occ.into_mut(), (key, abort));
                abort.abort();

                let _prev_hash = self.hashes_by_task.remove(&abort.id());
                debug_assert_eq!(Some(hash), _prev_hash);
            }
            Entry::Vacant(vac) => {
                vac.insert((key, abort));
            }
        };

        // Associate the key's hash with this task's ID, for looking up tasks by ID.
        let _prev = self.hashes_by_task.insert(id, hash);
        debug_assert!(_prev.is_none(), "no prior task should have had the same ID");
    }

    /// Waits until one of the tasks in the map completes and returns its
    /// output, along with the key corresponding to that task.
    ///
    /// Returns `None` if the map is empty.
    ///
    /// # Cancel Safety
    ///
    /// This method is cancel safe. If `join_next` is used as the event in a [`tokio::select!`]
    /// statement and some other branch completes first, it is guaranteed that no tasks were
    /// removed from this `JoinMap`.

Review comment (suggestion):
This drops the key. If the destructor of the key panics, we don't call `abort` or adjust `hashes_by_task`.

Code after:
    /// [`LocalSet`]: tokio::task::LocalSet
    /// [`join_next`]: Self::join_next
    #[track_caller]
    pub fn spawn_local_on<F>(&mut self, key: K, task: F, local_set: &LocalSet)
    where
        F: Future<Output = V>,
        F: 'static,
    {
        let task = self.tasks.spawn_local_on(task, local_set);
        self.insert(key, task)
    }

    fn insert(&mut self, mut key: K, mut abort: AbortHandle) {
        let hash_builder = self.hashes_by_task.hasher();
        let hash = hash_one(hash_builder, &key);
        let id = abort.id();

        // Insert the new key into the map of tasks by keys.
        let entry =
            self.tasks_by_key
                .entry(hash, |(k, _)| *k == key, |(k, _)| hash_one(hash_builder, k));
        match entry {
            Entry::Occupied(occ) => {
                // There was a previous task spawned with the same key! Cancel
                // that task, and remove its ID from the map of hashes by task IDs.
                (key, abort) = std::mem::replace(occ.into_mut(), (key, abort));

                // Remove the old task ID.
                let _prev_hash = self.hashes_by_task.remove(&abort.id());
                debug_assert_eq!(Some(hash), _prev_hash);

                // Associate the key's hash with the new task's ID, for looking up tasks by ID.
                let _prev = self.hashes_by_task.insert(id, hash);
                debug_assert!(_prev.is_none(), "no prior task should have had the same ID");

                // Note: it's important to drop `key` and abort the task here.
                // This defends against any panics during drop handling for causing inconsistent state.
                abort.abort();
                drop(key);
            }
            Entry::Vacant(vac) => {
                vac.insert((key, abort));

                // Associate the key's hash with this task's ID, for looking up tasks by ID.
                let _prev = self.hashes_by_task.insert(id, hash);
                debug_assert!(_prev.is_none(), "no prior task should have had the same ID");
            }
        };
    }

    /// Waits until one of the tasks in the map completes and returns its