Allow explicitly specifying :max_demand in DemandDispatcher (#291)

Author: lexmag

lib/gen_stage/dispatchers/demand_dispatcher.ex

@@ -11,9 +11,12 @@ defmodule GenStage.DemandDispatcher do
   The demand dispatcher accepts the following options
   on initialization:
 
-    * `:shuffle_demands_on_first_dispatch` - when true, shuffle the initial demands list
+    * `:shuffle_demands_on_first_dispatch` - when `true`, shuffle the initial demands list
       which is constructed on subscription before first dispatch. It prevents overloading
-      the first consumer on first dispatch.  Defaults to `false`.
+      the first consumer on first dispatch. Defaults to `false`.
+
+    * `:max_demand` - the maximum demand expected on `GenStage.ask/3`.
+      Defaults to the first demand asked.
 
   ### Examples
 
@@ -27,8 +30,9 @@ defmodule GenStage.DemandDispatcher do
   @doc false
   def init(opts) do
     shuffle_demand = Keyword.get(opts, :shuffle_demands_on_first_dispatch, false)
+    max_demand = Keyword.get(opts, :max_demand)
 
-    {:ok, {[], 0, nil, shuffle_demand}}
+    {:ok, {[], 0, max_demand, shuffle_demand}}
   end
 
   @doc false
@@ -102,7 +106,7 @@ defmodule GenStage.DemandDispatcher do
     {now, later, length - counter, 0}
   end
 
-  defp add_demand(counter, pid, ref, [{c, _, _} | _] = demands) when counter > c do
+  defp add_demand(counter, pid, ref, [{current, _, _} | _] = demands) when counter > current do
     [{counter, pid, ref} | demands]
   end
 

test/gen_stage/demand_dispatcher_test.exs

@@ -8,7 +8,8 @@ defmodule GenStage.DemandDispatcherTest do
 
   defp dispatcher(opts) do
     shuffle_demand = Keyword.get(opts, :shuffle_demands_on_first_dispatch, false)
-    {:ok, {[], 0, nil, ^shuffle_demand} = state} = D.init(opts)
+    max_demand = Keyword.get(opts, :max_demand)
+    {:ok, {[], 0, ^max_demand, ^shuffle_demand} = state} = D.init(opts)
     state
   end
 
@@ -191,14 +192,13 @@ defmodule GenStage.DemandDispatcherTest do
     pid = self()
     ref1 = make_ref()
     ref2 = make_ref()
-    disp = dispatcher([])
+    disp = dispatcher(max_demand: 3)
 
     {:ok, 0, disp} = D.subscribe([], {pid, ref1}, disp)
     {:ok, 0, disp} = D.subscribe([], {pid, ref2}, disp)
 
     log =
       capture_log(fn ->
-        {:ok, 3, disp} = D.ask(3, {pid, ref1}, disp)
         {:ok, 4, disp} = D.ask(4, {pid, ref2}, disp)
         disp
       end)