Report leak when context is released

Author: zhaomaosu

source/loader/layers/sanitizer/asan_interceptor.cpp

@@ -146,6 +146,8 @@ ur_result_t enqueueMemSetShadow(ur_context_handle_t Context,
 
 } // namespace
 
+bool SanitizerInterceptor::AbnormalExit = false;
+
 SanitizerInterceptor::SanitizerInterceptor(logger::Logger &logger)
     : logger(logger) {
     if (Options(logger).MaxQuarantineSizeMB) {
@@ -159,15 +161,7 @@ SanitizerInterceptor::~SanitizerInterceptor() {
     DestroyShadowMemoryOnCPU();
     DestroyShadowMemoryOnPVC();
 
-    std::shared_lock<ur_shared_mutex> Guard(m_AllocationMapMutex);
-    bool HasLeak = false;
-    for (const auto &[_, AI] : m_AllocationMap) {
-        if (!AI->IsReleased) {
-            HasLeak = true;
-            ReportMemoryLeak(AI);
-        }
-    }
-    if (HasLeak) {
+    if (AbnormalExit) {
         exit(1);
     }
 }
@@ -887,6 +881,40 @@ SanitizerInterceptor::findAllocInfoByAddress(uptr Address) {
     return It;
 }
 
+std::vector<AllocationIterator>
+SanitizerInterceptor::findAllocInfoByContext(ur_context_handle_t Context) {
+    std::shared_lock<ur_shared_mutex> Guard(m_AllocationMapMutex);
+    std::vector<AllocationIterator> AllocInfos;
+    for (auto It = m_AllocationMap.begin(); It != m_AllocationMap.end(); It++) {
+        const auto &[_, AI] = *It;
+        if (AI->Context == Context) {
+            AllocInfos.emplace_back(It);
+        }
+    }
+    return AllocInfos;
+}
+
+ContextInfo::~ContextInfo() {
+    [[maybe_unused]] auto Result =
+        getContext()->urDdiTable.Context.pfnRelease(Handle);
+    assert(Result == UR_RESULT_SUCCESS);
+
+    bool HasLeak = false;
+    std::vector<AllocationIterator> AllocInfos =
+        getContext()->interceptor->findAllocInfoByContext(Handle);
+    for (const auto &It : AllocInfos) {
+        const auto &[_, AI] = *It;
+        if (AI->IsReleased) {
+            ReportMemoryLeak(AI);
+            HasLeak = true;
+        }
+    }
+
+    if (HasLeak) {
+        SanitizerInterceptor::AbnormalExit = true;
+    }
+}
+
 ur_result_t USMLaunchInfo::initialize() {
     UR_CALL(getContext()->urDdiTable.Context.pfnRetain(Context));
     UR_CALL(getContext()->urDdiTable.Device.pfnRetain(Device));

source/loader/layers/sanitizer/asan_interceptor.hpp

@@ -119,11 +119,7 @@ struct ContextInfo {
         assert(Result == UR_RESULT_SUCCESS);
     }
 
-    ~ContextInfo() {
-        [[maybe_unused]] auto Result =
-            getContext()->urDdiTable.Context.pfnRelease(Handle);
-        assert(Result == UR_RESULT_SUCCESS);
-    }
+    ~ContextInfo();
 
     void insertAllocInfo(const std::vector<ur_device_handle_t> &Devices,
                          std::shared_ptr<AllocInfo> &AI) {
@@ -169,6 +165,8 @@ struct DeviceGlobalInfo {
 
 class SanitizerInterceptor {
   public:
+    static bool AbnormalExit;
+
     explicit SanitizerInterceptor(logger::Logger &logger);
 
     ~SanitizerInterceptor();
@@ -208,6 +206,9 @@ class SanitizerInterceptor {
 
     std::optional<AllocationIterator> findAllocInfoByAddress(uptr Address);
 
+    std::vector<AllocationIterator>
+    findAllocInfoByContext(ur_context_handle_t Context);
+
     std::shared_ptr<ContextInfo> getContextInfo(ur_context_handle_t Context) {
         std::shared_lock<ur_shared_mutex> Guard(m_ContextMapMutex);
         assert(m_ContextMap.find(Context) != m_ContextMap.end());

source/loader/layers/sanitizer/asan_report.cpp

@@ -77,10 +77,10 @@ void ReportDoubleFree(uptr Addr, const StackTrace &Stack,
 }
 
 void ReportMemoryLeak(const std::shared_ptr<AllocInfo> &AI) {
-    context.logger.always(
+    getContext()->logger.always(
         "\n====ERROR: DeviceSanitizer: detected memory leaks of {}",
         ToString(AI->Type));
-    context.logger.always("Direct leak of {} byte(s) at {} allocated from:",
+    getContext()->logger.always("Direct leak of {} byte(s) at {} allocated from:",
                           AI->UserEnd - AI->UserBegin, (void *)AI->UserBegin);
     AI->AllocStack.print();
 }