refactor(pubsub): replace custom Edgehog.PubSub with Ash.Notifier.PubSub (#1014)

- Removed the legacy Edgehog.PubSub module and the PublishNotification change.
- Updated Deployment and OTAOperation resources to use Ash.Notifier.PubSub.
- Added pub_sub blocks with proper transforms to map Ash actions to legacy event types.
- Updated code subscribing to events to use Phoenix.PubSub directly.
- Adjusted tests to match new notification structure.

By using as built in notification system we will be able to integrate more features in the edgehog backend. Namely Ash GraphQL Subscriptions and a better notification system to campaigns, reacting only to relevant updates and discarding other events.

Signed-off-by: ArnelaL <arnela.lisic@secomind.com>

Author: ArnelaL

backend/lib/edgehog/changes/publish_notification.ex

@@ -22,9 +22,9 @@ defmodule Edgehog.Containers.Deployment do
   @moduledoc false
   use Edgehog.MultitenantResource,
     domain: Edgehog.Containers,
-    extensions: [AshGraphql.Resource]
+    extensions: [AshGraphql.Resource],
+    notifiers: [Ash.Notifier.PubSub]
 
-  alias Edgehog.Changes.PublishNotification
   alias Edgehog.Containers.Deployment.Calculations
   alias Edgehog.Containers.Deployment.Changes
   alias Edgehog.Containers.Deployment.Types.DeploymentState
@@ -66,7 +66,6 @@ defmodule Edgehog.Containers.Deployment do
       change manage_relationship(:device_id, :device, type: :append)
       change Changes.Relate
       change Changes.SendRequest
-      change {PublishNotification, event_type: :deployment_created}
     end
 
     create :just_create do
@@ -85,7 +84,6 @@ defmodule Edgehog.Containers.Deployment do
 
       change manage_relationship(:device_id, :device, type: :append)
       change Changes.Relate
-      change {PublishNotification, event_type: :deployment_created}
     end
 
     if @testing do
@@ -157,7 +155,6 @@ defmodule Edgehog.Containers.Deployment do
       require_atomic? false
 
       change Changes.SendDeploymentToDevice
-      change {PublishNotification, event_type: :deployment_updated}
     end
 
     update :upgrade_release do
@@ -177,35 +174,30 @@ defmodule Edgehog.Containers.Deployment do
       change set_attribute(:state, :sent)
 
       require_atomic? false
-      change {PublishNotification, event_type: :deployment_updated}
     end
 
     update :mark_as_started do
       change set_attribute(:state, :started)
 
       require_atomic? false
-      change {PublishNotification, event_type: :deployment_updated}
     end
 
     update :mark_as_starting do
       require_atomic? false
 
       change Changes.MarkAsStarting
-      change {PublishNotification, event_type: :deployment_updated}
     end
 
     update :mark_as_stopped do
       change set_attribute(:state, :stopped)
 
       require_atomic? false
-      change {PublishNotification, event_type: :deployment_updated}
     end
 
     update :mark_as_stopping do
       require_atomic? false
 
       change Changes.MarkAsStopping
-      change {PublishNotification, event_type: :deployment_updated}
     end
 
     update :mark_as_errored do
@@ -217,22 +209,19 @@ defmodule Edgehog.Containers.Deployment do
       change set_attribute(:state, :error)
 
       require_atomic? false
-      change {PublishNotification, event_type: :deployment_error}
     end
 
     update :mark_as_deleting do
       change set_attribute(:state, :deleting)
 
       require_atomic? false
-      change {PublishNotification, event_type: :deployment_updated}
     end
 
     update :maybe_run_ready_actions do
       change Changes.MaybeRunReadyActions
       change Changes.MaybePublishDeploymentReady
 
       require_atomic? false
-      change {PublishNotification, event_type: :deployment_updated}
     end
 
     read :filter_by_release do
@@ -294,6 +283,56 @@ defmodule Edgehog.Containers.Deployment do
     identity :release_instance, [:device_id, :release_id]
   end
 
+  pub_sub do
+    prefix "deployments"
+    module EdgehogWeb.Endpoint
+
+    publish :deploy, [[:id, "*"]]
+    publish :just_create, [[:id, "*"]]
+
+    publish :mark_as_sent, [[:id, "*"]]
+    publish :mark_as_started, [[:id, "*"]]
+    publish :mark_as_starting, [[:id, "*"]]
+    publish :mark_as_stopped, [[:id, "*"]]
+    publish :mark_as_stopping, [[:id, "*"]]
+    publish :mark_as_errored, [[:id, "*"]]
+    publish :mark_as_deleting, [[:id, "*"]]
+    publish :maybe_run_ready_actions, [[:id, "*"]]
+
+    transform fn notification ->
+      deployment = notification.data
+      action = notification.action.name
+
+      event_type =
+        cond do
+          Map.get(notification.metadata || %{}, :custom_event) == :deployment_ready ->
+            :deployment_ready
+
+          action in [:deploy, :just_create] ->
+            :deployment_created
+
+          action in [
+            :mark_as_sent,
+            :mark_as_started,
+            :mark_as_starting,
+            :mark_as_stopped,
+            :mark_as_stopping,
+            :mark_as_deleting,
+            :maybe_run_ready_actions
+          ] ->
+            :deployment_updated
+
+          action == :mark_as_errored ->
+            :deployment_error
+
+          true ->
+            :unknown_event
+        end
+
+      {event_type, deployment}
+    end
+  end
+
   postgres do
     table "application_deployments"
 

backend/lib/edgehog/containers/deployment.ex

@@ -24,17 +24,22 @@ defmodule Edgehog.Containers.Deployment.Changes.MaybePublishDeploymentReady do
   """
   use Ash.Resource.Change
 
-  alias Edgehog.PubSub
-
   @impl Ash.Resource.Change
   def change(changeset, _opts, _context) do
     changeset
     |> Ash.Changeset.load(:is_ready)
     |> Ash.Changeset.after_transaction(fn _changeset, transaction_result ->
       with {:ok, deployment} <- transaction_result,
            {:ok, deployment} <- Ash.load(deployment, :is_ready) do
-        if deployment.is_ready,
-          do: PubSub.publish!(:deployment_ready, deployment)
+        if deployment.is_ready do
+          Ash.Notifier.notify(%Ash.Notifier.Notification{
+            data: deployment,
+            for: [Ash.Notifier.PubSub],
+            action: %{type: :update, name: :maybe_run_ready_actions},
+            resource: Edgehog.Containers.Deployment,
+            metadata: %{custom_event: :deployment_ready}
+          })
+        end
 
         {:ok, deployment}
       end

backend/lib/edgehog/containers/deployment/changes/maybe_publish_deployment_ready.ex

@@ -27,7 +27,6 @@ defmodule Edgehog.DeploymentCampaigns.DeploymentMechanism.Lazy.Core do
   alias Edgehog.DeploymentCampaigns.DeploymentCampaign
   alias Edgehog.DeploymentCampaigns.DeploymentTarget
   alias Edgehog.Error.AstarteAPIError
-  alias Edgehog.PubSub
 
   @doc """
   Fetches a deployment campaign by its ID and tenant ID, raising an error if not found.
@@ -247,7 +246,11 @@ defmodule Edgehog.DeploymentCampaigns.DeploymentMechanism.Lazy.Core do
     - :ok if the subscription is successful, otherwise raises an error.
   """
   def subscribe_to_deployment_updates!(deployment_id) do
-    with {:error, reason} <- PubSub.subscribe_to_events_for({:deployment, deployment_id}) do
+    with {:error, reason} <-
+           Phoenix.PubSub.subscribe(
+             Edgehog.PubSub,
+             "deployments:#{deployment_id}"
+           ) do
       raise reason
     end
   end
@@ -262,7 +265,10 @@ defmodule Edgehog.DeploymentCampaigns.DeploymentMechanism.Lazy.Core do
     - :ok if the subscription is successful, otherwise raises an error.
   """
   def unsubscribe_to_deployment_updates!(deployment_id) do
-    PubSub.unsubscribe_to_events_for({:deployment, deployment_id})
+    Phoenix.PubSub.unsubscribe(
+      Edgehog.PubSub,
+      "deployments:#{deployment_id}"
+    )
   end
 
   @doc """

backend/lib/edgehog/deployment_campaigns/deployment_mechanism/lazy/core.ex

@@ -449,7 +449,7 @@ defmodule Edgehog.DeploymentCampaigns.DeploymentMechanism.Lazy.Executor do
   # events enqueued with the :next_event action. This means that we can be sure an :info event
   # or a timeout won't be handled, e.g., between a rollout and the handling of its error
 
-  def handle_event(:info, {:deployment_ready, deployment}, _state, data) do
+  def handle_event(:info, %{payload: {:deployment_ready, deployment}}, _state, data) do
     # We always cancel the retry timeout for every kind of update we see on an Deployment.
     # This ensures we don't resend the request even if we accidentally miss the acknowledge.
     # If the timeout does not exist, this is a no-op anyway.
@@ -464,11 +464,11 @@ defmodule Edgehog.DeploymentCampaigns.DeploymentMechanism.Lazy.Executor do
   end
 
   # Ignore deployment_updated events
-  def handle_event(:info, {:deployment_updated, _deployment}, _state, _data) do
+  def handle_event(:info, %{payload: {:deployment_updated, _deployment}}, _state, _data) do
     :keep_state_and_data
   end
 
-  def handle_event(:info, {:deployment_error, deployment}, _state, data) do
+  def handle_event(:info, %{payload: {:deployment_error, deployment}}, _state, data) do
     # We always cancel the retry timeout for every kind of update we see on an Deployment.
     # This ensures we don't resend the request even if we accidentally miss the acknowledge.
     # If the timeout does not exist, this is a no-op anyway.

backend/lib/edgehog/deployment_campaigns/deployment_mechanism/lazy/executor.ex

@@ -24,9 +24,9 @@ defmodule Edgehog.OSManagement.OTAOperation do
     domain: Edgehog.OSManagement,
     extensions: [
       AshGraphql.Resource
-    ]
+    ],
+    notifiers: [Ash.Notifier.PubSub]
 
-  alias Edgehog.Changes.PublishNotification
   alias Edgehog.OSManagement.OTAOperation.Changes
   alias Edgehog.OSManagement.OTAOperation.ManualActions
   alias Edgehog.OSManagement.OTAOperation.Status
@@ -67,7 +67,6 @@ defmodule Edgehog.OSManagement.OTAOperation do
       accept [:base_image_url, :device_id]
 
       change Changes.SendUpdateRequest
-      change {PublishNotification, event_type: :ota_operation_created}
     end
 
     create :manual do
@@ -95,7 +94,6 @@ defmodule Edgehog.OSManagement.OTAOperation do
       change set_attribute(:manual?, true)
       change Changes.HandleEphemeralImageUpload
       change Changes.SendUpdateRequest
-      change {PublishNotification, event_type: :ota_operation_created}
     end
 
     destroy :destroy do
@@ -107,12 +105,11 @@ defmodule Edgehog.OSManagement.OTAOperation do
     end
 
     update :mark_as_timed_out do
-      # Needed because PublishNotification and HandleEphemeralImageDeletion are not atomic
+      # Needed because HandleEphemeralImageDeletion are not atomic
       require_atomic? false
 
       change set_attribute(:status, :failure)
       change set_attribute(:status_code, :request_timeout)
-      change {PublishNotification, event_type: :ota_operation_updated}
 
       change Changes.HandleEphemeralImageDeletion do
         where attribute_equals(:manual?, true)
@@ -126,11 +123,9 @@ defmodule Edgehog.OSManagement.OTAOperation do
     update :update_status do
       accept [:status, :status_progress, :status_code, :message]
 
-      # Needed because PublishNotification and HandleEphemeralImageDeletion are not atomic
+      # Needed because and HandleEphemeralImageDeletion are not atomic
       require_atomic? false
 
-      change {PublishNotification, event_type: :ota_operation_updated}
-
       change Changes.HandleEphemeralImageDeletion do
         where [attribute_equals(:manual?, true), attribute_in(:status, @terminal_statuses)]
       end
@@ -222,7 +217,37 @@ defmodule Edgehog.OSManagement.OTAOperation do
     calculate :finished?, :boolean, expr(status == :success or status == :failure)
   end
 
-  # TODO: notifiers to replace the old PubSub
+  pub_sub do
+    prefix "ota_operations"
+    module EdgehogWeb.Endpoint
+
+    publish :create_managed, [[:id, "*"]]
+    publish :manual, [[:id, "*"]]
+    publish :mark_as_timed_out, [[:id, "*"]]
+    publish :update_status, [[:id, "*"]]
+
+    transform fn notification ->
+      ota_operation = notification.data
+      action = notification.action.name
+
+      event_type =
+        cond do
+          action in [:create_managed, :manual] ->
+            :ota_operation_created
+
+          action in [
+            :mark_as_timed_out,
+            :update_status
+          ] ->
+            :ota_operation_updated
+
+          true ->
+            :unknown_event
+        end
+
+      {event_type, ota_operation}
+    end
+  end
 
   postgres do
     table "ota_operations"