From dd384bc9462b66fc3226f23c8f45ed62ab947623 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Fri, 22 Nov 2024 02:51:40 +0100
Subject: [PATCH] Use fixed port in desktop app and fallback to random if taken
 (#2867)

Co-authored-by: Wojtek Mach <wojtekmach@users.noreply.github.com>
---
 lib/livebook/application.ex | 60 ++++++++++++++++++++++++++++++++++---
 rel/app/env.bat.eex         |  2 +-
 rel/app/env.sh.eex          |  2 +-
 3 files changed, 58 insertions(+), 6 deletions(-)

diff --git a/lib/livebook/application.ex b/lib/livebook/application.ex
index 97460707e..c20ac35c0 100644
--- a/lib/livebook/application.ex
+++ b/lib/livebook/application.ex
@@ -1,6 +1,8 @@
 defmodule Livebook.Application do
   use Application
 
+  require Logger
+
   def start(_type, _args) do
     Livebook.Utils.HTTP.set_proxy_options()
 
@@ -69,7 +71,7 @@ defmodule Livebook.Application do
             [
               {module, name: LivebookWeb.ZTA, identity_key: key},
               # We skip the access url as we do our own logging below
-              {LivebookWeb.Endpoint, log_access_url: false}
+              endpoint_childspec(log_access_url: false)
             ] ++ app_specs()
         end
 
@@ -180,12 +182,42 @@ defmodule Livebook.Application do
     end
   end
 
-  if Mix.target() == :app do
+  @app? Mix.target() == :app
+
+  if @app? do
     defp app_specs, do: [LivebookApp]
   else
     defp app_specs, do: []
   end
 
+  # In order to provide good first experience with the desktop app,
+  # in case the endpoint or iframe port is taken, we automatically
+  # fallback to a random port.
+
+  if @app? do
+    defp endpoint_childspec(opts) do
+      %{start: start} = childspec = LivebookWeb.Endpoint.child_spec(opts)
+      %{childspec | start: {__MODULE__, :endpoint_start, [start]}}
+    end
+  else
+    defp endpoint_childspec(opts), do: LivebookWeb.Endpoint.child_spec(opts)
+  end
+
+  @doc false
+  def endpoint_start({mod, fun, args}) do
+    with {:error,
+          {:shutdown,
+           {:failed_to_start_child, {LivebookWeb.Endpoint, :http},
+            {:shutdown, {:failed_to_start_child, :listener, :eaddrinuse}}}}} <-
+           apply(mod, fun, args) do
+      config = Application.get_env(:livebook, LivebookWeb.Endpoint)
+      config = put_in(config[:http][:port], 0)
+      Application.put_env(:livebook, LivebookWeb.Endpoint, config, persistent: true)
+      Logger.warning("Starting server using a random port")
+      endpoint_start({mod, fun, args})
+    end
+  end
+
   defp iframe_server_specs() do
     server? = Phoenix.Endpoint.server?(:livebook, LivebookWeb.Endpoint)
     port = Livebook.Config.iframe_port()
@@ -193,7 +225,7 @@ defmodule Livebook.Application do
     if server? do
       http = Application.fetch_env!(:livebook, LivebookWeb.Endpoint)[:http]
 
-      iframe_opts =
+      opts =
         [
           scheme: :http,
           plug: LivebookWeb.IframeEndpoint,
@@ -201,12 +233,32 @@ defmodule Livebook.Application do
           thousand_island_options: [supervisor_options: [name: LivebookWeb.IframeEndpoint]]
         ] ++ Keyword.take(http, [:ip])
 
-      [{Bandit, iframe_opts}]
+      [iframe_endpoint_childspec(opts)]
     else
       []
     end
   end
 
+  if @app? do
+    defp iframe_endpoint_childspec(opts) do
+      %{start: start} = childspec = Bandit.child_spec(opts)
+      %{childspec | start: {__MODULE__, :iframe_endpoint_start, [start]}}
+    end
+  else
+    defp iframe_endpoint_childspec(opts), do: Bandit.child_spec(opts)
+  end
+
+  @doc false
+  def iframe_endpoint_start({mod, fun, [opts]}) do
+    with {:error, {:shutdown, {:failed_to_start_child, :listener, :eaddrinuse}}} <-
+           apply(mod, fun, [opts]) do
+      Application.put_env(:livebook, :iframe_port, 0, persistent: true)
+      opts = Keyword.replace!(opts, :port, 0)
+      Logger.warning("Starting iframe server using a random port")
+      iframe_endpoint_start({mod, fun, [opts]})
+    end
+  end
+
   defp load_lb_env_vars() do
     secrets =
       for {"LB_" <> name = var, value} <- System.get_env() do
diff --git a/rel/app/env.bat.eex b/rel/app/env.bat.eex
index 99ca89f06..a5559fc23 100644
--- a/rel/app/env.bat.eex
+++ b/rel/app/env.bat.eex
@@ -9,7 +9,7 @@ set vendor_dir=!RELEASE_ROOT!\vendor\livebook-!RELEASE_VSN!
 set MIX_ARCHIVES=!vendor_dir!\archives
 set MIX_REBAR3=!vendor_dir!\rebar3
 if not defined LIVEBOOK_SHUTDOWN_ENABLED set LIVEBOOK_SHUTDOWN_ENABLED=true
-if not defined LIVEBOOK_PORT (set LIVEBOOK_PORT=0)
+if not defined LIVEBOOK_PORT (set LIVEBOOK_PORT=32123)
 set PATH=!vendor_dir!\otp\erts-<%= @release.erts_version%>\bin;!vendor_dir!\otp\bin;!vendor_dir!\elixir\bin;!PATH!
 
 if defined LIVEBOOK_NODE set RELEASE_NODE=!LIVEBOOK_NODE!
diff --git a/rel/app/env.sh.eex b/rel/app/env.sh.eex
index e0d79170b..6f8619d8b 100644
--- a/rel/app/env.sh.eex
+++ b/rel/app/env.sh.eex
@@ -9,7 +9,7 @@ vendor_dir="${RELEASE_ROOT}/vendor/livebook-${RELEASE_VSN}"
 export MIX_ARCHIVES="${vendor_dir}/archives"
 export MIX_REBAR3="${vendor_dir}/rebar3"
 export LIVEBOOK_SHUTDOWN_ENABLED=${LIVEBOOK_SHUTDOWN_ENABLED:-true}
-[ -z "$LIVEBOOK_PORT" ] && export LIVEBOOK_PORT=0
+[ -z "$LIVEBOOK_PORT" ] && export LIVEBOOK_PORT=32123
 export PATH="${vendor_dir}/otp/erts-<%= @release.erts_version%>/bin:${vendor_dir}/otp/bin:${vendor_dir}/elixir/bin:$PATH"
 
 if [ ! -z "${LIVEBOOK_NODE}" ]; then export RELEASE_NODE=${LIVEBOOK_NODE}; fi