(Sources-4) Implement Event::SourceCompleted for installed package sources (#1289)

Progress towards #1234
Branched from #1288

This PR sets up the `main_loop.rs` handling for package sources.

Here is the set up:

- The database changes via a `revision` bump
- We `dispatch()` a series of `SourceRequest`s via
`lsp::spawn_blocking()`. One `SourceRequest` is made per package
currently used by the workspace (removing any packages we've already
requested sources for).
- On the tokio task, an implementer of `trait SourceProvider` is called
to `provide()` sources for a package. On success, it sends the
`directory` of the sources back to the main loop via a new
`Event::SourceCompleted` event.
- Note that this all happens on a task, not blocking the main loop while
we are creating these sources.
- The main loop picks this `Event` up and calls
`db.set_package_sources(package, &directory)` for that package, which
actually loads the files from the directory.
- The next `revision` of the `db` can now use these `package.files()` in
LSP requests

I purposefully have not hooked up a live production level implementer of
`trait SourceProvider`. Instead we have a `TestProvider` that is used to
verify that all of this infrastructure is working, but for production
this is all set to `None` and no `SourceRequest`s are ever sent out, so
it is safe to merge.

If we agree on the shape of all of this, then in the next PRs I can hook
up the real source provider.

Author: DavisVaughan

crates/ark/src/lsp.rs

@@ -32,6 +32,7 @@ pub mod inputs;
 pub mod main_loop;
 pub mod markdown;
 pub(crate) mod open_file;
+pub mod sources;
 
 pub mod find_references;
 pub mod rename;

crates/ark/src/lsp/main_loop.rs

@@ -53,6 +53,8 @@ use crate::lsp::diagnostics::generate_diagnostics;
 use crate::lsp::handlers;
 use crate::lsp::indexer;
 use crate::lsp::open_file::OpenFile;
+use crate::lsp::sources::SourceCompleted;
+use crate::lsp::sources::SourceScheduler;
 use crate::lsp::state::WorldState;
 use crate::lsp::state_handlers;
 use crate::lsp::state_handlers::ConsoleInputs;
@@ -98,6 +100,7 @@ pub(crate) enum Event {
     Lsp(LspMessage),
     Kernel(KernelNotification),
     OakScanCompleted(ScanCompleted),
+    SourceCompleted(SourceCompleted),
 }
 
 #[derive(Debug)]
@@ -189,6 +192,24 @@ pub(crate) struct LspState {
     /// main-loop handlers. Must be out of [`WorldState`] because the scheduler
     /// is not clonable.
     pub(crate) oak_scheduler: ScanScheduler,
+
+    /// Scheduler of [crate::lsp::sources::SourceRequest]s. Scheduling and source
+    /// consumption all happen from the main loop.
+    pub(crate) source_scheduler: SourceScheduler,
+}
+
+impl LspState {
+    pub(crate) fn new(
+        console_notification_tx: TokioUnboundedSender<ConsoleNotification>,
+        source_scheduler: SourceScheduler,
+    ) -> Self {
+        Self {
+            capabilities: Capabilities::default(),
+            console_notification_tx,
+            oak_scheduler: ScanScheduler::new(),
+            source_scheduler,
+        }
+    }
 }
 
 /// State for the auxiliary loop
@@ -246,29 +267,20 @@ impl GlobalState {
 
         Self::from_parts(
             client,
-            console_notification_tx,
             WorldState::new(db, library),
+            LspState::new(console_notification_tx, SourceScheduler::new(None)),
         )
     }
 
-    /// Assemble the state around an already-built `WorldState`. Splitting this
-    /// out from [`GlobalState::new`] lets tests construct a state without the
-    /// R `.libPaths()` lookup that `new` does.
-    fn from_parts(
-        client: Client,
-        console_notification_tx: TokioUnboundedSender<ConsoleNotification>,
-        world: WorldState,
-    ) -> Self {
+    /// Assemble the state around an already-built [`WorldState`] and
+    /// [`LspState`]. Splitting this out from [`GlobalState::new`] lets tests
+    /// construct a state without the R `.libPaths()` lookup that `new` does, and
+    /// with a db / provider configured up front.
+    pub(crate) fn from_parts(client: Client, world: WorldState, lsp_state: LspState) -> Self {
         // Transmission channel for the main loop events. Shared with the
         // tower-lsp backend and the Jupyter kernel.
         let (events_tx, events_rx) = tokio_unbounded_channel::<Event>();
 
-        let lsp_state = LspState {
-            capabilities: Capabilities::default(),
-            console_notification_tx,
-            oak_scheduler: ScanScheduler::new(),
-        };
-
         Self {
             world,
             lsp_state,
@@ -544,6 +556,12 @@ impl GlobalState {
                     warm_workspace_index(self.world.db.clone());
                 }
             },
+
+            Event::SourceCompleted(SourceCompleted { package, response }) => {
+                if let Some(directory) = self.lsp_state.source_scheduler.finish(package, response) {
+                    self.world.db.set_package_sources(package, &directory);
+                }
+            },
         }
         lsp::log_info!("Finished handling event in {}ms", loop_tick.elapsed().as_millis());
 
@@ -555,6 +573,9 @@ impl GlobalState {
         if salsa::plumbing::current_revision(&self.world.db) != old_revision {
             lsp::log_info!("World state revision advanced");
             diagnostics_refresh_all(&self.world);
+            self.lsp_state
+                .source_scheduler
+                .schedule(&self.world.db, &self.events_tx);
         }
 
         Ok(())
@@ -587,21 +608,15 @@ impl GlobalState {
 /// real `handle_event()` and the private channels rather than a reconstruction.
 #[cfg(test)]
 impl GlobalState {
-    /// Build a state with an empty db and no R library paths. Takes a `client`
-    /// because the struct holds one, but the event paths exercised in tests
-    /// never touch it.
-    pub(crate) fn new_test(client: Client) -> Self {
-        let (console_notification_tx, _) = tokio_unbounded_channel::<ConsoleNotification>();
-        let world = WorldState::new(OakDatabase::new(), Library::new(vec![]));
-        Self::from_parts(client, console_notification_tx, world)
-    }
-
-    /// Run `event` through the real `handle_event`, then pump the scan
-    /// completions it spawns until the scheduler goes idle. This is what
-    /// `main_loop()` does, minus the surrounding `loop`.
+    /// Run `event` through the real `handle_event`, then pump any pending
+    /// events until we reach quiescence. This includes:
+    /// - Pending oak scans
+    /// - Pending source requests
     pub(crate) async fn handle_event_to_quiescence(&mut self, event: Event) {
         self.handle_event(event).await.unwrap();
-        while self.lsp_state.oak_scheduler.has_pending_scans() {
+        while self.lsp_state.oak_scheduler.has_pending_scans() ||
+            self.lsp_state.source_scheduler.has_pending()
+        {
             let event = self.next_event().await;
             self.handle_event(event).await.unwrap();
         }

crates/ark/src/lsp/sources.rs

@@ -0,0 +1,184 @@
+use std::collections::HashMap;
+use std::path::Path;
+use std::path::PathBuf;
+use std::sync::Arc;
+
+use aether_path::FilePath;
+use oak_db::Db;
+use oak_db::Package;
+use stdext::result::ResultExt;
+
+use crate::lsp::main_loop::Event;
+use crate::lsp::main_loop::TokioUnboundedSender;
+
+/// Handle source requests to provide R source files for an installed package
+///
+/// Implementations live outside of oak, oak is only in charge of ingesting the
+/// returned directory.
+pub(crate) trait SourceHandler: Send + Sync {
+    fn handle(&self, request: &SourceRequest) -> SourceResponse;
+}
+
+#[derive(Debug, Clone)]
+pub(crate) struct SourceRequest {
+    name: String,
+    version: String,
+    library_path: PathBuf,
+}
+
+// TODO!: Remove when we have a production `SourceHandler`
+#[cfg_attr(not(test), expect(dead_code))]
+#[derive(Debug)]
+pub(crate) enum SourceResponse {
+    Success(PathBuf),
+    Failure,
+}
+
+#[derive(Debug)]
+pub(crate) struct SourceCompleted {
+    pub(crate) package: Package,
+    pub(crate) response: SourceResponse,
+}
+
+/// State of a particular [Package]'s [SourceRequest]
+///
+/// There is also a 3rd implied state of "we've never seen this package before" if it
+/// isn't in the `state` hash map.
+enum SourceState {
+    /// The [SourceRequest] is in flight
+    Pending,
+
+    /// We have received a [SourceResponse]. Regardless of [SourceResponse::Success] or
+    /// [SourceResponse::Failure], we mark the package as `Finished` so we never request
+    /// it again.
+    Finished,
+}
+
+pub(crate) struct SourceScheduler {
+    // TODO!: Remove the `Option<>` when we implement a production `SourceHandler`
+    handler: Option<Arc<dyn SourceHandler>>,
+    state: HashMap<Package, SourceState>,
+}
+
+impl SourceScheduler {
+    pub(crate) fn new(handler: Option<Arc<dyn SourceHandler>>) -> Self {
+        Self {
+            handler,
+            state: HashMap::new(),
+        }
+    }
+
+    pub(crate) fn schedule(&mut self, db: &dyn Db, events_tx: &TokioUnboundedSender<Event>) {
+        let Some(handler) = &self.handler else {
+            return;
+        };
+
+        // For each package used by the workspace, request its sources if we have never
+        // seen it before
+        for package in oak_db::workspace_dependencies(db) {
+            if self.state.contains_key(package) {
+                // If we've seen this package before, don't request sources again!
+                continue;
+            }
+
+            let package = *package;
+
+            let Some(request) = SourceRequest::from_package(db, &package).log_err() else {
+                // Go straight to `Finished` if we can't generate the source request,
+                // something is structurally wrong
+                self.state.insert(package, SourceState::Finished);
+                continue;
+            };
+
+            let handler = Arc::clone(handler);
+            let tx = events_tx.clone();
+
+            // Mark as `Pending` just before launching the tokio task
+            self.state.insert(package, SourceState::Pending);
+
+            crate::lsp::spawn_blocking(move || {
+                let response = handler.handle(&request);
+
+                tx.send(Event::SourceCompleted(SourceCompleted {
+                    package,
+                    response,
+                }))
+                .log_err();
+
+                Ok(None)
+            });
+        }
+    }
+
+    #[must_use]
+    pub(crate) fn finish(&mut self, package: Package, response: SourceResponse) -> Option<PathBuf> {
+        self.state.insert(package, SourceState::Finished);
+        match response {
+            SourceResponse::Success(directory) => Some(directory),
+            SourceResponse::Failure => None,
+        }
+    }
+
+    /// Whether any source request is in flight. Allows tests to deterministically "wait"
+    /// for pending source requests to finish.
+    #[cfg(test)]
+    pub(crate) fn has_pending(&self) -> bool {
+        self.state
+            .values()
+            .any(|state| matches!(state, SourceState::Pending))
+    }
+}
+
+impl SourceRequest {
+    fn from_package(db: &dyn Db, package: &Package) -> anyhow::Result<Self> {
+        let name = package.name(db).clone();
+
+        let Some(version) = package.version(db).to_owned() else {
+            return Err(anyhow::anyhow!(
+                "Package {name} is missing a version to provide sources for."
+            ));
+        };
+
+        let library_path = match package.description_path(db) {
+            FilePath::File(path) => {
+                match path.as_path().as_std_path().parent().and_then(Path::parent) {
+                    Some(library_path) => library_path.to_path_buf(),
+                    None => {
+                        return Err(anyhow::anyhow!(
+                            "Package {name} does not have an associated library path."
+                        ))
+                    },
+                }
+            },
+            FilePath::Virtual(uri) => {
+                return Err(anyhow::anyhow!(
+                    "Package {name} is unexpectedly a virtual uri {uri}."
+                ))
+            },
+        };
+
+        Ok(Self {
+            name,
+            version,
+            library_path,
+        })
+    }
+
+    // TODO!: Remove when we have a production `SourceHandler`
+    #[cfg_attr(not(test), expect(dead_code))]
+    pub(crate) fn name(&self) -> &str {
+        &self.name
+    }
+
+    // TODO!: Remove when we have a production `SourceHandler`
+    #[cfg_attr(not(test), expect(dead_code))]
+    pub(crate) fn version(&self) -> &str {
+        &self.version
+    }
+
+    // TODO!: Remove when we have a production `SourceHandler`
+    #[cfg_attr(not(test), expect(dead_code))]
+    pub(crate) fn library_path(&self) -> &Path {
+        &self.library_path
+    }
+}

crates/ark/src/lsp/tests.rs

@@ -4,6 +4,7 @@ mod find_references;
 mod goto_definition;
 mod main_loop;
 mod rename;
+mod sources;
 mod state;
 mod state_handlers;
 mod utils;