From caea8d611fa074d90ea52140224fd87598d2856b Mon Sep 17 00:00:00 2001
From: Dietmar Maurer <dietmar@proxmox.com>
Date: Thu, 2 Jul 2020 14:00:32 +0200
Subject: [PATCH] proxmox-backup-client: add benchmark command

This is just a start, We need to add more useful things here...
---
 src/bin/proxmox-backup-client.rs           | 16 +++--
 src/bin/proxmox_backup_client/benchmark.rs | 82 ++++++++++++++++++++++
 src/bin/proxmox_backup_client/mod.rs       |  2 +
 3 files changed, 96 insertions(+), 4 deletions(-)
 create mode 100644 src/bin/proxmox_backup_client/benchmark.rs
 create mode 100644 src/bin/proxmox_backup_client/mod.rs

diff --git a/src/bin/proxmox-backup-client.rs b/src/bin/proxmox-backup-client.rs
index 56ab9e11..dea13402 100644
--- a/src/bin/proxmox-backup-client.rs
+++ b/src/bin/proxmox-backup-client.rs
@@ -60,16 +60,19 @@ use proxmox_backup::backup::{
     Shell,
 };
 
+mod proxmox_backup_client;
+use proxmox_backup_client::*;
+
 const ENV_VAR_PBS_FINGERPRINT: &str = "PBS_FINGERPRINT";
 const ENV_VAR_PBS_PASSWORD: &str = "PBS_PASSWORD";
 
 
-const REPO_URL_SCHEMA: Schema = StringSchema::new("Repository URL.")
+pub const REPO_URL_SCHEMA: Schema = StringSchema::new("Repository URL.")
     .format(&BACKUP_REPO_URL)
     .max_length(256)
     .schema();
 
-const KEYFILE_SCHEMA: Schema = StringSchema::new(
+pub const KEYFILE_SCHEMA: Schema = StringSchema::new(
     "Path to encryption key. All data will be encrypted using this key.")
     .schema();
 
@@ -84,7 +87,7 @@ fn get_default_repository() -> Option<String> {
     std::env::var("PBS_REPOSITORY").ok()
 }
 
-fn extract_repository_from_value(
+pub fn extract_repository_from_value(
     param: &Value,
 ) -> Result<BackupRepository, Error> {
 
@@ -2429,6 +2432,10 @@ fn main() {
         .completion_cb("keyfile", tools::complete_file_name)
         .completion_cb("chunk-size", complete_chunk_size);
 
+    let benchmark_cmd_def = CliCommand::new(&API_METHOD_BENCHMARK)
+        .completion_cb("repository", complete_repository)
+        .completion_cb("keyfile", tools::complete_file_name);
+
     let upload_log_cmd_def = CliCommand::new(&API_METHOD_UPLOAD_LOG)
         .arg_param(&["snapshot", "logfile"])
         .completion_cb("snapshot", complete_backup_snapshot)
@@ -2518,7 +2525,8 @@ fn main() {
         .insert("key", key_mgmt_cli())
         .insert("mount", mount_cmd_def)
         .insert("catalog", catalog_mgmt_cli())
-        .insert("task", task_mgmt_cli());
+        .insert("task", task_mgmt_cli())
+        .insert("benchmark", benchmark_cmd_def);
 
     let rpcenv = CliEnvironment::new();
     run_cli_command(cmd_def, rpcenv, Some(|future| {
diff --git a/src/bin/proxmox_backup_client/benchmark.rs b/src/bin/proxmox_backup_client/benchmark.rs
new file mode 100644
index 00000000..4b1c8d4c
--- /dev/null
+++ b/src/bin/proxmox_backup_client/benchmark.rs
@@ -0,0 +1,82 @@
+use std::path::PathBuf;
+use std::sync::Arc;
+
+use anyhow::{Error};
+use serde_json::Value;
+use chrono::{TimeZone, Utc};
+
+use proxmox::api::{ApiMethod, RpcEnvironment};
+use proxmox::api::api;
+
+use proxmox_backup::backup::{
+   load_and_decrypt_key,
+   CryptConfig,
+
+};
+
+use proxmox_backup::client::*;
+
+use crate::{
+    KEYFILE_SCHEMA, REPO_URL_SCHEMA,
+    extract_repository_from_value,
+    get_encryption_key_password,
+    record_repository,
+    connect,
+};
+
+#[api(
+   input: {
+       properties: {
+           repository: {
+               schema: REPO_URL_SCHEMA,
+               optional: true,
+           },
+           keyfile: {
+               schema: KEYFILE_SCHEMA,
+               optional: true,
+           },
+       }
+   }
+)]
+/// Run benchmark tests
+pub async fn benchmark(
+    param: Value,
+    _info: &ApiMethod,
+    _rpcenv: &mut dyn RpcEnvironment,
+) -> Result<(), Error> {
+
+    let repo = extract_repository_from_value(&param)?;
+
+    let keyfile = param["keyfile"].as_str().map(PathBuf::from);
+
+    let crypt_config = match keyfile {
+        None => None,
+        Some(path) => {
+            let (key, _) = load_and_decrypt_key(&path, &get_encryption_key_password)?;
+            let crypt_config = CryptConfig::new(key)?;
+            Some(Arc::new(crypt_config))
+        }
+    };
+
+    let backup_time = Utc.timestamp(Utc::now().timestamp(), 0);
+
+    let client = connect(repo.host(), repo.user())?;
+    record_repository(&repo);
+
+    let client = BackupWriter::start(
+        client,
+        crypt_config.clone(),
+        repo.store(),
+        "host",
+        "benshmark",
+        backup_time,
+        false,
+    ).await?;
+
+    println!("Start upload speed test");
+    let speed = client.upload_speedtest().await?;
+
+    println!("Upload speed: {} MiB/s", speed);
+
+    Ok(())
+}
diff --git a/src/bin/proxmox_backup_client/mod.rs b/src/bin/proxmox_backup_client/mod.rs
new file mode 100644
index 00000000..d9d69c96
--- /dev/null
+++ b/src/bin/proxmox_backup_client/mod.rs
@@ -0,0 +1,2 @@
+mod benchmark;
+pub use benchmark::*;
-- 
2.39.2