From aaffe75c25c18bd171b60edc2d5a6f287a9f9808 Mon Sep 17 00:00:00 2001
From: Paul Adenot <paul@paul.cx>
Date: Fri, 15 Mar 2024 13:30:55 +0100
Subject: [PATCH] Fetch avrt.dll functions dynamically again.

Turns out this can crash sometimes, and we need to be more careful.

This reverts commit 661c653c8650b2b1c485e86b5a5b4f6d0ed07d1b.
---
 CMakeLists.txt       |  2 +-
 src/cubeb_wasapi.cpp | 58 ++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 530e3a4..75955be 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -264,7 +264,7 @@ if(USE_WASAPI)
   target_sources(cubeb PRIVATE
     src/cubeb_wasapi.cpp)
   target_compile_definitions(cubeb PRIVATE USE_WASAPI)
-  target_link_libraries(cubeb PRIVATE avrt ole32 ksuser)
+  target_link_libraries(cubeb PRIVATE ole32 ksuser)
 endif()
 
 check_include_files("windows.h;mmsystem.h" USE_WINMM)
diff --git a/src/cubeb_wasapi.cpp b/src/cubeb_wasapi.cpp
index 01417a5..b5546c3 100644
--- a/src/cubeb_wasapi.cpp
+++ b/src/cubeb_wasapi.cpp
@@ -204,6 +204,11 @@ struct auto_stream_ref {
   cubeb_stream * stm;
 };
 
+using set_mm_thread_characteristics_function =
+    decltype(&AvSetMmThreadCharacteristicsA);
+using revert_mm_thread_characteristics_function =
+    decltype(&AvRevertMmThreadCharacteristics);
+
 extern cubeb_ops const wasapi_ops;
 
 static com_heap_ptr<wchar_t>
@@ -301,6 +306,13 @@ struct cubeb {
       nullptr;
   void * output_collection_changed_user_ptr = nullptr;
   UINT64 performance_counter_frequency;
+  /* Library dynamically opened to increase the render thread priority, and
+     the two function pointers we need. */
+  HMODULE mmcss_module = nullptr;
+  set_mm_thread_characteristics_function set_mm_thread_characteristics =
+      nullptr;
+  revert_mm_thread_characteristics_function revert_mm_thread_characteristics =
+      nullptr;
 };
 
 class wasapi_endpoint_notification_client;
@@ -1401,7 +1413,8 @@ static unsigned int __stdcall wasapi_stream_render_loop(LPVOID stream)
 
   /* We could consider using "Pro Audio" here for WebAudio and
      maybe WebRTC. */
-  mmcss_handle = AvSetMmThreadCharacteristicsA("Audio", &mmcss_task_index);
+  mmcss_handle =
+      stm->context->set_mm_thread_characteristics("Audio", &mmcss_task_index);
   if (!mmcss_handle) {
     /* This is not fatal, but we might glitch under heavy load. */
     LOG("Unable to use mmcss to bump the render thread priority: %lx",
@@ -1509,7 +1522,7 @@ static unsigned int __stdcall wasapi_stream_render_loop(LPVOID stream)
   }
 
   if (mmcss_handle) {
-    AvRevertMmThreadCharacteristics(mmcss_handle);
+    stm->context->revert_mm_thread_characteristics(mmcss_handle);
   }
 
   if (FAILED(hr)) {
@@ -1522,6 +1535,18 @@ static unsigned int __stdcall wasapi_stream_render_loop(LPVOID stream)
 void
 wasapi_destroy(cubeb * context);
 
+HANDLE WINAPI
+set_mm_thread_characteristics_noop(const char *, LPDWORD mmcss_task_index)
+{
+  return (HANDLE)1;
+}
+
+BOOL WINAPI
+revert_mm_thread_characteristics_noop(HANDLE mmcss_handle)
+{
+  return true;
+}
+
 HRESULT
 register_notification_client(cubeb_stream * stm)
 {
@@ -1757,6 +1782,31 @@ wasapi_init(cubeb ** context, char const * context_name)
     ctx->performance_counter_frequency = 0;
   }
 
+  ctx->mmcss_module = LoadLibraryW(L"Avrt.dll");
+
+  bool success = false;
+  if (ctx->mmcss_module) {
+    ctx->set_mm_thread_characteristics =
+        reinterpret_cast<set_mm_thread_characteristics_function>(
+            GetProcAddress(ctx->mmcss_module, "AvSetMmThreadCharacteristicsW"));
+    ctx->revert_mm_thread_characteristics =
+        reinterpret_cast<revert_mm_thread_characteristics_function>(
+            GetProcAddress(ctx->mmcss_module,
+                           "AvRevertMmThreadCharacteristics"));
+    success = ctx->set_mm_thread_characteristics &&
+              ctx->revert_mm_thread_characteristics;
+  }
+  if (!success) {
+    // This is not a fatal error, but we might end up glitching when
+    // the system is under high load.
+    LOG("Could not load avrt.dll or fetch AvSetMmThreadCharacteristicsW "
+        "AvRevertMmThreadCharacteristics: %lx",
+        GetLastError());
+    ctx->set_mm_thread_characteristics = &set_mm_thread_characteristics_noop;
+    ctx->revert_mm_thread_characteristics =
+        &revert_mm_thread_characteristics_noop;
+  }
+
   *context = ctx;
 
   return CUBEB_OK;
@@ -1813,6 +1863,10 @@ wasapi_destroy(cubeb * context)
     }
   }
 
+  if (context->mmcss_module) {
+    FreeLibrary(context->mmcss_module);
+  }
+
   delete context;
 }