From 52e5273401b2d654f888ea249fb2e6517d30d309 Mon Sep 17 00:00:00 2001
From: Frank Hunleth <fhunleth@troodon-software.com>
Date: Thu, 10 Nov 2022 12:09:01 -0500
Subject: [PATCH] Support passing an MFA for customizing shutdown (#1519)

This commit lets you customize Livebook shutdown to gracefully power
off. One use case is for Nerves-based devices that need to do more than
call `System.stop/0` to power off.

This doesn't change the `LIVEBOOK_SHUTDOWN_ENABLED` environment
variable. It still works the same.

The `:shutdown_enabled` configuration is now `:shutdown_callback`. Valid
values are `nil` or an MFA. An unset or `nil` callback hides the
shutdown button in the UI.
---
 config/config.exs                           |  2 +-
 config/dev.exs                              |  2 +-
 lib/livebook.ex                             |  2 +-
 lib/livebook/config.ex                      |  8 ++++----
 lib/livebook_web/live/hooks/sidebar_hook.ex | 13 ++++++++-----
 lib/livebook_web/live/layout_helpers.ex     |  2 +-
 6 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/config/config.exs b/config/config.exs
index 595f4bf85..e39e1d291 100644
--- a/config/config.exs
+++ b/config/config.exs
@@ -27,7 +27,7 @@ config :livebook,
   force_ssl_host: nil,
   learn_notebooks: [],
   plugs: [],
-  shutdown_enabled: false,
+  shutdown_callback: nil,
   storage: Livebook.Storage.Ets,
   update_instructions_url: nil,
   within_iframe: false
diff --git a/config/dev.exs b/config/dev.exs
index f98d4496c..e0fa17b2c 100644
--- a/config/dev.exs
+++ b/config/dev.exs
@@ -24,7 +24,7 @@ config :livebook, LivebookWeb.Endpoint,
   ]
 
 config :livebook, :iframe_port, 4001
-config :livebook, :shutdown_enabled, true
+config :livebook, :shutdown_callback, {System, :stop, []}
 
 # Feature flags
 config :livebook, :feature_flags,
diff --git a/lib/livebook.ex b/lib/livebook.ex
index 9903853c7..e2ada4000 100644
--- a/lib/livebook.ex
+++ b/lib/livebook.ex
@@ -119,7 +119,7 @@ defmodule Livebook do
     end
 
     if Livebook.Config.boolean!("LIVEBOOK_SHUTDOWN_ENABLED", false) do
-      config :livebook, :shutdown_enabled, true
+      config :livebook, :shutdown_callback, {System, :stop, []}
     end
 
     if Livebook.Config.boolean!("LIVEBOOK_WITHIN_IFRAME", false) do
diff --git a/lib/livebook/config.ex b/lib/livebook/config.ex
index efbb8c61c..1e4f1edc2 100644
--- a/lib/livebook/config.ex
+++ b/lib/livebook/config.ex
@@ -107,11 +107,11 @@ defmodule Livebook.Config do
   end
 
   @doc """
-  Returns whether the shutdown feature is enabled.
+  Returns an mfa if there's a way to shut down the system.
   """
-  @spec shutdown_enabled?() :: boolean()
-  def shutdown_enabled?() do
-    Application.fetch_env!(:livebook, :shutdown_enabled)
+  @spec shutdown_callback() :: mfa() | nil
+  def shutdown_callback() do
+    Application.fetch_env!(:livebook, :shutdown_callback)
   end
 
   @doc """
diff --git a/lib/livebook_web/live/hooks/sidebar_hook.ex b/lib/livebook_web/live/hooks/sidebar_hook.ex
index ee42a4db6..ada8a4a71 100644
--- a/lib/livebook_web/live/hooks/sidebar_hook.ex
+++ b/lib/livebook_web/live/hooks/sidebar_hook.ex
@@ -23,11 +23,14 @@ defmodule LivebookWeb.SidebarHook do
   defp handle_info(_event, socket), do: {:cont, socket}
 
   defp handle_event("shutdown", _params, socket) do
-    if Livebook.Config.shutdown_enabled?() do
-      System.stop()
-      {:halt, put_flash(socket, :info, "Livebook is shutting down. You can close this page.")}
-    else
-      socket
+    case Livebook.Config.shutdown_callback() do
+      {m, f, a} ->
+        apply(m, f, a)
+
+        {:halt, put_flash(socket, :info, "Livebook is shutting down. You can close this page.")}
+
+      _ ->
+        socket
     end
   end
 
diff --git a/lib/livebook_web/live/layout_helpers.ex b/lib/livebook_web/live/layout_helpers.ex
index e57018a4e..d6c7d40f6 100644
--- a/lib/livebook_web/live/layout_helpers.ex
+++ b/lib/livebook_web/live/layout_helpers.ex
@@ -113,7 +113,7 @@ defmodule LivebookWeb.LayoutHelpers do
           <.hub_section socket={@socket} hubs={@saved_hubs} current_page={@current_page} />
         </div>
         <div class="flex flex-col">
-          <%= if Livebook.Config.shutdown_enabled?() do %>
+          <%= if Livebook.Config.shutdown_callback() do %>
             <button
               class="h-7 flex items-center text-gray-400 hover:text-white border-l-4 border-transparent hover:border-white"
               aria-label="shutdown"