Repository: plataformatec/broadway
Pull request #134 (2019-11-16 12:58:34)

Add support for rate limiting producers
Initial work on rate limiting producers. I think that this is the right direction, butI might be wrong so let me know :) Some questions:

  * we should probably limit the size of the buffer where we store messages. Should we have `:buffer_keep` and `:buffer_size` like GenStage has?

  * right now, the rate limiting is per producer. At work, we need rate limiting per pipeline, that is, the total number of messages consumed needs to be under a certain rate. A solution can be to just have each producer have a rate limit of `total_rate_limit / number_of_producers` per second. This practically works but might not be optimal because if a producer can produce a lot while another producer is slow, the fast producer will be penalized by its own rate limiting instead of sharing the rate limiting and improving throughput. Another thought is that Broadway doesn't have correlation between producers and processors, so I don't know if rate limiting per producer is every useful rather than rate limiting globally. What are your thoughts on this?

josevalim: I think we should go ahead and implement the shared counter. We can implement it with ETS:

1. On every operation, we bump the ETS counter using update_counter. Everything that is above the rate limiting limit, is queued.

2. We have a separate process that is clearing the ets row every X seconds. This process will issue a select_replace. If the value we read is greater than or equal to the limit, we send a message to all producers telling them they can go back to emit items.


Patch:
diff --git a/guides/examples/Apache Kafka.md b/guides/examples/Apache Kafka.md
new file mode 100644
index 0000000..a083722
--- /dev/null
+++ b/guides/examples/Apache Kafka.md
@@ -0,0 +1,225 @@
+# Apache Kafka
+
+Kafka is a distributed streaming platform that has three key capabilities:
+
+  * Publish and subscribe to streams of records
+  * Store streams of records in a fault-tolerant durable way
+  * Process streams of records as they occur
+
+## Getting Started
+
+In order to use Broadway with Kafka, we need to:
+
+  1. Create a stream of records (or use an existing one)
+  1. Configure your Elixir project to use Broadway
+  1. Define your pipeline configuration
+  1. Implement Broadway callbacks
+  1. Run the Broadway pipeline
+
+## Create a stream of records (or use an existing one)
+
+In case you don't have Kafka installed yet, please follow the instructions on Kafka's
+[Quickstart](https://kafka.apache.org/quickstart) for a clean installation. After
+initializing Kafka, you can create a new stream by running:
+
+    kafka-topics --create --zookeeper localhost:2181 --partitions 3 --topic test
+
+## Configure your Elixir project to use Broadway
+
+This guide describes the steps necessary to integrate Broadway with Kafka using
+[BroadwayKafka](https://github.com/plataformatec/broadway_kafka),
+which is a Broadway Kafka Connector provided by [Plataformatec](http://www.plataformatec.com).
+
+BroadwayKafka can subscribe to one or more topics and process streams of records
+using Kafka's [Consumer API](https://kafka.apache.org/documentation.html#consumerapi).
+
+Each GenStage producer initialized by BroadwayKafka will be available as a consumer,
+all registered using the same self-labeled **consumer group**. Each record published to a
+topic/partition will be delivered to one consumer instance within each consumer group.
+
+Bear in mind that a topic/partition can be assigned to any consumer instance that has
+been subscribed using the same consumer group, i.e, any Broadway instance or application
+running on any machine connected to the Kafka cluster.
+
+### Starting a new project
+
+Create a new project running:
+
+    mix new my_app --sup
+
+The `--sup` flag instructs Elixir to generate an application with a supervision tree.
+
+### Setting up dependencies
+
+Add `:broadway_kafka` to the list of dependencies in `mix.exs`:
+
+    def deps do
+      [
+        ...
+        {:broadway_kafka, "~> 0.1.0"}
+      ]
+    end
+
+Don't forget to check for the latest version of dependencies.
+
+## Define the pipeline configuration
+
+Broadway is a process-based behaviour and to define a Broadway pipeline,
+we need to define three functions: `start_link/1`, `handle_message/3`
+and optionally `handle_batch/4`. We will cover `start_link/1` in this
+section and the `handle_` callbacks in the next one.
+
+Similar to other process-based behaviours, `start_link/1` simply
+delegates to `Broadway.start_link/2`, which should define the
+producers, processors, and batchers in the Broadway pipeline.
+Assuming we want to consume messages from a topic called
+`test`, one possible configuration would be:
+
+    defmodule MyBroadway do
+      use Broadway
+
+      alias Broadway.Message
+
+      def start_link(_opts) do
+        Broadway.start_link(__MODULE__,
+          name: __MODULE__,
+          producer: [
+            module: {BroadwayKafka.Producer, [
+              hosts: [localhost: 9092],
+              group_id: "group_1",
+              topics: ["test"],
+            ]},
+            stages: 10
+          ],
+          processors: [
+            default: [
+              stages: 10
+            ]
+          ],
+          batchers: [
+            default: [
+              batch_size: 100,
+              batch_timeout: 200,
+              stages: 10
+            ]
+          ]
+        )
+      end
+
+      ...callbacks...
+    end
+
+> **Note**: Pipelines built on top of BroadwayKafka are automatically partitioned.
+So even though there are multiple processes (stages), these processes will preserve
+Kafka's ordering semantics when it comes to topics/partitions. Internally, this is
+achieved by making sure all messages from the same topic/partition will always be
+forwarded to the same processor and batch processor.
+
+For a full list of options for `BroadwayKafka.Producer`, refer to the
+official [BroadwayKafka](https://hexdocs.pm/broadway_kafka/) documentation.
+
+For general information about setting up Broadway, see `Broadway`
+module docs as well as `Broadway.start_link/2`.
+
+## Implement Broadway callbacks
+
+In order to process incoming messages, we need to implement the
+required callbacks. For the sake of simplicity, we're considering that
+all messages received from the queue are just numbers:
+
+    defmodule MyBroadway do
+      use Broadway
+
+      alias Broadway.Message
+
+      ...start_link...
+
+      @impl true
+      def handle_message(_, message, _) do
+        message
+        |> Message.update_data(fn data -> {data, String.to_integer(data) * 2} end)
+      end
+
+      @impl true
+      def handle_batch(_, messages, _, _) do
+        list = messages |> Enum.map(fn e -> e.data end)
+        IO.inspect(list, label: "Got batch")
+        messages
+      end
+    end
+
+We are not doing anything fancy here, but it should be enough for our
+purpose. First, we update the message's data individually inside
+`handle_message/3` and then we print each batch inside `handle_batch/4`.
+
+For more information, see `c:Broadway.handle_message/3` and
+`c:Broadway.handle_batch/4`.
+
+> Note: Since Broadway v0.2, batching is optional. In case you don't need to
+> group messages as batches for further processing/publishing, you can remove
+> the `:batchers` configuration along with the `handle_batch/4` callback.
+
+## Run the Broadway pipeline
+
+To run your `Broadway` pipeline, you just need to add as a child in
+a supervision tree. Most applications have a supervision tree defined
+at `lib/my_app/application.ex`. You can add Broadway as a child to a
+supervisor as follows:
+
+    children = [
+      {MyBroadway, []}
+    ]
+
+    Supervisor.start_link(children, strategy: :one_for_one)
+
+Now the Broadway pipeline should be started when your application starts.
+Also, if your Broadway has any dependency (for example, it needs to talk
+to the database), make sure that Broadway is listed *after* its dependencies
+in the supervision tree.
+
+You can now test your pipeline by entering an `iex` session:
+
+    iex -S mix
+
+If everything went fine, you should see lots of `info` log messages like this
+one coming from the `:brod` supervisors:
+
+    15:14:04.356 [info]  [supervisor: {:local, :brod_sup}, started: [pid: #PID<0.251.0>, id: :test_client, mfargs: {:brod_client, :start_link, [[localhost: 9092], :test_client, []]}, restart_type: {:permanent, 10}, shutdown: 5000, child_type: :worker]]
+
+ [Brod](https://github.com/klarna/brod/) is the client that BroadwayKafka uses
+ under the hood to communicate with Kafka.
+
+### Sending messages to Kafka
+
+Finally, we can send some sample messages to Kafka using using `:brod` with the following snippet:
+
+    topic = "test"
+    client_id = :my_client
+    hosts = [localhost: 9092]
+
+    :ok = :brod.start_client(hosts, client_id, _client_config=[])
+    :ok = :brod.start_producer(client_id, topic, _producer_config = [])
+
+    Enum.each(1..1000, fn i ->
+      partition = rem(i, 3)
+      :ok = :brod.produce_sync(client_id, topic, partition, _key="", "#{i}")
+    end)
+
+You should see the output showing the generated batches:
+
+    Got batch: [
+      {"2", 4},
+      {"5", 10},
+      {"8", 16},
+      {"11", 22},
+      {"14", 28},
+      ...
+    ]
+    Got batch: [
+      {"3", 6},
+      {"6", 12},
+      {"9", 18},
+      {"12", 24},
+      {"15", 30},
+      ...
+    ]
diff --git a/lib/broadway.ex b/lib/broadway.ex
index af7a68b..6c92a1e 100644
--- a/lib/broadway.ex
+++ b/lib/broadway.ex
@@ -600,6 +600,22 @@ defmodule Broadway do
        `:transformer` callback will cause the whole producer to terminate,
        possibly leaving unacknowledged messages along the way.
 
+    * `:rate_limiting` - Optional. A list of options to enable and configure
+      rate limiting for producing. If this option is present, rate limiting is
+      enabled, otherwise it isn't. Rate limiting refers to the rate at which
+      producers will forward messages to the rest of the pipeline. The rate
+      limiting is applied to and shared by all producers within the time limit.
+      The following options are supported:
+
+        * `:allowed_messages` - Required. An integer that describes how many
+          messages are allowed in the specified interval.
+
+        * `:interval` - Required. An integer that describes the interval
+         (in milliseconds) during which the number of allowed messages is
+         allowed. If the producer produces more than `allowed_messages`
+         in `interval`, only `allowed_messages` will be published until
+         the end of `interval`, and then more messages will be published.
+
     * `:hibernate_after` - Optional. Overrides the top-level `:hibernate_after`.
 
     * `:spawn_opt` - Optional. Overrides the top-level `:spawn_opt`.
@@ -751,6 +767,10 @@ defmodule Broadway do
     in is immediately delivered. When set to `:bulk`, batch is
     delivered when its size or timeout is reached. Defaults to `:flush`.
 
+  * `:metadata` - optionally a map of additional fields to add to the
+    message. This can be used, for example, when testing
+    `BroadwayRabbitMQ.Producer`.
+
   ## Examples
 
   For example, in your tests, you may do:
@@ -761,13 +781,19 @@ defmodule Broadway do
       assert length(failed) == 0
 
   """
-  @spec test_messages(GenServer.server(), data :: [term]) :: reference
+  @spec test_messages(GenServer.server(), data :: [term], opts :: Keyword.t()) :: reference
   def test_messages(broadway, data, opts \\ []) when is_list(data) and is_list(opts) do
     batch_mode = Keyword.get(opts, :batch_mode, :flush)
+    metadata = Map.new(Keyword.get(opts, :metadata, []))
+
     ref = make_ref()
     ack = {Broadway.CallerAcknowledger, {self(), ref}, :ok}
 
-    messages = Enum.map(data, &%Message{data: &1, acknowledger: ack, batch_mode: batch_mode})
+    messages =
+      Enum.map(
+        data,
+        &%Message{data: &1, acknowledger: ack, batch_mode: batch_mode, metadata: metadata}
+      )
 
     :ok = push_messages(broadway, messages)
     ref
@@ -789,7 +815,14 @@ defmodule Broadway do
           stages: [type: :pos_integer, default: 1],
           transformer: [type: :mfa, default: nil],
           spawn_opt: [type: :keyword_list],
-          hibernate_after: [type: :pos_integer]
+          hibernate_after: [type: :pos_integer],
+          rate_limiting: [
+            type: :non_empty_keyword_list,
+            keys: [
+              allowed_messages: [required: true, type: :pos_integer],
+              interval: [required: true, type: :pos_integer]
+            ]
+          ]
         ]
       ],
       processors: [
diff --git a/lib/broadway/producer.ex b/lib/broadway/producer.ex
index 7c8e25f..a4dc133 100644
--- a/lib/broadway/producer.ex
+++ b/lib/broadway/producer.ex
@@ -1,7 +1,8 @@
 defmodule Broadway.Producer do
   @moduledoc false
   use GenStage
-  alias Broadway.Message
+
+  alias Broadway.{Message, RateLimiter}
 
   @doc """
   Invoked once by Broadway during `Broadway.start_link/2`.
@@ -61,6 +62,8 @@ defmodule Broadway.Producer do
     {module, arg} = args[:module]
     transformer = args[:transformer]
     dispatcher = args[:dispatcher]
+    broadway_name = args[:broadway][:name]
+    rate_limiting_options = args[:rate_limiting]
 
     # Inject the topology index only if the args are a keyword list.
     arg =
@@ -70,11 +73,27 @@ defmodule Broadway.Producer do
         arg
       end
 
+    rate_limiting_state =
+      if rate_limiting_options do
+        %{
+          state: :open,
+          draining?: false,
+          table_name: RateLimiter.table_name(broadway_name),
+          # A queue of "batches" of messages that we buffered.
+          message_buffer: :queue.new(),
+          # A queue of demands (integers) that we buffered.
+          demand_buffer: :queue.new()
+        }
+      else
+        nil
+      end
+
     state = %{
       module: module,
       module_state: nil,
       transformer: transformer,
-      consumers: []
+      consumers: [],
+      rate_limiting: rate_limiting_state
     }
 
     case module.init(arg) do
@@ -104,22 +123,17 @@ defmodule Broadway.Producer do
     {:noreply, [], update_in(state.consumers, &List.delete(&1, from))}
   end
 
+  # If we're rate limited, we store the demand in the buffer instead of forwarding it.
+  # We'll forward it once the rate limit is lifted.
   @impl true
-  def handle_demand(demand, state) do
-    %{module: module, transformer: transformer, module_state: module_state} = state
-
-    case module.handle_demand(demand, module_state) do
-      {:noreply, events, new_module_state} when is_list(events) ->
-        messages = transform_events(events, transformer)
-        {:noreply, messages, %{state | module_state: new_module_state}}
-
-      {:noreply, events, new_module_state, :hibernate} ->
-        messages = transform_events(events, transformer)
-        {:noreply, messages, %{state | module_state: new_module_state}, :hibernate}
+  def handle_demand(demand, %{rate_limiting: %{state: :closed}} = state) do
+    state = update_in(state.rate_limiting.demand_buffer, &:queue.in(demand, &1))
+    {:noreply, [], state}
+  end
 
-      {:stop, reason, new_module_state} ->
-        {:stop, reason, %{state | module_state: new_module_state}}
-    end
+  def handle_demand(demand, state) do
+    %{module: module, module_state: module_state} = state
+    handle_no_reply(module.handle_demand(demand, module_state), state)
   end
 
   @impl true
@@ -135,10 +149,12 @@ defmodule Broadway.Producer do
     |> case do
       {:reply, reply, events, new_module_state} ->
         messages = transform_events(events, state.transformer)
+        {state, messages} = maybe_rate_limit_and_buffer_messages(state, messages)
         {:reply, reply, messages, %{state | module_state: new_module_state}}
 
       {:reply, reply, events, new_module_state, :hibernate} ->
         messages = transform_events(events, state.transformer)
+        {state, messages} = maybe_rate_limit_and_buffer_messages(state, messages)
         {:reply, reply, messages, %{state | module_state: new_module_state}, :hibernate}
 
       {:stop, reason, reply, new_module_state} ->
@@ -171,14 +187,70 @@ defmodule Broadway.Producer do
   end
 
   @impl true
+  def handle_info({__MODULE__, :cancel_consumers}, %{rate_limiting: %{} = rate_limiting} = state) do
+    rate_limiting = %{rate_limiting | draining?: true}
+
+    if :queue.is_empty(rate_limiting.message_buffer) do
+      cancel_consumers(state)
+    end
+
+    {:noreply, [], %{state | rate_limiting: rate_limiting}}
+  end
+
   def handle_info({__MODULE__, :cancel_consumers}, state) do
-    for from <- state.consumers do
-      send(self(), {:"$gen_producer", from, {:cancel, :shutdown}})
+    cancel_consumers(state)
+    {:noreply, [], state}
+  end
+
+  # Don't forward buffered demand when we're draining or when the rate limiting is closed.
+  def handle_info(
+        {__MODULE__, :handle_next_demand},
+        %{rate_limiting: %{draining?: draining?, state: rl_state}} = state
+      )
+      when draining? or rl_state == :closed do
+    {:noreply, [], state}
+  end
+
+  def handle_info({__MODULE__, :handle_next_demand}, state) do
+    case get_and_update_in(state.rate_limiting.demand_buffer, &:queue.out/1) do
+      {{:value, demand}, state} ->
+        case handle_demand(demand, state) do
+          {:noreply, messages, state} ->
+            schedule_next_handle_demand_if_any(state)
+            {:noreply, messages, state}
+
+          {:noreply, messages, state, :hibernate} ->
+            schedule_next_handle_demand_if_any(state)
+            {:noreply, messages, state, :hibernate}
+
+          {:stop, reason, state} ->
+            {:stop, reason, state}
+        end
+
+      {:empty, state} ->
+        {:noreply, [], state}
     end
+  end
 
+  # If the rate limit is lifted but our rate limiting state was "open",
+  # we don't need to do anything since we don't have anything in the buffer.
+  def handle_info({RateLimiter, :reset_rate_limiting}, %{rate_limiting: %{state: :open}} = state) do
     {:noreply, [], state}
   end
 
+  def handle_info({RateLimiter, :reset_rate_limiting}, state) do
+    state = put_in(state.rate_limiting.state, :open)
+
+    {state, messages} = rate_limit_and_buffer_messages(state)
+
+    # We'll schedule to handle the buffered demand regardless of
+    # the state of rate limiting. We'll check if we can forward it
+    # when handling the message.
+    schedule_next_handle_demand_if_any(state)
+
+    {:noreply, messages, state}
+  end
+
   def handle_info(message, state) do
     %{module: module, module_state: module_state} = state
 
@@ -200,10 +272,12 @@ defmodule Broadway.Producer do
     case reply do
       {:noreply, events, new_module_state} when is_list(events) ->
         messages = transform_events(events, transformer)
+        {state, messages} = maybe_rate_limit_and_buffer_messages(state, messages)
         {:noreply, messages, %{state | module_state: new_module_state}}
 
       {:noreply, events, new_module_state, :hibernate} ->
         messages = transform_events(events, transformer)
+        {state, messages} = maybe_rate_limit_and_buffer_messages(state, messages)
         {:noreply, messages, %{state | module_state: new_module_state}, :hibernate}
 
       {:stop, reason, new_module_state} ->
@@ -236,4 +310,128 @@ defmodule Broadway.Producer do
             "struct. In case you're using a standard GenStage producer, please set the " <>
             ":transformer option to transform produced events into message structs"
   end
+
+  defp maybe_rate_limit_and_buffer_messages(state, messages) do
+    if state.rate_limiting && messages != [] do
+      state = update_in(state.rate_limiting.message_buffer, &enqueue_batch(&1, messages))
+      rate_limit_and_buffer_messages(state)
+    else
+      {state, messages}
+    end
+  end
+
+  defp rate_limit_and_buffer_messages(%{rate_limiting: %{state: :closed}} = state) do
+    {state, []}
+  end
+
+  defp rate_limit_and_buffer_messages(%{rate_limiting: rate_limiting} = state) do
+    %{message_buffer: buffer, table_name: table_name, draining?: draining?} = rate_limiting
+
+    {rate_limiting, messages_to_emit} =
+      case RateLimiter.get_currently_allowed(table_name) do
+        # No point in trying to emit messages if no messages are allowed. In that case,
+        # we close the rate limiting and don't emit anything.
+        allowed when allowed <= 0 ->
+          {%{rate_limiting | state: :closed}, []}
+
+        allowed ->
+          {allowed_left, probably_emittable, buffer} = dequeue_many(buffer, allowed, [])
+
+          {rate_limiting_state, messages_to_emit, messages_to_buffer} =
+            rate_limit_messages(
+              table_name,
+              probably_emittable,
+              _probably_emittable_count = allowed - allowed_left
+            )
+
+          new_buffer = enqueue_batch_r(buffer, messages_to_buffer)
+
+          rate_limiting = %{
+            rate_limiting
+            | message_buffer: new_buffer,
+              state: rate_limiting_state
+          }
+
+          if draining? and :queue.is_empty(new_buffer) do
+            cancel_consumers(state)
+          end
+
+          {rate_limiting, messages_to_emit}
+      end
+
+    {%{state | rate_limiting: rate_limiting}, messages_to_emit}
+  end
+
+  defp reverse_split_demand(rest, 0, acc) do
+    {0, acc, rest}
+  end
+
+  defp reverse_split_demand([], demand, acc) do
+    {demand, acc, []}
+  end
+
+  defp reverse_split_demand([head | tail], demand, acc) do
+    reverse_split_demand(tail, demand - 1, [head | acc])
+  end
+
+  defp dequeue_many(queue, demand, acc) do
+    case :queue.out(queue) do
+      {{:value, list}, queue} ->
+        case reverse_split_demand(list, demand, acc) do
+          {0, acc, []} ->
+            {0, Enum.reverse(acc), queue}
+
+          {0, acc, rest} ->
+            {0, Enum.reverse(acc), :queue.in_r(rest, queue)}
+
+          {demand, acc, []} ->
+            dequeue_many(queue, demand, acc)
+        end
+
+      {:empty, queue} ->
+        {demand, Enum.reverse(acc), queue}
+    end
+  end
+
+  defp enqueue_batch(queue, _list = []), do: queue
+  defp enqueue_batch(queue, list), do: :queue.in(list, queue)
+
+  defp enqueue_batch_r(queue, _list = []), do: queue
+  defp enqueue_batch_r(queue, list), do: :queue.in_r(list, queue)
+
+  defp rate_limit_messages(_state, [], _count) do
+    {[], []}
+  end
+
+  defp rate_limit_messages(table_name, messages, message_count) do
+    case RateLimiter.rate_limit(table_name, message_count) do
+      # If no more messages are allowed, we're rate limited but we're able
+      # to emit all messages that we have.
+      0 ->
+        {:closed, messages, _to_buffer = []}
+
+      # We were able to emit all messages and still more messages are allowed,
+      # so the rate limiting is "open".
+      left when left > 0 ->
+        {:open, messages, _to_buffer = []}
+
+      # We went over the rate limit, so we split (on negative index) the messages
+      # we were able to emit and close the rate limiting.
+      overflow when overflow < 0 ->
+        {emittable, to_buffer} = Enum.split(messages, overflow)
+        {:closed, emittable, to_buffer}
+    end
+  end
+
+  defp schedule_next_handle_demand_if_any(state) do
+    if not :queue.is_empty(state.rate_limiting.demand_buffer) do
+      send(self(), {__MODULE__, :handle_next_demand})
+    end
+  end
+
+  defp cancel_consumers(state) do
+    for from <- state.consumers do
+      send(self(), {:"$gen_producer", from, {:cancel, :shutdown}})
+    end
+  end
 end
diff --git a/lib/broadway/rate_limiter.ex b/lib/broadway/rate_limiter.ex
new file mode 100644
index 0000000..8b415ac
--- /dev/null
+++ b/lib/broadway/rate_limiter.ex
@@ -0,0 +1,76 @@
+defmodule Broadway.RateLimiter do
+  # TODO: Use :counters once we require Erlang/OTP 22+
+  @moduledoc false
+
+  use GenServer
+  @row_name :rate_limit_counter
+
+  def start_link(opts) do
+    case Keyword.fetch!(opts, :rate_limiting) do
+      # If we don't have rate limiting options, we don't even need to start this rate
+      # limiter process.
+      nil ->
+        :ignore
+
+      rate_limiting_opts ->
+        name = Keyword.fetch!(opts, :name)
+        producers_names = Keyword.fetch!(opts, :producers_names)
+        args = {name, rate_limiting_opts, producers_names}
+        GenServer.start_link(__MODULE__, args, name: table_name(name))
+    end
+  end
+
+  def rate_limit(table_name, amount)
+      when is_atom(table_name) and is_integer(amount) and amount > 0 do
+    :ets.update_counter(table_name, @row_name, -amount)
+  end
+
+  def get_currently_allowed(table_name) when is_atom(table_name) do
+    :ets.lookup_element(table_name, @row_name, 2)
+  end
+
+  def table_name(broadway_name) when is_atom(broadway_name) do
+    Module.concat(broadway_name, RateLimiter)
+  end
+
+  @impl true
+  def init({broadway_name, rate_limiting_opts, producers_names}) do
+    interval = Keyword.fetch!(rate_limiting_opts, :interval)
+    allowed = Keyword.fetch!(rate_limiting_opts, :allowed_messages)
+
+    table = :ets.new(table_name(broadway_name), [:named_table, :public, :set])
+    :ets.insert(table, {@row_name, allowed})
+
+    _ = schedule_next_reset(interval)
+
+    state = %{
+      interval: interval,
+      allowed: allowed,
+      producers_names: producers_names,
+      table: table
+    }
+
+    {:ok, state}
+  end
+
+  @impl true
+  def handle_info(:reset_limit, state) do
+    %{producers_names: producers_names, interval: interval, allowed: allowed, table: table} =
+      state
+
+    true = :ets.insert(table, {@row_name, allowed})
+
+    for name <- producers_names,
+        pid = Process.whereis(name),
+        is_pid(pid),
+        do: send(pid, {__MODULE__, :reset_rate_limiting})
+
+    _ = schedule_next_reset(interval)
+
+    {:noreply, state}
+  end
+
+  defp schedule_next_reset(interval) do
+    _ref = Process.send_after(self(), :reset_limit, interval)
+  end
+end
diff --git a/lib/broadway/server.ex b/lib/broadway/server.ex
index 489b8c7..a0dc426 100644
--- a/lib/broadway/server.ex
+++ b/lib/broadway/server.ex
@@ -2,7 +2,7 @@ defmodule Broadway.Server do
   @moduledoc false
   @behaviour GenServer
 
-  alias Broadway.{Producer, Processor, Batcher, Consumer, Terminator}
+  alias Broadway.{Producer, Processor, Batcher, Consumer, Terminator, RateLimiter}
 
   def start_link(module, opts) do
     GenServer.start_link(__MODULE__, {module, opts}, opts)
@@ -78,6 +78,7 @@ defmodule Broadway.Server do
 
     children =
       [
+        build_rate_limiter_spec(config, producers_names),
         build_producer_supervisor_spec(config, producers_specs),
         build_processor_supervisor_spec(config, processors_specs)
       ] ++
@@ -113,6 +114,18 @@ defmodule Broadway.Server do
     [name: name] ++ Keyword.take(config, [:spawn_opt, :hibernate_after])
   end
 
+  defp build_rate_limiter_spec(config, producers_names) do
+    %{name: broadway_name, producer_config: producer_config} = config
+
+    opts = [
+      name: broadway_name,
+      rate_limiting: producer_config[:rate_limiting],
+      producers_names: producers_names
+    ]
+
+    {RateLimiter, opts}
+  end
+
   defp build_producers_specs(config, opts) do
     %{
       name: broadway_name,


Test patch:
diff --git a/test/broadway/producer_test.exs b/test/broadway/producer_test.exs
index da06a62..4cc0108 100644
--- a/test/broadway/producer_test.exs
+++ b/test/broadway/producer_test.exs
@@ -66,7 +66,14 @@ defmodule Broadway.ProducerTest do
   end
 
   setup do
-    %{state: %{module: FakeProducer, transformer: nil, module_state: nil}}
+    %{
+      state: %{
+        module: FakeProducer,
+        transformer: nil,
+        module_state: nil,
+        rate_limiting: nil
+      }
+    }
   end
 
   test "init with bad return" do
diff --git a/test/broadway_test.exs b/test/broadway_test.exs
index d9cd129..21ff4a6 100644
--- a/test/broadway_test.exs
+++ b/test/broadway_test.exs
@@ -27,7 +27,11 @@ defmodule BroadwayTest do
     end
 
     @impl true
-    def handle_demand(_demand, state) do
+    def handle_demand(demand, state) do
+      if test_pid = state[:test_pid] do
+        send(test_pid, {:handle_demand_called, demand, System.monotonic_time()})
+      end
+
       {:noreply, [], state}
     end
 
@@ -314,7 +318,7 @@ defmodule BroadwayTest do
       assert Agent.get(Module.concat(broadway, "Agent"), & &1) == %{}
 
       # An that is a child of the main supervisor
-      assert [_, _, _, _, {Agent, _, _, _}] =
+      assert [_, _, _, _, _, {Agent, _, _, _}] =
                Supervisor.which_children(Module.concat(broadway, "Broadway.Supervisor"))
     end
 
@@ -404,6 +408,48 @@ defmodule BroadwayTest do
     assert_receive {:ack, :ref, [_, _], []}
   end
 
+  describe "test_messages/3" do
+    setup do
+      handle_message = fn message, _ ->
+        message
+      end
+
+      context = %{
+        handle_message: handle_message
+      }
+
+      broadway_name = new_unique_name()
+
+      {:ok, broadway} =
+        Broadway.start_link(CustomHandlers,
+          name: broadway_name,
+          context: context,
+          producer: [module: {ManualProducer, []}],
+          processors: [default: [stages: 1, min_demand: 1, max_demand: 2]]
+        )
+
+      %{broadway: broadway}
+    end
+
+    test "metadata field is added to message, if specified in test_messages",
+         %{broadway: broadway} do
+      ref =
+        Broadway.test_messages(broadway, [:message],
+          batch_mode: :flush,
+          metadata: %{field: :value}
+        )
+
+      assert_receive {:ack, ^ref, [%Message{data: :message, metadata: %{field: :value}}], []}
+    end
+
+    test "metadata field in message defaults to %{}, if not specified in test_messages",
+         %{broadway: broadway} do
+      ref = Broadway.test_messages(broadway, [:message], batch_mode: :flush)
+
+      assert_receive {:ack, ^ref, [%Message{data: :message, metadata: %{}}], []}
+    end
+  end
+
   describe "processor" do
     setup do
       test_pid = self()
@@ -652,7 +698,7 @@ defmodule BroadwayTest do
       assert_receive {:batch_handled, :even, messages, %BatchInfo{batcher: :even, size: 5}}
                      when length(messages) == 5
 
-      refute_receive {:batch_handled, _, _}
+      refute_received {:batch_handled, _, _}
     end
 
     test "generate batches with the remaining messages after :batch_timeout is reached",
@@ -661,7 +707,7 @@ defmodule BroadwayTest do
 
       assert_receive {:batch_handled, :odd, messages, _} when length(messages) == 3
       assert_receive {:batch_handled, :even, messages, _} when length(messages) == 2
-      refute_receive {:batch_handled, _, _, _}
+      refute_received {:batch_handled, _, _, _}
     end
   end
 
@@ -933,7 +979,7 @@ defmodule BroadwayTest do
 
       assert_receive {:ack, ^ref, [%Message{data: 1}], []}
       assert_receive :manually_acked
-      refute_receive {:ack, ^ref, _successful, _failed}
+      refute_received {:ack, ^ref, _successful, _failed}
     end
 
     test "acks multiple messages" do
@@ -969,7 +1015,7 @@ defmodule BroadwayTest do
 
       assert_receive {:ack, ^ref, [_, _, _], [_]}
       assert_receive :manually_acked
-      refute_receive {:ack, ^ref, _successful, _failed}
+      refute_received {:ack, ^ref, _successful, _failed}
     end
   end
 
@@ -1333,10 +1379,10 @@ defmodule BroadwayTest do
       GenStage.stop(producer)
 
       assert_receive {:DOWN, ^ref_producer, _, _, _}
-      refute_receive {:DOWN, ^ref_processor, _, _, _}
-      refute_receive {:DOWN, ^ref_batcher, _, _, _}
-      refute_receive {:DOWN, ^ref_consumer, _, _, _}
       assert_receive {:producer_initialized, ^producer}
+      refute_received {:DOWN, ^ref_processor, _, _, _}
+      refute_received {:DOWN, ^ref_batcher, _, _, _}
+      refute_received {:DOWN, ^ref_consumer, _, _, _}
     end
 
     test "processors resubscribe to the restarted producers and keep processing messages",
@@ -1443,12 +1489,11 @@ defmodule BroadwayTest do
 
       assert_receive {:message_handled, %{data: 1}}
       assert_receive {:message_handled, %{data: 2}}
-
-      refute_receive {:message_handled, %{data: :kill_processor}}
-      refute_receive {:message_handled, %{data: 3}}
-
       assert_receive {:message_handled, %{data: 4}}
       assert_receive {:message_handled, %{data: 5}}
+
+      refute_received {:message_handled, %{data: :kill_processor}}
+      refute_received {:message_handled, %{data: 3}}
     end
 
     test "batches are created normally (without the lost messages)", %{broadway: broadway} do
@@ -1559,7 +1604,7 @@ defmodule BroadwayTest do
       values = Enum.map(successful, & &1.data)
       assert values == [6, 7]
 
-      refute_receive {:ack, _, _successful, _failed}
+      refute_received {:ack, _, _successful, _failed}
     end
   end
 
@@ -1748,8 +1793,8 @@ defmodule BroadwayTest do
       Process.exit(broadway, :shutdown)
 
       assert_receive {:ack, _, [%{data: :message_during_cancel}], []}
-      refute_receive {:ack, _, [%{data: :message_after_cancel}], []}
       assert_receive {:EXIT, ^broadway, :shutdown}
+      refute_received {:ack, _, [%{data: :message_after_cancel}], []}
     end
 
     @tag shutdown: 1
@@ -1760,6 +1805,170 @@ defmodule BroadwayTest do
     end
   end
 
+  describe "rate limiting" do
+    test "with an interval and a number of allowed messages in that interval" do
+      broadway_name = new_unique_name()
+      test_pid = self()
+
+      handle_message = fn message, _ ->
+        send(test_pid, {:handle_message_called, message, System.monotonic_time()})
+        message
+      end
+
+      {:ok, _broadway} =
+        Broadway.start_link(CustomHandlers,
+          name: broadway_name,
+          producer: [
+            module: {ManualProducer, []},
+            stages: 2,
+            rate_limiting: [allowed_messages: 2, interval: 50]
+          ],
+          processors: [default: []],
+          context: %{handle_message: handle_message}
+        )
+
+      send(
+        get_producer(broadway_name),
+        {:push_messages,
+         [
+           %Message{data: 1, acknowledger: {CallerAcknowledger, {self(), :ref}, :unused}},
+           %Message{data: 2, acknowledger: {CallerAcknowledger, {self(), :ref}, :unused}},
+           %Message{data: 3, acknowledger: {CallerAcknowledger, {self(), :ref}, :unused}},
+           %Message{data: 4, acknowledger: {CallerAcknowledger, {self(), :ref}, :unused}},
+           %Message{data: 5, acknowledger: {CallerAcknowledger, {self(), :ref}, :unused}}
+         ]}
+      )
+
+      assert_receive {:handle_message_called, %Message{data: 1}, timestamp1}
+      assert_receive {:handle_message_called, %Message{data: 2}, _timestamp2}
+      assert_receive {:handle_message_called, %Message{data: 3}, timestamp3}
+      assert_receive {:handle_message_called, %Message{data: 4}, _timestamp4}
+      assert_receive {:handle_message_called, %Message{data: 5}, timestamp5}
+
+      assert_receive {:ack, :ref, [_, _], []}
+      assert_receive {:ack, :ref, [_, _], []}
+      assert_receive {:ack, :ref, [_], []}
+
+      # To be safe, we're saying that these messages have to be at least 10ms apart
+      # even if the interval is 50ms. This is because the time when the processor
+      # receives the message is not the time when the producer produced it.
+      assert System.convert_time_unit(timestamp3 - timestamp1, :native, :millisecond) >= 10
+      assert System.convert_time_unit(timestamp5 - timestamp3, :native, :millisecond) >= 10
+    end
+
+    test "buffers demand and flushes it when the rate limiting is lifted" do
+      broadway_name = new_unique_name()
+      test_pid = self()
+
+      handle_message = fn message, _ ->
+        send(test_pid, {:handle_message_called, message, System.monotonic_time()})
+        message
+      end
+
+      {:ok, _broadway} =
+        Broadway.start_link(CustomHandlers,
+          name: broadway_name,
+          producer: [
+            module: {ManualProducer, %{test_pid: test_pid}},
+            stages: 1,
+            # We use a long rate limiting interval so that we can trigger the rate limiting
+            # manually from this test.
+            rate_limiting: [allowed_messages: 2, interval: 10000]
+          ],
+          processors: [default: [stages: 1, max_demand: 2]],
+          context: %{handle_message: handle_message}
+        )
+
+      send(
+        get_producer(broadway_name),
+        {:push_messages,
+         [
+           %Message{data: 1, acknowledger: {CallerAcknowledger, {self(), :ref}, :unused}},
+           %Message{data: 2, acknowledger: {CallerAcknowledger, {self(), :ref}, :unused}},
+           %Message{data: 3, acknowledger: {CallerAcknowledger, {self(), :ref}, :unused}}
+         ]}
+      )
+
+      # First, we assert that handle_demand is called and that two events are emitted.
+      assert_receive {:handle_demand_called, _demand, demand_timestamp1}
+      assert_receive {:handle_message_called, %Message{data: 1}, timestamp1}
+      assert_receive {:handle_message_called, %Message{data: 2}, _timestamp2}
+
+      assert demand_timestamp1 < timestamp1
+
+      # Now, since the rate limiting interval is long, we can assert that
+      # handle_demand is not called and that the third message is not emitted yet.
+      refute_received {:handle_demand_called, _demand, _timestamp}
+      refute_received {:handle_message_called, %Message{data: 3}, _timestamp3}
+
+      # We "cheat" and manually tell the rate limiter to reset the limit.
+      send(get_rate_limiter(broadway_name), :reset_limit)
+
+      assert_receive {:handle_demand_called, _demand, demand_timestamp2}
+      assert_receive {:handle_message_called, %Message{data: 3}, timestamp3}
+
+      assert demand_timestamp2 < timestamp3
+      assert demand_timestamp2 > timestamp1
+    end
+
+    test "works when stopping a pipeline and draining producers" do
+      Process.flag(:trap_exit, true)
+
+      broadway_name = new_unique_name()
+      test_pid = self()
+
+      handle_message = fn message, _ ->
+        send(test_pid, {:handle_message_called, message, System.monotonic_time()})
+        message
+      end
+
+      {:ok, broadway} =
+        Broadway.start_link(CustomHandlers,
+          name: broadway_name,
+          producer: [
+            module: {ManualProducer, %{test_pid: test_pid}},
+            stages: 1,
+            # We use a long rate limiting interval so that we can trigger the rate limiting
+            # manually from this test.
+            rate_limiting: [allowed_messages: 2, interval: 10000]
+          ],
+          processors: [default: [stages: 1, max_demand: 2]],
+          context: %{handle_message: handle_message}
+        )
+
+      # First, we use all the rate limiting we have available.
+      send(
+        get_producer(broadway_name),
+        {:push_messages,
+         [
+           %Message{data: 1, acknowledger: {CallerAcknowledger, {self(), :ref}, :unused}},
+           %Message{data: 2, acknowledger: {CallerAcknowledger, {self(), :ref}, :unused}},
+           %Message{data: 3, acknowledger: {CallerAcknowledger, {self(), :ref}, :unused}}
+         ]}
+      )
+
+      assert_receive {:handle_message_called, %Broadway.Message{data: 1}, timestamp1}
+      assert_receive {:handle_message_called, %Broadway.Message{data: 2}, _timestamp2}
+
+      # Then we stop the pipeline, triggering the :message_during_cancel message of
+      # the ManualProducer.
+      Process.exit(broadway, :shutdown)
+
+      refute_received {:handle_message_called, %Broadway.Message{data: :message_during_cancel},
+                       _timestamp}
+
+      send(get_rate_limiter(broadway_name), :reset_limit)
+
+      assert_receive {:handle_message_called, %Broadway.Message{data: 3}, timestamp3}
+
+      assert_receive {:handle_message_called, %Broadway.Message{data: :message_during_cancel},
+                      timestamp_cancel}
+
+      assert timestamp_cancel > timestamp1
+      assert timestamp3 < timestamp_cancel
+    end
+  end
+
   defp new_unique_name() do
     :"Elixir.Broadway#{System.unique_integer([:positive, :monotonic])}"
   end
@@ -1780,6 +1989,10 @@ defmodule BroadwayTest do
     :"#{broadway_name}.Broadway.Consumer_#{key}_#{index}"
   end
 
+  defp get_rate_limiter(broadway_name) do
+    :"#{broadway_name}.RateLimiter"
+  end
+
   defp get_n_producers(broadway_name) do
     Supervisor.count_children(:"#{broadway_name}.Broadway.ProducerSupervisor").workers
   end