better queue implementation, add locks and make actor futures send

Cargo.toml

@@ -13,6 +13,7 @@ path = "tests/websocket.rs"
 async-trait = "0.1.68"
 flume = "0.10.14"
 futures = "0.3.28"
+parking_lot = "0.12.1"
 pin-project = "1.1.0"
 thiserror = "1.0.40"
 tokio = { version = "1.28.2", features = [

src/lib.rs

@@ -2,7 +2,8 @@ use crate::runtime::{ActorRuntime, Runtime};
 use async_trait::async_trait;
 use flume::{SendError, Sender};
 use message::{Envelope, Enveloper, Message, MessageRequest};
-use std::fmt::Debug;
+use parking_lot::Mutex;
+use std::{fmt::Debug, sync::Arc};
 use tokio::sync::oneshot;
 pub mod debug;
 pub mod message;
@@ -17,7 +18,7 @@ pub trait Actor {
         Self: Sized + Send + 'static,
     {
         println!("Starting actor");
-        ActorRuntime::run(self)
+        ActorRuntime::run(Arc::new(Mutex::new(self)))
     }
 }
 
@@ -27,7 +28,7 @@ pub trait Handler<M>: Actor
 where
     M: Message,
 {
-    async fn handle(&mut self, message: M) -> Result<M::Response, Error>;
+    async fn handle(this: Arc<Mutex<Self>>, message: M) -> Result<M::Response, Error>;
 }
 
 /// A handle to a spawned actor. Obtained when calling `start` on an [Actor] and is used to send messages
@@ -221,6 +222,8 @@ mod tests {
 
     use std::{sync::atomic::AtomicUsize, time::Duration};
 
+    use tokio::task::LocalSet;
+
     use super::*;
 
     #[tokio::test]
@@ -246,7 +249,7 @@ mod tests {
 
         #[async_trait]
         impl Handler<Foo> for Testor {
-            async fn handle(&mut self, _: Foo) -> Result<usize, Error> {
+            async fn handle(this: Arc<Mutex<Testor>>, _: Foo) -> Result<usize, Error> {
                 println!("Handling Foo");
                 Ok(10)
             }
@@ -254,16 +257,19 @@ mod tests {
 
         #[async_trait]
         impl Handler<Bar> for Testor {
-            async fn handle(&mut self, _: Bar) -> Result<isize, Error> {
-                println!("Handling Bar");
+            async fn handle(this: Arc<Mutex<Testor>>, _: Bar) -> Result<isize, Error> {
+                for _ in 0..10_000 {
+                    println!("Handling Bar");
+                }
                 Ok(10)
             }
         }
 
-        let handle = Testor {}.start();
-
         let mut res = 0;
         let mut res2 = 0;
+
+        let handle = Testor {}.start();
+        println!("HELLO WORLDS");
         for _ in 0..100 {
             res += handle.send_wait(Foo {}).unwrap().await.unwrap();
             res2 += handle.send_wait(Bar {}).unwrap().await.unwrap();
@@ -280,8 +286,8 @@ mod tests {
         assert_eq!(res2, 1000);
     }
 
-    #[tokio::test]
-    async fn it_works_yolo() {
+    #[test]
+    fn it_works_yolo() {
         #[derive(Debug)]
         struct Testor {}
 
@@ -305,7 +311,7 @@ mod tests {
 
         #[async_trait]
         impl Handler<Foo> for Testor {
-            async fn handle(&mut self, _: Foo) -> Result<usize, Error> {
+            async fn handle(this: Arc<Mutex<Testor>>, _: Foo) -> Result<usize, Error> {
                 println!("INCREMENTING COUNT FOO");
                 COUNT.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                 Ok(10)
@@ -314,25 +320,37 @@ mod tests {
 
         #[async_trait]
         impl Handler<Bar> for Testor {
-            async fn handle(&mut self, _: Bar) -> Result<isize, Error> {
+            async fn handle(this: Arc<Mutex<Testor>>, _: Bar) -> Result<isize, Error> {
                 println!("INCREMENTING COUNT BAR");
                 COUNT.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                 Ok(10)
             }
         }
 
-        let handle = Testor {}.start();
+        let rt = tokio::runtime::Builder::new_current_thread()
+            .enable_all()
+            .build()
+            .unwrap();
+        let local_set = LocalSet::new();
 
-        handle.send_wait(Bar {}).unwrap().await.unwrap();
-        handle.send(Foo {}).unwrap();
-        handle.send_forget(Bar {});
+        let task = async {
+            let handle = Testor {}.start();
 
-        for _ in 0..100 {
-            let _ = handle.send(Foo {});
+            handle.send_wait(Bar {}).unwrap().await.unwrap();
+            handle.send(Foo {}).unwrap();
             handle.send_forget(Bar {});
-            tokio::time::sleep(Duration::from_micros(100)).await
-        }
 
-        assert_eq!(COUNT.load(std::sync::atomic::Ordering::Relaxed), 203);
+            for _ in 0..100 {
+                let _ = handle.send(Foo {});
+                handle.send_forget(Bar {});
+                tokio::time::sleep(Duration::from_micros(100)).await
+            }
+
+            assert_eq!(COUNT.load(std::sync::atomic::Ordering::Relaxed), 203);
+            handle.send_cmd(ActorCommand::Stop)
+        };
+
+        local_set.spawn_local(task);
+        rt.block_on(local_set)
     }
 }

src/message.rs

@@ -1,5 +1,8 @@
+use std::sync::Arc;
+
 use crate::{Actor, Error, Handler};
 use async_trait::async_trait;
+use parking_lot::Mutex;
 use tokio::sync::oneshot;
 
 /// Represents a message that can be sent to an actor. The response type is what the actor must return in its handler implementation.
@@ -11,7 +14,7 @@ pub trait Message {
 /// type when creating an actor handle, otherwise we would only be able to send a single message type to the actor.
 #[async_trait(?Send)]
 pub trait ActorMessage<A: Actor> {
-    async fn handle(&mut self, actor: &mut A);
+    async fn handle(&mut self, actor: Arc<Mutex<A>>);
 }
 
 /// Used by [ActorHandle][super::ActorHandle]s to pack [Message]s into [Envelope]s so we have a type erased message to send to the actor.
@@ -32,7 +35,7 @@ where
 {
     pub fn new<M>(message: M, tx: Option<oneshot::Sender<M::Response>>) -> Self
     where
-        A: Handler<M>,
+        A: Handler<M> + 'static,
         M: Message + Send + 'static,
         M::Response: Send,
     {
@@ -57,7 +60,7 @@ impl<A> ActorMessage<A> for Envelope<A>
 where
     A: Actor,
 {
-    async fn handle(&mut self, actor: &mut A) {
+    async fn handle(&mut self, actor: Arc<Mutex<A>>) {
         self.message.handle(actor).await
     }
 }
@@ -66,11 +69,11 @@ where
 impl<A, M> ActorMessage<A> for EnvelopeInner<M>
 where
     M: Message,
-    A: Actor + Handler<M>,
+    A: Actor + Handler<M> + 'static,
 {
-    async fn handle(&mut self, actor: &mut A) {
+    async fn handle(&mut self, actor: Arc<Mutex<A>>) {
         let message = self.message.take().expect("Message already processed");
-        let result = actor.handle(message).await;
+        let result = A::handle(actor, message).await;
         if let Some(res_tx) = self.tx.take() {
             let _ = res_tx.send(result.unwrap());
         }
@@ -79,7 +82,7 @@ where
 
 impl<A, M> Enveloper<A, M> for A
 where
-    A: Actor + Handler<M>,
+    A: Actor + Handler<M> + 'static,
     M: Message + Send + 'static,
     M::Response: Send,
 {

src/runtime.rs

@@ -4,15 +4,19 @@ use crate::{
 };
 use flume::Receiver;
 use futures::Future;
+use parking_lot::Mutex;
 use pin_project::pin_project;
 use std::{
     collections::VecDeque,
     pin::Pin,
+    sync::Arc,
     task::{Context, Poll},
 };
 
+const DEFAULT_QUEUE_CAPACITY: usize = 16;
+
 pub trait Runtime<A> {
-    fn run(actor: A) -> ActorHandle<A>
+    fn run(actor: Arc<Mutex<A>>) -> ActorHandle<A>
     where
         A: Actor + Send + 'static,
     {
@@ -26,19 +30,19 @@ pub trait Runtime<A> {
 #[pin_project]
 pub struct ActorRuntime<A>
 where
-    A: Actor,
+    A: Actor + Send + 'static,
 {
-    actor: A,
+    actor: Arc<Mutex<A>>,
     command_rx: Receiver<ActorCommand>,
     message_rx: Receiver<Envelope<A>>,
-    message_queue: VecDeque<Envelope<A>>,
+    process_queue: VecDeque<ActorJob<A>>,
 }
 
-impl<A> Runtime<A> for ActorRuntime<A> where A: Actor {}
+impl<A> Runtime<A> for ActorRuntime<A> where A: Actor + Send + 'static {}
 
 impl<A> Future for ActorRuntime<A>
 where
-    A: Actor,
+    A: Actor + Send + 'static,
 {
     type Output = Result<(), Error>;
 
@@ -60,21 +64,28 @@ where
                 Poll::Pending => {}
             };
 
-            // Process all pending messages
-            this.message_queue
-                .retain_mut(|message| message.handle(this.actor).as_mut().poll(cx).is_pending());
+            // Process the pending futures
+            this.process_queue
+                .retain_mut(|job| job.poll(&mut this.actor.lock(), cx).is_pending());
 
-            // Poll message receiver and continue to process if anything comes up
+            // Poll message receiver
             while let Poll::Ready(Ok(message)) =
                 Pin::new(&mut this.message_rx.recv_async()).poll(cx)
             {
-                this.message_queue.push_back(message);
+                // this.process_queue.push_back(ActorJob::new(message));
+                if this.process_queue.len() >= DEFAULT_QUEUE_CAPACITY {
+                    break;
+                }
             }
 
+            // Process pending futures again after we've potentially received some
+            this.process_queue
+                .retain_mut(|job| job.poll(&mut this.actor.lock(), cx).is_pending());
+
             // Poll again and process new messages if any
             match Pin::new(&mut this.message_rx.recv_async()).poll(cx) {
                 Poll::Ready(Ok(message)) => {
-                    this.message_queue.push_back(message);
+                    // this.process_queue.push_back(ActorJob::new(message));
                     continue;
                 }
                 Poll::Ready(Err(_)) => {
@@ -82,11 +93,12 @@ where
                     break Poll::Ready(Err(Error::ActorChannelClosed));
                 }
                 Poll::Pending => {
-                    if !this.message_queue.is_empty() {
+                    if !this.process_queue.is_empty() {
                         continue;
                     }
                 }
             };
+
             cx.waker().wake_by_ref();
             return Poll::Pending;
         }
@@ -95,10 +107,10 @@ where
 
 impl<A> ActorRuntime<A>
 where
-    A: Actor,
+    A: Actor + 'static + Send,
 {
     pub fn new(
-        actor: A,
+        actor: Arc<Mutex<A>>,
         command_rx: Receiver<ActorCommand>,
         message_rx: Receiver<Envelope<A>>,
     ) -> Self {
@@ -107,7 +119,52 @@ where
             actor,
             command_rx,
             message_rx,
-            message_queue: VecDeque::new(),
+            process_queue: VecDeque::with_capacity(DEFAULT_QUEUE_CAPACITY),
         }
     }
 }
+
+pub trait ActorFuture<A>
+where
+    A: Actor,
+{
+    type Output;
+    fn poll(
+        self: Pin<&mut Self>,
+        actor: &mut A,
+        cx: &mut std::task::Context<'_>,
+    ) -> Poll<Self::Output>;
+}
+
+struct ActorJob<A>(pub Pin<Box<dyn ActorFuture<A, Output = ()> + Send + 'static>>);
+
+impl<A> ActorJob<A>
+where
+    A: Actor,
+{
+    fn new<F>(future: F) -> Self
+    where
+        F: ActorFuture<A, Output = ()> + Send + 'static,
+    {
+        Self(Box::pin(future))
+    }
+
+    fn poll(&mut self, actor: &mut A, cx: &mut std::task::Context<'_>) -> Poll<()> {
+        self.0.as_mut().poll(actor, cx)
+    }
+}
+
+/* impl<A> ActorFuture<A> for Envelope<A>
+where
+    A: Actor,
+{
+    type Output = ();
+
+    fn poll(
+        self: Pin<&mut Self>,
+        actor: &mut A,
+        cx: &mut std::task::Context<'_>,
+    ) -> Poll<Self::Output> {
+        self.get_mut().handle(actor).as_mut().poll(cx)
+    }
+} */

src/ws.rs

@@ -6,9 +6,11 @@ use crate::{
 use async_trait::async_trait;
 use flume::Receiver;
 use futures::{
+    sink::Feed,
     stream::{SplitSink, SplitStream},
     Future, SinkExt, Stream, StreamExt,
 };
+use parking_lot::Mutex;
 use pin_project::pin_project;
 use std::{
     collections::VecDeque,
@@ -16,46 +18,134 @@ use std::{
     sync::atomic::AtomicUsize,
     task::{Context, Poll},
 };
-use warp::ws::WebSocket;
+use std::{sync::Arc, time::Duration};
+use warp::ws::{Message, WebSocket};
 
 pub struct WebsocketActor {
     websocket: Option<WebSocket>,
 }
 
+impl WebsocketActor {
+    pub fn new(ws: WebSocket) -> Self {
+        Self {
+            websocket: Some(ws),
+        }
+    }
+}
+
 impl Actor for WebsocketActor {
-    fn start(self) -> ActorHandle<Self>
-    where
-        Self: Sized + Send + 'static,
-    {
-        println!("Starting websocket actor");
-        WebsocketRuntime::run(self)
+    fn start(self) -> ActorHandle<Self> {
+        WebsocketRuntime::run(Arc::new(Mutex::new(self)))
     }
 }
+type WsFuture = Pin<Box<dyn Future<Output = Result<Option<Message>, Error>> + Send>>;
 
-impl WebsocketActor {
-    pub fn new(ws: WebSocket) -> Self {
+struct ActorItem {
+    message: Option<Message>,
+    future: Option<WsFuture>,
+}
+
+impl ActorItem {
+    fn poll(
+        mut self: Pin<&mut Self>,
+        actor: Arc<Mutex<WebsocketActor>>,
+        cx: &mut std::task::Context<'_>,
+    ) -> Poll<Result<Option<Message>, warp::Error>> {
+        //         let mut this = self.project();
+        let message = self.as_mut().message.take();
+
+        match message {
+            Some(message) => {
+                let fut = WebsocketActor::handle(actor, message);
+                self.future = Some(fut);
+                // let Some(ref mut fut) = self.future else {panic!("Rust no longer works")};
+                match self.future.as_mut().unwrap().as_mut().poll(cx) {
+                    Poll::Ready(result) => match result {
+                        Ok(response) => Poll::Ready(Ok(response)),
+                        Err(e) => {
+                            println!("Shit's fucked son {e}");
+                            Poll::Ready(Ok(None))
+                        }
+                    },
+                    Poll::Pending => Poll::Pending,
+                }
+            }
+            None => match self.future {
+                Some(ref mut fut) => match fut.as_mut().poll(cx) {
+                    Poll::Ready(result) => match result {
+                        Ok(response) => Poll::Ready(Ok(response)),
+                        Err(e) => {
+                            println!("Shit's fucked son {e}");
+                            Poll::Ready(Ok(None))
+                        }
+                    },
+                    Poll::Pending => Poll::Pending,
+                },
+                None => Poll::Pending,
+            },
+        }
+    }
+}
+
+impl ActorItem {
+    pub fn new(message: Message) -> Self {
         Self {
-            websocket: Some(ws),
+            message: Some(message),
+            future: None,
         }
     }
 }
+/*
+trait WsActorFuture {
+    fn poll(
+        self: Pin<&mut Self>,
+        actor: &mut WebsocketActor,
+        sink: Pin<&mut SplitSink<WebSocket, Message>>,
+        cx: &mut std::task::Context<'_>,
+    ) -> Poll<Result<(), warp::Error>>;
+}
 
+impl WsActorFuture for Message {
+    fn poll(
+        mut self: Pin<&mut Self>,
+        actor: &mut WebsocketActor,
+        mut sink: Pin<&mut SplitSink<WebSocket, Message>>,
+        cx: &mut std::task::Context<'_>,
+    ) -> Poll<Result<(), warp::Error>> {
+        match actor.handle(self.as_mut().clone()).as_mut().poll(cx) {
+            Poll::Ready(result) => match result {
+                Ok(response) => {
+                    if let Some(response) = response {
+                        Pin::new(&mut sink.feed(response)).poll(cx)
+                    } else {
+                        Poll::Pending
+                    }
+                }
+                Err(e) => {
+                    println!("Shit's fucked son {e}");
+                    Poll::Ready(Ok(()))
+                }
+            },
+            Poll::Pending => Poll::Pending,
+        }
+    }
+}
+ */
 static PROCESSED: AtomicUsize = AtomicUsize::new(0);
 
 #[pin_project]
 pub struct WebsocketRuntime {
-    actor: WebsocketActor,
+    actor: Arc<Mutex<WebsocketActor>>,
 
     status: ActorStatus,
 
-    // Pin these 2 as we are polling them directly so we know they never move
     /// The receiving end of the websocket
     #[pin]
     ws_stream: SplitStream<WebSocket>,
 
     /// The sending end of the websocket
     #[pin]
-    ws_sink: SplitSink<WebSocket, warp::ws::Message>,
+    ws_sink: SplitSink<WebSocket, Message>,
 
     /// Actor message receiver
     message_rx: Receiver<Envelope<WebsocketActor>>,
@@ -67,19 +157,20 @@ pub struct WebsocketRuntime {
     message_queue: VecDeque<Envelope<WebsocketActor>>,
 
     /// Received, but not yet processed websocket messages
-    request_queue: VecDeque<warp::ws::Message>,
+    processing_queue: VecDeque<ActorItem>,
 
-    /// Processed websocket messages ready to be flushed in the sink
-    response_queue: VecDeque<warp::ws::Message>,
+    /// Processed websocket messages being flushed to the sink
+    response_queue: VecDeque<Message>,
 }
 
 impl WebsocketRuntime {
     pub fn new(
-        mut actor: WebsocketActor,
+        actor: Arc<Mutex<WebsocketActor>>,
         command_rx: Receiver<ActorCommand>,
         message_rx: Receiver<Envelope<WebsocketActor>>,
     ) -> Self {
         let (ws_sink, ws_stream) = actor
+            .lock()
             .websocket
             .take()
             .expect("Websocket runtime already started")
@@ -92,7 +183,7 @@ impl WebsocketRuntime {
             message_rx,
             command_rx,
             message_queue: VecDeque::new(),
-            request_queue: VecDeque::new(),
+            processing_queue: VecDeque::new(),
             response_queue: VecDeque::new(),
             status: ActorStatus::Starting,
         }
@@ -106,7 +197,7 @@ impl Future for WebsocketRuntime {
         let mut this = self.project();
 
         loop {
-            // Poll command receiver
+            // Poll command receiver and immediatelly process it
             match Pin::new(&mut this.command_rx.recv_async()).poll(cx) {
                 Poll::Ready(Ok(message)) => match message {
                     ActorCommand::Stop => {
@@ -124,48 +215,54 @@ impl Future for WebsocketRuntime {
             // Poll the websocket stream for any messages and store them to the queue
             while let Poll::Ready(Some(ws_message)) = this.ws_stream.as_mut().poll_next(cx) {
                 match ws_message {
-                    Ok(message) => this.request_queue.push_back(message),
+                    Ok(message) => {
+                        this.processing_queue.push_back(ActorItem::new(message));
+                        if this.processing_queue.len() > 500 {
+                            break;
+                        }
+                    }
                     Err(e) => {
                         eprintln!("WS error occurred {e}")
                     }
                 }
             }
 
-            // Respond to any queued and processed websocket messages
             let mut idx = 0;
-            while idx < this.request_queue.len() {
-                let ws_message = &this.request_queue[idx];
-                match this.actor.handle(ws_message.to_owned()).as_mut().poll(cx) {
+            //             let actor = &mut this.actor.lock();
+            while idx < this.processing_queue.len() {
+                let job = &mut this.processing_queue[idx];
+                match ActorItem::poll(Pin::new(job), this.actor.clone(), cx) {
                     Poll::Ready(result) => match result {
                         Ok(response) => {
                             if let Some(response) = response {
-                                match Pin::new(&mut this.ws_sink.feed(response)).poll(cx) {
-                                    Poll::Ready(result) => {
-                                        result?;
-                                        this.request_queue.swap_remove_front(idx);
-                                        PROCESSED
-                                            .fetch_add(1, std::sync::atomic::Ordering::Acquire);
-                                    }
-                                    Poll::Pending => idx += 1,
+                                let feed = &mut this.ws_sink.feed(response);
+                                let mut feed = Pin::new(feed);
+                                while feed.as_mut().poll(cx).is_pending() {
+                                    let _ = feed.as_mut().poll(cx);
                                 }
                             }
+                            PROCESSED.fetch_add(1, std::sync::atomic::Ordering::Acquire);
+                            this.processing_queue.swap_remove_front(idx);
+                        }
+                        Err(e) => {
+                            println!("Shit's fucked my dude {e}")
                         }
-                        Err(e) => return Poll::Ready(Err(e)),
                     },
                     Poll::Pending => idx += 1,
                 }
             }
 
             println!(
-                "PROCESSED {}",
-                PROCESSED.load(std::sync::atomic::Ordering::Acquire)
+                "PROCESSED {} CURRENT IN QUEUE {}",
+                PROCESSED.load(std::sync::atomic::Ordering::Acquire),
+                this.processing_queue.len(),
             );
 
             let _ = Pin::new(&mut this.ws_sink.flush()).poll(cx);
 
             // Process all messages
-            this.message_queue
-                .retain_mut(|message| message.handle(this.actor).as_mut().poll(cx).is_pending());
+            /*             this.message_queue
+            .retain_mut(|message| message.handle(actor).as_mut().poll(cx).is_pending()); */
 
             // Poll message receiver and continue to process if anything comes up
             while let Poll::Ready(Ok(message)) =
@@ -198,7 +295,7 @@ impl Future for WebsocketRuntime {
 }
 
 impl Runtime<WebsocketActor> for WebsocketRuntime {
-    fn run(actor: WebsocketActor) -> ActorHandle<WebsocketActor> {
+    fn run(actor: Arc<Mutex<WebsocketActor>>) -> ActorHandle<WebsocketActor> {
         let (message_tx, message_rx) = flume::unbounded();
         let (command_tx, command_rx) = flume::unbounded();
         tokio::spawn(WebsocketRuntime::new(actor, command_rx, message_rx));
@@ -206,17 +303,18 @@ impl Runtime<WebsocketActor> for WebsocketRuntime {
     }
 }
 
-impl crate::Message for warp::ws::Message {
-    type Response = Option<warp::ws::Message>;
+impl crate::Message for Message {
+    type Response = Option<Message>;
 }
 
 #[async_trait]
-impl Handler<warp::ws::Message> for WebsocketActor {
+impl Handler<Message> for WebsocketActor {
     async fn handle(
-        &mut self,
-        message: warp::ws::Message,
-    ) -> Result<<warp::ws::Message as crate::message::Message>::Response, crate::Error> {
-        // println!("Actor received message {message:?}");
+        this: Arc<Mutex<Self>>,
+        message: Message,
+    ) -> Result<<Message as crate::message::Message>::Response, crate::Error> {
+        println!("Actor received message {message:?}");
+        // tokio::time::sleep(Duration::from_micros(500)).await;
         if message.is_text() {
             Ok(Some(message))
         } else {

tests/websocket.rs

@@ -1,27 +1,32 @@
-use std::collections::HashMap;
-use std::sync::{Arc, RwLock};
-
 use actors::ws::WebsocketActor;
 use actors::{Actor, ActorHandle};
+use std::collections::HashMap;
+use std::sync::atomic::AtomicUsize;
+use std::sync::{Arc, RwLock};
 use warp::Filter;
 
 type Arbiter = Arc<RwLock<HashMap<usize, ActorHandle<WebsocketActor>>>>;
 
+static ID: AtomicUsize = AtomicUsize::new(0);
+
 #[tokio::main]
 async fn main() {
-    let arbiter = Arc::new(RwLock::new(HashMap::new()));
-    let arbiter = warp::any().map(move || arbiter.clone());
+    let pool = Arc::new(RwLock::new(HashMap::new()));
+    let pool = warp::any().map(move || pool.clone());
+
     // GET /chat -> websocket upgrade
     let chat = warp::path("chat")
         // The `ws()` filter will prepare Websocket handshake...
         .and(warp::ws())
-        .and(arbiter)
-        .map(|ws: warp::ws::Ws, arbiter: Arbiter| {
+        .and(pool)
+        .map(|ws: warp::ws::Ws, pool: Arbiter| {
             // This will call our function if the handshake succeeds.
             ws.on_upgrade(move |socket| {
                 let actor = WebsocketActor::new(socket);
                 let handle = actor.start();
-                arbiter.write().unwrap().insert(0, handle);
+                // let id = ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+                // println!("Adding actor {id}");
+                pool.write().unwrap().insert(0, handle);
                 futures::future::ready(())
             })
         });
@@ -32,7 +37,7 @@ async fn main() {
 
     let routes = index.or(chat);
 
-    warp::serve(routes).run(([127, 0, 0, 1], 3030)).await;
+    warp::serve(routes).run(([127, 0, 0, 1], 3030)).await
 }
 
 static INDEX_HTML: &str = r#"<!DOCTYPE html>
@@ -53,8 +58,10 @@ static INDEX_HTML: &str = r#"<!DOCTYPE html>
         const uri = 'ws://' + location.host + '/chat';
         const ws = new WebSocket(uri);
 
-        function message(data) {
+        let num = 0;
 
+        function message(data) {
+            if (num % 10000 === 0) chat.innerHTML = `${num}`
         }
 
         ws.onopen = function() {
@@ -63,6 +70,7 @@ static INDEX_HTML: &str = r#"<!DOCTYPE html>
 
         ws.onmessage = function(msg) {
             message(msg.data);
+            num += 1;
         };
 
         ws.onclose = function() {
@@ -72,7 +80,7 @@ static INDEX_HTML: &str = r#"<!DOCTYPE html>
         send.onclick = function() {
             const msg = text.value;
             let i = 0;
-            while (i < 10000) {
+            while (i < 100000) {
                 ws.send(msg);
                 i += 1;
             }