From 709d7005ef7c3abcc2f4e896c76a153ec179d2a6 Mon Sep 17 00:00:00 2001
From: Nick Banks <nibanks@microsoft.com>
Date: Wed, 6 Nov 2024 09:43:02 +0000
Subject: [PATCH] Support Thread Affinity for Execution Config (#4614)

---
 src/core/worker.c               | 9 +++++++--
 src/cs/lib/msquic_generated.cs  | 1 +
 src/inc/msquic.h                | 1 +
 src/perf/lib/PerfClient.cpp     | 3 +--
 src/perf/lib/PerfClient.h       | 1 -
 src/perf/lib/SecNetPerf.h       | 1 +
 src/perf/lib/SecNetPerfMain.cpp | 7 +++++++
 src/platform/platform_worker.c  | 3 +++
 8 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/src/core/worker.c b/src/core/worker.c
index 24c687933a..6338834413 100644
--- a/src/core/worker.c
+++ b/src/core/worker.c
@@ -118,8 +118,13 @@ QuicWorkerInitialize(
             break;
         }
 
-        if (MsQuicLib.ExecutionConfig && MsQuicLib.ExecutionConfig->Flags & QUIC_EXECUTION_CONFIG_FLAG_HIGH_PRIORITY) {
-            ThreadFlags |= CXPLAT_THREAD_FLAG_HIGH_PRIORITY;
+        if (MsQuicLib.ExecutionConfig) {
+            if (MsQuicLib.ExecutionConfig->Flags & QUIC_EXECUTION_CONFIG_FLAG_HIGH_PRIORITY) {
+                ThreadFlags |= CXPLAT_THREAD_FLAG_HIGH_PRIORITY;
+            }
+            if (MsQuicLib.ExecutionConfig->Flags & QUIC_EXECUTION_CONFIG_FLAG_AFFINITIZE) {
+                ThreadFlags |= CXPLAT_THREAD_FLAG_SET_AFFINITIZE;
+            }
         }
 
         CXPLAT_THREAD_CONFIG ThreadConfig = {
diff --git a/src/cs/lib/msquic_generated.cs b/src/cs/lib/msquic_generated.cs
index cf1477ea76..264add5c3e 100644
--- a/src/cs/lib/msquic_generated.cs
+++ b/src/cs/lib/msquic_generated.cs
@@ -218,6 +218,7 @@ internal enum QUIC_EXECUTION_CONFIG_FLAGS
         XDP = 0x0004,
         NO_IDEAL_PROC = 0x0008,
         HIGH_PRIORITY = 0x0010,
+        AFFINITIZE = 0x0020,
     }
 
     internal unsafe partial struct QUIC_EXECUTION_CONFIG
diff --git a/src/inc/msquic.h b/src/inc/msquic.h
index d9c25b8933..5582f83c2e 100644
--- a/src/inc/msquic.h
+++ b/src/inc/msquic.h
@@ -272,6 +272,7 @@ typedef enum QUIC_EXECUTION_CONFIG_FLAGS {
     QUIC_EXECUTION_CONFIG_FLAG_XDP              = 0x0004,
     QUIC_EXECUTION_CONFIG_FLAG_NO_IDEAL_PROC    = 0x0008,
     QUIC_EXECUTION_CONFIG_FLAG_HIGH_PRIORITY    = 0x0010,
+    QUIC_EXECUTION_CONFIG_FLAG_AFFINITIZE       = 0x0020,
 #endif
 } QUIC_EXECUTION_CONFIG_FLAGS;
 
diff --git a/src/perf/lib/PerfClient.cpp b/src/perf/lib/PerfClient.cpp
index 6cd808fe3a..5ae2d86897 100644
--- a/src/perf/lib/PerfClient.cpp
+++ b/src/perf/lib/PerfClient.cpp
@@ -146,7 +146,6 @@ PerfClient::Init(
     WorkerCount = CxPlatProcCount();
     TryGetVariableUnitValue(argc, argv, "threads", &WorkerCount);
     TryGetVariableUnitValue(argc, argv, "workers", &WorkerCount);
-    TryGetValue(argc, argv, "affinitize", &AffinitizeWorkers);
 
 #ifdef QUIC_COMPARTMENT_ID
     TryGetValue(argc, argv, "comp", &CompartmentId);
@@ -326,7 +325,7 @@ PerfClient::Start(
     // Configure and start all the workers.
     //
     uint16_t ThreadFlags =
-        AffinitizeWorkers ?
+        PerfDefaultAffinitizeThreads ?
             (uint16_t)CXPLAT_THREAD_FLAG_SET_AFFINITIZE :
             (uint16_t)CXPLAT_THREAD_FLAG_SET_IDEAL_PROC;
     if (PerfDefaultHighPriority) {
diff --git a/src/perf/lib/PerfClient.h b/src/perf/lib/PerfClient.h
index 3522365683..f9fdd9bc2a 100644
--- a/src/perf/lib/PerfClient.h
+++ b/src/perf/lib/PerfClient.h
@@ -179,7 +179,6 @@ struct PerfClient {
     uint8_t IncrementTarget {TRUE};
     // Local execution parameters
     uint32_t WorkerCount;
-    uint8_t AffinitizeWorkers {FALSE};
     uint8_t SpecificLocalAddresses {FALSE};
 #ifdef QUIC_COMPARTMENT_ID
     uint16_t CompartmentId {UINT16_MAX};
diff --git a/src/perf/lib/SecNetPerf.h b/src/perf/lib/SecNetPerf.h
index 9590503fda..22f1d4f4da 100644
--- a/src/perf/lib/SecNetPerf.h
+++ b/src/perf/lib/SecNetPerf.h
@@ -53,6 +53,7 @@ extern QUIC_CONGESTION_CONTROL_ALGORITHM PerfDefaultCongestionControl;
 extern uint8_t PerfDefaultEcnEnabled;
 extern uint8_t PerfDefaultQeoAllowed;
 extern uint8_t PerfDefaultHighPriority;
+extern uint8_t PerfDefaultAffinitizeThreads;
 
 extern CXPLAT_DATAPATH* Datapath;
 
diff --git a/src/perf/lib/SecNetPerfMain.cpp b/src/perf/lib/SecNetPerfMain.cpp
index 3ba2f5486f..7e6f3ba3a8 100644
--- a/src/perf/lib/SecNetPerfMain.cpp
+++ b/src/perf/lib/SecNetPerfMain.cpp
@@ -28,6 +28,7 @@ QUIC_CONGESTION_CONTROL_ALGORITHM PerfDefaultCongestionControl = QUIC_CONGESTION
 uint8_t PerfDefaultEcnEnabled = false;
 uint8_t PerfDefaultQeoAllowed = false;
 uint8_t PerfDefaultHighPriority = false;
+uint8_t PerfDefaultAffinitizeThreads = false;
 
 #ifdef _KERNEL_MODE
 volatile int BufferCurrent;
@@ -181,6 +182,12 @@ QuicMainStart(
         Config->Flags |= QUIC_EXECUTION_CONFIG_FLAG_HIGH_PRIORITY;
         SetConfig = true;
     }
+
+    TryGetValue(argc, argv, "affinitize", &PerfDefaultAffinitizeThreads);
+    if (PerfDefaultHighPriority) {
+        Config->Flags |= QUIC_EXECUTION_CONFIG_FLAG_AFFINITIZE;
+        SetConfig = true;
+    }
 #endif // _KERNEL_MODE
 
     if (TryGetValue(argc, argv, "pollidle", &Config->PollingIdleTimeoutUs)) {
diff --git a/src/platform/platform_worker.c b/src/platform/platform_worker.c
index 00b7937cfa..ffe4c73d22 100644
--- a/src/platform/platform_worker.c
+++ b/src/platform/platform_worker.c
@@ -162,6 +162,9 @@ CxPlatWorkerPoolLazyStart(
         if (Config->Flags & QUIC_EXECUTION_CONFIG_FLAG_HIGH_PRIORITY) {
             ThreadFlags |= CXPLAT_THREAD_FLAG_HIGH_PRIORITY;
         }
+        if (Config->Flags & QUIC_EXECUTION_CONFIG_FLAG_AFFINITIZE) {
+            ThreadFlags |= CXPLAT_THREAD_FLAG_SET_AFFINITIZE;
+        }
     }
 
     CXPLAT_THREAD_CONFIG ThreadConfig = {