Make drop notification timing for plugin custom values more consistent (#12341)

# Description
This keeps plugin custom values that have requested drop notification
around during the lifetime of a plugin call / stream by sending them to
a channel that gets persisted during the lifetime of the call.

Before this change, it was very likely that the drop notification would
be sent before the plugin ever had a chance to handle the value it
received.

Tests have been added to make sure this works - see the `custom_values`
plugin.

cc @ayax79 

# User-Facing Changes
This is basically just a bugfix, just a slightly big one.

However, I did add an `as_mut_any()` function for custom values, to
avoid having to clone them. This is a breaking change.

crates/nu_plugin_custom_values/src/cool_custom_value.rs

@@ -144,4 +144,8 @@ impl CustomValue for CoolCustomValue {
     fn as_any(&self) -> &dyn std::any::Any {
         self
     }
+
+    fn as_mut_any(&mut self) -> &mut dyn std::any::Any {
+        self
+    }
 }

crates/nu_plugin_custom_values/src/drop_check.rs

@@ -47,6 +47,10 @@ impl CustomValue for DropCheckValue {
         self
     }
 
+    fn as_mut_any(&mut self) -> &mut dyn std::any::Any {
+        self
+    }
+
     fn notify_plugin_on_drop(&self) -> bool {
         // This is what causes Nushell to let us know when the value is dropped
         true

crates/nu_plugin_custom_values/src/generate.rs

@@ -42,6 +42,6 @@ impl SimplePluginCommand for Generate {
 fn test_examples() -> Result<(), nu_protocol::ShellError> {
     use nu_plugin_test_support::PluginTest;
 
-    PluginTest::new("custom_values", crate::CustomValuePlugin.into())?
+    PluginTest::new("custom_values", CustomValuePlugin::new().into())?
         .test_command_examples(&Generate)
 }

crates/nu_plugin_custom_values/src/generate2.rs

@@ -66,6 +66,6 @@ impl SimplePluginCommand for Generate2 {
 fn test_examples() -> Result<(), nu_protocol::ShellError> {
     use nu_plugin_test_support::PluginTest;
 
-    PluginTest::new("custom_values", crate::CustomValuePlugin.into())?
+    PluginTest::new("custom_values", crate::CustomValuePlugin::new().into())?
         .test_command_examples(&Generate2)
 }

crates/nu_plugin_custom_values/src/handle_custom_value.rs

@@ -0,0 +1,42 @@
+use nu_protocol::{CustomValue, LabeledError, ShellError, Span, Value};
+use serde::{Deserialize, Serialize};
+
+/// References a stored handle within the plugin
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct HandleCustomValue(pub u64);
+
+impl HandleCustomValue {
+    pub fn into_value(self, span: Span) -> Value {
+        Value::custom(Box::new(self), span)
+    }
+}
+
+#[typetag::serde]
+impl CustomValue for HandleCustomValue {
+    fn clone_value(&self, span: Span) -> Value {
+        self.clone().into_value(span)
+    }
+
+    fn type_name(&self) -> String {
+        "HandleCustomValue".into()
+    }
+
+    fn to_base_value(&self, span: Span) -> Result<Value, ShellError> {
+        Err(LabeledError::new("Unsupported operation")
+            .with_label("can't call to_base_value() directly on this", span)
+            .with_help("HandleCustomValue uses custom_value_to_base_value() on the plugin instead")
+            .into())
+    }
+
+    fn notify_plugin_on_drop(&self) -> bool {
+        true
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    fn as_mut_any(&mut self) -> &mut dyn std::any::Any {
+        self
+    }
+}

crates/nu_plugin_custom_values/src/handle_get.rs

@@ -0,0 +1,61 @@
+use nu_plugin::{EngineInterface, EvaluatedCall, SimplePluginCommand};
+use nu_protocol::{LabeledError, ShellError, Signature, Type, Value};
+
+use crate::{handle_custom_value::HandleCustomValue, CustomValuePlugin};
+
+pub struct HandleGet;
+
+impl SimplePluginCommand for HandleGet {
+    type Plugin = CustomValuePlugin;
+
+    fn name(&self) -> &str {
+        "custom-value handle get"
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::build(self.name())
+            .input_output_type(Type::Custom("HandleCustomValue".into()), Type::Any)
+    }
+
+    fn usage(&self) -> &str {
+        "Get a value previously stored in a handle"
+    }
+
+    fn run(
+        &self,
+        plugin: &Self::Plugin,
+        _engine: &EngineInterface,
+        call: &EvaluatedCall,
+        input: &Value,
+    ) -> Result<Value, LabeledError> {
+        if let Some(handle) = input
+            .as_custom_value()?
+            .as_any()
+            .downcast_ref::<HandleCustomValue>()
+        {
+            // Find the handle
+            let value = plugin
+                .handles
+                .lock()
+                .map_err(|err| LabeledError::new(err.to_string()))?
+                .get(&handle.0)
+                .cloned();
+
+            if let Some(value) = value {
+                Ok(value)
+            } else {
+                Err(LabeledError::new("Handle expired")
+                    .with_label("this handle is no longer valid", input.span())
+                    .with_help("the plugin may have exited, or there was a bug"))
+            }
+        } else {
+            Err(ShellError::UnsupportedInput {
+                msg: "requires HandleCustomValue".into(),
+                input: format!("got {}", input.get_type()),
+                msg_span: call.head,
+                input_span: input.span(),
+            }
+            .into())
+        }
+    }
+}

crates/nu_plugin_custom_values/src/handle_make.rs

@@ -0,0 +1,47 @@
+use std::sync::atomic;
+
+use nu_plugin::{EngineInterface, EvaluatedCall, SimplePluginCommand};
+use nu_protocol::{LabeledError, Signature, Type, Value};
+
+use crate::{handle_custom_value::HandleCustomValue, CustomValuePlugin};
+
+pub struct HandleMake;
+
+impl SimplePluginCommand for HandleMake {
+    type Plugin = CustomValuePlugin;
+
+    fn name(&self) -> &str {
+        "custom-value handle make"
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::build(self.name())
+            .input_output_type(Type::Any, Type::Custom("HandleCustomValue".into()))
+    }
+
+    fn usage(&self) -> &str {
+        "Store a value in plugin memory and return a handle to it"
+    }
+
+    fn run(
+        &self,
+        plugin: &Self::Plugin,
+        _engine: &EngineInterface,
+        call: &EvaluatedCall,
+        input: &Value,
+    ) -> Result<Value, LabeledError> {
+        // Generate an id and store in the plugin.
+        let new_id = plugin.counter.fetch_add(1, atomic::Ordering::Relaxed);
+
+        plugin
+            .handles
+            .lock()
+            .map_err(|err| LabeledError::new(err.to_string()))?
+            .insert(new_id, input.clone());
+
+        Ok(Value::custom(
+            Box::new(HandleCustomValue(new_id)),
+            call.head,
+        ))
+    }
+}

crates/nu_plugin_custom_values/src/main.rs

@@ -1,22 +1,43 @@
+use std::{
+    collections::BTreeMap,
+    sync::{atomic::AtomicU64, Mutex},
+};
+
+use handle_custom_value::HandleCustomValue;
 use nu_plugin::{serve_plugin, EngineInterface, MsgPackSerializer, Plugin, PluginCommand};
 
 mod cool_custom_value;
+mod handle_custom_value;
 mod second_custom_value;
 
 mod drop_check;
 mod generate;
 mod generate2;
+mod handle_get;
+mod handle_make;
 mod update;
 mod update_arg;
 
 use drop_check::{DropCheck, DropCheckValue};
 use generate::Generate;
 use generate2::Generate2;
-use nu_protocol::{CustomValue, LabeledError};
+use handle_get::HandleGet;
+use handle_make::HandleMake;
+use nu_protocol::{CustomValue, LabeledError, Spanned, Value};
 use update::Update;
 use update_arg::UpdateArg;
 
-pub struct CustomValuePlugin;
+#[derive(Default)]
+pub struct CustomValuePlugin {
+    counter: AtomicU64,
+    handles: Mutex<BTreeMap<u64, Value>>,
+}
+
+impl CustomValuePlugin {
+    pub fn new() -> Self {
+        Self::default()
+    }
+}
 
 impl Plugin for CustomValuePlugin {
     fn commands(&self) -> Vec<Box<dyn PluginCommand<Plugin = Self>>> {
@@ -26,22 +47,54 @@ impl Plugin for CustomValuePlugin {
             Box::new(Update),
             Box::new(UpdateArg),
             Box::new(DropCheck),
+            Box::new(HandleGet),
+            Box::new(HandleMake),
         ]
     }
 
+    fn custom_value_to_base_value(
+        &self,
+        _engine: &EngineInterface,
+        custom_value: Spanned<Box<dyn CustomValue>>,
+    ) -> Result<Value, LabeledError> {
+        // HandleCustomValue depends on the plugin state to get.
+        if let Some(handle) = custom_value
+            .item
+            .as_any()
+            .downcast_ref::<HandleCustomValue>()
+        {
+            Ok(self
+                .handles
+                .lock()
+                .map_err(|err| LabeledError::new(err.to_string()))?
+                .get(&handle.0)
+                .cloned()
+                .unwrap_or_else(|| Value::nothing(custom_value.span)))
+        } else {
+            custom_value
+                .item
+                .to_base_value(custom_value.span)
+                .map_err(|err| err.into())
+        }
+    }
+
     fn custom_value_dropped(
         &self,
         _engine: &EngineInterface,
         custom_value: Box<dyn CustomValue>,
     ) -> Result<(), LabeledError> {
-        // This is how we implement our drop behavior for DropCheck.
+        // This is how we implement our drop behavior.
         if let Some(drop_check) = custom_value.as_any().downcast_ref::<DropCheckValue>() {
             drop_check.notify();
+        } else if let Some(handle) = custom_value.as_any().downcast_ref::<HandleCustomValue>() {
+            if let Ok(mut handles) = self.handles.lock() {
+                handles.remove(&handle.0);
+            }
         }
         Ok(())
     }
 }
 
 fn main() {
-    serve_plugin(&CustomValuePlugin, MsgPackSerializer {})
+    serve_plugin(&CustomValuePlugin::default(), MsgPackSerializer {})
 }

crates/nu_plugin_custom_values/src/second_custom_value.rs

@@ -74,4 +74,8 @@ impl CustomValue for SecondCustomValue {
     fn as_any(&self) -> &dyn std::any::Any {
         self
     }
+
+    fn as_mut_any(&mut self) -> &mut dyn std::any::Any {
+        self
+    }
 }

crates/nu_plugin_custom_values/src/update.rs

@@ -67,6 +67,6 @@ impl SimplePluginCommand for Update {
 fn test_examples() -> Result<(), nu_protocol::ShellError> {
     use nu_plugin_test_support::PluginTest;
 
-    PluginTest::new("custom_values", crate::CustomValuePlugin.into())?
+    PluginTest::new("custom_values", crate::CustomValuePlugin::new().into())?
         .test_command_examples(&Update)
 }