Add missing multi-threaded stress tests for extension restart scenarios (#4526)

* Initial plan

* Add new extension restart stress test scenarios

Co-authored-by: mikeagun <[email protected]>

* Apply clang-format to stress test code for consistency

Co-authored-by: mikeagun <[email protected]>

* Fix critical logic bug in stress test thread configuration

Co-authored-by: mikeagun <[email protected]>

* Fix thread creation timing and object table initialization in stress tests

Co-authored-by: mikeagun <[email protected]>

* Update code comments to remove scenario references and make them self-descriptive

Co-authored-by: mikeagun <[email protected]>

* Fix comment punctuation consistency - ensure all comments end with periods

Co-authored-by: dthaler <[email protected]>

* Revert unrelated formatting changes to keep PR focused on stress tests

Co-authored-by: dthaler <[email protected]>

* Wrap JIT stress tests with CONFIG_BPF_JIT_DISABLED ifdef guards

Co-authored-by: mikeagun <[email protected]>

---------

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: mikeagun <[email protected]>
Co-authored-by: D. Michael Agun <[email protected]>
Co-authored-by: dthaler <[email protected]>

Author: 3 people

tests/stress/km/stress_tests_km.cpp

@@ -30,9 +30,9 @@ static const std::map<std::string, test_program_attributes> _test_program_info =
 struct object_table_entry
 {
     std::unique_ptr<std::mutex> lock{nullptr};
-    _Guarded_by_(lock) bool available { true };
-    _Guarded_by_(lock) bpf_object_ptr object { nullptr };
-    _Guarded_by_(lock) bool loaded { false };
+    _Guarded_by_(lock) bool available{true};
+    _Guarded_by_(lock) bpf_object_ptr object{nullptr};
+    _Guarded_by_(lock) bool loaded{false};
     bool attach{false};
 
     // The following fields are for debugging this test itself.
@@ -1634,6 +1634,7 @@ _mt_bindmonitor_tail_call_invoke_program_test(
     WSACleanup();
 }
 
+#if !defined(CONFIG_BPF_JIT_DISABLED)
 TEST_CASE("jit_load_attach_detach_unload_random_v4_test", "[jit_mt_stress_test]")
 {
     // This test attempts to load the same JIT'ed ebpf program multiple times in different threads.  This test
@@ -1652,6 +1653,7 @@ TEST_CASE("jit_load_attach_detach_unload_random_v4_test", "[jit_mt_stress_test]"
     _print_test_control_info(local_test_control_info);
     _mt_prog_load_stress_test(EBPF_EXECUTION_JIT, local_test_control_info);
 }
+#endif // !defined(CONFIG_BPF_JIT_DISABLED)
 
 TEST_CASE("native_load_attach_detach_unload_random_v4_test", "[native_mt_stress_test]")
 {
@@ -1694,13 +1696,16 @@ TEST_CASE("native_unique_load_attach_detach_unload_random_v4_test", "[native_mt_
 
 TEST_CASE("native_invoke_v4_v6_programs_restart_extension_test", "[native_mt_stress_test]")
 {
+    // Multi-threaded stress test where each thread loads different programs with extension restart.
     // Test layout:
-    // 1. Create 2 'monitor' threads:
-    //    - Thread #1 loads a native ebpf SOCK_ADDR program that attaches to CGROUP/CONNECT4.
+    // 1. Create 2 'monitor' threads, each loading a DIFFERENT native eBPF program:
+    //    - Thread #1 loads a native eBPF SOCK_ADDR program (cgroup_count_connect4.sys) that attaches to
+    //    CGROUP/CONNECT4.
     //      > This program monitors an IPv4 endpoint, 127.0.0.1:<target_port>. On every invocation, the program updates
     //        the count (TCP) 'connect' attempts in the 'connect4_count_map' map at its port.
-    //    - Thread #2 loads another native ebpf SOCK_ADDR program that attaches to CGROUP/CONNECT6.
-    //      > The behavior of this program is identical to that of the v4 program (loaded by thread #1), except it is
+    //    - Thread #2 loads a DIFFERENT native eBPF SOCK_ADDR program (cgroup_count_connect6.sys) that attaches to
+    //    CGROUP/CONNECT6.
+    //      > The behavior of this program is similar to the v4 program (loaded by thread #1), except it is
     //        invoked for IPv6 connection attempts ([::1]:<target_port>).
     //
     // 2  Until the end of test, each test thread will:
@@ -1714,13 +1719,14 @@ TEST_CASE("native_invoke_v4_v6_programs_restart_extension_test", "[native_mt_str
     // 3. In parallel, start the 'extension restart' thread to continuously restart the netebpf extension
     //    (if specified on the command line).
     //
+    // NOTE: Each thread loads different programs to test eBPF component resiliency with different program instances.
     // NOTE: The '-tt', '-er' and the '-erd' command line parameters are not used by this test.
 
     _km_test_init();
     LOG_INFO("\nStarting test *** native_invoke_v4_v6_programs_restart_extension_test ***");
     test_control_info local_test_control_info = _global_test_control_info;
 
-    // This test needs only 2 threads (one per program).
+    // This test needs only 2 threads (one per different program).
     local_test_control_info.threads_count = 2;
 
     _print_test_control_info(local_test_control_info);
@@ -1772,6 +1778,7 @@ TEST_CASE("bindmonitor_tail_call_invoke_program_test", "[native_mt_stress_test]"
     _mt_bindmonitor_tail_call_invoke_program_test(EBPF_EXECUTION_NATIVE, local_test_control_info);
 }
 
+#if !defined(CONFIG_BPF_JIT_DISABLED)
 TEST_CASE("jit_unique_load_attach_detach_unload_random_v4_test", "[jit_mt_stress_test]")
 {
     // This test attempts to load a unique JIT ebpf program multiple times in different threads. Specifically:
@@ -1872,3 +1879,300 @@ TEST_CASE("jit_bindmonitor_tail_call_invoke_program_test", "[jit_mt_stress_test]
     _print_test_control_info(local_test_control_info);
     _mt_bindmonitor_tail_call_invoke_program_test(EBPF_EXECUTION_JIT, local_test_control_info);
 }
+#endif // !defined(CONFIG_BPF_JIT_DISABLED)
+
+static void
+_mt_load_stress_test_with_restart_timing(
+    ebpf_execution_type_t program_type, const test_control_info& test_control_info, bool start_restart_before_load)
+{
+    constexpr uint32_t OBJECT_TABLE_SIZE{64};
+    std::vector<object_table_entry> object_table(OBJECT_TABLE_SIZE);
+    for (uint32_t index = 0; auto& entry : object_table) {
+        entry.available = true;
+        entry.lock = std::make_unique<std::mutex>();
+        entry.object.reset();
+        entry.attach = !(index % 2) ? true : false;
+        entry.index = index++;
+        entry.reuse_count = 0;
+        entry.tag = 0xC001DEA1;
+    }
+
+    // We have 3 types of threads, so we need (test_control_info.threads_count * 3) total threads.
+    size_t total_threads = ((size_t)test_control_info.threads_count * 3);
+    std::vector<thread_context> thread_context_table(
+        total_threads, {{}, {}, false, {}, thread_role_type::ROLE_NOT_SET, 0, 0, 0, false, 0, 0, object_table});
+    std::vector<std::thread> test_thread_table(total_threads);
+
+    // Extension restart thread setup.
+    std::vector<std::string> extension_names;
+    std::vector<std::thread> extension_restart_thread_table;
+    std::vector<thread_context> extension_restart_thread_context_table;
+
+    // An incrementing 'compartment Id' to ensure that _each_ 'Attacher' thread gets a unique compartment id.
+    uint32_t compartment_id{1};
+
+    // Get the single program info - we only have one program in km tests.
+    const auto& program_info = _test_program_info.begin();
+    const auto& program_name = program_info->first;
+    const auto& program_attribs = program_info->second;
+    extension_names.push_back(program_attribs.extension_name);
+
+    // Configure context for each thread.
+    for (size_t i = 0; i < total_threads; i++) {
+        // First, prepare the context for this thread.
+        auto& context_entry = thread_context_table[i];
+        context_entry.program_name = program_name;
+
+        if (!(compartment_id % 3)) {
+            context_entry.role = thread_role_type::DESTROYER;
+        } else if (!(compartment_id % 2)) {
+            context_entry.role = thread_role_type::ATTACHER;
+        } else {
+            context_entry.role = thread_role_type::CREATOR;
+        }
+
+        if (program_type == EBPF_EXECUTION_NATIVE) {
+            context_entry.is_native_program = true;
+            if (test_control_info.use_unique_native_programs && context_entry.role == thread_role_type::CREATOR) {
+                // Create unique native programs for 'creator' threads only.
+                context_entry.file_name = _make_unique_file_copy(program_attribs.native_file_name);
+            } else {
+                context_entry.file_name = program_attribs.native_file_name;
+            }
+        } else {
+            context_entry.file_name = program_attribs.jit_file_name;
+        }
+        context_entry.thread_index = (uint32_t)i;
+        context_entry.compartment_id = compartment_id++;
+        context_entry.duration_minutes = test_control_info.duration_minutes;
+        context_entry.extension_restart_enabled = test_control_info.extension_restart_enabled;
+    }
+
+    // Handle extension restart timing.
+    if (test_control_info.extension_restart_enabled) {
+        if (start_restart_before_load) {
+            // Start extension restart immediately, then start load threads.
+            LOG_INFO("Starting extension restart BEFORE program loading threads");
+            configure_extension_restart(
+                test_control_info,
+                extension_names,
+                extension_restart_thread_table,
+                extension_restart_thread_context_table,
+                object_table);
+            // Small delay to ensure restart thread is running before we start load threads.
+            std::this_thread::sleep_for(std::chrono::milliseconds(500));
+        }
+    }
+
+    // Now create the program loading threads.
+    for (size_t i = 0; i < total_threads; i++) {
+        auto& context_entry = thread_context_table[i];
+        auto& thread_entry = test_thread_table[i];
+        thread_entry = std::move(std::thread(_test_thread_function, std::ref(context_entry)));
+    }
+
+    if (test_control_info.extension_restart_enabled && !start_restart_before_load) {
+        // Wait for programs to load, then start extension restart.
+        LOG_INFO("Waiting for programs to load, then starting extension restart");
+        std::this_thread::sleep_for(std::chrono::seconds(5)); // Give programs time to load and attach
+        configure_extension_restart(
+            test_control_info,
+            extension_names,
+            extension_restart_thread_table,
+            extension_restart_thread_context_table,
+            object_table);
+    }
+
+    wait_and_verify_test_threads(
+        test_control_info,
+        test_thread_table,
+        thread_context_table,
+        extension_restart_thread_table,
+        extension_restart_thread_context_table);
+}
+
+static void
+_mt_invoke_stress_test_multiple_programs(ebpf_execution_type_t program_type, const test_control_info& test_control_info)
+{
+    WSAData data{};
+    auto error = WSAStartup(MAKEWORD(2, 2), &data);
+    REQUIRE(error == 0);
+
+    // For testing different programs per thread, create multiple copies of programs.
+    constexpr uint32_t MAX_PROGRAM_COPIES = 4;
+
+    size_t actual_threads = std::min((size_t)test_control_info.threads_count, (size_t)MAX_PROGRAM_COPIES);
+    std::vector<object_table_entry> object_table(actual_threads);
+    for (uint32_t index = 0; auto& entry : object_table) {
+        entry.available = true;
+        entry.lock = std::make_unique<std::mutex>();
+        entry.object.reset();
+        entry.attach = !(index % 2) ? true : false;
+        entry.index = index++;
+        entry.reuse_count = 0;
+        entry.tag = 0xC001DEA1;
+    }
+
+    std::vector<thread_context> thread_context_table(
+        actual_threads, {{}, {}, false, {}, thread_role_type::ROLE_NOT_SET, 0, 0, 0, false, 0, 0, object_table});
+    std::vector<std::thread> test_thread_table(actual_threads);
+
+    for (uint32_t i = 0; i < actual_threads; i++) {
+        // First, prepare the context for this thread.
+        auto& context_entry = thread_context_table[i];
+
+        if (program_type == EBPF_EXECUTION_NATIVE) {
+            // For native programs, create unique file copies so each thread has a different program.
+            context_entry.file_name = _make_unique_file_copy("cgroup_sock_addr.sys");
+            context_entry.is_native_program = true;
+        } else {
+            // For JIT programs, all threads can load the same file (but get different instances).
+            context_entry.file_name = "cgroup_sock_addr.o";
+            context_entry.is_native_program = false;
+        }
+
+        context_entry.program_name = "cgroup_sock_addr";
+        context_entry.role = thread_role_type::CREATOR; // All threads are creators for this test
+        context_entry.thread_index = i;
+        context_entry.compartment_id = i + 1; // Unique compartment IDs
+        context_entry.duration_minutes = test_control_info.duration_minutes;
+        context_entry.extension_restart_enabled = test_control_info.extension_restart_enabled;
+    }
+
+    // Now create all the threads after context setup is complete.
+    for (uint32_t i = 0; i < actual_threads; i++) {
+        auto& context_entry = thread_context_table[i];
+        auto& thread_entry = test_thread_table[i];
+        thread_entry = std::move(std::thread(_test_thread_function, std::ref(context_entry)));
+    }
+
+    // If requested, start the 'extension stop-and-restart' thread for extension for this program type.
+    std::vector<std::string> extension_names = {"netebpfext"};
+    std::vector<std::thread> extension_restart_thread_table;
+    std::vector<thread_context> extension_restart_thread_context_table;
+    if (test_control_info.extension_restart_enabled) {
+        configure_extension_restart(
+            test_control_info,
+            extension_names,
+            extension_restart_thread_table,
+            extension_restart_thread_context_table,
+            object_table);
+    }
+
+    wait_and_verify_test_threads(
+        test_control_info,
+        test_thread_table,
+        thread_context_table,
+        extension_restart_thread_table,
+        extension_restart_thread_context_table);
+}
+
+#if !defined(CONFIG_BPF_JIT_DISABLED)
+TEST_CASE("load_attach_stress_test_restart_during_load_jit", "[jit_mt_stress_test]")
+{
+    // Test resiliency during program 'open + load + attach' sequence with extension restart.
+    // Starts extension restart immediately and then begins program loading in multiple threads.
+    // Tests JIT programs with multiple threads loading the same program.
+
+    _km_test_init();
+    LOG_INFO("\nStarting test *** load_attach_stress_test_restart_during_load_jit ***");
+    test_control_info local_test_control_info = _global_test_control_info;
+
+    // Enable extension restart if not already enabled.
+    local_test_control_info.extension_restart_enabled = true;
+
+    _print_test_control_info(local_test_control_info);
+    _mt_load_stress_test_with_restart_timing(EBPF_EXECUTION_JIT, local_test_control_info, true);
+}
+#endif // !defined(CONFIG_BPF_JIT_DISABLED)
+
+TEST_CASE("load_attach_stress_test_restart_during_load_native", "[native_mt_stress_test]")
+{
+    // Test resiliency during program 'open + load + attach' sequence with extension restart.
+    // Starts extension restart immediately and then begins program loading in multiple threads.
+    // Tests native programs with multiple threads loading copies of programs.
+
+    _km_test_init();
+    LOG_INFO("\nStarting test *** load_attach_stress_test_restart_during_load_native ***");
+    test_control_info local_test_control_info = _global_test_control_info;
+
+    // Enable extension restart and unique native programs for this test.
+    local_test_control_info.extension_restart_enabled = true;
+    local_test_control_info.use_unique_native_programs = true;
+
+    _print_test_control_info(local_test_control_info);
+    _mt_load_stress_test_with_restart_timing(EBPF_EXECUTION_NATIVE, local_test_control_info, true);
+}
+
+#if !defined(CONFIG_BPF_JIT_DISABLED)
+TEST_CASE("load_attach_stress_test_restart_after_load_jit", "[jit_mt_stress_test]")
+{
+    // Test resiliency after program 'open + load + attach' sequence with extension restart.
+    // Completes program loading first, then starts extension restart.
+    // Ensures loaded + attached programs continue to be invoked after extension restart.
+    // Tests JIT programs with multiple threads loading the same program.
+
+    _km_test_init();
+    LOG_INFO("\nStarting test *** load_attach_stress_test_restart_after_load_jit ***");
+    test_control_info local_test_control_info = _global_test_control_info;
+
+    // Enable extension restart if not already enabled.
+    local_test_control_info.extension_restart_enabled = true;
+
+    _print_test_control_info(local_test_control_info);
+    _mt_load_stress_test_with_restart_timing(EBPF_EXECUTION_JIT, local_test_control_info, false);
+}
+#endif // !defined(CONFIG_BPF_JIT_DISABLED)
+
+TEST_CASE("load_attach_stress_test_restart_after_load_native", "[native_mt_stress_test]")
+{
+    // Test resiliency after program 'open + load + attach' sequence with extension restart.
+    // Completes program loading first, then starts extension restart.
+    // Ensures loaded + attached programs continue to be invoked after extension restart.
+    // Tests native programs with multiple threads loading copies of programs.
+
+    _km_test_init();
+    LOG_INFO("\nStarting test *** load_attach_stress_test_restart_after_load_native ***");
+    test_control_info local_test_control_info = _global_test_control_info;
+
+    // Enable extension restart and unique native programs for this test.
+    local_test_control_info.extension_restart_enabled = true;
+    local_test_control_info.use_unique_native_programs = true;
+
+    _print_test_control_info(local_test_control_info);
+    _mt_load_stress_test_with_restart_timing(EBPF_EXECUTION_NATIVE, local_test_control_info, false);
+}
+
+#if !defined(CONFIG_BPF_JIT_DISABLED)
+TEST_CASE("invoke_different_programs_restart_extension_test_jit", "[jit_mt_stress_test]")
+{
+    // Multi-threaded stress test where each thread loads different programs with extension restart.
+    // Tests JIT programs with different programs in each thread.
+
+    _km_test_init();
+    LOG_INFO("\nStarting test *** invoke_different_programs_restart_extension_test_jit ***");
+    test_control_info local_test_control_info = _global_test_control_info;
+
+    // This test needs multiple threads for different programs.
+    local_test_control_info.threads_count = std::max(local_test_control_info.threads_count, 4u);
+
+    _print_test_control_info(local_test_control_info);
+    _mt_invoke_stress_test_multiple_programs(EBPF_EXECUTION_JIT, local_test_control_info);
+}
+#endif // !defined(CONFIG_BPF_JIT_DISABLED)
+
+TEST_CASE("invoke_different_programs_restart_extension_test_native", "[native_mt_stress_test]")
+{
+    // Multi-threaded stress test where each thread loads different programs with extension restart.
+    // Tests native programs with different programs in each thread.
+
+    _km_test_init();
+    LOG_INFO("\nStarting test *** invoke_different_programs_restart_extension_test_native ***");
+    test_control_info local_test_control_info = _global_test_control_info;
+
+    // This test needs multiple threads for different programs.
+    local_test_control_info.threads_count = std::max(local_test_control_info.threads_count, 4u);
+
+    _print_test_control_info(local_test_control_info);
+    _mt_invoke_stress_test_multiple_programs(EBPF_EXECUTION_NATIVE, local_test_control_info);
+}