Merge pull request #108 from zmstone/251122-ensure-cleanup-after-shutdown

fix: ensure cleanup after shutdown

Author: zmstone

changelog.md

@@ -1,3 +1,7 @@
+* 4.1.3
+  - Ensure `wolff_client_sup:ensure_absence` and `wolff_producers_sup:ensure_absence` will perform shutdown and cleanup atomically.
+    Previously, if the caller process is killed while waiting for shutdown, a terminated child may leak under the supervisor.
+
 * 4.1.2
   - Made sure `wolff_client:check_topic_exists_with_client_pid/2` triggers topic creation when `allow_auto_topic_creation` is set to `true`.
 

src/wolff_client.erl

@@ -57,7 +57,8 @@
         %% but we need to find connection by {topic(), partition()}
         leaders => #{{topic(), partition()} => connection()},
         %% Reference counting so we may drop connection metadata when no longer required.
-        known_topics := #{topic() => #{producer_group() => true}}
+        known_topics := #{topic() => #{producer_group() => true}},
+        owner := pid()
        }.
 
 -define(DEFAULT_METADATA_TIMEOUT, 10000).
@@ -78,6 +79,7 @@
 start_link(ClientId, Hosts, Config) ->
   {ConnCfg0, MyCfg} = split_config(Config),
   ConnCfg = ConnCfg0#{client_id => ClientId},
+  Owner = self(),
   St = #{client_id => ClientId,
          seed_hosts => Hosts,
          config => MyCfg,
@@ -86,7 +88,8 @@ start_link(ClientId, Hosts, Config) ->
          metadata_conn => not_initialized,
          metadata_ts => #{},
          leaders => #{},
-         known_topics => #{}
+         known_topics => #{},
+         owner => Owner
         },
   case maps:get(reg_name, Config, false) of
     false -> gen_server:start_link(?MODULE, St, []);
@@ -199,7 +202,13 @@ handle_call(Call, _From, St) ->
   {reply, {error, {unknown_call, Call}}, St}.
 
 handle_info({'EXIT', Pid, Reason}, St) ->
-  {noreply, flush_exit_signals(St, Pid, Reason)};
+  %% Check if this is a shutdown signal from the supervisor
+  case erlang:whereis(wolff_client_sup) of
+    Pid when Reason =:= shutdown ->
+      {stop, shutdown, St};
+    _ ->
+      {noreply, flush_exit_signals(St, Pid, Reason)}
+  end;
 handle_info(_Info, St) ->
   {noreply, upgrade(St)}.
 
@@ -243,7 +252,7 @@ handle_cast(_Cast, St) ->
 code_change(_OldVsn, St, _Extra) ->
   {ok, St}.
 
-terminate(_, #{client_id := ClientID, conns := Conns} = St) ->
+terminate(_Reason, #{client_id := ClientID, conns := Conns} = St) ->
   ok = wolff_client_sup:deregister_client(ClientID),
   MetadataConn = maps:get(metadata_conn, St, none),
   ok = close_connections(Conns),
@@ -307,9 +316,9 @@ close_connections(Conns, Topic) ->
          end,
   do_close_connections(maps:to_list(Conns), Pred, #{}).
 
-flush_exit_signals(St0) ->
+flush_exit_signals(#{owner := Owner} = St0) ->
   receive
-    {'EXIT', Pid, Reason} ->
+    {'EXIT', Pid, Reason} when Pid =/= Owner ->
       flush_exit_signals(St0, Pid, Reason)
   after
     0 ->

src/wolff_client_sup.erl

@@ -50,6 +50,16 @@ ensure_present(ClientId, Hosts, Config) ->
 %% @doc Ensure client stopped and deleted under supervisor.
 -spec ensure_absence(wolff:client_id()) -> ok.
 ensure_absence(ClientId) ->
+  %% This receive may take 5s because the child spec's shutdown policy is to kill after 5s
+  %% To ensure the cleanup is complete, spawn a process to do it, in case the caller
+  %% gets killed while waiting for shutdown.
+  {Pid, Mref} = erlang:spawn_monitor(fun() -> do_ensure_absence(ClientId) end),
+  receive
+      {'DOWN', Mref, process, Pid, _} ->
+          ok
+  end.
+
+do_ensure_absence(ClientId) ->
   case supervisor:terminate_child(?SUPERVISOR, ClientId) of
     ok -> ok = supervisor:delete_child(?SUPERVISOR, ClientId);
     {error, not_found} -> ok
@@ -79,7 +89,8 @@ child_spec(ClientId, Hosts, Config) ->
     start => {wolff_client, start_link, [ClientId, Hosts, Config]},
     restart => transient,
     type => worker,
-    modules => [wolff_client]
+    modules => [wolff_client],
+    shutdown => 5000
    }.
 
 %% @doc Create a ets table which is used for client registration.

src/wolff_producers_sup.erl

@@ -48,6 +48,16 @@ ensure_present(ClientId, ProducerId, Config) ->
 %% ensure client stopped and deleted under supervisor
 -spec ensure_absence(wolff_producers:id()) -> ok.
 ensure_absence(ProducerId) ->
+  %% This receive may take 5s because the child spec's shutdown policy is to kill after 5s
+  %% To ensure the cleanup is complete, spawn a process to do it, in case the caller
+  %% gets killed while waiting for shutdown.
+  {Pid, Mref} = erlang:spawn_monitor(fun() -> do_ensure_absence(ProducerId) end),
+  receive
+      {'DOWN', Mref, process, Pid, _} ->
+          ok
+  end.
+
+do_ensure_absence(ProducerId) ->
   case supervisor:terminate_child(?SUPERVISOR, ProducerId) of
     ok ->
       ok = wolff_producers:cleanup_workers_table(ProducerId),
@@ -60,7 +70,9 @@ child_spec(ClientId, ProducerId, Config) ->
   #{id => ProducerId,
     start => {wolff_producers, start_link, [ClientId, ProducerId, Config]},
     restart => transient,
-    type => worker
+    type => worker,
+    modules => [wolff_producers],
+    shutdown => 5000
    }.
 
 %% Find the running process (gen_server of wolff_producers) from thie child ID.