RUST-1396 Respawn mongocryptd on server selection failure (#730)

Author: abr-egn

Cargo.toml

@@ -31,13 +31,15 @@ default = ["tokio-runtime"]
 tokio-runtime = [
     "tokio/macros",
     "tokio/net",
+    "tokio/process",
     "tokio/rt",
     "tokio/time",
     "serde_bytes",
 ]
 async-std-runtime = [
     "async-std",
     "async-std/attributes",
+    "async-std/unstable",
     "async-std-resolver",
     "tokio-util/compat",
 ]
@@ -66,7 +68,7 @@ zlib-compression = ["flate2"]
 snappy-compression = ["snap"]
 
 # DO NOT USE; see https://jira.mongodb.org/browse/RUST-569 for the status of CSFLE support in the Rust driver.
-csfle = ["mongocrypt", "which", "rayon"]
+csfle = ["mongocrypt", "rayon"]
 
 [dependencies]
 async-trait = "0.1.42"
@@ -106,7 +108,6 @@ trust-dns-proto = "0.21.2"
 trust-dns-resolver = "0.21.2"
 typed-builder = "0.10.0"
 webpki-roots = "0.22.4"
-which = { version = "4.2.5", optional = true }
 zstd = { version = "0.11.2", optional = true }
 
 [dependencies.async-std]

src/client/csfle.rs

@@ -2,10 +2,7 @@ pub mod client_encryption;
 pub mod options;
 mod state_machine;
 
-use std::{
-    path::Path,
-    process::{Command, Stdio},
-};
+use std::path::Path;
 
 use derivative::Derivative;
 use mongocrypt::Crypt;
@@ -26,7 +23,7 @@ use options::{
     EO_MONGOCRYPTD_URI,
 };
 
-use self::state_machine::CryptExecutor;
+use self::state_machine::{CryptExecutor, MongocryptdOptions};
 
 use super::WeakClient;
 
@@ -49,15 +46,16 @@ struct AuxClients {
 impl ClientState {
     pub(super) async fn new(client: &Client, mut opts: AutoEncryptionOptions) -> Result<Self> {
         let crypt = Self::make_crypt(&opts)?;
-        let mongocryptd_client = Self::spawn_mongocryptd_if_needed(&opts, &crypt).await?;
+        let mongocryptd_opts = Self::make_mongocryptd_opts(&opts, &crypt)?;
         let aux_clients = Self::make_aux_clients(client, &opts)?;
-        let exec = CryptExecutor::new(
+        let exec = CryptExecutor::new_implicit(
             aux_clients.key_vault_client,
             opts.key_vault_namespace.clone(),
-            mongocryptd_client,
-            aux_clients.metadata_client,
             opts.tls_options.take(),
-        )?;
+            mongocryptd_opts,
+            aux_clients.metadata_client,
+        )
+        .await?;
 
         Ok(Self {
             crypt,
@@ -98,57 +96,37 @@ impl ClientState {
         Ok(crypt)
     }
 
-    /// If crypt_shared is unavailable and options have not disabled it, spawn mongocryptd.  Returns
-    /// a `Client` connected to the mongocryptd if one was spawned.
-    async fn spawn_mongocryptd_if_needed(
+    fn make_mongocryptd_opts(
         opts: &AutoEncryptionOptions,
         crypt: &Crypt,
-    ) -> Result<Option<Client>> {
+    ) -> Result<Option<MongocryptdOptions>> {
         if opts.bypass_auto_encryption == Some(true)
             || opts.extra_option(&EO_MONGOCRYPTD_BYPASS_SPAWN)? == Some(true)
             || crypt.shared_lib_version().is_some()
             || opts.extra_option(&EO_CRYPT_SHARED_REQUIRED)? == Some(true)
         {
             return Ok(None);
         }
-        let which_path;
-        let bin_path = match opts.extra_option(&EO_MONGOCRYPTD_SPAWN_PATH)? {
-            Some(s) => Path::new(s),
-            None => {
-                which_path = which::which("mongocryptd")
-                    .map_err(|e| Error::invalid_argument(format!("{}", e)))?;
-                &which_path
-            }
-        };
-        let mut args: Vec<&str> = vec![];
-        let has_idle = if let Some(spawn_args) = opts.extra_option(&EO_MONGOCRYPTD_SPAWN_ARGS)? {
-            let mut has_idle = false;
-            for arg in spawn_args {
+        let spawn_path = opts
+            .extra_option(&EO_MONGOCRYPTD_SPAWN_PATH)?
+            .map(std::path::PathBuf::from);
+        let mut spawn_args = vec![];
+        if let Some(args) = opts.extra_option(&EO_MONGOCRYPTD_SPAWN_ARGS)? {
+            for arg in args {
                 let str_arg = arg.as_str().ok_or_else(|| {
                     Error::invalid_argument("non-string entry in mongocryptdSpawnArgs")
                 })?;
-                has_idle |= str_arg.starts_with("--idleShutdownTimeoutSecs");
-                args.push(str_arg);
+                spawn_args.push(str_arg.to_string());
             }
-            has_idle
-        } else {
-            false
-        };
-        if !has_idle {
-            args.push("--idleShutdownTimeoutSecs=60");
         }
-        Command::new(bin_path)
-            .args(&args)
-            .stdout(Stdio::null())
-            .stderr(Stdio::null())
-            .spawn()?;
-
         let uri = opts
             .extra_option(&EO_MONGOCRYPTD_URI)?
-            .unwrap_or("mongodb://localhost:27020");
-        let mut options = super::options::ClientOptions::parse_uri(uri, None).await?;
-        options.server_selection_timeout = Some(std::time::Duration::from_millis(10_000));
-        Ok(Some(Client::with_options(options)?))
+            .map(|s| s.to_string());
+        Ok(Some(MongocryptdOptions {
+            spawn_path,
+            spawn_args,
+            uri,
+        }))
     }
 
     fn make_aux_clients(

src/client/csfle/client_encryption.rs

@@ -36,11 +36,9 @@ impl ClientEncryption {
     /// Create a new key vault handle with the given options.
     pub fn new(options: ClientEncryptionOptions) -> Result<Self> {
         let crypt = Crypt::builder().build()?;
-        let exec = CryptExecutor::new(
+        let exec = CryptExecutor::new_explicit(
             options.key_vault_client.weak(),
             options.key_vault_namespace.clone(),
-            None,
-            None,
             options.tls_options,
         )?;
         let key_vault = options

src/client/csfle/state_machine.rs

@@ -1,20 +1,24 @@
-use std::convert::TryInto;
+use std::{
+    convert::TryInto,
+    path::{Path, PathBuf},
+    time::Duration,
+};
 
 use bson::{Document, RawDocument, RawDocumentBuf};
 use futures_util::{stream, TryStreamExt};
 use mongocrypt::ctx::{Ctx, State};
 use rayon::ThreadPool;
 use tokio::{
     io::{AsyncReadExt, AsyncWriteExt},
-    sync::oneshot,
+    sync::{oneshot, Mutex},
 };
 
 use crate::{
     client::{options::ServerAddress, WeakClient},
     cmap::options::StreamOptions,
     error::{Error, Result},
     operation::{RawOutput, RunCommand},
-    runtime::AsyncStream,
+    runtime::{AsyncStream, Process},
     Client,
     Namespace,
 };
@@ -25,18 +29,16 @@ use super::options::KmsProvidersTlsOptions;
 pub(crate) struct CryptExecutor {
     key_vault_client: WeakClient,
     key_vault_namespace: Namespace,
-    mongocryptd_client: Option<Client>,
-    metadata_client: Option<WeakClient>,
     tls_options: Option<KmsProvidersTlsOptions>,
     crypto_threads: ThreadPool,
+    mongocryptd: Option<Mongocryptd>,
+    metadata_client: Option<WeakClient>,
 }
 
 impl CryptExecutor {
-    pub(crate) fn new(
+    pub(crate) fn new_explicit(
         key_vault_client: WeakClient,
         key_vault_namespace: Namespace,
-        mongocryptd_client: Option<Client>,
-        metadata_client: Option<WeakClient>,
         tls_options: Option<KmsProvidersTlsOptions>,
     ) -> Result<Self> {
         let num_cpus = std::thread::available_parallelism()?.get();
@@ -47,13 +49,30 @@ impl CryptExecutor {
         Ok(Self {
             key_vault_client,
             key_vault_namespace,
-            mongocryptd_client,
-            metadata_client,
             tls_options,
             crypto_threads,
+            mongocryptd: None,
+            metadata_client: None,
         })
     }
 
+    pub(crate) async fn new_implicit(
+        key_vault_client: WeakClient,
+        key_vault_namespace: Namespace,
+        tls_options: Option<KmsProvidersTlsOptions>,
+        mongocryptd_opts: Option<MongocryptdOptions>,
+        metadata_client: Option<WeakClient>,
+    ) -> Result<Self> {
+        let mongocryptd = match mongocryptd_opts {
+            Some(opts) => Some(Mongocryptd::new(opts).await?),
+            None => None,
+        };
+        let mut exec = Self::new_explicit(key_vault_client, key_vault_namespace, tls_options)?;
+        exec.mongocryptd = mongocryptd;
+        exec.metadata_client = metadata_client;
+        Ok(exec)
+    }
+
     pub(crate) async fn run_ctx(&self, ctx: Ctx, db: Option<&str>) -> Result<RawDocumentBuf> {
         let mut result = None;
         // This needs to be a `Result` so that the `Ctx` can be temporarily owned by the processing
@@ -89,11 +108,25 @@ impl CryptExecutor {
                         Error::internal("db required for NeedMongoMarkings state")
                     })?;
                     let op = RawOutput(RunCommand::new_raw(db.to_string(), command, None, None)?);
-                    let mongocryptd_client = self
-                        .mongocryptd_client
+                    let mongocryptd = self
+                        .mongocryptd
                         .as_ref()
                         .ok_or_else(|| Error::internal("mongocryptd client not found"))?;
-                    let response = mongocryptd_client.execute_operation(op, None).await?;
+                    let result = mongocryptd.client.execute_operation(op.clone(), None).await;
+                    let response = match result {
+                        Ok(r) => r,
+                        Err(e) if e.is_server_selection_error() => {
+                            mongocryptd.respawn().await?;
+                            match mongocryptd.client.execute_operation(op, None).await {
+                                Ok(r) => r,
+                                Err(new_e) if !new_e.is_server_selection_error() => {
+                                    return Err(new_e)
+                                }
+                                Err(_) => return Err(e),
+                            }
+                        }
+                        Err(e) => return Err(e),
+                    };
                     ctx.mongo_feed(response.raw_body())?;
                     ctx.mongo_done()?;
                 }
@@ -181,6 +214,80 @@ impl CryptExecutor {
     }
 }
 
+#[derive(Debug)]
+struct Mongocryptd {
+    opts: MongocryptdOptions,
+    client: Client,
+    child: Mutex<Result<Process>>,
+}
+
+impl Mongocryptd {
+    const DEFAULT_URI: &'static str = "mongodb://localhost:27020";
+    const SERVER_SELECTION_TIMEOUT: Duration = Duration::from_millis(10_000);
+
+    async fn new(opts: MongocryptdOptions) -> Result<Self> {
+        let child = Mutex::new(Ok(Self::spawn(&opts)?));
+        let uri = opts.uri.as_deref().unwrap_or(Self::DEFAULT_URI);
+        let mut options = crate::options::ClientOptions::parse_uri(uri, None).await?;
+        options.server_selection_timeout = Some(Self::SERVER_SELECTION_TIMEOUT);
+        let client = Client::with_options(options)?;
+        Ok(Self {
+            opts,
+            client,
+            child,
+        })
+    }
+
+    async fn respawn(&self) -> Result<()> {
+        let mut child = match self.child.try_lock() {
+            Ok(l) => l,
+            _ => {
+                // Another respawn is in progress.  Lock to wait for it.
+                return unit_err(&*self.child.lock().await);
+            }
+        };
+        let new_child = Self::spawn(&self.opts);
+        if new_child.is_ok() {
+            if let Ok(old_child) = child.as_mut() {
+                let _ = old_child.wait().await;
+            }
+        }
+        *child = new_child;
+        unit_err(&*child)
+    }
+
+    fn spawn(opts: &MongocryptdOptions) -> Result<Process> {
+        let bin_path = match &opts.spawn_path {
+            Some(s) => s,
+            None => Path::new("mongocryptd"),
+        };
+        let mut args: Vec<&str> = vec![];
+        let mut has_idle = false;
+        for arg in &opts.spawn_args {
+            has_idle |= arg.starts_with("--idleShutdownTimeoutSecs");
+            args.push(arg);
+        }
+        if !has_idle {
+            args.push("--idleShutdownTimeoutSecs=60");
+        }
+        Process::spawn(bin_path, &args)
+    }
+}
+
+fn unit_err<T>(r: &Result<T>) -> Result<()> {
+    match r {
+        Ok(_) => Ok(()),
+        Err(e) => Err(e.clone()),
+    }
+}
+
+#[derive(Debug)]
+pub(crate) struct MongocryptdOptions {
+    pub(crate) spawn_path: Option<PathBuf>,
+    pub(crate) spawn_args: Vec<String>,
+    pub(crate) uri: Option<String>,
+}
+
 fn result_ref<T>(r: &Result<T>) -> Result<&T> {
     r.as_ref().map_err(Error::clone)
 }

src/operation/raw_output.rs

@@ -7,6 +7,7 @@ use super::Operation;
 
 /// Forwards all implementation to the wrapped `Operation`, but returns the response unparsed and
 /// unvalidated as a `RawCommandResponse`.
+#[derive(Clone)]
 pub(crate) struct RawOutput<Op>(pub(crate) Op);
 
 impl<Op: Operation> Operation for RawOutput<Op> {

src/operation/run_command/mod.rs

@@ -15,7 +15,7 @@ use crate::{
     selection_criteria::SelectionCriteria,
 };
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub(crate) struct RunCommand<'conn> {
     db: String,
     command: RawDocumentBuf,

src/runtime/mod.rs

@@ -3,6 +3,8 @@ mod http;
 #[cfg(feature = "async-std-runtime")]
 mod interval;
 mod join_handle;
+#[cfg(feature = "csfle")]
+mod process;
 mod resolver;
 mod stream;
 mod sync_read_ext;
@@ -14,6 +16,8 @@ mod worker_handle;
 
 use std::{future::Future, net::SocketAddr, time::Duration};
 
+#[cfg(feature = "csfle")]
+pub(crate) use self::process::Process;
 pub(crate) use self::{
     acknowledged_message::AcknowledgedMessage,
     join_handle::AsyncJoinHandle,