altworx_runbox

defmodule Runbox.Runtime.Timezip.MultiQueue do
  @moduledoc group: :internal
  @moduledoc """
  Buffering multi-queue data structure for Timezip.

  If the structure would be heap, the complexity would be Olog(N) where N is number of messages
  (topped by max demand).
  With multi-queue operation is ~ O(M) where M is number of producers.

  The produces / queues are ordered. Messages in the same timestamp across different queues are
  produced in this order.
  """

  alias __MODULE__

  defstruct queues: []

  @type queue_id :: any()
  @type t() :: %MultiQueue{
          queues: [{queue_id(), :queue.queue()}]
        }
  @type demand_distribution() :: %{optional(queue_id()) => pos_integer()}

  def new do
    %MultiQueue{}
  end

  @doc """
  Adds new queue to MultiQueue.

  Adds a new empty queue at the end.
  """
  @spec add_queue(t(), queue_id()) :: t()
  def add_queue(%MultiQueue{queues: queues} = multi_queue, queue_id) do
    %MultiQueue{multi_queue | queues: queues ++ [{queue_id, :queue.new()}]}
  end

  @doc "Removes queue and its content from multi_queue"
  @spec remove_queue(t(), queue_id()) :: t()
  def remove_queue(%MultiQueue{queues: queues} = multi_queue, queue_id) do
    %MultiQueue{multi_queue | queues: Enum.reject(queues, fn {id, _} -> id == queue_id end)}
  end

  @doc "Adds messages for queue id to multiqueue."
  @spec enqueue(t(), queue_id(), [message :: term()]) :: {:ok, t()} | {:error, term()}
  def enqueue(%MultiQueue{queues: queues} = buffer, id, messages) do
    {new_queues, found} =
      Enum.map_reduce(queues, false, fn
        {^id, queue}, _res ->
          new_queue = enqueue_many(queue, messages)
          {{id, new_queue}, true}

        other, res ->
          {other, res}
      end)

    if found do
      {:ok, %MultiQueue{buffer | queues: new_queues}}
    else
      {:error, {:bad_queue_id, id}}
    end
  end

  @doc """
  Main logic, will return messages from all queues zipped/sorted by min timestamp.

  Will return as many messages as possible, until some queue is empty.

  In returned tuple, returned values means:
   - *multi_queue:* multi-queue without emitted messages

   - *demand_distribution:* map with `id -> integer()`, meaning how many messages in output
     originates from particular queue id, can be used for manual demand handling in genstage

   - *emitted:* returned messages from multiqueue front
  """
  @spec dequeue_all(t(), comparator :: (msg, msg -> boolean())) ::
          {:ok, t(), demand_distribution(), emitted :: [msg]} | {:error, term()}
        when msg: term()
  def dequeue_all(multi_queue, comparator \\ &default_ts_tuple_comparator/2) do
    dequeue_all(multi_queue, [], %{}, comparator)
  end

  defp dequeue_all(multi_queue, acc_msgs, acc_demand_dist, comparator) do
    # find any message with the lowest timestamp
    case get_min_message(multi_queue, comparator) do
      :empty ->
        {:ok, multi_queue, acc_demand_dist, Enum.reverse(List.flatten(acc_msgs))}

      {:value, min_msg} ->
        # emit all messages with that timestamp
        {multi_queue, acc_msgs, acc_demand_dist} =
          dequeue_all_at(multi_queue, min_msg, acc_msgs, acc_demand_dist, comparator)

        # try again, maybe there is another timestamp to emit
        dequeue_all(multi_queue, acc_msgs, acc_demand_dist, comparator)
    end
  end

  # Dequeue all messages at a specified time. The time is specified by a message at that time (so we
  # can use comparator).
  defp dequeue_all_at(multi_queue, min_msg, acc_msgs, acc_demand_dist, comparator) do
    {new_queues, {acc_msgs, acc_demand_dist}} =
      Enum.map_reduce(
        multi_queue.queues,
        {acc_msgs, acc_demand_dist},
        fn {queue_id, queue}, {acc_msgs, acc_demand_dist} ->
          # for every queue pop all messages at the min_msg time
          {queue, msgs} = pop_all_at(queue, min_msg, comparator)
          len = length(msgs)

          {
            {queue_id, queue},
            {
              [msgs, acc_msgs],
              Map.update(acc_demand_dist, queue_id, len, &(&1 + len))
            }
          }
        end
      )

    {%MultiQueue{multi_queue | queues: new_queues}, acc_msgs, acc_demand_dist}
  end

  defp enqueue_many(queue, items) do
    Enum.reduce(items, queue, fn item, q_acc -> :queue.in(item, q_acc) end)
  end

  # Find any message with minimal time as `{:value, msg}`.  If any queue is empty `:empty` is
  # returned instead.
  defp get_min_message(%MultiQueue{queues: queues}, comparator) do
    queues
    |> Enum.map(fn {_id, queue} -> :queue.peek(queue) end)
    |> Enum.reduce(&queue_min_reducer(&1, &2, comparator))
  end

  defp queue_min_reducer(_item, :empty = acc, _comparator), do: acc
  defp queue_min_reducer(:empty = item, _acc, _comparator), do: item

  defp queue_min_reducer({:value, msg1} = item, {:value, msg2} = acc, comparator) do
    if comparator.(msg1, msg2), do: item, else: acc
  end

  # Pop all messages at time specified by minimal message. Note the provided message must be
  # minimal!
  defp pop_all_at(queue, min_msg, acc_msgs \\ [], comparator) do
    case :queue.out(queue) do
      {:empty, new_queue} ->
        {new_queue, acc_msgs}

      {{:value, value}, new_queue} ->
        # this checks !(min_msg.timestamp < value.timestamp) (the else clause is the main one)
        # which is equivalent to value.timestamp <= min_msg.timestamp
        # given the min_msg is minimal it can only be true that value.timestamp == min_msg.timestamp
        if comparator.(min_msg, value) do
          # value.timestamp > min_msg.timestamp => we cannot remove this yet
          {queue, acc_msgs}
        else
          # value.timestamp == min_msg.timestamp => we can pop this msg
          pop_all_at(new_queue, min_msg, [value | acc_msgs], comparator)
        end
    end
  end

  # return true if first msg is before second (smaller ts)
  defp default_ts_tuple_comparator({ts1, _val}, {ts2, _val2}) do
    ts1 < ts2
  end
end