From e66666d8ebab0f7ea92c1c060270c02bff7d324f Mon Sep 17 00:00:00 2001
From: Zou Xinyi <zouxinyi@selectdb.com>
Date: Tue, 7 Jan 2025 20:25:36 +0800
Subject: [PATCH] 1

---
 be/src/runtime/thread_context.h               | 38 +++++++++++--------
 .../runtime/shared_hash_table_controller.h    |  9 +++++
 2 files changed, 32 insertions(+), 15 deletions(-)

diff --git a/be/src/runtime/thread_context.h b/be/src/runtime/thread_context.h
index 9ba7949ec5afad2..a5d7ee917f46060 100644
--- a/be/src/runtime/thread_context.h
+++ b/be/src/runtime/thread_context.h
@@ -575,21 +575,29 @@ class ScopeSkipMemoryCheck {
 
 // if use mem hook, avoid repeated consume.
 // must call create_thread_local_if_not_exits() before use thread_context().
-#define CONSUME_THREAD_MEM_TRACKER(size)                                                       \
-    do {                                                                                       \
-        if (size == 0 || doris::use_mem_hook) {                                                \
-            break;                                                                             \
-        }                                                                                      \
-        if (doris::pthread_context_ptr_init) {                                                 \
-            DCHECK(bthread_self() == 0);                                                       \
-            doris::thread_context_ptr->consume_memory(size);                                   \
-        } else if (bthread_self() != 0) {                                                      \
-            static_cast<doris::ThreadContext*>(bthread_getspecific(doris::btls_key))           \
-                    ->consume_memory(size);                                                    \
-        } else if (doris::ExecEnv::ready()) {                                                  \
-            MEMORY_ORPHAN_CHECK();                                                             \
-            doris::ExecEnv::GetInstance()->orphan_mem_tracker()->consume_no_update_peak(size); \
-        }                                                                                      \
+#define CONSUME_THREAD_MEM_TRACKER(size)                                                           \
+    do {                                                                                           \
+        if (size == 0 || doris::use_mem_hook) {                                                    \
+            break;                                                                                 \
+        }                                                                                          \
+        if (doris::pthread_context_ptr_init) {                                                     \
+            DCHECK(bthread_self() == 0);                                                           \
+            doris::thread_context_ptr->consume_memory(size);                                       \
+        } else if (bthread_self() != 0) {                                                          \
+            auto* bthread_context =                                                                \
+                    static_cast<doris::ThreadContext*>(bthread_getspecific(doris::btls_key));      \
+            DCHECK(bthread_context != nullptr);                                                    \
+            if (bthread_context != nullptr) {                                                      \
+                bthread_context->consume_memory(size);                                             \
+            } else {                                                                               \
+                doris::ExecEnv::GetInstance()->orphan_mem_tracker()->consume_no_update_peak(size); \
+            }                                                                                      \
+            static_cast<doris::ThreadContext*>(bthread_getspecific(doris::btls_key))               \
+                    ->consume_memory(size);                                                        \
+        } else if (doris::ExecEnv::ready()) {                                                      \
+            MEMORY_ORPHAN_CHECK();                                                                 \
+            doris::ExecEnv::GetInstance()->orphan_mem_tracker()->consume_no_update_peak(size);     \
+        }                                                                                          \
     } while (0)
 #define RELEASE_THREAD_MEM_TRACKER(size) CONSUME_THREAD_MEM_TRACKER(-size)
 
diff --git a/be/src/vec/runtime/shared_hash_table_controller.h b/be/src/vec/runtime/shared_hash_table_controller.h
index ff9ad4d0ef48c2e..0926cf78af2cd9f 100644
--- a/be/src/vec/runtime/shared_hash_table_controller.h
+++ b/be/src/vec/runtime/shared_hash_table_controller.h
@@ -26,6 +26,7 @@
 #include <vector>
 
 #include "common/status.h"
+#include "runtime/thread_context.h"
 #include "vec/core/block.h"
 
 namespace doris {
@@ -49,6 +50,14 @@ struct RuntimeFilterContext {
     bool ignored = false;
     bool disabled = false;
     std::string err_msg;
+    QueryThreadContext query_thread_context;
+
+    RuntimeFilterContext() { query_thread_context.init_unlocked(); }
+
+    ~RuntimeFilterContext() {
+        SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(query_thread_context.query_mem_tracker);
+        hybrid_set.reset();
+    }
 };
 
 using RuntimeFilterContextSPtr = std::shared_ptr<RuntimeFilterContext>;