fix: Selectively run migrations (#1353)

* Prevents migrations from running if the tenant no longer requires.
* Setup Partition size for Connect module
* Reduce test flakiness and usage of `dev_tenant` in integration tests

Author: filipecabaco

config/runtime.exs

@@ -15,6 +15,9 @@ slot_name_suffix = System.get_env("SLOT_NAME_SUFFIX")
 migration_partition_slots =
   System.get_env("MIGRATION_PARTITION_SLOTS", "#{System.schedulers_online() * 2}") |> String.to_integer()
 
+connect_partition_slots =
+  System.get_env("CONNECT_PARTITION_SLOTS", "#{System.schedulers_online() * 2}") |> String.to_integer()
+
 if !(db_version in [nil, "ipv6", "ipv4"]),
   do: raise("Invalid IP version, please set either ipv6 or ipv4")
 
@@ -35,6 +38,7 @@ socket_options =
 
 config :realtime,
   migration_partition_slots: migration_partition_slots,
+  connect_partition_slots: connect_partition_slots,
   tenant_max_bytes_per_second: System.get_env("TENANT_MAX_BYTES_PER_SECOND", "100000") |> String.to_integer(),
   tenant_max_channels_per_client: System.get_env("TENANT_MAX_CHANNELS_PER_CLIENT", "100") |> String.to_integer(),
   tenant_max_concurrent_users: System.get_env("TENANT_MAX_CONCURRENT_USERS", "200") |> String.to_integer(),

config/test.exs

@@ -35,13 +35,10 @@ config :realtime,
   secure_channels: true,
   db_enc_key: "1234567890123456",
   jwt_claim_validators: System.get_env("JWT_CLAIM_VALIDATORS", "{}"),
-  api_jwt_secret: System.get_env("API_JWT_SECRET"),
+  api_jwt_secret: System.get_env("API_JWT_SECRET", "secret"),
   metrics_jwt_secret: "test",
   prom_poll_rate: 5_000
 
-config :joken,
-  current_time_adapter: RealtimeWeb.Joken.CurrentTime.Mock
-
 # Print only errors during test
 config :logger,
   compile_time_purge_matching: [[module: Postgrex], [module: DBConnection]],

lib/realtime/api/tenant.ex

@@ -26,6 +26,7 @@ defmodule Realtime.Api.Tenant do
     field(:events_per_second_rolling, :float, virtual: true)
     field(:events_per_second_now, :integer, virtual: true)
     field(:private_only, :boolean, default: false)
+    field(:migrations_ran, :integer, default: 0)
 
     has_many(:extensions, Realtime.Api.Extensions,
       foreign_key: :tenant_external_id,
@@ -73,7 +74,8 @@ defmodule Realtime.Api.Tenant do
       :max_channels_per_client,
       :max_joins_per_second,
       :suspend,
-      :private_only
+      :private_only,
+      :migrations_ran
     ])
     |> validate_required([
       :external_id,

lib/realtime/application.ex

@@ -51,6 +51,7 @@ defmodule Realtime.Application do
     region = Application.get_env(:realtime, :region)
     :syn.join(RegionNodes, region, self(), node: node())
     migration_partition_slots = Application.get_env(:realtime, :migration_partition_slots)
+    connect_partition_slots = Application.get_env(:realtime, :connect_partition_slots)
 
     children =
       [
@@ -75,7 +76,10 @@ defmodule Realtime.Application do
          name: Realtime.Tenants.Migrations.DynamicSupervisor,
          partitions: migration_partition_slots},
         {PartitionSupervisor,
-         child_spec: DynamicSupervisor, strategy: :one_for_one, name: Realtime.Tenants.Connect.DynamicSupervisor},
+         child_spec: DynamicSupervisor,
+         strategy: :one_for_one,
+         name: Realtime.Tenants.Connect.DynamicSupervisor,
+         partitions: connect_partition_slots},
         {PartitionSupervisor,
          child_spec: DynamicSupervisor,
          strategy: :one_for_one,

lib/realtime/tenants.ex

@@ -295,6 +295,27 @@ defmodule Realtime.Tenants do
     |> tap(fn _ -> Cache.distributed_invalidate_tenant_cache(external_id) end)
   end
 
+  @doc """
+  Checks if migrations for a given tenant need to run.
+  """
+  @spec run_migrations?(binary()) :: boolean()
+  def run_migrations?(external_id) do
+    tenant = Cache.get_tenant_by_external_id(external_id)
+    tenant.migrations_ran < Enum.count(Migrations.migrations())
+  end
+
+  @doc """
+  Updates the migrations_ran field for a tenant.
+  """
+  @spec update_migrations_ran(binary(), integer()) :: {:ok, Tenant.t()} | {:error, term()}
+  def update_migrations_ran(external_id, count) do
+    external_id
+    |> Cache.get_tenant_by_external_id()
+    |> Tenant.changeset(%{migrations_ran: count})
+    |> Repo.update!()
+    |> tap(fn _ -> Cache.distributed_invalidate_tenant_cache(external_id) end)
+  end
+
   defp broadcast_operation_event(action, external_id),
     do: Phoenix.PubSub.broadcast!(Realtime.PubSub, "realtime:operations:" <> external_id, action)
 end

lib/realtime/tenants/connect.ex

@@ -125,13 +125,13 @@ defmodule Realtime.Tenants.Connect do
   end
 
   @doc """
-  Returns the pid of the tenant Connection process
+  Returns the pid of the tenant Connection process and db_conn pid
   """
-  @spec whereis(binary()) :: pid | nil
+  @spec whereis(binary()) :: pid() | nil
   def whereis(tenant_id) do
     case :syn.lookup(__MODULE__, tenant_id) do
-      {pid, _} -> pid
-      :undefined -> nil
+      {pid, _} when is_pid(pid) -> pid
+      _ -> nil
     end
   end
 
@@ -141,8 +141,12 @@ defmodule Realtime.Tenants.Connect do
   @spec shutdown(binary()) :: :ok | nil
   def shutdown(tenant_id) do
     case whereis(tenant_id) do
-      pid when is_pid(pid) -> GenServer.stop(pid)
-      _ -> :ok
+      pid when is_pid(pid) ->
+        send(pid, :shutdown_connect)
+        :ok
+
+      _ ->
+        :ok
     end
   end
 
@@ -196,7 +200,7 @@ defmodule Realtime.Tenants.Connect do
   def handle_continue(:run_migrations, state) do
     %{tenant: tenant, db_conn_pid: db_conn_pid} = state
 
-    with :ok <- Migrations.run_migrations(tenant),
+    with res when res in [:ok, :noop] <- Migrations.run_migrations(tenant),
          :ok <- Migrations.create_partitions(db_conn_pid) do
       {:noreply, state, {:continue, :start_listen_and_replication}}
     else
@@ -273,42 +277,19 @@ defmodule Realtime.Tenants.Connect do
     {:noreply, %{state | connected_users_bucket: connected_users_bucket}}
   end
 
-  def handle_info(:shutdown, state) do
-    %{
-      db_conn_pid: db_conn_pid,
-      replication_connection_pid: replication_connection_pid,
-      listen_pid: listen_pid
-    } = state
-
+  def handle_info(:shutdown_no_connected_users, state) do
     Logger.info("Tenant has no connected users, database connection will be terminated")
-    :ok = GenServer.stop(db_conn_pid, :normal, 500)
-
-    replication_connection_pid && Process.alive?(replication_connection_pid) &&
-      GenServer.stop(replication_connection_pid, :normal, 500)
-
-    listen_pid && Process.alive?(listen_pid) &&
-      GenServer.stop(listen_pid, :normal, 500)
-
-    {:stop, :normal, state}
+    shutdown_connect_process(state)
   end
 
   def handle_info(:suspend_tenant, state) do
-    %{
-      db_conn_pid: db_conn_pid,
-      replication_connection_pid: replication_connection_pid,
-      listen_pid: listen_pid
-    } = state
-
     Logger.warning("Tenant was suspended, database connection will be terminated")
-    :ok = GenServer.stop(db_conn_pid, :normal, 500)
-
-    replication_connection_pid && Process.alive?(replication_connection_pid) &&
-      GenServer.stop(replication_connection_pid, :normal, 500)
-
-    listen_pid && Process.alive?(listen_pid) &&
-      GenServer.stop(listen_pid, :normal, 500)
+    shutdown_connect_process(state)
+  end
 
-    {:stop, :normal, state}
+  def handle_info(:shutdown_connect, state) do
+    Logger.warning("Shutdowning tenant connection")
+    shutdown_connect_process(state)
   end
 
   # Handle database connection termination
@@ -371,7 +352,7 @@ defmodule Realtime.Tenants.Connect do
          @connected_users_bucket_shutdown,
          check_connected_user_interval
        ) do
-    Process.send_after(self(), :shutdown, check_connected_user_interval)
+    Process.send_after(self(), :shutdown_no_connected_users, check_connected_user_interval)
   end
 
   defp send_connected_user_check_message(connected_users_bucket, check_connected_user_interval) do
@@ -381,4 +362,22 @@ defmodule Realtime.Tenants.Connect do
 
   defp tenant_suspended?(%Tenant{suspend: true}), do: {:error, :tenant_suspended}
   defp tenant_suspended?(_), do: :ok
+
+  defp shutdown_connect_process(state) do
+    %{
+      db_conn_pid: db_conn_pid,
+      replication_connection_pid: replication_connection_pid,
+      listen_pid: listen_pid
+    } = state
+
+    :ok = GenServer.stop(db_conn_pid, :shutdown, 500)
+
+    replication_connection_pid && Process.alive?(replication_connection_pid) &&
+      GenServer.stop(replication_connection_pid, :normal, 500)
+
+    listen_pid && Process.alive?(listen_pid) &&
+      GenServer.stop(listen_pid, :normal, 500)
+
+    {:stop, :normal, state}
+  end
 end

lib/realtime/tenants/connect/check_connection.ex

@@ -10,8 +10,13 @@ defmodule Realtime.Tenants.Connect.CheckConnection do
     %{tenant: tenant} = acc
 
     case Database.check_tenant_connection(tenant) do
-      {:ok, conn} -> {:ok, %{acc | db_conn_pid: conn, db_conn_reference: Process.monitor(conn)}}
-      {:error, error} -> {:error, error}
+      {:ok, conn} ->
+        Process.link(conn)
+        db_conn_reference = Process.monitor(conn)
+        {:ok, %{acc | db_conn_pid: conn, db_conn_reference: db_conn_reference}}
+
+      {:error, error} ->
+        {:error, error}
     end
   end
 end

lib/realtime/tenants/migrations.ex

@@ -8,6 +8,7 @@ defmodule Realtime.Tenants.Migrations do
 
   import Realtime.Logs
 
+  alias Realtime.Tenants
   alias Realtime.Database
   alias Realtime.Registry.Unique
   alias Realtime.Repo
@@ -149,7 +150,7 @@ defmodule Realtime.Tenants.Migrations do
   @doc """
   Run migrations for the given tenant.
   """
-  @spec run_migrations(Tenant.t()) :: :ok | {:error, any()}
+  @spec run_migrations(Tenant.t()) :: :ok | :noop | {:error, any()}
   def run_migrations(%Tenant{} = tenant) do
     %{extensions: [%{settings: settings} | _]} = tenant
     attrs = %__MODULE__{tenant_external_id: tenant.external_id, settings: settings}
@@ -159,9 +160,13 @@ defmodule Realtime.Tenants.Migrations do
 
     spec = {__MODULE__, attrs}
 
-    case DynamicSupervisor.start_child(supervisor, spec) do
-      :ignore -> :ok
-      error -> error
+    if Tenants.run_migrations?(tenant.external_id) do
+      case DynamicSupervisor.start_child(supervisor, spec) do
+        :ignore -> :ok
+        error -> error
+      end
+    else
+      :noop
     end
   end
 
@@ -174,8 +179,12 @@ defmodule Realtime.Tenants.Migrations do
     Logger.metadata(external_id: tenant_external_id, project: tenant_external_id)
 
     case migrate(settings) do
-      {:ok, _} -> :ignore
-      {:error, error} -> {:stop, error}
+      :ok ->
+        Tenants.update_migrations_ran(tenant_external_id, Enum.count(@migrations))
+        :ignore
+
+      {:error, error} ->
+        {:stop, error}
     end
   end
 
@@ -199,9 +208,9 @@ defmodule Realtime.Tenants.Migrations do
 
       try do
         opts = [all: true, prefix: "realtime", dynamic_repo: repo]
-        res = Ecto.Migrator.run(Repo, @migrations, :up, opts)
+        Ecto.Migrator.run(Repo, @migrations, :up, opts)
 
-        {:ok, res}
+        :ok
       rescue
         error ->
           log_error("MigrationsFailedToRun", error)
@@ -244,4 +253,6 @@ defmodule Realtime.Tenants.Migrations do
 
     :ok
   end
+
+  def migrations(), do: @migrations
 end

lib/realtime_web/channels/user_socket.ex

@@ -30,7 +30,7 @@ defmodule RealtimeWeb.UserSocket do
       %{uri: %{host: host}, x_headers: headers} = opts
 
       {:ok, external_id} = Database.get_external_id(host)
-
+      external_id = String.upcase(external_id)
       Logger.metadata(external_id: external_id, project: external_id)
       Logger.put_process_level(self(), :error)
 

lib/realtime_web/controllers/tenant_controller.ex

@@ -147,9 +147,8 @@ defmodule RealtimeWeb.TenantController do
             _e, acc -> acc
           end)
 
-        with {:ok, %Tenant{} = tenant} <-
-               Api.create_tenant(%{tenant_params | "extensions" => extensions}),
-             :ok <- Migrations.run_migrations(tenant) do
+        with {:ok, %Tenant{} = tenant} <- Api.create_tenant(%{tenant_params | "extensions" => extensions}),
+             res when res in [:ok, :noop] <- Migrations.run_migrations(tenant) do
           Logger.metadata(external_id: tenant.external_id, project: tenant.external_id)
 
           conn