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

   1 # Pleroma: A lightweight social networking server
   2 # Copyright © 2017-2020 Pleroma Authors <https://pleroma.social/>
   3 # SPDX-License-Identifier: AGPL-3.0-only
   4
   5 defmodule Pleroma.Web.Streamer.State do
   6   use GenServer
   7   require Logger
   8
   9   alias Pleroma.Web.Streamer.StreamerSocket
  10
  11   @env Mix.env()
  12
  13   def start_link(_) do
  14     GenServer.start_link(__MODULE__, %{sockets: %{}}, name: __MODULE__)
  15   end
  16
  17   def add_socket(topic, socket) do
  18     GenServer.call(__MODULE__, {:add, topic, socket})
  19   end
  20
  21   def remove_socket(topic, socket) do
  22     do_remove_socket(@env, topic, socket)
  23   end
  24
  25   def get_sockets do
  26     %{sockets: stream_sockets} = GenServer.call(__MODULE__, :get_state)
  27     stream_sockets
  28   end
  29
  30   def init(init_arg) do
  31     {:ok, init_arg}
  32   end
  33
  34   def handle_call(:get_state, _from, state) do
  35     {:reply, state, state}
  36   end
  37
  38   def handle_call({:add, topic, socket}, _from, %{sockets: sockets} = state) do
  39     internal_topic = internal_topic(topic, socket)
  40     stream_socket = StreamerSocket.from_socket(socket)
  41
  42     sockets_for_topic =
  43       sockets
  44       |> Map.get(internal_topic, [])
  45       |> List.insert_at(0, stream_socket)
  46       |> Enum.uniq()
  47
  48     state = put_in(state, [:sockets, internal_topic], sockets_for_topic)
  49     Logger.debug("Got new conn for #{topic}")
  50     {:reply, state, state}
  51   end
  52
  53   def handle_call({:remove, topic, socket}, _from, %{sockets: sockets} = state) do
  54     internal_topic = internal_topic(topic, socket)
  55     stream_socket = StreamerSocket.from_socket(socket)
  56
  57     sockets_for_topic =
  58       sockets
  59       |> Map.get(internal_topic, [])
  60       |> List.delete(stream_socket)
  61
  62     state = Kernel.put_in(state, [:sockets, internal_topic], sockets_for_topic)
  63     {:reply, state, state}
  64   end
  65
  66   defp do_remove_socket(:test, _, _) do
  67     :ok
  68   end
  69
  70   defp do_remove_socket(_env, topic, socket) do
  71     GenServer.call(__MODULE__, {:remove, topic, socket})
  72   end
  73
  74   defp internal_topic(topic, socket)
  75        when topic in ~w[user user:notification direct] do
  76     "#{topic}:#{socket.assigns[:user].id}"
  77   end
  78
  79   defp internal_topic(topic, _) do
  80     topic
  81   end
  82 end