Repository: general-CbIC/poolex
Pull request #98 (2024-08-12 11:59:52)

Implement `remove_idle_workers!/2`


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6cb1201..52515b3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,7 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
-- Added function `add_idle_workers!/2` for adding some idle workers to the initialized pool in runtime.
+- Added functions `add_idle_workers!/2` and `remove_idle_workers!/2` for changing count of idle workers in runtime.
 
 ### Changed
 
diff --git a/lib/poolex.ex b/lib/poolex.ex
index 0e45c01..943eea4 100644
--- a/lib/poolex.ex
+++ b/lib/poolex.ex
@@ -268,6 +268,22 @@ defmodule Poolex do
     GenServer.call(pool_id, {:add_idle_workers, workers_count})
   end
 
+  @doc """
+  Removes some idle workers from existing pool.
+  If the number of workers to remove is greater than the number of idle workers, all idle workers will be removed.
+  """
+  @spec remove_idle_workers!(pool_id(), pos_integer()) :: :ok | no_return()
+  def remove_idle_workers!(_pool_id, workers_count) when workers_count < 1 do
+    message = "workers_count must be positive number, received: #{inspect(workers_count)}"
+
+    raise ArgumentError, message
+  end
+
+  def remove_idle_workers!(pool_id, workers_count)
+      when is_atom(pool_id) and is_integer(workers_count) do
+    GenServer.call(pool_id, {:remove_idle_workers, workers_count})
+  end
+
   @impl GenServer
   def init(opts) do
     Process.flag(:trap_exit, true)
@@ -414,6 +430,19 @@ defmodule Poolex do
     {:reply, :ok, new_state}
   end
 
+  @impl GenServer
+  def handle_call({:remove_idle_workers, workers_count}, _from, %State{} = state) do
+    new_state =
+      state
+      |> IdleWorkers.to_list()
+      |> Enum.take(workers_count)
+      |> Enum.reduce(state, fn worker, acc_state ->
+        IdleWorkers.remove(acc_state, worker)
+      end)
+
+    {:reply, :ok, new_state}
+  end
+
   @impl GenServer
   def handle_cast({:release_busy_worker, worker}, %State{} = state) do
     if WaitingCallers.empty?(state) do


Test patch:
diff --git a/test/poolex_test.exs b/test/poolex_test.exs
index 76030af..4f5bde1 100644
--- a/test/poolex_test.exs
+++ b/test/poolex_test.exs
@@ -610,4 +610,40 @@ defmodule PoolexTest do
       end)
     end
   end
+
+  describe "remove_idle_workers!/2" do
+    test "removes idle workers from pool" do
+      initial_fun = fn -> 0 end
+
+      pool_name = start_pool(worker_module: Agent, worker_args: [initial_fun], workers_count: 5)
+
+      assert %DebugInfo{idle_workers_count: 5} = Poolex.get_debug_info(pool_name)
+      assert :ok = Poolex.remove_idle_workers!(pool_name, 2)
+      assert %DebugInfo{idle_workers_count: 3} = Poolex.get_debug_info(pool_name)
+    end
+
+    test "removes all idle workers when argument is bigger than idle_workers count" do
+      initial_fun = fn -> 0 end
+
+      pool_name = start_pool(worker_module: Agent, worker_args: [initial_fun], workers_count: 3)
+
+      assert %DebugInfo{idle_workers_count: 3} = Poolex.get_debug_info(pool_name)
+      assert :ok = Poolex.remove_idle_workers!(pool_name, 5)
+      assert %DebugInfo{idle_workers_count: 0} = Poolex.get_debug_info(pool_name)
+    end
+
+    test "raises error on non positive workers_count" do
+      initial_fun = fn -> 0 end
+
+      pool_name = start_pool(worker_module: Agent, worker_args: [initial_fun], workers_count: 5)
+
+      assert_raise(ArgumentError, fn ->
+        Poolex.remove_idle_workers!(pool_name, -1)
+      end)
+
+      assert_raise(ArgumentError, fn ->
+        Poolex.remove_idle_workers!(pool_name, 0)
+      end)
+    end
+  end
 end