Add :limit option to Task.yield_many/2

Author: josevalim

lib/elixir/lib/task.ex

@@ -1109,7 +1109,8 @@ defmodule Task do
  * `{:ok, term}` if the task has successfully reported its
  result back in the given time interval
  * `{:exit, reason}` if the task has died
- * `nil` if the task keeps running past the timeout
+ * `nil` if the task keeps running, either because a limit
+ has been reached or past the timeout
 
  Check `yield/2` for more information.
 
@@ -1162,6 +1163,10 @@ defmodule Task do
  The second argument is either a timeout or options, which defaults
  to this:
 
+ * `:limit` - the maximum amount of tasks to wait for.
+ If the limit is reached before the timeout, this function
+ returns immediately without triggering the `:on_timeout` behaviour
+
  * `:timeout` - the maximum amount of time (in milliseconds or `:infinity`)
  each task is allowed to execute for. Defaults to `5000`.
 
@@ -1173,7 +1178,11 @@ defmodule Task do
  * `:kill_task` - the task that timed out is killed.
  """
  @spec yield_many([t], timeout) :: [{t, {:ok, term} | {:exit, term} | nil}]
- @spec yield_many([t], timeout: timeout, on_timeout: :nothing | :ignore | :kill_task) ::
+ @spec yield_many([t],
+ limit: pos_integer(),
+ timeout: timeout,
+ on_timeout: :nothing | :ignore | :kill_task
+ ) ::
  [{t, {:ok, term} | {:exit, term} | nil}]
  def yield_many(tasks, opts \\ [])
 
@@ -1182,9 +1191,6 @@ defmodule Task do
  end
 
  def yield_many(tasks, opts) when is_list(opts) do
- on_timeout = Keyword.get(opts, :on_timeout, :nothing)
- timeout = Keyword.get(opts, :timeout, 5_000)
-
  refs =
  Map.new(tasks, fn %Task{ref: ref, owner: owner} = task ->
  if owner != self() do
@@ -1194,6 +1200,9 @@ defmodule Task do
  {ref, nil}
  end)
 
+ on_timeout = Keyword.get(opts, :on_timeout, :nothing)
+ timeout = Keyword.get(opts, :timeout, 5_000)
+ limit = Keyword.get(opts, :limit, map_size(refs))
  timeout_ref = make_ref()
 
  timer_ref =
@@ -1202,16 +1211,17 @@ defmodule Task do
  end
 
  try do
- yield_many(map_size(refs), refs, timeout_ref, timer_ref)
+ yield_many(limit, refs, timeout_ref, timer_ref)
  catch
  {:noconnection, reason} ->
  exit({reason, {__MODULE__, :yield_many, [tasks, timeout]}})
  else
- refs ->
+ {timed_out?, refs} ->
  for task <- tasks do
  value =
  with nil <- Map.fetch!(refs, task.ref) do
  case on_timeout do
+ _ when not timed_out? -> nil
  :nothing -> nil
  :kill_task -> shutdown(task, :brutal_kill)
  :ignore -> ignore(task)
@@ -1226,7 +1236,7 @@ defmodule Task do
  defp yield_many(0, refs, timeout_ref, timer_ref) do
  timer_ref && Process.cancel_timer(timer_ref)
  receive do: (^timeout_ref -> :ok), after: (0 -> :ok)
- refs
+ {false, refs}
  end
 
  defp yield_many(limit, refs, timeout_ref, timer_ref) do
@@ -1243,7 +1253,7 @@ defmodule Task do
  end
 
  ^timeout_ref ->
- refs
+ {true, refs}
  end
  end
 

lib/elixir/test/elixir/task_test.exs

@@ -601,6 +601,37 @@ defmodule TaskTest do
  [{task1, {:ok, :result}}, {task2, nil}, {task3, {:exit, :normal}}]
  end
 
+ test "returns results from multiple tasks with limit" do
+ task1 = %Task{ref: make_ref(), owner: self(), pid: nil, mfa: {__MODULE__, :test, 1}}
+ task2 = %Task{ref: make_ref(), owner: self(), pid: nil, mfa: {__MODULE__, :test, 1}}
+ task3 = %Task{ref: make_ref(), owner: self(), pid: nil, mfa: {__MODULE__, :test, 1}}
+
+ send(self(), {task1.ref, :result})
+ send(self(), {:DOWN, task3.ref, :process, self(), :normal})
+
+ assert Task.yield_many([task1, task2, task3], limit: 1, timeout: :infinity) ==
+ [{task1, {:ok, :result}}, {task2, nil}, {task3, nil}]
+
+ assert Task.yield_many([task2, task3], limit: 1, timeout: :infinity) ==
+ [{task2, nil}, {task3, {:exit, :normal}}]
+ end
+
+ test "returns results from multiple tasks with limit and on timeout" do
+ Process.flag(:trap_exit, true)
+ task1 = Task.async(fn -> Process.sleep(:infinity) end)
+ task2 = Task.async(fn -> :done end)
+
+ assert Task.yield_many([task1, task2], timeout: :infinity, on_timeout: :kill_task, limit: 1) ==
+ [{task1, nil}, {task2, {:ok, :done}}]
+
+ assert Process.alive?(task1.pid)
+
+ assert Task.yield_many([task1], timeout: 0, on_timeout: :kill_task, limit: 1) ==
+ [{task1, nil}]
+
+ refute Process.alive?(task1.pid)
+ end
+
  test "returns results on infinity timeout" do
  task1 = %Task{ref: make_ref(), owner: self(), pid: nil, mfa: {__MODULE__, :test, 1}}
  task2 = %Task{ref: make_ref(), owner: self(), pid: nil, mfa: {__MODULE__, :test, 1}}