From 503210c3bffbc90a5fbcabf4c12a6e5fa9036d8a Mon Sep 17 00:00:00 2001
From: Jason Volk <jason@zemos.net>
Date: Thu, 19 Dec 2024 07:33:29 +0000
Subject: [PATCH] toggle worker_affinity feature from program argument

Signed-off-by: Jason Volk <jason@zemos.net>
---
 src/main/clap.rs    | 15 ++++++++++++++-
 src/main/runtime.rs | 19 +++++++++++++++----
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/src/main/clap.rs b/src/main/clap.rs
index 28577356..c447deae 100644
--- a/src/main/clap.rs
+++ b/src/main/clap.rs
@@ -2,7 +2,7 @@
 
 use std::path::PathBuf;
 
-use clap::Parser;
+use clap::{ArgAction, Parser};
 use conduwuit::{
 	config::{Figment, FigmentValue},
 	err, toml,
@@ -50,6 +50,19 @@ pub(crate) struct Args {
 	/// Set the histogram bucket count (tokio_unstable). Default is 20.
 	#[arg(long, hide(true), env = "CONDUWUIT_RUNTIME_HISTOGRAM_BUCKETS", default_value = "20")]
 	pub(crate) worker_histogram_buckets: usize,
+
+	/// Toggles worker affinity feature.
+	#[arg(
+		long,
+		hide(true),
+		env = "CONDUWUIT_RUNTIME_WORKER_AFFINITY",
+		action = ArgAction::Set,
+		num_args = 0..=1,
+		require_equals(false),
+		default_value = "true",
+		default_missing_value = "true",
+	)]
+	pub(crate) worker_affinity: bool,
 }
 
 /// Parse commandline arguments into structured data
diff --git a/src/main/runtime.rs b/src/main/runtime.rs
index 6fc405a7..ff8ab322 100644
--- a/src/main/runtime.rs
+++ b/src/main/runtime.rs
@@ -1,4 +1,4 @@
-use std::{thread, time::Duration};
+use std::{sync::OnceLock, thread, time::Duration};
 
 use conduwuit::Result;
 use tokio::runtime::Builder;
@@ -12,9 +12,14 @@ const GLOBAL_QUEUE_INTERVAL: u32 = 192;
 const KERNEL_QUEUE_INTERVAL: u32 = 256;
 const KERNEL_EVENTS_PER_TICK: usize = 512;
 
-pub(super) fn new(args: &Args) -> Result<tokio::runtime::Runtime> {
-	let mut builder = Builder::new_multi_thread();
+static WORKER_AFFINITY: OnceLock<bool> = OnceLock::new();
 
+pub(super) fn new(args: &Args) -> Result<tokio::runtime::Runtime> {
+	WORKER_AFFINITY
+		.set(args.worker_affinity)
+		.expect("set WORKER_AFFINITY from program argument");
+
+	let mut builder = Builder::new_multi_thread();
 	builder
 		.enable_io()
 		.enable_time()
@@ -63,7 +68,13 @@ fn enable_histogram(builder: &mut Builder, args: &Args) {
 )]
 fn thread_start() {
 	#[cfg(feature = "worker_affinity")]
-	set_worker_affinity();
+	if WORKER_AFFINITY
+		.get()
+		.copied()
+		.expect("WORKER_AFFINITY initialized by runtime::new()")
+	{
+		set_worker_affinity();
+	}
 }
 
 #[cfg(feature = "worker_affinity")]