goblin

defmodule Goblin.DiskTables do
  @moduledoc false
  use GenServer
  alias Goblin.DiskTables.{DiskTable, BinarySearchIterator, StreamIterator}
  alias Goblin.BloomFilter
  alias Goblin.Manifest
  alias Goblin.Compactor

  @file_suffix ".goblin"

  defstruct [
    :dir,
    :table,
    :manifest_server,
    :compactor_server,
    :opts,
    file_number: 0
  ]

  @spec start_link(keyword()) :: GenServer.on_start()
  def start_link(opts) do
    name = opts[:name]
    local_name = opts[:local_name] || name

    args =
      opts
      |> Keyword.take([
        :db_dir,
        :manifest_server,
        :compactor_server,
        :bf_fpp,
        :max_sst_size
      ])
      |> Keyword.put(:local_name, local_name)

    GenServer.start_link(__MODULE__, args, name: name)
  end

  @spec new(Goblin.server(), Enumerable.t(Goblin.triple()), keyword()) ::
          {:ok, [Goblin.db_file()]} | {:error, term()}
  def new(server, stream, opts) do
    next_file_f = fn ->
      file = new_file(server)
      {tmp_file(file), file}
    end

    opts = Keyword.merge(opts, GenServer.call(server, :get_opts))

    with {:ok, disk_tables} <- DiskTable.write_new(stream, next_file_f, opts),
         :ok <- GenServer.call(server, {:put, disk_tables}) do
      {:ok, Enum.map(disk_tables, & &1.file)}
    end
  end

  @spec remove(Goblin.server(), Goblin.db_file()) :: :ok | {:error, term()}
  def remove(server, file) do
    GenServer.call(server, {:remove, file})
  end

  @spec search_iterators(Goblin.table(), [Goblin.db_key()], Goblin.seq_no()) :: [
          Goblin.Iterable.t()
        ]
  def search_iterators(ets, keys, seq) do
    Enum.reduce(keys, %{}, fn key, acc ->
      guard = [
        {:andalso, {:"=<", :"$1", key}, {:"=<", key, :"$2"}}
      ]

      ms = [{{:_, {:"$1", :"$2"}, :"$3"}, guard, [:"$3"]}]

      :ets.select(ets, ms)
      |> Enum.filter(fn disk_table ->
        BloomFilter.is_member(disk_table.bloom_filter, key)
      end)
      |> Enum.reduce(acc, fn disk_table, acc ->
        Map.update(acc, disk_table, [key], &[key | &1])
      end)
    end)
    |> Enum.map(fn {disk_table, keys} ->
      BinarySearchIterator.new(disk_table, keys, seq)
    end)
  end

  @spec stream_iterators(Goblin.table(), Goblin.db_key(), Goblin.db_key(), Goblin.seq_no()) :: [
          Goblin.Iterable.t()
        ]
  def stream_iterators(ets, min, max, seq) do
    guard =
      cond do
        is_nil(min) and is_nil(max) -> []
        is_nil(min) -> [{:"=<", :"$2", max}]
        is_nil(max) -> [{:"=<", min, :"$3"}]
        true -> [{:andalso, {:"=<", :"$2", max}, {:"=<", min, :"$3"}}]
      end

    ms = [{{:"$1", {:"$2", :"$3"}, :_}, guard, [:"$1"]}]

    :ets.select(ets, ms)
    |> Enum.map(&StreamIterator.new(&1, seq))
  end

  @spec iterator(Goblin.db_file()) :: Goblin.Iterable.t()
  def iterator(file) do
    StreamIterator.new(file)
  end

  @spec new_file(Goblin.server()) :: {Goblin.db_file(), Goblin.db_file()}
  def new_file(server) do
    GenServer.call(server, :new_file)
  end

  @impl GenServer
  def init(args) do
    table = :ets.new(args[:local_name], [:named_table])

    {:ok,
     %__MODULE__{
       dir: args[:db_dir],
       table: table,
       manifest_server: args[:manifest_server],
       compactor_server: args[:compactor_server],
       opts: [bf_fpp: args[:bf_fpp], max_sst_size: args[:max_sst_size]]
     }, {:continue, :recover_state}}
  end

  @impl GenServer
  def handle_call({:put, disk_tables}, _from, state) do
    Enum.each(disk_tables, fn disk_table ->
      :ets.insert(
        state.table,
        {disk_table.file, disk_table.key_range, disk_table}
      )

      Compactor.put(state.compactor_server, disk_table)
    end)

    {:reply, :ok, state}
  end

  def handle_call({:remove, file}, _from, state) do
    :ets.delete(state.table, file)
    {:reply, :ok, state}
  end

  def handle_call(:new_file, _from, state) do
    new_file = file_path(state.dir, state.file_number)
    {:reply, new_file, %{state | file_number: state.file_number + 1}}
  end

  def handle_call(:get_opts, _from, state) do
    {:reply, state.opts, state}
  end

  @impl GenServer
  def handle_continue(:recover_state, state) do
    %{disk_tables: disk_table_names, count: count} =
      Manifest.snapshot(state.manifest_server, [:disk_tables, :count])

    case recover_disk_tables(disk_table_names) do
      {:ok, disk_tables} ->
        Enum.each(disk_tables, fn disk_table ->
          :ets.insert(
            state.table,
            {disk_table.file, disk_table.key_range, disk_table}
          )

          Compactor.put(state.compactor_server, disk_table)
        end)

        :ets.insert(state.table, {:ready})
        {:noreply, %{state | file_number: count}}

      {:error, reason} ->
        {:stop, reason, state}
    end
  end

  defp recover_disk_tables(disk_table_names, acc \\ [])
  defp recover_disk_tables([], acc), do: {:ok, acc}

  defp recover_disk_tables([disk_table_name | disk_table_names], acc) do
    with {:ok, disk_table} <- DiskTable.parse(disk_table_name) do
      recover_disk_tables(disk_table_names, [disk_table | acc])
    end
  end

  defp tmp_file(file), do: "#{file}.tmp"
  defp file_path(dir, number), do: Path.join(dir, "#{to_string(number)}#{@file_suffix}")
end