defmodule Kvasir.AgentServer.Control.Server do @commands ~W(REBUILD STATUS CONNECT METRICS VERIFY_CACHE) def commands, do: @commands def init(socket, transport, server) do # Join for status updates :pg.join({server, :status}, spawn_link(fn -> status_loop(socket, transport, server) end)) :ok end def close(_server) do :ok end ### The Commands alias Kvasir.AgentServer.Config @spec handle_command([String.t()], term, %{required(String.t()) => fun}) :: :ok | {:reply, iodata()} def handle_command(command, server, _custom_commands) def handle_command(["REBUILD", topic | opts], server, _custom_commands) do [key | part] = :lists.reverse(opts) partition = part |> :lists.reverse() |> Enum.join(" ") if agent = Enum.find_value( Config.agents(server, topic), &if(&1.partition == partition, do: &1.agent) ) do with {:ok, p_key} <- agent.__agent__(:key).parse(key), true <- match_partition?(partition, p_key) || {:error, :key_partition_mismatch}, :ok <- agent.rebuild(p_key) do {:reply, "OK\n"} else {:error, err} -> {:reply, "ERR #{inspect(err)}\n"} end else {:reply, "ERR unknown_agent\n"} end end def handle_command(["VERIFY_CACHE", restore, topic | opts], server, _custom_commands) do case String.downcase(restore) do r when r in ~W(true yes y 1) -> check_cache(topic, opts, server, true) r when r in ~W(false no n 0) -> check_cache(topic, opts, server, false) _ -> {:reply, "ERR invalid restore argument required\n"} end end def handle_command(["CONNECT" | opts], server, _custom_commands) do if id = List.first(opts) do a = server |> Config.agents(id) |> Enum.map( &[id, ?\ , ip!(&1.opts[:ip]), ?\ , to_string(&1.opts[:port]), ?\ , &1.partition, ?\n] ) {:reply, [ ["LIST ", id, ?\n], a, ["DONE ", id, ?\n] ]} else a = server |> Config.agents() |> Enum.map( &[&1.id, ?\ , ip!(&1.opts[:ip]), ?\ , to_string(&1.opts[:port]), ?\ , &1.partition, ?\n] ) {:reply, [ ["LIST", ?\n], a, ["DONE", ?\n] ]} end end def handle_command(["METRICS"], server, _custom_commands) do a = server |> Config.agents() |> Enum.map(&[&1.id, ?\ , &1.partition, ?\ , to_string(:counters.get(&1.counter, 1)), ?\n]) {:reply, [ "METRICS\n", a, "DONE METRICS\n" ]} end def handle_command(["STATUS"], server, _custom_commands), do: {:reply, ["STATUS ", to_string(Config.status(server)), ?\n]} def handle_command([cmd | opts], server, custom_commands) do if c = custom_commands[cmd], do: c.(opts), else: unknown!([cmd | opts], server) end def handle_command(unknown, server, _custom_commands), do: unknown!(unknown, server) defp unknown!(unknown, server) do require Logger Logger.warn(fn -> "Kvasir AgentServer<#{inspect(server)}>: Received unknown control command: #{ inspect(unknown) } " end) if List.first(unknown) in @commands do {:reply, "ERR missing command arguments\n"} else {:reply, "ERR unknown command\n"} end end defp ip!(ip), do: :inet.ntoa(ip) defp status_loop(socket, transport, server) do receive do {:status, ^server, new} -> transport.send(socket, ["STATUS ", to_string(new), ?\n]) end status_loop(socket, transport, server) end ### Cache ### defp check_cache(topic, opts, server, restore) do if opts == [] do server |> Config.agents(topic) |> Enum.map(& &1.agent) |> check_cache_integrity(topic, restore) else partition = Enum.join(opts, " ") server |> Config.agents(topic) |> Enum.filter(&(&1.partition == partition)) |> Enum.map(& &1.agent) |> check_cache_integrity(topic, restore) end end defp check_cache_integrity(agents, topic, restore) defp check_cache_integrity([], _topic, _restore), do: {:reply, "ERR unknown_agent\n"} defp check_cache_integrity(agents, topic, restore) do {:async, "OK\n", fn sender -> total = agents |> Enum.map(&Task.async(&1, :check_cache_integrity, [restore])) |> Enum.map(&Task.await(&1, :infinity)) |> Enum.sum() sender.("VERIFY_CACHE #{topic} #{total}\n") end} end ### Partition Matching ### @spec match_partition?(String.t(), term) :: boolean defp match_partition?(partition, key) defp match_partition?("*", _key), do: true defp match_partition?(partition, key) do [match, compare, value] = String.split(partition, " ") match = match_partition_match(match, key) value = match_partition_value(value) match_partition_compare(compare, match, value) end defp match_partition_match(".", key), do: key defp match_partition_match(pattern, key) do pattern |> String.split(".", trim: true) |> Enum.reduce(key, &Map.get(&2, String.to_atom(&1))) end defp match_partition_compare("=", match, value), do: match == value defp match_partition_compare("<", match, value), do: match < value defp match_partition_compare("<=", match, value), do: match <= value defp match_partition_compare(">", match, value), do: match > value defp match_partition_compare(">=", match, value), do: match >= value defp match_partition_value(":" <> atom), do: String.to_atom(atom) defp match_partition_value("\"" <> string), do: String.slice(string, 0..-2) defp match_partition_value("true"), do: true defp match_partition_value("false"), do: false defp match_partition_value(value), do: String.to_integer(value) end