Add Dockerfile variant for clustered deployment on Fly.io (#2286)

Co-authored-by: José Valim <jose.valim@dashbit.co>

README.md

@@ -212,6 +212,11 @@ The following environment variables can be used to configure Livebook on boot:
     Those certificates are used during for server authentication when Livebook
     accesses files from external sources.
 
+  * LIVEBOOK_CLUSTER - configures clustering strategy when running multiple
+    instances of Livebook. Currently the only supported value is `dns:QUERY`,
+    in which case nodes ask DNS for A/AAAA records using the given query and
+    try to connect to peer nodes on the discovered IPs.
+
   * LIVEBOOK_COOKIE - sets the cookie for running Livebook in a cluster.
     Defaults to a random string that is generated on boot.
 

lib/livebook.ex

@@ -212,6 +212,10 @@ defmodule Livebook do
            :identity_provider,
            Livebook.Config.identity_provider!("LIVEBOOK_IDENTITY_PROVIDER") ||
              {LivebookWeb.SessionIdentity, :unused}
+
+    if dns_cluster_query = Livebook.Config.dns_cluster_query!("LIVEBOOK_CLUSTER") do
+      config :livebook, :dns_cluster_query, dns_cluster_query
+    end
   end
 
   @doc """

lib/livebook/application.ex

@@ -49,6 +49,7 @@ defmodule Livebook.Application do
           iframe_server_specs() ++
             identity_provider() ++
             [
+              {DNSCluster, query: Application.get_env(:livebook, :dns_cluster_query) || :ignore},
               # Start the Endpoint (http/https)
               # We skip the access url as we do our own logging below
               {LivebookWeb.Endpoint, log_access_url: false}

lib/livebook/config.ex

@@ -628,6 +628,21 @@ defmodule Livebook.Config do
     end
   end
 
+  @doc """
+  Parses and validates DNS cluster query from env.
+  """
+  def dns_cluster_query!(env) do
+    if cluster_config = System.get_env(env) do
+      case cluster_config do
+        "dns:" <> query ->
+          query
+
+        other ->
+          abort!(~s{expected #{env} to be "dns:query", got: #{inspect(other)}})
+      end
+    end
+  end
+
   @app_version Mix.Project.config()[:version]
 
   @doc """

lib/livebook/hubs/dockerfile.ex

@@ -8,6 +8,7 @@ defmodule Livebook.Hubs.Dockerfile do
   @type config :: %{
           deploy_all: boolean(),
           docker_tag: String.t(),
+          clustering: nil | :fly_io,
           zta_provider: atom() | nil,
           zta_key: String.t() | nil
         }
@@ -19,7 +20,13 @@ defmodule Livebook.Hubs.Dockerfile do
   def config_changeset(attrs \\ %{}) do
     default_image = Livebook.Config.docker_images() |> hd()
 
-    data = %{deploy_all: false, docker_tag: default_image.tag, zta_provider: nil, zta_key: nil}
+    data = %{
+      deploy_all: false,
+      docker_tag: default_image.tag,
+      clustering: nil,
+      zta_provider: nil,
+      zta_key: nil
+    }
 
     zta_types =
       for provider <- Livebook.Config.identity_providers(),
@@ -29,11 +36,12 @@ defmodule Livebook.Hubs.Dockerfile do
     types = %{
       deploy_all: :boolean,
       docker_tag: :string,
+      clustering: Ecto.ParameterizedType.init(Ecto.Enum, values: [:fly_io]),
       zta_provider: Ecto.ParameterizedType.init(Ecto.Enum, values: zta_types),
       zta_key: :string
     }
 
-    cast({data, types}, attrs, [:deploy_all, :docker_tag, :zta_provider, :zta_key])
+    cast({data, types}, attrs, [:deploy_all, :docker_tag, :clustering, :zta_provider, :zta_key])
     |> validate_required([:deploy_all, :docker_tag])
   end
 
@@ -108,13 +116,31 @@ defmodule Livebook.Hubs.Dockerfile do
     RUN /app/bin/warmup_apps.sh
     """
 
+    startup =
+      if config.clustering == :fly_io do
+        ~S"""
+        # Custom startup script to cluster multiple Livebook nodes on Fly.io
+        RUN printf '\
+        #!/bin/bash\n\
+        export ERL_AFLAGS="-proto_dist inet6_tcp"\n\
+        export LIVEBOOK_DISTRIBUTION="name"\n\
+        export LIVEBOOK_NODE="${FLY_APP_NAME}-${FLY_IMAGE_REF##*-}@${FLY_PRIVATE_IP}"\n\
+        export LIVEBOOK_CLUSTER="dns:${FLY_APP_NAME}.internal"\n\
+        /app/bin/livebook start\n\
+        ' > /app/bin/start.sh && chmod +x /app/bin/start.sh
+
+        CMD [ "/app/bin/start.sh" ]
+        """
+      end
+
     [
       image,
       image_envs,
       hub_config,
       apps_config,
       notebook,
-      apps_warmup
+      apps_warmup,
+      startup
     ]
     |> Enum.reject(&is_nil/1)
     |> Enum.join("\n")

lib/livebook_web/live/app_helpers.ex

@@ -102,13 +102,41 @@ defmodule LivebookWeb.AppHelpers do
         ]}
       />
       <.radio_field label="Base image" field={@form[:docker_tag]} options={docker_tag_options()} />
+      <div class="grid grid-cols-1 md:grid-cols-2">
+        <.select_field
+          label="Clustering"
+          help={
+            ~S'''
+            When running multiple
+            instances of Livebook,
+            they need to be connected
+            into a single cluster.
+            You must either deploy
+            it as a single instance
+            or choose a platform to
+            enable clustering on.
+            '''
+          }
+          field={@form[:clustering]}
+          options={[
+            {"Single instance", ""},
+            {"Fly.io", "fly_io"}
+          ]}
+        />
+      </div>
       <%= if Hubs.Provider.type(@hub) == "team" do %>
         <div class="flex flex-col">
           <div class="grid grid-cols-1 md:grid-cols-2 gap-3">
             <.select_field
               label="Zero Trust Authentication provider"
               field={@form[:zta_provider]}
-              help="Enable this option if you want to deploy your notebooks behind an authentication proxy"
+              help={
+                ~S'''
+                Enable this option if you want
+                to deploy your notebooks behind
+                an authentication proxy
+                '''
+              }
               prompt="None"
               options={zta_options()}
             />

mix.exs

@@ -110,6 +110,7 @@ defmodule Livebook.MixProject do
       {:aws_signature, "~> 0.3.0"},
       {:mint_web_socket, "~> 1.0.0"},
       {:protobuf, "~> 0.8.0"},
+      {:dns_cluster, "~> 0.1.1"},
       {:phoenix_live_reload, "~> 1.2", only: :dev},
       {:floki, ">= 0.27.0", only: :test},
       {:bypass, "~> 2.1", only: :test},

mix.lock

@@ -7,6 +7,7 @@
   "cowboy_telemetry": {:hex, :cowboy_telemetry, "0.4.0", "f239f68b588efa7707abce16a84d0d2acf3a0f50571f8bb7f56a15865aae820c", [:rebar3], [{:cowboy, "~> 2.7", [hex: :cowboy, repo: "hexpm", optional: false]}, {:telemetry, "~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "7d98bac1ee4565d31b62d59f8823dfd8356a169e7fcbb83831b8a5397404c9de"},
   "cowlib": {:hex, :cowlib, "2.12.1", "a9fa9a625f1d2025fe6b462cb865881329b5caff8f1854d1cbc9f9533f00e1e1", [:make, :rebar3], [], "hexpm", "163b73f6367a7341b33c794c4e88e7dbfe6498ac42dcd69ef44c5bc5507c8db0"},
   "decimal": {:hex, :decimal, "2.1.1", "5611dca5d4b2c3dd497dec8f68751f1f1a54755e8ed2a966c2633cf885973ad6", [:mix], [], "hexpm", "53cfe5f497ed0e7771ae1a475575603d77425099ba5faef9394932b35020ffcc"},
+  "dns_cluster": {:hex, :dns_cluster, "0.1.1", "73b4b2c3ec692f8a64276c43f8c929733a9ab9ac48c34e4c0b3d9d1b5cd69155", [:mix], [], "hexpm", "03a3f6ff16dcbb53e219b99c7af6aab29eb6b88acf80164b4bd76ac18dc890b3"},
   "earmark_parser": {:hex, :earmark_parser, "1.4.37", "2ad73550e27c8946648b06905a57e4d454e4d7229c2dafa72a0348c99d8be5f7", [:mix], [], "hexpm", "6b19783f2802f039806f375610faa22da130b8edc21209d0bff47918bb48360e"},
   "ecto": {:hex, :ecto, "3.10.3", "eb2ae2eecd210b4eb8bece1217b297ad4ff824b4384c0e3fdd28aaf96edd6135", [:mix], [{:decimal, "~> 1.6 or ~> 2.0", [hex: :decimal, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: true]}, {:telemetry, "~> 0.4 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "44bec74e2364d491d70f7e42cd0d690922659d329f6465e89feb8a34e8cd3433"},
   "ex_doc": {:hex, :ex_doc, "0.30.6", "5f8b54854b240a2b55c9734c4b1d0dd7bdd41f71a095d42a70445c03cf05a281", [:mix], [{:earmark_parser, "~> 1.4.31", [hex: :earmark_parser, repo: "hexpm", optional: false]}, {:makeup_elixir, "~> 0.14", [hex: :makeup_elixir, repo: "hexpm", optional: false]}, {:makeup_erlang, "~> 0.1", [hex: :makeup_erlang, repo: "hexpm", optional: false]}], "hexpm", "bd48f2ddacf4e482c727f9293d9498e0881597eae6ddc3d9562bd7923375109f"},

test/livebook/hubs/dockerfile_test.exs

@@ -195,6 +195,16 @@ defmodule Livebook.Hubs.DockerfileTest do
                COPY files/data.csv files/image.jpeg /apps/files/
                """
     end
+
+    test "deploying with fly.io cluster setup" do
+      config = dockerfile_config(%{clustering: :fly_io})
+      hub = personal_hub()
+      file = Livebook.FileSystem.File.local(p("/notebook.livemd"))
+
+      dockerfile = Dockerfile.build_dockerfile(config, hub, [], [], file, [], %{})
+
+      assert dockerfile =~ ~s/export LIVEBOOK_CLUSTER="dns:${FLY_APP_NAME}.internal"/
+    end
   end
 
   describe "warnings/6" do