From mboxrd@z Thu Jan 1 00:00:00 1970 Return-Path: Received: from firstgate.proxmox.com (firstgate.proxmox.com [212.224.123.68]) (using TLSv1.3 with cipher TLS_AES_256_GCM_SHA384 (256/256 bits) key-exchange X25519 server-signature RSA-PSS (2048 bits)) (No client certificate requested) by lists.proxmox.com (Postfix) with ESMTPS id 233BD69A5B for ; Tue, 14 Sep 2021 12:50:18 +0200 (CEST) Received: from firstgate.proxmox.com (localhost [127.0.0.1]) by firstgate.proxmox.com (Proxmox) with ESMTP id 1AA06C175 for ; Tue, 14 Sep 2021 12:50:18 +0200 (CEST) Received: from proxmox-new.maurer-it.com (proxmox-new.maurer-it.com [94.136.29.106]) (using TLSv1.3 with cipher TLS_AES_256_GCM_SHA384 (256/256 bits) key-exchange X25519 server-signature RSA-PSS (2048 bits) server-digest SHA256) (No client certificate requested) by firstgate.proxmox.com (Proxmox) with ESMTPS id 7877BC16C for ; Tue, 14 Sep 2021 12:50:17 +0200 (CEST) Received: from proxmox-new.maurer-it.com (localhost.localdomain [127.0.0.1]) by proxmox-new.maurer-it.com (Proxmox) with ESMTP id 4A8AF4484F; Tue, 14 Sep 2021 12:50:17 +0200 (CEST) From: Dietmar Maurer To: pbs-devel@lists.proxmox.com Date: Tue, 14 Sep 2021 12:50:14 +0200 Message-Id: <20210914105014.3903432-1-dietmar@proxmox.com> X-Mailer: git-send-email 2.30.2 MIME-Version: 1.0 Content-Transfer-Encoding: 8bit X-SPAM-LEVEL: Spam detection results: 0 AWL 0.676 Adjusted score from AWL reputation of From: address BAYES_00 -1.9 Bayes spam probability is 0 to 1% KAM_DMARC_STATUS 0.01 Test Rule for DKIM or SPF Failure with Strict Alignment SPF_HELO_NONE 0.001 SPF: HELO does not publish an SPF Record SPF_PASS -0.001 SPF: sender matches SPF record URIBL_BLOCKED 0.001 ADMINISTRATOR NOTICE: The query to URIBL was blocked. See http://wiki.apache.org/spamassassin/DnsBlocklists#dnsbl-block for more information. [upid.pid] Subject: [pbs-devel] [PATCH] use tokio::task_local to store/access the current worker X-BeenThere: pbs-devel@lists.proxmox.com X-Mailman-Version: 2.1.29 Precedence: list List-Id: Proxmox Backup Server development discussion List-Unsubscribe: , List-Archive: List-Post: List-Help: List-Subscribe: , X-List-Received-Date: Tue, 14 Sep 2021 10:50:18 -0000 This is used to simplify logging: worker_log!() and worker_warn!() --- src/server/worker_task.rs | 79 +++++++++++++++++++++++++++++++-------- 1 file changed, 64 insertions(+), 15 deletions(-) diff --git a/src/server/worker_task.rs b/src/server/worker_task.rs index 2ef8ba9d..9a0173b3 100644 --- a/src/server/worker_task.rs +++ b/src/server/worker_task.rs @@ -39,6 +39,10 @@ lazy_static! { static ref WORKER_TASK_LIST: Mutex>> = Mutex::new(HashMap::new()); } +tokio::task_local! { + static CURRENT_WORKER: Arc; +} + /// checks if the task UPID refers to a worker from this process fn is_local_worker(upid: &UPID) -> bool { upid.pid == server::pid() && upid.pstart == server::pstart() @@ -664,10 +668,10 @@ impl WorkerTask { let worker = WorkerTask::new(worker_type, worker_id, auth_id, to_stdout)?; let upid_str = worker.upid.to_string(); let f = f(worker.clone()); - tokio::spawn(async move { + tokio::spawn(CURRENT_WORKER.scope(Arc::clone(&worker), async move { let result = f.await; worker.log_result(&result); - }); + })); Ok(upid_str) } @@ -686,22 +690,24 @@ impl WorkerTask { let upid_str = worker.upid.to_string(); let _child = std::thread::Builder::new().name(upid_str.clone()).spawn(move || { - let worker1 = worker.clone(); - let result = match std::panic::catch_unwind(move || f(worker1)) { - Ok(r) => r, - Err(panic) => { - match panic.downcast::<&str>() { - Ok(panic_msg) => { - Err(format_err!("worker panicked: {}", panic_msg)) - } - Err(_) => { - Err(format_err!("worker panicked: unknown type.")) + let worker1 = Arc::clone(&worker); + CURRENT_WORKER.sync_scope(Arc::clone(&worker), || { + let result = match std::panic::catch_unwind(move || f(worker1)) { + Ok(r) => r, + Err(panic) => { + match panic.downcast::<&str>() { + Ok(panic_msg) => { + Err(format_err!("worker panicked: {}", panic_msg)) + } + Err(_) => { + Err(format_err!("worker panicked: unknown type.")) + } } } - } - }; + }; - worker.log_result(&result); + worker.log_result(&result); + }); }); Ok(upid_str) @@ -841,3 +847,46 @@ pub async fn wait_for_local_worker(upid_str: &str) -> Result<(), Error> { } Ok(()) } + +/// Call worker.log() for the current worker +/// +/// This will simply use `println()` if called from ouside a worker. +pub fn worker_log_str>(msg: S) { + let msg: &str = msg.as_ref(); + CURRENT_WORKER.try_with(|worker| worker.log(msg)).unwrap_or_else(|_| { + println!("{}", msg); + }); +} + +/// Call worker.warn() for the current worker +/// +/// This will simply use `eprintln()` if called from ouside a worker. +pub fn worker_warn_str>(msg: S) { + let msg: &str = msg.as_ref(); + CURRENT_WORKER.try_with(|worker| worker.log(msg)).unwrap_or_else(|_| { + eprintln!("WARN: {}", msg); + }); +} + +/// Macro to call worker_log_str() with format +#[macro_export] +macro_rules! worker_log { + ($($arg:tt)*) => ({ + $crate::server::worker_log_str(format!($($arg)*)); + }) +} + +/// Macro to call worker_warn_str with format +#[macro_export] +macro_rules! worker_warn { + ($($arg:tt)*) => ({ + $crate::server::worker_warn_str(format!($($arg)*)); + }) +} + +pub fn worker_fail_on_abort() -> Result<(), Error> { + match CURRENT_WORKER.try_with(|worker| worker.fail_on_abort()) { + Ok(result) => result, + Err(_) => bail!("worker abort failed: not inside a worker!"), + } +} -- 2.30.2