electric

defmodule Electric.StackSupervisor do
  @moduledoc """
  Root supervisor that starts a stack of processes to serve shapes.

  Full supervision tree looks roughly like this:

  First, we start 2 registries, `Electric.ProcessRegistry`, and a registry for shape subscriptions. Both are named using the provided `stack_id` variable.

  1. `Electric.Postgres.Inspector.EtsInspector` is started with a pool name as a config option, module that is passed from the base config is __ignored__
  2. `Electric.Connection.Supervisor` takes a LOT of options to configure replication and start the rest of the tree. It starts (3) and then (4) in `rest-for-one` mode
  3. `Electric.Connection.Manager` takes all the connection/replication options and starts the db pools. It goes through the following steps:
      - start_admin_pool
        Starts an admin connection pool (3.1) using replication credentials before the lock is acquired.
        This pool is available for operations like lock breaking, shutdown cleanup, and ETS inspector cache population.
      - start_replication_client
        This starts a replication client (3.2) with no auto-reconnection, because manager is expected to restart this client in case something goes wrong. The streaming of WAL does not start automatically and has to be started explicitly by the manager.
        The replication client acquires the advisory lock, creates/verifies the publication and slot.
      - start_snapshot_pool
        Starts a snapshot connection pool (3.3) using pooled credentials after the replication client is configured.
        If it's ok, we then do a bunch of checks, then ask (3) to finally start (4), and start streaming.

      1. `Postgrex` admin pool - small pool (1-4 connections) for metadata/admin operations
      2. `Electric.Postgres.ReplicationClient` - connects to PG in replication mode, sets up slots, _does not start streaming_ until requested
      3. `Postgrex` snapshot pool - larger pool for querying initial snapshots & info about the DB
  4. `Electric.Shapes.Supervisor` is a supervisor responsible for taking the replication log from the replication client and shoving it into storage appropriately. It starts 3 things in one-for-all mode:
      1. `Electric.Shapes.DynamicConsumerSupervisor` is DynamicSupervisor. It oversees various per-shape processes
          1. `Electric.Shapes.Consumer` is a consumer subscribing to `LogCollector`, which acts a shared producer for all shapes. It passes any incoming operation along to the storage.
          2. `Electric.Shapes.Consumer.Snapshotter` is a temporary GenServer that executes initial snapshot query and writes that to storage
          3. `Electric.Shapes.Consumer.Materializer` monitors a sub-shape in order to invalidate dependent shapes
      2. `Electric.Replication.PublicationManager` manages all filters on the publication for the replication
      3. `Electric.Replication.ShapeLogCollector` collects transactions from the replication connection, fanning them out to `Electric.Shapes.Consumer` (4.1.1.2)
      4. `Electric.ShapeCache` coordinates shape creation and handle allocation, shape metadata
  """

  opts =
    if Application.compile_env(:electric, :start_in_library_mode, true) do
      [
        # Setting `restart: :transient` is required for passing the `:auto_shutdown` to `Supervisor.init()` below.
        restart: :transient
      ]
    else
      [
        restart: :transient,
        # Make StackSupervisor `significant` so that in the case that electric is in single-stack mode, the stack stopping
        # will stop the entire Electric application (since `auto_shutdown` is set to `:any_significant` in `Application`).
        significant: true
      ]
    end

  use Supervisor, opts

  alias Electric.ShapeCache.LogChunker

  require Logger

  @opts_schema NimbleOptions.new!(
                 name: [type: :any, required: false],
                 stack_id: [type: :string, required: true],
                 persistent_kv: [type: :any, required: true],
                 stack_events_registry: [type: :atom, required: true],
                 connection_opts: [
                   type: :keyword_list,
                   required: true,
                   keys: Electric.connection_opts_schema()
                 ],
                 max_shapes: [type: {:or, [:non_neg_integer, nil]}, default: nil],
                 max_concurrent_requests: [
                   type: :map,
                   keys: [
                     initial: [type: :integer, required: true],
                     existing: [type: :integer, required: true]
                   ],
                   default: Electric.Config.default(:max_concurrent_requests)
                 ],
                 replication_opts: [
                   type: :keyword_list,
                   required: true,
                   keys: [
                     connection_opts: [
                       type: :keyword_list,
                       required: true,
                       keys: Electric.connection_opts_schema()
                     ],
                     publication_name: [type: :string, required: true],
                     slot_name: [type: :string, required: true],
                     slot_temporary?: [type: :boolean, default: false],
                     try_creating_publication?: [type: :boolean, default: true],
                     max_txn_size: [type: {:or, [:non_neg_integer, nil]}, default: nil],
                     max_batch_size: [
                       type: :non_neg_integer,
                       default: Electric.Config.default(:max_batch_size)
                     ],
                     replication_idle_timeout: [
                       type: :non_neg_integer,
                       default: Electric.Config.default(:replication_idle_timeout)
                     ]
                   ]
                 ],
                 pool_opts: [
                   type: :keyword_list,
                   required: false,
                   doc:
                     "will be passed on to the Postgrex connection pool. See `t:Postgrex.start_option()`, apart from the connection options."
                 ],
                 storage: [type: :mod_arg, required: true],
                 storage_dir: [type: :string, required: true],
                 chunk_bytes_threshold: [
                   type: :pos_integer,
                   default: LogChunker.default_chunk_size_threshold()
                 ],
                 feature_flags: [type: {:list, :string}, default: []],
                 tweaks: [
                   type: :keyword_list,
                   required: false,
                   doc:
                     "tweaks to the behaviour of parts of the supervision tree, used mostly for tests",
                   default: [],
                   keys: [
                     publication_alter_debounce_ms: [type: :non_neg_integer, default: 0],
                     cleanup_interval_ms: [type: :non_neg_integer, default: 10_000],
                     registry_partitions: [type: :non_neg_integer, required: false],
                     shape_cleaner_opts: [
                       type: :keyword_list,
                       required: false,
                       keys: [
                         on_cleanup: [type: {:fun, 1}]
                       ]
                     ],
                     publication_refresh_period: [type: :non_neg_integer, default: 60_000],
                     schema_reconciler_period: [type: :non_neg_integer, default: 60_000],
                     shape_hibernate_after: [
                       type: :integer,
                       default: Electric.Config.default(:shape_hibernate_after)
                     ],
                     shape_enable_suspend?: [
                       type: :boolean,
                       default: Electric.Config.default(:shape_enable_suspend?)
                     ],
                     shape_suspend_after: [
                       type: :non_neg_integer,
                       default: Electric.Config.default(:shape_suspend_after)
                     ],
                     snapshot_timeout_to_first_data: [
                       type: :pos_integer,
                       default: Electric.Config.default(:snapshot_timeout_to_first_data)
                     ],
                     conn_max_requests: [
                       # allow 0 as a setting which means max_requests=infinity
                       type: :non_neg_integer,
                       default: Electric.Config.default(:conn_max_requests)
                     ],
                     handler_fullsweep_after: [
                       type: {:or, [:non_neg_integer, nil]},
                       default: nil
                     ],
                     process_spawn_opts: [type: :map, default: %{}],
                     consumer_gc_heap_threshold: [
                       type: {:or, [:non_neg_integer, nil]},
                       default: Electric.Config.default(:consumer_gc_heap_threshold)
                     ],
                     consumer_partitions: [type: {:or, [:pos_integer, nil]}, default: nil]
                   ]
                 ],
                 lock_breaker_guard: [
                   type: {:or, [{:fun, 0}, nil]},
                   default: nil,
                   doc:
                     "Optional guard callback for the lock breaker. When set, the lock breaker will only run if this callback returns true."
                 ],
                 manual_table_publishing?: [
                   type: :boolean,
                   required: false,
                   doc:
                     "Specify whether tables are to be added to the Postgres publication automatically or by hand",
                   default: false
                 ],
                 shape_db_opts: [
                   type: :keyword_list,
                   required: true,
                   doc: "Configuration of the shape db sub-system",
                   keys: [
                     storage_dir: [type: :string, required: true],
                     exclusive_mode: [type: :boolean],
                     synchronous: [type: :string],
                     cache_size: [type: :integer],
                     connection_idle_timeout: [type: :integer],
                     read_pool_size: [type: :integer]
                   ]
                 ],
                 telemetry_opts: [type: :keyword_list, default: []],
                 telemetry_span_attrs: [
                   # Validates the OpenTelemetry.attributes_map() type
                   # cf. https://github.com/open-telemetry/opentelemetry-erlang/blob/9f7affe630676d2803b04f69d0c759effb6e0245/apps/opentelemetry_api/src/opentelemetry.erl#L118
                   type:
                     {:or,
                      [
                        {:map, {:or, [:atom, :string]},
                         {:or,
                          [
                            :atom,
                            :string,
                            :integer,
                            :float,
                            :boolean,
                            {:list, {:or, [:atom, :string, :integer, :float, :boolean]}},
                            :map
                          ]}},
                        {:list,
                         {:tuple,
                          [
                            {:or, [:atom, :string]},
                            {:or,
                             [
                               :atom,
                               :string,
                               :integer,
                               :float,
                               :boolean,
                               {:list, {:or, [:atom, :string, :integer, :float, :boolean]}},
                               :map
                             ]}
                          ]}}
                      ]},
                   required: false
                 ]
               )

  def opts_schema do
    @opts_schema
  end

  def start_link(opts) do
    opts = obfuscate_password(opts)

    with {:ok, config} <- NimbleOptions.validate(Map.new(opts), @opts_schema) do
      Supervisor.start_link(__MODULE__, config, Keyword.take(opts, [:name]))
    end
  end

  defp obfuscate_password(opts) when is_list(opts) do
    opts
    |> Keyword.update(:connection_opts, [], &Electric.Utils.obfuscate_password/1)
    |> Keyword.update(:replication_opts, [], fn repl_opts ->
      Keyword.update(repl_opts, :connection_opts, [], &Electric.Utils.obfuscate_password/1)
    end)
  end

  def subscribe_to_stack_events(
        registry \\ Electric.stack_events_registry(),
        stack_id,
        ref \\ make_ref()
      )

  def subscribe_to_stack_events(registry, stack_id, ref) do
    {:ok, _pid} = Registry.register(registry, {:stack_status, stack_id}, ref)
    ref
  end

  def dispatch_stack_event(registry \\ Electric.stack_events_registry(), stack_id, event)

  # noop if there's no registry running
  def dispatch_stack_event(nil, _stack_id, _event) do
    :ok
  end

  def dispatch_stack_event(registry, stack_id, event) do
    Registry.dispatch(registry, {:stack_status, stack_id}, fn entries ->
      for {pid, ref} <- entries do
        send(pid, {:stack_status, ref, event})
      end
    end)
  end

  def subscribe_to_shape_events(stack_id, handle, ref \\ make_ref()) do
    registry = Electric.StackSupervisor.registry_name(stack_id)
    Registry.register(registry, handle, ref)
    ref
  end

  def build_shared_opts(opts) do
    # needs validation
    opts = Map.new(opts)
    stack_id = opts[:stack_id]

    shape_changes_registry_name = registry_name(stack_id)

    shape_cache =
      Access.get(
        opts,
        :shape_cache,
        {Electric.ShapeCache, stack_id: stack_id, server: Electric.ShapeCache.name(stack_id)}
      )

    publication_manager =
      Access.get(
        opts,
        :publication_manager,
        {Electric.Replication.PublicationManager, stack_id: stack_id}
      )

    persistent_kv = Access.fetch!(opts, :persistent_kv)

    [
      shape_cache: shape_cache,
      publication_manager: publication_manager,
      registry: shape_changes_registry_name,
      stack_events_registry: opts[:stack_events_registry],
      storage: shared_storage_opts(opts),
      inspector: shared_inspector_opts(opts),
      stack_id: stack_id,
      persistent_kv: persistent_kv,
      feature_flags: Map.get(opts, :feature_flags, []),
      max_concurrent_requests:
        Map.get(opts, :max_concurrent_requests, Electric.Config.default(:max_concurrent_requests))
    ]
  end

  def registry_name(stack_id) do
    :"Electric.Registry.ShapeChanges:#{stack_id}"
  end

  defp shared_storage_opts(config) do
    {mod, storage_opts} = config.storage

    storage_opts =
      storage_opts
      |> Keyword.put(:stack_id, config.stack_id)
      |> Keyword.put(:chunk_bytes_threshold, config[:chunk_bytes_threshold])

    Electric.ShapeCache.Storage.shared_opts({mod, storage_opts})
  end

  defp shared_inspector_opts(config) do
    Map.get_lazy(
      config,
      :inspector,
      fn ->
        {Electric.Postgres.Inspector.EtsInspector,
         stack_id: config.stack_id,
         server: Electric.Postgres.Inspector.EtsInspector.name(config.stack_id)}
      end
    )
  end

  @impl true
  def init(%{stack_id: stack_id} = config) do
    Logger.debug("StackSupervisor for stack #{inspect(stack_id)} is initializing...")

    Process.set_label({:stack_supervisor, stack_id})
    Logger.metadata(stack_id: stack_id)
    Electric.Telemetry.Sentry.set_tags_context(stack_id: stack_id)

    storage = shared_storage_opts(config)
    inspector = shared_inspector_opts(config)

    metadata_db_pool = Electric.Connection.Manager.admin_pool(stack_id)

    shape_changes_registry_name = registry_name(stack_id)

    shape_hibernate_after = Keyword.fetch!(config.tweaks, :shape_hibernate_after)
    shape_enable_suspend? = Keyword.fetch!(config.tweaks, :shape_enable_suspend?)
    shape_suspend_after = Keyword.fetch!(config.tweaks, :shape_suspend_after)
    process_spawn_opts = Keyword.fetch!(config.tweaks, :process_spawn_opts)
    consumer_gc_heap_threshold = Keyword.fetch!(config.tweaks, :consumer_gc_heap_threshold)

    shape_cache_opts = [
      stack_id: stack_id
    ]

    connection_manager_opts = [
      stack_id: stack_id,
      # Coming from the outside, need validation
      connection_opts: config.connection_opts,
      stack_events_registry: config.stack_events_registry,
      replication_opts:
        [
          stack_id: stack_id,
          handle_event: {Electric.Replication.ShapeLogCollector, :handle_event_async, [stack_id]}
        ] ++ config.replication_opts,
      pool_opts: [types: PgInterop.Postgrex.Types] ++ config.pool_opts,
      timeline_opts: [
        stack_id: stack_id,
        persistent_kv: config.persistent_kv
      ],
      persistent_kv: config.persistent_kv,
      shape_cache_opts: shape_cache_opts,
      inspector: inspector,
      max_shapes: config.max_shapes,
      tweaks: config.tweaks,
      manual_table_publishing?: config.manual_table_publishing?,
      lock_breaker_guard: config.lock_breaker_guard
    ]

    registry_partitions =
      Keyword.get(config.tweaks, :registry_partitions, System.schedulers_online())

    telemetry_spec = Electric.StackSupervisor.Telemetry.configure(config)

    children =
      [
        # put telemetry processes first so that they are torn down last and can
        # continue reporting while the stack terminates.
        telemetry_spec,
        {Electric.ProcessRegistry, partitions: registry_partitions, stack_id: stack_id},
        {Electric.StackConfig,
         stack_id: stack_id,
         seed_config: [
           chunk_bytes_threshold: config.chunk_bytes_threshold,
           snapshot_timeout_to_first_data: config.tweaks[:snapshot_timeout_to_first_data],
           inspector: inspector,
           shape_hibernate_after: shape_hibernate_after,
           shape_enable_suspend?: shape_enable_suspend?,
           shape_suspend_after: shape_suspend_after,
           process_spawn_opts: process_spawn_opts,
           consumer_gc_heap_threshold: consumer_gc_heap_threshold,
           feature_flags: Map.get(config, :feature_flags, [])
         ]},
        {Electric.AsyncDeleter,
         stack_id: stack_id,
         storage_dir: config.storage_dir,
         cleanup_interval_ms: config.tweaks[:cleanup_interval_ms]},
        {Registry,
         name: shape_changes_registry_name, keys: :duplicate, partitions: registry_partitions},
        Electric.ShapeCache.Storage.stack_child_spec(storage),
        {Electric.Postgres.Inspector.EtsInspector,
         stack_id: stack_id, pool: metadata_db_pool, persistent_kv: config.persistent_kv},
        {Electric.MonitoredCoreSupervisor,
         stack_id: stack_id,
         connection_manager_opts: connection_manager_opts,
         shape_db_opts: config.shape_db_opts}
      ]
      |> Enum.reject(&is_nil/1)

    Supervisor.init(children, strategy: :one_for_one, auto_shutdown: :any_significant)
  end
end