Add rate limiter to helm agent fetches

Pretty sure these can get overwhelmed at large scale, and a simple rate limit should improve the load shedding here.

Author: michaeljguarino

charts/console/templates/deployment.yaml

@@ -8,6 +8,10 @@ metadata:
 {{ include "console.labels" . | indent 4 }}
 spec:
   replicas: {{ .Values.replicaCount }}
+  {{- if .Values.strategy }}
+  strategy:
+    {{- toYaml .Values.strategy | nindent 4 }}
+  {{- end }}
   selector:
     matchLabels:
       app.kubernetes.io/name: console

charts/console/templates/migration.yaml

@@ -43,6 +43,8 @@ spec:
             name: console-migration-env
         env:
         {{ include "console.env" . | nindent 8 }}
+        - name: DB_POOL_TARGET
+          value: '1000'
         {{- if  .Values.extraEnv }}
         {{- toYaml .Values.extraEnv | nindent 8 }}
         {{- end }}

charts/console/templates/secrets.yaml

@@ -58,6 +58,12 @@ data:
 {{ if .Values.console.config.tarballQps }}
   CONSOLE_TARBALL_QPS: {{ .Values.console.config.tarballQps | toString | b64enc | quote }}
 {{ end }}
+{{ if .Values.console.config.cacheAgentQps }}
+  CONSOLE_CACHE_AGENT_QPS: {{ .Values.console.config.cacheAgentQps | toString | b64enc | quote }}
+{{ end }}
+{{ if .Values.console.config.cacheAgentQueueLimit }}
+  CONSOLE_CACHE_AGENT_QUEUE_LIMIT: {{ .Values.console.config.cacheAgentQueueLimit | toString | b64enc | quote }}
+{{ end }}
 {{ if .Values.console.config.dbPoolSize }}
   DB_POOL_SIZE: {{ .Values.console.config.dbPoolSize | toString | b64enc | quote }}
 {{ end }}

charts/console/values.yaml

@@ -32,6 +32,10 @@ disableAdditionalVolumes: false
 # replicaCount is the replica count for console deployment
 replicaCount: 2
 
+# strategy configures the main console Deployment strategy.
+# For rolling updates, set type to RollingUpdate and tune rollingUpdate.maxSurge/maxUnavailable.
+strategy: {}
+
 # homeDir is the prefix for the mount path for console-conf volume in console container. The actual config will go in {homeDir}/.plural
 homeDir: /home/console
 
@@ -72,6 +76,12 @@ console:
     # tarballQps configures the global QPS limit for digest and git tarball endpoints.
     tarballQps: 100
 
+    # cacheAgentQps configures the global QPS limit for pod-local cache agents.
+    cacheAgentQps: ~
+
+    # cacheAgentQueueLimit configures the global queue limit for pod-local cache agents.
+    cacheAgentQueueLimit: ~
+
     # dbPoolSize configures the Ecto database connection pool size.
     dbPoolSize: 50
 

config/config.exs

@@ -65,6 +65,9 @@ config :console,
   kas_dns: "https://kas.example.com",
   qps: 1_000,
   tarball_qps: 100,
+  cache_agent_qps: 50,
+  cache_agent_queue_limit: 10,
+  cache_agent_queue_shed: 5,
   nowatchers: false,
   default_project_name: "default",
   prom_plugins: [Console.Prom.Plugin],
@@ -122,7 +125,6 @@ config :console, Console.Guardian,
 
 config :console, Console.Repo,
   types: Console.PostgrexTypes,
-  queue_target: 1000,
   migration_timestamps: [type: :utc_datetime_usec]
 
 config :tzdata, :autoupdate, :disabled

lib/console/deployments/git/agent.ex

@@ -18,7 +18,6 @@ defmodule Console.Deployments.Git.Agent do
 
   @poll :timer.seconds(120)
   @timeout :timer.seconds(10)
-  @limit 50
   @limit_interval :timer.seconds(1)
 
   defmodule State, do: defstruct [:git, :cache, :last_pull, :url]
@@ -269,15 +268,40 @@ defmodule Console.Deployments.Git.Agent do
 
   defp touch(pid, line), do: send(pid, {:touch, line})
 
-  defp rate_limited(key, fun) when is_function(fun, 0) do
-    :erlang.term_to_binary(key)
-    |> Hammer.check_rate(@limit_interval, @limit)
-    |> case do
-      {:allow, _} -> fun.()
-      {:deny, _} -> {:error, :rate_limited}
+  @doc false
+  def rate_limited({_, pid} = key, fun) when is_function(fun, 0) and is_pid(pid) do
+    hammer_key = :erlang.term_to_binary(key)
+    with {:allow, _} <- Hammer.check_rate(hammer_key, @limit_interval, cache_agent_qps()),
+         {:q, :ok} <- {:q, queue_limit(key, pid)} do
+      fun.()
+    else
+      {:deny, _} ->
+        Logger.warning "rate limiting git/helm agent fetch"
+        {:error, :rate_limited}
+      {:q, _} ->
+        Logger.warning "rate limiting git/helm agent fetch due to message queue length"
+        {:error, :rate_limited}
     end
   end
 
+  defp queue_limit(key, pid) when is_pid(pid) do
+    lim = cache_agent_queue_limit()
+    shed = cache_agent_queue_shed()
+    case {:q, Process.info(pid, :message_queue_len)} do
+      {:q, {_, ^lim}} ->
+        case Hammer.check_rate(:erlang.term_to_binary({:shed, key}), @limit_interval, shed) do
+          {:allow, _} -> :ok
+          {:deny, _} -> :error
+        end
+      {:q, {_, len}} when len < lim -> :ok
+      _ -> :error
+    end
+  end
+
+  defp cache_agent_queue_limit(), do: Console.conf(:cache_agent_queue_limit)
+  defp cache_agent_queue_shed(), do: Console.conf(:cache_agent_queue_shed)
+  defp cache_agent_qps(), do: Console.conf(:cache_agent_qps)
+
   defp refresh(%GitRepository{health: :pullable} = git, cache), do: Cache.refresh(%{cache | git: git})
   defp refresh(_, cache), do: cache
 

lib/console/deployments/helm/agent.ex

@@ -1,5 +1,6 @@
 defmodule Console.Deployments.Helm.Agent do
   use GenServer, restart: :temporary
+  import Console.Deployments.Git.Agent, only: [rate_limited: 2]
   alias Console.Repo
   alias Console.Deployments.Git
   alias Console.Deployments.Helm.{AgentCache, Discovery, Supervisor}
@@ -21,7 +22,8 @@ defmodule Console.Deployments.Helm.Agent do
          _ <- touch(pid, line) do
       {:ok, opener(pid, f), d, i}
     else
-      _ -> GenServer.call(pid, {:fetch, chart, vsn}, @timeout)
+      _ ->
+        rate_limited({:helm, pid}, fn -> GenServer.call(pid, {:fetch, chart, vsn}, @timeout) end)
     end
   end
 
@@ -32,7 +34,8 @@ defmodule Console.Deployments.Helm.Agent do
          _ <- touch(pid, line) do
       {:ok, d}
     else
-      _ -> GenServer.call(pid, {:digest, chart, vsn}, @timeout)
+      _ ->
+        rate_limited({:helm, pid}, fn -> GenServer.call(pid, {:digest, chart, vsn}, @timeout) end)
     end
   end
 

lib/console/pipelines/ai/cluster/pipeline.ex

@@ -1,5 +1,5 @@
 defmodule Console.Pipelines.AI.Cluster.Pipeline do
-  use Console.Pipelines.Consumer
+  use Console.Pipelines.Consumer, demand: 10
   import Console.Pipelines.AI.Base
   alias Console.PubSub
 

lib/console/pipelines/ai/service/pipeline.ex

@@ -1,5 +1,5 @@
 defmodule Console.Pipelines.AI.Service.Pipeline do
-  use Console.Pipelines.Consumer
+  use Console.Pipelines.Consumer, demand: 20
   import Console.Pipelines.AI.Base
   alias Console.PubSub
 

lib/console_web/exceptions.ex

@@ -0,0 +1,6 @@
+defimpl Plug.Exception, for: DBConnection.ConnectionError do
+  def status(%{reason: :queue_timeout}), do: 429
+  def status(_), do: 500
+
+  def actions(_), do: []
+end