dprint_core/plugins/process/

communicator.rs

use serde::de::DeserializeOwned;
use std::cell::RefCell;
use std::io::BufRead;
use std::io::ErrorKind;
use std::io::Read;
use std::io::Write;
use std::path::Path;
use std::path::PathBuf;
use std::process::Child;
use std::process::ChildStderr;
use std::process::Command;
use std::process::Stdio;
use std::rc::Rc;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::oneshot;
use tokio_util::sync::CancellationToken;

use super::PLUGIN_SCHEMA_VERSION;
use super::messages::CheckConfigUpdatesMessageBody;
use super::messages::CheckConfigUpdatesResponseBody;
use super::messages::FormatMessageBody;
use super::messages::HostFormatMessageBody;
use super::messages::MessageBody;
use super::messages::ProcessPluginMessage;
use super::messages::RegisterConfigMessageBody;
use super::messages::ResponseBody;
use crate::async_runtime::DropGuardAction;
use crate::async_runtime::LocalBoxFuture;
use crate::communication::AtomicFlag;
use crate::communication::IdGenerator;
use crate::communication::MessageReader;
use crate::communication::MessageWriter;
use crate::communication::RcIdStore;
use crate::communication::SingleThreadMessageWriter;
use crate::configuration::ConfigKeyMap;
use crate::configuration::ConfigurationDiagnostic;
use crate::configuration::GlobalConfiguration;
use crate::plugins::ConfigChange;
use crate::plugins::CriticalFormatError;
use crate::plugins::FileMatchingInfo;
use crate::plugins::FormatConfigId;
use crate::plugins::FormatError;
use crate::plugins::FormatRange;
use crate::plugins::FormatResult;
use crate::plugins::HostFormatRequest;
use crate::plugins::NullCancellationToken;
use crate::plugins::PluginInfo;
use crate::plugins::error_to_string;

type Result<T> = std::result::Result<T, FormatError>;

type DprintCancellationToken = Arc<dyn super::super::CancellationToken>;

/// The detailed reason a schema version negotiation step failed.
#[derive(Debug, thiserror::Error)]
enum SchemaVersionError {
  #[error("Failed asking for schema version: {0}")]
  Ask(std::io::Error),
  #[error("Failed flushing schema version request: {0}")]
  Flush(std::io::Error),
  #[error("Could not read success response: {0}")]
  ReadAcknowledgement(std::io::Error),
  #[error("Plugin response was unexpected ({0}).")]
  UnexpectedAcknowledgement(u32),
  #[error("Could not read schema version: {0}")]
  ReadVersion(std::io::Error),
}

/// Errors that can occur while communicating with a process plugin.
#[derive(Debug, thiserror::Error)]
enum CommunicatorError {
  #[error("Error starting {executable} with args [{args}]. {error}")]
  StartProcess { executable: String, args: String, error: std::io::Error },
  #[error("Failed plugin schema verification. This may indicate you are using an old version of the dprint CLI or plugin and should upgrade")]
  SchemaVerification(#[from] SchemaVersionError),
  #[error(
    "This plugin is too old to run in the dprint CLI and you will need to manually upgrade it (version was {actual}, but expected {expected}).\n\nUpgrade instructions: https://github.com/dprint/dprint/issues/731"
  )]
  PluginTooOld { actual: u32, expected: u32 },
  #[error("Your dprint CLI is too old to run this plugin (version was {actual}, but expected {expected}). Try running: dprint upgrade")]
  CliTooOld { actual: u32, expected: u32 },
  #[error("Error waiting on message ({message_id}). {error}")]
  WaitOnMessage {
    message_id: u32,
    error: tokio::sync::oneshot::error::RecvError,
  },
  #[error("Unexpected data channel for success response: {0}")]
  UnexpectedDataForSuccess(u32),
  #[error("Unexpected format channel for success response: {0}")]
  UnexpectedFormatForSuccess(u32),
  #[error("Unexpected success channel for data response: {0}")]
  UnexpectedSuccessForData(u32),
  #[error("Unexpected format channel for data response: {0}")]
  UnexpectedFormatForData(u32),
  #[error("Unexpected success channel for format response: {0}")]
  UnexpectedSuccessForFormat(u32),
  #[error("Unexpected data channel for format response: {0}")]
  UnexpectedDataForFormat(u32),
  #[error("Unknown message kind: {0}")]
  UnknownMessageKind(u32),
  #[error("Could not find host format callback for message id: {0}")]
  HostFormatCallbackNotFound(u32),
}

impl From<CommunicatorError> for FormatError {
  fn from(err: CommunicatorError) -> Self {
    FormatError::new(err)
  }
}

pub type HostFormatCallback = Rc<dyn Fn(HostFormatRequest) -> LocalBoxFuture<'static, FormatResult>>;

pub struct ProcessPluginCommunicatorFormatRequest {
  pub file_path: PathBuf,
  pub file_bytes: Vec<u8>,
  pub range: FormatRange,
  pub config_id: FormatConfigId,
  pub override_config: ConfigKeyMap,
  pub on_host_format: HostFormatCallback,
  pub token: DprintCancellationToken,
}

enum MessageResponseChannel {
  Acknowledgement(oneshot::Sender<Result<()>>),
  Data(oneshot::Sender<Result<Vec<u8>>>),
  Format(oneshot::Sender<Result<Option<Vec<u8>>>>),
}

struct Context {
  stdin_writer: SingleThreadMessageWriter<ProcessPluginMessage>,
  shutdown_flag: Arc<AtomicFlag>,
  id_generator: IdGenerator,
  messages: RcIdStore<MessageResponseChannel>,
  format_request_tokens: RcIdStore<Arc<CancellationToken>>,
  host_format_callbacks: RcIdStore<HostFormatCallback>,
}

/// Communicates with a process plugin.
pub struct ProcessPluginCommunicator {
  child: RefCell<Option<Child>>,
  context: Rc<Context>,
}

impl Drop for ProcessPluginCommunicator {
  fn drop(&mut self) {
    self.kill();
  }
}

impl ProcessPluginCommunicator {
  pub async fn new(executable_file_path: &Path, on_std_err: impl Fn(String) + Clone + Send + Sync + 'static) -> Result<Self> {
    ProcessPluginCommunicator::new_internal(executable_file_path, false, on_std_err).await
  }

  /// Provides the `--init` CLI flag to tell the process plugin to do any initialization necessary
  pub async fn new_with_init(executable_file_path: &Path, on_std_err: impl Fn(String) + Clone + Send + Sync + 'static) -> Result<Self> {
    ProcessPluginCommunicator::new_internal(executable_file_path, true, on_std_err).await
  }

  async fn new_internal(executable_file_path: &Path, is_init: bool, on_std_err: impl Fn(String) + Clone + Send + Sync + 'static) -> Result<Self> {
    let mut args = vec!["--parent-pid".to_string(), std::process::id().to_string()];
    if is_init {
      args.push("--init".to_string());
    }

    let shutdown_flag = Arc::new(AtomicFlag::default());
    let mut child = Command::new(executable_file_path)
      .args(&args)
      .stdin(Stdio::piped())
      .stderr(Stdio::piped())
      .stdout(Stdio::piped())
      .spawn()
      .map_err(|err| CommunicatorError::StartProcess {
        executable: executable_file_path.display().to_string(),
        args: args.join(" "),
        error: err,
      })?;

    // read and output stderr prefixed
    let stderr = child.stderr.take().unwrap();
    crate::async_runtime::spawn_blocking({
      let shutdown_flag = shutdown_flag.clone();
      let on_std_err = on_std_err.clone();
      move || {
        std_err_redirect(shutdown_flag, stderr, on_std_err);
      }
    });

    // verify the schema version
    let mut stdout_reader = MessageReader::new(child.stdout.take().unwrap());
    let mut stdin_writer = MessageWriter::new(child.stdin.take().unwrap());

    let (mut stdout_reader, stdin_writer, schema_version) = crate::async_runtime::spawn_blocking(move || {
      let schema_version = get_plugin_schema_version(&mut stdout_reader, &mut stdin_writer).map_err(CommunicatorError::SchemaVerification)?;
      Ok::<_, FormatError>((stdout_reader, stdin_writer, schema_version))
    })
    .await??;

    if schema_version != PLUGIN_SCHEMA_VERSION {
      // kill the child to prevent it from ouputting to stderr
      let _ = child.kill();
      let err = if schema_version < PLUGIN_SCHEMA_VERSION {
        CommunicatorError::PluginTooOld {
          actual: schema_version,
          expected: PLUGIN_SCHEMA_VERSION,
        }
      } else {
        CommunicatorError::CliTooOld {
          actual: schema_version,
          expected: PLUGIN_SCHEMA_VERSION,
        }
      };
      return Err(err.into());
    }

    let stdin_writer = SingleThreadMessageWriter::for_stdin(stdin_writer);
    let context = Rc::new(Context {
      id_generator: Default::default(),
      shutdown_flag,
      stdin_writer,
      messages: Default::default(),
      format_request_tokens: Default::default(),
      host_format_callbacks: Default::default(),
    });

    // read from stdout
    let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
    crate::async_runtime::spawn_blocking({
      let shutdown_flag = context.shutdown_flag.clone();
      let on_std_err = on_std_err.clone();
      move || {
        loop {
          match ProcessPluginMessage::read(&mut stdout_reader) {
            Ok(message) => {
              if tx.send(message).is_err() {
                break; // closed
              }
            }
            Err(err) if err.kind() == ErrorKind::BrokenPipe => {
              break;
            }
            Err(err) => {
              if !shutdown_flag.is_raised() {
                on_std_err(format!("Error reading stdout message: {:#}", err));
              }
              break;
            }
          }
        }
      }
    });
    crate::async_runtime::spawn({
      let context = context.clone();
      async move {
        while let Some(message) = rx.recv().await {
          if let Err(err) = handle_stdout_message(message, &context) {
            if !context.shutdown_flag.is_raised() {
              on_std_err(format!("Error reading stdout message: {:#}", err));
            }
            break;
          }
        }
        // clear out all the messages
        context.messages.take_all();
      }
    });

    Ok(Self {
      child: RefCell::new(Some(child)),
      context,
    })
  }

  /// Perform a graceful shutdown.
  pub async fn shutdown(&self) {
    if self.context.shutdown_flag.raise() {
      // attempt to exit nicely
      tokio::select! {
        // we wait for acknowledgement in order to give the process
        // plugin a chance to clean up (ex. in case it has spawned
        // any processes it needs to kill or something like that)
        _ = self.send_with_acknowledgement(MessageBody::Close) => {}
        _ = tokio::time::sleep(Duration::from_millis(250)) => {
          self.kill();
        }
      }
    } else {
      self.kill();
    }
  }

  pub fn kill(&self) {
    self.context.shutdown_flag.raise();
    if let Some(mut child) = self.child.borrow_mut().take() {
      let _ignore = child.kill();
    }
  }

  pub async fn register_config(&self, config_id: FormatConfigId, global_config: &GlobalConfiguration, plugin_config: &ConfigKeyMap) -> Result<()> {
    let global_config = serde_json::to_vec(global_config)?;
    let plugin_config = serde_json::to_vec(plugin_config)?;
    self
      .send_with_acknowledgement(MessageBody::RegisterConfig(RegisterConfigMessageBody {
        config_id,
        global_config,
        plugin_config,
      }))
      .await?;
    Ok(())
  }

  pub async fn release_config(&self, config_id: FormatConfigId) -> Result<()> {
    self.send_with_acknowledgement(MessageBody::ReleaseConfig(config_id)).await?;
    Ok(())
  }

  pub async fn ask_is_alive(&self) -> bool {
    self.send_with_acknowledgement(MessageBody::IsAlive).await.is_ok()
  }

  pub async fn plugin_info(&self) -> Result<PluginInfo> {
    self.send_receiving_data(MessageBody::GetPluginInfo).await
  }

  pub async fn license_text(&self) -> Result<String> {
    self.send_receiving_string(MessageBody::GetLicenseText).await
  }

  pub async fn resolved_config(&self, config_id: FormatConfigId) -> Result<String> {
    self.send_receiving_string(MessageBody::GetResolvedConfig(config_id)).await
  }

  pub async fn file_matching_info(&self, config_id: FormatConfigId) -> Result<FileMatchingInfo> {
    self.send_receiving_data(MessageBody::GetFileMatchingInfo(config_id)).await
  }

  pub async fn config_diagnostics(&self, config_id: FormatConfigId) -> Result<Vec<ConfigurationDiagnostic>> {
    self.send_receiving_data(MessageBody::GetConfigDiagnostics(config_id)).await
  }

  pub async fn check_config_updates(&self, message: &CheckConfigUpdatesMessageBody) -> Result<Vec<ConfigChange>> {
    let bytes = serde_json::to_vec(&message)?;
    let response: CheckConfigUpdatesResponseBody = self.send_receiving_data(MessageBody::CheckConfigUpdates(bytes)).await?;
    Ok(response.changes)
  }

  pub async fn format_text(&self, request: ProcessPluginCommunicatorFormatRequest) -> FormatResult {
    let (tx, rx) = oneshot::channel::<Result<Option<Vec<u8>>>>();

    let message_id = self.context.id_generator.next();
    let store_guard = self.context.host_format_callbacks.store_with_guard(message_id, request.on_host_format);
    let maybe_result = self
      .send_message_with_id(
        message_id,
        MessageBody::Format(FormatMessageBody {
          file_path: request.file_path,
          file_bytes: request.file_bytes,
          range: request.range,
          config_id: request.config_id,
          override_config: serde_json::to_vec(&request.override_config).unwrap(),
        }),
        MessageResponseChannel::Format(tx),
        rx,
        request.token.clone(),
      )
      .await;

    drop(store_guard); // explicit for clarity

    if request.token.is_cancelled() {
      Ok(None)
    } else {
      match maybe_result {
        Ok(result) => result,
        Err(err) => Err(CriticalFormatError(err).into()),
      }
    }
  }

  /// Checks if the process is functioning.
  pub async fn is_process_alive(&self) -> bool {
    if self.context.shutdown_flag.is_raised() {
      false
    } else {
      self.ask_is_alive().await
    }
  }

  async fn send_with_acknowledgement(&self, body: MessageBody) -> Result<()> {
    let (tx, rx) = oneshot::channel::<Result<()>>();
    self
      .send_message(body, MessageResponseChannel::Acknowledgement(tx), rx, Arc::new(NullCancellationToken))
      .await?
  }

  async fn send_receiving_string(&self, body: MessageBody) -> Result<String> {
    let data = self.send_receiving_bytes(body).await??;
    Ok(String::from_utf8(data)?)
  }

  async fn send_receiving_data<T: DeserializeOwned>(&self, body: MessageBody) -> Result<T> {
    let data = self.send_receiving_bytes(body).await??;
    Ok(serde_json::from_slice(&data)?)
  }

  async fn send_receiving_bytes(&self, body: MessageBody) -> Result<Result<Vec<u8>>> {
    let (tx, rx) = oneshot::channel::<Result<Vec<u8>>>();
    self
      .send_message(body, MessageResponseChannel::Data(tx), rx, Arc::new(NullCancellationToken))
      .await
  }

  async fn send_message<T: Default>(
    &self,
    body: MessageBody,
    response_channel: MessageResponseChannel,
    receiver: oneshot::Receiver<Result<T>>,
    token: Arc<dyn super::super::CancellationToken>,
  ) -> Result<Result<T>> {
    let message_id = self.context.id_generator.next();
    self.send_message_with_id(message_id, body, response_channel, receiver, token).await
  }

  async fn send_message_with_id<T: Default>(
    &self,
    message_id: u32,
    body: MessageBody,
    response_channel: MessageResponseChannel,
    receiver: oneshot::Receiver<Result<T>>,
    token: Arc<dyn super::super::CancellationToken>,
  ) -> Result<Result<T>> {
    let mut drop_guard = DropGuardAction::new(|| {
      // clear up memory
      self.context.messages.take(message_id);
      // send cancellation to the client
      let _ = self.context.stdin_writer.send(ProcessPluginMessage {
        id: self.context.id_generator.next(),
        body: MessageBody::CancelFormat(message_id),
      });
    });

    self.context.messages.store(message_id, response_channel);
    self.context.stdin_writer.send(ProcessPluginMessage { id: message_id, body })?;
    tokio::select! {
      _ = token.wait_cancellation() => {
        drop(drop_guard); // explicit
        Ok(Ok(Default::default()))
      }
      response = receiver => {
        drop_guard.forget(); // we completed, so don't run the drop guard
        match response {
          Ok(data) => Ok(data),
          Err(err) => Err(CommunicatorError::WaitOnMessage { message_id, error: err }.into()),
        }
      }
    }
  }
}

fn get_plugin_schema_version<TRead: Read + Unpin, TWrite: Write + Unpin>(
  reader: &mut MessageReader<TRead>,
  writer: &mut MessageWriter<TWrite>,
) -> std::result::Result<u32, SchemaVersionError> {
  // since this is the setup, identify exactly where it failed
  writer.send_u32(0).map_err(SchemaVersionError::Ask)?; // ask for schema version
  writer.flush().map_err(SchemaVersionError::Flush)?;
  let acknowledgement_response = reader.read_u32().map_err(SchemaVersionError::ReadAcknowledgement)?;
  if acknowledgement_response != 0 {
    return Err(SchemaVersionError::UnexpectedAcknowledgement(acknowledgement_response));
  }
  reader.read_u32().map_err(SchemaVersionError::ReadVersion)
}

fn std_err_redirect(shutdown_flag: Arc<AtomicFlag>, stderr: ChildStderr, on_std_err: impl Fn(String) + Send + Sync + 'static) {
  let reader = std::io::BufReader::new(stderr);
  for line in reader.lines() {
    match line {
      Ok(line) => on_std_err(line),
      Err(err) => {
        if shutdown_flag.is_raised() || err.kind() == ErrorKind::BrokenPipe {
          return;
        } else {
          on_std_err(format!("Error reading line from process plugin stderr. {:#}", err));
        }
      }
    }
  }
}

fn handle_stdout_message(message: ProcessPluginMessage, context: &Rc<Context>) -> Result<()> {
  match message.body {
    MessageBody::Success(message_id) => match context.messages.take(message_id) {
      Some(MessageResponseChannel::Acknowledgement(channel)) => {
        let _ignore = channel.send(Ok(()));
      }
      Some(MessageResponseChannel::Data(channel)) => {
        let _ignore = channel.send(Err(CommunicatorError::UnexpectedDataForSuccess(message_id).into()));
      }
      Some(MessageResponseChannel::Format(channel)) => {
        let _ignore = channel.send(Err(CommunicatorError::UnexpectedFormatForSuccess(message_id).into()));
      }
      None => {}
    },
    MessageBody::DataResponse(response) => match context.messages.take(response.message_id) {
      Some(MessageResponseChannel::Acknowledgement(channel)) => {
        let _ignore = channel.send(Err(CommunicatorError::UnexpectedSuccessForData(response.message_id).into()));
      }
      Some(MessageResponseChannel::Data(channel)) => {
        let _ignore = channel.send(Ok(response.data));
      }
      Some(MessageResponseChannel::Format(channel)) => {
        let _ignore = channel.send(Err(CommunicatorError::UnexpectedFormatForData(response.message_id).into()));
      }
      None => {}
    },
    MessageBody::Error(response) => {
      let err: FormatError = String::from_utf8_lossy(&response.data).into_owned().into();
      match context.messages.take(response.message_id) {
        Some(MessageResponseChannel::Acknowledgement(channel)) => {
          let _ignore = channel.send(Err(err));
        }
        Some(MessageResponseChannel::Data(channel)) => {
          let _ignore = channel.send(Err(err));
        }
        Some(MessageResponseChannel::Format(channel)) => {
          let _ignore = channel.send(Err(err));
        }
        None => {}
      }
    }
    MessageBody::FormatResponse(response) => match context.messages.take(response.message_id) {
      Some(MessageResponseChannel::Acknowledgement(channel)) => {
        let _ignore = channel.send(Err(CommunicatorError::UnexpectedSuccessForFormat(response.message_id).into()));
      }
      Some(MessageResponseChannel::Data(channel)) => {
        let _ignore = channel.send(Err(CommunicatorError::UnexpectedDataForFormat(response.message_id).into()));
      }
      Some(MessageResponseChannel::Format(channel)) => {
        let _ignore = channel.send(Ok(response.data));
      }
      None => {}
    },
    MessageBody::CancelFormat(message_id) => {
      if let Some(token) = context.format_request_tokens.take(message_id) {
        token.cancel();
      }
      context.host_format_callbacks.take(message_id);
      // do not clear from context.messages here because the cancellation will do that
    }
    MessageBody::HostFormat(body) => {
      // spawn a task to do the host formatting, then send a message back to the
      // plugin with the result
      let context = context.clone();
      crate::async_runtime::spawn(async move {
        let result = host_format(context.clone(), message.id, body).await;

        // ignore failure, as this means that the process shut down
        // at which point handling would have occurred elsewhere
        let _ignore = context.stdin_writer.send(ProcessPluginMessage {
          id: context.id_generator.next(),
          body: match result {
            Ok(result) => MessageBody::FormatResponse(ResponseBody {
              message_id: message.id,
              data: result,
            }),
            Err(err) => MessageBody::Error(ResponseBody {
              message_id: message.id,
              data: error_to_string(&err).into_bytes(),
            }),
          },
        });
      });
    }
    MessageBody::IsAlive => {
      // the CLI is not documented as supporting this, but we might as well respond
      let _ = context.stdin_writer.send(ProcessPluginMessage {
        id: context.id_generator.next(),
        body: MessageBody::Success(message.id),
      });
    }
    MessageBody::Format(_)
    | MessageBody::Close
    | MessageBody::GetPluginInfo
    | MessageBody::GetLicenseText
    | MessageBody::RegisterConfig(_)
    | MessageBody::ReleaseConfig(_)
    | MessageBody::GetConfigDiagnostics(_)
    | MessageBody::GetFileMatchingInfo(_)
    | MessageBody::GetResolvedConfig(_)
    | MessageBody::CheckConfigUpdates(_) => {
      let _ = context.stdin_writer.send(ProcessPluginMessage {
        id: context.id_generator.next(),
        body: MessageBody::Error(ResponseBody {
          message_id: message.id,
          data: "Unsupported plugin to CLI message.".as_bytes().to_vec(),
        }),
      });
    }
    // If encountered, process plugin should exit and
    // the CLI should kill the process plugin.
    MessageBody::Unknown(message_kind) => {
      return Err(CommunicatorError::UnknownMessageKind(message_kind).into());
    }
  }

  Ok(())
}

async fn host_format(context: Rc<Context>, message_id: u32, body: HostFormatMessageBody) -> FormatResult {
  let Some(callback) = context.host_format_callbacks.get_cloned(body.original_message_id) else {
    return FormatResult::Err(CommunicatorError::HostFormatCallbackNotFound(body.original_message_id).into());
  };

  let token = Arc::new(CancellationToken::new());
  let store_guard = context.format_request_tokens.store_with_guard(message_id, token.clone());
  let result = callback(HostFormatRequest {
    file_path: body.file_path,
    file_bytes: body.file_text,
    range: body.range,
    override_config: serde_json::from_slice(&body.override_config).unwrap(),
    token,
  })
  .await;
  drop(store_guard); // explicit for clarity
  result
}