Handle when user stops the App Deployment (#2554)

lib/livebook/hubs/team_client.ex

@@ -336,9 +336,7 @@ defmodule Livebook.Hubs.TeamClient do
   defp remove_deployment_group(state, deployment_group) do
     %{
       state
-      | deployment_groups: Enum.reject(state.deployment_groups, &(&1.id == deployment_group.id)),
+      | deployment_groups: Enum.reject(state.deployment_groups, &(&1.id == deployment_group.id))
-        app_deployments:
-          Enum.reject(state.app_deployments, &(&1.deployment_group_id == deployment_group.id))
     }
   end
 
@@ -554,6 +552,7 @@ defmodule Livebook.Hubs.TeamClient do
     |> dispatch_deployment_groups(user_connected)
     |> dispatch_app_deployments(user_connected)
     |> dispatch_agents(user_connected)
+    |> dispatch_connection()
   end
 
   defp handle_event(:agent_connected, agent_connected, state) do
@@ -564,31 +563,49 @@ defmodule Livebook.Hubs.TeamClient do
     |> dispatch_deployment_groups(agent_connected)
     |> dispatch_app_deployments(agent_connected)
     |> dispatch_agents(agent_connected)
+    |> dispatch_connection()
   end
 
-  defp handle_event(:app_deployment_created, %Teams.AppDeployment{} = app_deployment, state) do
+  defp handle_event(:app_deployment_started, %Teams.AppDeployment{} = app_deployment, state) do
     deployment_group_id = app_deployment.deployment_group_id
 
     with {:ok, deployment_group} <- fetch_deployment_group(deployment_group_id, state) do
-      Teams.Broadcasts.app_deployment_created(app_deployment)
-      state = put_app_deployment(state, app_deployment)
-
       if deployment_group.id == state.deployment_group_id do
         manager_sync()
       end
-
-      state
-    end
     end
 
-  defp handle_event(:app_deployment_created, app_deployment_created, state) do
+    Teams.Broadcasts.app_deployment_started(app_deployment)
+    put_app_deployment(state, app_deployment)
+  end
+
+  defp handle_event(:app_deployment_started, app_deployment_started, state) do
     handle_event(
-      :app_deployment_created,
+      :app_deployment_started,
-      build_app_deployment(state, app_deployment_created.app_deployment),
+      build_app_deployment(state, app_deployment_started.app_deployment),
       state
     )
   end
 
+  defp handle_event(:app_deployment_stopped, %Teams.AppDeployment{} = app_deployment, state) do
+    deployment_group_id = app_deployment.deployment_group_id
+
+    with {:ok, deployment_group} <- fetch_deployment_group(deployment_group_id, state) do
+      if deployment_group.id == state.deployment_group_id do
+        manager_sync()
+      end
+    end
+
+    Teams.Broadcasts.app_deployment_stopped(app_deployment)
+    remove_app_deployment(state, app_deployment)
+  end
+
+  defp handle_event(:app_deployment_stopped, %{id: id}, state) do
+    with {:ok, app_deployment} <- fetch_app_deployment(id, state) do
+      handle_event(:app_deployment_stopped, app_deployment, state)
+    end
+  end
+
   defp handle_event(:agent_joined, %Teams.Agent{} = agent, state) do
     Teams.Broadcasts.agent_joined(agent)
     put_agent(state, agent)
@@ -660,9 +677,9 @@ defmodule Livebook.Hubs.TeamClient do
 
   defp dispatch_app_deployments(state, %{app_deployments: app_deployments}) do
     app_deployments = Enum.map(app_deployments, &build_app_deployment(state, &1))
-    {created, _, _} = diff(state.app_deployments, app_deployments, &(&1.id == &2.id))
+    {started, stopped, _} = diff(state.app_deployments, app_deployments, &(&1.id == &2.id))
 
-    dispatch_events(state, app_deployment_created: created)
+    dispatch_events(state, app_deployment_started: started, app_deployment_stopped: stopped)
   end
 
   defp dispatch_agents(state, %{agents: agents}) do
@@ -672,6 +689,11 @@ defmodule Livebook.Hubs.TeamClient do
     dispatch_events(state, agent_joined: joined, agent_left: left)
   end
 
+  defp dispatch_connection(%{hub: %{id: id}} = state) do
+    Teams.Broadcasts.client_connected(id)
+    state
+  end
+
   defp update_hub(state, %{public_key: org_public_key}) do
     hub = %{state.hub | org_public_key: org_public_key}
 
@@ -713,6 +735,9 @@ defmodule Livebook.Hubs.TeamClient do
 
   defp fetch_agent(id, state), do: fetch_entry(state.agents, &(&1.id == id), state)
 
+  defp fetch_app_deployment(id, state),
+    do: fetch_entry(state.app_deployments, &(&1.id == id), state)
+
   defp fetch_entry(entries, fun, state) do
     if entry = Enum.find(entries, fun) do
       {:ok, entry}

lib/livebook/teams/broadcasts.ex

@@ -3,30 +3,36 @@ defmodule Livebook.Teams.Broadcasts do
 
   @type broadcast :: :ok | {:error, term()}
 
-  @deployment_groups_topic "teams:deployment_groups"
-  @app_deployments_topic "teams:app_deployments"
   @agents_topic "teams:agents"
+  @app_deployments_topic "teams:app_deployments"
+  @clients_topic "teams:clients"
+  @deployment_groups_topic "teams:deployment_groups"
 
   @doc """
   Subscribes to one or more subtopics in `"teams"`.
 
   ## Messages
 
-  Topic `#{@deployment_groups_topic}`:
-
-    * `{:deployment_group_created, DeploymentGroup.t()}`
-    * `{:deployment_group_updated, DeploymentGroup.t()}`
-    * `{:deployment_group_deleted, DeploymentGroup.t()}`
-
-  Topic `#{@app_deployments_topic}`:
-
-    * `{:app_deployment_created, AppDeployment.t()}`
-
   Topic `#{@agents_topic}`:
 
     * `{:agent_joined, Agent.t()}`
     * `{:agent_left, Agent.t()}`
 
+  Topic `#{@app_deployments_topic}`:
+
+    * `{:app_deployment_started, AppDeployment.t()}`
+    * `{:app_deployment_stopped, AppDeployment.t()}`
+
+  Topic `#{@clients_topic}`:
+
+    * `{:client_connected, id}`
+
+  Topic `#{@deployment_groups_topic}`:
+
+    * `{:deployment_group_created, DeploymentGroup.t()}`
+    * `{:deployment_group_update, DeploymentGroup.t()}`
+    * `{:deployment_group_deleted, DeploymentGroup.t()}`
+
   """
   @spec subscribe(atom() | list(atom())) :: :ok | {:error, term()}
   def subscribe(topics) when is_list(topics) do
@@ -53,6 +59,14 @@ defmodule Livebook.Teams.Broadcasts do
     Phoenix.PubSub.unsubscribe(Livebook.PubSub, "teams:#{topic}")
   end
 
+  @doc """
+  Broadcasts under `#{@clients_topic}` topic when hub received a new client connection.
+  """
+  @spec client_connected(String.t()) :: broadcast()
+  def client_connected(id) do
+    broadcast(@clients_topic, {:client_connected, id})
+  end
+
   @doc """
   Broadcasts under `#{@deployment_groups_topic}` topic when hub received a new deployment group.
   """
@@ -78,11 +92,19 @@ defmodule Livebook.Teams.Broadcasts do
   end
 
   @doc """
-  Broadcasts under `#{@app_deployments_topic}` topic when hub received a new app deployment.
+  Broadcasts under `#{@app_deployments_topic}` topic when hub received to start a new app deployment.
   """
-  @spec app_deployment_created(AppDeployment.t()) :: broadcast()
+  @spec app_deployment_started(AppDeployment.t()) :: broadcast()
-  def app_deployment_created(%AppDeployment{} = app_deployment) do
+  def app_deployment_started(%AppDeployment{} = app_deployment) do
-    broadcast(@app_deployments_topic, {:app_deployment_created, app_deployment})
+    broadcast(@app_deployments_topic, {:app_deployment_started, app_deployment})
+  end
+
+  @doc """
+  Broadcasts under `#{@app_deployments_topic}` topic when hub received to stop an app deployment.
+  """
+  @spec app_deployment_stopped(AppDeployment.t()) :: broadcast()
+  def app_deployment_stopped(%AppDeployment{} = app_deployment) do
+    broadcast(@app_deployments_topic, {:app_deployment_stopped, app_deployment})
   end
 
   @doc """

lib/livebook_web/live/hub/teams/deployment_group_component.ex

@@ -65,7 +65,9 @@ defmodule LivebookWeb.Hub.Teams.DeploymentGroupComponent do
             </.link>
           </.labeled_text>
           <.labeled_text class="grow mt-6 lg:border-l lg:pl-4" label="Apps deployed">
-            <span class="text-lg font-normal"><%= @app_deployments_count %></span>
+            <span class="text-lg font-normal" aria-label="apps deployed">
+              <%= @app_deployments_count %>
+            </span>
             <.link
               patch={~p"/hub/#{@hub.id}/groups/#{@deployment_group.id}/apps/new"}
               class="pl-2 text-blue-600"

proto/lib/livebook_proto.ex

@@ -8,7 +8,8 @@ defmodule LivebookProto do
           LivebookProto.AgentConnected.t()
           | LivebookProto.AgentJoined.t()
           | LivebookProto.AgentLeft.t()
-          | LivebookProto.AppDeploymentCreated.t()
+          | LivebookProto.AppDeploymentStarted.t()
+          | LivebookProto.AppDeploymentStopped.t()
           | LivebookProto.FileSystemCreated.t()
           | LivebookProto.FileSystemDeleted.t()
           | LivebookProto.FileSystemUpdated.t()

proto/lib/livebook_proto/app_deployment_started.pb.ex

@@ -1,4 +1,4 @@
-defmodule LivebookProto.AppDeploymentCreated do
+defmodule LivebookProto.AppDeploymentStarted do
   use Protobuf, syntax: :proto3, protoc_gen_elixir_version: "0.12.0"
 
   field :app_deployment, 1, type: LivebookProto.AppDeployment, json_name: "appDeployment"

proto/lib/livebook_proto/app_deployment_stopped.pb.ex

@@ -0,0 +1,5 @@
+defmodule LivebookProto.AppDeploymentStopped do
+  use Protobuf, syntax: :proto3, protoc_gen_elixir_version: "0.12.0"
+
+  field :id, 1, type: :string
+end

proto/lib/livebook_proto/event.pb.ex

@@ -58,12 +58,17 @@ defmodule LivebookProto.Event do
     json_name: "agentConnected",
     oneof: 0
 
-  field :app_deployment_created, 12,
+  field :app_deployment_started, 12,
-    type: LivebookProto.AppDeploymentCreated,
+    type: LivebookProto.AppDeploymentStarted,
-    json_name: "appDeploymentCreated",
+    json_name: "appDeploymentStarted",
     oneof: 0
 
   field :user_deleted, 13, type: LivebookProto.UserDeleted, json_name: "userDeleted", oneof: 0
   field :agent_joined, 14, type: LivebookProto.AgentJoined, json_name: "agentJoined", oneof: 0
   field :agent_left, 15, type: LivebookProto.AgentLeft, json_name: "agentLeft", oneof: 0
+
+  field :app_deployment_stopped, 16,
+    type: LivebookProto.AppDeploymentStopped,
+    json_name: "appDeploymentStopped",
+    oneof: 0
 end

proto/messages.proto

@@ -126,10 +126,14 @@ message AppDeployment {
   int64 deployed_at = 8;
 }
 
-message AppDeploymentCreated {
+message AppDeploymentStarted {
   AppDeployment app_deployment = 1;
 }
 
+message AppDeploymentStopped {
+  string id = 1;
+}
+
 message UserDeleted {
   string id = 1;
 }
@@ -162,9 +166,10 @@ message Event {
     DeploymentGroupUpdated deployment_group_updated = 9;
     DeploymentGroupDeleted deployment_group_deleted = 10;
     AgentConnected agent_connected = 11;
-    AppDeploymentCreated app_deployment_created = 12;
+    AppDeploymentStarted app_deployment_started = 12;
     UserDeleted user_deleted = 13;
     AgentJoined agent_joined = 14;
     AgentLeft agent_left = 15;
+    AppDeploymentStopped app_deployment_stopped = 16;
   }
 }

test/livebook_teams/hubs/team_client_test.exs

@@ -5,7 +5,7 @@ defmodule Livebook.Hubs.TeamClientTest do
 
   setup do
     Livebook.Hubs.Broadcasts.subscribe([:connection, :file_systems, :secrets])
-    Livebook.Teams.Broadcasts.subscribe([:deployment_groups, :app_deployments, :agents])
+    Livebook.Teams.Broadcasts.subscribe([:clients, :deployment_groups, :app_deployments, :agents])
     :ok
   end
 
@@ -110,7 +110,7 @@ defmodule Livebook.Hubs.TeamClientTest do
     end
 
     @tag :tmp_dir
-    test "receives app events", %{team: team, tmp_dir: tmp_dir} do
+    test "receives app events", %{team: team, node: node, tmp_dir: tmp_dir} do
       deployment_group = build(:deployment_group, name: team.id, mode: :online)
       assert {:ok, id} = Livebook.Teams.create_deployment_group(team, deployment_group)
 
@@ -136,13 +136,18 @@ defmodule Livebook.Hubs.TeamClientTest do
 
       sha = app_deployment.sha
 
-      assert_receive {:app_deployment_created,
+      assert_receive {:app_deployment_started,
                       %Livebook.Teams.AppDeployment{
                         slug: ^slug,
                         sha: ^sha,
                         title: ^title,
                         deployment_group_id: ^id
-                      }}
+                      } = app_deployment}
+
+      # force app deployment to be deleted
+      erpc_call(node, :stop_app_deployment, [app_deployment.id, team.org_id])
+
+      assert_receive {:app_deployment_stopped, ^app_deployment}
     end
   end
 
@@ -313,6 +318,7 @@ defmodule Livebook.Hubs.TeamClientTest do
     end
 
     test "dispatches the app deployments list", %{team: team, user_connected: user_connected} do
+      hub_id = team.id
       pid = connect_to_teams(team)
 
       deployment_group =
@@ -360,8 +366,16 @@ defmodule Livebook.Hubs.TeamClientTest do
       }
 
       send(pid, {:event, :user_connected, user_connected})
-      assert_receive {:app_deployment_created, ^app_deployment}, 5000
+      assert_receive {:client_connected, ^hub_id}
+      assert_receive {:app_deployment_started, ^app_deployment}, 5000
       assert app_deployment in TeamClient.get_app_deployments(team.id)
+
+      # deletes the app deployment
+      user_connected = %{user_connected | app_deployments: []}
+      send(pid, {:event, :user_connected, user_connected})
+      assert_receive {:client_connected, ^hub_id}
+      assert_receive {:app_deployment_stopped, ^app_deployment}
+      refute app_deployment in TeamClient.get_app_deployments(team.id)
     end
 
     test "dispatches the agents list", %{team: team, user_connected: user_connected} do
@@ -515,21 +529,41 @@ defmodule Livebook.Hubs.TeamClientTest do
     end
 
     test "dispatches the deployment groups list",
-         %{team: team, deployment_group: teams_deployment_group, agent_connected: agent_connected} do
+         %{
+           team: team,
+           deployment_group: teams_deployment_group,
+           agent_key: teams_agent_key,
+           agent_connected: agent_connected
+         } do
+      agent_key =
+        build(:agent_key,
+          id: to_string(teams_agent_key.id),
+          key: teams_agent_key.key,
+          deployment_group_id: to_string(teams_agent_key.deployment_group_id)
+        )
+
       deployment_group =
         build(:deployment_group,
           id: to_string(teams_deployment_group.id),
           name: teams_deployment_group.name,
           mode: teams_deployment_group.mode,
           hub_id: team.id,
-          secrets: []
+          agent_keys: [agent_key]
         )
 
+      livebook_proto_agent_key =
+        %LivebookProto.AgentKey{
+          id: agent_key.id,
+          key: agent_key.key,
+          deployment_group_id: agent_key.deployment_group_id
+        }
+
       livebook_proto_deployment_group =
         %LivebookProto.DeploymentGroup{
           id: to_string(deployment_group.id),
           name: deployment_group.name,
           mode: to_string(deployment_group.mode),
+          agent_keys: [livebook_proto_agent_key],
           secrets: []
         }
 
@@ -743,16 +777,29 @@ defmodule Livebook.Hubs.TeamClientTest do
       Livebook.Apps.subscribe()
 
       send(pid, {:event, :agent_connected, agent_connected})
-      assert_receive {:app_deployment_created, ^app_deployment}
+      assert_receive {:app_deployment_started, ^app_deployment}
-      assert_receive {:app_created, %{pid: app_pid, slug: ^slug}}
+      assert_receive {:app_created, %{slug: ^slug}}
 
       assert_receive {:app_updated,
-                      %{slug: ^slug, sessions: [%{app_status: %{execution: :executed}}]}}
+                      %{
+                        slug: ^slug,
+                        warnings: [],
+                        sessions: [%{app_status: %{execution: :executed}}]
+                      }}
 
       assert app_deployment in TeamClient.get_app_deployments(team.id)
 
-      Livebook.Hubs.delete_hub(team.id)
+      agent_connected = %{agent_connected | app_deployments: []}
-      Livebook.App.close(app_pid)
+      send(pid, {:event, :agent_connected, agent_connected})
+      assert_receive {:app_deployment_stopped, ^app_deployment}
+      refute app_deployment in TeamClient.get_app_deployments(team.id)
+
+      assert_receive {:app_closed,
+                      %{
+                        slug: ^slug,
+                        warnings: [],
+                        sessions: [%{app_status: %{execution: :executed}}]
+                      }}
     end
 
     test "dispatches the agents list", %{team: team, agent_connected: agent_connected} do
@@ -789,6 +836,8 @@ defmodule Livebook.Hubs.TeamClientTest do
   defp connect_to_teams(%{id: id} = team) do
     Livebook.Hubs.save_hub(team)
     assert_receive {:hub_connected, ^id}
+    assert_receive {:client_connected, ^id}
+
     TeamClient.get_pid(team.id)
   end
 end

test/livebook_teams/web/hub/deployment_group_test.exs

@@ -8,11 +8,12 @@ defmodule LivebookWeb.Integration.Hub.DeploymentGroupTest do
 
   setup %{user: user, node: node} do
     Livebook.Hubs.Broadcasts.subscribe([:crud, :connection, :secrets, :file_systems])
-    Livebook.Teams.Broadcasts.subscribe([:deployment_groups, :agents])
+    Livebook.Teams.Broadcasts.subscribe([:clients, :app_deployments, :deployment_groups, :agents])
     hub = create_team_hub(user, node)
     id = hub.id
 
     assert_receive {:hub_connected, ^id}
+    assert_receive {:client_connected, ^id}
 
     {:ok, hub: hub}
   end
@@ -214,7 +215,7 @@ defmodule LivebookWeb.Integration.Hub.DeploymentGroupTest do
   end
 
   test "shows the agent count", %{conn: conn, hub: hub} do
-    name = "TEAMS_EDIT_DEPLOYMENT_GROUP"
+    name = "TEAMS_EDIT_DEPLOYMENT_GROUP2"
     %{id: id} = insert_deployment_group(name: name, mode: :online, hub_id: hub.id)
 
     {:ok, view, _html} = live(conn, ~p"/hub/#{hub.id}")
@@ -251,4 +252,44 @@ defmodule LivebookWeb.Integration.Hub.DeploymentGroupTest do
            |> Floki.text()
            |> String.trim() == "1"
   end
+
+  @tag :tmp_dir
+  test "shows the app deployed count", %{conn: conn, hub: hub, tmp_dir: tmp_dir} do
+    name = "TEAMS_EDIT_DEPLOYMENT_GROUP3"
+    %{id: id} = insert_deployment_group(name: name, mode: :online, hub_id: hub.id)
+
+    {:ok, view, _html} = live(conn, ~p"/hub/#{hub.id}")
+
+    assert view
+           |> element("#hub-deployment-group-#{id} [aria-label=\"apps deployed\"]")
+           |> render()
+           |> Floki.parse_fragment!()
+           |> Floki.text()
+           |> String.trim() == "0"
+
+    notebook = %{
+      Livebook.Notebook.new()
+      | app_settings: %{
+          Livebook.Notebook.AppSettings.new()
+          | slug: Livebook.Utils.random_short_id()
+        },
+        name: "MyNotebook",
+        hub_id: hub.id,
+        deployment_group_id: to_string(id)
+    }
+
+    files_dir = Livebook.FileSystem.File.local(tmp_dir)
+
+    {:ok, app_deployment} = Livebook.Teams.AppDeployment.new(notebook, files_dir)
+    :ok = Livebook.Teams.deploy_app(hub, app_deployment)
+
+    assert_receive {:app_deployment_started, _}
+
+    assert view
+           |> element("#hub-deployment-group-#{id} [aria-label=\"apps deployed\"]")
+           |> render()
+           |> Floki.parse_fragment!()
+           |> Floki.text()
+           |> String.trim() == "1"
+  end
 end