MuonTrap.Daemon new features

lib/muontrap/daemon.ex

@@ -23,23 +23,24 @@ defmodule MuonTrap.Daemon do
   is a list of options. The same options as `MuonTrap.cmd/3` are available with
   the following additions:
 
-  * `:name` - Name the Daemon GenServer
-  * `:log_output` - When set, send output from the command to the Logger. Specify the log level (e.g., `:debug`)
-  * `:log_prefix` - Prefix each log message with this string (defaults to the program's path)
+  * `:name` - Name the Daemon GenServer.
+  * `:msg_callback` - When set, it sends the output from the command to the callback. Only funtions with /1 arity are allowed.
+  * `:log_output` - When set, it sends the output from the command to the Logger. Specify the log level (e.g., `:debug`).
+  * `:log_prefix` - Prefix each log message with this string (defaults to the program's path).
   * `:stderr_to_stdout` - When set to `true`, redirect stderr to stdout. Defaults to `false`.
 
   If you want to run multiple `MuonTrap.Daemon`s under one supervisor, they'll
   all need unique IDs. Use `Supervisor.child_spec/2` like this:
 
   ```elixir
-  Supervisor.child_spec({MuonTrap.Daemon, ["my_server"), []]}, id: :server1)
+  Supervisor.child_spec({MuonTrap.Daemon, ["my_server", []]}, id: :server1)
   ```
   """
 
   defmodule State do
     @moduledoc false
 
-    defstruct [:command, :port, :cgroup_path, :log_output, :log_prefix]
+    defstruct [:command, :port, :cgroup_path, :log_output, :log_prefix, :msg_callback]
   end
 
   def child_spec([command, args]) do
@@ -107,6 +108,7 @@ defmodule MuonTrap.Daemon do
        command: command,
        port: port,
        cgroup_path: Map.get(options, :cgroup_path),
+       msg_callback: Map.get(options, :msg_callback),
        log_output: Map.get(options, :log_output),
        log_prefix: Map.get(options, :log_prefix, command <> ": ")
      }}
@@ -139,17 +141,22 @@ defmodule MuonTrap.Daemon do
   end
 
   @impl true
-  def handle_info({_port, {:data, _}}, %State{log_output: nil} = state) do
-    # Ignore output
+  def handle_info(
+        {port, {:data, {_, message}}},
+        %State{port: port, log_output: nil, msg_callback: msg_callback} = state
+      ) do
+    dispatch_message(msg_callback, message)
     {:noreply, state}
   end
 
   @impl true
   def handle_info(
         {port, {:data, {_, message}}},
-        %State{port: port, log_output: log_level, log_prefix: prefix} = state
+        %State{port: port, log_output: log_level, log_prefix: prefix, msg_callback: msg_callback} =
+          state
       ) do
     Logger.log(log_level, [prefix, message])
+    dispatch_message(msg_callback, message)
     {:noreply, state}
   end
 
@@ -168,4 +175,13 @@ defmodule MuonTrap.Daemon do
 
     {:stop, reason, state}
   end
+
+  @impl true
+  def handle_info(unhandled_msg, state) do
+    Logger.warn("Unhandled message: #{inspect(unhandled_msg)}")
+    {:noreply, state}
+  end
+
+  defp dispatch_message(nil, _message), do: :ok
+  defp dispatch_message(msg_callback, message), do: msg_callback.(message)
 end

lib/muontrap/options.ex

@@ -22,6 +22,7 @@ defmodule MuonTrap.Options do
   * `:parallelism`
   * `:env`
   * `:name` - `MuonTrap.Daemon`-only
+  * `:msg_callback` - `MuonTrap.Daemon`-only
   * `:log_output` - `MuonTrap.Daemon`-only
   * `:log_prefix` - `MuonTrap.Daemon`-only
   * `:cgroup_controllers`
@@ -106,6 +107,18 @@ defmodule MuonTrap.Options do
   defp validate_option(:daemon, {:log_prefix, prefix}, opts) when is_binary(prefix),
     do: Map.put(opts, :log_prefix, prefix)
 
+  defp validate_option(:daemon, {:msg_callback, nil}, opts), do: opts
+
+  defp validate_option(:daemon, {:msg_callback, function}, opts) when is_function(function) do
+    with function_info <- Function.info(function),
+         true <- function_info[:arity] == 1 do
+      Map.put(opts, :msg_callback, function)
+    else
+      _arity_match_error ->
+        raise(ArgumentError, "Invalid :msg_callback, only functions with /1 arity are allowed")
+    end
+  end
+
   # MuonTrap common options
   defp validate_option(_any, {:cgroup_controllers, controllers}, opts) when is_list(controllers),
     do: Map.put(opts, :cgroup_controllers, controllers)

test/daemon_test.exs

@@ -77,6 +77,31 @@ defmodule DaemonTest do
     assert capture_log(fun) =~ "echo says: hello"
   end
 
+  test "daemon logs unhandled messages" do
+    fun = fn ->
+      {:ok, _pid} = start_supervised(daemon_spec("echo", ["hello"], name: UnhandledMsg))
+
+      send(UnhandledMsg, "this is an unhandled msg")
+
+      wait_for_close_check()
+      Logger.flush()
+    end
+
+    assert capture_log(fun) =~ "Unhandled message: \"this is an unhandled msg\""
+  end
+
+  test "daemon dispatch the message to msg_callback" do
+    fun = fn ->
+      {:ok, _pid} =
+        start_supervised(daemon_spec("echo", ["hello"], msg_callback: &msg_test_callback/1))
+
+      wait_for_close_check()
+      Logger.flush()
+    end
+
+    assert capture_log(fun) =~ "msg_callback echo says: hello"
+  end
+
   test "can pass environment variables to the daemon" do
     fun = fn ->
       {:ok, _pid} =
@@ -175,4 +200,9 @@ defmodule DaemonTest do
     memory = Integer.parse(memory_str)
     assert memory > 1000
   end
+
+  def msg_test_callback(msg) do
+    require Logger
+    Logger.log(:info, ["msg_callback echo says: ", msg])
+  end
 end

test/options_test.exs

@@ -59,6 +59,24 @@ defmodule MuonTrap.OptionsTest do
     assert_raise ArgumentError, fn ->
       Options.validate(:daemon, "echo", [], log_output: :bad_level)
     end
+
+    assert_raise ArgumentError, fn ->
+      Options.validate(:daemon, "echo", [], msg_callback: false)
+    end
+
+    raise_msg = "Invalid :msg_callback, only functions with /1 arity are allowed"
+
+    assert_raise ArgumentError, raise_msg, fn ->
+      Options.validate(:daemon, "echo", [], msg_callback: &Kernel.+/2)
+    end
+
+    :daemon
+    |> Options.validate("echo", [], msg_callback: &inspect/1)
+    |> Map.get(:msg_callback)
+    |> Kernel.==(&inspect/1)
+    |> assert()
+
+    assert Map.get(Options.validate(:daemon, "echo", [], msg_callback: nil), :msg_callback) == nil
   end
 
   test "common commands basically work" do