websocket.rs 2.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798
  1. use drama::ws::WebsocketActor;
  2. use drama::{Actor, ActorHandle, Hello};
  3. use std::collections::HashMap;
  4. use std::sync::atomic::AtomicUsize;
  5. use std::sync::{Arc, RwLock};
  6. use warp::Filter;
  7. type Arbiter = Arc<RwLock<HashMap<usize, ActorHandle<WebsocketActor>>>>;
  8. static ID: AtomicUsize = AtomicUsize::new(0);
  9. #[tokio::main]
  10. async fn main() {
  11. let pool = Arc::new(RwLock::new(HashMap::new()));
  12. let pool = warp::any().map(move || pool.clone());
  13. let hello = Hello {}.start().await;
  14. let hello = warp::any().map(move || hello.clone());
  15. // GET /chat -> websocket upgrade
  16. let chat = warp::path("chat")
  17. // The `ws()` filter will prepare Websocket handshake...
  18. .and(warp::ws())
  19. .and(pool)
  20. .and(hello)
  21. .map(
  22. |ws: warp::ws::Ws, pool: Arbiter, hello: ActorHandle<Hello>| {
  23. // This will call our function if the handshake succeeds.
  24. ws.on_upgrade(|socket| async move {
  25. let actor = WebsocketActor::new(socket, hello);
  26. let handle = actor.start().await;
  27. let id = ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
  28. println!("Adding actor {id}");
  29. pool.write().unwrap().insert(id, handle);
  30. })
  31. },
  32. );
  33. // GET / -> index html
  34. let index = warp::path::end().map(|| warp::reply::html(INDEX_HTML));
  35. let routes = index.or(chat);
  36. warp::serve(routes).run(([127, 0, 0, 1], 3030)).await
  37. }
  38. static INDEX_HTML: &str = r#"<!DOCTYPE html>
  39. <html lang="en">
  40. <head>
  41. <title>Warp Chat</title>
  42. </head>
  43. <body>
  44. <h1>Warp chat</h1>
  45. <div id="chat">
  46. <p><em>Connecting...</em></p>
  47. </div>
  48. <input type="text" id="text" />
  49. <button type="button" id="send">Send</button>
  50. <script type="text/javascript">
  51. const chat = document.getElementById('chat');
  52. const text = document.getElementById('text');
  53. const uri = 'ws://' + location.host + '/chat';
  54. const ws = new WebSocket(uri);
  55. let num = 0;
  56. function message(data) {
  57. if (num % 10000 === 0) chat.innerHTML = `${num}`
  58. }
  59. ws.onopen = function() {
  60. chat.innerHTML = '<p><em>Connected!</em></p>';
  61. };
  62. ws.onmessage = function(msg) {
  63. num += 1;
  64. message(msg.data);
  65. };
  66. ws.onclose = function() {
  67. chat.getElementsByTagName('em')[0].innerText = 'Disconnected!';
  68. };
  69. send.onclick = function() {
  70. const msg = text.value;
  71. let i = 0;
  72. while (i < 100000) {
  73. ws.send(msg);
  74. i += 1;
  75. }
  76. text.value = '';
  77. message('<You>: ' + msg);
  78. };
  79. </script>
  80. </body>
  81. </html>
  82. "#;