[red-knot] Refactor program.check scheduling (#11202)
This commit is contained in:
Micha Reiser
@@ -16,7 +16,7 @@ use tracing_tree::time::Uptime;
|
||||
use red_knot::db::{HasJar, ParallelDatabase, QueryError, SemanticDb, SourceDb, SourceJar};
|
||||
use red_knot::files::FileId;
|
||||
use red_knot::module::{ModuleSearchPath, ModuleSearchPathKind};
|
||||
use red_knot::program::check::RayonCheckScheduler;
|
||||
use red_knot::program::check::ExecutionMode;
|
||||
use red_knot::program::{FileChange, FileChangeKind, Program};
|
||||
use red_knot::watch::FileWatcher;
|
||||
use red_knot::Workspace;
|
||||
@@ -143,22 +143,16 @@ impl MainLoop {
|
||||
|
||||
// Spawn a new task that checks the program. This needs to be done in a separate thread
|
||||
// to prevent blocking the main loop here.
|
||||
rayon::spawn(move || {
|
||||
rayon::in_place_scope(|scope| {
|
||||
let scheduler = RayonCheckScheduler::new(&program, scope);
|
||||
|
||||
match program.check(&scheduler) {
|
||||
Ok(result) => {
|
||||
sender
|
||||
.send(OrchestratorMessage::CheckProgramCompleted {
|
||||
diagnostics: result,
|
||||
revision,
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
Err(QueryError::Cancelled) => {}
|
||||
}
|
||||
});
|
||||
rayon::spawn(move || match program.check(ExecutionMode::ThreadPool) {
|
||||
Ok(result) => {
|
||||
sender
|
||||
.send(OrchestratorMessage::CheckProgramCompleted {
|
||||
diagnostics: result,
|
||||
revision,
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
Err(QueryError::Cancelled) => {}
|
||||
});
|
||||
}
|
||||
MainLoopMessage::ApplyChanges(changes) => {
|
||||
|
||||
@@ -1,5 +1,3 @@
|
||||
use std::num::NonZeroUsize;
|
||||
|
||||
use rayon::{current_num_threads, yield_local};
|
||||
use rustc_hash::FxHashSet;
|
||||
|
||||
@@ -12,30 +10,21 @@ use crate::symbols::Dependency;
|
||||
impl Program {
|
||||
/// Checks all open files in the workspace and its dependencies.
|
||||
#[tracing::instrument(level = "debug", skip_all)]
|
||||
pub fn check(&self, scheduler: &dyn CheckScheduler) -> QueryResult<Vec<String>> {
|
||||
pub fn check(&self, mode: ExecutionMode) -> QueryResult<Vec<String>> {
|
||||
self.cancelled()?;
|
||||
|
||||
let check_loop = CheckFilesLoop::new(scheduler);
|
||||
let mut context = CheckContext::new(self);
|
||||
|
||||
check_loop.run(self.workspace().open_files.iter().copied())
|
||||
}
|
||||
match mode {
|
||||
ExecutionMode::SingleThreaded => SingleThreadedExecutor.run(&mut context)?,
|
||||
ExecutionMode::ThreadPool => ThreadPoolExecutor.run(&mut context)?,
|
||||
};
|
||||
|
||||
/// Checks a single file and its dependencies.
|
||||
#[tracing::instrument(level = "debug", skip(self, scheduler))]
|
||||
pub fn check_file(
|
||||
&self,
|
||||
file: FileId,
|
||||
scheduler: &dyn CheckScheduler,
|
||||
) -> QueryResult<Vec<String>> {
|
||||
self.cancelled()?;
|
||||
|
||||
let check_loop = CheckFilesLoop::new(scheduler);
|
||||
|
||||
check_loop.run([file].into_iter())
|
||||
Ok(context.finish())
|
||||
}
|
||||
|
||||
#[tracing::instrument(level = "debug", skip(self, context))]
|
||||
fn do_check_file(&self, file: FileId, context: &CheckContext) -> QueryResult<Diagnostics> {
|
||||
fn check_file(&self, file: FileId, context: &CheckFileContext) -> QueryResult<Diagnostics> {
|
||||
self.cancelled()?;
|
||||
|
||||
let symbol_table = self.symbol_table(file)?;
|
||||
@@ -63,7 +52,7 @@ impl Program {
|
||||
// Supporting non-first-party code also requires supporting typing stubs.
|
||||
if let Some(dependency) = self.resolve_module(dependency_name)? {
|
||||
if dependency.path(self)?.root().kind().is_first_party() {
|
||||
context.schedule_check_file(dependency.path(self)?.file());
|
||||
context.schedule_dependency(dependency.path(self)?.file());
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -81,212 +70,343 @@ impl Program {
|
||||
}
|
||||
}
|
||||
|
||||
/// Schedules checks for files.
|
||||
pub trait CheckScheduler {
|
||||
/// Schedules a check for a file.
|
||||
///
|
||||
/// The check can either be run immediately on the current thread or the check can be queued
|
||||
/// in a thread pool and ran asynchronously.
|
||||
///
|
||||
/// The order in which scheduled checks are executed is not guaranteed.
|
||||
///
|
||||
/// The implementation should call [`CheckFileTask::run`] to execute the check.
|
||||
fn check_file(&self, file_task: CheckFileTask);
|
||||
|
||||
/// The maximum number of checks that can be run concurrently.
|
||||
///
|
||||
/// Returns `None` if the checks run on the current thread (no concurrency).
|
||||
fn max_concurrency(&self) -> Option<NonZeroUsize>;
|
||||
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
|
||||
pub enum ExecutionMode {
|
||||
SingleThreaded,
|
||||
ThreadPool,
|
||||
}
|
||||
|
||||
/// Scheduler that runs checks on a rayon thread pool.
|
||||
pub struct RayonCheckScheduler<'program, 'scope_ref, 'scope> {
|
||||
program: &'program Program,
|
||||
scope: &'scope_ref rayon::Scope<'scope>,
|
||||
/// Context that stores state information about the entire check operation.
|
||||
struct CheckContext<'a> {
|
||||
/// IDs of the files that have been queued for checking.
|
||||
///
|
||||
/// Used to avoid queuing the same file twice.
|
||||
scheduled_files: FxHashSet<FileId>,
|
||||
|
||||
/// Reference to the program that is checked.
|
||||
program: &'a Program,
|
||||
|
||||
/// The aggregated diagnostics
|
||||
diagnostics: Vec<String>,
|
||||
}
|
||||
|
||||
impl<'program, 'scope_ref, 'scope> RayonCheckScheduler<'program, 'scope_ref, 'scope> {
|
||||
pub fn new(program: &'program Program, scope: &'scope_ref rayon::Scope<'scope>) -> Self {
|
||||
Self { program, scope }
|
||||
impl<'a> CheckContext<'a> {
|
||||
fn new(program: &'a Program) -> Self {
|
||||
Self {
|
||||
scheduled_files: FxHashSet::default(),
|
||||
program,
|
||||
diagnostics: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the tasks to check all open files in the workspace.
|
||||
fn check_open_files(&mut self) -> Vec<CheckOpenFileTask> {
|
||||
self.scheduled_files
|
||||
.extend(self.program.workspace().open_files());
|
||||
|
||||
self.program
|
||||
.workspace()
|
||||
.open_files()
|
||||
.map(|file_id| CheckOpenFileTask { file_id })
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Returns the task to check a dependency.
|
||||
fn check_dependency(&mut self, file_id: FileId) -> Option<CheckDependencyTask> {
|
||||
if self.scheduled_files.insert(file_id) {
|
||||
Some(CheckDependencyTask { file_id })
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Pushes the result for a single file check operation
|
||||
fn push_diagnostics(&mut self, diagnostics: &Diagnostics) {
|
||||
self.diagnostics.extend_from_slice(diagnostics);
|
||||
}
|
||||
|
||||
/// Returns a reference to the program that is being checked.
|
||||
fn program(&self) -> &'a Program {
|
||||
self.program
|
||||
}
|
||||
|
||||
/// Creates a task context that is used to check a single file.
|
||||
fn task_context<'b, S>(&self, dependency_scheduler: &'b S) -> CheckTaskContext<'a, 'b, S>
|
||||
where
|
||||
S: ScheduleDependency,
|
||||
{
|
||||
CheckTaskContext {
|
||||
program: self.program,
|
||||
dependency_scheduler,
|
||||
}
|
||||
}
|
||||
|
||||
fn finish(self) -> Vec<String> {
|
||||
self.diagnostics
|
||||
}
|
||||
}
|
||||
|
||||
impl<'program, 'scope_ref, 'scope> CheckScheduler
|
||||
for RayonCheckScheduler<'program, 'scope_ref, 'scope>
|
||||
/// Trait that abstracts away how a dependency of a file gets scheduled for checking.
|
||||
trait ScheduleDependency {
|
||||
/// Schedules the file with the given ID for checking.
|
||||
fn schedule(&self, file_id: FileId);
|
||||
}
|
||||
|
||||
impl<T> ScheduleDependency for T
|
||||
where
|
||||
'program: 'scope,
|
||||
T: Fn(FileId),
|
||||
{
|
||||
fn check_file(&self, check_file_task: CheckFileTask) {
|
||||
let child_span =
|
||||
tracing::trace_span!("check_file", file_id = check_file_task.file_id.as_u32());
|
||||
let program = self.program;
|
||||
|
||||
self.scope
|
||||
.spawn(move |_| child_span.in_scope(|| check_file_task.run(program)));
|
||||
|
||||
if current_num_threads() == 1 {
|
||||
yield_local();
|
||||
}
|
||||
}
|
||||
|
||||
fn max_concurrency(&self) -> Option<NonZeroUsize> {
|
||||
if current_num_threads() == 1 {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(NonZeroUsize::new(current_num_threads()).unwrap_or(NonZeroUsize::MIN))
|
||||
fn schedule(&self, file_id: FileId) {
|
||||
let f = self;
|
||||
f(file_id);
|
||||
}
|
||||
}
|
||||
|
||||
/// Scheduler that runs all checks on the current thread.
|
||||
pub struct SameThreadCheckScheduler<'a> {
|
||||
/// Context that is used to run a single file check task.
|
||||
///
|
||||
/// The task is generic over `S` because it is passed across thread boundaries and
|
||||
/// we don't want to add the requirement that [`ScheduleDependency`] must be [`Send`].
|
||||
struct CheckTaskContext<'a, 'scheduler, S>
|
||||
where
|
||||
S: ScheduleDependency,
|
||||
{
|
||||
dependency_scheduler: &'scheduler S,
|
||||
program: &'a Program,
|
||||
}
|
||||
|
||||
impl<'a> SameThreadCheckScheduler<'a> {
|
||||
pub fn new(program: &'a Program) -> Self {
|
||||
Self { program }
|
||||
impl<'a, 'scheduler, S> CheckTaskContext<'a, 'scheduler, S>
|
||||
where
|
||||
S: ScheduleDependency,
|
||||
{
|
||||
fn as_file_context(&self) -> CheckFileContext<'scheduler> {
|
||||
CheckFileContext {
|
||||
dependency_scheduler: self.dependency_scheduler,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl CheckScheduler for SameThreadCheckScheduler<'_> {
|
||||
fn check_file(&self, task: CheckFileTask) {
|
||||
task.run(self.program);
|
||||
}
|
||||
/// Context passed when checking a single file.
|
||||
///
|
||||
/// This is a trimmed down version of [`CheckTaskContext`] with the type parameter `S` erased
|
||||
/// to avoid monomorphization of [`Program:check_file`].
|
||||
struct CheckFileContext<'a> {
|
||||
dependency_scheduler: &'a dyn ScheduleDependency,
|
||||
}
|
||||
|
||||
fn max_concurrency(&self) -> Option<NonZeroUsize> {
|
||||
None
|
||||
impl<'a> CheckFileContext<'a> {
|
||||
fn schedule_dependency(&self, file_id: FileId) {
|
||||
self.dependency_scheduler.schedule(file_id);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct CheckFileTask {
|
||||
file_id: FileId,
|
||||
context: CheckContext,
|
||||
enum CheckFileTask {
|
||||
OpenFile(CheckOpenFileTask),
|
||||
Dependency(CheckDependencyTask),
|
||||
}
|
||||
|
||||
impl CheckFileTask {
|
||||
/// Runs the check and communicates the result to the orchestrator.
|
||||
pub fn run(self, program: &Program) {
|
||||
match program.do_check_file(self.file_id, &self.context) {
|
||||
Ok(diagnostics) => self
|
||||
.context
|
||||
.sender
|
||||
.send(CheckFileMessage::Completed(diagnostics))
|
||||
.unwrap(),
|
||||
Err(QueryError::Cancelled) => self
|
||||
.context
|
||||
.sender
|
||||
.send(CheckFileMessage::Cancelled)
|
||||
.unwrap(),
|
||||
/// Runs the task and returns the results for checking this file.
|
||||
fn run<S>(&self, context: &CheckTaskContext<S>) -> QueryResult<Diagnostics>
|
||||
where
|
||||
S: ScheduleDependency,
|
||||
{
|
||||
match self {
|
||||
Self::OpenFile(task) => task.run(context),
|
||||
Self::Dependency(task) => task.run(context),
|
||||
}
|
||||
}
|
||||
|
||||
fn file_id(&self) -> FileId {
|
||||
match self {
|
||||
CheckFileTask::OpenFile(task) => task.file_id,
|
||||
CheckFileTask::Dependency(task) => task.file_id,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct CheckContext {
|
||||
sender: crossbeam::channel::Sender<CheckFileMessage>,
|
||||
/// Task to check an open file.
|
||||
|
||||
#[derive(Debug)]
|
||||
struct CheckOpenFileTask {
|
||||
file_id: FileId,
|
||||
}
|
||||
|
||||
impl CheckContext {
|
||||
fn new(sender: crossbeam::channel::Sender<CheckFileMessage>) -> Self {
|
||||
Self { sender }
|
||||
}
|
||||
|
||||
/// Queues a new file for checking using the [`CheckScheduler`].
|
||||
#[allow(unused)]
|
||||
fn schedule_check_file(&self, file_id: FileId) {
|
||||
self.sender.send(CheckFileMessage::Queue(file_id)).unwrap();
|
||||
impl CheckOpenFileTask {
|
||||
fn run<S>(&self, context: &CheckTaskContext<S>) -> QueryResult<Diagnostics>
|
||||
where
|
||||
S: ScheduleDependency,
|
||||
{
|
||||
context
|
||||
.program
|
||||
.check_file(self.file_id, &context.as_file_context())
|
||||
}
|
||||
}
|
||||
|
||||
struct CheckFilesLoop<'a> {
|
||||
scheduler: &'a dyn CheckScheduler,
|
||||
pending: usize,
|
||||
queued_files: FxHashSet<FileId>,
|
||||
/// Task to check a dependency file.
|
||||
#[derive(Debug)]
|
||||
struct CheckDependencyTask {
|
||||
file_id: FileId,
|
||||
}
|
||||
|
||||
impl<'a> CheckFilesLoop<'a> {
|
||||
fn new(scheduler: &'a dyn CheckScheduler) -> Self {
|
||||
Self {
|
||||
scheduler,
|
||||
queued_files: FxHashSet::default(),
|
||||
pending: 0,
|
||||
impl CheckDependencyTask {
|
||||
fn run<S>(&self, context: &CheckTaskContext<S>) -> QueryResult<Diagnostics>
|
||||
where
|
||||
S: ScheduleDependency,
|
||||
{
|
||||
context
|
||||
.program
|
||||
.check_file(self.file_id, &context.as_file_context())
|
||||
}
|
||||
}
|
||||
|
||||
/// Executor that schedules the checking of individual program files.
|
||||
trait CheckExecutor {
|
||||
fn run(self, context: &mut CheckContext) -> QueryResult<()>;
|
||||
}
|
||||
|
||||
/// Executor that runs all check operations on the current thread.
|
||||
///
|
||||
/// The executor does not schedule dependencies for checking.
|
||||
/// The main motivation for scheduling dependencies
|
||||
/// in a multithreaded environment is to parse and index the dependencies concurrently.
|
||||
/// However, that doesn't make sense in a single threaded environment, because the dependencies then compute
|
||||
/// with checking the open files. Checking dependencies in a single threaded environment is more likely
|
||||
/// to hurt performance because we end up analyzing files in their entirety, even if we only need to type check parts of them.
|
||||
#[derive(Debug, Default)]
|
||||
struct SingleThreadedExecutor;
|
||||
|
||||
impl CheckExecutor for SingleThreadedExecutor {
|
||||
fn run(self, context: &mut CheckContext) -> QueryResult<()> {
|
||||
let mut queue = context.check_open_files();
|
||||
|
||||
let noop_schedule_dependency = |_| {};
|
||||
|
||||
while let Some(file) = queue.pop() {
|
||||
context.program().cancelled()?;
|
||||
|
||||
let task_context = context.task_context(&noop_schedule_dependency);
|
||||
context.push_diagnostics(&file.run(&task_context)?);
|
||||
}
|
||||
}
|
||||
|
||||
fn run(mut self, files: impl Iterator<Item = FileId>) -> QueryResult<Vec<String>> {
|
||||
let (sender, receiver) = if let Some(max_concurrency) = self.scheduler.max_concurrency() {
|
||||
crossbeam::channel::bounded(max_concurrency.get())
|
||||
} else {
|
||||
// The checks run on the current thread. That means it is necessary to store all messages
|
||||
// or we risk deadlocking when the main loop never gets a chance to read the messages.
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Executor that runs the check operations on a thread pool.
|
||||
///
|
||||
/// The executor runs each check operation as its own task using a thread pool.
|
||||
///
|
||||
/// Other than [`SingleThreadedExecutor`], this executor schedules dependencies for checking. It
|
||||
/// even schedules dependencies for checking when the thread pool size is 1 for a better debugging experience.
|
||||
#[derive(Debug, Default)]
|
||||
struct ThreadPoolExecutor;
|
||||
|
||||
impl CheckExecutor for ThreadPoolExecutor {
|
||||
fn run(self, context: &mut CheckContext) -> QueryResult<()> {
|
||||
let num_threads = current_num_threads();
|
||||
let single_threaded = num_threads == 1;
|
||||
let span = tracing::trace_span!("ThreadPoolExecutor::run", num_threads);
|
||||
let _ = span.enter();
|
||||
|
||||
let mut queue: Vec<_> = context
|
||||
.check_open_files()
|
||||
.into_iter()
|
||||
.map(CheckFileTask::OpenFile)
|
||||
.collect();
|
||||
|
||||
let (sender, receiver) = if single_threaded {
|
||||
// Use an unbounded queue for single threaded execution to prevent deadlocks
|
||||
// when a single file schedules multiple dependencies.
|
||||
crossbeam::channel::unbounded()
|
||||
} else {
|
||||
// Use a bounded queue to apply backpressure when the orchestration thread isn't able to keep
|
||||
// up processing messages from the worker threads.
|
||||
crossbeam::channel::bounded(num_threads)
|
||||
};
|
||||
|
||||
let context = CheckContext::new(sender.clone());
|
||||
let schedule_sender = sender.clone();
|
||||
let schedule_dependency = move |file_id| {
|
||||
schedule_sender
|
||||
.send(ThreadPoolMessage::ScheduleDependency(file_id))
|
||||
.unwrap();
|
||||
};
|
||||
|
||||
for file in files {
|
||||
self.queue_file(file, context.clone());
|
||||
}
|
||||
let result = rayon::in_place_scope(|scope| {
|
||||
let mut pending = 0usize;
|
||||
|
||||
self.run_impl(receiver, &context)
|
||||
}
|
||||
loop {
|
||||
context.program().cancelled()?;
|
||||
|
||||
fn run_impl(
|
||||
mut self,
|
||||
receiver: crossbeam::channel::Receiver<CheckFileMessage>,
|
||||
context: &CheckContext,
|
||||
) -> QueryResult<Vec<String>> {
|
||||
let mut result = Vec::default();
|
||||
let mut cancelled = false;
|
||||
// 1. Try to get a queued message to ensure that we have always remaining space in the channel to prevent blocking the worker threads.
|
||||
// 2. Try to process a queued file
|
||||
// 3. If there's no queued file wait for the next incoming message.
|
||||
// 4. Exit if there are no more messages and no senders.
|
||||
let message = if let Ok(message) = receiver.try_recv() {
|
||||
message
|
||||
} else if let Some(task) = queue.pop() {
|
||||
pending += 1;
|
||||
|
||||
for message in receiver {
|
||||
match message {
|
||||
CheckFileMessage::Completed(diagnostics) => {
|
||||
result.extend_from_slice(&diagnostics);
|
||||
let task_context = context.task_context(&schedule_dependency);
|
||||
let sender = sender.clone();
|
||||
let task_span = tracing::trace_span!(
|
||||
parent: &span,
|
||||
"CheckFileTask::run",
|
||||
file_id = task.file_id().as_u32(),
|
||||
);
|
||||
|
||||
self.pending -= 1;
|
||||
scope.spawn(move |_| {
|
||||
task_span.in_scope(|| match task.run(&task_context) {
|
||||
Ok(result) => {
|
||||
sender.send(ThreadPoolMessage::Completed(result)).unwrap();
|
||||
}
|
||||
Err(err) => sender.send(ThreadPoolMessage::Errored(err)).unwrap(),
|
||||
});
|
||||
});
|
||||
|
||||
if self.pending == 0 {
|
||||
break;
|
||||
// If this is a single threaded rayon thread pool, yield the current thread
|
||||
// or we never start processing the work items.
|
||||
if single_threaded {
|
||||
yield_local();
|
||||
}
|
||||
}
|
||||
CheckFileMessage::Queue(id) => {
|
||||
if !cancelled {
|
||||
self.queue_file(id, context.clone());
|
||||
}
|
||||
}
|
||||
CheckFileMessage::Cancelled => {
|
||||
self.pending -= 1;
|
||||
cancelled = true;
|
||||
|
||||
if self.pending == 0 {
|
||||
break;
|
||||
continue;
|
||||
} else if let Ok(message) = receiver.recv() {
|
||||
message
|
||||
} else {
|
||||
break;
|
||||
};
|
||||
|
||||
match message {
|
||||
ThreadPoolMessage::ScheduleDependency(dependency) => {
|
||||
if let Some(task) = context.check_dependency(dependency) {
|
||||
queue.push(CheckFileTask::Dependency(task));
|
||||
}
|
||||
}
|
||||
ThreadPoolMessage::Completed(diagnostics) => {
|
||||
context.push_diagnostics(&diagnostics);
|
||||
pending -= 1;
|
||||
|
||||
if pending == 0 && queue.is_empty() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
ThreadPoolMessage::Errored(err) => {
|
||||
return Err(err);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if cancelled {
|
||||
Err(QueryError::Cancelled)
|
||||
} else {
|
||||
Ok(result)
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
});
|
||||
|
||||
fn queue_file(&mut self, file_id: FileId, context: CheckContext) {
|
||||
if self.queued_files.insert(file_id) {
|
||||
self.pending += 1;
|
||||
|
||||
self.scheduler
|
||||
.check_file(CheckFileTask { file_id, context });
|
||||
}
|
||||
result
|
||||
}
|
||||
}
|
||||
|
||||
enum CheckFileMessage {
|
||||
#[derive(Debug)]
|
||||
enum ThreadPoolMessage {
|
||||
ScheduleDependency(FileId),
|
||||
Completed(Diagnostics),
|
||||
Queue(FileId),
|
||||
Cancelled,
|
||||
Errored(QueryError),
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user