ex_limiter

defmodule ExLimiter.Storage.Memcache do
  @moduledoc """
  Token bucket backend written for memcache. Stores the last timestamp
  and amount in separate keys, and utilizes memcache increments for consumption
  """
  use ExLimiter.Storage
  alias ExLimiter.Utils

  def fetch(%Bucket{key: key}) do
    key_map = keys(key)

    try do
      Map.keys(key_map)
      |> Memcachir.mget(cas: true)
      |> case do
        {:ok, result} -> from_memcached(result, key, key_map)
        _ -> Bucket.new(key)
      end
    catch
      :exit, _ -> Bucket.new(key)
    end
  end

  def refresh(%Bucket{key: key} = bucket, _type \\ :soft) do
    try do
      keys(key)
      |> Enum.map(&mset_command(&1, bucket))
      |> Memcachir.mset_cas()
      |> case do
        {:ok, _} -> {:ok, bucket}
        {:error, error} -> {:error, error}
      end
    catch
      :exit, _ -> {:error, :memcached}
    end
  end

  def delete(%Bucket{key: key}) do
    try do
      keys(key)
      |> Map.keys()
      |> Enum.map(&Memcachir.delete/1)

      Bucket.new(key)
    catch
      :exit, _ -> Bucket.new(key)
    end
  end

  def update(key, update_fun) do
    bucket = %Bucket{key: key}

    fetch(bucket)
    |> update_fun.()
    |> refresh()
    |> case do
      {:ok, b} -> b
      {:error, _} -> fetch(bucket) # memcached latency is short enough that there probably wasn't much leaked here
    end
  end

  def consume(%Bucket{key: key} = bucket, inc) do
    try do
      case Memcachir.incr("amount_#{key}", inc) do
        {:ok, result} -> {:ok, %{bucket | value: result}}
        _ -> {:ok, Bucket.new(key)}
      end
    catch
      :exit, _ -> {:ok, Bucket.new(key)}
    end
  end

  defp keys(key), do: %{"amount_#{key}" => :value, "last_#{key}" => :last}

  defp from_memcached(map, key, key_map) when is_map(map) do
    key_map
    |> Enum.reduce(%{version: %{}}, &reduce_memcached(&1, &2, map))
    |> Enum.into(%{})
    |> Bucket.new(key)
  end

  defp mset_command({key, bucket_key}, %Bucket{version: versions} = b) do
    value = Map.get(b, bucket_key) |> to_string()
    {key, value, Map.get(versions, bucket_key, 0)}
  end

  defp reduce_memcached({key, bucket_key}, acc, map), do: add_result(acc, bucket_key, map[key])

  defp add_result(%{version: versions} = acc, bucket_key, {val, cas}) do
    acc
    |> Map.put(bucket_key, Utils.parse_integer(val))
    |> Map.put(:version, Map.put(versions, bucket_key, cas))
  end
  defp add_result(acc, bucket_key, _), do: add_result(acc, bucket_key, default(bucket_key))

  defp default(:value), do: {0, 0}
  defp default(:last), do: {Utils.now(), 0}
end