feat: make ctrlc available to plugins (#13181)

# Description

This PR adds a new method to `EngineInterface`: `register_ctrlc_handler`
which takes a closure to run when the plugin's driving engine receives a
ctrlc-signal. It also adds a mirror of the `signals` attribute from the
main shell `EngineState`.

This is an example of how a plugin which makes a long poll http request
can end the request on ctrlc:
https://github.com/cablehead/nu_plugin_http/blob/main/src/commands/request.rs#L68-L77

To facilitate the feature, a new attribute has been added to
`EngineState`: `ctrlc_handlers`. This is a Vec of closures that will be
run when the engine's process receives a ctrlc signal.

When plugins are added to an `engine_state` during a `merge_delta`, the
engine passes the ctrlc_handlers to the plugin's
`.configure_ctrlc_handler` method, which gives the plugin a chance to
register a handler that sends a ctrlc packet through the
`PluginInterface`, if an instance of the plugin is currently running.

On the plugin side: `EngineInterface` also has a ctrlc_handlers Vec of
closures. Plugin calls can use `register_ctrlc_handler` to register a
closure that will be called in the plugin process when the
PluginInput::Ctrlc command is received.

For future reference these are some alternate places that were
investigated for tying the ctrlc trigger to transmitting a Ctrlc packet
through the `PluginInterface`:

- Directly from `src/signals.rs`: the handler there would need a
reference to the Vec<Arc<RegisteredPlugins>>, which would require us to
wrap the plugins in a Mutex, which we don't want to do.

- have `PersistentPlugin.get_plugin` pass down the engine's
CtrlcHandlers to .get and then to .spawn (if the plugin isn't already
running). Once we have CtrlcHandlers in spawn, we can register a handler
to write directly to PluginInterface. We don't want to double down on
passing engine_state to spawn this way though, as it's unpredictable
because it would depend on whether the plugin has already been spawned
or not.

- pass `ctrlc_handlers` to PersistentPlugin::new so it can store it on
itself so it's available to spawn.

- in `PersistentPlugin.spawn`, create a handler that sends to a clone of
the GC event loop's tx. this has the same issues with regards to how to
get CtrlcHandlers to the spawn method, and is more complicated than a
handler that writes directly to PluginInterface

# User-Facing Changes

No breaking changes

---------

Co-authored-by: Ian Manske <ian.manske@pm.me>
This commit is contained in:
Andy Gayton 2024-07-30 09:29:18 -04:00 committed by GitHub
parent e3f78b8793
commit 7b82c6b482
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
17 changed files with 310 additions and 21 deletions

View File

@ -1,7 +1,10 @@
//! Implements the stream multiplexing interface for both the plugin side and the engine side.
use nu_plugin_protocol::{ByteStreamInfo, ListStreamInfo, PipelineDataHeader, StreamMessage};
use nu_protocol::{ByteStream, IntoSpanned, ListStream, PipelineData, Reader, ShellError, Signals};
use nu_protocol::{
engine::Sequence, ByteStream, IntoSpanned, ListStream, PipelineData, Reader, ShellError,
Signals,
};
use std::{
io::{Read, Write},
sync::Mutex,
@ -10,7 +13,7 @@ use std::{
pub mod stream;
use crate::{util::Sequence, Encoder};
use crate::Encoder;
use self::stream::{StreamManager, StreamManagerHandle, StreamWriter, WriteStreamMessage};

View File

@ -1,5 +1,3 @@
use crate::util::Sequence;
use super::{
stream::{StreamManager, StreamManagerHandle},
test_util::TestCase,
@ -10,8 +8,8 @@ use nu_plugin_protocol::{
StreamMessage,
};
use nu_protocol::{
ByteStream, ByteStreamSource, ByteStreamType, DataSource, ListStream, PipelineData,
PipelineMetadata, ShellError, Signals, Span, Value,
engine::Sequence, ByteStream, ByteStreamSource, ByteStreamType, DataSource, ListStream,
PipelineData, PipelineMetadata, ShellError, Signals, Span, Value,
};
use std::{path::Path, sync::Arc};

View File

@ -1,7 +1,5 @@
mod sequence;
mod waitable;
mod with_custom_values_in;
pub use sequence::Sequence;
pub use waitable::*;
pub use with_custom_values_in::with_custom_values_in;

View File

@ -1,7 +1,7 @@
//! Interface used by the engine to communicate with the plugin.
use nu_plugin_core::{
util::{with_custom_values_in, Sequence, Waitable, WaitableMut},
util::{with_custom_values_in, Waitable, WaitableMut},
Interface, InterfaceManager, PipelineDataWriter, PluginRead, PluginWrite, StreamManager,
StreamManagerHandle,
};
@ -11,8 +11,8 @@ use nu_plugin_protocol::{
PluginOutput, ProtocolInfo, StreamId, StreamMessage,
};
use nu_protocol::{
ast::Operator, CustomValue, IntoSpanned, PipelineData, PluginMetadata, PluginSignature,
ShellError, Signals, Span, Spanned, Value,
ast::Operator, engine::Sequence, CustomValue, IntoSpanned, PipelineData, PluginMetadata,
PluginSignature, ShellError, Signals, Span, Spanned, Value,
};
use nu_utils::SharedCow;
use std::{
@ -664,6 +664,12 @@ impl PluginInterface {
self.flush()
}
/// Send the plugin a ctrl-c signal.
pub fn ctrlc(&self) -> Result<(), ShellError> {
self.write(PluginInput::Ctrlc)?;
self.flush()
}
/// Write an [`EngineCallResponse`]. Writes the full stream contained in any [`PipelineData`]
/// before returning.
pub fn write_engine_call_response(

View File

@ -6,7 +6,7 @@ use crate::{
use super::{PluginInterface, PluginSource};
use nu_plugin_core::CommunicationMode;
use nu_protocol::{
engine::{EngineState, Stack},
engine::{ctrlc, EngineState, Stack},
PluginGcConfig, PluginIdentity, PluginMetadata, RegisteredPlugin, ShellError,
};
use std::{
@ -37,6 +37,8 @@ struct MutableState {
preferred_mode: Option<PreferredCommunicationMode>,
/// Garbage collector config
gc_config: PluginGcConfig,
/// RAII guard for this plugin's ctrl-c handler
ctrlc_guard: Option<ctrlc::Guard>,
}
#[derive(Debug, Clone, Copy)]
@ -64,6 +66,7 @@ impl PersistentPlugin {
metadata: None,
preferred_mode: None,
gc_config,
ctrlc_guard: None,
}),
}
}
@ -299,6 +302,34 @@ impl RegisteredPlugin for PersistentPlugin {
fn as_any(self: Arc<Self>) -> Arc<dyn std::any::Any + Send + Sync> {
self
}
fn configure_ctrlc_handler(
self: Arc<Self>,
handlers: &ctrlc::Handlers,
) -> Result<(), ShellError> {
let guard = {
// We take a weakref to the plugin so that we don't create a cycle to the
// RAII guard that will be stored on the plugin.
let plugin = Arc::downgrade(&self);
handlers.register(Box::new(move || {
// write a Ctrl-C packet through the PluginInterface if the plugin is alive and
// running
if let Some(plugin) = plugin.upgrade() {
if let Ok(mutable) = plugin.mutable.lock() {
if let Some(ref running) = mutable.running {
let _ = running.interface.ctrlc();
}
}
}
}))?
};
if let Ok(mut mutable) = self.mutable.lock() {
mutable.ctrlc_guard = Some(guard);
}
Ok(())
}
}
/// Anything that can produce a plugin interface.

View File

@ -208,6 +208,8 @@ pub enum PluginInput {
Drop(StreamId),
/// See [`StreamMessage::Ack`].
Ack(StreamId),
/// Signal a ctrlc event
Ctrlc,
}
impl TryFrom<PluginInput> for StreamMessage {

View File

@ -1,7 +1,7 @@
//! Interface used by the plugin to communicate with the engine.
use nu_plugin_core::{
util::{Sequence, Waitable, WaitableMut},
util::{Waitable, WaitableMut},
Interface, InterfaceManager, PipelineDataWriter, PluginRead, PluginWrite, StreamManager,
StreamManagerHandle,
};
@ -11,13 +11,14 @@ use nu_plugin_protocol::{
PluginOutput, ProtocolInfo,
};
use nu_protocol::{
engine::Closure, Config, DeclId, LabeledError, PipelineData, PluginMetadata, PluginSignature,
ShellError, Signals, Span, Spanned, Value,
engine::{ctrlc, Closure, Sequence},
Config, DeclId, LabeledError, PipelineData, PluginMetadata, PluginSignature, ShellError,
Signals, Span, Spanned, Value,
};
use nu_utils::SharedCow;
use std::{
collections::{btree_map, BTreeMap, HashMap},
sync::{mpsc, Arc},
sync::{atomic::AtomicBool, mpsc, Arc},
};
/// Plugin calls that are received by the [`EngineInterfaceManager`] for handling.
@ -63,6 +64,10 @@ struct EngineInterfaceState {
mpsc::Sender<(EngineCallId, mpsc::Sender<EngineCallResponse<PipelineData>>)>,
/// The synchronized output writer
writer: Box<dyn PluginWrite<PluginOutput>>,
// Mirror signals from `EngineState`
signals: Signals,
/// Registered Ctrl-C handlers
ctrlc_handlers: ctrlc::Handlers,
}
impl std::fmt::Debug for EngineInterfaceState {
@ -116,6 +121,8 @@ impl EngineInterfaceManager {
stream_id_sequence: Sequence::default(),
engine_call_subscription_sender: subscription_tx,
writer: Box::new(writer),
signals: Signals::new(Arc::new(AtomicBool::new(false))),
ctrlc_handlers: ctrlc::Handlers::new(),
}),
protocol_info_mut,
plugin_call_sender: Some(plug_tx),
@ -235,7 +242,6 @@ impl InterfaceManager for EngineInterfaceManager {
fn consume(&mut self, input: Self::Input) -> Result<(), ShellError> {
log::trace!("from engine: {:?}", input);
match input {
PluginInput::Hello(info) => {
let info = Arc::new(info);
@ -331,6 +337,11 @@ impl InterfaceManager for EngineInterfaceManager {
});
self.send_engine_call_response(id, response)
}
PluginInput::Ctrlc => {
self.state.signals.trigger();
self.state.ctrlc_handlers.run();
Ok(())
}
}
}
@ -510,6 +521,15 @@ impl EngineInterface {
self.state.writer.is_stdout()
}
/// Register a closure which will be called when the engine receives a Ctrl-C signal. Returns a
/// RAII guard that will keep the closure alive until it is dropped.
pub fn register_ctrlc_handler(
&self,
handler: ctrlc::Handler,
) -> Result<ctrlc::Guard, ShellError> {
self.state.ctrlc_handlers.register(handler)
}
/// Get the full shell configuration from the engine. As this is quite a large object, it is
/// provided on request only.
///
@ -959,6 +979,10 @@ impl EngineInterface {
self.write(PluginOutput::CallResponse(self.context()?, response))?;
self.flush()
}
pub fn signals(&self) -> &Signals {
&self.state.signals
}
}
impl Interface for EngineInterface {

View File

@ -0,0 +1,139 @@
use std::fmt::Debug;
use std::sync::{Arc, Mutex};
use crate::{engine::Sequence, ShellError};
/// Handler is a closure that can be sent across threads and shared.
pub type Handler = Box<dyn Fn() + Send + Sync>;
/// Manages a collection of handlers.
#[derive(Clone)]
pub struct Handlers {
/// List of handler tuples containing an ID and the handler itself.
handlers: Arc<Mutex<Vec<(usize, Handler)>>>,
/// Sequence generator for unique IDs.
next_id: Arc<Sequence>,
}
impl Debug for Handlers {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Handlers")
.field("next_id", &self.next_id)
.finish()
}
}
/// Guard that unregisters a handler when dropped.
#[derive(Clone)]
pub struct Guard {
/// Unique ID of the handler.
id: usize,
/// Reference to the handlers list.
handlers: Arc<Mutex<Vec<(usize, Handler)>>>,
}
impl Drop for Guard {
/// Drops the `Guard`, removing the associated handler from the list.
fn drop(&mut self) {
if let Ok(mut handlers) = self.handlers.lock() {
handlers.retain(|(id, _)| *id != self.id);
}
}
}
impl Debug for Guard {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Guard").field("id", &self.id).finish()
}
}
impl Handlers {
pub fn new() -> Handlers {
let handlers = Arc::new(Mutex::new(vec![]));
let next_id = Arc::new(Sequence::default());
Handlers { handlers, next_id }
}
/// Registers a new handler and returns an RAII guard which will unregister the handler when
/// dropped.
pub fn register(&self, handler: Handler) -> Result<Guard, ShellError> {
let id = self.next_id.next()?;
if let Ok(mut handlers) = self.handlers.lock() {
handlers.push((id, handler));
}
Ok(Guard {
id,
handlers: Arc::clone(&self.handlers),
})
}
/// Runs all registered handlers.
pub fn run(&self) {
if let Ok(handlers) = self.handlers.lock() {
for (_, handler) in handlers.iter() {
handler();
}
}
}
}
impl Default for Handlers {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicBool, Ordering};
#[test]
/// Tests registering and running multiple handlers.
fn test_multiple_handlers() {
let handlers = Handlers::new();
let called1 = Arc::new(AtomicBool::new(false));
let called2 = Arc::new(AtomicBool::new(false));
let called1_clone = Arc::clone(&called1);
let called2_clone = Arc::clone(&called2);
let _guard1 = handlers.register(Box::new(move || {
called1_clone.store(true, Ordering::SeqCst);
}));
let _guard2 = handlers.register(Box::new(move || {
called2_clone.store(true, Ordering::SeqCst);
}));
handlers.run();
assert!(called1.load(Ordering::SeqCst));
assert!(called2.load(Ordering::SeqCst));
}
#[test]
/// Tests the dropping of a guard and ensuring the handler is unregistered.
fn test_guard_drop() {
let handlers = Handlers::new();
let called = Arc::new(AtomicBool::new(false));
let called_clone = Arc::clone(&called);
let guard = handlers.register(Box::new(move || {
called_clone.store(true, Ordering::Relaxed);
}));
// Ensure the handler is registered
assert_eq!(handlers.handlers.lock().unwrap().len(), 1);
drop(guard);
// Ensure the handler is removed after dropping the guard
assert_eq!(handlers.handlers.lock().unwrap().len(), 0);
handlers.run();
// Ensure the handler is not called after being dropped
assert!(!called.load(Ordering::Relaxed));
}
}

View File

@ -2,6 +2,7 @@ use crate::{
ast::Block,
debugger::{Debugger, NoopDebugger},
engine::{
ctrlc,
usage::{build_usage, Usage},
CachedFile, Command, CommandType, EnvVars, OverlayFrame, ScopeFrame, Stack, StateDelta,
Variable, Visibility, DEFAULT_OVERLAY_NAME,
@ -85,6 +86,7 @@ pub struct EngineState {
pub spans: Vec<Span>,
usage: Usage,
pub scope: ScopeFrame,
pub ctrlc_handlers: Option<ctrlc::Handlers>,
signals: Signals,
pub env_vars: Arc<EnvVars>,
pub previous_env_vars: Arc<HashMap<String, Value>>,
@ -145,6 +147,7 @@ impl EngineState {
0,
false,
),
ctrlc_handlers: None,
signals: Signals::empty(),
env_vars: Arc::new(
[(DEFAULT_OVERLAY_NAME.to_string(), HashMap::new())]
@ -268,8 +271,13 @@ impl EngineState {
#[cfg(feature = "plugin")]
if !delta.plugins.is_empty() {
// Replace plugins that overlap in identity.
for plugin in std::mem::take(&mut delta.plugins) {
// Connect plugins to the ctrlc handlers
if let Some(handlers) = &self.ctrlc_handlers {
plugin.clone().configure_ctrlc_handler(handlers)?;
}
// Replace plugins that overlap in identity.
if let Some(existing) = self
.plugins
.iter_mut()

View File

@ -9,6 +9,7 @@ mod engine_state;
mod error_handler;
mod overlay;
mod pattern_match;
mod sequence;
mod stack;
mod stack_out_dest;
mod state_delta;
@ -27,8 +28,11 @@ pub use engine_state::*;
pub use error_handler::*;
pub use overlay::*;
pub use pattern_match::*;
pub use sequence::*;
pub use stack::*;
pub use stack_out_dest::*;
pub use state_delta::*;
pub use state_working_set::*;
pub use variable::*;
pub mod ctrlc;

View File

@ -1,4 +1,4 @@
use nu_protocol::ShellError;
use crate::ShellError;
use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};
/// Implements an atomically incrementing sequential series of numbers

View File

@ -56,6 +56,13 @@ impl Signals {
}
}
/// Triggers an interrupt.
pub fn trigger(&self) {
if let Some(signals) = &self.signals {
signals.store(true, Ordering::Relaxed);
}
}
/// Returns whether an interrupt has been triggered.
#[inline]
pub fn interrupted(&self) -> bool {

View File

@ -1,6 +1,6 @@
use std::{any::Any, sync::Arc};
use crate::{PluginGcConfig, PluginIdentity, PluginMetadata, ShellError};
use crate::{engine::ctrlc, PluginGcConfig, PluginIdentity, PluginMetadata, ShellError};
/// Trait for plugins registered in the [`EngineState`](crate::engine::EngineState).
pub trait RegisteredPlugin: Send + Sync {
@ -34,4 +34,12 @@ pub trait RegisteredPlugin: Send + Sync {
/// This is necessary in order to allow `nu_plugin` to handle the implementation details of
/// plugins.
fn as_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync>;
/// Give this plugin a chance to register for Ctrl-C signals.
fn configure_ctrlc_handler(
self: Arc<Self>,
_handler: &ctrlc::Handlers,
) -> Result<(), ShellError> {
Ok(())
}
}

View File

@ -0,0 +1,50 @@
use std::sync::mpsc;
use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
use nu_protocol::{Category, LabeledError, PipelineData, Signature};
use crate::ExamplePlugin;
/// `example ctrlc`
pub struct Ctrlc;
impl PluginCommand for Ctrlc {
type Plugin = ExamplePlugin;
fn name(&self) -> &str {
"example ctrlc"
}
fn usage(&self) -> &str {
"Example command that demonstrates registering a ctrl-c handler"
}
fn signature(&self) -> Signature {
Signature::build(self.name()).category(Category::Experimental)
}
fn search_terms(&self) -> Vec<&str> {
vec!["example"]
}
fn run(
&self,
_plugin: &ExamplePlugin,
engine: &EngineInterface,
_call: &EvaluatedCall,
_input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let (sender, receiver) = mpsc::channel::<()>();
let _guard = engine.register_ctrlc_handler(Box::new(move || {
let _ = sender.send(());
}));
eprintln!("interrupt status: {:?}", engine.signals().interrupted());
eprintln!("waiting for ctrl-c signal...");
receiver.recv().expect("handler went away");
eprintln!("interrupt status: {:?}", engine.signals().interrupted());
eprintln!("peace.");
Ok(PipelineData::Empty)
}
}

View File

@ -15,12 +15,14 @@ pub use two::Two;
// Engine interface demos
mod call_decl;
mod config;
mod ctrlc;
mod disable_gc;
mod env;
mod view_span;
pub use call_decl::CallDecl;
pub use config::Config;
pub use ctrlc::Ctrlc;
pub use disable_gc::DisableGc;
pub use env::Env;
pub use view_span::ViewSpan;

View File

@ -27,6 +27,7 @@ impl Plugin for ExamplePlugin {
Box::new(Env),
Box::new(ViewSpan),
Box::new(DisableGc),
Box::new(Ctrlc),
Box::new(CallDecl),
// Stream demos
Box::new(CollectBytes),

View File

@ -1,4 +1,7 @@
use nu_protocol::{engine::EngineState, Signals};
use nu_protocol::{
engine::{ctrlc::Handlers, EngineState},
Signals,
};
use std::sync::{
atomic::{AtomicBool, Ordering},
Arc,
@ -7,8 +10,13 @@ use std::sync::{
pub(crate) fn ctrlc_protection(engine_state: &mut EngineState) {
let interrupt = Arc::new(AtomicBool::new(false));
engine_state.set_signals(Signals::new(interrupt.clone()));
let ctrlc_handlers = Handlers::new();
engine_state.ctrlc_handlers = Some(ctrlc_handlers.clone());
ctrlc::set_handler(move || {
interrupt.store(true, Ordering::Relaxed);
ctrlc_handlers.run();
})
.expect("Error setting Ctrl-C handler");
}