selftest/mpmc_ring: don't starve other queues

1 files changed, 28 insertions, 21 deletions

diff --git a/src/selftest/mpmc_ring.h b/src/selftest/mpmc_ring.h
index 265fbc7..4d59b54 100644
--- a/src/selftest/mpmc_ring.h
+++ b/src/selftest/mpmc_ring.h
@@ -12,39 +12,44 @@
 #include <linux/workqueue.h>
 #include <linux/wait.h>
 
-#define THREADS_PRODUCER 2
-#define THREADS_CONSUMER 2
+#define THREADS_PRODUCER 20
+#define THREADS_CONSUMER 20
 #define ELEMENT_COUNT 100000000LL /* divisible by threads_{consumer,producer} */
 #define QUEUE_SIZE 1024
 
 #define EXPECTED_TOTAL ((ELEMENT_COUNT * (ELEMENT_COUNT + 1)) / 2)
-#define PER_PRODUCER (ELEMENT_COUNT/THREADS_PRODUCER)
-#define PER_CONSUMER (ELEMENT_COUNT/THREADS_CONSUMER)
+#define PER_PRODUCER (ELEMENT_COUNT / THREADS_PRODUCER)
+#define PER_CONSUMER (ELEMENT_COUNT / THREADS_CONSUMER)
 #define THREADS_TOTAL (THREADS_PRODUCER + THREADS_CONSUMER)
 
 struct worker_producer {
 	struct work_struct work;
+	struct workqueue_struct *wq;
 	struct mpmc_ptr_ring *ring;
+	uint64_t i;
 	int thread_num;
 };
 
 struct worker_consumer {
 	struct work_struct work;
+	struct workqueue_struct *wq;
 	struct mpmc_ptr_ring *ring;
-	int thread_num;
-	int64_t total;
-	int64_t count;
+	uint64_t total;
+	uint64_t count;
+	uint64_t i;
 };
 
 static __init void producer_function(struct work_struct *work)
 {
 	struct worker_producer *td = container_of(work, struct worker_producer, work);
-	uintptr_t i;
 
-	for (i = td->thread_num * PER_PRODUCER + 1; i <= (td->thread_num + 1) * PER_PRODUCER; ++i) {
-		while (mpmc_ptr_ring_produce(td->ring, (void *)i)) {
-			schedule();
-			/*pr_info("We have awoken (producer)");*/
+	if (!td->i)
+		td->i = td->thread_num * PER_PRODUCER + 1;
+
+	for (; td->i <= (td->thread_num + 1) * PER_PRODUCER; ++td->i) {
+		if (mpmc_ptr_ring_produce(td->ring, (void *)td->i)) {
+			queue_work_on(smp_processor_id(), td->wq, work);
+			return;
 		}
 	}
 }
@@ -52,14 +57,12 @@ static __init void producer_function(struct work_struct *work)
 static __init void consumer_function(struct work_struct *work)
 {
 	struct worker_consumer *td = container_of(work, struct worker_consumer, work);
-	int i;
 
-	for (i = 0; i < PER_CONSUMER; ++i) {
+	for (td->i = 0; td->i < PER_CONSUMER; ++td->i) {
 		uintptr_t value;
-		while (!(value = (uintptr_t)mpmc_ptr_ring_consume(td->ring))) {
-			schedule();
-			/*cpu_relax();*/
-			/*pr_info("We have awoken (consumer)");*/
+		if (!(value = (uintptr_t)mpmc_ptr_ring_consume(td->ring))) {
+			queue_work_on(smp_processor_id(), td->wq, work);
+			return;
 		}
 
 		td->total += value;
@@ -84,19 +87,23 @@ bool __init mpmc_ring_selftest(void)
 	BUG_ON(mpmc_ptr_ring_init(&ring, QUEUE_SIZE, GFP_KERNEL));
 
 	wq = alloc_workqueue("mpmc_ring_selftest", WQ_UNBOUND, 0);
+	BUG_ON(!wq);
 
 	for (i = 0; i < THREADS_PRODUCER; ++i) {
 		producers[i].ring = &ring;
+		producers[i].wq = wq;
 		producers[i].thread_num = i;
+		producers[i].i = 0;
 		INIT_WORK(&producers[i].work, producer_function);
 		queue_work_on(cpumask_next_online(&cpu), wq, &producers[i].work);
 	}
 
 	for (i = 0; i < THREADS_CONSUMER; ++i) {
 		consumers[i].ring = &ring;
-		consumers[i].thread_num = i;
+		consumers[i].wq = wq;
 		consumers[i].total = 0;
 		consumers[i].count = 0;
+		consumers[i].i = 0;
 		INIT_WORK(&consumers[i].work, consumer_function);
 		queue_work_on(cpumask_next_online(&cpu), wq, &consumers[i].work);
 	}
@@ -119,8 +126,8 @@ bool __init mpmc_ring_selftest(void)
 	}
 
 	pr_info("mpmc_ring self-test failed:");
-	pr_info("Count: %lld, expected: %lld", count, ELEMENT_COUNT);
-	pr_info("Total: %lld, expected: %lld", total, EXPECTED_TOTAL);
+	pr_info("Count: %llu, expected: %llu", count, ELEMENT_COUNT);
+	pr_info("Total: %llu, expected: %llu", total, EXPECTED_TOTAL);
 
 	return false;
 }