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