// SPDX-License-Identifier: GPL-2.0 /* * Test module for stress and performance analysis of workqueue. * * Benchmarks queue_work() throughput on an unbound workqueue to measure * pool->lock contention under different affinity scope configurations * (e.g., cache vs cache_shard). * * The affinity scope is changed between runs via the workqueue's sysfs * affinity_scope attribute (WQ_SYSFS). * * Copyright (c) 2026 Meta Platforms, Inc. and affiliates * Copyright (c) 2026 Breno Leitao <leitao@debian.org> * */#include <linux/init.h>#include <linux/kernel.h>#include <linux/module.h>#include <linux/workqueue.h>#include <linux/kthread.h>#include <linux/moduleparam.h>#include <linux/completion.h>#include <linux/atomic.h>#include <linux/slab.h>#include <linux/ktime.h>#include <linux/cpumask.h>#include <linux/sched.h>#include <linux/sort.h>#include <linux/fs.h> #define WQ_NAME "bench_wq"#define SCOPE_PATH "/sys/bus/workqueue/devices/" WQ_NAME "/affinity_scope" static int nr_threads;module_param(nr_threads, int, 0444);MODULE_PARM_DESC(nr_threads,		 "Number of threads to spawn (default: 0 = num_online_cpus())"); static int wq_items = 50000;module_param(wq_items, int, 0444);MODULE_PARM_DESC(wq_items,		 "Number of work items each thread queues (default: 50000)"); static struct workqueue_struct *bench_wq;static atomic_t threads_done;static DECLARE_COMPLETION(start_comp);static DECLARE_COMPLETION(all_done_comp); struct thread_ctx {	struct completion work_done;	struct work_struct work;	u64 *latencies;	int cpu;	int items;}; static void bench_work_fn(struct work_struct *work){	struct thread_ctx *ctx = container_of(work, struct thread_ctx, work); 	complete(&ctx->work_done);} static int bench_kthread_fn(void *data){	struct thread_ctx *ctx = data;	ktime_t t_start, t_end;	int i; 	/* Wait for all threads to be ready */	wait_for_completion(&start_comp); 	if (kthread_should_stop())		return 0; 	for (i = 0; i < ctx->items; i++) {		reinit_completion(&ctx->work_done);		INIT_WORK(&ctx->work, bench_work_fn); 		t_start = ktime_get();		queue_work(bench_wq, &ctx->work);		t_end = ktime_get(); 		ctx->latencies[i] = ktime_to_ns(ktime_sub(t_end, t_start));		wait_for_completion(&ctx->work_done);	} 	if (atomic_dec_and_test(&threads_done))		complete(&all_done_comp); 	/*	 * Wait for kthread_stop() so the module text isn't freed	 * while we're still executing.	 */	while (!kthread_should_stop())		schedule(); 	return 0;} static int cmp_u64(const void *a, const void *b){	u64 va = *(const u64 *)a;	u64 vb = *(const u64 *)b; 	if (va < vb)		return -1;	if (va > vb)		return 1;	return 0;} static int __init set_affn_scope(const char *scope){	struct file *f;	loff_t pos = 0;	ssize_t ret; 	f = filp_open(SCOPE_PATH, O_WRONLY, 0);	if (IS_ERR(f)) {		pr_err("test_workqueue: open %s failed: %ld\n",		       SCOPE_PATH, PTR_ERR(f));		return PTR_ERR(f);	} 	ret = kernel_write(f, scope, strlen(scope), &pos);	filp_close(f, NULL); 	if (ret < 0) {		pr_err("test_workqueue: write '%s' failed: %zd\n", scope, ret);		return ret;	} 	return 0;} static int __init run_bench(int n_threads, const char *scope, const char *label){	struct task_struct **tasks;	unsigned long total_items;	struct thread_ctx *ctxs;	u64 *all_latencies;	ktime_t start, end;	int cpu, i, j, ret;	s64 elapsed_us; 	ret = set_affn_scope(scope);	if (ret)		return ret; 	ctxs = kcalloc(n_threads, sizeof(*ctxs), GFP_KERNEL);	if (!ctxs)		return -ENOMEM; 	tasks = kcalloc(n_threads, sizeof(*tasks), GFP_KERNEL);	if (!tasks) {		kfree(ctxs);		return -ENOMEM;	} 	total_items = (unsigned long)n_threads * wq_items;	all_latencies = kvmalloc_array(total_items, sizeof(u64), GFP_KERNEL);	if (!all_latencies) {		kfree(tasks);		kfree(ctxs);		return -ENOMEM;	} 	/* Allocate per-thread latency arrays */	for (i = 0; i < n_threads; i++) {		ctxs[i].latencies = kvmalloc_array(wq_items, sizeof(u64),						   GFP_KERNEL);		if (!ctxs[i].latencies) {			while (--i >= 0)				kvfree(ctxs[i].latencies);			kvfree(all_latencies);			kfree(tasks);			kfree(ctxs);			return -ENOMEM;		}	} 	atomic_set(&threads_done, n_threads);	reinit_completion(&all_done_comp);	reinit_completion(&start_comp); 	/* Create kthreads, each bound to a different online CPU */	i = 0;	for_each_online_cpu(cpu) {		if (i >= n_threads)			break; 		ctxs[i].cpu = cpu;		ctxs[i].items = wq_items;		init_completion(&ctxs[i].work_done); 		tasks[i] = kthread_create(bench_kthread_fn, &ctxs[i],					  "wq_bench/%d", cpu);		if (IS_ERR(tasks[i])) {			ret = PTR_ERR(tasks[i]);			pr_err("test_workqueue: failed to create kthread %d: %d\n",			       i, ret);			/* Unblock threads waiting on start_comp before stopping them */			complete_all(&start_comp);			while (--i >= 0)				kthread_stop(tasks[i]);			goto out_free;		} 		kthread_bind(tasks[i], cpu);		wake_up_process(tasks[i]);		i++;	} 	/* Start timing and release all threads */	start = ktime_get();	complete_all(&start_comp); 	/* Wait for all threads to finish the benchmark */	wait_for_completion(&all_done_comp); 	/* Drain any remaining work */	flush_workqueue(bench_wq); 	/* Ensure all kthreads have fully exited before module memory is freed */	for (i = 0; i < n_threads; i++)		kthread_stop(tasks[i]); 	end = ktime_get();	elapsed_us = ktime_us_delta(end, start); 	/* Merge all per-thread latencies and sort for percentile calculation */	j = 0;	for (i = 0; i < n_threads; i++) {		memcpy(&all_latencies[j], ctxs[i].latencies,		       wq_items * sizeof(u64));		j += wq_items;	} 	sort(all_latencies, total_items, sizeof(u64), cmp_u64, NULL); 	pr_info("test_workqueue:   %-16s %llu items/sec\tp50=%llu\tp90=%llu\tp95=%llu ns\n",		label,		elapsed_us ? div_u64(total_items * 1000000ULL, elapsed_us) : 0,		all_latencies[total_items * 50 / 100],		all_latencies[total_items * 90 / 100],		all_latencies[total_items * 95 / 100]); 	ret = 0;out_free:	for (i = 0; i < n_threads; i++)		kvfree(ctxs[i].latencies);	kvfree(all_latencies);	kfree(tasks);	kfree(ctxs); 	return ret;} static const char * const bench_scopes[] = {	"cpu", "smt", "cache_shard", "cache", "numa", "system",}; static int __init test_workqueue_init(void){	int n_threads = min(nr_threads ?: num_online_cpus(), num_online_cpus());	int i; 	if (wq_items <= 0) {		pr_err("test_workqueue: wq_items must be > 0\n");		return -EINVAL;	} 	bench_wq = alloc_workqueue(WQ_NAME, WQ_UNBOUND | WQ_SYSFS, 0);	if (!bench_wq)		return -ENOMEM; 	pr_info("test_workqueue: running %d threads, %d items/thread\n",		n_threads, wq_items); 	for (i = 0; i < ARRAY_SIZE(bench_scopes); i++)		run_bench(n_threads, bench_scopes[i], bench_scopes[i]); 	destroy_workqueue(bench_wq); 	/* Return -EAGAIN so the module doesn't stay loaded after the benchmark */	return -EAGAIN;} module_init(test_workqueue_init);MODULE_AUTHOR("Breno Leitao <leitao@debian.org>");MODULE_DESCRIPTION("Stress/performance benchmark for workqueue subsystem");MODULE_LICENSE("GPL");