From 8316a1597e12651d86ab85313269dd4fad595fd5 Mon Sep 17 00:00:00 2001 From: Jack Wright <56345+ayax79@users.noreply.github.com> Date: Wed, 3 Jul 2024 04:44:26 -0700 Subject: [PATCH] Polars: Check to see if the cache is empty before enabling GC. More logging (#13286) There was a bug where anytime the plugin cache remove was called, the plugin gc was turned back on. This probably happened when I added the reference counter logic. --- crates/nu_plugin_polars/src/cache/list.rs | 16 +++- crates/nu_plugin_polars/src/cache/mod.rs | 90 ++++++++++++++++++- .../src/dataframe/values/mod.rs | 8 ++ crates/nu_plugin_polars/src/lib.rs | 3 + 4 files changed, 112 insertions(+), 5 deletions(-) diff --git a/crates/nu_plugin_polars/src/cache/list.rs b/crates/nu_plugin_polars/src/cache/list.rs index da434901e4..88fc1769ba 100644 --- a/crates/nu_plugin_polars/src/cache/list.rs +++ b/crates/nu_plugin_polars/src/cache/list.rs @@ -16,7 +16,7 @@ impl PluginCommand for ListDF { } fn usage(&self) -> &str { - "Lists stored dataframes." + "Lists stored polars objects." } fn signature(&self) -> Signature { @@ -119,6 +119,20 @@ impl PluginCommand for ListDF { }, call.head, ))), + PolarsPluginObject::NuPolarsTestData(_, _) => Ok(Some(Value::record( + record! { + "key" => Value::string(key.to_string(), call.head), + "columns" => Value::nothing(call.head), + "rows" => Value::nothing(call.head), + "type" => Value::string("When", call.head), + "estimated_size" => Value::nothing(call.head), + "span_contents" => Value::string(span_contents.to_string(), call.head), + "span_start" => Value::int(call.head.start as i64, call.head), + "span_end" => Value::int(call.head.end as i64, call.head), + "reference_count" => Value::int(value.reference_count as i64, call.head), + }, + call.head, + ))) } })?; let vals = vals.into_iter().flatten().collect(); diff --git a/crates/nu_plugin_polars/src/cache/mod.rs b/crates/nu_plugin_polars/src/cache/mod.rs index a7e566d327..0fbf49f46f 100644 --- a/crates/nu_plugin_polars/src/cache/mod.rs +++ b/crates/nu_plugin_polars/src/cache/mod.rs @@ -69,10 +69,12 @@ impl Cache { None }; - // Once there are no more entries in the cache - // we can turn plugin gc back on - debug!("PolarsPlugin: Cache is empty enabling GC"); - engine.set_gc_disabled(false).map_err(LabeledError::from)?; + if lock.is_empty() { + // Once there are no more entries in the cache + // we can turn plugin gc back on + debug!("PolarsPlugin: Cache is empty enabling GC"); + engine.set_gc_disabled(false).map_err(LabeledError::from)?; + } drop(lock); Ok(removed) } @@ -170,3 +172,83 @@ pub(crate) fn cache_commands() -> Vec>, + } + + impl MockEngineWrapper { + fn new(gc_enabled: bool) -> Self { + Self { + gc_enabled: Rc::new(RefCell::new(gc_enabled)), + } + } + + fn gc_enabled(&self) -> bool { + *self.gc_enabled.borrow() + } + } + + impl EngineWrapper for &MockEngineWrapper { + fn get_env_var(&self, _key: &str) -> Option { + todo!() + } + + fn use_color(&self) -> bool { + todo!() + } + + fn set_gc_disabled(&self, disabled: bool) -> Result<(), ShellError> { + let _ = self.gc_enabled.replace(!disabled); + Ok(()) + } + } + + #[test] + pub fn test_remove_plugin_cache_enable() { + let mock_engine = MockEngineWrapper::new(false); + + let cache = Cache::default(); + let mut lock = cache.cache.lock().expect("should be able to acquire lock"); + let key0 = Uuid::new_v4(); + lock.insert( + key0, + CacheValue { + uuid: Uuid::new_v4(), + value: PolarsPluginObject::NuPolarsTestData(Uuid::new_v4(), "object_0".into()), + created: Local::now().into(), + span: Span::unknown(), + reference_count: 1, + }, + ); + + let key1 = Uuid::new_v4(); + lock.insert( + key1, + CacheValue { + uuid: Uuid::new_v4(), + value: PolarsPluginObject::NuPolarsTestData(Uuid::new_v4(), "object_1".into()), + created: Local::now().into(), + span: Span::unknown(), + reference_count: 1, + }, + ); + drop(lock); + + let _ = cache + .remove(&mock_engine, &key0, false) + .expect("should be able to remove key0"); + assert!(!mock_engine.gc_enabled()); + + let _ = cache + .remove(&mock_engine, &key1, false) + .expect("should be able to remove key1"); + assert!(mock_engine.gc_enabled()); + } +} diff --git a/crates/nu_plugin_polars/src/dataframe/values/mod.rs b/crates/nu_plugin_polars/src/dataframe/values/mod.rs index d52ee25203..ec59304eab 100644 --- a/crates/nu_plugin_polars/src/dataframe/values/mod.rs +++ b/crates/nu_plugin_polars/src/dataframe/values/mod.rs @@ -27,6 +27,7 @@ pub enum PolarsPluginType { NuExpression, NuLazyGroupBy, NuWhen, + NuPolarsTestData, } impl fmt::Display for PolarsPluginType { @@ -37,6 +38,7 @@ impl fmt::Display for PolarsPluginType { Self::NuExpression => write!(f, "NuExpression"), Self::NuLazyGroupBy => write!(f, "NuLazyGroupBy"), Self::NuWhen => write!(f, "NuWhen"), + Self::NuPolarsTestData => write!(f, "NuPolarsTestData"), } } } @@ -48,6 +50,7 @@ pub enum PolarsPluginObject { NuExpression(NuExpression), NuLazyGroupBy(NuLazyGroupBy), NuWhen(NuWhen), + NuPolarsTestData(Uuid, String), } impl PolarsPluginObject { @@ -95,6 +98,7 @@ impl PolarsPluginObject { Self::NuExpression(_) => PolarsPluginType::NuExpression, Self::NuLazyGroupBy(_) => PolarsPluginType::NuLazyGroupBy, Self::NuWhen(_) => PolarsPluginType::NuWhen, + Self::NuPolarsTestData(_, _) => PolarsPluginType::NuPolarsTestData, } } @@ -105,6 +109,7 @@ impl PolarsPluginObject { PolarsPluginObject::NuExpression(e) => e.id, PolarsPluginObject::NuLazyGroupBy(lg) => lg.id, PolarsPluginObject::NuWhen(w) => w.id, + PolarsPluginObject::NuPolarsTestData(id, _) => *id, } } @@ -115,6 +120,9 @@ impl PolarsPluginObject { PolarsPluginObject::NuExpression(e) => e.into_value(span), PolarsPluginObject::NuLazyGroupBy(lg) => lg.into_value(span), PolarsPluginObject::NuWhen(w) => w.into_value(span), + PolarsPluginObject::NuPolarsTestData(id, s) => { + Value::string(format!("{id}:{s}"), Span::test_data()) + } } } } diff --git a/crates/nu_plugin_polars/src/lib.rs b/crates/nu_plugin_polars/src/lib.rs index b5868e13d6..e4d627e4db 100644 --- a/crates/nu_plugin_polars/src/lib.rs +++ b/crates/nu_plugin_polars/src/lib.rs @@ -3,6 +3,7 @@ use std::cmp::Ordering; use cache::cache_commands; pub use cache::{Cache, Cacheable}; use dataframe::{stub::PolarsCmd, values::CustomValueType}; +use log::debug; use nu_plugin::{EngineInterface, Plugin, PluginCommand}; mod cache; @@ -41,6 +42,7 @@ impl EngineWrapper for &EngineInterface { } fn set_gc_disabled(&self, disabled: bool) -> Result<(), ShellError> { + debug!("set_gc_disabled called with {disabled}"); EngineInterface::set_gc_disabled(self, disabled) } } @@ -72,6 +74,7 @@ impl Plugin for PolarsPlugin { engine: &EngineInterface, custom_value: Box, ) -> Result<(), LabeledError> { + debug!("custom_value_dropped called {:?}", custom_value); if !self.disable_cache_drop { let id = CustomValueType::try_from_custom_value(custom_value)?.id(); let _ = self.cache.remove(engine, &id, false);