From 2f0f3e9979fd0bbd9a9a124697d190817d8c17f9 Mon Sep 17 00:00:00 2001
From: Thomas Lamprecht <t.lamprecht@proxmox.com>
Date: Mon, 14 Nov 2022 16:07:10 +0100
Subject: [PATCH] file restore: allow to pass dimm size via env

Signed-off-by: Thomas Lamprecht <t.lamprecht@proxmox.com>
---
 proxmox-file-restore/src/block_driver_qemu.rs | 12 ++++++++++--
 proxmox-file-restore/src/qemu_helper.rs       | 14 +++++---------
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/proxmox-file-restore/src/block_driver_qemu.rs b/proxmox-file-restore/src/block_driver_qemu.rs
index c915f3cf..f73e5938 100644
--- a/proxmox-file-restore/src/block_driver_qemu.rs
+++ b/proxmox-file-restore/src/block_driver_qemu.rs
@@ -19,7 +19,7 @@ use pbs_datastore::catalog::ArchiveEntry;
 
 use super::block_driver::*;
 use crate::get_user_run_dir;
-use crate::qemu_helper;
+use crate::qemu_helper::{self, MAX_MEMORY_DIMM_SIZE};
 
 const RESTORE_VM_MAP: &str = "restore-vm-map.json";
 
@@ -204,9 +204,17 @@ async fn handle_extra_guest_memory_needs(cid: i32, path: &[u8]) {
         Some("true") => (),
         _ => return, // this is opt-in
     }
+    let size = match var("PBS_FILE_RESTORE_MEM_HOTPLUG_SIZE_MB").map(|v| v.parse::<usize>()) {
+        Ok(Ok(size)) if size > MAX_MEMORY_DIMM_SIZE => {
+            log::warn!("limit memory request of {size} to {MAX_MEMORY_DIMM_SIZE}");
+            MAX_MEMORY_DIMM_SIZE
+        }
+        Ok(Ok(size)) => size,
+        _ => 256, // in practice this means a total of ~ 512 MB depending on disk count
+    };
 
     if path_is_zfs(path) {
-        if let Err(err) = qemu_helper::set_dynamic_memory(cid, None).await {
+        if let Err(err) = qemu_helper::hotplug_memory(cid, size).await {
             log::error!("could not increase memory: {err}");
         }
     }
diff --git a/proxmox-file-restore/src/qemu_helper.rs b/proxmox-file-restore/src/qemu_helper.rs
index d186a26b..e5d23f40 100644
--- a/proxmox-file-restore/src/qemu_helper.rs
+++ b/proxmox-file-restore/src/qemu_helper.rs
@@ -135,14 +135,10 @@ async fn send_qmp_request<T: AsyncBufRead + AsyncWrite + Unpin>(
     Ok(buf)
 }
 
-pub async fn set_dynamic_memory(cid: i32, target_memory: Option<usize>) -> Result<(), Error> {
-    let target_memory = match target_memory {
-        Some(size) if size > MAX_MEMORY_DIMM_SIZE => {
-            bail!("cannot set to {size}M, maximum is {DYNAMIC_MEMORY_MB}M")
-        }
-        Some(size) => size,
-        None => MAX_MEMORY_DIMM_SIZE,
-    };
+pub(crate) async fn hotplug_memory(cid: i32, dimm_mb: usize) -> Result<(), Error> {
+    if dimm_mb > MAX_MEMORY_DIMM_SIZE {
+        bail!("cannot set to {dimm_mb}M, maximum is {MAX_MEMORY_DIMM_SIZE}M");
+    }
 
     let path = format!("{QMP_SOCKET_PREFIX}{cid}.sock");
     let mut stream = tokio::io::BufStream::new(tokio::net::UnixStream::connect(path).await?);
@@ -155,7 +151,7 @@ pub async fn set_dynamic_memory(cid: i32, target_memory: Option<usize>) -> Resul
         "arguments": {
             "qom-type": "memory-backend-ram",
             "id": "mem0",
-            "size": target_memory * 1024 * 1024,
+            "size": dimm_mb * 1024 * 1024,
         }
     });
     let _ = send_qmp_request(&mut stream, &serde_json::to_string(&request)?).await?;