tools/unitctl: Enable Multi Socket Support

This commit refactors the CLI code to accept multiple instances of the control socket flag. All subcommands except for edit and save now support being run against multiple specified instances of unitd. * control_socket_addresses CLI field is now a vector * centralize error related logic into the error module * wait_for_socket now returns a vector of sockets. all sockets in vector are waited upon and validated * extraneous code is removed * applications, execute, import, listeners, and status commands all run against N control sockets now * edit and save commands return error when run against a single control socket Signed-off-by: Ava Hahn <[email protected]>
kohanaworld · Jul 8, 2024 · 706ea1a · 706ea1a
1 parent b5fe3ea
commit 706ea1a
Show file tree

Hide file tree

Showing 12 changed files with 282 additions and 202 deletions.
diff --git a/tools/unitctl/README.md b/tools/unitctl/README.md
@@ -111,6 +111,14 @@ $ unitctl app reload wasm
 }
 ```
 
+*Note:* Both of the above commands support operating on multiple instances
+of Unit at once. To do this, pass multiple values for the `-s` flag as
+shown below:
+
+```
+$ unitctl -s '127.0.0.1:8001' -s /run/nginx-unit.control.sock app list
+```
+
 ### Lists active listeners from running Unit processes
 ```
 unitctl listeners
@@ -122,6 +130,13 @@ No socket path provided - attempting to detect from running instance
 }
 ```
 
+*Note:* This command supports operating on multiple instances of Unit at once.
+To do this, pass multiple values for the `-s` flag as shown below:
+
+```
+$ unitctl -s '127.0.0.1:8001' -s /run/nginx-unit.control.sock listeners
+```
+
 ### Get the current status of NGINX Unit processes
 ```
 $ unitctl status -t yaml
@@ -136,6 +151,13 @@ requests:
 applications: {}
 ```
 
+*Note:* This command supports operating on multiple instances of Unit at once.
+To do this, pass multiple values for the `-s` flag as shown below:
+
+```
+$ unitctl -s '127.0.0.1:8001' -s /run/nginx-unit.control.sock status
+```
+
 ### Send arbitrary configuration payloads to Unit
 ```
 $ echo '{
@@ -158,6 +180,13 @@ $ echo '{
 }
 ```
 
+*Note:* This command supports operating on multiple instances of Unit at once.
+To do this, pass multiple values for the `-s` flag as shown below:
+
+```
+$ unitctl -s '127.0.0.1:8001' -s /run/nginx-unit.control.sock execute ...
+```
+
 ### Edit current configuration in your favorite editor
 ```
 $ unitctl edit
@@ -168,6 +197,8 @@ $ unitctl edit
 }
 ```
 
+*Note:* This command does not support operating on multiple instances of Unit at once.
+
 ### Import configuration, certificates, and NJS modules from directory
 ```
 $ unitctl import /opt/unit/config
@@ -191,6 +222,8 @@ $ unitctl export -f - > config.tar
 
 *Note:* The exported configuration omits certificates.
 
+*Note:* This command does not support operating on multiple instances of Unit at once.
+
 ### Wait for socket to become available
 ```
 $ unitctl --wait-timeout-seconds=3 --wait-max-tries=4 import /opt/unit/config`

diff --git a/tools/unitctl/unitctl/src/cmd/applications.rs b/tools/unitctl/unitctl/src/cmd/applications.rs
@@ -1,36 +1,46 @@
 use crate::unitctl::{ApplicationArgs, ApplicationCommands, UnitCtl};
-use crate::{wait, UnitctlError};
+use crate::{wait, UnitctlError, eprint_error};
 use crate::requests::send_empty_body_deserialize_response;
 use unit_client_rs::unit_client::UnitClient;
 
 pub(crate) async fn cmd(cli: &UnitCtl, args: &ApplicationArgs) -> Result<(), UnitctlError> {
-    let control_socket = wait::wait_for_socket(cli).await?;
-    let client = UnitClient::new(control_socket);
+    let clients: Vec<UnitClient> = wait::wait_for_sockets(cli)
+        .await?
+        .into_iter()
+        .map(|sock| UnitClient::new(sock))
+        .collect();
 
-    match &args.command {
-        ApplicationCommands::Reload { ref name } => client
-            .restart_application(name)
-            .await
-            .map_err(|e| UnitctlError::UnitClientError { source: *e })
-            .and_then(|r| args.output_format.write_to_stdout(&r)),
+    for client in clients {
+        let _ = match &args.command {
+            ApplicationCommands::Reload { ref name } => client
+                .restart_application(name)
+                .await
+                .map_err(|e| UnitctlError::UnitClientError { source: *e })
+                .and_then(|r| args.output_format.write_to_stdout(&r)),
 
-        /* we should be able to use this but the openapi generator library
-         * is fundamentally incorrect and provides a broken API for the
-         * applications endpoint.
-        ApplicationCommands::List {} => client
-            .applications()
-            .await
-            .map_err(|e| UnitctlError::UnitClientError { source: *e })
-            .and_then(|response| args.output_format.write_to_stdout(&response)),*/
+            /* we should be able to use this but the openapi generator library
+             * is fundamentally incorrect and provides a broken API for the
+             * applications endpoint.
+            ApplicationCommands::List {} => client
+                .applications()
+                .await
+                .map_err(|e| UnitctlError::UnitClientError { source: *e })
+                .and_then(|response| args.output_format.write_to_stdout(&response)),*/
 
-        ApplicationCommands::List {} => {
-            args.output_format.write_to_stdout(
-                &send_empty_body_deserialize_response(
-                    &client,
-                    "GET",
-                    "/config/applications",
-                ).await?
-            )
-        },
+            ApplicationCommands::List {} => {
+                args.output_format.write_to_stdout(
+                    &send_empty_body_deserialize_response(
+                        &client,
+                        "GET",
+                        "/config/applications",
+                    ).await?
+                )
+            },
+        }.map_err(|error| {
+            eprint_error(&error);
+            std::process::exit(error.exit_code());
+        });
     }
+
+    Ok(())
 }
diff --git a/tools/unitctl/unitctl/src/cmd/edit.rs b/tools/unitctl/unitctl/src/cmd/edit.rs
@@ -1,6 +1,7 @@
 use crate::inputfile::{InputFile, InputFormat};
 use crate::requests::{send_and_validate_config_deserialize_response, send_empty_body_deserialize_response};
 use crate::unitctl::UnitCtl;
+use crate::unitctl_error::ControlSocketErrorKind;
 use crate::{wait, OutputFormat, UnitctlError};
 use std::path::{Path, PathBuf};
 use unit_client_rs::unit_client::UnitClient;
@@ -19,8 +20,16 @@ const EDITOR_KNOWN_LIST: [&str; 8] = [
 ];
 
 pub(crate) async fn cmd(cli: &UnitCtl, output_format: OutputFormat) -> Result<(), UnitctlError> {
-    let control_socket = wait::wait_for_socket(cli).await?;
-    let client = UnitClient::new(control_socket);
+    if cli.control_socket_addresses.is_some() &&
+        cli.control_socket_addresses.clone().unwrap().len() > 1 {
+            return Err(UnitctlError::ControlSocketError{
+                kind: ControlSocketErrorKind::General,
+                message: "too many control sockets. specify at most one.".to_string(),
+            });
+        }
+
+    let mut control_sockets = wait::wait_for_sockets(cli).await?;
+    let client = UnitClient::new(control_sockets.pop().unwrap());
     // Get latest configuration
     let current_config = send_empty_body_deserialize_response(&client, "GET", "/config").await?;
 

diff --git a/tools/unitctl/unitctl/src/cmd/execute.rs b/tools/unitctl/unitctl/src/cmd/execute.rs
@@ -5,7 +5,7 @@ use crate::requests::{
 };
 use crate::unitctl::UnitCtl;
 use crate::wait;
-use crate::{OutputFormat, UnitctlError};
+use crate::{OutputFormat, UnitctlError, eprint_error};
 use unit_client_rs::unit_client::UnitClient;
 
 pub(crate) async fn cmd(
@@ -15,8 +15,11 @@ pub(crate) async fn cmd(
     method: &str,
     path: &str,
 ) -> Result<(), UnitctlError> {
-    let control_socket = wait::wait_for_socket(cli).await?;
-    let client = UnitClient::new(control_socket);
+    let clients: Vec<_> = wait::wait_for_sockets(cli)
+        .await?
+        .into_iter()
+        .map(|sock| UnitClient::new(sock))
+        .collect();
 
     let path_trimmed = path.trim();
     let method_upper = method.to_uppercase();
@@ -28,7 +31,21 @@ pub(crate) async fn cmd(
         eprintln!("Cannot use GET method with input file - ignoring input file");
     }
 
-    send_and_deserialize(client, method_upper, input_file_arg, path_trimmed, output_format).await
+    for client in clients {
+        let _ = send_and_deserialize(
+            client,
+            method_upper.clone(),
+            input_file_arg.clone(),
+            path_trimmed,
+            output_format
+        ).await
+            .map_err(|e| {
+                eprint_error(&e);
+                std::process::exit(e.exit_code());
+            });
+    }
+
+    Ok(())
 }
 
 async fn send_and_deserialize(

diff --git a/tools/unitctl/unitctl/src/cmd/import.rs b/tools/unitctl/unitctl/src/cmd/import.rs
@@ -50,8 +50,12 @@ pub async fn cmd(cli: &UnitCtl, directory: &PathBuf) -> Result<(), UnitctlError>
         });
     }
 
-    let control_socket = wait::wait_for_socket(cli).await?;
-    let client = UnitClient::new(control_socket);
+    let clients: Vec<_> = wait::wait_for_sockets(cli)
+        .await?
+        .into_iter()
+        .map(|sock| UnitClient::new(sock))
+        .collect();
+
     let mut results = vec![];
     for i in WalkDir::new(directory)
         .follow_links(true)
@@ -60,7 +64,9 @@ pub async fn cmd(cli: &UnitCtl, directory: &PathBuf) -> Result<(), UnitctlError>
         .filter_map(Result::ok)
         .filter(|e| !e.path().is_dir())
     {
-        results.push(process_entry(i, &client).await);
+        for client in &clients {
+            results.push(process_entry(i.clone(), client).await);
+        }
     }
 
     if results.iter().filter(|r| r.is_err()).count() == results.len() {

diff --git a/tools/unitctl/unitctl/src/cmd/listeners.rs b/tools/unitctl/unitctl/src/cmd/listeners.rs
@@ -1,14 +1,23 @@
 use crate::unitctl::UnitCtl;
 use crate::wait;
-use crate::{OutputFormat, UnitctlError};
+use crate::{OutputFormat, UnitctlError, eprint_error};
 use unit_client_rs::unit_client::UnitClient;
 
 pub async fn cmd(cli: &UnitCtl, output_format: OutputFormat) -> Result<(), UnitctlError> {
-    let control_socket = wait::wait_for_socket(cli).await?;
-    let client = UnitClient::new(control_socket);
-    client
-        .listeners()
-        .await
-        .map_err(|e| UnitctlError::UnitClientError { source: *e })
-        .and_then(|response| output_format.write_to_stdout(&response))
+    let socks = wait::wait_for_sockets(cli)
+        .await?;
+    let clients = socks.iter()
+        .map(|sock| UnitClient::new(sock.clone()));
+
+    for client in clients {
+        let _ = client.listeners()
+            .await
+            .map_err(|e| {
+                let err = UnitctlError::UnitClientError { source: *e };
+                eprint_error(&err);
+                std::process::exit(err.exit_code());
+            })
+            .and_then(|response| output_format.write_to_stdout(&response));
+    }
+    Ok(())
 }
diff --git a/tools/unitctl/unitctl/src/cmd/save.rs b/tools/unitctl/unitctl/src/cmd/save.rs
@@ -2,6 +2,7 @@ use crate::unitctl::UnitCtl;
 use crate::wait;
 use crate::UnitctlError;
 use crate::requests::send_empty_body_deserialize_response;
+use crate::unitctl_error::ControlSocketErrorKind;
 use unit_client_rs::unit_client::UnitClient;
 use tar::{Builder, Header};
 use std::fs::File;
@@ -12,13 +13,21 @@ pub async fn cmd(
     cli: &UnitCtl,
     filename: &String
 ) -> Result<(), UnitctlError> {
+    if cli.control_socket_addresses.is_some() &&
+        cli.control_socket_addresses.clone().unwrap().len() > 1 {
+            return Err(UnitctlError::ControlSocketError{
+                kind: ControlSocketErrorKind::General,
+                message: "too many control sockets. specify at most one.".to_string(),
+            });
+        }
+
+    let mut control_sockets = wait::wait_for_sockets(cli).await?;
+    let client = UnitClient::new(control_sockets.pop().unwrap());
+
     if !filename.ends_with(".tar") {
         eprintln!("Warning: writing uncompressed tarball to {}", filename);
     }
 
-    let control_socket = wait::wait_for_socket(cli).await?;
-    let client = UnitClient::new(control_socket);
-
     let config_res = serde_json::to_string_pretty(
         &send_empty_body_deserialize_response(&client, "GET", "/config").await?
     );

diff --git a/tools/unitctl/unitctl/src/cmd/status.rs b/tools/unitctl/unitctl/src/cmd/status.rs
@@ -1,14 +1,23 @@
 use crate::unitctl::UnitCtl;
 use crate::wait;
-use crate::{OutputFormat, UnitctlError};
+use crate::{OutputFormat, UnitctlError, eprint_error};
 use unit_client_rs::unit_client::UnitClient;
 
 pub async fn cmd(cli: &UnitCtl, output_format: OutputFormat) -> Result<(), UnitctlError> {
-    let control_socket = wait::wait_for_socket(cli).await?;
-    let client = UnitClient::new(control_socket);
-    client
-        .status()
-        .await
-        .map_err(|e| UnitctlError::UnitClientError { source: *e })
-        .and_then(|response| output_format.write_to_stdout(&response))
+    let socks = wait::wait_for_sockets(cli)
+        .await?;
+    let clients = socks.iter()
+        .map(|sock| UnitClient::new(sock.clone()));
+
+    for client in clients {
+        let _ = client.status()
+            .await
+            .map_err(|e| {
+                let err = UnitctlError::UnitClientError { source: *e };
+                eprint_error(&err);
+                std::process::exit(err.exit_code());
+            })
+            .and_then(|response| output_format.write_to_stdout(&response));
+    }
+    Ok(())
 }