diff --git a/crates/red_knot/src/main.rs b/crates/red_knot/src/main.rs index 36a3e4ea7a..d7e02749f2 100644 --- a/crates/red_knot/src/main.rs +++ b/crates/red_knot/src/main.rs @@ -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) => { diff --git a/crates/red_knot/src/program/check.rs b/crates/red_knot/src/program/check.rs index 258f9d0062..af23a6eb42 100644 --- a/crates/red_knot/src/program/check.rs +++ b/crates/red_knot/src/program/check.rs @@ -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> { + pub fn check(&self, mode: ExecutionMode) -> QueryResult> { 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> { - 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 { + fn check_file(&self, file: FileId, context: &CheckFileContext) -> QueryResult { 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; +#[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, + + /// Reference to the program that is checked. + program: &'a Program, + + /// The aggregated diagnostics + diagnostics: Vec, } -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 { + 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 { + 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 { + 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 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 { - 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 { - 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(&self, context: &CheckTaskContext) -> QueryResult + 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, +/// Task to check an open file. + +#[derive(Debug)] +struct CheckOpenFileTask { + file_id: FileId, } -impl CheckContext { - fn new(sender: crossbeam::channel::Sender) -> 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(&self, context: &CheckTaskContext) -> QueryResult + 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, +/// 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(&self, context: &CheckTaskContext) -> QueryResult + 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) -> QueryResult> { - 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, - context: &CheckContext, - ) -> QueryResult> { - 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), }