Skip to content

Commit

Permalink
Faster env snapshotting in proc-macro-srv
Browse files Browse the repository at this point in the history
  • Loading branch information
@Veykril
Veykril committed Jun 30, 2024
1 parent 678420e commit da7e07f
Show file tree
Hide file tree
Showing 6 changed files with 146 additions and 117 deletions.
31 changes: 17 additions & 14 deletions crates/proc-macro-api/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -156,23 +156,26 @@ impl ProcMacro {
let call_site = span_data_table.insert_full(call_site).0;
let mixed_site = span_data_table.insert_full(mixed_site).0;
let task = ExpandMacro {
macro_body: FlatTree::new(subtree, version, &mut span_data_table),
macro_name: self.name.to_string(),
attributes: attr.map(|subtree| FlatTree::new(subtree, version, &mut span_data_table)),
data: msg::ExpandMacroData {
macro_body: FlatTree::new(subtree, version, &mut span_data_table),
macro_name: self.name.to_string(),
attributes: attr
.map(|subtree| FlatTree::new(subtree, version, &mut span_data_table)),
has_global_spans: ExpnGlobals {
serialize: version >= HAS_GLOBAL_SPANS,
def_site,
call_site,
mixed_site,
},
span_data_table: if version >= RUST_ANALYZER_SPAN_SUPPORT {
serialize_span_data_index_map(&span_data_table)
} else {
Vec::new()
},
},
lib: self.dylib_path.to_path_buf().into(),
env: env.into(),
current_dir,
has_global_spans: ExpnGlobals {
serialize: version >= HAS_GLOBAL_SPANS,
def_site,
call_site,
mixed_site,
},
span_data_table: if version >= RUST_ANALYZER_SPAN_SUPPORT {
serialize_span_data_index_map(&span_data_table)
} else {
Vec::new()
},
};

let response = self.process.send_task(msg::Request::ExpandMacro(Box::new(task)))?;
Expand Down
44 changes: 26 additions & 18 deletions crates/proc-macro-api/src/msg.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,16 @@ pub struct PanicMessage(pub String);

#[derive(Debug, Serialize, Deserialize)]
pub struct ExpandMacro {
pub lib: Utf8PathBuf,
/// Environment variables to set during macro expansion.
pub env: Vec<(String, String)>,
pub current_dir: Option<String>,
#[serde(flatten)]
pub data: ExpandMacroData,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct ExpandMacroData {
/// Argument of macro call.
///
/// In custom derive this will be a struct or enum; in attribute-like macro - underlying
Expand All @@ -86,13 +96,6 @@ pub struct ExpandMacro {

/// Possible attributes for the attribute-like macros.
pub attributes: Option<FlatTree>,

pub lib: Utf8PathBuf,

/// Environment variables to set during macro expansion.
pub env: Vec<(String, String)>,

pub current_dir: Option<String>,
/// marker for serde skip stuff
#[serde(skip_serializing_if = "ExpnGlobals::skip_serializing_if")]
#[serde(default)]
Expand Down Expand Up @@ -268,25 +271,30 @@ mod tests {
let tt = fixture_token_tree();
let mut span_data_table = Default::default();
let task = ExpandMacro {
macro_body: FlatTree::new(&tt, CURRENT_API_VERSION, &mut span_data_table),
macro_name: Default::default(),
attributes: None,
data: ExpandMacroData {
macro_body: FlatTree::new(&tt, CURRENT_API_VERSION, &mut span_data_table),
macro_name: Default::default(),
attributes: None,
has_global_spans: ExpnGlobals {
serialize: true,
def_site: 0,
call_site: 0,
mixed_site: 0,
},
span_data_table: Vec::new(),
},
lib: Utf8PathBuf::from_path_buf(std::env::current_dir().unwrap()).unwrap(),
env: Default::default(),
current_dir: Default::default(),
has_global_spans: ExpnGlobals {
serialize: true,
def_site: 0,
call_site: 0,
mixed_site: 0,
},
span_data_table: Vec::new(),
};

let json = serde_json::to_string(&task).unwrap();
// println!("{}", json);
let back: ExpandMacro = serde_json::from_str(&json).unwrap();

assert_eq!(tt, back.macro_body.to_subtree_resolved(CURRENT_API_VERSION, &span_data_table));
assert_eq!(
tt,
back.data.macro_body.to_subtree_resolved(CURRENT_API_VERSION, &span_data_table)
);
}
}
3 changes: 2 additions & 1 deletion crates/proc-macro-srv-cli/src/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,8 @@ fn run() -> io::Result<()> {

let write_response = |msg: msg::Response| msg.write(&mut io::stdout().lock());

let mut srv = proc_macro_srv::ProcMacroSrv::default();
let env = EnvSnapshot::new();

Check failure on line 41 in crates/proc-macro-srv-cli/src/main.rs

View workflow job for this annotation

GitHub Actions / Rust (ubuntu-latest) 

failed to resolve: use of undeclared type `EnvSnapshot`
let mut srv = proc_macro_srv::ProcMacroSrv::new(&env);
let mut buf = String::new();

while let Some(req) = read_request(&mut buf)? {
Expand Down
177 changes: 96 additions & 81 deletions crates/proc-macro-srv/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,9 @@ use std::{
collections::{hash_map::Entry, HashMap},
env,
ffi::OsString,
fs, thread,
fs,
path::{Path, PathBuf},
thread,
time::SystemTime,
};

Expand All @@ -50,15 +52,21 @@ use crate::server_impl::TokenStream;

pub const RUSTC_VERSION_STRING: &str = env!("RUSTC_VERSION");

#[derive(Default)]
pub struct ProcMacroSrv {
pub struct ProcMacroSrv<'env> {
expanders: HashMap<(Utf8PathBuf, SystemTime), dylib::Expander>,
span_mode: SpanMode,
env: &'env EnvSnapshot,
}

impl<'env> ProcMacroSrv<'env> {
pub fn new(env: &'env EnvSnapshot) -> Self {
Self { expanders: Default::default(), span_mode: Default::default(), env }
}
}

const EXPANDER_STACK_SIZE: usize = 8 * 1024 * 1024;

impl ProcMacroSrv {
impl<'env> ProcMacroSrv<'env> {
pub fn set_span_mode(&mut self, span_mode: SpanMode) {
self.span_mode = span_mode;
}
Expand All @@ -69,52 +77,24 @@ impl ProcMacroSrv {

pub fn expand(
&mut self,
task: msg::ExpandMacro,
msg::ExpandMacro { lib, env, current_dir, data }: msg::ExpandMacro,
) -> Result<(msg::FlatTree, Vec<u32>), msg::PanicMessage> {
let span_mode = self.span_mode;
let expander = self.expander(task.lib.as_ref()).map_err(|err| {
let snapped_env = self.env;
let expander = self.expander(lib.as_ref()).map_err(|err| {
debug_assert!(false, "should list macros before asking to expand");
msg::PanicMessage(format!("failed to load macro: {err}"))
})?;

let prev_env = EnvSnapshot::new();
for (k, v) in &task.env {
env::set_var(k, v);
}
let prev_working_dir = match &task.current_dir {
Some(dir) => {
let prev_working_dir = std::env::current_dir().ok();
if let Err(err) = std::env::set_current_dir(dir) {
eprintln!("Failed to set the current working dir to {dir}. Error: {err:?}")
}
prev_working_dir
}
None => None,
};

let ExpnGlobals { def_site, call_site, mixed_site, .. } = task.has_global_spans;
let prev_env = EnvChange::apply(snapped_env, env, current_dir.as_ref().map(<_>::as_ref));

let result = match span_mode {
SpanMode::Id => {
expand_id(task, expander, def_site, call_site, mixed_site).map(|it| (it, vec![]))
}
SpanMode::RustAnalyzer => {
expand_ra_span(task, expander, def_site, call_site, mixed_site)
}
SpanMode::Id => expand_id(data, expander).map(|it| (it, vec![])),
SpanMode::RustAnalyzer => expand_ra_span(data, expander),
};

prev_env.rollback();

if let Some(dir) = prev_working_dir {
if let Err(err) = std::env::set_current_dir(&dir) {
eprintln!(
"Failed to set the current working dir to {}. Error: {:?}",
dir.display(),
err
)
}
}

result.map_err(msg::PanicMessage)
}

Expand Down Expand Up @@ -168,32 +148,28 @@ impl ProcMacroSrvSpan for Span {
}

fn expand_id(
task: msg::ExpandMacro,
msg::ExpandMacroData {
macro_body,
macro_name,
attributes,
has_global_spans: ExpnGlobals { serialize: _, def_site, call_site, mixed_site },
span_data_table: _,
}: msg::ExpandMacroData,
expander: &dylib::Expander,
def_site: usize,
call_site: usize,
mixed_site: usize,
) -> Result<msg::FlatTree, String> {
let def_site = TokenId(def_site as u32);
let call_site = TokenId(call_site as u32);
let mixed_site = TokenId(mixed_site as u32);

let macro_body = task.macro_body.to_subtree_unresolved(CURRENT_API_VERSION);
let attributes = task.attributes.map(|it| it.to_subtree_unresolved(CURRENT_API_VERSION));
let macro_body = macro_body.to_subtree_unresolved(CURRENT_API_VERSION);
let attributes = attributes.map(|it| it.to_subtree_unresolved(CURRENT_API_VERSION));
let result = thread::scope(|s| {
let thread = thread::Builder::new()
.stack_size(EXPANDER_STACK_SIZE)
.name(task.macro_name.clone())
.name(macro_name.clone())
.spawn_scoped(s, || {
expander
.expand(
&task.macro_name,
macro_body,
attributes,
def_site,
call_site,
mixed_site,
)
.expand(&macro_name, macro_body, attributes, def_site, call_site, mixed_site)
.map(|it| msg::FlatTree::new_raw(&it, CURRENT_API_VERSION))
});
let res = match thread {
Expand All @@ -210,35 +186,31 @@ fn expand_id(
}

fn expand_ra_span(
task: msg::ExpandMacro,
msg::ExpandMacroData {
macro_body,
macro_name,
attributes,
has_global_spans: ExpnGlobals { serialize: _, def_site, call_site, mixed_site },
span_data_table,
}: msg::ExpandMacroData,
expander: &dylib::Expander,
def_site: usize,
call_site: usize,
mixed_site: usize,
) -> Result<(msg::FlatTree, Vec<u32>), String> {
let mut span_data_table = deserialize_span_data_index_map(&task.span_data_table);
let mut span_data_table = deserialize_span_data_index_map(&span_data_table);

let def_site = span_data_table[def_site];
let call_site = span_data_table[call_site];
let mixed_site = span_data_table[mixed_site];

let macro_body = task.macro_body.to_subtree_resolved(CURRENT_API_VERSION, &span_data_table);
let macro_body = macro_body.to_subtree_resolved(CURRENT_API_VERSION, &span_data_table);
let attributes =
task.attributes.map(|it| it.to_subtree_resolved(CURRENT_API_VERSION, &span_data_table));
attributes.map(|it| it.to_subtree_resolved(CURRENT_API_VERSION, &span_data_table));
let result = thread::scope(|s| {
let thread = thread::Builder::new()
.stack_size(EXPANDER_STACK_SIZE)
.name(task.macro_name.clone())
.name(macro_name.clone())
.spawn_scoped(s, || {
expander
.expand(
&task.macro_name,
macro_body,
attributes,
def_site,
call_site,
mixed_site,
)
.expand(&macro_name, macro_body, attributes, def_site, call_site, mixed_site)
.map(|it| {
(
msg::FlatTree::new(&it, CURRENT_API_VERSION, &mut span_data_table),
Expand Down Expand Up @@ -269,31 +241,74 @@ impl PanicMessage {
}
}

struct EnvSnapshot {
pub struct EnvSnapshot {
vars: HashMap<OsString, OsString>,
}

impl EnvSnapshot {
fn new() -> EnvSnapshot {
pub fn new() -> EnvSnapshot {
EnvSnapshot { vars: env::vars_os().collect() }
}
}

struct EnvChange<'snap> {
changed_vars: Vec<String>,
prev_working_dir: Option<PathBuf>,
snap: &'snap EnvSnapshot,
}

impl<'snap> EnvChange<'snap> {
fn apply(
snap: &'snap EnvSnapshot,
new_vars: Vec<(String, String)>,
current_dir: Option<&Path>,
) -> EnvChange<'snap> {
let prev_working_dir = match current_dir {
Some(dir) => {
let prev_working_dir = std::env::current_dir().ok();
if let Err(err) = std::env::set_current_dir(dir) {
eprintln!(
"Failed to set the current working dir to {}. Error: {err:?}",
dir.display()
)
}
prev_working_dir
}
None => None,
};
EnvChange {
snap,
changed_vars: new_vars
.into_iter()
.map(|(k, v)| {
env::set_var(&k, v);
k
})
.collect(),
prev_working_dir,
}
}

fn rollback(self) {}
}

impl Drop for EnvSnapshot {
impl Drop for EnvChange<'_> {
fn drop(&mut self) {
for (name, value) in env::vars_os() {
let old_value = self.vars.remove(&name);
if old_value != Some(value) {
match old_value {
None => env::remove_var(name),
Some(old_value) => env::set_var(name, old_value),
}
for name in self.changed_vars.drain(..) {
match self.snap.vars.get::<std::ffi::OsStr>(name.as_ref()) {
Some(prev_val) => env::set_var(name, prev_val),
None => env::remove_var(name),
}
}
for (name, old_value) in self.vars.drain() {
env::set_var(name, old_value)

if let Some(dir) = &self.prev_working_dir {
if let Err(err) = std::env::set_current_dir(&dir) {
eprintln!(
"Failed to set the current working dir to {}. Error: {:?}",
dir.display(),
err
)
}
}
}
}
Expand Down
Loading

0 comments on commit da7e07f

Please sign in to comment.