[pbs-devel] [PATCH backup 5/7] proxy: implement 'reload-certificate' command

[Dietmar Maurer]

> I wish there was some nice form of a `select_loop!`-like helper...

Another way would be to avoid the select inside the loop, for example
by using an Atomic counter (cert is loaded on next accept, not immediately)

--- Use AtomicUsize ---

diff --git a/src/bin/proxmox-backup-proxy.rs b/src/bin/proxmox-backup-proxy.rs
index fc773459..8ecdacec 100644
--- a/src/bin/proxmox-backup-proxy.rs
+++ b/src/bin/proxmox-backup-proxy.rs
@@ -2,6 +2,7 @@ use std::sync::Arc;
 use std::path::{Path, PathBuf};
 use std::pin::Pin;
 use std::os::unix::io::AsRawFd;
+use std::sync::atomic::{AtomicUsize, Ordering};
 
 use anyhow::{bail, format_err, Error};
 use futures::*;
@@ -122,13 +123,13 @@ async fn run() -> Result<(), Error> {
     let acceptor = make_tls_acceptor()?;
 
     // to renew the acceptor we just let a command-socket handler trigger a Notify:
-    let notify_tls_cert_reload = Arc::new(tokio::sync::Notify::new());
+    let notify_tls_cert_reload = Arc::new(AtomicUsize::new(0));
     commando_sock.register_command(
         "reload-certificate".to_string(),
         {
             let notify_tls_cert_reload = Arc::clone(&notify_tls_cert_reload);
             move |_value| -> Result<_, Error> {
-                notify_tls_cert_reload.notify_one();
+                notify_tls_cert_reload.fetch_add(1, Ordering::SeqCst);
                 Ok(Value::Null)
             }
         },
@@ -201,7 +202,7 @@ fn accept_connections(
     listener: tokio::net::TcpListener,
     acceptor: Arc<openssl::ssl::SslAcceptor>,
     debug: bool,
-    notify_tls_cert_reload: Arc<tokio::sync::Notify>,
+    notify_tls_cert_reload: Arc<AtomicUsize>,
 ) -> tokio::sync::mpsc::Receiver<ClientStreamResult> {
 
     let (sender, receiver) = tokio::sync::mpsc::channel(MAX_PENDING_ACCEPTS);
@@ -216,47 +217,26 @@ async fn accept_connection(
     mut acceptor: Arc<openssl::ssl::SslAcceptor>,
     debug: bool,
     sender: tokio::sync::mpsc::Sender<ClientStreamResult>,
-    notify_tls_cert_reload: Arc<tokio::sync::Notify>,
+    notify_tls_cert_reload: Arc<AtomicUsize>,
 ) {
     let accept_counter = Arc::new(());
 
-    // Note that these must not be moved out/modified directly, they get pinned in the loop and
-    // "rearmed" after waking up:
-    let mut reload_tls = notify_tls_cert_reload.notified();
-    let mut accept = listener.accept();
-
     loop {
-        let sock;
-
-        // normally we'd use `tokio::pin!()` but we need this to happen outside the loop and we
-        // need to be able to "rearm" the futures:
-        let reload_tls_pin = unsafe { Pin::new_unchecked(&mut reload_tls) };
-        let accept_pin = unsafe { Pin::new_unchecked(&mut accept) };
-        tokio::select! {
-            _ = reload_tls_pin => {
-                // rearm the notification:
-                reload_tls = notify_tls_cert_reload.notified();
-
-                log::info!("reloading certificate");
-                match make_tls_acceptor() {
-                    Err(err) => eprintln!("error reloading certificate: {}", err),
-                    Ok(new_acceptor) => acceptor = new_acceptor,
-                }
+        let (sock, _addr) = match listener.accept().await {
+            Ok(conn) => conn,
+            Err(err) =>  {
+                eprintln!("error accepting tcp connection: {}", err);
                 continue;
             }
-            res = accept_pin => match res {
-                Err(err) => {
-                    eprintln!("error accepting tcp connection: {}", err);
-                    continue;
-                }
-                Ok((new_sock, _addr)) => {
-                    // rearm the accept future:
-                    accept = listener.accept();
+        };
 
-                    sock = new_sock;
-                }
+        if notify_tls_cert_reload.swap(0, Ordering::SeqCst) > 0 {
+            log::info!("reloading certificate");
+            match make_tls_acceptor() {
+                Err(err) =>  eprintln!("error reloading certificate: {}", err),
+                Ok(new_acceptor) => acceptor = new_acceptor,
             }
-        };
+        }
 
         sock.set_nodelay(true).unwrap();
         let _ = set_tcp_keepalive(sock.as_raw_fd(), PROXMOX_BACKUP_TCP_KEEPALIVE_TIME);