feat: websocket max heap size configuration (#1538)

* fix: set max process heap size to 500MB instead of 8GB
* feat: set websocket transport max heap size

WEBSOCKET_MAX_HEAP_SIZE can be used to configure it

Author: edgurgel

README.md

@@ -169,6 +169,7 @@ If you're using the default tenant, the URL is `ws://realtime-dev.localhost:4000
 | CONNECT_PARTITION_SLOTS                         | number  | Number of dynamic supervisor partitions used by the Connect, ReplicationConnect processes                                                                                                                                                                                                                                       |
 | METRICS_CLEANER_SCHEDULE_TIMER_IN_MS            | number  | Time in ms to run the Metric Cleaner task                                                                                                                                                                                                                                                                                       |
 | METRICS_RPC_TIMEOUT_IN_MS                       | number  | Time in ms to wait for RPC call to fetch Metric per node                                                                                                                                                                                                                                                                        |
+| WEBSOCKET_MAX_HEAP_SIZE                         | number  | Max number of bytes to be allocated as heap for the WebSocket transport process. If the limit is reached the process is brutally killed. Defaults to 50MB.                                                                                                                                                                      |
 | REQUEST_ID_BAGGAGE_KEY                          | string  | OTEL Baggage key to be used as request id                                                                                                                                                                                                                                                                                       |
 | OTEL_SDK_DISABLED                               | boolean | Disable OpenTelemetry tracing completely when 'true'                                                                                                                                                                                                                                                                            |
 | OTEL_TRACES_EXPORTER                            | string  | Possible values: `otlp` or `none`. See [https://github.com/open-telemetry/opentelemetry-erlang/tree/v1.4.0/apps#os-environment] for more details on how to configure the traces exporter.                                                                                                                                       |

config/runtime.exs

@@ -68,6 +68,7 @@ janitor_children_timeout = Env.get_integer("JANITOR_CHILDREN_TIMEOUT", :timer.se
 janitor_schedule_timer = Env.get_integer("JANITOR_SCHEDULE_TIMER_IN_MS", :timer.hours(4))
 platform = if System.get_env("AWS_EXECUTION_ENV") == "AWS_ECS_FARGATE", do: :aws, else: :fly
 broadcast_pool_size = Env.get_integer("BROADCAST_POOL_SIZE", 10)
+websocket_max_heap_size = div(Env.get_integer("WEBSOCKET_MAX_HEAP_SIZE", 50_000_000), :erlang.system_info(:wordsize))
 
 no_channel_timeout_in_ms =
   if config_env() == :test,
@@ -107,6 +108,7 @@ config :realtime, Realtime.Repo,
   ssl: ssl_opts
 
 config :realtime,
+  websocket_max_heap_size: websocket_max_heap_size,
   migration_partition_slots: migration_partition_slots,
   connect_partition_slots: connect_partition_slots,
   rebalance_check_interval_in_ms: rebalance_check_interval_in_ms,

lib/realtime_web/channels/user_socket.ex

@@ -1,4 +1,12 @@
 defmodule RealtimeWeb.UserSocket do
+  # This is defined up here before `use Phoenix.Socket` is called so that we can define `Phoenix.Socket.init/1`
+  # It has to be overridden because we need to set the `max_heap_size` flag from the transport process context
+  @impl true
+  def init(state) when is_tuple(state) do
+    Process.flag(:max_heap_size, max_heap_size())
+    Phoenix.Socket.__init__(state)
+  end
+
   use Phoenix.Socket
   use Realtime.Logs
 
@@ -122,4 +130,6 @@ defmodule RealtimeWeb.UserSocket do
       _ -> @default_log_level
     end
   end
+
+  defp max_heap_size(), do: Application.fetch_env!(:realtime, :websocket_max_heap_size)
 end

mix.exs

@@ -4,7 +4,7 @@ defmodule Realtime.MixProject do
   def project do
     [
       app: :realtime,
-      version: "2.50.2",
+      version: "2.51.0",
       elixir: "~> 1.17.3",
       elixirc_paths: elixirc_paths(Mix.env()),
       start_permanent: Mix.env() == :prod,

rel/vm.args.eex

@@ -10,13 +10,13 @@
 ## Tweak GC to run more often
 ##-env ERL_FULLSWEEP_AFTER 10
 
-## Limit process heap for all procs to 1000 MB
-+hmax 1000000000
+## Limit process heap for all procs to 500 MB. The number here is the number of words
++hmax <%= div(500_000_000, :erlang.system_info(:wordsize)) %>
 
 ## Set distribution buffer busy limit (default is 1024)
 +zdbbl 100000
 
 ## Disable Busy Wait
 +sbwt none
 +sbwtdio none
-+sbwtdcpu none
++sbwtdcpu none

test/realtime_web/channels/realtime_channel_test.exs

@@ -28,6 +28,14 @@ defmodule RealtimeWeb.RealtimeChannelTest do
 
   setup :rls_context
 
+  test "max heap size is set", %{tenant: tenant} do
+    jwt = Generators.generate_jwt_token(tenant)
+    {:ok, %Socket{} = socket} = connect(UserSocket, %{}, conn_opts(tenant, jwt))
+
+    assert Process.info(socket.transport_pid, :max_heap_size) ==
+             {:max_heap_size, %{error_logger: true, include_shared_binaries: false, kill: true, size: 6_250_000}}
+  end
+
   describe "broadcast" do
     @describetag policies: [:authenticated_all_topic_read]