defmodule Elsa.Util do
@moduledoc """
Provides functions for simplifying first-class interactions (consuming and
producing) such as connecting to a cluster and establishing a persistent
client process for interacting with a cluster.
"""
require Logger
alias Elsa.RetryConfig
@default_max_chunk_size 900_000
@timestamp_size_in_bytes 10
@doc """
Wrap establishing a connection to a cluster for performing an operation.
"""
@spec with_connection(Elsa.endpoints(), atom(), fun()) :: term()
def with_connection(endpoints, type \\ :any, fun) when is_function(fun) do
endpoints
|> reformat_endpoints()
|> connect(type)
|> do_with_connection(fun)
end
@doc """
Retrieves the appropriate registry for the given value and validates it exists.
Executes the function with the registry name if it successfully locates one.
"""
@spec with_registry(atom() | String.t(), (atom() -> term())) :: term() | {:error, String.t()}
def with_registry(connection, function) when is_function(function, 1) do
registry = Elsa.ElsaSupervisor.registry(connection)
case Process.whereis(registry) do
nil -> {:error, "Elsa with connection #{connection} has not been started correctly"}
_pid -> function.(registry)
end
end
@doc """
Retrieves the pid of a brod client process if it exists and executes the
given function against the client.
"""
@spec with_client(atom(), (pid() -> term())) :: term() | {:error, String.t()}
def with_client(registry, function) when is_function(function, 1) do
case Elsa.ElsaRegistry.whereis_name({registry, :brod_client}) do
:undefined -> {:error, "Unable to find brod_client in registry(#{registry})"}
pid -> function.(pid)
end
end
@doc """
Convert supplied cluster endpoints from common keyword list format to
brod-compatible tuple.
"""
@spec reformat_endpoints(keyword()) :: [{charlist(), integer()}]
def reformat_endpoints(endpoints) do
Enum.map(endpoints, fn {key, value} -> {to_charlist(key), value} end)
end
@doc """
Retrieve the api version of the desired operation supported by the
connected cluster.
"""
@spec get_api_version(pid(), atom()) :: non_neg_integer()
def get_api_version(connection, api) do
{:ok, api_versions} = :kpro.get_api_versions(connection)
{_, version} = Map.get(api_versions, api)
version
end
@doc """
Determines if client pid is alive
"""
@spec client?(pid() | atom()) :: boolean()
def client?(pid) when is_pid(pid) do
Process.alive?(pid)
end
def client?(client) when is_atom(client) do
case Process.whereis(client) do
pid when is_pid(pid) -> client?(pid)
nil -> false
end
end
@doc """
Create a named client connection process for managing interactions
with the connected cluster.
"""
@spec start_client(keyword(), atom()) :: :ok | {:error, term()}
def start_client(endpoints, name, config \\ []) do
:brod.start_client(endpoints, name, config)
end
@spec chunk_by_byte_size(any(), integer()) :: list()
@doc """
Process messages into chunks of size up to the size specified by the calling function in bytes,
and determined by the function argument. If no chunk size is specified the default maximum
size a chunk will be is approximately 1 megabyte. If no sizing function is provided to construct
the appropriately sized chunks, the internal function based on Kernel.byte_size/1 is used.
"""
@spec chunk_by_byte_size(term(), integer(), fun()) :: [term()]
def chunk_by_byte_size(collection, chunk_byte_size \\ @default_max_chunk_size, byte_size_function \\ &get_byte_size/1) do
collection
|> Enum.chunk_while({0, []}, &chunk(&1, &2, chunk_byte_size, byte_size_function), &after_chunk/1)
end
@doc """
Return the number of partitions for a given topic. Bypasses the need for a persistent client
for lighter weight interactions from one-off calls.
Note that when passing in a connection/pid for the first argument, the underlying brod calls
will cache the results, including the non-existence of a topic! So if this function is being
used to wait for a recently created topic to be fully online, pass in a list of endpoints
rather than a connection. For convenience, the `get_endpoints/1` is provided to extract
the endpoints from a connection.
"""
@spec partition_count(list() | Elsa.connection() | pid(), String.t(), RetryConfig.t()) ::
{:ok, integer()} | {:error, any()}
def partition_count(connection_or_endpoints, topic, retry_config \\ RetryConfig.no_retry())
def partition_count(_client_or_endpoints, _topic, %RetryConfig{tries: 0}) do
{:error, :out_of_tries}
end
def partition_count(endpoints, topic, %RetryConfig{} = retry_config) when is_list(endpoints) do
case :brod.get_metadata(reformat_endpoints(endpoints), [topic]) do
{:error, reason} -> maybe_retry(reason, endpoints, topic, retry_config)
{:ok, metadata} -> {:ok, count_partitions(metadata)}
end
end
def partition_count(client, topic, %RetryConfig{} = retry_config) when is_atom(client) or is_pid(client) do
case :brod_client.get_partitions_count(client, topic) do
{:error, reason} -> maybe_retry(reason, client, topic, retry_config)
success -> success
end
end
@doc """
Like `partition_count/3` but returns only the partition count, and raises on failure
"""
@spec partition_count!(list() | Elsa.connection() | pid, String.t(), RetryConfig.t()) :: integer()
def partition_count!(connection_or_endpoints, topic, retry_config \\ RetryConfig.no_retry())
def partition_count!(connection_or_endpoints, topic, %RetryConfig{} = retry_config) do
case partition_count(connection_or_endpoints, topic, retry_config) do
{:ok, partition_count} ->
partition_count
{:error, error} ->
raise "#{__MODULE__}.partition_count! failed: #{error}"
end
end
@spec get_endpoints(Elsa.connection()) :: {:ok, list()} | {:error, any}
def get_endpoints(connection) do
case :brod_client.get_bootstrap(connection) do
{:ok, {endpoints, _}} -> {:ok, endpoints}
{:ok, endpoints} -> {:ok, endpoints}
error -> error
end
end
defp maybe_retry(reason, endpoints_or_client, topic, retry_config) do
Logger.info(
"#{__MODULE__}: error getting partition metadata for #{topic}: #{reason}. #{retry_config.tries - 1} tries remaining."
)
:timer.sleep(retry_config.dwell_ms)
partition_count(endpoints_or_client, topic, RetryConfig.decrement(retry_config))
end
# Handle brod < 3.16
defp count_partitions(%{topic_metadata: topic_metadatas}) do
[count | _] = for %{partition_metadata: metadata} <- topic_metadatas, do: Enum.count(metadata)
count
end
# Handle brod 3.16+
defp count_partitions(%{topics: topics}) do
[count | _] = for %{partitions: partitions} <- topics, do: Enum.count(partitions)
count
end
defp connect(endpoints, :controller), do: :kpro.connect_controller(endpoints, [])
defp connect(endpoints, _type), do: :kpro.connect_any(endpoints, [])
defp do_with_connection({:ok, connection}, fun) do
fun.(connection)
after
:kpro.close_connection(connection)
end
defp do_with_connection({:error, reason}, _fun) do
raise Elsa.ConnectError, message: format_reason(reason)
end
defp format_reason(reason) do
cond do
is_binary(reason) -> reason
Exception.exception?(reason) -> Exception.format(:error, reason)
true -> inspect(reason)
end
end
defp chunk(item, {current_size, current_batch}, chunk_byte_size, byte_size_function) do
item_size = byte_size_function.(item) + @timestamp_size_in_bytes
new_total = current_size + item_size
case new_total < chunk_byte_size do
true -> add_item_to_batch(new_total, item, current_batch)
false -> finish_batch(item_size, item, current_batch)
end
end
defp add_item_to_batch(total, item, batch) do
{:cont, {total, [item | batch]}}
end
defp finish_batch(total, item, batch) do
{:cont, Enum.reverse(batch), {total, [item]}}
end
defp after_chunk({_size, []}) do
{:cont, {0, []}}
end
defp after_chunk({_size, current_batch}) do
finish_batch(0, nil, current_batch)
end
defp get_byte_size(%{key: key, value: value} = msg) do
header_size =
Map.get(msg, :headers, [])
|> Enum.map(fn {key, value} -> byte_size(key) + byte_size(value) end)
|> Enum.sum()
byte_size(key) + byte_size(value) + header_size
end
end
