Repository: tokio-rs/tokio
File: tokio/src/runtime/task/trace/mod.rs (Rust)

Code before:
        while let Some(notified) = injection.pop(&mut synced.inject) {
            drop(notified);
        }

        drop(synced);

        // precondition: we have drained the tasks from the local and injection
        // queues.
        trace_owned(owned)
    }
}

/// Trace the `OwnedTasks`.
///
/// # Preconditions
///
/// This helper presumes exclusive access to each task. The tasks must not exist
/// in any other queue.
fn trace_owned<S>(owned: &OwnedTasks<Arc<S>>) -> Vec<Trace> {
    // notify each task
    let mut tasks = vec![];
    owned.for_each(|task| {
        // set the notified bit
        task.as_raw().state().transition_to_notified_for_tracing();
        // store the tasks into a vec
        tasks.push(task.clone());
        // do NOT poll `task` here, since we hold a lock on `owned` and the task
        // may complete and need to remove itself from `owned`.
    });

    tasks
        .into_iter()
        .map(|task| {
            let ((), trace) = Trace::capture(|| task.as_raw().poll());
            trace
        })
        .collect()
}


Review comment (suggestion):
Okay, so, I've figured out how I think we should handle this refcount business:

 * First, we wrap `transition_to_notified_for_tracing` with a method on `Task` that returns a `Notified`. This `Notified` will own the refcount that we created inside of `transition_to_notified_for_tracing`.
 * Next, when you want to poll the task, instead of calling `poll` directly, you instead use [`OwnedTasks::assert_owner`][1] to convert the `Notified` into a `LocalNotified`.
 * To poll the task, you call [`LocalNotified::run`](https://github.com/tokio-rs/tokio/blob/eaba9712e8a4b35a1a2cdeb3b27845d6239c12ea/tokio/src/runtime/task/mod.rs#L395-L402).

This way, we only increment the refcount once (so there's no call to `Task::clone`), and all refcounts we create are owned by `Notified` object, so their destructor will clear it if it gets dropped before it is polled, e.g. due to a panic.

[1]: https://github.com/tokio-rs/tokio/blob/eaba9712e8a4b35a1a2cdeb3b27845d6239c12ea/tokio/src/runtime/task/list.rs#L127-L139

Code after:
        while let Some(notified) = injection.pop(&mut synced.inject) {
            drop(notified);
        }

        drop(synced);

        // precondition: we have drained the tasks from the local and injection
        // queues.
        trace_owned(owned)
    }
}

/// Trace the `OwnedTasks`.
///
/// # Preconditions
///
/// This helper presumes exclusive access to each task. The tasks must not exist
/// in any other queue.
fn trace_owned<S: Schedule>(owned: &OwnedTasks<S>) -> Vec<Trace> {
    // notify each task
    let mut tasks = vec![];
    owned.for_each(|task| {
        // notify the task (and thus make it poll-able) and stash it
        tasks.push(task.notify_for_tracing());
        // we do not poll it here since we hold a lock on `owned` and the task
        // may complete and need to remove itself from `owned`.
    });

    tasks
        .into_iter()
        .map(|task| {
            let local_notified = owned.assert_owner(task);
            let ((), trace) = Trace::capture(|| local_notified.run());
            trace
        })
        .collect()
}