defmodule Poolder.Batcher do defmacro __using__(opts) do name = Keyword.get(opts, :name, __MODULE__) batch_limit = Keyword.get(opts, :limit, 5) batch_timeout = Keyword.get(opts, :timeout, :infinity) batch_reverse = Keyword.get(opts, :reverse, false) retry = Keyword.get(opts, :retry, count: 0, backoff: 0) retries = Keyword.get(retry, :count) backoff = Keyword.get(retry, :backoff, 0) quote bind_quoted: [ name: name, batch_limit: batch_limit, batch_timeout: batch_timeout, batch_reverse: batch_reverse, retries: retries, backoff: backoff ] do @name name @batch_limit batch_limit @batch_timeout batch_timeout @batch_reverse batch_reverse @retries retries @backoff backoff @infinity batch_timeout == :infinity use GenServer @behaviour Poolder.Batcher @compile {:inline, add_queue: 2, add_many_queue: 2} # API def start_link(opts) do GenServer.start_link(__MODULE__, opts, name: @name) end def push(item) do GenServer.cast(__MODULE__, {:push, item}) end def push(pid, item) do GenServer.cast(pid, {:push, item}) end def push_many(items) do GenServer.cast(__MODULE__, {:push_many, items}) end def push_many(pid, items) do GenServer.cast(pid, {:push_many, items}) end def force_flush() do send(__MODULE__, :flush) end def force_flush(pid) do send(pid, :flush) end @impl true def init(_opts) do state = %{queue: :queue.new(), counter: :counters.new(1, []), timer: nil} handle_init(state) end @impl true if not @infinity do def handle_cast({:push, item}, %{queue: q, counter: c, timer: nil} = state) do :counters.add(c, 1, 1) timer = Process.send_after(self(), :flush, @batch_timeout) {:noreply, %{state | queue: add_queue(q, item), timer: timer}} end end def handle_cast({:push, item}, %{queue: q, counter: c} = state) do :counters.add(c, 1, 1) q = add_queue(q, item) if :counters.get(c, 1) >= @batch_limit do updated_state = %{state | queue: :queue.new(), timer: nil} flush(q, updated_state) cancel_timer(state.timer) :counters.sub(c, 1, @batch_limit) {:noreply, updated_state} else {:noreply, %{state | queue: q}} end end if not @infinity do def handle_cast({:push_many, items}, %{queue: q, counter: c, timer: nil} = state) do :counters.add(c, 1, length(items)) timer = Process.send_after(self(), :flush, @batch_timeout) {:noreply, %{state | queue: add_many_queue(q, items), timer: timer}} end end def handle_cast({:push_many, items}, %{queue: q, counter: c} = state) do :counters.add(c, 1, length(items)) q = add_many_queue(q, items) if :counters.get(c, 1) >= @batch_limit do updated_state = %{state | queue: :queue.new(), timer: nil} flush(q, updated_state) cancel_timer(state.timer) :counters.sub(c, 1, @batch_limit) {:noreply, updated_state} else {:noreply, %{state | queue: q}} end end @impl true def handle_info(:flush, %{queue: q, counter: c, timer: timer} = state) do cancel_timer(timer) updated_state = %{state | queue: :queue.new(), timer: nil} :counters.put(c, 1, 0) flush(q, updated_state) {:noreply, updated_state} end def handle_info({:state, state}, _state) do {:noreply, state} end if @batch_reverse do defp add_queue(q, item) do :queue.in(item, q) end defp add_many_queue(q, items) do Enum.reduce(items, q, fn item, acc -> :queue.in(item, acc) end) end else defp add_queue(q, item) do :queue.in_r(item, q) end defp add_many_queue(q, items) do Enum.reduce(items, q, fn item, acc -> :queue.in_r(item, acc) end) end end defp flush(q, state) do batch = :queue.to_list(q) pid = self() spawn_link(__MODULE__, :try_handle_batch, [pid, batch, 0, state]) end def handle_init(state), do: {:ok, state} def handle_batch(_batch, state), do: {:ok, state} defp cancel_timer(nil), do: :ok defp cancel_timer(ref), do: Process.cancel_timer(ref) def try_handle_batch(pid, batch, attempt, state) do try do handle_batch(batch, state) rescue error -> if attempt < @retries do Process.sleep(@backoff) try_handle_batch(pid, batch, attempt + 1, state) else {:error, error} end end end defoverridable handle_batch: 2, handle_init: 1 end end @callback handle_init(state :: term) :: term @callback handle_batch(batch :: list(), state :: term) :: term end