git.squeep.com Git - akkoma/blob - lib/pleroma/web/streamer/state.ex

   1 defmodule Pleroma.Web.Streamer.State do
   2   use GenServer
   3   require Logger
   4
   5   alias Pleroma.Web.Streamer.StreamerSocket
   6
   7   def start_link(_) do
   8     GenServer.start_link(__MODULE__, %{sockets: %{}}, name: __MODULE__)
   9   end
  10
  11   def add_socket(topic, socket) do
  12     GenServer.call(__MODULE__, {:add, socket, topic})
  13   end
  14
  15   def remove_socket(topic, socket) do
  16     GenServer.call(__MODULE__, {:remove, socket, topic})
  17   end
  18
  19   def get_sockets do
  20     %{sockets: stream_sockets} = GenServer.call(__MODULE__, :get_state)
  21     stream_sockets
  22   end
  23
  24   def init(init_arg) do
  25     {:ok, init_arg}
  26   end
  27
  28   def handle_call(:get_state, _from, state) do
  29     {:reply, state, state}
  30   end
  31
  32   def handle_call({:add, socket, topic}, _from, %{sockets: sockets} = state) do
  33     internal_topic = internal_topic(topic, socket)
  34     stream_socket = StreamerSocket.from_socket(socket)
  35
  36     sockets_for_topic =
  37       sockets
  38       |> Map.get(internal_topic, [])
  39       |> List.insert_at(0, stream_socket)
  40       |> Enum.uniq()
  41
  42     state = put_in(state, [:sockets, internal_topic], sockets_for_topic)
  43     Logger.debug("Got new conn for #{topic}")
  44     {:reply, state, state}
  45   end
  46
  47   def handle_call({:remove, socket, topic}, _from, %{sockets: sockets} = state) do
  48     internal_topic = internal_topic(topic, socket)
  49     stream_socket = StreamerSocket.from_socket(socket)
  50
  51     sockets_for_topic =
  52       sockets
  53       |> Map.get(internal_topic, [])
  54       |> List.delete(stream_socket)
  55
  56     state = Kernel.put_in(state, [:sockets, internal_topic], sockets_for_topic)
  57     {:reply, state, state}
  58   end
  59
  60   defp internal_topic(topic, socket)
  61        when topic in ~w[user user:notification direct] do
  62     "#{topic}:#{socket.assigns[:user].id}"
  63   end
  64
  65   defp internal_topic(topic, _) do
  66     topic
  67   end
  68 end