holm/nushell
1
1
Fork 0
forked from extern/nushell
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Make EngineState clone cheaper with Arc on all of the heavy objects (#12229)

Browse Source 
# Description
This makes many of the larger objects in `EngineState` into `Arc`, and
uses `Arc::make_mut` to do clone-on-write if the reference is not
unique. This is generally very cheap, giving us the best of both worlds
- allowing us to mutate without cloning if we have an exclusive
reference, and cloning if we don't.

This started as more of a curiosity for me after remembering that
`Arc::make_mut` exists and can make using `Arc` for mostly immutable
data that sometimes needs to be changed very convenient, and also after
hearing someone complain about memory usage on Discord - this is a
somewhat significant win for that.

The exact objects that were wrapped in `Arc`:

- `files`, `file_contents` - the strings and byte buffers
- `decls` - the whole `Vec`, but mostly to avoid lots of individual
`malloc()` calls on Clone rather than for memory usage
- `blocks` - the blocks themselves, rather than the outer Vec
- `modules` - the modules themselves, rather than the outer Vec
- `env_vars`, `previous_env_vars` - the entire maps
- `config`

The changes required were relatively minimal, but this is a breaking API
change. In particular, blocks are added as Arcs, to allow the parser
cache functionality to work.

With my normal nu config, running on Linux, this saves me about 15 MiB
of process memory usage when running interactively (65 MiB → 50 MiB).

This also makes quick command executions cheaper, particularly since
every REPL loop now involves a clone of the engine state so that we can
recover from a panic. It also reduces memory usage where engine state
needs to be cloned and sent to another thread or kept within an
iterator.

# User-Facing Changes
Shouldn't be any, since it's all internal stuff, but it does change some
public interfaces so it's a breaking change
This commit is contained in:
Devyn Cairns
2024-03-19 11:07:00 -07:00
committed by GitHub
parent a29efe28f7
commit cf321ab510
21 changed files with 139 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
crates/nu-protocol/src/ast/expression.rs
View File

@ -1,3 +1,5 @@
use std::sync::Arc;
use serde::{Deserialize, Serialize};
use super::{Argument, Expr, ExternalArgument, RecordItem};
@ -325,7 +327,7 @@ impl Expression {
expr.replace_span(working_set, replaced, new_span);
}
Expr::Block(block_id) => {
let mut block = working_set.get_block(*block_id).clone();
let mut block = (**working_set.get_block(*block_id)).clone();
for pipeline in block.pipelines.iter_mut() {
for element in pipeline.elements.iter_mut() {
@ -333,10 +335,10 @@ impl Expression {
}
}
*block_id = working_set.add_block(block);
*block_id = working_set.add_block(Arc::new(block));
}
Expr::Closure(block_id) => {
let mut block = working_set.get_block(*block_id).clone();
let mut block = (**working_set.get_block(*block_id)).clone();
for pipeline in block.pipelines.iter_mut() {
for element in pipeline.elements.iter_mut() {
@ -344,7 +346,7 @@ impl Expression {
}
}
*block_id = working_set.add_block(block);
*block_id = working_set.add_block(Arc::new(block));
}
Expr::Binary(_) => {}
Expr::Bool(_) => {}
@ -429,7 +431,7 @@ impl Expression {
}
}
Expr::RowCondition(block_id) | Expr::Subexpression(block_id) => {
let mut block = working_set.get_block(*block_id).clone();
let mut block = (**working_set.get_block(*block_id)).clone();
for pipeline in block.pipelines.iter_mut() {
for element in pipeline.elements.iter_mut() {
@ -437,7 +439,7 @@ impl Expression {
}
}
*block_id = working_set.add_block(block);
*block_id = working_set.add_block(Arc::new(block));
}
Expr::Table(headers, cells) => {
for header in headers {

75
crates/nu-protocol/src/engine/engine_state.rs
View File

@ -64,23 +64,29 @@ impl Clone for IsDebugging {
/// will refer to the corresponding IDs rather than their definitions directly. At runtime, this means
/// less copying and smaller structures.
///
/// Many of the larger objects in this structure are stored within `Arc` to decrease the cost of
/// cloning `EngineState`. While `Arc`s are generally immutable, they can be modified using
/// `Arc::make_mut`, which automatically clones to a new allocation if there are other copies of
/// the `Arc` already in use, but will let us modify the `Arc` directly if we have the only
/// reference to it.
///
/// Note that the runtime stack is not part of this global state. Runtime stacks are handled differently,
/// but they also rely on using IDs rather than full definitions.
#[derive(Clone)]
pub struct EngineState {
files: Vec<(String, usize, usize)>,
file_contents: Vec<(Vec<u8>, usize, usize)>,
files: Vec<(Arc<String>, usize, usize)>,
file_contents: Vec<(Arc<Vec<u8>>, usize, usize)>,
pub(super) virtual_paths: Vec<(String, VirtualPath)>,
vars: Vec<Variable>,
decls: Vec<Box<dyn Command + 'static>>,
pub(super) blocks: Vec<Block>,
pub(super) modules: Vec<Module>,
decls: Arc<Vec<Box<dyn Command + 'static>>>,
pub(super) blocks: Vec<Arc<Block>>,
pub(super) modules: Vec<Arc<Module>>,
usage: Usage,
pub scope: ScopeFrame,
pub ctrlc: Option<Arc<AtomicBool>>,
pub env_vars: EnvVars,
pub previous_env_vars: HashMap<String, Value>,
pub config: Config,
pub env_vars: Arc<EnvVars>,
pub previous_env_vars: Arc<HashMap<String, Value>>,
pub config: Arc<Config>,
pub pipeline_externals_state: Arc<(AtomicU32, AtomicU32)>,
pub repl_state: Arc<Mutex<ReplState>>,
pub table_decl_id: Option<usize>,
@ -122,9 +128,11 @@ impl EngineState {
Variable::new(Span::new(0, 0), Type::Any, false),
Variable::new(Span::new(0, 0), Type::Any, false),
],
decls: vec![],
decls: Arc::new(vec![]),
blocks: vec![],
modules: vec![Module::new(DEFAULT_OVERLAY_NAME.as_bytes().to_vec())],
modules: vec![Arc::new(Module::new(
DEFAULT_OVERLAY_NAME.as_bytes().to_vec(),
))],
usage: Usage::new(),
// make sure we have some default overlay:
scope: ScopeFrame::with_empty_overlay(
@ -133,11 +141,13 @@ impl EngineState {
false,
),
ctrlc: None,
env_vars: [(DEFAULT_OVERLAY_NAME.to_string(), HashMap::new())]
.into_iter()
.collect(),
previous_env_vars: HashMap::new(),
config: Config::default(),
env_vars: Arc::new(
[(DEFAULT_OVERLAY_NAME.to_string(), HashMap::new())]
.into_iter()
.collect(),
),
previous_env_vars: Arc::new(HashMap::new()),
config: Arc::new(Config::default()),
pipeline_externals_state: Arc::new((AtomicU32::new(0), AtomicU32::new(0))),
repl_state: Arc::new(Mutex::new(ReplState {
buffer: "".to_string(),
@ -175,12 +185,16 @@ impl EngineState {
self.files.extend(delta.files);
self.file_contents.extend(delta.file_contents);
self.virtual_paths.extend(delta.virtual_paths);
self.decls.extend(delta.decls);
self.vars.extend(delta.vars);
self.blocks.extend(delta.blocks);
self.modules.extend(delta.modules);
self.usage.merge_with(delta.usage);
// Avoid potentially cloning the Arc if we aren't adding anything
if !delta.decls.is_empty() {
Arc::make_mut(&mut self.decls).extend(delta.decls);
}
let first = delta.scope.remove(0);
for (delta_name, delta_overlay) in first.clone().overlays {
@ -268,7 +282,7 @@ impl EngineState {
for mut scope in stack.env_vars.drain(..) {
for (overlay_name, mut env) in scope.drain() {
if let Some(env_vars) = self.env_vars.get_mut(&overlay_name) {
if let Some(env_vars) = Arc::make_mut(&mut self.env_vars).get_mut(&overlay_name) {
// Updating existing overlay
for (k, v) in env.drain() {
if k == "config" {
@ -276,7 +290,7 @@ impl EngineState {
// Instead, mutate a clone of it with into_config(), and put THAT in env_vars.
let mut new_record = v.clone();
let (config, error) = new_record.into_config(&self.config);
self.config = config;
self.config = Arc::new(config);
config_updated = true;
env_vars.insert(k, new_record);
if let Some(e) = error {
@ -288,7 +302,7 @@ impl EngineState {
}
} else {
// Pushing a new overlay
self.env_vars.insert(overlay_name, env);
Arc::make_mut(&mut self.env_vars).insert(overlay_name, env);
}
}
}
@ -422,10 +436,10 @@ impl EngineState {
pub fn add_env_var(&mut self, name: String, val: Value) {
let overlay_name = String::from_utf8_lossy(self.last_overlay_name(&[])).to_string();
if let Some(env_vars) = self.env_vars.get_mut(&overlay_name) {
if let Some(env_vars) = Arc::make_mut(&mut self.env_vars).get_mut(&overlay_name) {
env_vars.insert(name, val);
} else {
self.env_vars
Arc::make_mut(&mut self.env_vars)
.insert(overlay_name, [(name, val)].into_iter().collect());
}
}
@ -752,7 +766,7 @@ impl EngineState {
self.update_plugin_gc_configs(&conf.plugin_gc);
}
self.config = conf;
self.config = Arc::new(conf);
}
/// Fetch the configuration for a plugin
@ -867,7 +881,7 @@ impl EngineState {
.collect()
}
pub fn get_block(&self, block_id: BlockId) -> &Block {
pub fn get_block(&self, block_id: BlockId) -> &Arc<Block> {
self.blocks
.get(block_id)
.expect("internal error: missing block")
@ -878,7 +892,7 @@ impl EngineState {
/// Prefer to use [`.get_block()`] in most cases - `BlockId`s that don't exist are normally a
/// compiler error. This only exists to stop plugins from crashing the engine if they send us
/// something invalid.
pub fn try_get_block(&self, block_id: BlockId) -> Option<&Block> {
pub fn try_get_block(&self, block_id: BlockId) -> Option<&Arc<Block>> {
self.blocks.get(block_id)
}
@ -902,7 +916,7 @@ impl EngineState {
}
}
pub fn files(&self) -> impl Iterator<Item = &(String, usize, usize)> {
pub fn files(&self) -> impl Iterator<Item = &(Arc<String>, usize, usize)> {
self.files.iter()
}
@ -911,9 +925,10 @@ impl EngineState {
let next_span_end = next_span_start + contents.len();
self.file_contents
.push((contents, next_span_start, next_span_end));
.push((Arc::new(contents), next_span_start, next_span_end));
self.files.push((filename, next_span_start, next_span_end));
self.files
.push((Arc::new(filename), next_span_start, next_span_end));
self.num_files() - 1
}
@ -953,7 +968,7 @@ impl EngineState {
.unwrap_or_default()
}
pub fn get_file_contents(&self) -> &[(Vec<u8>, usize, usize)] {
pub fn get_file_contents(&self) -> &[(Arc<Vec<u8>>, usize, usize)] {
&self.file_contents
}
@ -1051,8 +1066,8 @@ mod engine_state_tests {
engine_state.merge_delta(delta)?;
assert_eq!(engine_state.num_files(), 2);
assert_eq!(&engine_state.files[0].0, "test.nu");
assert_eq!(&engine_state.files[1].0, "child.nu");
assert_eq!(&*engine_state.files[0].0, "test.nu");
assert_eq!(&*engine_state.files[1].0, "child.nu");
Ok(())
}

10
crates/nu-protocol/src/engine/state_delta.rs
View File

@ -12,13 +12,13 @@ use crate::RegisteredPlugin;
/// can be applied to the global state to update it to contain both previous state and the state held
/// within the delta.
pub struct StateDelta {
pub(super) files: Vec<(String, usize, usize)>,
pub(crate) file_contents: Vec<(Vec<u8>, usize, usize)>,
pub(super) files: Vec<(Arc<String>, usize, usize)>,
pub(crate) file_contents: Vec<(Arc<Vec<u8>>, usize, usize)>,
pub(super) virtual_paths: Vec<(String, VirtualPath)>,
pub(super) vars: Vec<Variable>, // indexed by VarId
pub(super) decls: Vec<Box<dyn Command>>, // indexed by DeclId
pub blocks: Vec<Block>, // indexed by BlockId
pub(super) modules: Vec<Module>, // indexed by ModuleId
pub blocks: Vec<Arc<Block>>, // indexed by BlockId
pub(super) modules: Vec<Arc<Module>>, // indexed by ModuleId
pub(super) usage: Usage,
pub scope: Vec<ScopeFrame>,
#[cfg(feature = "plugin")]
@ -131,7 +131,7 @@ impl StateDelta {
self.scope.pop();
}
pub fn get_file_contents(&self) -> &[(Vec<u8>, usize, usize)] {
pub fn get_file_contents(&self) -> &[(Arc<Vec<u8>>, usize, usize)] {
&self.file_contents
}
}

33
crates/nu-protocol/src/engine/state_working_set.rs
View File

@ -254,7 +254,7 @@ impl<'a> StateWorkingSet<'a> {
}
}
pub fn add_block(&mut self, block: Block) -> BlockId {
pub fn add_block(&mut self, block: Arc<Block>) -> BlockId {
self.delta.blocks.push(block);
self.num_blocks() - 1
@ -263,7 +263,7 @@ impl<'a> StateWorkingSet<'a> {
pub fn add_module(&mut self, name: &str, module: Module, comments: Vec<Span>) -> ModuleId {
let name = name.as_bytes().to_vec();
self.delta.modules.push(module);
self.delta.modules.push(Arc::new(module));
let module_id = self.num_modules() - 1;
if !comments.is_empty() {
@ -296,7 +296,7 @@ impl<'a> StateWorkingSet<'a> {
self.permanent_state.next_span_start()
}
pub fn files(&'a self) -> impl Iterator<Item = &(String, usize, usize)> {
pub fn files(&'a self) -> impl Iterator<Item = &(Arc<String>, usize, usize)> {
self.permanent_state.files().chain(self.delta.files.iter())
}
@ -320,7 +320,7 @@ impl<'a> StateWorkingSet<'a> {
pub fn add_file(&mut self, filename: String, contents: &[u8]) -> FileId {
// First, look for the file to see if we already have it
for (idx, (fname, file_start, file_end)) in self.files().enumerate() {
if fname == &filename {
if **fname == filename {
let prev_contents = self.get_span_contents(Span::new(*file_start, *file_end));
if prev_contents == contents {
return idx;
@ -331,13 +331,15 @@ impl<'a> StateWorkingSet<'a> {
let next_span_start = self.next_span_start();
let next_span_end = next_span_start + contents.len();
self.delta
.file_contents
.push((contents.to_vec(), next_span_start, next_span_end));
self.delta.file_contents.push((
Arc::new(contents.to_vec()),
next_span_start,
next_span_end,
));
self.delta
.files
.push((filename, next_span_start, next_span_end));
.push((Arc::new(filename), next_span_start, next_span_end));
self.num_files() - 1
}
@ -353,7 +355,7 @@ impl<'a> StateWorkingSet<'a> {
let (file_id, ..) = self
.files()
.enumerate()
.find(|(_, (fname, _, _))| fname == filename)?;
.find(|(_, (fname, _, _))| **fname == filename)?;
Some(self.get_span_for_file(file_id))
}
@ -626,8 +628,8 @@ impl<'a> StateWorkingSet<'a> {
pub fn list_env(&self) -> Vec<String> {
let mut env_vars = vec![];
for env_var in self.permanent_state.env_vars.clone().into_iter() {
env_vars.push(env_var.0)
for env_var in self.permanent_state.env_vars.iter() {
env_vars.push(env_var.0.clone());
}
env_vars
@ -742,7 +744,7 @@ impl<'a> StateWorkingSet<'a> {
output
}
pub fn get_block(&self, block_id: BlockId) -> &Block {
pub fn get_block(&self, block_id: BlockId) -> &Arc<Block> {
let num_permanent_blocks = self.permanent_state.num_blocks();
if block_id < num_permanent_blocks {
self.permanent_state.get_block(block_id)
@ -774,6 +776,7 @@ impl<'a> StateWorkingSet<'a> {
self.delta
.blocks
.get_mut(block_id - num_permanent_blocks)
.map(Arc::make_mut)
.expect("internal error: missing block")
}
}
@ -957,7 +960,7 @@ impl<'a> StateWorkingSet<'a> {
build_usage(&comment_lines)
}
pub fn find_block_by_span(&self, span: Span) -> Option<Block> {
pub fn find_block_by_span(&self, span: Span) -> Option<Arc<Block>> {
for block in &self.delta.blocks {
if Some(span) == block.span {
return Some(block.clone());
@ -1063,7 +1066,7 @@ impl<'a> miette::SourceCode for &StateWorkingSet<'a> {
}
let data = span_contents.data();
if filename == "<cli>" {
if **filename == "<cli>" {
if debugging {
let success_cli = "Successfully read CLI span";
dbg!(success_cli, String::from_utf8_lossy(data));
@ -1081,7 +1084,7 @@ impl<'a> miette::SourceCode for &StateWorkingSet<'a> {
dbg!(success_file);
}
return Ok(Box::new(miette::MietteSpanContents::new_named(
filename.clone(),
(**filename).clone(),
data,
retranslated,
span_contents.line(),