Merge pull request #44 from Logflare/feat/request-throttling

feat: request throttling

lib/counter.ex

@@ -20,6 +20,9 @@ defmodule Loadfest.Counter do
   def add(batch_size) do
     GenServer.cast(__MODULE__, {:add, batch_size})
   end
+  def requests() do
+    GenServer.call(__MODULE__, :requests)
+  end
 
   def handle_cast({:add, size}, state) do
     :counters.add(state.ref, idx(:requests), 1)
@@ -29,6 +32,11 @@ defmodule Loadfest.Counter do
     {:noreply, state}
   end
 
+  def handle_call(:requests, _caller, state) do
+    value = :counters.get(state.ref, idx(:requests))
+  {:reply, value, state}
+  end
+
   def handle_info(:log, state) do
     stats = %{
       requests: :counters.get(state.ref, idx(:requests)),

lib/pipeline.ex

@@ -4,6 +4,7 @@ defmodule Loadfest.Pipeline do
   alias Broadway.Message
   require Logger
   @source_names Application.get_env(:loadfest, :source_names)
+  @max_rps Application.get_env(:loadfest, :max_rps, 100_000)
   defmodule Producer do
     use GenStage
 
@@ -32,7 +33,7 @@ defmodule Loadfest.Pipeline do
         module: {Loadfest.Pipeline.Producer, []}
       ],
       processors: [
-        default: [concurrency: 50, max_demand: 1]
+        default: [concurrency: min(@max_rps, 50), max_demand: 1]
       ]
     )
   end
@@ -45,14 +46,17 @@ defmodule Loadfest.Pipeline do
       batch: message.data
     }
 
-    request = Loadfest.Client.send(name, body)
+    rps = Loadfest.Counter.requests()
+    if rps < @max_rps do
+      request = Loadfest.Client.send(name, body)
 
-    if request.status == 200 do
-      Loadfest.Counter.add(length(message.data))
-    end
+      if request.status == 200 do
+        Loadfest.Counter.add(length(message.data))
+      end
 
-    if request.status >= 400 do
-      Logger.warning("#{request.status} | #{inspect(request.body)}")
+      if request.status >= 400 do
+        Logger.warning("#{request.status} | #{inspect(request.body)}")
+      end
     end
 
     message