diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index 715037c806bf2..aec4bb0cb2274 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -3134,4 +3134,26 @@ if (onnxruntime_BUILD_SHARED_LIB AND NOT CMAKE_SYSTEM_NAME STREQUAL "Emscripten" ) endif() +if(WIN32 AND onnxruntime_USE_1DS_TELEMETRY) + if(NOT GDK_PLATFORM AND NOT CMAKE_SYSTEM_NAME STREQUAL "WindowsStore") + set(telemetry_lifecycle_test_dir "${TEST_SRC_DIR}/platform/windows_telemetry/lifecycle") + onnxruntime_add_shared_library(onnxruntime_telemetry_shutdown_test_library + "${telemetry_lifecycle_test_dir}/telemetry_shutdown_test_library.cc" + "${telemetry_lifecycle_test_dir}/telemetry_shutdown_test_library.def") + target_link_libraries(onnxruntime_telemetry_shutdown_test_library PRIVATE onnxruntime_common mat) + onnxruntime_add_executable(onnxruntime_telemetry_shutdown_test + "${telemetry_lifecycle_test_dir}/telemetry_shutdown_test.cc") + add_dependencies(onnxruntime_telemetry_shutdown_test onnxruntime_telemetry_shutdown_test_library) + foreach(mode process_exit normal_destruction explicit_unload) + add_test(NAME onnxruntime_telemetry_${mode}_test + COMMAND onnxruntime_telemetry_shutdown_test + "$" ${mode} + WORKING_DIRECTORY "$") + set_tests_properties(onnxruntime_telemetry_${mode}_test PROPERTIES TIMEOUT 15) + endforeach() + set_target_properties(onnxruntime_telemetry_shutdown_test_library onnxruntime_telemetry_shutdown_test + PROPERTIES FOLDER "ONNXRuntimeTest") + endif() +endif() + include(onnxruntime_fuzz_test.cmake) diff --git a/onnxruntime/core/platform/telemetry_1ds.cc b/onnxruntime/core/platform/telemetry_1ds.cc index 3d9a9d04cbbd8..fa11a2f1bde68 100644 --- a/onnxruntime/core/platform/telemetry_1ds.cc +++ b/onnxruntime/core/platform/telemetry_1ds.cc @@ -342,6 +342,14 @@ OneDsTelemetry::OneDsTelemetry(const Telemetry& local_telemetry) : OneDsTelemetr } OneDsTelemetry::~OneDsTelemetry() { +#if defined(_WIN32) + // ExitProcess terminates the other threads before running DLL destructors. + // The uploader cannot drain WinHTTP callbacks here, and a terminated thread + // may still own one of our locks. Normal destruction/FreeLibrary still cleans up. + if (telemetry_internal::IsDllShutdownInProgress()) { + return; + } +#endif std::lock_guard lock(global_mutex_); global_register_count_--; diff --git a/onnxruntime/core/platform/telemetry_1ds.h b/onnxruntime/core/platform/telemetry_1ds.h index 2d0eabfd48dc2..665d38acf2ffe 100644 --- a/onnxruntime/core/platform/telemetry_1ds.h +++ b/onnxruntime/core/platform/telemetry_1ds.h @@ -140,6 +140,8 @@ class OneDsTelemetry : public Telemetry { const std::string& lib_path) const override; private: + friend class OneDsTelemetryTestPeer; + static void ReportFailure(const char* operation_name, const char* message); template diff --git a/onnxruntime/core/platform/telemetry_1ds_platform.h b/onnxruntime/core/platform/telemetry_1ds_platform.h index 2cfdbd99dd0a8..311b80b19c553 100644 --- a/onnxruntime/core/platform/telemetry_1ds_platform.h +++ b/onnxruntime/core/platform/telemetry_1ds_platform.h @@ -13,4 +13,8 @@ const char* GetDefaultEncodedToken(); int32_t GetProcessorCount(); std::string GetCachePath(const std::string& directory); +#if defined(_WIN32) +bool IsDllShutdownInProgress(); +#endif + } // namespace onnxruntime::telemetry_internal diff --git a/onnxruntime/core/platform/windows/telemetry_1ds.cc b/onnxruntime/core/platform/windows/telemetry_1ds.cc index 88ad276a796a0..5fb8bd45bacb5 100644 --- a/onnxruntime/core/platform/windows/telemetry_1ds.cc +++ b/onnxruntime/core/platform/windows/telemetry_1ds.cc @@ -54,6 +54,15 @@ bool PrepareSampledProcessEvent(EventProperties& event) { namespace telemetry_internal { +bool IsDllShutdownInProgress() { + using RtlDllShutdownInProgressFn = BOOLEAN(NTAPI*)(); + const HMODULE ntdll = ::GetModuleHandleW(L"ntdll.dll"); + const auto is_shutting_down = + ntdll == nullptr ? nullptr + : reinterpret_cast(::GetProcAddress(ntdll, "RtlDllShutdownInProgress")); + return is_shutting_down != nullptr && is_shutting_down(); +} + const char* GetDefaultEncodedToken() { return "fllXSmJHWBYMX1cqWldIYU0KTA0IVnpWD0lkRloSWwhILF5WGmBDDUdEC1csC0NMZhZYWVEPBH9DWEhlR15MCl1UfgwPVWVDX0U="; } diff --git a/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test.cc b/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test.cc new file mode 100644 index 0000000000000..e04fd6357fc22 --- /dev/null +++ b/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test.cc @@ -0,0 +1,110 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +struct CloseWin32Handle { + void operator()(HANDLE handle) const { CloseHandle(handle); } +}; +using UniqueHandle = std::unique_ptr; + +int Fail(const char* message) { + std::fprintf(stderr, "%s (Win32 error: %lu)\n", message, GetLastError()); + return EXIT_FAILURE; +} + +int RunChild(const wchar_t* library_path, std::wstring_view mode, HANDLE verified) { + if (!SetEnvironmentVariableW(L"ORT_DISABLE_TELEMETRY", L"1")) { + return Fail("Could not disable network telemetry"); + } + const HMODULE library = LoadLibraryW(library_path); + if (!library) { + return Fail("LoadLibraryW failed"); + } + using RunTest = void (*)(bool, bool, HANDLE); + const auto run_test = reinterpret_cast(GetProcAddress(library, "OrtTestTelemetryShutdown")); + if (!run_test) { + return Fail("GetProcAddress failed"); + } + const bool process_exit = mode == L"process_exit"; + run_test(process_exit, mode == L"normal_destruction", verified); + const DWORD expected_signal = mode == L"normal_destruction" ? WAIT_OBJECT_0 : WAIT_TIMEOUT; + if (WaitForSingleObject(verified, 0) != expected_signal) { + return Fail("Telemetry destruction happened at the wrong lifecycle stage"); + } + if (process_exit) { + ExitProcess(EXIT_SUCCESS); + } + if (!FreeLibrary(library)) { + return Fail("FreeLibrary failed"); + } + return EXIT_SUCCESS; +} + +} // namespace + +int wmain(int argc, wchar_t** argv) { + if (argc == 5 && std::wstring_view(argv[1]) == L"--child") { + const auto verified = reinterpret_cast(static_cast(std::wcstoull(argv[4], nullptr, 10))); + return RunChild(argv[2], argv[3], verified); + } + if (argc != 3) { + return Fail("Expected test DLL path and lifecycle mode"); + } + const std::wstring_view mode(argv[2]); + if (mode != L"process_exit" && mode != L"normal_destruction" && mode != L"explicit_unload") { + return Fail("Unknown lifecycle mode"); + } + + // The parent checks a signal from the DLL destructor, so an early child exit + // cannot pass without actually exercising and verifying telemetry teardown. + SECURITY_ATTRIBUTES attributes{sizeof(SECURITY_ATTRIBUTES), nullptr, TRUE}; + UniqueHandle verified(CreateEventW(&attributes, TRUE, FALSE, nullptr)); + if (!verified) { + return Fail("CreateEventW failed"); + } + std::wstring executable(32768, L'\0'); + const DWORD length = GetModuleFileNameW(nullptr, executable.data(), static_cast(executable.size())); + if (length == 0 || length >= executable.size()) { + return Fail("GetModuleFileNameW failed"); + } + executable.resize(length); + std::wstring command = L"\"" + executable + L"\" --child \"" + argv[1] + L"\" " + argv[2] + L" " + + std::to_wstring(reinterpret_cast(verified.get())); + STARTUPINFOW startup{}; + startup.cb = sizeof(startup); + PROCESS_INFORMATION info{}; + if (!CreateProcessW(executable.c_str(), command.data(), nullptr, nullptr, TRUE, CREATE_NO_WINDOW, + nullptr, nullptr, &startup, &info)) { + return Fail("CreateProcessW failed"); + } + UniqueHandle process(info.hProcess); + UniqueHandle thread(info.hThread); + if (WaitForSingleObject(process.get(), 10000) != WAIT_OBJECT_0) { + TerminateProcess(process.get(), EXIT_FAILURE); + WaitForSingleObject(process.get(), 2000); + return Fail("Child process did not exit within 10 seconds"); + } + DWORD exit_code = EXIT_FAILURE; + if (!GetExitCodeProcess(process.get(), &exit_code)) { + return Fail("GetExitCodeProcess failed"); + } + if (exit_code != EXIT_SUCCESS) { + std::fprintf(stderr, "Child process failed with exit code %lu\n", exit_code); + return EXIT_FAILURE; + } + if (WaitForSingleObject(verified.get(), 0) != WAIT_OBJECT_0) { + return Fail("Child did not verify telemetry destruction"); + } + return EXIT_SUCCESS; +} diff --git a/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test_library.cc b/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test_library.cc new file mode 100644 index 0000000000000..0b7b15235ff48 --- /dev/null +++ b/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test_library.cc @@ -0,0 +1,123 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include + +#include +#include +#include +#include + +#include +#include + +#include "core/platform/telemetry_1ds.h" +#include "core/platform/telemetry_1ds_platform.h" + +namespace onnxruntime { +namespace { + +void Check(bool condition, const char* message) { + if (!condition) { + std::fprintf(stderr, "%s (Win32 error: %lu)\n", message, GetLastError()); + std::fflush(stderr); + // Also fail reliably from a DLL process-detach destructor. + TerminateProcess(GetCurrentProcess(), EXIT_FAILURE); + std::abort(); + } +} + +} // namespace + +class OneDsTelemetryTestPeer { + public: + static void Install(Microsoft::Applications::Events::ILogManager& manager) { + Check(OneDsTelemetry::global_register_count_ == 1, "Expected exactly one telemetry instance"); + Check(!OneDsTelemetry::log_manager_ && !OneDsTelemetry::config_, "Real uploader must be disabled"); + OneDsTelemetry::config_ = std::make_unique(); + OneDsTelemetry::log_manager_ = &manager; + } + + static bool IsClean() { + return OneDsTelemetry::global_register_count_ == 0 && + !OneDsTelemetry::log_manager_ && !OneDsTelemetry::config_; + } +}; + +namespace { + +class MockLogManager : public Microsoft::Applications::Events::NullLogManager { + public: + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(MockLogManager); + + explicit MockLogManager(bool pending) : callback_(CreateEventW(nullptr, TRUE, !pending, nullptr)) { + Check(callback_ != nullptr, "CreateEventW failed"); + } + + ~MockLogManager() override { CloseHandle(callback_); } + + Microsoft::Applications::Events::status_t Flush() override { + ++flush_count; + return Microsoft::Applications::Events::STATUS_SUCCESS; + } + + void FlushAndTeardown() override { + Check(flush_count == 1, "Expected Flush before FlushAndTeardown"); + ++teardown_count; + // A pending callback cannot arrive after ExitProcess has killed the other + // threads. Bound the mock's wait so a regression does not leave a stuck DLL. + Check(WaitForSingleObject(callback_, 2000) == WAIT_OBJECT_0, "Timed out draining the pending uploader"); + } + + int flush_count = 0; + int teardown_count = 0; + + private: + HANDLE callback_; +}; + +class TelemetryLifetime { + public: + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(TelemetryLifetime); + + TelemetryLifetime(bool process_exit, HANDLE verified) + : manager_(process_exit), process_exit_(process_exit), verified_(verified) { + telemetry_.emplace(); + OneDsTelemetryTestPeer::Install(manager_); + } + + ~TelemetryLifetime() { + if (telemetry_) { + Destroy(); + } + } + + void Destroy() { + Check(telemetry_internal::IsDllShutdownInProgress() == process_exit_, "Unexpected DLL shutdown state"); + telemetry_.reset(); + const int expected_calls = process_exit_ ? 0 : 1; + Check(manager_.flush_count == expected_calls, "Unexpected Flush count"); + Check(manager_.teardown_count == expected_calls, "Unexpected FlushAndTeardown count"); + if (!process_exit_) { + Check(OneDsTelemetryTestPeer::IsClean(), "Telemetry state was not cleaned up"); + } + Check(SetEvent(verified_) != FALSE, "Could not signal verified destruction"); + } + + private: + MockLogManager manager_; + std::optional telemetry_; + bool process_exit_; + HANDLE verified_; +}; + +} // namespace +} // namespace onnxruntime + +extern "C" void OrtTestTelemetryShutdown(bool process_exit, bool destroy_now, HANDLE verified) { + // A DLL-owned static exercises CRT destruction during real Windows shutdown. + static onnxruntime::TelemetryLifetime lifetime(process_exit, verified); + if (destroy_now) { + lifetime.Destroy(); + } +} diff --git a/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test_library.def b/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test_library.def new file mode 100644 index 0000000000000..517b8cf037088 --- /dev/null +++ b/onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test_library.def @@ -0,0 +1,3 @@ +LIBRARY onnxruntime_telemetry_shutdown_test_library +EXPORTS + OrtTestTelemetryShutdown