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,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();

crates/nu_plugin_polars/src/cache/mod.rs

@@ -69,10 +69,12 @@ impl Cache {
             None
         };
 
+        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<Box<dyn PluginCommand<Plugin = PolarsPlugi
         Box::new(get::CacheGet),
     ]
 }
+
+#[cfg(test)]
+mod test {
+    use std::{cell::RefCell, rc::Rc};
+
+    use super::*;
+
+    struct MockEngineWrapper {
+        gc_enabled: Rc<RefCell<bool>>,
+    }
+
+    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<String> {
+            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());
+    }
+}

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())
+            }
         }
     }
 }

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<dyn CustomValue>,
     ) -> 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);