defmodule Poolder.Batcher do defmacro __using__(opts) do name = Keyword.get(opts, :name) batch_limit = Keyword.get(opts, :limit, 5) batch_timeout = Keyword.get(opts, :timeout, :infinity) batch_reverse = Keyword.get(opts, :reverse, false) batch_indexed = Keyword.get(opts, :indexed, false) stream = Keyword.get(opts, :stream, true) retry = Keyword.get(opts, :retry, count: 0, backoff: 0) retries = Keyword.get(retry, :count) backoff = Keyword.get(retry, :backoff, 0) hibernate_after = Keyword.get(opts, :hibernate_after, :infinity) priority = Keyword.get(opts, :priority, :normal) quote bind_quoted: [ name: name, batch_limit: batch_limit, batch_timeout: batch_timeout, batch_reverse: batch_reverse, batch_indexed: batch_indexed, stream: stream, retries: retries, backoff: backoff, hibernate_after: hibernate_after, priority: priority ] do @name name || __MODULE__ @batch_limit batch_limit @batch_timeout batch_timeout @batch_reverse batch_reverse @batch_indexed batch_indexed @stream stream @retries retries @backoff backoff @infinity batch_timeout == :infinity @hibernate_after hibernate_after @priority priority @abnormal_priority priority != :normal @result_fun if @stream, do: &:edeque.stream/3, else: &:edeque.to_list/3 @behaviour Poolder.Batcher @compile {:inline, add: 4} if @hibernate_after <= @batch_timeout do raise ArgumentError, "#{__MODULE__} hibernate_after must be greater than batch_timeout" end # API def child_spec(opts \\ []) do %{ id: __MODULE__, start: {__MODULE__, :start_link, [opts]}, restart: :permanent, type: :worker, shutdown: 500 } end def start(opts \\ []) do pid = spawn(__MODULE__, :run, [opts]) Process.register(pid, @name) {:ok, pid} end def start_link(opts \\ []) do pid = spawn_link(__MODULE__, :run, [opts]) Process.register(pid, @name) {:ok, pid} end def run(opts) do {t, c} = :edeque.new() pid = self() monitor(pid, opts) if @abnormal_priority do Process.flag(:priority, @priority) end handle_init(opts) loop(t, c, nil) end def push(item) do send(__MODULE__, {:push, item}) end def push(pid, item) do send(pid, {:push, item}) end def push_front(pid, item) do send(pid, {:first, item}) end def push_front(item) do send(__MODULE__, {:first, item}) end def pop_at(pid, index) do send(pid, {:pop, index}) end def flush() do send(__MODULE__, :flush) end def flush(pid) do send(pid, :flush) end def loop(t, c, tref, hibernate_after \\ @hibernate_after) do receive do {x, item} -> add(x, t, c, item) case :edeque.size(c) do 1 -> tref = send_after() loop(t, c, tref) size when size < @batch_limit -> loop(t, c, tref) _ -> flush(t, c, tref) end if :edeque.size(c) == @batch_limit do flush(t, c, tref) else loop(t, c, tref) end :flush -> flush(t, c, tref) _ -> loop(t, c, tref) after hibernate_after -> handle_hibernate({t, c, tref}) :erlang.hibernate(__MODULE__, :loop, [t, c, tref]) end end defp do_result(t, c) do @result_fun.(t, c, reverse: @batch_reverse, indexed: @batch_indexed) end defp flush(t, c, tref) do tref = cancel_timer(tref) if :edeque.size(c) == 0 do loop(t, c, tref) else pid = self() messages = do_result(t, c) state = %{table: t} spawn_link(__MODULE__, :try_handle_batch, [pid, messages, 1, state]) :edeque.reset_counters(c) loop(:edeque.new_table(), c, tref) end end defp monitor(pid, args) when is_list(args) do case Keyword.get(args, :monitor) do nil -> :ok monitor -> monitor.monitor(pid) end end defp monitor(pid, args) when is_map(args) do case Map.get(args, :monitor) do nil -> :ok monitor -> monitor.monitor(pid) end end defp monitor(_pid, _args), do: :ok if @infinity do defp send_after, do: nil else defp send_after, do: Process.send_after(self(), :flush, @batch_timeout) end defp cancel_timer(nil), do: nil defp cancel_timer(tref) do Process.cancel_timer(tref) nil end defp add(:push, t, c, item) do :edeque.push(t, c, item) end defp add(:first, t, c, item) do :edeque.push_front(t, c, item) end defp add(:pop, t, c, index) do :edeque.pop_at(t, c, index) end def handle_init(state), do: {:ok, state} def handle_batch(_batch, _state), do: :ok def handle_hibernate(_state), do: :ok 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 after :ets.delete(state.table) end end defoverridable handle_batch: 2, handle_init: 1, handle_hibernate: 1 end end @callback handle_init(state :: term) :: term @callback handle_batch(batch_stream :: Enumerable.t(), state :: term) :: term @callback handle_hibernate(state :: term) :: term end