feat: rate limit support (#67)

Co-authored-by: Darwin D Wu <darwin67@users.noreply.github.com>

lib/inngest/error.ex

@@ -6,10 +6,14 @@ defmodule Inngest.RetryAfterError do
   defexception [:message, seconds: 3]
 end
 
-defmodule Inngest.InvalidDebounceConfigError do
-  defexception message: "a 'period' must be set for debounce"
+defmodule Inngest.DebounceConfigError do
+  defexception [:message]
+end
+
+defmodule Inngest.BatchEventConfigError do
+  defexception [:message]
 end
 
-defmodule Inngest.InvalidBatchEventConfigError do
+defmodule Inngest.RateLimitConfigError do
   defexception [:message]
 end

lib/inngest/function.ex

@@ -167,6 +167,7 @@ defmodule Inngest.Function do
           }
           |> maybe_debounce()
           |> maybe_batch_events()
+          |> maybe_rate_limit()
         ] ++ handler
       end
 
@@ -182,6 +183,11 @@ defmodule Inngest.Function do
         |> Inngest.FnOpts.validate_batch_events(config)
       end
 
+      defp maybe_rate_limit(config) do
+        fn_opts()
+        |> Inngest.FnOpts.validate_rate_limit(config)
+      end
+
       defp fn_opts() do
         case __MODULE__.__info__(:attributes) |> Keyword.get(:func) |> List.first() do
           nil -> %Inngest.FnOpts{}

lib/inngest/function/config.ex

@@ -1,11 +1,14 @@
 defmodule Inngest.FnOpts do
-  @moduledoc false
+  @moduledoc """
+  Function configuration options
+  """
 
   defstruct [
     :id,
     :name,
     :debounce,
     :batch_events,
+    :rate_limit,
     retries: 3
   ]
 
@@ -16,11 +19,12 @@ defmodule Inngest.FnOpts do
           name: binary(),
           retries: number() | nil,
           debounce: debounce() | nil,
-          batch_events: batch_events() | nil
+          batch_events: batch_events() | nil,
+          rate_limit: rate_limit() | nil
         }
 
   @type debounce() :: %{
-          key: nil | binary(),
+          key: binary() | nil,
           period: binary()
         }
 
@@ -29,6 +33,12 @@ defmodule Inngest.FnOpts do
           timeout: binary()
         }
 
+  @type rate_limit() :: %{
+          limit: number(),
+          period: binary(),
+          key: binary() | nil
+        }
+
   @doc """
   Validate the debounce configuration
   """
@@ -42,17 +52,17 @@ defmodule Inngest.FnOpts do
         period = Map.get(debounce, :period)
 
         if is_nil(period) do
-          raise Inngest.InvalidDebounceConfigError
+          raise Inngest.DebounceConfigError, message: "'period' must be set for debounce"
         end
 
         case Util.parse_duration(period) do
           {:error, error} ->
-            raise Inngest.InvalidDebounceConfigError, message: error
+            raise Inngest.DebounceConfigError, message: error
 
           {:ok, seconds} ->
             # credo:disable-for-next-line
             if seconds > 7 * Util.day_in_seconds() do
-              raise Inngest.InvalidDebounceConfigError,
+              raise Inngest.DebounceConfigError,
                 message: "cannot specify period for more than 7 days"
             end
         end
@@ -74,30 +84,59 @@ defmodule Inngest.FnOpts do
         max_size = Map.get(batch, :max_size)
         timeout = Map.get(batch, :timeout)
 
-        if is_nil(max_size) do
-          raise Inngest.InvalidBatchEventConfigError,
-            message: "'max_size' must be set for batch_events"
-        end
-
-        if is_nil(timeout) do
-          raise Inngest.InvalidBatchEventConfigError,
-            message: "'timeout' must be set for batch_events"
+        if is_nil(max_size) || is_nil(timeout) do
+          raise Inngest.BatchEventConfigError,
+            message: "'max_size' and 'timeout' must be set for batch_events"
         end
 
         case Util.parse_duration(timeout) do
           {:error, error} ->
-            raise Inngest.InvalidBatchEventConfigError, message: error
+            raise Inngest.BatchEventConfigError, message: error
 
           {:ok, seconds} ->
             # credo:disable-for-next-line
             if seconds < 1 || seconds > 60 do
-              raise Inngest.InvalidBatchEventConfigError,
+              raise Inngest.BatchEventConfigError,
                 message: "'timeout' duration set to '#{timeout}', needs to be 1s - 60s"
             end
         end
 
-        batch = batch |> Map.put(:maxSize, max_size) |> Map.drop([:max_size])
+        batch = %{maxSize: max_size, timeout: timeout}
         Map.put(config, :batchEvents, batch)
     end
   end
+
+  @doc """
+  Validate the rate limit config
+  """
+  @spec validate_rate_limit(t(), map()) :: map()
+  def validate_rate_limit(fnopts, config) do
+    case fnopts |> Map.get(:rate_limit) do
+      nil ->
+        config
+
+      rate_limit ->
+        limit = Map.get(rate_limit, :limit)
+        period = Map.get(rate_limit, :period)
+
+        if is_nil(limit) || is_nil(period) do
+          raise Inngest.RateLimitConfigError,
+            message: "'limit' and 'period' must be set for rate_limit"
+        end
+
+        case Util.parse_duration(period) do
+          {:error, error} ->
+            raise Inngest.RateLimitConfigError, message: error
+
+          {:ok, seconds} ->
+            # credo:disable-for-next-line
+            if seconds < 1 || seconds > 60 do
+              raise Inngest.RateLimitConfigError,
+                message: "'period' duration set to '#{period}', needs to be 1s - 60s"
+            end
+        end
+
+        Map.put(config, :rateLimit, rate_limit)
+    end
+  end
 end

mix.exs

@@ -49,8 +49,10 @@ defmodule Inngest.MixProject do
           Function: [
             Inngest.Event,
             Inngest.Function,
+            Inngest.FnOpts,
             Inngest.Trigger,
             Inngest.Function.Step,
+            Inngest.Function.Context,
             Inngest.Function.Input
           ],
           Router: [

test/inngest/function/cases/batch_events_test.exs

@@ -1,8 +1,7 @@
 defmodule Inngest.Function.Cases.BatchEventsTest do
   use ExUnit.Case, async: true
 
-  alias Inngest.Test.DevServer
-  import Inngest.Test.Helper
-
   # TODO: Add test after moving batching logic to OSS
+  # alias Inngest.Test.DevServer
+  # import Inngest.Test.Helper
 end

test/inngest/function/cases/rate_limit_test.exs

@@ -0,0 +1,36 @@
+defmodule Inngest.Function.Cases.RateLimitTest do
+  use ExUnit.Case, async: true
+
+  alias Inngest.Test.DevServer
+  import Inngest.Test.Helper
+
+  @default_sleep 10_000
+
+  @tag :integration
+  test "should only run 2 out of 10" do
+    event_ids = Enum.map(1..10, fn _ -> send_test_event("test/plug.ratelimit") end)
+
+    Process.sleep(@default_sleep)
+
+    fn_runs =
+      event_ids
+      |> Enum.map(fn id ->
+        {:ok, %{"data" => data}} = DevServer.run_ids(id)
+
+        if Enum.count(data) == 1 do
+          assert [
+                   %{
+                     "output" => "Rate Limited",
+                     "status" => "Completed",
+                     "run_id" => run_id
+                   }
+                 ] = data
+        else
+          nil
+        end
+      end)
+      |> Enum.filter(&(!is_nil(&1)))
+
+    assert Enum.count(fn_runs) <= 2
+  end
+end

test/inngest/function/config_test.exs

@@ -29,27 +29,29 @@ defmodule Inngest.FnOptsTest do
       assert %{debounce: %{period: "5s"}} = FnOpts.validate_debounce(@fn_opts, @config)
     end
 
-    ## Invalid configs
+    ##  configs
     test "should raise when period is missing" do
       opts = drop_at(@fn_opts, [:debounce, :period])
 
-      assert_raise Inngest.InvalidDebounceConfigError, fn ->
-        FnOpts.validate_debounce(opts, @config)
-      end
+      assert_raise Inngest.DebounceConfigError,
+                   "'period' must be set for debounce",
+                   fn ->
+                     FnOpts.validate_debounce(opts, @config)
+                   end
     end
 
     test "should raise with invalid period" do
       opts = update_at(@fn_opts, [:debounce, :period], "yolo")
 
-      assert_raise Inngest.InvalidDebounceConfigError, fn ->
+      assert_raise Inngest.DebounceConfigError, "invalid duration: 'yolo'", fn ->
         FnOpts.validate_debounce(opts, @config)
       end
     end
 
     test "should raise with period longer than 7 days" do
       opts = update_at(@fn_opts, [:debounce, :period], "8d")
 
-      assert_raise Inngest.InvalidDebounceConfigError, fn ->
+      assert_raise Inngest.DebounceConfigError, fn ->
         FnOpts.validate_debounce(opts, @config)
       end
     end
@@ -77,8 +79,8 @@ defmodule Inngest.FnOptsTest do
     test "should raise if max_size is missing" do
       opts = drop_at(@fn_opts, [:batch_events, :max_size])
 
-      assert_raise Inngest.InvalidBatchEventConfigError,
-                   "'max_size' must be set for batch_events",
+      assert_raise Inngest.BatchEventConfigError,
+                   "'max_size' and 'timeout' must be set for batch_events",
                    fn ->
                      FnOpts.validate_batch_events(opts, @config)
                    end
@@ -87,8 +89,8 @@ defmodule Inngest.FnOptsTest do
     test "should raise if timeout is missing" do
       opts = drop_at(@fn_opts, [:batch_events, :timeout])
 
-      assert_raise Inngest.InvalidBatchEventConfigError,
-                   "'timeout' must be set for batch_events",
+      assert_raise Inngest.BatchEventConfigError,
+                   "'max_size' and 'timeout' must be set for batch_events",
                    fn ->
                      FnOpts.validate_batch_events(opts, @config)
                    end
@@ -97,7 +99,7 @@ defmodule Inngest.FnOptsTest do
     test "should raise if timeout is invalid" do
       opts = update_at(@fn_opts, [:batch_events, :timeout], "hello")
 
-      assert_raise Inngest.InvalidBatchEventConfigError,
+      assert_raise Inngest.BatchEventConfigError,
                    "invalid duration: 'hello'",
                    fn ->
                      FnOpts.validate_batch_events(opts, @config)
@@ -107,11 +109,61 @@ defmodule Inngest.FnOptsTest do
     test "should raise if timeout is out of range" do
       opts = update_at(@fn_opts, [:batch_events, :timeout], "2m")
 
-      assert_raise Inngest.InvalidBatchEventConfigError,
+      assert_raise Inngest.BatchEventConfigError,
                    "'timeout' duration set to '2m', needs to be 1s - 60s",
                    fn ->
                      FnOpts.validate_batch_events(opts, @config)
                    end
     end
   end
+
+  describe "validate_rate_limit/2" do
+    @fn_opts %FnOpts{
+      id: "foobar",
+      name: "Foobar",
+      rate_limit: %{
+        limit: 10,
+        period: "5s"
+      }
+    }
+
+    test "should succeed with valid config" do
+      assert %{
+               rateLimit: %{
+                 limit: 10,
+                 period: "5s"
+               }
+             } = FnOpts.validate_rate_limit(@fn_opts, @config)
+    end
+
+    test "should raise when limit is missing" do
+      opts = drop_at(@fn_opts, [:rate_limit, :limit])
+
+      assert_raise Inngest.RateLimitConfigError,
+                   "'limit' and 'period' must be set for rate_limit",
+                   fn ->
+                     FnOpts.validate_rate_limit(opts, @config)
+                   end
+    end
+
+    test "should raise when period is missing" do
+      opts = drop_at(@fn_opts, [:rate_limit, :period])
+
+      assert_raise Inngest.RateLimitConfigError,
+                   "'limit' and 'period' must be set for rate_limit",
+                   fn ->
+                     FnOpts.validate_rate_limit(opts, @config)
+                   end
+    end
+
+    test "should raise if timeout is out of range" do
+      opts = update_at(@fn_opts, [:rate_limit, :period], "2m")
+
+      assert_raise Inngest.RateLimitConfigError,
+                   "'period' duration set to '2m', needs to be 1s - 60s",
+                   fn ->
+                     FnOpts.validate_rate_limit(opts, @config)
+                   end
+    end
+  end
 end

test/support/cases/rate_limit_fn.ex

@@ -0,0 +1,21 @@
+defmodule Inngest.Test.Case.RateLimitFn do
+  @moduledoc false
+
+  use Inngest.Function
+  alias Inngest.{FnOpts, Trigger}
+
+  @func %FnOpts{
+    id: "ratelimit-fn",
+    name: "RateLimit Function",
+    rate_limit: %{
+      limit: 2,
+      period: "5s"
+    }
+  }
+  @trigger %Trigger{event: "test/plug.ratelimit"}
+
+  @impl true
+  def exec(_ctx, _args) do
+    {:ok, "Rate Limited"}
+  end
+end