583 lines · 20.8 KB · Rust
RawDownload
use anyhow::{Context as _, Result};use futures::{    AsyncBufReadExt as _, FutureExt as _, StreamExt as _,    channel::mpsc::{self},    io::BufReader,    stream::FuturesUnordered,};use gpui::{App, AppContext as _, ClipboardItem, Entity, EntityId, Task, Window};use jupyter_protocol::{    ExecutionState, JupyterKernelspec, JupyterMessage, JupyterMessageContent, KernelInfoReply,    connection_info::{ConnectionInfo, Transport},};use project::Fs;use runtimelib::{RuntimeError, dirs};use smol::net::TcpListener;use std::{    env,    fmt::Debug,    net::{IpAddr, Ipv4Addr, SocketAddr},    path::PathBuf,    sync::Arc,};use util::command::Command;use uuid::Uuid; use super::{KernelSession, RunningKernel}; #[derive(Debug, Clone)]pub struct LocalKernelSpecification {    pub name: String,    pub path: PathBuf,    pub kernelspec: JupyterKernelspec,} impl PartialEq for LocalKernelSpecification {    fn eq(&self, other: &Self) -> bool {        self.name == other.name && self.path == other.path    }} impl Eq for LocalKernelSpecification {} impl LocalKernelSpecification {    #[must_use]    fn command(&self, connection_path: &PathBuf) -> Result<Command> {        let argv = &self.kernelspec.argv;         anyhow::ensure!(!argv.is_empty(), "Empty argv in kernelspec {}", self.name);        anyhow::ensure!(argv.len() >= 2, "Invalid argv in kernelspec {}", self.name);        anyhow::ensure!(            argv.iter().any(|arg| arg == "{connection_file}"),            "Missing 'connection_file' in argv in kernelspec {}",            self.name        );         let mut cmd = util::command::new_command(&argv[0]);         for arg in &argv[1..] {            if arg == "{connection_file}" {                cmd.arg(connection_path);            } else {                cmd.arg(arg);            }        }         if let Some(env) = &self.kernelspec.env {            cmd.envs(env);        }         Ok(cmd)    }} // Find a set of open ports. This creates a listener with port set to 0. The listener will be closed at the end when it goes out of scope.// There's a race condition between closing the ports and usage by a kernel, but it's inherent to the Jupyter protocol.async fn peek_ports(ip: IpAddr) -> Result<[u16; 5]> {    let mut addr_zeroport: SocketAddr = SocketAddr::new(ip, 0);    addr_zeroport.set_port(0);    let mut ports: [u16; 5] = [0; 5];    for i in 0..5 {        let listener = TcpListener::bind(addr_zeroport).await?;        let addr = listener.local_addr()?;        ports[i] = addr.port();    }    Ok(ports)} pub struct NativeRunningKernel {    pub process: util::command::Child,    connection_path: PathBuf,    _process_status_task: Option<Task<()>>,    pub working_directory: PathBuf,    pub request_tx: mpsc::Sender<JupyterMessage>,    pub stdin_tx: mpsc::Sender<JupyterMessage>,    pub execution_state: ExecutionState,    pub kernel_info: Option<KernelInfoReply>,} impl Debug for NativeRunningKernel {    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {        f.debug_struct("RunningKernel")            .field("process", &self.process)            .finish()    }} impl NativeRunningKernel {    pub fn new<S: KernelSession + 'static>(        kernel_specification: LocalKernelSpecification,        entity_id: EntityId,        working_directory: PathBuf,        fs: Arc<dyn Fs>,        // todo: convert to weak view        session: Entity<S>,        window: &mut Window,        cx: &mut App,    ) -> Task<Result<Box<dyn RunningKernel>>> {        window.spawn(cx, async move |cx| {            let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));            let ports = peek_ports(ip).await?;             let connection_info = ConnectionInfo {                transport: Transport::TCP,                ip: ip.to_string(),                stdin_port: ports[0],                control_port: ports[1],                hb_port: ports[2],                shell_port: ports[3],                iopub_port: ports[4],                signature_scheme: "hmac-sha256".to_string(),                key: uuid::Uuid::new_v4().to_string(),                kernel_name: Some(format!("zed-{}", kernel_specification.name)),            };             let runtime_dir = dirs::runtime_dir();            fs.create_dir(&runtime_dir)                .await                .with_context(|| format!("Failed to create jupyter runtime dir {runtime_dir:?}"))?;            let connection_path = runtime_dir.join(format!("kernel-zed-{entity_id}.json"));            let content = serde_json::to_string(&connection_info)?;            fs.atomic_write(connection_path.clone(), content).await?;             let mut cmd = kernel_specification.command(&connection_path)?;             let mut process = cmd                .current_dir(&working_directory)                .stdout(util::command::Stdio::piped())                .stderr(util::command::Stdio::piped())                .stdin(util::command::Stdio::piped())                .kill_on_drop(true)                .spawn()                .context("failed to start the kernel process")?;             let session_id = Uuid::new_v4().to_string();             let iopub_socket =                runtimelib::create_client_iopub_connection(&connection_info, "", &session_id)                    .await?;            let control_socket =                runtimelib::create_client_control_connection(&connection_info, &session_id).await?;             let peer_identity = runtimelib::peer_identity_for_session(&session_id)?;            let shell_socket =                runtimelib::create_client_shell_connection_with_identity(                    &connection_info,                    &session_id,                    peer_identity.clone(),                )                .await?;            let stdin_socket = runtimelib::create_client_stdin_connection_with_identity(                &connection_info,                &session_id,                peer_identity,            )            .await?;             let (mut shell_send, shell_recv) = shell_socket.split();            let (mut control_send, control_recv) = control_socket.split();            let (mut stdin_send, stdin_recv) = stdin_socket.split();             let (request_tx, mut request_rx) =                futures::channel::mpsc::channel::<JupyterMessage>(100);            let (stdin_tx, mut stdin_rx) =                futures::channel::mpsc::channel::<JupyterMessage>(100);             let recv_task = cx.spawn({                let session = session.clone();                let mut iopub = iopub_socket;                let mut shell = shell_recv;                let mut control = control_recv;                let mut stdin = stdin_recv;                 async move |cx| -> anyhow::Result<()> {                    loop {                        let (channel, result) = futures::select! {                            msg = iopub.read().fuse() => ("iopub", msg),                            msg = shell.read().fuse() => ("shell", msg),                            msg = control.read().fuse() => ("control", msg),                            msg = stdin.read().fuse() => ("stdin", msg),                        };                        match result {                            Ok(message) => {                                session                                    .update_in(cx, |session, window, cx| {                                        session.route(&message, window, cx);                                    })                                    .ok();                            }                            Err(                                ref err @ (RuntimeError::ParseError { .. }                                | RuntimeError::SerdeError(_)),                            ) => {                                let error_detail =                                    format!("Kernel issue on {channel} channel\n\n{err}");                                log::warn!("kernel: {error_detail}");                                let workspace_window = session                                    .update_in(cx, |_, window, _cx| {                                        window                                            .window_handle()                                            .downcast::<workspace::Workspace>()                                    })                                    .ok()                                    .flatten();                                if let Some(workspace_window) = workspace_window {                                    workspace_window                                        .update(cx, |workspace, _window, cx| {                                            struct KernelReadError;                                            workspace.show_toast(                                                workspace::Toast::new(                                                    workspace::notifications::NotificationId::unique::<KernelReadError>(),                                                    error_detail.clone(),                                                )                                                .on_click(                                                    "Copy Error",                                                    move |_window, cx| {                                                        cx.write_to_clipboard(                                                            ClipboardItem::new_string(                                                                error_detail.clone(),                                                            ),                                                        );                                                    },                                                ),                                                cx,                                            );                                        })                                        .ok();                                }                            }                            Err(err) => {                                anyhow::bail!("{channel} recv: {err}");                            }                        }                    }                }            });             let routing_task = cx.background_spawn({                async move {                    while let Some(message) = request_rx.next().await {                        match message.content {                            JupyterMessageContent::DebugRequest(_)                            | JupyterMessageContent::InterruptRequest(_)                            | JupyterMessageContent::ShutdownRequest(_) => {                                control_send.send(message).await?;                            }                            _ => {                                shell_send.send(message).await?;                            }                        }                    }                    anyhow::Ok(())                }            });             let stdin_routing_task = cx.background_spawn({                async move {                    while let Some(message) = stdin_rx.next().await {                        stdin_send.send(message).await?;                    }                    anyhow::Ok(())                }            });             let stderr = process.stderr.take();            let stdout = process.stdout.take();             cx.spawn(async move |_cx| {                use futures::future::Either;                 let stderr_lines = match stderr {                    Some(s) => Either::Left(                        BufReader::new(s)                            .lines()                            .map(|line| (log::Level::Error, line)),                    ),                    None => Either::Right(futures::stream::empty()),                };                let stdout_lines = match stdout {                    Some(s) => Either::Left(                        BufReader::new(s)                            .lines()                            .map(|line| (log::Level::Info, line)),                    ),                    None => Either::Right(futures::stream::empty()),                };                let mut lines = futures::stream::select(stderr_lines, stdout_lines);                while let Some((level, Ok(line))) = lines.next().await {                    log::log!(level, "kernel: {}", line);                }            })            .detach();             cx.spawn({                let session = session.clone();                async move |cx| {                    async fn with_name(                        name: &'static str,                        task: Task<Result<()>>,                    ) -> (&'static str, Result<()>) {                        (name, task.await)                    }                     let mut tasks = FuturesUnordered::new();                    tasks.push(with_name("recv task", recv_task));                    tasks.push(with_name("routing task", routing_task));                    tasks.push(with_name("stdin routing task", stdin_routing_task));                     while let Some((name, result)) = tasks.next().await {                        if let Err(err) = result {                            log::error!("kernel: handling failed for {name}: {err:?}");                             session.update(cx, |session, cx| {                                session.kernel_errored(                                    format!("handling failed for {name}: {err}"),                                    cx,                                );                                cx.notify();                            });                        }                    }                }            })            .detach();             let status = process.status();             let process_status_task = cx.spawn(async move |cx| {                let error_message = match status.await {                    Ok(status) => {                        if status.success() {                            log::info!("kernel process exited successfully");                            return;                        }                         format!("kernel process exited with status: {:?}", status)                    }                    Err(err) => {                        format!("kernel process exited with error: {:?}", err)                    }                };                 log::error!("{}", error_message);                 session.update(cx, |session, cx| {                    session.kernel_errored(error_message, cx);                     cx.notify();                });            });             anyhow::Ok(Box::new(Self {                process,                request_tx,                stdin_tx,                working_directory,                _process_status_task: Some(process_status_task),                connection_path,                execution_state: ExecutionState::Idle,                kernel_info: None,            }) as Box<dyn RunningKernel>)        })    }} impl RunningKernel for NativeRunningKernel {    fn request_tx(&self) -> mpsc::Sender<JupyterMessage> {        self.request_tx.clone()    }     fn stdin_tx(&self) -> mpsc::Sender<JupyterMessage> {        self.stdin_tx.clone()    }     fn working_directory(&self) -> &PathBuf {        &self.working_directory    }     fn execution_state(&self) -> &ExecutionState {        &self.execution_state    }     fn set_execution_state(&mut self, state: ExecutionState) {        self.execution_state = state;    }     fn kernel_info(&self) -> Option<&KernelInfoReply> {        self.kernel_info.as_ref()    }     fn set_kernel_info(&mut self, info: KernelInfoReply) {        self.kernel_info = Some(info);    }     fn force_shutdown(&mut self, _window: &mut Window, _cx: &mut App) -> Task<anyhow::Result<()>> {        self.kill();        Task::ready(Ok(()))    }     fn kill(&mut self) {        self._process_status_task.take();        self.request_tx.close_channel();        self.stdin_tx.close_channel();        self.process.kill().ok();    }} impl Drop for NativeRunningKernel {    fn drop(&mut self) {        std::fs::remove_file(&self.connection_path).ok();        self.kill();    }} async fn read_kernelspec_at(    // Path should be a directory to a jupyter kernelspec, as in    // /usr/local/share/jupyter/kernels/python3    kernel_dir: PathBuf,    fs: &dyn Fs,) -> Result<LocalKernelSpecification> {    let path = kernel_dir;    let kernel_name = if let Some(kernel_name) = path.file_name() {        kernel_name.to_string_lossy().into_owned()    } else {        anyhow::bail!("Invalid kernelspec directory: {path:?}");    };     if !fs.is_dir(path.as_path()).await {        anyhow::bail!("Not a directory: {path:?}");    }     let expected_kernel_json = path.join("kernel.json");    let spec = fs.load(expected_kernel_json.as_path()).await?;    let spec = serde_json::from_str::<JupyterKernelspec>(&spec)?;     Ok(LocalKernelSpecification {        name: kernel_name,        path,        kernelspec: spec,    })} /// Read a directory of kernelspec directoriesasync fn read_kernels_dir(path: PathBuf, fs: &dyn Fs) -> Result<Vec<LocalKernelSpecification>> {    let mut kernelspec_dirs = fs.read_dir(&path).await?;     let mut valid_kernelspecs = Vec::new();    while let Some(path) = kernelspec_dirs.next().await {        match path {            Ok(path) => {                if fs.is_dir(path.as_path()).await                    && let Ok(kernelspec) = read_kernelspec_at(path, fs).await                {                    valid_kernelspecs.push(kernelspec);                }            }            Err(err) => log::warn!("Error reading kernelspec directory: {err:?}"),        }    }     Ok(valid_kernelspecs)} pub async fn local_kernel_specifications(fs: Arc<dyn Fs>) -> Result<Vec<LocalKernelSpecification>> {    let mut data_dirs = dirs::data_dirs();     // Pick up any kernels from conda or conda environment    if let Ok(conda_prefix) = env::var("CONDA_PREFIX") {        let conda_prefix = PathBuf::from(conda_prefix);        let conda_data_dir = conda_prefix.join("share").join("jupyter");        data_dirs.push(conda_data_dir);    }     // Search for kernels inside the base python environment    let command = util::command::new_command("python")        .arg("-c")        .arg("import sys; print(sys.prefix)")        .output()        .await;     if let Ok(command) = command

Lines 1–500 of 583.