Repository: general-CbIC/poolex (elixir)

Issue:
Poolex doesn't verify if there are callers after a worker dies
I didn't create a test, but I think when a worker dies Poolex doesn't verify if there are new callers waiting for a worker:
https://github.com/general-CbIC/poolex/blob/main/lib/poolex.ex#L449-L482

We do here for example:
https://github.com/general-CbIC/poolex/blob/main/lib/poolex.ex#L401

I also think that Poolex wouldn't need to make a difference between a "temporary worker" and a "regular worker", and it could make the code easier (less complex). We only need to verify the number the workers and the configuration params:

If a worker dies, we remove it from the list of workers, and then:
  * We check if there are pending callers, we create a new worker
  * If not we check if the number of current workers is less than minimum, then we create a new worker.
  * Else, we remove don't create a new one.

Similarly, when a worker is released:
  * If we have pending callers, we send the worker to the caller
  * If the number of the current workers is in overflow, we kill it.
  * If not, we put it on the idle list.

I hope it helps

Pull request:
Check pending callers after worker restart
Fixed a bug where a restarted worker was not automatically dispatched to pending callers.

Additional improvements:
- Simplified code by removing the `temporary_worker` conception.
- Prevented possible memory leaks by changing workers `restart` strategy to `temporary`.

Closes #53 

Patch:
diff --git a/lib/poolex.ex b/lib/poolex.ex
index 536c244..171893f 100644
--- a/lib/poolex.ex
+++ b/lib/poolex.ex
@@ -321,7 +321,8 @@ defmodule Poolex do
   defp start_worker(%State{} = state) do
     DynamicSupervisor.start_child(state.supervisor, %{
       id: make_ref(),
-      start: {state.worker_module, state.worker_start_fun, state.worker_args}
+      start: {state.worker_module, state.worker_start_fun, state.worker_args},
+      restart: :temporary
     })
   end
 
@@ -336,16 +337,11 @@ defmodule Poolex do
       if state.overflow < state.max_overflow do
         {:ok, new_worker} = start_worker(state)
 
-        Monitoring.add(state.monitor_id, new_worker, :temporary_worker)
+        Monitoring.add(state.monitor_id, new_worker, :worker)
 
-        new_state = %State{
-          state
-          | busy_workers_state:
-              BusyWorkers.add(state.busy_workers_impl, state.busy_workers_state, new_worker),
-            overflow: state.overflow + 1
-        }
+        state = add_worker_to_busy_workers(state, new_worker)
 
-        {:reply, {:ok, new_worker}, new_state}
+        {:reply, {:ok, new_worker}, %State{state | overflow: state.overflow + 1}}
       else
         Monitoring.add(state.monitor_id, from_pid, :caller)
 
@@ -358,14 +354,9 @@ defmodule Poolex do
       {idle_worker_pid, new_idle_workers_state} =
         IdleWorkers.pop(state.idle_workers_impl, state.idle_workers_state)
 
-      new_busy_workers_state =
-        BusyWorkers.add(state.busy_workers_impl, state.busy_workers_state, idle_worker_pid)
+      state = add_worker_to_busy_workers(state, idle_worker_pid)
 
-      new_state = %State{
-        state
-        | idle_workers_state: new_idle_workers_state,
-          busy_workers_state: new_busy_workers_state
-      }
+      new_state = %State{state | idle_workers_state: new_idle_workers_state}
 
       {:reply, {:ok, idle_worker_pid}, new_state}
     end
@@ -407,23 +398,31 @@ defmodule Poolex do
     end
   end
 
+  @impl GenServer
+  def handle_info(
+        {:DOWN, monitoring_reference, _process, dead_process_pid, _reason},
+        %State{} = state
+      ) do
+    case Monitoring.remove(state.monitor_id, monitoring_reference) do
+      :worker ->
+        {:noreply, handle_down_worker(state, dead_process_pid)}
+
+      :caller ->
+        {:noreply, handle_down_caller(state, dead_process_pid)}
+    end
+  end
+
   @spec release_busy_worker(State.t(), worker()) :: State.t()
   defp release_busy_worker(%State{} = state, worker) do
     if BusyWorkers.member?(state.busy_workers_impl, state.busy_workers_state, worker) do
-      busy_workers_state =
-        BusyWorkers.remove(state.busy_workers_impl, state.busy_workers_state, worker)
+      state = remove_worker_from_busy_workers(state, worker)
 
       if state.overflow > 0 do
         stop_worker(state.supervisor, worker)
 
-        %State{state | busy_workers_state: busy_workers_state}
+        state
       else
-        %State{
-          state
-          | busy_workers_state: busy_workers_state,
-            idle_workers_state:
-              IdleWorkers.add(state.idle_workers_impl, state.idle_workers_state, worker)
-        }
+        add_worker_to_idle_workers(state, worker)
       end
     else
       state
@@ -440,61 +439,97 @@ defmodule Poolex do
     %{state | waiting_callers_state: new_waiting_callers_state}
   end
 
-  @impl GenServer
-  def handle_info(
-        {:DOWN, monitoring_reference, _process, dead_process_pid, _reason},
-        %State{} = state
-      ) do
-    case Monitoring.remove(state.monitor_id, monitoring_reference) do
-      :temporary_worker ->
-        {:noreply,
-         %State{
-           state
-           | overflow: state.overflow - 1,
-             idle_workers_state:
-               IdleWorkers.remove(
-                 state.idle_workers_impl,
-                 state.idle_workers_state,
-                 dead_process_pid
-               )
-         }}
+  @spec add_worker_to_busy_workers(State.t(), worker()) :: State.t()
+  defp add_worker_to_busy_workers(%State{} = state, worker) do
+    %State{
+      state
+      | busy_workers_state:
+          BusyWorkers.add(
+            state.busy_workers_impl,
+            state.busy_workers_state,
+            worker
+          )
+    }
+  end
 
-      :worker ->
-        {:ok, new_worker} = start_worker(state)
+  @spec add_worker_to_idle_workers(State.t(), worker()) :: State.t()
+  defp add_worker_to_idle_workers(%State{} = state, worker) do
+    %State{
+      state
+      | idle_workers_state:
+          IdleWorkers.add(
+            state.idle_workers_impl,
+            state.busy_workers_state,
+            worker
+          )
+    }
+  end
 
-        Monitoring.add(state.monitor_id, new_worker, :worker)
+  @spec remove_worker_from_busy_workers(State.t(), worker()) :: State.t()
+  defp remove_worker_from_busy_workers(%State{} = state, worker) do
+    %State{
+      state
+      | busy_workers_state:
+          BusyWorkers.remove(
+            state.busy_workers_impl,
+            state.busy_workers_state,
+            worker
+          )
+    }
+  end
 
-        temp_idle_workers_state =
-          IdleWorkers.remove(state.idle_workers_impl, state.idle_workers_state, dead_process_pid)
+  @spec remove_worker_from_idle_workers(State.t(), worker()) :: State.t()
+  defp remove_worker_from_idle_workers(%State{} = state, worker) do
+    %State{
+      state
+      | idle_workers_state:
+          IdleWorkers.remove(
+            state.idle_workers_impl,
+            state.idle_workers_state,
+            worker
+          )
+    }
+  end
 
-        new_idle_workers_state =
-          IdleWorkers.add(state.idle_workers_impl, temp_idle_workers_state, new_worker)
+  @spec handle_down_worker(State.t(), pid()) :: State.t()
+  defp handle_down_worker(%State{} = state, dead_process_pid) do
+    state =
+      state
+      |> remove_worker_from_idle_workers(dead_process_pid)
+      |> remove_worker_from_busy_workers(dead_process_pid)
 
-        state = %State{
-          state
-          | idle_workers_state: new_idle_workers_state,
-            busy_workers_state:
-              BusyWorkers.remove(
-                state.busy_workers_impl,
-                state.busy_workers_state,
-                dead_process_pid
-              )
-        }
+    if WaitingCallers.empty?(state.waiting_callers_impl, state.waiting_callers_state) do
+      if state.overflow > 0 do
+        %State{state | overflow: state.overflow - 1}
+      else
+        {:ok, new_worker} = start_worker(state)
 
-        {:noreply, state}
+        Monitoring.add(state.monitor_id, new_worker, :worker)
 
-      :caller ->
-        new_waiting_callers_state =
-          WaitingCallers.remove_by_pid(
-            state.waiting_callers_impl,
-            state.waiting_callers_state,
-            dead_process_pid
-          )
+        add_worker_to_idle_workers(state, new_worker)
+      end
+    else
+      {:ok, new_worker} = start_worker(state)
+      Monitoring.add(state.monitor_id, new_worker, :worker)
 
-        {:noreply, %{state | waiting_callers_state: new_waiting_callers_state}}
+      state
+      |> add_worker_to_busy_workers(new_worker)
+      |> provide_worker_to_waiting_caller(new_worker)
     end
   end
 
+  @spec handle_down_caller(State.t(), pid()) :: State.t()
+  defp handle_down_caller(%State{} = state, dead_process_pid) do
+    new_waiting_callers_state =
+      WaitingCallers.remove_by_pid(
+        state.waiting_callers_impl,
+        state.waiting_callers_state,
+        dead_process_pid
+      )
+
+    %State{state | waiting_callers_state: new_waiting_callers_state}
+  end
+
   @impl GenServer
   def terminate(reason, %State{} = state) do
     DynamicSupervisor.stop(state.supervisor, reason)
diff --git a/lib/poolex/monitoring.ex b/lib/poolex/monitoring.ex
index 3f4a682..e3b731d 100644
--- a/lib/poolex/monitoring.ex
+++ b/lib/poolex/monitoring.ex
@@ -1,7 +1,7 @@
 defmodule Poolex.Monitoring do
   @moduledoc false
   @type monitor_id() :: atom() | reference()
-  @type kind_of_process() :: :worker | :caller | :temporary_worker
+  @type kind_of_process() :: :worker | :caller
 
   @spec init(Poolex.pool_id()) :: {:ok, monitor_id()}
   @doc false


Test patch:
diff --git a/test/poolex_test.exs b/test/poolex_test.exs
index 23b5175..ebf3a37 100644
--- a/test/poolex_test.exs
+++ b/test/poolex_test.exs
@@ -155,6 +155,31 @@ defmodule PoolexTest do
       assert agent_pid != new_agent_pid
     end
 
+    test "restart busy workers when pending callers" do
+      pool_name = start_pool(worker_module: SomeWorker, workers_count: 1)
+
+      # test_process = self()
+      launch_long_tasks(pool_name, 2)
+
+      debug_info = Poolex.get_debug_info(pool_name)
+      assert debug_info.busy_workers_count == 1
+      assert length(debug_info.waiting_callers) == 1
+
+      [busy_worker_pid] = debug_info.busy_workers_pids
+      Process.exit(busy_worker_pid, :kill)
+
+      # To be sure that DOWN message will be handed
+      :timer.sleep(1)
+
+      debug_info = Poolex.get_debug_info(pool_name)
+      assert debug_info.busy_workers_count == 1
+      assert Enum.empty?(debug_info.waiting_callers)
+
+      [new_worker_pid] = debug_info.busy_workers_pids
+
+      assert busy_worker_pid != new_worker_pid
+    end
+
     test "works on callers" do
       pool_name = start_pool(worker_module: SomeWorker, workers_count: 1)