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   @env Mix.env()
   8
   9   def start_link(_) do
  10     GenServer.start_link(__MODULE__, %{sockets: %{}}, name: __MODULE__)
  11   end
  12
  13   def add_socket(topic, socket) do
  14     GenServer.call(__MODULE__, {:add, topic, socket})
  15   end
  16
  17   def remove_socket(topic, socket) do
  18     do_remove_socket(@env, topic, socket)
  19   end
  20
  21   def get_sockets do
  22     %{sockets: stream_sockets} = GenServer.call(__MODULE__, :get_state)
  23     stream_sockets
  24   end
  25
  26   def init(init_arg) do
  27     {:ok, init_arg}
  28   end
  29
  30   def handle_call(:get_state, _from, state) do
  31     {:reply, state, state}
  32   end
  33
  34   def handle_call({:add, topic, socket}, _from, %{sockets: sockets} = state) do
  35     internal_topic = internal_topic(topic, socket)
  36     stream_socket = StreamerSocket.from_socket(socket)
  37
  38     sockets_for_topic =
  39       sockets
  40       |> Map.get(internal_topic, [])
  41       |> List.insert_at(0, stream_socket)
  42       |> Enum.uniq()
  43
  44     state = put_in(state, [:sockets, internal_topic], sockets_for_topic)
  45     Logger.debug("Got new conn for #{topic}")
  46     {:reply, state, state}
  47   end
  48
  49   def handle_call({:remove, topic, socket}, _from, %{sockets: sockets} = state) do
  50     internal_topic = internal_topic(topic, socket)
  51     stream_socket = StreamerSocket.from_socket(socket)
  52
  53     sockets_for_topic =
  54       sockets
  55       |> Map.get(internal_topic, [])
  56       |> List.delete(stream_socket)
  57
  58     state = Kernel.put_in(state, [:sockets, internal_topic], sockets_for_topic)
  59     {:reply, state, state}
  60   end
  61
  62   defp do_remove_socket(:test, _, _) do
  63     :ok
  64   end
  65
  66   defp do_remove_socket(_env, topic, socket) do
  67     GenServer.call(__MODULE__, {:remove, topic, socket})
  68   end
  69
  70   defp internal_topic(topic, socket)
  71        when topic in ~w[user user:notification direct] do
  72     "#{topic}:#{socket.assigns[:user].id}"
  73   end
  74
  75   defp internal_topic(topic, _) do
  76     topic
  77   end
  78 end