crates/proc_macro_api/src/process.rs

   1 //! Handle process life-time and message passing for proc-macro client
   2
   3 use std::{
   4     convert::{TryFrom, TryInto},
   5     ffi::{OsStr, OsString},
   6     io::{self, BufRead, BufReader, Write},
   7     path::{Path, PathBuf},
   8     process::{Child, Command, Stdio},
   9     sync::{Arc, Weak},
  10 };
  11
  12 use crossbeam_channel::{bounded, Receiver, Sender};
  13
  14 use crate::{
  15     msg::{ErrorCode, Message, Request, Response, ResponseError},
  16     rpc::{ListMacrosResult, ListMacrosTask, ProcMacroKind},
  17 };
  18
  19 #[derive(Debug, Default)]
  20 pub(crate) struct ProcMacroProcessSrv {
  21     inner: Weak<Sender<Task>>,
  22 }
  23
  24 #[derive(Debug)]
  25 pub(crate) struct ProcMacroProcessThread {
  26     // XXX: drop order is significant
  27     sender: Arc<Sender<Task>>,
  28     handle: jod_thread::JoinHandle<()>,
  29 }
  30
  31 impl ProcMacroProcessSrv {
  32     pub(crate) fn run(
  33         process_path: PathBuf,
  34         args: impl IntoIterator<Item = impl AsRef<OsStr>>,
  35     ) -> io::Result<(ProcMacroProcessThread, ProcMacroProcessSrv)> {
  36         let process = Process::run(process_path, args)?;
  37
  38         let (task_tx, task_rx) = bounded(0);
  39         let handle = jod_thread::spawn(move || {
  40             client_loop(task_rx, process);
  41         });
  42
  43         let task_tx = Arc::new(task_tx);
  44         let srv = ProcMacroProcessSrv { inner: Arc::downgrade(&task_tx) };
  45         let thread = ProcMacroProcessThread { handle, sender: task_tx };
  46
  47         Ok((thread, srv))
  48     }
  49
  50     pub(crate) fn find_proc_macros(
  51         &self,
  52         dylib_path: &Path,
  53     ) -> Result<Vec<(String, ProcMacroKind)>, tt::ExpansionError> {
  54         let task = ListMacrosTask { lib: dylib_path.to_path_buf() };
  55
  56         let result: ListMacrosResult = self.send_task(Request::ListMacro(task))?;
  57         Ok(result.macros)
  58     }
  59
  60     pub(crate) fn send_task<R>(&self, req: Request) -> Result<R, tt::ExpansionError>
  61     where
  62         R: TryFrom<Response, Error = &'static str>,
  63     {
  64         let (result_tx, result_rx) = bounded(0);
  65         let sender = match self.inner.upgrade() {
  66             None => return Err(tt::ExpansionError::Unknown("proc macro process is closed".into())),
  67             Some(it) => it,
  68         };
  69         sender
  70             .send(Task { req, result_tx })
  71             .map_err(|_| tt::ExpansionError::Unknown("proc macro server crashed".into()))?;
  72
  73         let res = result_rx
  74             .recv()
  75             .map_err(|_| tt::ExpansionError::Unknown("proc macro server crashed".into()))?;
  76
  77         match res {
  78             Some(Response::Error(err)) => {
  79                 return Err(tt::ExpansionError::ExpansionError(err.message));
  80             }
  81             Some(res) => Ok(res.try_into().map_err(|err| {
  82                 tt::ExpansionError::Unknown(format!("Fail to get response, reason : {:#?} ", err))
  83             })?),
  84             None => Err(tt::ExpansionError::Unknown("Empty result".into())),
  85         }
  86     }
  87 }
  88
  89 fn client_loop(task_rx: Receiver<Task>, mut process: Process) {
  90     let (mut stdin, mut stdout) = process.stdio().expect("couldn't access child stdio");
  91
  92     for Task { req, result_tx } in task_rx {
  93         match send_request(&mut stdin, &mut stdout, req) {
  94             Ok(res) => result_tx.send(res).unwrap(),
  95             Err(_err) => {
  96                 log::error!(
  97                     "proc macro server crashed, server process state: {:?}",
  98                     process.child.try_wait()
  99                 );
 100                 let res = Response::Error(ResponseError {
 101                     code: ErrorCode::ServerErrorEnd,
 102                     message: "proc macro server crashed".into(),
 103                 });
 104                 result_tx.send(res.into()).unwrap();
 105                 // Exit the thread.
 106                 break;
 107             }
 108         }
 109     }
 110 }
 111
 112 struct Task {
 113     req: Request,
 114     result_tx: Sender<Option<Response>>,
 115 }
 116
 117 struct Process {
 118     child: Child,
 119 }
 120
 121 impl Drop for Process {
 122     fn drop(&mut self) {
 123         let _ = self.child.kill();
 124     }
 125 }
 126
 127 impl Process {
 128     fn run(
 129         path: PathBuf,
 130         args: impl IntoIterator<Item = impl AsRef<OsStr>>,
 131     ) -> io::Result<Process> {
 132         let args: Vec<OsString> = args.into_iter().map(|s| s.as_ref().into()).collect();
 133         let child = mk_child(&path, &args)?;
 134         Ok(Process { child })
 135     }
 136
 137     fn stdio(&mut self) -> Option<(impl Write, impl BufRead)> {
 138         let stdin = self.child.stdin.take()?;
 139         let stdout = self.child.stdout.take()?;
 140         let read = BufReader::new(stdout);
 141
 142         Some((stdin, read))
 143     }
 144 }
 145
 146 fn mk_child(path: &Path, args: impl IntoIterator<Item = impl AsRef<OsStr>>) -> io::Result<Child> {
 147     Command::new(&path)
 148         .args(args)
 149         .stdin(Stdio::piped())
 150         .stdout(Stdio::piped())
 151         .stderr(Stdio::inherit())
 152         .spawn()
 153 }
 154
 155 fn send_request(
 156     mut writer: &mut impl Write,
 157     mut reader: &mut impl BufRead,
 158     req: Request,
 159 ) -> io::Result<Option<Response>> {
 160     req.write(&mut writer)?;
 161     Ok(Response::read(&mut reader)?)
 162 }