Add GenStage.demand/1

Author: José Valim

lib/gen_stage.ex

@@ -1231,6 +1231,16 @@ defmodule GenStage do
     cast(stage, {:"$info", msg})
   end
 
+  @doc """
+  Returns the demand mode for a producer.
+
+  It is either `:forward` or `:accumulate`. See `demand/2`.
+  """
+  @spec demand(stage) :: :forward | :accumulate
+  def demand(stage) do
+    call(stage, :"$demand")
+  end
+
   @doc """
   Sets the demand mode for a producer.
 
@@ -1784,6 +1794,10 @@ defmodule GenStage do
     producer_info(msg, stage)
   end
 
+  def handle_call(:"$demand", _from, stage) do
+    producer_demand(stage)
+  end
+
   def handle_call({:"$subscribe", current, to, opts}, _from, stage) do
     consumer_subscribe(current, to, opts, stage)
   end
@@ -2089,6 +2103,14 @@ defmodule GenStage do
 
   ## Producer helpers
 
+  defp producer_demand(%{events: :forward} = stage) do
+    {:reply, :forward, stage}
+  end
+
+  defp producer_demand(%{events: events} = stage) when is_list(events) do
+    {:reply, :accumulate, stage}
+  end
+
   defp producer_demand(:forward, %{type: :producer_consumer} = stage) do
     # That's the only mode on producer consumers.
     {:noreply, stage}

test/gen_stage_test.exs

@@ -627,26 +627,33 @@ defmodule GenStageTest do
   describe "demand" do
     test "can be set to :accumulate on init" do
       {:ok, producer} = Counter.start_link({:producer, 0, demand: :accumulate})
+      assert GenStage.demand(producer) == :accumulate
       {:ok, consumer} = Forwarder.start_link({:consumer, self()})
       GenStage.sync_subscribe(consumer, to: producer, max_demand: 4, min_demand: 0)
       refute_received {:consumed, [0, 1, 2, 3]}
       GenStage.demand(producer, :forward)
+      assert GenStage.demand(producer) == :forward
       assert_receive {:consumed, [0, 1, 2, 3]}
     end
 
     test "can be set to :accumulate via API" do
       {:ok, producer} = Counter.start_link({:producer, 0})
+      assert GenStage.demand(producer) == :forward
       GenStage.demand(producer, :accumulate)
+      assert GenStage.demand(producer) == :accumulate
       {:ok, consumer} = Forwarder.start_link({:consumer, self()})
       GenStage.sync_subscribe(consumer, to: producer, max_demand: 4, min_demand: 0)
       refute_received {:consumed, [0, 1, 2, 3]}
       GenStage.demand(producer, :forward)
+      assert GenStage.demand(producer) == :forward
       assert_receive {:consumed, [0, 1, 2, 3]}
     end
 
     test "can be set to :forward via API before subscriptions" do
       {:ok, producer} = Counter.start_link({:producer, 0, demand: :accumulate})
+      assert GenStage.demand(producer) == :accumulate
       GenStage.demand(producer, :forward)
+      assert GenStage.demand(producer) == :forward
       {:ok, consumer} = Forwarder.start_link({:consumer, self()})
       GenStage.sync_subscribe(consumer, to: producer, max_demand: 4, min_demand: 0)
       assert_receive {:consumed, [0, 1, 2, 3]}