From a555f28bf81dbc92b236cce4ec35884b906ed38b Mon Sep 17 00:00:00 2001
From: Piotr Zdunowski <piotr.zdunowski@intel.com>
Date: Tue, 16 Feb 2021 14:30:07 +0000
Subject: [PATCH] Don't unregister trim callback during process shutdown.

Resolves: NEO-4668

Signed-off-by: Piotr Zdunowski <piotr.zdunowski@intel.com>
---
 opencl/test/unit_test/mocks/mock_wddm.h              |  5 +++++
 .../windows/wddm_residency_controller_tests.cpp      | 12 +++++++++++-
 shared/source/os_interface/windows/wddm/wddm.cpp     | 12 +++++++++++-
 shared/source/os_interface/windows/wddm/wddm.h       |  2 ++
 4 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/opencl/test/unit_test/mocks/mock_wddm.h b/opencl/test/unit_test/mocks/mock_wddm.h
index 2cf80a0b3f..66b7869e19 100644
--- a/opencl/test/unit_test/mocks/mock_wddm.h
+++ b/opencl/test/unit_test/mocks/mock_wddm.h
@@ -118,6 +118,10 @@ class WddmMock : public Wddm {
         return Wddm::verifySharedHandle(osHandle);
     }
 
+    bool isShutdownInProgress() override {
+        return shutdownStatus;
+    };
+
     void resetGdi(Gdi *gdi);
 
     WddmMockHelpers::MakeResidentCall makeResidentResult;
@@ -158,6 +162,7 @@ class WddmMock : public Wddm {
     bool makeResidentStatus = true;
     bool callBaseMakeResident = true;
     bool callBaseCreatePagingLogger = true;
+    bool shutdownStatus = false;
 };
 
 struct GmockWddm : WddmMock {
diff --git a/opencl/test/unit_test/os_interface/windows/wddm_residency_controller_tests.cpp b/opencl/test/unit_test/os_interface/windows/wddm_residency_controller_tests.cpp
index 1d2b6862aa..cfabac19d5 100644
--- a/opencl/test/unit_test/os_interface/windows/wddm_residency_controller_tests.cpp
+++ b/opencl/test/unit_test/os_interface/windows/wddm_residency_controller_tests.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2018-2020 Intel Corporation
+ * Copyright (C) 2018-2021 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -262,6 +262,16 @@ TEST_F(WddmResidencyControllerWithGdiTest, givenWddmResidencyControllerWhenItIsD
     EXPECT_EQ(trimCallbackHandle, gdi->getUnregisterTrimNotificationArg().Handle);
 }
 
+TEST_F(WddmResidencyControllerWithGdiTest, givenWddmResidencyControllerWhenItIsDestructedDuringProcessShutdownThenDontUnregisterTrimCallback) {
+    wddm->shutdownStatus = true;
+
+    std::memset(&gdi->getUnregisterTrimNotificationArg(), 0, sizeof(D3DKMT_UNREGISTERTRIMNOTIFICATION));
+    mockOsContextWin.reset();
+
+    EXPECT_EQ(nullptr, gdi->getUnregisterTrimNotificationArg().Callback);
+    EXPECT_EQ(nullptr, gdi->getUnregisterTrimNotificationArg().Handle);
+}
+
 TEST_F(WddmResidencyControllerTest, givenUsedAllocationWhenCallingRemoveFromTrimCandidateListIfUsedThenRemoveIt) {
     MockWddmAllocation allocation;
     residencyController->addToTrimCandidateList(&allocation);
diff --git a/shared/source/os_interface/windows/wddm/wddm.cpp b/shared/source/os_interface/windows/wddm/wddm.cpp
index 5ec3dc75bb..4b78bafe48 100644
--- a/shared/source/os_interface/windows/wddm/wddm.cpp
+++ b/shared/source/os_interface/windows/wddm/wddm.cpp
@@ -959,10 +959,20 @@ VOID *Wddm::registerTrimCallback(PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback, Wd
     }
     return nullptr;
 }
+bool Wddm::isShutdownInProgress() {
+    auto handle = GetModuleHandleA("ntdll.dll");
+
+    if (!handle) {
+        return true;
+    }
+
+    auto RtlDllShutdownInProgress = reinterpret_cast<BOOLEAN(WINAPI *)()>(GetProcAddress(handle, "RtlDllShutdownInProgress"));
+    return RtlDllShutdownInProgress();
+}
 
 void Wddm::unregisterTrimCallback(PFND3DKMT_TRIMNOTIFICATIONCALLBACK callback, VOID *trimCallbackHandle) {
     DEBUG_BREAK_IF(callback == nullptr);
-    if (trimCallbackHandle == nullptr) {
+    if (trimCallbackHandle == nullptr || isShutdownInProgress()) {
         return;
     }
     D3DKMT_UNREGISTERTRIMNOTIFICATION unregisterTrimNotification;
diff --git a/shared/source/os_interface/windows/wddm/wddm.h b/shared/source/os_interface/windows/wddm/wddm.h
index 13ce15093a..44c4f8d96f 100644
--- a/shared/source/os_interface/windows/wddm/wddm.h
+++ b/shared/source/os_interface/windows/wddm/wddm.h
@@ -90,6 +90,8 @@ class Wddm {
     MOCKABLE_VIRTUAL void *virtualAlloc(void *inPtr, size_t size, unsigned long flags, unsigned long type);
     MOCKABLE_VIRTUAL int virtualFree(void *ptr, size_t size, unsigned long flags);
 
+    MOCKABLE_VIRTUAL bool isShutdownInProgress();
+
     bool configureDeviceAddressSpace();
 
     GT_SYSTEM_INFO *getGtSysInfo() const {