Merge pull request #18660 from Veykril/push-snumrtvzwqvw

fix: copied proc-macros not being cleaned up on exit

Author: Veykril

src/tools/rust-analyzer/crates/proc-macro-srv/src/dylib.rs

@@ -3,12 +3,12 @@
 mod version;
 
 use proc_macro::bridge;
-use std::{fmt, fs::File, io};
+use std::{fmt, fs, io, time::SystemTime};
 
 use libloading::Library;
 use memmap2::Mmap;
 use object::Object;
-use paths::{AbsPath, Utf8Path, Utf8PathBuf};
+use paths::{Utf8Path, Utf8PathBuf};
 use proc_macro_api::ProcMacroKind;
 
 use crate::ProcMacroSrvSpan;
@@ -23,14 +23,9 @@ fn is_derive_registrar_symbol(symbol: &str) -> bool {
     symbol.contains(NEW_REGISTRAR_SYMBOL)
 }
 
-fn find_registrar_symbol(file: &Utf8Path) -> io::Result<Option<String>> {
-    let file = File::open(file)?;
-    let buffer = unsafe { Mmap::map(&file)? };
-
-    Ok(object::File::parse(&*buffer)
-        .map_err(invalid_data_err)?
-        .exports()
-        .map_err(invalid_data_err)?
+fn find_registrar_symbol(buffer: &[u8]) -> object::Result<Option<String>> {
+    Ok(object::File::parse(buffer)?
+        .exports()?
         .into_iter()
         .map(|export| export.name())
         .filter_map(|sym| String::from_utf8(sym.into()).ok())
@@ -113,17 +108,17 @@ struct ProcMacroLibraryLibloading {
 }
 
 impl ProcMacroLibraryLibloading {
-    fn open(file: &Utf8Path) -> Result<Self, LoadProcMacroDylibError> {
-        let symbol_name = find_registrar_symbol(file)?.ok_or_else(|| {
-            invalid_data_err(format!("Cannot find registrar symbol in file {file}"))
-        })?;
+    fn open(path: &Utf8Path) -> Result<Self, LoadProcMacroDylibError> {
+        let buffer = unsafe { Mmap::map(&fs::File::open(path)?)? };
+        let symbol_name =
+            find_registrar_symbol(&buffer).map_err(invalid_data_err)?.ok_or_else(|| {
+                invalid_data_err(format!("Cannot find registrar symbol in file {path}"))
+            })?;
 
-        let abs_file: &AbsPath = file
-            .try_into()
-            .map_err(|_| invalid_data_err(format!("expected an absolute path, got {file}")))?;
-        let version_info = version::read_dylib_info(abs_file)?;
+        let version_info = version::read_dylib_info(&buffer)?;
+        drop(buffer);
 
-        let lib = load_library(file).map_err(invalid_data_err)?;
+        let lib = load_library(path).map_err(invalid_data_err)?;
         let proc_macros = crate::proc_macros::ProcMacros::from_lib(
             &lib,
             symbol_name,
@@ -133,30 +128,33 @@ impl ProcMacroLibraryLibloading {
     }
 }
 
-pub(crate) struct Expander {
-    inner: ProcMacroLibraryLibloading,
-    path: Utf8PathBuf,
-}
-
-impl Drop for Expander {
+struct RemoveFileOnDrop(Utf8PathBuf);
+impl Drop for RemoveFileOnDrop {
     fn drop(&mut self) {
         #[cfg(windows)]
-        std::fs::remove_file(&self.path).ok();
-        _ = self.path;
+        std::fs::remove_file(&self.0).unwrap();
+        _ = self.0;
     }
 }
 
+// Drop order matters as we can't remove the dylib before the library is unloaded
+pub(crate) struct Expander {
+    inner: ProcMacroLibraryLibloading,
+    _remove_on_drop: RemoveFileOnDrop,
+    modified_time: SystemTime,
+}
+
 impl Expander {
     pub(crate) fn new(lib: &Utf8Path) -> Result<Expander, LoadProcMacroDylibError> {
         // Some libraries for dynamic loading require canonicalized path even when it is
         // already absolute
         let lib = lib.canonicalize_utf8()?;
+        let modified_time = fs::metadata(&lib).and_then(|it| it.modified())?;
 
         let path = ensure_file_with_lock_free_access(&lib)?;
-
         let library = ProcMacroLibraryLibloading::open(path.as_ref())?;
 
-        Ok(Expander { inner: library, path })
+        Ok(Expander { inner: library, _remove_on_drop: RemoveFileOnDrop(path), modified_time })
     }
 
     pub(crate) fn expand<S: ProcMacroSrvSpan>(
@@ -181,6 +179,10 @@ impl Expander {
     pub(crate) fn list_macros(&self) -> Vec<(String, ProcMacroKind)> {
         self.inner.proc_macros.list_macros()
     }
+
+    pub(crate) fn modified_time(&self) -> SystemTime {
+        self.modified_time
+    }
 }
 
 /// Copy the dylib to temp directory to prevent locking in Windows
@@ -194,20 +196,23 @@ fn ensure_file_with_lock_free_access(path: &Utf8Path) -> io::Result<Utf8PathBuf>
     }
 
     let mut to = Utf8PathBuf::from_path_buf(std::env::temp_dir()).unwrap();
+    to.push("rust-analyzer-proc-macros");
+    _ = fs::create_dir(&to);
 
     let file_name = path.file_name().ok_or_else(|| {
         io::Error::new(io::ErrorKind::InvalidInput, format!("File path is invalid: {path}"))
     })?;
 
-    // Generate a unique number by abusing `HashMap`'s hasher.
-    // Maybe this will also "inspire" a libs team member to finally put `rand` in libstd.
-    let t = RandomState::new().build_hasher().finish();
-
-    let mut unique_name = t.to_string();
-    unique_name.push_str(file_name);
-
-    to.push(unique_name);
-    std::fs::copy(path, &to)?;
+    to.push({
+        // Generate a unique number by abusing `HashMap`'s hasher.
+        // Maybe this will also "inspire" a libs team member to finally put `rand` in libstd.
+        let t = RandomState::new().build_hasher().finish();
+        let mut unique_name = t.to_string();
+        unique_name.push_str(file_name);
+        unique_name.push('-');
+        unique_name
+    });
+    fs::copy(path, &to)?;
     Ok(to)
 }
 

src/tools/rust-analyzer/crates/proc-macro-srv/src/dylib/version.rs

@@ -1,13 +1,8 @@
 //! Reading proc-macro rustc version information from binary data
 
-use std::{
-    fs::File,
-    io::{self, Read},
-};
+use std::io::{self, Read};
 
-use memmap2::Mmap;
 use object::read::{File as BinaryFile, Object, ObjectSection};
-use paths::AbsPath;
 
 #[derive(Debug)]
 #[allow(dead_code)]
@@ -21,14 +16,14 @@ pub struct RustCInfo {
 }
 
 /// Read rustc dylib information
-pub fn read_dylib_info(dylib_path: &AbsPath) -> io::Result<RustCInfo> {
+pub fn read_dylib_info(buffer: &[u8]) -> io::Result<RustCInfo> {
     macro_rules! err {
         ($e:literal) => {
             io::Error::new(io::ErrorKind::InvalidData, $e)
         };
     }
 
-    let ver_str = read_version(dylib_path)?;
+    let ver_str = read_version(buffer)?;
     let mut items = ver_str.split_whitespace();
     let tag = items.next().ok_or_else(|| err!("version format error"))?;
     if tag != "rustc" {
@@ -106,11 +101,8 @@ fn read_section<'a>(dylib_binary: &'a [u8], section_name: &str) -> io::Result<&'
 ///
 /// Check this issue for more about the bytes layout:
 /// <https://github.com/rust-lang/rust-analyzer/issues/6174>
-pub fn read_version(dylib_path: &AbsPath) -> io::Result<String> {
-    let dylib_file = File::open(dylib_path)?;
-    let dylib_mmapped = unsafe { Mmap::map(&dylib_file) }?;
-
-    let dot_rustc = read_section(&dylib_mmapped, ".rustc")?;
+pub fn read_version(buffer: &[u8]) -> io::Result<String> {
+    let dot_rustc = read_section(buffer, ".rustc")?;
 
     // check if magic is valid
     if &dot_rustc[0..4] != b"rust" {
@@ -159,8 +151,12 @@ pub fn read_version(dylib_path: &AbsPath) -> io::Result<String> {
 
 #[test]
 fn test_version_check() {
-    let path = paths::AbsPathBuf::assert(crate::proc_macro_test_dylib_path());
-    let info = read_dylib_info(&path).unwrap();
+    let info = read_dylib_info(&unsafe {
+        memmap2::Mmap::map(&std::fs::File::open(crate::proc_macro_test_dylib_path()).unwrap())
+            .unwrap()
+    })
+    .unwrap();
+
     assert_eq!(
         info.version_string,
         crate::RUSTC_VERSION_STRING,

src/tools/rust-analyzer/crates/proc-macro-srv/src/lib.rs

@@ -35,7 +35,6 @@ use std::{
     fs,
     path::{Path, PathBuf},
     thread,
-    time::SystemTime,
 };
 
 use paths::{Utf8Path, Utf8PathBuf};
@@ -53,7 +52,7 @@ use crate::server_impl::TokenStream;
 pub const RUSTC_VERSION_STRING: &str = env!("RUSTC_VERSION");
 
 pub struct ProcMacroSrv<'env> {
-    expanders: HashMap<(Utf8PathBuf, SystemTime), dylib::Expander>,
+    expanders: HashMap<Utf8PathBuf, dylib::Expander>,
     span_mode: SpanMode,
     env: &'env EnvSnapshot,
 }
@@ -81,10 +80,9 @@ impl<'env> ProcMacroSrv<'env> {
     ) -> Result<(msg::FlatTree, Vec<u32>), msg::PanicMessage> {
         let span_mode = self.span_mode;
         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 expander = self
+            .expander(lib.as_ref())
+            .map_err(|err| msg::PanicMessage(format!("failed to load macro: {err}")))?;
 
         let prev_env = EnvChange::apply(snapped_env, env, current_dir.as_ref().map(<_>::as_ref));
 
@@ -107,16 +105,20 @@ impl<'env> ProcMacroSrv<'env> {
     }
 
     fn expander(&mut self, path: &Utf8Path) -> Result<&dylib::Expander, String> {
-        let time = fs::metadata(path)
-            .and_then(|it| it.modified())
-            .map_err(|err| format!("Failed to get file metadata for {path}: {err}",))?;
-
-        Ok(match self.expanders.entry((path.to_path_buf(), time)) {
-            Entry::Vacant(v) => v.insert(
-                dylib::Expander::new(path)
-                    .map_err(|err| format!("Cannot create expander for {path}: {err}",))?,
-            ),
-            Entry::Occupied(e) => e.into_mut(),
+        let expander = || {
+            dylib::Expander::new(path)
+                .map_err(|err| format!("Cannot create expander for {path}: {err}",))
+        };
+
+        Ok(match self.expanders.entry(path.to_path_buf()) {
+            Entry::Vacant(v) => v.insert(expander()?),
+            Entry::Occupied(mut e) => {
+                let time = fs::metadata(path).and_then(|it| it.modified()).ok();
+                if Some(e.get().modified_time()) != time {
+                    e.insert(expander()?);
+                }
+                e.into_mut()
+            }
         })
     }
 }