Ensure suspensions are closed on failures

Author: José Valim

lib/elixir/lib/stream.ex

@@ -109,8 +109,22 @@ defmodule Stream do
               last && { last, fun }, :lists.reverse(accs), acc)
     end
 
+    defp do_each(reduce, [], last, accs, { command, acc }) do
+      handle_reduce reduce.({ command, [acc|accs] }), [], last
+    end
+
     defp do_each(reduce, after_funs, last, accs, { command, acc }) do
-      case reduce.({ command, [acc|accs] }) do
+      try do
+        handle_reduce reduce.({ command, [acc|accs] }), after_funs, last
+      catch
+        kind, reason ->
+          lc fun inlist after_funs, do: fun.()
+          :erlang.raise(kind, reason, :erlang.get_stacktrace)
+      end
+    end
+
+    defp handle_reduce(res, after_funs, last) do
+      case res do
         { :suspended, [acc|accs], continuation } ->
           { :suspended, acc, &do_each(continuation, after_funs, last, accs, &1) }
         { :halted, [acc|_] } ->
@@ -456,8 +470,7 @@ defmodule Stream do
   defp do_flat_map(next_acc, next, mapper, acc, fun) do
     case next.({ :cont, next_acc }) do
       { :suspended, [val|next_acc], next } ->
-        enum = mapper.(val)
-        do_flat_map(next_acc, next, mapper, acc, fun, &Enumerable.reduce(enum, &1, fun))
+        do_flat_map(next_acc, next, mapper, acc, fun, &Enumerable.reduce(mapper.(val), &1, fun))
       { reason, _ } ->
         { reason, elem(acc, 1) }
     end
@@ -468,9 +481,11 @@ defmodule Stream do
       reduce.(acc)
     catch
       { :stream_flat_map, h } ->
-        reduce.({ :halt, elem(acc, 1) })
         next.({ :halt, next_acc })
         { :halted, h }
+      kind, reason ->
+        next.({ :halt, next_acc })
+        :erlang.raise(kind, reason, :erlang.get_stacktrace)
     else
       { _, acc }             -> do_flat_map(next_acc, next, mapper, { :cont, acc }, fun)
       { :suspended, acc, c } -> { :suspended, acc, &do_flat_map(next_acc, next, mapper, &1, fun, c) }
@@ -734,7 +749,18 @@ defmodule Stream do
   end
 
   defp do_zip(zips, { :cont, acc }, callback) do
-    do_zip(zips, acc, callback, [], [])
+    try do
+      do_zip(zips, acc, callback, [], [])
+    catch
+      kind, reason ->
+        do_zip_close(zips)
+        :erlang.raise(kind, reason, :erlang.get_stacktrace)
+    else
+      { :next, buffer, acc } ->
+        do_zip(buffer, acc, callback)
+      { :done, _ } = o ->
+        o
+    end
   end
 
   defp do_zip([{ fun, fun_acc }|t], acc, callback, list, buffer) do
@@ -749,7 +775,7 @@ defmodule Stream do
 
   defp do_zip([], acc, callback, list, buffer) do
     zipped = list_to_tuple(:lists.reverse(list))
-    do_zip(:lists.reverse(buffer), callback.(zipped, acc), callback)
+    { :next, :lists.reverse(buffer), callback.(zipped, acc) }
   end
 
   defp do_zip_close([]), do: :ok
@@ -892,7 +918,10 @@ defmodule Stream do
 
   defp do_resource(next_acc, next_fun, { :cont, acc }, fun, after_fun) do
     try do
-      next_fun.(next_acc)
+      case next_fun.(next_acc) do
+        nil -> nil
+        { v, next_acc } -> { fun.(v, acc), next_acc }
+      end
     catch
       kind, reason ->
         after_fun.(next_acc)
@@ -901,8 +930,8 @@ defmodule Stream do
       nil ->
         after_fun.(next_acc)
         { :done, acc }
-      { v, next_acc } ->
-        do_resource(next_acc, next_fun, fun.(v, acc), fun, after_fun)
+      { acc, next_acc } ->
+        do_resource(next_acc, next_fun, acc, fun, after_fun)
     end
   end
 

lib/elixir/test/elixir/stream_test.exs

@@ -26,17 +26,26 @@ defmodule StreamTest do
   end
 
   test "after" do
-    stream = Stream.after([1,2,3], fn -> Process.put(:after, true) end)
+    stream = Stream.after([1,2,3], fn -> Process.put(:stream_after, true) end)
 
     # Done
-    Process.put(:after, false)
+    Process.put(:stream_after, false)
     assert Enum.to_list(stream) == [1,2,3]
-    assert Process.get(:after)
+    assert Process.get(:stream_after)
 
     # Halted
-    Process.put(:after, false)
+    Process.put(:stream_after, false)
     assert Enum.take(stream, 1) == [1]
-    assert Process.get(:after)
+    assert Process.get(:stream_after)
+  end
+
+  test "after closes on errors" do
+    stream = Stream.after([1,2,3], fn -> Process.put(:stream_after, true) end)
+
+    Process.put(:stream_after, false)
+    stream = Stream.map(stream, fn x -> if x > 2, do: throw(:error), else: x end)
+    assert catch_throw(Enum.to_list(stream)) == :error
+    assert Process.get(:stream_after)
   end
 
   test "chunk" do
@@ -256,7 +265,7 @@ defmodule StreamTest do
     assert Enum.zip(list, list) == Enum.zip(stream, stream)
   end
 
-  test "flat_map does not leave stream suspended" do
+  test "flat_map does not leave inner stream suspended" do
     stream = Stream.flat_map [1,2,3],
       fn i ->
         Stream.resource(fn -> i end,
@@ -267,7 +276,9 @@ defmodule StreamTest do
     Process.put(:stream_flat_map, false)
     assert stream |> Enum.take(3) == [1,2,3]
     assert Process.get(:stream_flat_map)
+  end
 
+  test "flat_map does not leave outer stream suspended" do
     stream = Stream.resource(fn -> 1 end,
                              fn acc -> { acc, acc + 1 } end,
                              fn _ -> Process.put(:stream_flat_map, true) end)
@@ -278,6 +289,17 @@ defmodule StreamTest do
     assert Process.get(:stream_flat_map)
   end
 
+  test "flat_map closes on error" do
+    stream = Stream.resource(fn -> 1 end,
+                             fn acc -> { acc, acc + 1 } end,
+                             fn _ -> Process.put(:stream_flat_map, true) end)
+    stream = Stream.flat_map(stream, fn _ -> throw(:error) end)
+
+    Process.put(:stream_flat_map, false)
+    assert catch_throw(Enum.to_list(stream)) == :error
+    assert Process.get(:stream_flat_map)
+  end
+
   test "iterate" do
     stream = Stream.iterate(0, &(&1+2))
     assert Enum.take(stream, 5) == [0,2,4,6,8]
@@ -402,6 +424,17 @@ defmodule StreamTest do
     assert Enum.to_list(stream) == [5, 4, 3, 2, 1]
   end
 
+  test "resource closes on errors" do
+    stream = Stream.resource(fn -> 1 end,
+                             fn acc -> { acc, acc + 1 } end,
+                             fn _ -> Process.put(:stream_resource, true) end)
+
+    Process.put(:stream_resource, false)
+    stream = Stream.map(stream, fn x -> if x > 2, do: throw(:error), else: x end)
+    assert catch_throw(Enum.to_list(stream)) == :error
+    assert Process.get(:stream_resource)
+  end
+
   test "resource is zippable" do
     # File.stream! uses Stream.resource underneath
     stream = File.stream!(__FILE__)
@@ -479,6 +512,25 @@ defmodule StreamTest do
     assert Process.get(:stream_zip) == :done
   end
 
+  test "zip/2 closes on inner error" do
+    stream = Stream.after([1, 2, 3], fn -> Process.put(:stream_zip, true) end)
+    stream = Stream.zip(stream, Stream.map([:a, :b, :c], fn _ -> throw(:error) end))
+
+    Process.put(:stream_zip, false)
+    assert catch_throw(Enum.to_list(stream)) == :error
+    assert Process.get(:stream_zip)
+  end
+
+  test "zip/2 closes on outer error" do
+    stream = Stream.after([1, 2, 3], fn -> Process.put(:stream_zip, true) end)
+             |> Stream.zip([:a, :b, :c])
+             |> Stream.map(fn _ -> throw(:error) end)
+
+    Process.put(:stream_zip, false)
+    assert catch_throw(Enum.to_list(stream)) == :error
+    assert Process.get(:stream_zip)
+  end
+
   test "with_index" do
     stream = Stream.with_index([1,2,3])
     assert is_lazy(stream)