retry_queue.ex 6.63 KB
Newer Older
1
# Pleroma: A lightweight social networking server
kaniini's avatar
kaniini committed
2
# Copyright © 2017-2019 Pleroma Authors <https://pleroma.social/>
3
4
# SPDX-License-Identifier: AGPL-3.0-only

eal's avatar
eal committed
5
6
defmodule Pleroma.Web.Federator.RetryQueue do
  use GenServer
Maksim's avatar
Maksim committed
7

eal's avatar
eal committed
8
9
10
  require Logger

  def init(args) do
Michael Loftis's avatar
Michael Loftis committed
11
12
13
    queue_table = :ets.new(:pleroma_retry_queue, [:bag, :protected])

    {:ok, %{args | queue_table: queue_table, running_jobs: :sets.new()}}
eal's avatar
eal committed
14
15
16
  end

  def start_link() do
Michael Loftis's avatar
Michael Loftis committed
17
    enabled =
lain's avatar
lain committed
18
      if Mix.env() == :test, do: true, else: Pleroma.Config.get([__MODULE__, :enabled], false)
lain's avatar
lain committed
19
20
21

    if enabled do
      Logger.info("Starting retry queue")
Michael Loftis's avatar
Michael Loftis committed
22
23
24
25
26
27
28
29
30
31

      linkres =
        GenServer.start_link(
          __MODULE__,
          %{delivered: 0, dropped: 0, queue_table: nil, running_jobs: nil},
          name: __MODULE__
        )

      maybe_kickoff_timer()
      linkres
lain's avatar
lain committed
32
33
34
35
    else
      Logger.info("Retry queue disabled")
      :ignore
    end
eal's avatar
eal committed
36
37
38
39
40
41
  end

  def enqueue(data, transport, retries \\ 0) do
    GenServer.cast(__MODULE__, {:maybe_enqueue, data, transport, retries + 1})
  end

Michael Loftis's avatar
Michael Loftis committed
42
43
44
45
46
47
48
49
  def get_stats() do
    GenServer.call(__MODULE__, :get_stats)
  end

  def reset_stats() do
    GenServer.call(__MODULE__, :reset_stats)
  end

eal's avatar
eal committed
50
  def get_retry_params(retries) do
lain's avatar
lain committed
51
    if retries > Pleroma.Config.get([__MODULE__, :max_retries]) do
eal's avatar
eal committed
52
      {:drop, "Max retries reached"}
eal's avatar
eal committed
53
    else
eal's avatar
eal committed
54
      {:retry, growth_function(retries)}
eal's avatar
eal committed
55
56
57
    end
  end

Michael Loftis's avatar
Michael Loftis committed
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
  def get_retry_timer_interval() do
    Pleroma.Config.get([:retry_queue, :interval], 1000)
  end

  defp ets_count_expires(table, current_time) do
    :ets.select_count(
      table,
      [
        {
          {:"$1", :"$2"},
          [{:"=<", :"$1", {:const, current_time}}],
          [true]
        }
      ]
    )
  end

  defp ets_pop_n_expired(table, current_time, desired) do
    {popped, _continuation} =
      :ets.select(
        table,
        [
          {
            {:"$1", :"$2"},
            [{:"=<", :"$1", {:const, current_time}}],
            [:"$_"]
          }
        ],
        desired
      )

    popped
    |> List.foldl(true, fn e, acc ->
      :ets.delete_object(table, e)
      acc
    end)

    popped
  end

  def maybe_start_job(running_jobs, queue_table) do
    # we don't want to hit the ets or the DateTime more times than we have to
    # could optimize slightly further by not using the count, and instead grabbing
    # up to N objects early...
    current_time = DateTime.to_unix(DateTime.utc_now())
    n_running_jobs = :sets.size(running_jobs)

lain's avatar
lain committed
105
    if n_running_jobs < Pleroma.Config.get([__MODULE__, :max_jobs]) do
Michael Loftis's avatar
Michael Loftis committed
106
107
108
109
      n_ready_jobs = ets_count_expires(queue_table, current_time)

      if n_ready_jobs > 0 do
        # figure out how many we could start
lain's avatar
lain committed
110
        available_job_slots = Pleroma.Config.get([__MODULE__, :max_jobs]) - n_running_jobs
Michael Loftis's avatar
Michael Loftis committed
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
        start_n_jobs(running_jobs, queue_table, current_time, available_job_slots)
      else
        running_jobs
      end
    else
      running_jobs
    end
  end

  defp start_n_jobs(running_jobs, _queue_table, _current_time, 0) do
    running_jobs
  end

  defp start_n_jobs(running_jobs, queue_table, current_time, available_job_slots)
       when available_job_slots > 0 do
    candidates = ets_pop_n_expired(queue_table, current_time, available_job_slots)

    candidates
    |> List.foldl(running_jobs, fn {_, e}, rj ->
      {:ok, pid} = Task.start(fn -> worker(e) end)
      mref = Process.monitor(pid)
      :sets.add_element(mref, rj)
    end)
  end

  def worker({:send, data, transport, retries}) do
    case transport.publish_one(data) do
      {:ok, _} ->
        GenServer.cast(__MODULE__, :inc_delivered)
        :delivered

      {:error, _reason} ->
        enqueue(data, transport, retries)
        :retry
    end
  end

  def handle_call(:get_stats, _from, %{delivered: delivery_count, dropped: drop_count} = state) do
    {:reply, %{delivered: delivery_count, dropped: drop_count}, state}
  end

  def handle_call(:reset_stats, _from, %{delivered: delivery_count, dropped: drop_count} = state) do
    {:reply, %{delivered: delivery_count, dropped: drop_count},
     %{state | delivered: 0, dropped: 0}}
  end

  def handle_cast(:reset_stats, state) do
    {:noreply, %{state | delivered: 0, dropped: 0}}
  end

  def handle_cast(
        {:maybe_enqueue, data, transport, retries},
        %{dropped: drop_count, queue_table: queue_table, running_jobs: running_jobs} = state
      ) do
eal's avatar
eal committed
165
    case get_retry_params(retries) do
Maksim's avatar
updates    
Maksim committed
166
      {:retry, timeout} ->
Michael Loftis's avatar
Michael Loftis committed
167
168
169
        :ets.insert(queue_table, {timeout, {:send, data, transport, retries}})
        running_jobs = maybe_start_job(running_jobs, queue_table)
        {:noreply, %{state | running_jobs: running_jobs}}
eal's avatar
eal committed
170

eal's avatar
eal committed
171
172
173
      {:drop, message} ->
        Logger.debug(message)
        {:noreply, %{state | dropped: drop_count + 1}}
eal's avatar
eal committed
174
175
176
    end
  end

Michael Loftis's avatar
Michael Loftis committed
177
178
179
180
181
182
183
184
185
186
187
188
189
190
  def handle_cast(:kickoff_timer, state) do
    retry_interval = get_retry_timer_interval()
    Process.send_after(__MODULE__, :retry_timer_run, retry_interval)
    {:noreply, state}
  end

  def handle_cast(:inc_delivered, %{delivered: delivery_count} = state) do
    {:noreply, %{state | delivered: delivery_count + 1}}
  end

  def handle_cast(:inc_dropped, %{dropped: drop_count} = state) do
    {:noreply, %{state | dropped: drop_count + 1}}
  end

eal's avatar
eal committed
191
192
  def handle_info({:send, data, transport, retries}, %{delivered: delivery_count} = state) do
    case transport.publish_one(data) do
eal's avatar
eal committed
193
      {:ok, _} ->
eal's avatar
eal committed
194
        {:noreply, %{state | delivered: delivery_count + 1}}
eal's avatar
eal committed
195

Maksim's avatar
Maksim committed
196
      {:error, _reason} ->
eal's avatar
eal committed
197
        enqueue(data, transport, retries)
eal's avatar
eal committed
198
199
200
201
        {:noreply, state}
    end
  end

Michael Loftis's avatar
Michael Loftis committed
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
  def handle_info(
        :retry_timer_run,
        %{queue_table: queue_table, running_jobs: running_jobs} = state
      ) do
    maybe_kickoff_timer()
    running_jobs = maybe_start_job(running_jobs, queue_table)
    {:noreply, %{state | running_jobs: running_jobs}}
  end

  def handle_info({:DOWN, ref, :process, _pid, _reason}, state) do
    %{running_jobs: running_jobs, queue_table: queue_table} = state
    running_jobs = :sets.del_element(ref, running_jobs)
    running_jobs = maybe_start_job(running_jobs, queue_table)
    {:noreply, %{state | running_jobs: running_jobs}}
  end

eal's avatar
eal committed
218
219
220
221
222
  def handle_info(unknown, state) do
    Logger.debug("RetryQueue: don't know what to do with #{inspect(unknown)}, ignoring")
    {:noreply, state}
  end

Michael Loftis's avatar
Michael Loftis committed
223
224
  if Mix.env() == :test do
    defp growth_function(_retries) do
lain's avatar
lain committed
225
      _shutit = Pleroma.Config.get([__MODULE__, :initial_timeout])
Michael Loftis's avatar
Michael Loftis committed
226
227
228
229
      DateTime.to_unix(DateTime.utc_now()) - 1
    end
  else
    defp growth_function(retries) do
lain's avatar
lain committed
230
231
      round(Pleroma.Config.get([__MODULE__, :initial_timeout]) * :math.pow(retries, 3)) +
        DateTime.to_unix(DateTime.utc_now())
Michael Loftis's avatar
Michael Loftis committed
232
233
234
235
236
    end
  end

  defp maybe_kickoff_timer() do
    GenServer.cast(__MODULE__, :kickoff_timer)
eal's avatar
eal committed
237
238
  end
end