From 61d7b5013cb69c51d00656cd226cf3c49889b376 Mon Sep 17 00:00:00 2001
From: Hannes Laimer <h.laimer@proxmox.com>
Date: Wed, 2 Sep 2020 11:41:22 +0200
Subject: [PATCH] add benchmark flag to backup creation for proper cleanup when
 running a benchmark

Signed-off-by: Hannes Laimer <h.laimer@proxmox.com>
Signed-off-by: Dietmar Maurer <dietmar@proxmox.com>
---
 src/api2/backup.rs                         | 28 ++++++++++++++++++----
 src/bin/proxmox-backup-client.rs           |  1 +
 src/bin/proxmox_backup_client/benchmark.rs |  1 +
 src/client/backup_writer.rs                |  4 +++-
 4 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/src/api2/backup.rs b/src/api2/backup.rs
index ad608d85..59f2792b 100644
--- a/src/api2/backup.rs
+++ b/src/api2/backup.rs
@@ -38,6 +38,7 @@ pub const API_METHOD_UPGRADE_BACKUP: ApiMethod = ApiMethod::new(
             ("backup-id", false, &BACKUP_ID_SCHEMA),
             ("backup-time", false, &BACKUP_TIME_SCHEMA),
             ("debug", true, &BooleanSchema::new("Enable verbose debug logging.").schema()),
+            ("benchmark", true, &BooleanSchema::new("Job is a benchmark (do not keep data).").schema()),
         ]),
     )
 ).access(
@@ -56,6 +57,7 @@ fn upgrade_to_backup_protocol(
 
 async move {
     let debug = param["debug"].as_bool().unwrap_or(false);
+    let benchmark = param["benchmark"].as_bool().unwrap_or(false);
 
     let userid: Userid = rpcenv.get_user().unwrap().parse()?;
 
@@ -90,11 +92,24 @@ async move {
 
     let backup_group = BackupGroup::new(backup_type, backup_id);
 
+    let worker_type = if backup_type == "host" && backup_id == "benchmark" {
+        if !benchmark {
+            bail!("unable to run benchmark without --benchmark flags");
+        }
+        "benchmark"
+    } else {
+        if benchmark {
+            bail!("benchmark flags is only allowed on 'host/benchmark'");
+        }
+        "backup"
+    };
+
     // lock backup group to only allow one backup per group at a time
     let (owner, _group_guard) = datastore.create_locked_backup_group(&backup_group, &userid)?;
 
     // permission check
-    if owner != userid { // only the owner is allowed to create additional snapshots
+    if owner != userid && worker_type != "benchmark" {
+        // only the owner is allowed to create additional snapshots
         bail!("backup owner check failed ({} != {})", userid, owner);
     }
 
@@ -116,14 +131,15 @@ async move {
     let (path, is_new, _snap_guard) = datastore.create_locked_backup_dir(&backup_dir)?;
     if !is_new { bail!("backup directory already exists."); }
 
-    WorkerTask::spawn("backup", Some(worker_id), userid.clone(), true, move |worker| {
+
+    WorkerTask::spawn(worker_type, Some(worker_id), userid.clone(), true, move |worker| {
         let mut env = BackupEnvironment::new(
             env_type, userid, worker.clone(), datastore, backup_dir);
 
         env.debug = debug;
         env.last_backup = last_backup;
 
-        env.log(format!("starting new backup on datastore '{}': {:?}", store, path));
+        env.log(format!("starting new {} on datastore '{}': {:?}", worker_type, store, path));
 
         let service = H2Service::new(env.clone(), worker.clone(), &BACKUP_API_ROUTER, debug);
 
@@ -160,7 +176,11 @@ async move {
                 req = req_fut => req,
                 abrt = abort_future => abrt,
             };
-
+            if benchmark {
+                env.log("benchmark finished successfully");
+                env.remove_backup()?;
+                return Ok(());
+            }
             match (res, env.ensure_finished()) {
                 (Ok(_), Ok(())) => {
                     env.log("backup finished successfully");
diff --git a/src/bin/proxmox-backup-client.rs b/src/bin/proxmox-backup-client.rs
index 9a6f309d..1e9cc680 100644
--- a/src/bin/proxmox-backup-client.rs
+++ b/src/bin/proxmox-backup-client.rs
@@ -1026,6 +1026,7 @@ async fn create_backup(
         &backup_id,
         backup_time,
         verbose,
+        false
     ).await?;
 
     let previous_manifest = if let Ok(previous_manifest) = client.download_previous_manifest().await {
diff --git a/src/bin/proxmox_backup_client/benchmark.rs b/src/bin/proxmox_backup_client/benchmark.rs
index 9355d39e..b7c67e7f 100644
--- a/src/bin/proxmox_backup_client/benchmark.rs
+++ b/src/bin/proxmox_backup_client/benchmark.rs
@@ -226,6 +226,7 @@ async fn test_upload_speed(
         "benchmark",
         backup_time,
         false,
+        true
     ).await?;
 
     if verbose { eprintln!("Start TLS speed test"); }
diff --git a/src/client/backup_writer.rs b/src/client/backup_writer.rs
index 35e9c0bd..64c3cf27 100644
--- a/src/client/backup_writer.rs
+++ b/src/client/backup_writer.rs
@@ -53,6 +53,7 @@ impl BackupWriter {
         backup_id: &str,
         backup_time: DateTime<Utc>,
         debug: bool,
+        benchmark: bool
     ) -> Result<Arc<BackupWriter>, Error> {
 
         let param = json!({
@@ -60,7 +61,8 @@ impl BackupWriter {
             "backup-id": backup_id,
             "backup-time": backup_time.timestamp(),
             "store": datastore,
-            "debug": debug
+            "debug": debug,
+            "benchmark": benchmark
         });
 
         let req = HttpClient::request_builder(