From 488dff6433ea3841acc5a6059bb34422952a410a Mon Sep 17 00:00:00 2001 From: ZeWaka Date: Sun, 4 Oct 2026 02:17:35 -0700 Subject: [PATCH] Allocate proc bodies in a separate mimalloc heap Parse proc bodies into their own mimalloc heap, then drop them and collect that heap once both background jobs finish, so their pages actually return to the OS. --- crates/dm-langserver/src/jrpc_io.rs | 26 ++++++++++++++++ crates/dm-langserver/src/main.rs | 43 +++++++++++++++++++++------ crates/dm-langserver/src/proc_heap.rs | 29 ++++++++++++++++++ crates/dreammaker/src/parser.rs | 23 ++++++++++++-- crates/dreammaker/tests/ast_tests.rs | 20 +++++++++++++ 5 files changed, 130 insertions(+), 11 deletions(-) create mode 100644 crates/dm-langserver/src/proc_heap.rs diff --git a/crates/dm-langserver/src/jrpc_io.rs b/crates/dm-langserver/src/jrpc_io.rs index 0e905f28..53ec4364 100644 --- a/crates/dm-langserver/src/jrpc_io.rs +++ b/crates/dm-langserver/src/jrpc_io.rs @@ -3,6 +3,7 @@ //! JSON-RPC over stdin/stdout with Content-Length headers. use std::io::{self, BufRead, Write}; +use std::sync::mpsc; pub fn run_until_stdin_eof(mut f: F) { let stdin = io::stdin(); @@ -12,6 +13,31 @@ pub fn run_until_stdin_eof(mut f: F) { } } +pub enum Event { + /// A message from stdin, or `None` at EOF. + Read(Result, String>), + /// A background job finished. + JobFinished, +} + +/// Read stdin on its own thread, so background threads can also send events. +pub fn stdin_events() -> (mpsc::Sender, mpsc::Receiver) { + let (tx, rx) = mpsc::channel(); + let reader_tx = tx.clone(); + std::thread::spawn(move || { + let mut stdin = io::stdin().lock(); + loop { + let read = read(&mut stdin).map_err(|error| error.to_string()); + let done = !matches!(read, Ok(Some(_))); + let _ = reader_tx.send(Event::Read(read)); + if done { + break; + } + } + }); + (tx, rx) +} + pub fn run_with_read(input: &mut R, mut f: F) { while let Some(message) = read(input).expect("JSON-RPC read error") { f(&message); diff --git a/crates/dm-langserver/src/main.rs b/crates/dm-langserver/src/main.rs index 60689145..b8e3f66b 100644 --- a/crates/dm-langserver/src/main.rs +++ b/crates/dm-langserver/src/main.rs @@ -26,7 +26,7 @@ use std::collections::VecDeque; use std::collections::hash_map::Entry; use std::path::PathBuf; use std::rc::Rc; -use std::sync::{Arc, Mutex}; +use std::sync::{Arc, Mutex, mpsc}; use clap::Parser; use dm::annotation::{Annotation, AnnotationTree}; @@ -48,6 +48,7 @@ mod document; mod extras; mod find_references; mod jrpc_io; +mod proc_heap; mod symbol_search; use extras::{QueryObjectTree, Reparse, SetTraceVsc, StartDebugger}; @@ -131,8 +132,17 @@ impl DmLangserverCli { match &self.command { Some(Command::Debugger(args)) => debugger::debugger_main(args), None => { - let mut engine = Engine::new(); - jrpc_io::run_until_stdin_eof(|message| engine.handle_input(message)); + let (events_tx, events) = jrpc_io::stdin_events(); + let mut engine = Engine::new(events_tx); + for event in events { + match event { + jrpc_io::Event::Read(read) => match read.expect("JSON-RPC read error") { + Some(message) => engine.handle_input(&message), + None => break, + }, + jrpc_io::Event::JobFinished => engine.release_proc_bodies(), + } + } engine.exit(0); }, } @@ -316,6 +326,8 @@ struct Engine { context: dm::Context, defines: Option, objtree: Arc, + /// Lets background jobs tell the main loop when they finish. + events_tx: mpsc::Sender, references_table: background::Background, annotations: HashMap)>, @@ -327,7 +339,7 @@ struct Engine { } impl Engine { - fn new() -> Self { + fn new(events_tx: mpsc::Sender) -> Self { Engine { docs: Default::default(), @@ -340,6 +352,7 @@ impl Engine { context: dm::Context::default(), defines: None, objtree: Default::default(), + events_tx, references_table: Default::default(), annotations: Default::default(), @@ -608,6 +621,7 @@ impl Engine { { let mut parser = dm::Parser::new(&self.context, &mut pp); parser.enable_procs(); + parser.set_proc_body_scope(proc_heap::enter, proc_heap::exit); let (fatal_errored_2, objtree) = parser.parse_object_tree_2(); fatal_errored = fatal_errored_2; self.objtree = Arc::new(objtree); @@ -628,6 +642,7 @@ impl Engine { // Background thread: prepare the Find All References database. let references_objtree = self.objtree.clone(); + let events_tx = self.events_tx.clone(); self.references_table.spawn(move || { let table = find_references::ReferencesTable::new(&references_objtree); let elapsed = start.elapsed(); @@ -637,7 +652,9 @@ impl Engine { elapsed.subsec_millis() ); print_thread_total(); + drop(references_objtree); collect_freed_memory(); + let _ = events_tx.send(jrpc_io::Event::JobFinished); table }); @@ -652,6 +669,7 @@ impl Engine { let root = self.root.clone(); let related_info = self.client_caps.related_info; let diagnostics_tracker = self.diagnostics_tracker.clone(); + let events_tx = self.events_tx.clone(); std::thread::spawn(move || { dreamchecker::run(&context, &objtree); let elapsed = start.elapsed(); @@ -672,8 +690,9 @@ impl Engine { diagnostics_tracker.lock().unwrap().send(map); issue_notification::(Default::default()); - drop(context); + drop((context, objtree)); collect_freed_memory(); + let _ = events_tx.send(jrpc_io::Event::JobFinished); }); } else { self.issue_notification::(Default::default()); @@ -713,10 +732,6 @@ impl Engine { ); } - /*if let Some(objtree) = Arc::get_mut(&mut self.objtree) { - objtree.drop_code(); - }*/ - // Print the total time. print_thread_total(); @@ -1282,6 +1297,16 @@ impl Engine { jrpc_io::write(&serde_json::to_string(&response).expect("response bad to_string")); } + /// Drop proc bodies once the background jobs are done with them. + fn release_proc_bodies(&mut self) { + if let Some(objtree) = Arc::get_mut(&mut self.objtree) { + objtree.drop_code(); + proc_heap::collect(); + // the heaps share segments, so the main heap needs collecting too + collect_freed_memory(); + } + } + fn handle_call(&mut self, call: Call) -> Option { match call { Call::Invalid { id } => Some(Output::invalid_request(id, JSONRPC_VERSION)), diff --git a/crates/dm-langserver/src/proc_heap.rs b/crates/dm-langserver/src/proc_heap.rs new file mode 100644 index 00000000..4f483f3b --- /dev/null +++ b/crates/dm-langserver/src/proc_heap.rs @@ -0,0 +1,29 @@ +//! Proc bodies are only needed until dreamchecker and Find All References are done with them. Keeping them in their own heap means their memory can actually be given back once they're dropped. +#![allow(unsafe_code)] + +use libmimalloc_sys::{ + mi_heap_collect, mi_heap_get_backing, mi_heap_new, mi_heap_set_default, mi_heap_t, +}; + +thread_local! { + // SAFETY: mi_heap_new is a ffi call, no safety contract + static PROC_HEAP: *mut mi_heap_t = unsafe { mi_heap_new() }; +} + +/// Start allocating in the proc heap. +pub fn enter() { + // SAFETY: the heap belongs to this thread + PROC_HEAP.with(|&heap| unsafe { mi_heap_set_default(heap) }); +} + +/// Go back to allocating in this thread's normal heap. +pub fn exit() { + // SAFETY: the backing heap belongs to this thread + unsafe { mi_heap_set_default(mi_heap_get_backing()) }; +} + +/// Give the memory of dropped proc bodies back to the OS. +pub fn collect() { + // SAFETY: the heap belongs to this thread + PROC_HEAP.with(|&heap| unsafe { mi_heap_collect(heap, true) }); +} diff --git a/crates/dreammaker/src/parser.rs b/crates/dreammaker/src/parser.rs index cbf7c1b4..00e9ec9f 100644 --- a/crates/dreammaker/src/parser.rs +++ b/crates/dreammaker/src/parser.rs @@ -336,6 +336,8 @@ pub struct Parser<'ctx, 'an, 'inp> { procs: bool, procs_bad: u64, procs_good: u64, + proc_body_enter: fn(), + proc_body_exit: fn(), } impl<'ctx, 'an, 'inp> HasLocation for Parser<'ctx, 'an, 'inp> { @@ -385,6 +387,8 @@ impl<'ctx, 'an, 'inp> Parser<'ctx, 'an, 'inp> { procs: false, procs_bad: 0, procs_good: 0, + proc_body_enter: || {}, + proc_body_exit: || {}, } } @@ -392,6 +396,12 @@ impl<'ctx, 'an, 'inp> Parser<'ctx, 'an, 'inp> { self.procs = true; } + /// Run `enter` before and `exit` after parsing each proc body so we can keep proc bodies in their own heap. + pub fn set_proc_body_scope(&mut self, enter: fn(), exit: fn()) { + self.proc_body_enter = enter; + self.proc_body_exit = exit; + } + pub fn annotate_to(&mut self, annotations: &'an mut AnnotationTree) { self.annotations = Some(annotations); self.procs = true; @@ -1270,6 +1280,8 @@ impl<'ctx, 'an, 'inp> Parser<'ctx, 'an, 'inp> { } let code = if self.procs { + (self.proc_body_enter)(); + let outer_errors = std::mem::take(&mut *self.context.errors_mut()); let result = { let mut subparser: Parser<'ctx, '_, '_> = Parser::new(self.context, body_tt); if let Some(a) = self.annotations.as_mut() { @@ -1283,13 +1295,20 @@ impl<'ctx, 'an, 'inp> Parser<'ctx, 'an, 'inp> { } else { self.procs_bad += 1; } - match result { + let code = match result { Err(err) => { self.context.register_error(err); None }, Ok(code) => Some(code), - } + }; + (self.proc_body_exit)(); + // errors outlive the proc body, so copy them out of its scope + let body_errors = std::mem::replace(&mut *self.context.errors_mut(), outer_errors); + self.context + .errors_mut() + .extend(body_errors.iter().cloned()); + code } else { None }; diff --git a/crates/dreammaker/tests/ast_tests.rs b/crates/dreammaker/tests/ast_tests.rs index 45a26df6..d3dde4ab 100644 --- a/crates/dreammaker/tests/ast_tests.rs +++ b/crates/dreammaker/tests/ast_tests.rs @@ -374,3 +374,23 @@ fn for_key_value_with_as() { }, ); } + +#[test] +fn proc_body_errors_keep_their_order() { + // one diagnostic per line, from both proc parameters and proc bodies, including a body that fails to parse + with_code( + " +/proc/one(var/a) + var/tmp/x = 1 + src. +/proc/two(var/b) + var/z = 1 + +/proc/three(var/c) + var/tmp/y = 1 + ", + |context, _| { + let lines: Vec = context.errors().iter().map(|e| e.location().line).collect(); + assert_eq!(lines, [1, 2, 3, 4, 5, 6, 7]); + }, + ); +}