From c0ca7f8beab997d89fb4708c8ffc2b55370540ca Mon Sep 17 00:00:00 2001
From: Turtle <deng@emqx.io>
Date: Sun, 26 Sep 2021 18:21:34 +0800
Subject: [PATCH] fix(force_shutdown): cannot suicide if the process hangs up

---
 src/emqx.appup.src         | 48 ++++++++++++++++++++++++++++----------
 src/emqx_alarm_handler.erl | 12 +++++-----
 src/emqx_misc.erl          | 22 +++++++++++++++--
 3 files changed, 62 insertions(+), 20 deletions(-)

diff --git a/src/emqx.appup.src b/src/emqx.appup.src
index 442ffe200..0c3ab3da8 100644
--- a/src/emqx.appup.src
+++ b/src/emqx.appup.src
@@ -21,7 +21,9 @@
      {load_module, emqx_trie, soft_purge, soft_purge, []},
      {load_module, emqx_router, soft_purge, soft_purge, [emqx_trie]},
      {load_module,emqx_pqueue,brutal_purge,soft_purge,[]},
-     {load_module,emqx_mqueue,brutal_purge,soft_purge,[]}
+     {load_module,emqx_mqueue,brutal_purge,soft_purge,[]},
+     {load_module,emqx_alarm_handler,brutal_purge,soft_purge,[]},
+     {load_module,emqx_misc,brutal_purge,soft_purge,[]}
    ]},
    {"4.2.1", [
      {add_module, emqx_congestion},
@@ -41,7 +43,9 @@
      {load_module, emqx_trie, soft_purge, soft_purge, []},
      {load_module, emqx_router, soft_purge, soft_purge, [emqx_trie]},
      {load_module,emqx_pqueue,brutal_purge,soft_purge,[]},
-     {load_module,emqx_mqueue,brutal_purge,soft_purge,[]}
+     {load_module,emqx_mqueue,brutal_purge,soft_purge,[]},
+     {load_module,emqx_alarm_handler,brutal_purge,soft_purge,[]},
+     {load_module,emqx_misc,brutal_purge,soft_purge,[]}
    ]},
    {<<"4.2.[23]">>, [
      {add_module, emqx_congestion},
@@ -58,7 +62,9 @@
      {load_module, emqx_trie, soft_purge, soft_purge, []},
      {load_module, emqx_router, soft_purge, soft_purge, [emqx_trie]},
      {load_module,emqx_pqueue,brutal_purge,soft_purge,[]},
-     {load_module,emqx_mqueue,brutal_purge,soft_purge,[]}
+     {load_module,emqx_mqueue,brutal_purge,soft_purge,[]},
+     {load_module,emqx_alarm_handler,brutal_purge,soft_purge,[]},
+     {load_module,emqx_misc,brutal_purge,soft_purge,[]}
    ]},
    {<<"4.2.4">>, [
      {load_module, emqx_frame, brutal_purge, soft_purge, []},
@@ -73,7 +79,9 @@
      {load_module, emqx_trie, soft_purge, soft_purge, []},
      {load_module, emqx_router, soft_purge, soft_purge, [emqx_trie]},
      {load_module,emqx_pqueue,brutal_purge,soft_purge,[]},
-     {load_module,emqx_mqueue,brutal_purge,soft_purge,[]}
+     {load_module,emqx_mqueue,brutal_purge,soft_purge,[]},
+     {load_module,emqx_alarm_handler,brutal_purge,soft_purge,[]},
+     {load_module,emqx_misc,brutal_purge,soft_purge,[]}
    ]},
    {<<"4.2.5">>, [
      {load_module, emqx_channel, brutal_purge, soft_purge, []},
@@ -87,11 +95,15 @@
      {load_module, emqx_trie, soft_purge, soft_purge, []},
      {load_module, emqx_router, soft_purge, soft_purge, [emqx_trie]},
      {load_module,emqx_pqueue,brutal_purge,soft_purge,[]},
-     {load_module,emqx_mqueue,brutal_purge,soft_purge,[]}
+     {load_module,emqx_mqueue,brutal_purge,soft_purge,[]},
+     {load_module,emqx_alarm_handler,brutal_purge,soft_purge,[]},
+     {load_module,emqx_misc,brutal_purge,soft_purge,[]}
    ]},
    {<<"4.2.[6-7]">>, [
      {load_module,emqx_pqueue,brutal_purge,soft_purge,[]},
-     {load_module,emqx_mqueue,brutal_purge,soft_purge,[]}
+     {load_module,emqx_mqueue,brutal_purge,soft_purge,[]},
+     {load_module,emqx_alarm_handler,brutal_purge,soft_purge,[]},
+     {load_module,emqx_misc,brutal_purge,soft_purge,[]}
    ]},
    {<<".*">>, []}
  ],
@@ -115,7 +127,9 @@
      {load_module, emqx_trie, soft_purge, soft_purge, [emqx_router]},
      {load_module, emqx_router, soft_purge, soft_purge, []},
      {load_module,emqx_pqueue,brutal_purge,soft_purge,[]},
-     {load_module,emqx_mqueue,brutal_purge,soft_purge,[]}
+     {load_module,emqx_mqueue,brutal_purge,soft_purge,[]},
+     {load_module,emqx_alarm_handler,brutal_purge,soft_purge,[]},
+     {load_module,emqx_misc,brutal_purge,soft_purge,[]}
    ]},
    {"4.2.1", [
      {load_module, emqx_shared_sub, brutal_purge, soft_purge, []},
@@ -135,7 +149,9 @@
      {load_module, emqx_trie, soft_purge, soft_purge, [emqx_router]},
      {load_module, emqx_router, soft_purge, soft_purge, []},
      {load_module,emqx_pqueue,brutal_purge,soft_purge,[]},
-     {load_module,emqx_mqueue,brutal_purge,soft_purge,[]}
+     {load_module,emqx_mqueue,brutal_purge,soft_purge,[]},
+     {load_module,emqx_alarm_handler,brutal_purge,soft_purge,[]},
+     {load_module,emqx_misc,brutal_purge,soft_purge,[]}
    ]},
    {<<"4.2.[23]">>, [
      {load_module, emqx_shared_sub, brutal_purge, soft_purge, []},
@@ -152,7 +168,9 @@
      {load_module, emqx_trie, soft_purge, soft_purge, [emqx_router]},
      {load_module, emqx_router, soft_purge, soft_purge, []},
      {load_module,emqx_pqueue,brutal_purge,soft_purge,[]},
-     {load_module,emqx_mqueue,brutal_purge,soft_purge,[]}
+     {load_module,emqx_mqueue,brutal_purge,soft_purge,[]},
+     {load_module,emqx_alarm_handler,brutal_purge,soft_purge,[]},
+     {load_module,emqx_misc,brutal_purge,soft_purge,[]}
    ]},
    {<<"4.2.4">>, [
      {load_module, emqx_frame, brutal_purge, soft_purge, []},
@@ -167,7 +185,9 @@
      {load_module, emqx_trie, soft_purge, soft_purge, [emqx_router]},
      {load_module, emqx_router, soft_purge, soft_purge, []},
      {load_module,emqx_pqueue,brutal_purge,soft_purge,[]},
-     {load_module,emqx_mqueue,brutal_purge,soft_purge,[]}
+     {load_module,emqx_mqueue,brutal_purge,soft_purge,[]},
+     {load_module,emqx_alarm_handler,brutal_purge,soft_purge,[]},
+     {load_module,emqx_misc,brutal_purge,soft_purge,[]}
    ]},
    {<<"4.2.5">>, [
      {load_module, emqx_channel, brutal_purge, soft_purge, []},
@@ -181,11 +201,15 @@
      {load_module, emqx_trie, soft_purge, soft_purge, [emqx_router]},
      {load_module, emqx_router, soft_purge, soft_purge, []},
      {load_module,emqx_pqueue,brutal_purge,soft_purge,[]},
-     {load_module,emqx_mqueue,brutal_purge,soft_purge,[]}
+     {load_module,emqx_mqueue,brutal_purge,soft_purge,[]},
+     {load_module,emqx_alarm_handler,brutal_purge,soft_purge,[]},
+     {load_module,emqx_misc,brutal_purge,soft_purge,[]}
    ]},
    {<<"4.2.[6-7]">>, [
      {load_module,emqx_pqueue,brutal_purge,soft_purge,[]},
-     {load_module,emqx_mqueue,brutal_purge,soft_purge,[]}
+     {load_module,emqx_mqueue,brutal_purge,soft_purge,[]},
+     {load_module,emqx_alarm_handler,brutal_purge,soft_purge,[]},
+     {load_module,emqx_misc,brutal_purge,soft_purge,[]}
    ]},
    {<<".*">>, []}
  ]
diff --git a/src/emqx_alarm_handler.erl b/src/emqx_alarm_handler.erl
index 7f66f6eb1..cd52e586f 100644
--- a/src/emqx_alarm_handler.erl
+++ b/src/emqx_alarm_handler.erl
@@ -56,20 +56,20 @@ init({_Args, {alarm_handler, _ExistingAlarms}}) ->
 init(_) ->
     {ok, []}.
 
-handle_event({set_alarm, {system_memory_high_watermark, []}}, State) -> 
+handle_event({set_alarm, {system_memory_high_watermark, []}}, State) ->
     emqx_alarm:activate(high_system_memory_usage, #{high_watermark => emqx_os_mon:get_sysmem_high_watermark()}),
     {ok, State};
 
-handle_event({set_alarm, {process_memory_high_watermark, Pid}}, State) -> 
-    emqx_alarm:activate(high_process_memory_usage, #{pid => Pid,
+handle_event({set_alarm, {process_memory_high_watermark, Pid}}, State) ->
+    emqx_alarm:activate(high_process_memory_usage, #{pid => list_to_binary(pid_to_list(Pid)),
                                                      high_watermark => emqx_os_mon:get_procmem_high_watermark()}),
     {ok, State};
 
-handle_event({clear_alarm, system_memory_high_watermark}, State) -> 
+handle_event({clear_alarm, system_memory_high_watermark}, State) ->
     emqx_alarm:deactivate(high_system_memory_usage),
     {ok, State};
 
-handle_event({clear_alarm, process_memory_high_watermark}, State) -> 
+handle_event({clear_alarm, process_memory_high_watermark}, State) ->
     emqx_alarm:deactivate(high_process_memory_usage),
     {ok, State};
 
@@ -85,4 +85,4 @@ handle_call(_Query, State) ->
 terminate(swap, _State) ->
     {emqx_alarm_handler, []};
 terminate(_, _) ->
-    ok.
\ No newline at end of file
+    ok.
diff --git a/src/emqx_misc.erl b/src/emqx_misc.erl
index 88e3c91a6..55e8b28f9 100644
--- a/src/emqx_misc.erl
+++ b/src/emqx_misc.erl
@@ -45,6 +45,8 @@
         , index_of/2
         ]).
 
+-define(OOM_FACTOR, 1.25).
+
 %% @doc Merge options
 -spec(merge_opts(Opts, Opts) -> Opts when Opts :: proplists:proplist()).
 merge_opts(Defaults, Options) ->
@@ -185,8 +187,8 @@ do_check_oom([{Val, Max, Reason}|Rest]) ->
 
 tune_heap_size(#{max_heap_size := MaxHeapSize}) ->
     %% If set to zero, the limit is disabled.
-    erlang:process_flag(max_heap_size, #{size => MaxHeapSize,
-                                         kill => false,
+    erlang:process_flag(max_heap_size, #{size => must_kill_heap_size(MaxHeapSize),
+                                         kill => true,
                                          error_logger => true
                                         });
 tune_heap_size(undefined) -> ok.
@@ -233,3 +235,19 @@ index_of(E, I, [E|_]) ->
 index_of(E, I, [_|L]) ->
     index_of(E, I+1, L).
 
+must_kill_heap_size(Size) ->
+    %% We set the max allowed heap size by `erlang:process_flag(max_heap_size, #{size => Size})`,
+    %% where the `Size` cannot be set to an integer lager than `(1 bsl 59) - 1` on a 64-bit system,
+    %% or `(1 bsl 27) - 1` on a 32-bit system.
+    MaxAllowedSize = case erlang:system_info(wordsize) of
+        8 -> % arch_64
+            (1 bsl 59) - 1;
+        4 -> % arch_32
+            (1 bsl 27) - 1
+    end,
+    %% We multiply the size with factor ?OOM_FACTOR, to give the
+    %% process a chance to suicide by `check_oom/1`
+    case ceil(Size * ?OOM_FACTOR) of
+        Size0 when Size0 >= MaxAllowedSize -> MaxAllowedSize;
+        Size0 -> Size0
+    end.