Add :ordered option to async_stream (#6234)

Author: ericmj

lib/elixir/lib/task.ex

@@ -344,9 +344,10 @@ defmodule Task do
   terminates the current process and a failure in the current process
   terminates all tasks.
 
-  When streamed, each task will emit `{:ok, val}` upon successful
-  completion or `{:exit, val}` if the caller is trapping exits. Results
-  are emitted in the same order as the original `enumerable`.
+  When streamed, each task will emit `{:ok, value}` upon successful
+  completion or `{:exit, reason, element}` if the caller is trapping
+  exits, where `element` is the stream element. Results are emitted
+  in the same order as the original `enumerable`.
 
   The level of concurrency can be controlled via the `:max_concurrency`
   option and defaults to `System.schedulers_online/0`. A timeout
@@ -362,13 +363,18 @@ defmodule Task do
 
     * `:max_concurrency` - sets the maximum number of tasks to run
       at the same time. Defaults to `System.schedulers_online/0`.
+    * `:ordered` - whether the results should be returned in the same order
+      as the input stream. This option is useful when you have large
+      streams and don't want to buffer results before they are delivered.
+      Defaults to `true`.
     * `:timeout` - the maximum amount of time (in milliseconds) each
       task is allowed to execute for. Defaults to `5000`.
     * `:on_timeout` - what do to when a task times out. The possible
       values are:
       * `:exit` (default) - the process that spawned the tasks exits.
       * `:kill_task` - the task that timed out is killed. The value
-        emitted for that task is `{:exit, :timeout}`.
+        emitted for that task is `{:exit, :timeout, element}`, where
+        `element` is the element it timed out on.
 
   ## Example
 

lib/elixir/lib/task/supervised.ex

@@ -145,6 +145,7 @@ defmodule Task.Supervised do
   def stream(enumerable, acc, reducer, mfa, options, spawn) do
     next = &Enumerable.reduce(enumerable, &1, fn x, acc -> {:suspend, [x | acc]} end)
     max_concurrency = Keyword.get(options, :max_concurrency, System.schedulers_online)
+    ordered? = Keyword.get(options, :ordered, true)
     timeout = Keyword.get(options, :timeout, 5000)
     on_timeout = Keyword.get(options, :on_timeout, :exit)
     parent = self()
@@ -168,6 +169,7 @@ defmodule Task.Supervised do
       reducer: reducer,
       monitor_pid: monitor_pid,
       monitor_ref: monitor_ref,
+      ordered: ordered?,
       timeout: timeout,
       on_timeout: on_timeout,
     }
@@ -200,44 +202,51 @@ defmodule Task.Supervised do
        when max == 0
        when next == :done do
     %{monitor_pid: monitor_pid, monitor_ref: monitor_ref,
-      timeout: timeout, on_timeout: on_timeout} = config
+      timeout: timeout, on_timeout: on_timeout, ordered: ordered?} = config
 
     receive do
       # The task at position "position" replied with "value". We put the
       # response in the "waiting" map and do nothing, since we'll only act on
       # this response when the replying task dies (we'll notice in the :down
       # message).
-      {{^monitor_ref, position}, value} ->
-        %{^position => {pid, :running}} = waiting
-        waiting = Map.put(waiting, position, {pid, {:ok, value}})
+      {{^monitor_ref, position}, reply} ->
+        %{^position => {pid, :running, element}} = waiting
+        waiting = Map.put(waiting, position, {pid, {:ok, reply}, element})
         stream_reduce({:cont, acc}, max, spawned, delivered, waiting, next, config)
 
       # The task at position "position" died for some reason. We check if it
       # replied already (then the death is peaceful) or if it's still running
       # (then the reply from this task will be {:exit, reason}). This message is
       # sent to us by the monitor process, not by the dying task directly.
       {kind, {^monitor_ref, position}, reason} when kind in [:down , :timed_out] ->
-        waiting =
+        result =
           case waiting do
             # If the task replied, we don't care whether it went down for timeout
             # or for normal reasons.
-            %{^position => {_, {:ok, _} = ok}} ->
-              Map.put(waiting, position, {nil, ok})
+            %{^position => {_, {:ok, _} = ok, _}} ->
+              ok
             # If the task exited by itself before replying, we emit {:exit, reason}.
-            %{^position => {_, :running}} when kind == :down ->
-              Map.put(waiting, position, {nil, {:exit, reason}})
+            %{^position => {_, :running, element}} when kind == :down ->
+              {:exit, reason, element}
             # If the task timed out before replying, we either exit (on_timeout: :exit)
             # or emit {:exit, :timeout} (on_timeout: :kill_task) (note the task is already
             # dead at this point).
-            %{^position => {_, :running}} when kind == :timed_out ->
+            %{^position => {_, :running, element}} when kind == :timed_out ->
               if on_timeout == :exit do
                 stream_cleanup_inbox(monitor_pid, monitor_ref)
-                exit({:timeout, {__MODULE__, :stream, [timeout]}})
+                exit({:timeout, {__MODULE__, :stream, [element, timeout]}})
               else
-                Map.put(waiting, position, {nil, {:exit, :timeout}})
+                {:exit, :timeout, element}
               end
           end
-        stream_deliver({:cont, acc}, max + 1, spawned, delivered, waiting, next, config)
+
+        if ordered? do
+          waiting = Map.put(waiting, position, {:done, result})
+          stream_deliver({:cont, acc}, max + 1, spawned, delivered, waiting, next, config)
+        else
+          pair = deliver_now(result, acc, next, config)
+          stream_reduce(pair, max + 1, spawned, delivered + 1, waiting, next, config)
+        end
 
       # The monitor process died. We just cleanup the messages from the monitor
       # process and exit.
@@ -269,6 +278,20 @@ defmodule Task.Supervised do
     end
   end
 
+  defp deliver_now(reply, acc, next, config) do
+    %{reducer: reducer, monitor_pid: monitor_pid,
+      monitor_ref: monitor_ref, timeout: timeout} = config
+    try do
+      reducer.(reply, acc)
+    catch
+      kind, reason ->
+        stacktrace = System.stacktrace()
+        is_function(next) && next.({:halt, []})
+        stream_close(monitor_pid, monitor_ref, timeout)
+        :erlang.raise(kind, reason, stacktrace)
+    end
+  end
+
   defp stream_deliver({:suspend, acc}, max, spawned, delivered, waiting, next, config) do
     continuation = &stream_deliver(&1, max, spawned, delivered, waiting, next, config)
     {:suspended, acc, continuation}
@@ -283,7 +306,7 @@ defmodule Task.Supervised do
       monitor_ref: monitor_ref, timeout: timeout} = config
 
     case waiting do
-      %{^delivered => {nil, reply}} ->
+      %{^delivered => {:done, reply}} ->
         try do
           reducer.(reply, acc)
         catch
@@ -341,7 +364,7 @@ defmodule Task.Supervised do
     receive do
       {:spawned, {^monitor_ref, ^spawned}, pid} ->
         send(pid, {self(), {monitor_ref, spawned}})
-        Map.put(waiting, spawned, {pid, :running})
+        Map.put(waiting, spawned, {pid, :running, value})
       {:DOWN, ^monitor_ref, _, ^monitor_pid, reason} ->
         stream_cleanup_inbox(monitor_pid, monitor_ref)
         exit({reason, {__MODULE__, :stream, [timeout]}})

lib/elixir/lib/task/supervisor.ex

@@ -135,9 +135,10 @@ defmodule Task.Supervisor do
   own task. The tasks will be spawned under the given `supervisor` and
   linked to the current process, similarly to `async/4`.
 
-  When streamed, each task will emit `{:ok, val}` upon successful
-  completion or `{:exit, val}` if the caller is trapping exits. Results
-  are emitted in the same order as the original `enumerable`.
+  When streamed, each task will emit `{:ok, value}` upon successful
+  completion or `{:exit, reason, element}` if the caller is trapping
+  exits, where `element` is the stream element. Results are emitted
+  in the same order as the original `enumerable`.
 
   The level of concurrency can be controlled via the `:max_concurrency`
   option and defaults to `System.schedulers_online/0`. A timeout
@@ -152,9 +153,19 @@ defmodule Task.Supervisor do
 
     * `:max_concurrency` - sets the maximum number of tasks to run
       at the same time. Defaults to `System.schedulers_online/0`.
+    * `:ordered` - whether the results should be returned in the same order
+      as the input stream. This option is useful when you have large
+      streams and don't want to buffer results before they are delivered.
+      Defaults to `true`.
     * `:timeout` - the maximum amount of time to wait (in milliseconds)
       without receiving a task reply (across all running tasks).
       Defaults to `5000`.
+    * `:on_timeout` - what do to when a task times out. The possible
+      values are:
+      * `:exit` (default) - the process that spawned the tasks exits.
+      * `:kill_task` - the task that timed out is killed. The value
+        emitted for that task is `{:exit, :timeout, element}`, where
+        `element` is the element it timed out on.
 
   ## Examples
 

lib/elixir/test/elixir/task/supervisor_test.exs

@@ -219,9 +219,9 @@ defmodule Task.SupervisorTest do
     test "streams an enumerable with exits", %{supervisor: supervisor} do
       Process.flag(:trap_exit, true)
       assert supervisor
-             |> Task.Supervisor.async_stream(1..4, &exit/1, @opts)
+             |> Task.Supervisor.async_stream(1..4, &exit(Integer.to_string(&1)), @opts)
              |> Enum.to_list ==
-             [exit: 1, exit: 2, exit: 3, exit: 4]
+             [{:exit, "1", 1}, {:exit, "2", 2}, {:exit, "3", 3}, {:exit, "4", 4}]
     end
 
     test "shuts down unused tasks", %{supervisor: supervisor} do
@@ -277,7 +277,7 @@ defmodule Task.SupervisorTest do
       assert supervisor
              |> Task.Supervisor.async_stream_nolink(1..4, &exit/1, @opts)
              |> Enum.to_list ==
-             [exit: 1, exit: 2, exit: 3, exit: 4]
+             [{:exit, 1, 1}, {:exit, 2, 2}, {:exit, 3, 3}, {:exit, 4, 4}]
     end
 
     test "shuts down unused tasks", %{supervisor: supervisor} do

lib/elixir/test/elixir/task_test.exs

@@ -492,7 +492,24 @@ defmodule TaskTest do
   describe "async_stream/2" do
     test "timeout" do
       assert catch_exit([:infinity] |> Task.async_stream(&sleep/1, [timeout: 0]) |> Enum.to_list) ==
-             {:timeout, {Task.Supervised, :stream, [0]}}
+             {:timeout, {Task.Supervised, :stream, [:infinity, 0]}}
+      refute_received _
+    end
+
+    test "streams an enumerable with ordered: false" do
+      opts = [max_concurrency: 1, ordered: false]
+      assert 4..1 |> Task.async_stream(&sleep(&1 * 100), opts) |> Enum.to_list ==
+             [ok: 400, ok: 300, ok: 200, ok: 100]
+
+      opts = [max_concurrency: 4, ordered: false]
+      assert 4..1 |> Task.async_stream(&sleep(&1 * 100), opts) |> Enum.to_list ==
+             [ok: 100, ok: 200, ok: 300, ok: 400]
+    end
+
+    test "streams an enumerable with ordered: false, on_timeout: :kill_task" do
+      opts = [max_concurrency: 4, ordered: false, on_timeout: :kill_task, timeout: 50]
+      assert [100, 1, 100, 1] |> Task.async_stream(&sleep/1, opts) |> Enum.to_list() ==
+             [{:ok, 1}, {:ok, 1}, {:exit, :timeout, 100}, {:exit, :timeout, 100}]
       refute_received _
     end
   end
@@ -526,7 +543,7 @@ defmodule TaskTest do
       test "streams an enumerable with exits" do
         Process.flag(:trap_exit, true)
         assert 1..4 |> Task.async_stream(&exit/1, @opts) |> Enum.to_list ==
-               [exit: 1, exit: 2, exit: 3, exit: 4]
+               [{:exit, 1, 1}, {:exit, 2, 2}, {:exit, 3, 3}, {:exit, 4, 4}]
         refute_received {:EXIT, _, _}
       end
 
@@ -569,7 +586,7 @@ defmodule TaskTest do
       test "with inner halt on failure" do
         Process.flag(:trap_exit, true)
         assert 1..8 |> Stream.take(4) |> Task.async_stream(&exit/1, @opts) |> Enum.to_list ==
-               [exit: 1, exit: 2, exit: 3, exit: 4]
+               [{:exit, 1, 1}, {:exit, 2, 2}, {:exit, 3, 3}, {:exit, 4, 4}]
       end
 
       test "with inner halt and slowest first" do
@@ -585,7 +602,7 @@ defmodule TaskTest do
       test "with outer halt on failure" do
         Process.flag(:trap_exit, true)
         assert 1..8 |> Task.async_stream(&exit/1, @opts) |> Enum.take(4) ==
-               [exit: 1, exit: 2, exit: 3, exit: 4]
+               [{:exit, 1, 1}, {:exit, 2, 2}, {:exit, 3, 3}, {:exit, 4, 4}]
       end
 
       test "with outer halt and slowest first" do
@@ -622,7 +639,7 @@ defmodule TaskTest do
       test "with :on_timeout set to :kill_task" do
         opts = Keyword.merge(@opts, on_timeout: :kill_task, timeout: 50)
         assert [100, 1, 100, 1] |> Task.async_stream(&sleep/1, opts) |> Enum.to_list() ==
-               [exit: :timeout, ok: 1, exit: :timeout, ok: 1]
+               [{:exit, :timeout, 100}, {:ok, 1}, {:exit, :timeout, 100}, {:ok, 1}]
         refute_received _
       end
     end