add method to clean idle connection while oom

src/common/backend/utils/adt/misc.cpp

@@ -94,11 +94,11 @@ Datum current_query(PG_FUNCTION_ARGS)
 #define SIGNAL_BACKEND_SUCCESS 0
 #define SIGNAL_BACKEND_ERROR 1
 #define SIGNAL_BACKEND_NOPERMISSION 2
-static int pg_signal_backend(ThreadId pid, int sig)
+static int pg_signal_backend(ThreadId pid, int sig, bool checkPermission)
 {
     PGPROC* proc = NULL;
 
-    if (!superuser()) {
+    if (checkPermission && !superuser()) {
         /*
          * Since the user is not superuser, check for matching roles. Trust
          * that BackendPidGetProc will return NULL if the pid isn't valid,
@@ -202,7 +202,7 @@ Datum pg_cancel_backend(PG_FUNCTION_ARGS)
         ereport(ERROR,
             (errcode(ERRCODE_INSUFFICIENT_PRIVILEGE), (errmsg("kill backend is prohibited during online expansion."))));
 
-    r = pg_signal_backend(tid, SIGINT);
+    r = pg_signal_backend(tid, SIGINT, true);
 
     if (r == SIGNAL_BACKEND_NOPERMISSION)
         ereport(ERROR,
@@ -243,7 +243,7 @@ Datum pg_cancel_invalid_query(PG_FUNCTION_ARGS)
 #endif
 }
 
-static int kill_backend(ThreadId tid)
+int kill_backend(ThreadId tid, bool checkPermission)
 {
     /*
      * It is forbidden to kill backend in the online expansion to protect
@@ -253,17 +253,17 @@ static int kill_backend(ThreadId tid)
         ereport(ERROR,
             (errcode(ERRCODE_INSUFFICIENT_PRIVILEGE), (errmsg("kill backend is prohibited during online expansion."))));
 
-    int r = pg_signal_backend(tid, SIGTERM);
+    int r = pg_signal_backend(tid, SIGTERM, checkPermission);
     if (r == SIGNAL_BACKEND_NOPERMISSION)
         ereport(ERROR,
             (errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
                 (errmsg("must be system admin or have the same role to terminate other backend"))));
-    
+#ifdef ENABLE_MULTIPLE_NODES
     if (t_thrd.proc && t_thrd.proc->workingVersionNum >= 92060) {
         uint64 query_id = get_query_id_beentry(tid);
         (void)gs_close_all_stream_by_debug_id(query_id);
     }
-
+#endif
     return r;
 }
 
@@ -274,7 +274,7 @@ static int kill_backend(ThreadId tid)
 Datum pg_terminate_backend(PG_FUNCTION_ARGS)
 {
     ThreadId tid = PG_GETARG_INT64(0);
-    int r = kill_backend(tid);
+    int r = kill_backend(tid, true);
     PG_RETURN_BOOL(r == SIGNAL_BACKEND_SUCCESS);
 }
 
@@ -285,7 +285,7 @@ Datum pg_terminate_session(PG_FUNCTION_ARGS)
     int r = -1;
 
     if (tid == sid) {
-        r = kill_backend(tid);
+        r = kill_backend(tid, true);
     } else if (ENABLE_THREAD_POOL) {
         ThreadPoolSessControl *sess_ctrl = g_threadPoolControler->GetSessionCtrl();
         int ctrl_idx = sess_ctrl->FindCtrlIdxBySessId(sid);
@@ -952,7 +952,7 @@ void cancel_backend(ThreadId       pid)
 {
     int sig_return = 0;
 
-    sig_return = pg_signal_backend(pid, SIGINT);
+    sig_return = pg_signal_backend(pid, SIGINT, true);
 
     if (sig_return == SIGNAL_BACKEND_NOPERMISSION) {
         ereport(ERROR,
@@ -961,3 +961,15 @@ void cancel_backend(ThreadId       pid)
                 errhint("fail to cancel backend process for privilege")));
     }
 }
+
+Datum gs_clean_connection_by_memory(PG_FUNCTION_ARGS)
+{
+    int targetMemory = PG_GETARG_INT32(0);
+    if (targetMemory < 0) {
+        targetMemory = 0;
+    }
+    CleanConnectionByMemory(targetMemory);
+
+    PG_RETURN_VOID();
+}
+

src/common/backend/utils/misc/guc.cpp

@@ -4012,6 +4012,20 @@ static void InitConfigureNamesInt()
             NULL,
             NULL},
 
+        {{"min_dynamic_memory",
+            PGC_SIGHUP,
+            RESOURCES_MEM,
+            gettext_noop("Sets the min number of dynamic memory hold by the process."),
+            NULL,
+            GUC_UNIT_KB},
+            &g_instance.attr.attr_memory.min_dynamic_memory,
+            -1,
+            -1,
+            INT_MAX,
+            NULL,
+            NULL,
+            NULL},
+
         {{"local_syscache_threshold",
              PGC_POSTMASTER,
              RESOURCES_MEM,

src/common/backend/utils/mmgr/memprot.cpp

@@ -41,6 +41,9 @@
 #include "replication/walreceiver.h"
 #include "replication/walsender.h"
 
+#define OOM_TIME_THRESHOLD 50
+#define OOM_TIME_INTERVAL_MSEC 1000
+
 /* Track memory usage by all shared memory context */
 int32 shareTrackedMemChunks = 0;
 int64 shareTrackedBytes = 0;
@@ -599,6 +602,19 @@ static bool memTracker_ReserveMemChunks(int32 numChunksToReserve, bool needProte
     /* Query memory quota is exhausted. Reset the counter then return false. */
     if (MemoryIsNotEnough(total, *maxSize, needProtect)) {
         gs_atomic_add_32(currSize, -numChunksToReserve);
+        if (g_instance.attr.attr_memory.min_dynamic_memory >= 0) {
+            int current = pg_atomic_add_fetch_u32(&g_instance.exec_cxt.oomTimes, 1);
+            if (current == 1) {
+                g_instance.exec_cxt.firstTime = GetCurrentTimestamp();
+            }
+            if (current == OOM_TIME_THRESHOLD) {
+                if (!TimestampDifferenceExceeds(g_instance.exec_cxt.firstTime, GetCurrentTimestamp(),
+                    OOM_TIME_INTERVAL_MSEC)) {
+                    CleanConnectionByMemory(g_instance.attr.attr_memory.min_dynamic_memory);
+                }
+		g_instance.exec_cxt.oomTimes = 0;
+            }
+        }
         return false;
     }
 
@@ -1001,3 +1017,71 @@ void gs_memprot_process_gpu_memory(uint32 size)
         maxChunksPerProcess -= i;
     }
 }
+
+extern int kill_backend(ThreadId tid, bool checkPermission);
+
+static int memoryUsageCompare(const void* p1, const void* p2)
+{
+    SessMemoryUsage* m1 = (SessMemoryUsage*)p1;
+    SessMemoryUsage* m2 = (SessMemoryUsage*)p2;
+    /* release idle session first */
+    if (m1->state != m2->state) {
+        return m1->state < m2->state;
+    }
+    return m1->usedSize > m2->usedSize;
+}
+
+static void TerminateConnectionByMemory(int targetMemoryKBytes)
+{
+    int sessNum = 0;
+    SessMemoryUsage* sessMemory = NULL;
+    if (ENABLE_THREAD_POOL) {
+        sessMemory = g_threadPoolControler->GetSessionCtrl()->getSessionMemoryUsage(&sessNum);
+    } else {
+        sessMemory = getThreadMemoryUsage(&sessNum);
+    }
+
+    /* sort the session list by memory and state */
+    qsort(sessMemory, sessNum, sizeof(SessMemoryUsage), memoryUsageCompare);
+
+    uint64 currentMemory = (uint64)processMemInChunks << chunkSizeInBits;
+    const int BYTE_PER_KB = 1024;
+    uint64 targetMemory = (uint64)targetMemoryKBytes * BYTE_PER_KB;
+    int i = 0;
+    for (i = 0; i < sessNum; i++) {
+        if (unlikely(currentMemory < targetMemory)) {
+            break;
+        }
+
+        if (ENABLE_THREAD_POOL) {
+            ThreadPoolSessControl *sess_ctrl = g_threadPoolControler->GetSessionCtrl();
+            int ctrl_idx = sess_ctrl->FindCtrlIdxBySessId(sessMemory[i].sessid);
+            sess_ctrl->SendSignal((int)ctrl_idx, SIGTERM);
+        } else {
+            kill_backend(sessMemory[i].sessid, false);
+        }
+
+        currentMemory -= sessMemory[i].usedSize;
+    }
+
+    write_stderr("[OutOfMemoryFlee] Clean connection to release memory. current Memory is: %d MB, "
+                 " total connection: %d, clean connection: %d\n",
+        processMemInChunks >> (chunkSizeInBits - BITS_IN_MB), sessNum, i);
+
+    free(sessMemory);
+}
+
+void CleanConnectionByMemory(int targetMemory)
+{
+    PG_TRY();
+    {
+        g_instance.comm_cxt.rejectRequest = true;
+        TerminateConnectionByMemory(targetMemory);
+        g_instance.comm_cxt.rejectRequest = false;
+    }
+    PG_CATCH();
+    {
+        g_instance.comm_cxt.rejectRequest = false;
+    }
+    PG_END_TRY();
+}

src/gausskernel/process/postmaster/pgstat.cpp

@@ -8370,6 +8370,37 @@ void getThreadMemoryDetail(Tuplestorestate* tupStore, TupleDesc tupDesc, uint32*
     PG_END_TRY();
 }
 
+SessMemoryUsage* getThreadMemoryUsage(int* num)
+{
+    uint32 max_thread_count = g_instance.proc_base->allProcCount - g_instance.attr.attr_storage.max_prepared_xacts;
+    volatile PGPROC* proc = NULL;
+    uint32 idx = 0;
+
+    HOLD_INTERRUPTS();
+
+    SessMemoryUsage* result = (SessMemoryUsage*)malloc(max_thread_count * sizeof(SessMemoryUsage));
+    int index = 0;
+
+    for (idx = 0; idx < max_thread_count; idx++) {
+        proc = g_instance.proc_base_all_procs[idx];
+
+        /* lock this proc's delete MemoryContext action */
+        (void)syscalllockAcquire(&((PGPROC*)proc)->deleMemContextMutex);
+        if (proc->usedMemory != NULL) {
+            result[index].sessid = proc->pid;
+            result[index].usedSize = *proc->usedMemory;
+            result[index].state = (int)t_thrd.shemem_ptr_cxt.BackendStatusArray[t_thrd.proc_cxt.MyBackendId - 1].st_state;
+            index++;
+        }
+        (void)syscalllockRelease(&((PGPROC*)proc)->deleMemContextMutex);
+    }
+
+    RESUME_INTERRUPTS();
+
+    *num = index;
+    return result;
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 #ifdef MEMORY_CONTEXT_CHECKING
 

src/gausskernel/process/postmaster/postmaster.cpp

@@ -3706,6 +3706,11 @@ CAC_state canAcceptConnections(bool isSession)
             return CAC_RECOVERY; /* else must be crash recovery */
     }
 
+    if (g_instance.comm_cxt.rejectRequest) {
+        result = CAC_OOM;
+    }
+
+
     if (isSession)
         return result;
 
@@ -6499,11 +6504,15 @@ static int BackendStartup(Port* port, bool isConnectHaPort)
     port->canAcceptConnections = canAcceptConnections(false);
     if (port->canAcceptConnections != CAC_OK && port->canAcceptConnections != CAC_WAITBACKUP) {
         (void)ReleasePostmasterChildSlot(childSlot);
-        if (port->canAcceptConnections == CAC_TOOMANY)
+        if (port->canAcceptConnections == CAC_TOOMANY) {
             ereport(WARNING, (errmsg("could not fork new process for connection due to too many connections")));
-        else
+        } else if (port->canAcceptConnections == CAC_OOM) {
+            ereport(
+                WARNING, (errmsg("The server is cleaning connection to reduce memory usage, please retry later")));
+        } else {
             ereport(
                 WARNING, (errmsg("could not fork new process for connection due to PMstate %s", GetPMState(pmState))));
+        }
 
         return STATUS_ERROR;
     }

src/gausskernel/process/threadpool/knl_instance.cpp

@@ -282,6 +282,7 @@ static void knl_g_comm_init(knl_g_comm_context* comm_cxt)
     comm_cxt->cal_all_space_info_in_progress = false;
     comm_cxt->current_gsrewind_count = 0;
     comm_cxt->usedDnSpace = NULL;
+    comm_cxt->rejectRequest = false;
 
     knl_g_quota_init(&g_instance.comm_cxt.quota_cxt);
     knl_g_localinfo_init(&g_instance.comm_cxt.localinfo_cxt);

src/gausskernel/process/threadpool/threadpool_sessctl.cpp

@@ -550,6 +550,39 @@ void ThreadPoolSessControl::getSessionMemoryDetail(Tuplestorestate* tupStore,
     PG_END_TRY();
 }
 
+SessMemoryUsage* ThreadPoolSessControl::getSessionMemoryUsage(int* num)
+{
+    AutoMutexLock alock(&m_sessCtrlock);
+    knl_sess_control* ctrl = NULL;
+    Dlelem* elem = NULL;
+
+    HOLD_INTERRUPTS();
+    alock.lock();
+
+    SessMemoryUsage* result = (SessMemoryUsage*)malloc(m_activeSessionCount * sizeof(SessMemoryUsage));
+    int index = 0;
+
+    elem = DLGetHead(&m_activelist);
+    while (elem != NULL) {
+        ctrl = (knl_sess_control*)DLE_VAL(elem);
+        knl_session_context* sess = ctrl->sess;
+        if (sess) {
+            result[index].sessid = sess->session_id;
+            result[index].usedSize = sess->stat_cxt.trackedBytes;
+            result[index].state = (int)t_thrd.shemem_ptr_cxt.BackendStatusArray[sess->session_ctr_index].st_state;
+            index++;
+        }
+        elem = DLGetSucc(elem);
+    }
+
+    alock.unLock();
+    RESUME_INTERRUPTS();
+
+    *num = index;
+    return result;
+}
+
+
 knl_session_context* ThreadPoolSessControl::GetSessionByIdx(int idx)
 {
     if (IsValidCtrlIndex(idx)) {

src/gausskernel/storage/lmgr/proc.cpp

@@ -591,6 +591,9 @@ void InitProcess(void)
         ereport(ERROR,
                 (errcode(ERRCODE_LOCK_NOT_AVAILABLE), errmsg("failed to acquire mutex lock for deleMemContextMutex.")));
     t_thrd.proc->topmcxt = t_thrd.top_mem_cxt;
+    if (t_thrd.role == WORKER || t_thrd.role == THREADPOOL_WORKER) {
+        t_thrd.proc->usedMemory = &t_thrd.utils_cxt.trackedBytes;
+    }
     if (syscalllockRelease(&t_thrd.proc->deleMemContextMutex) != 0)
         ereport(ERROR,
                 (errcode(ERRCODE_LOCK_NOT_AVAILABLE), errmsg("failed to release mutex lock for deleMemContextMutex.")));
@@ -1134,6 +1137,7 @@ static void ProcKill(int code, Datum arg)
         ereport(ERROR,
                 (errcode(ERRCODE_LOCK_NOT_AVAILABLE), errmsg("failed to acquire mutex lock for deleMemContextMutex.")));
     t_thrd.proc->topmcxt = NULL;
+    t_thrd.proc->usedMemory = NULL;
     if (syscalllockRelease(&t_thrd.proc->deleMemContextMutex) != 0)
         ereport(ERROR,
                 (errcode(ERRCODE_LOCK_NOT_AVAILABLE), errmsg("failed to release mutex lock for deleMemContextMutex.")));

src/include/knl/knl_guc/knl_instance_attr_memory.h

@@ -45,6 +45,7 @@ typedef struct knl_instance_attr_memory {
     bool enable_memory_limit;
     int memorypool_size;
     int max_process_memory;
+    int min_dynamic_memory;
     int local_syscache_threshold;
 } knl_instance_attr_memory;
 

src/include/knl/knl_instance.h

@@ -590,6 +590,8 @@ typedef struct knl_g_comm_context {
 
     HTAB* usedDnSpace;
     uint32 current_gsrewind_count;
+
+    bool rejectRequest;
 } knl_g_comm_context;
 
 typedef struct knl_g_libpq_context {
@@ -626,6 +628,8 @@ typedef struct knl_g_executor_context {
 #ifndef ENABLE_MULTIPLE_NODES
     char* nodeName;
 #endif
+    volatile uint32 oomTimes;
+    TimestampTz firstTime;
 } knl_g_executor_context;
 
 typedef struct knl_g_xlog_context {

src/include/libpq/libpq-be.h

@@ -70,7 +70,8 @@ typedef struct {
 #include "libpq/sha2.h"
 #include "libcomm/libcomm.h"
 
-typedef enum CAC_state { CAC_OK, CAC_STARTUP, CAC_SHUTDOWN, CAC_RECOVERY, CAC_TOOMANY, CAC_WAITBACKUP } CAC_state;
+typedef enum CAC_state {
+    CAC_OK, CAC_STARTUP, CAC_SHUTDOWN, CAC_RECOVERY, CAC_TOOMANY, CAC_WAITBACKUP, CAC_OOM } CAC_state;
 
 /*
  * GSSAPI specific state information

src/include/nodes/memnodes.h

@@ -264,5 +264,12 @@ extern MemoryProtectFuncDef SharedFunctions;
             IsA((context), MemalignSharedAllocSetContext)))
 
 #define AllocSetContextUsedSpace(aset) ((aset)->totalSpace - (aset)->freeSpace)
+
+typedef struct SessMemoryUsage {
+    uint64 sessid;
+    int64 usedSize;
+    int state;
+} SessMemoryUsage;
+
 #endif /* MEMNODES_H */
 

src/include/pgstat.h

@@ -2009,6 +2009,7 @@ extern void DumpMemoryContext(DUMP_TYPE type);
 
 extern void getThreadMemoryDetail(Tuplestorestate* tupStore, TupleDesc tupDesc, uint32* procIdx);
 extern void getSharedMemoryDetail(Tuplestorestate* tupStore, TupleDesc tupDesc);
+extern SessMemoryUsage* getThreadMemoryUsage(int* num);
 
 typedef enum TimeInfoType {
     DB_TIME = 0, /*total elapsed time while dealing user command.*/

src/include/storage/proc.h

@@ -182,6 +182,7 @@ struct PGPROC {
     char myProgName[64];
     pg_time_t myStartTime;
     syscalllock deleMemContextMutex;
+    int64* usedMemory;
 
     /*
      * All PROCLOCK objects for locks held or awaited by this backend are

src/include/threadpool/threadpool_sessctl.h

@@ -62,6 +62,7 @@ public:
     void CheckSessionTimeout();
     void CheckPermissionForSendSignal(knl_session_context* sess, sig_atomic_t* lock);
     void getSessionMemoryDetail(Tuplestorestate* tupStore, TupleDesc tupDesc, knl_sess_control** sess);
+    SessMemoryUsage* getSessionMemoryUsage(int* num);
     knl_session_context* GetSessionByIdx(int idx);
     int FindCtrlIdxBySessId(uint64 id);
     TransactionId ListAllSessionGttFrozenxids(int maxSize, ThreadId *pids, TransactionId *xids, int *n);

src/include/utils/memprot.h

@@ -103,4 +103,6 @@ extern void gs_memprot_reset_beyondchunk(void);
 
 extern int getSessionMemoryUsageMB();
 
+extern void CleanConnectionByMemory(int targetMemory);
+
 #endif