mentat

defmodule Mentat do
  @external_resource readme = "README.md"
  @moduledoc readme
             |> File.read!()
             |> String.split("<!--MDOC !-->")
             |> Enum.fetch!(1)

  use Supervisor
  use Oath

  @cache_opts NimbleOptions.new!([
    name: [
      type: :atom,
      doc: "The cache name as an atom",
      required: true
    ],
    cleanup_interval: [
      type: :pos_integer,
      doc: "How often the janitor process will remove old keys.",
      default: 5_000
    ],
    ets_args: [
      type: {:list, :any},
      doc: "Additional arguments to pass to `:ets.new/2`",
      default: []
    ],
    ttl: [
      type: :timeout,
      doc: "The default ttl for all keys. Defaults to `:infinity`.",
      default: :infinity
    ],
    limit: [
      type: :keyword_list,
      doc: "Limits to the number of keys a cache will store.",
      keys: [
        size: [
          type: :pos_integer,
          required: true,
          doc: "The maximum number of values to store in the cache"
        ],
        reclaim: [
          type: :float,
          default: 0.1,
          doc: "The percentage of keys to reclaim if the limit is exceeded."
        ]
      ]
    ],
    clock: [
      type: :any,
      doc: false,
      type_doc: false,
      default: System
    ]
  ])

  @default_limit %{reclaim: 0.1}

  @type cache_opt() :: unquote(NimbleOptions.option_typespec(@cache_opts))
  @type cache_opts() :: [cache_opt()]
  @type name :: atom()
  @type key :: term()
  @type value :: term()
  @type put_opts :: [
    {:ttl, pos_integer() | :infinity},
  ]

  alias Mentat.Janitor

  @doc false
  def child_spec(opts) do
    name = opts[:name] || raise ArgumentError, ":name is required"

    %{
      id: name,
      type: :supervisor,
      start: {__MODULE__, :start_link, [opts]},
    }
  end

  @doc """
  Starts a new cache.

  Options:
  #{NimbleOptions.docs(@cache_opts)}
  """
  @spec start_link(cache_opts()) :: Supervisor.on_start()
  def start_link(args) do
    args = NimbleOptions.validate!(args, @cache_opts)
    name = args[:name]
    Supervisor.start_link(__MODULE__, args, name: name)
  end

  @doc """
  Fetches a value or executes the fallback function. The function can return
  either `{:commit, term()}` or `{:ignore, term()}`. If `{:commit, term()}` is
  returned, the value will be stored in the cache before its returned. See the
  "TTLs" section for a list of options.

  Since `fetch` is logically the same as a `get` (and also a `put` if
  `{:commit, term()}` is returned from the fallback function), it shares the
  same Telemetry events (`[:mentat, :get]` and `[:mentat, :put]`).

  ## Example

  ```
  Mentat.fetch(:cache, user_id, fn user_id ->
    case get_user(user_id) do
      {:ok, user} ->
        {:commit, user}

      error ->
        {:ignore, error}
    end
  end)
  ```
  """
  @spec fetch(name(), key(), put_opts(), (key() -> {:commit, value()} | {:ignore, value()})) :: value()
  def fetch(cache, key, opts \\ [], fallback) do
    {status, value} = lookup(cache, key)

    if status == :hit do
      value
    else
      case fallback.(key) do
        {:commit, value} -> put(cache, key, value, opts)
        {:ignore, value} -> value
      end
    end
  end

  @doc """
  Retrieves a value from a the cache. Returns `nil` if the key is not found.
  """
  @spec get(name(), key()) :: value()
  def get(cache, key) do
    {_status, value} = lookup(cache, key)
    value
  end


  @doc """
  Puts a new key into the cache. See the "TTLs" section for a list of
  options.
  """
  @spec put(name(), key(), value(), put_opts()) :: value() | no_return()
  @decorate pre("ttls are positive", fn _, _, _, opts ->
    if opts[:ttl], do: opts[:ttl] > 0, else: true
  end)
  @decorate post("value is returned", fn _, _, value, _, return ->
    value == return
  end)
  def put(cache, key, value, opts \\ [])
  def put(cache, key, value, opts) do
    :telemetry.span([:mentat, :put], %{key: key, cache: cache}, fn ->
      config = get_config(cache)

      # Leaving this telemetry event for legacy support. We should remove eventually.
      :telemetry.execute([:mentat, :put], %{}, %{key: key, cache: cache})

      now = ms_time(config.clock)
      ttl = opts[:ttl] || config.ttl

      if ttl < 0 do
        raise ArgumentError, "`:ttl` must be greater than 0"
      end

      true = :ets.insert(cache, {key, value, now, ttl})

      # If we've reached the limit on the table, we need to purge a number of old
      # keys. We do this by calling the janitor process and telling it to purge.
      # This will, in turn call immediately back into the remove_oldest function.
      # The back and forth here is confusing to follow, but its necessary because
      # we want to do the purging in a different process.
      if config.limit != :none && :ets.info(cache, :size) > config.limit.size do
        Janitor.reclaim(janitor(cache))
      end

      {value, %{key: key, cache: cache}}
    end)
  end

  @doc """
  Updates a keys inserted at time. This is useful in conjunction with limits
  when you want to evict the oldest keys. Returns `true` if the key was found
  and `false` if it was not.
  """
  @spec touch(name(), key()) :: boolean()
  def touch(cache, key) do
    config = get_config(cache)
    now    = ms_time(config.clock)
    :ets.update_element(cache, key, {3, now})
  end

  @doc """
  Deletes a key from the cache
  """
  @spec delete(name(), key()) :: true
  def delete(cache, key) do
    :ets.delete(cache, key)
  end

  @doc """
  Returns the current size of the cache
  """
  @spec size(name()) :: pos_integer()
  def size(cache) do
    :ets.info(cache, :size)
  end

  @doc """
  Returns a list of all keys. By default this function only returns keys
  that have no exceeded their TTL. You can pass the `all: true` option to the function
  in order to return all present keys, which may include keys that have exceeded
  their TTL but have not been purged yet.
  """
  @spec keys(name()) :: [key()]
  def keys(cache, opts \\ []) do
    ms = if opts[:all] == true do
      [{{:"$1", :_, :_, :_}, [], [:"$1"]}]
    else
      config = get_config(cache)
      now    = ms_time(config.clock)
      [
        {{:"$1", :_, :"$2", :"$3"},
         [
           {:orelse,
            {:andalso, {:is_integer, :"$3"}, {:>, {:+, :"$2", :"$3"}, now}},
            {:==, :"$3", :infinity}}
         ], [:"$1"]}
      ]
    end

    :ets.select(cache, ms)
  end

  @doc """
  Removes all keys from the cache.
  """
  @spec purge(name()) :: true
  def purge(cache) do
    :ets.delete_all_objects(cache)
  end

  @doc false
  def remove_expired(cache) do
    config = get_config(cache)
    now    = ms_time(config.clock)

    # Find all expired keys by selecting the timestamp and ttl, adding them together
    # and finding the keys that are lower than the current time
    ms = [
      {{:_, :_, :"$1", :"$2"},
        [{:andalso, {:is_integer, :"$2"}, {:<, {:+, :"$1", :"$2"}, now}}], [true]}
    ]

    :ets.select_delete(cache, ms)
  end

  @doc false
  def remove_oldest(cache, count) do
    ms = [{{:"$1", :_, :"$2", :_}, [], [{{:"$2", :"$1"}}]}]
    entries = :ets.select(cache, ms)

    entries
    |> Enum.sort()
    |> Enum.take(count)
    |> Enum.map(fn {ts, key} -> :ets.match_delete(cache, {key, :_, ts, :_}) end)
    |> Enum.count(& &1)
  end

  def init(args) do
    name     = args[:name]
    interval = args[:cleanup_interval]
    limit    = args[:limit] || :none
    limit    = if limit != :none, do: Map.merge(@default_limit, Map.new(limit)), else: limit
    ets_args = args[:ets_args] || []
    clock    = args[:clock] || System
    ttl      = args[:ttl]
    ^name    = :ets.new(name, [:set, :named_table, :public] ++ ets_args)

    put_config(name, %{limit: limit, ttl: ttl, clock: clock})

    janitor_opts = [
      name: janitor(name),
      interval: interval,
      cache: name
    ]

    children = [
      {Mentat.Janitor, janitor_opts}
    ]

    Supervisor.init(children, strategy: :one_for_one)
  end

  def stop(name) do
    Supervisor.stop(name)
  end

  defp lookup(cache, key) do
    :telemetry.span([:mentat, :get], %{cache: cache, key: key}, fn ->
      config = get_config(cache)
      now    = ms_time(config.clock)

      {status, value} =
        case :ets.lookup(cache, key) do
          [] -> {:miss, nil}
          [{^key, _val, ts, ttl}] when is_integer(ttl) and ts + ttl <= now -> {:miss, nil}
          [{^key, val, _ts, _expire_at}] -> {:hit, val}
        end

      # Leaving this telemetry event for legacy support. We should remove eventually.
      :telemetry.execute([:mentat, :get], %{status: status}, %{key: key, cache: cache})

      {{status, value}, %{key: key, cache: cache, status: status}}
    end)
  end

  defp put_config(cache, config) do
    :persistent_term.put({__MODULE__, cache}, config)
  end

  def get_config(cache) do
    :persistent_term.get({__MODULE__, cache})
  end

  defp ms_time(clock) do
    # Clock is going `System` in most cases and is set inside the init function
    clock.monotonic_time(:millisecond)
  end

  defp janitor(name) do
    :"#{name}_janitor"
  end
end