Fix 1DS telemetry deadlock during Windows process termination

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
+          "$<TARGET_FILE:onnxruntime_telemetry_shutdown_test_library>" ${mode}
+        WORKING_DIRECTORY "$<TARGET_FILE_DIR:onnxruntime_telemetry_shutdown_test>")
+      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)

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<std::mutex> lock(global_mutex_);
 
   global_register_count_--;

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 <typename Operation>

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

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<RtlDllShutdownInProgressFn>(::GetProcAddress(ntdll, "RtlDllShutdownInProgress"));
+  return is_shutting_down != nullptr && is_shutting_down();
+}
+
 const char* GetDefaultEncodedToken() {
   return "fllXSmJHWBYMX1cqWldIYU0KTA0IVnpWD0lkRloSWwhILF5WGmBDDUdEC1csC0NMZhZYWVEPBH9DWEhlR15MCl1UfgwPVWVDX0U=";
 }

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 <Windows.h>
+
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <cwchar>
+#include <memory>
+#include <string>
+#include <string_view>
+
+namespace {
+
+struct CloseWin32Handle {
+  void operator()(HANDLE handle) const { CloseHandle(handle); }
+};
+using UniqueHandle = std::unique_ptr<void, CloseWin32Handle>;
+
+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<RunTest>(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<HANDLE>(static_cast<uintptr_t>(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<DWORD>(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<uintptr_t>(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;
+}

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 <Windows.h>
+
+#include <cstdio>
+#include <cstdlib>
+#include <memory>
+#include <optional>
+
+#include <ILogConfiguration.hpp>
+#include <NullObjects.hpp>
+
+#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<Microsoft::Applications::Events::ILogConfiguration>();
+    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<OneDsTelemetry> 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();
+  }
+}

onnxruntime/test/platform/windows_telemetry/lifecycle/telemetry_shutdown_test_library.def

@@ -0,0 +1,3 @@
+LIBRARY onnxruntime_telemetry_shutdown_test_library
+EXPORTS
+  OrtTestTelemetryShutdown