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_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.