[pbs-devel] [RFC PATCH proxmox-backup 2/2] server/state: add spawn_internal_task and use it for websockets

[Dominik Csapak <d.csapak@proxmox.com>]

is a helper to spawn an internal tokio task without it showing up
in the task list

it is still tracked for reload and notifies the last_worker_listeners

this enables the console to survive a reload of proxmox-backup-proxy

Signed-off-by: Dominik Csapak <d.csapak@proxmox.com>
---
i am not completely sure if this is the right way...
alternatively, we could of course have WorkerTasks that simply do not
show up in the task list... though for what we are doing here,
it is much overhead thats not really needed

 src/api2/node.rs    |  2 +-
 src/server/state.rs | 36 +++++++++++++++++++++++++++++-------
 2 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/src/api2/node.rs b/src/api2/node.rs
index b572a9c..b80a661 100644
--- a/src/api2/node.rs
+++ b/src/api2/node.rs
@@ -267,7 +267,7 @@ fn upgrade_to_websocket(
 
         let (ws, response) = WebSocket::new(parts.headers)?;
 
-        tokio::spawn(async move {
+        crate::server::spawn_internal_task(async move {
             let conn: Upgraded = match req_body.on_upgrade().map_err(Error::from).await {
                 Ok(upgraded) => upgraded,
                 _ => bail!("error"),
diff --git a/src/server/state.rs b/src/server/state.rs
index 41e910d..d073243 100644
--- a/src/server/state.rs
+++ b/src/server/state.rs
@@ -19,6 +19,7 @@ pub struct ServerState {
     pub shutdown_listeners: BroadcastData<()>,
     pub last_worker_listeners: BroadcastData<()>,
     pub worker_count: usize,
+    pub task_count: usize,
     pub reload_request: bool,
 }
 
@@ -28,6 +29,7 @@ lazy_static! {
         shutdown_listeners: BroadcastData::new(),
         last_worker_listeners: BroadcastData::new(),
         worker_count: 0,
+        task_count: 0,
         reload_request: false,
     });
 }
@@ -101,20 +103,40 @@ pub fn last_worker_future() ->  impl Future<Output = Result<(), Error>> {
 }
 
 pub fn set_worker_count(count: usize) {
+    SERVER_STATE.lock().unwrap().worker_count = count;
+
+    check_last_worker();
+}
+
+pub fn check_last_worker() {
     let mut data = SERVER_STATE.lock().unwrap();
-    data.worker_count = count;
 
-    if !(data.mode == ServerMode::Shutdown && data.worker_count == 0) { return; }
+    if !(data.mode == ServerMode::Shutdown && data.worker_count == 0 && data.task_count == 0) { return; }
 
     data.last_worker_listeners.notify_listeners(Ok(()));
 }
 
-
-pub fn check_last_worker() {
-
+/// Spawns a tokio task that will be tracked for reload
+/// and if it is finished, notify the last_worker_listener if we
+/// are in shutdown mode
+pub fn spawn_internal_task<T>(task: T)
+where
+    T: Future + Send + 'static,
+    T::Output: Send + 'static,
+{
     let mut data = SERVER_STATE.lock().unwrap();
+    data.task_count += 1;
 
-    if !(data.mode == ServerMode::Shutdown && data.worker_count == 0) { return; }
+    tokio::spawn(async move {
+        let _ = tokio::spawn(task).await; // ignore errors
 
-    data.last_worker_listeners.notify_listeners(Ok(()));
+        { // drop mutex
+            let mut data = SERVER_STATE.lock().unwrap();
+            if data.task_count > 0 {
+                data.task_count -= 1;
+            }
+        }
+
+        check_last_worker();
+    });
 }
-- 
2.20.1