Fix live view actor leak

cbrewster · Jul 13, 2020 · a2f7409 · a2f7409
1 parent c582894
commit a2f7409
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 16 deletions.
diff --git a/example/src/top.rs b/example/src/top.rs
@@ -1,5 +1,5 @@
 use crate::RootTemplate;
-use actix::{Arbiter, AsyncContext, Message};
+use actix::{AsyncContext, Message};
 use actix_web::{web, Responder};
 use alcova_macros::LiveTemplate;
 use liveview::{
@@ -54,7 +54,7 @@ impl LiveHandler<Tick> for TopLive {
             addr.do_send(TopResult(res));
         };
 
-        Arbiter::spawn(execution);
+        ctx.spawn(actix::fut::wrap_future(execution));
     }
 }
 

diff --git a/liveview/src/live_socket.rs b/liveview/src/live_socket.rs
@@ -1,5 +1,5 @@
 use crate::{
-    live_view::{LiveViewAction, LiveViewId},
+    live_view::{LiveViewAction, LiveViewId, LiveViewMessage},
     Changes, LiveViewRegistry, RenderedTemplate,
 };
 use actix::{Actor, ActorContext, AsyncContext, Handler, Message, Recipient, StreamHandler};
@@ -53,7 +53,7 @@ pub struct SocketViewMessage {
 pub struct LiveSocket {
     registry: LiveViewRegistry,
     context: LiveSocketContext,
-    live_views: Vec<Recipient<LiveViewAction>>,
+    live_views: Vec<Recipient<LiveViewMessage>>,
     heart_beat: Instant,
 }
 
@@ -64,6 +64,14 @@ impl Actor for LiveSocket {
     fn started(&mut self, ctx: &mut Self::Context) {
         self.hb(ctx);
     }
+
+    fn stopped(&mut self, _ctx: &mut Self::Context) {
+        for view in &self.live_views {
+            if view.do_send(LiveViewMessage::Stop).is_err() {
+                warn!("An error occurred shutting down a live view");
+            }
+        }
+    }
 }
 
 impl Handler<SocketViewMessage> for LiveSocket {
@@ -108,7 +116,10 @@ impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for LiveSocket {
                         self.live_views.push(live_view);
                     }
                     ServerMessage::LiveView { id, action } => {
-                        if self.live_views[id.0].do_send(action).is_err() {
+                        if self.live_views[id.0]
+                            .do_send(LiveViewMessage::ClientAction(action))
+                            .is_err()
+                        {
                             warn!("An error occurred handling message");
                         }
                     }

diff --git a/liveview/src/live_view.rs b/liveview/src/live_view.rs
@@ -2,7 +2,7 @@ use crate::{
     live_socket::{ClientMessage, LiveSocketContext, SocketViewMessage},
     LiveSocket, LiveTemplate,
 };
-use actix::{Actor, Addr, Context, Handler, Message};
+use actix::{Actor, ActorContext, Addr, Context, Handler, Message};
 use actix_web::{HttpRequest, HttpResponse, Responder};
 use jsonwebtoken::{encode, EncodingKey, Header};
 use serde::{de::DeserializeOwned, Deserialize, Serialize};
@@ -99,6 +99,12 @@ where
 
 #[derive(Message, Debug, Deserialize)]
 #[rtype(result = "()")]
+pub enum LiveViewMessage {
+    ClientAction(LiveViewAction),
+    Stop,
+}
+
+#[derive(Debug, Deserialize)]
 pub struct LiveViewAction {
     action: String,
     value: Option<String>,
@@ -152,16 +158,21 @@ where
     }
 }
 
-impl<T> Handler<LiveViewAction> for LiveViewActor<T>
+impl<T> Handler<LiveViewMessage> for LiveViewActor<T>
 where
     T: LiveView + Unpin + 'static,
 {
     type Result = ();
 
-    fn handle(&mut self, msg: LiveViewAction, ctx: &mut Self::Context) -> Self::Result {
-        let value = msg.value.unwrap_or(String::new());
-        self.view.handle_event(&msg.action, &value, ctx);
-        self.send_changes();
+    fn handle(&mut self, msg: LiveViewMessage, ctx: &mut Self::Context) -> Self::Result {
+        match msg {
+            LiveViewMessage::ClientAction(LiveViewAction { action, value }) => {
+                let value = value.unwrap_or(String::new());
+                self.view.handle_event(&action, &value, ctx);
+                self.send_changes();
+            }
+            LiveViewMessage::Stop => ctx.stop(),
+        }
     }
 }
 

diff --git a/liveview/src/registry.rs b/liveview/src/registry.rs
@@ -1,6 +1,6 @@
 use crate::{
     live_socket::LiveSocketContext,
-    live_view::{signing_secret, Claims, LiveViewAction, LiveViewActor, LiveViewId},
+    live_view::{signing_secret, Claims, LiveViewActor, LiveViewId, LiveViewMessage},
     LiveSocket, LiveView,
 };
 use actix::{Actor, Addr, Recipient};
@@ -14,7 +14,7 @@ trait LiveViewSpawner: Sync + Send {
         socket: Addr<LiveSocket>,
         context: &LiveSocketContext,
         session_data: &str,
-    ) -> Recipient<LiveViewAction>;
+    ) -> Recipient<LiveViewMessage>;
 }
 
 impl<F> LiveViewSpawner for F
@@ -26,15 +26,15 @@ where
             Addr<LiveSocket>,
             &'r LiveSocketContext,
             &'r str,
-        ) -> Recipient<LiveViewAction>,
+        ) -> Recipient<LiveViewMessage>,
 {
     fn spawn<'a>(
         &self,
         id: LiveViewId,
         socket: Addr<LiveSocket>,
         ctx: &'a LiveSocketContext,
         session_data: &'a str,
-    ) -> Recipient<LiveViewAction> {
+    ) -> Recipient<LiveViewMessage> {
         self(id, socket, ctx, session_data)
     }
 }
@@ -91,7 +91,7 @@ impl LiveViewRegistry {
         socket: Addr<LiveSocket>,
         ctx: &LiveSocketContext,
         session_data: &str,
-    ) -> Option<Recipient<LiveViewAction>> {
+    ) -> Option<Recipient<LiveViewMessage>> {
         self.live_views
             .get(name.as_ref())
             .map(|spawner| spawner.spawn(id, socket, ctx, session_data))