From 750426b60f0677621cd7721094831d83ccdc65a2 Mon Sep 17 00:00:00 2001
From: xamfy <spratyaksh22@gmail.com>
Date: Wed, 5 Oct 2022 16:13:58 +0530
Subject: [PATCH] feature: websocket stats endpoint

---
 Cargo.lock               | 54 ++++++++++++++++++++++++++++++++++++++++
 Cargo.toml               |  3 +++
 src/handlers/handlers.rs | 39 ++++++++++++++++++++++++++++-
 src/main.rs              |  3 +++
 4 files changed, 98 insertions(+), 1 deletion(-)

diff --git a/Cargo.lock b/Cargo.lock
index af0e866..0e6944a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2,6 +2,30 @@
 # It is not intended for manual editing.
 version = 3
 
+[[package]]
+name = "actix"
+version = "0.13.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f728064aca1c318585bf4bb04ffcfac9e75e508ab4e8b1bd9ba5dfe04e2cbed5"
+dependencies = [
+ "actix-rt",
+ "actix_derive",
+ "bitflags",
+ "bytes",
+ "crossbeam-channel",
+ "futures-core",
+ "futures-sink",
+ "futures-task",
+ "futures-util",
+ "log",
+ "once_cell",
+ "parking_lot",
+ "pin-project-lite",
+ "smallvec",
+ "tokio",
+ "tokio-util",
+]
+
 [[package]]
 name = "actix-codec"
 version = "0.5.0"
@@ -204,6 +228,23 @@ dependencies = [
  "url",
 ]
 
+[[package]]
+name = "actix-web-actors"
+version = "4.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "31efe7896f3933ce03dd4710be560254272334bb321a18fd8ff62b1a557d9d19"
+dependencies = [
+ "actix",
+ "actix-codec",
+ "actix-http",
+ "actix-web",
+ "bytes",
+ "bytestring",
+ "futures-core",
+ "pin-project-lite",
+ "tokio",
+]
+
 [[package]]
 name = "actix-web-codegen"
 version = "4.1.0"
@@ -266,6 +307,17 @@ dependencies = [
  "syn",
 ]
 
+[[package]]
+name = "actix_derive"
+version = "0.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6d44b8fee1ced9671ba043476deddef739dd0959bf77030b26b738cc591737a7"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "adler"
 version = "1.0.2"
@@ -2855,8 +2907,10 @@ checksum = "c811ca4a8c853ef420abd8592ba53ddbbac90410fab6903b3e79972a631f7680"
 name = "x-server-stats"
 version = "0.1.1"
 dependencies = [
+ "actix",
  "actix-governor",
  "actix-web",
+ "actix-web-actors",
  "actix-web-lab",
  "askama",
  "askama_actix",
diff --git a/Cargo.toml b/Cargo.toml
index 6ed6421..1f69dcd 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -24,7 +24,9 @@ path = "src/main.rs"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+actix = "0.13.0"
 actix-web = "4"
+actix-web-actors = "4.1.0"
 #config = "0.13.1"
 config = "0.12"
 deadpool-postgres = { version = "0.10.2", features = ["serde"] }
@@ -49,3 +51,4 @@ systemstat = "0.2.0"
 
 nats = "0.23.0" #https://github.com/nats-io/nats.rs
 kafka = "0.9" #https://github.com/kafka-rust/kafka-rust
+
diff --git a/src/handlers/handlers.rs b/src/handlers/handlers.rs
index 2dc5883..68176a0 100644
--- a/src/handlers/handlers.rs
+++ b/src/handlers/handlers.rs
@@ -1,9 +1,13 @@
 use actix_web::http::StatusCode;
-use actix_web::{get, web, Error, HttpResponse};
+use actix_web::{get, web, App, Error, HttpRequest, HttpResponse, HttpServer};
 use askama::Template;
 use deadpool_postgres::{Pool};
 use minify::html::minify;
 
+
+use actix::{Actor, StreamHandler};
+use actix_web_actors::ws;
+
 use crate::Stats;
 
 extern crate minify;
@@ -86,3 +90,36 @@ pub async fn index_page() -> Result<HttpResponse, Error> {
         .content_type("text/html; charset=utf-8")
         .body(html_str))
 }
+
+/**
+==============================
+     WebSocket Endpoints
+==============================
+ */
+
+/// Define HTTP actor
+struct MyWs;
+
+impl Actor for MyWs {
+    type Context = ws::WebsocketContext<Self>;
+}
+
+/// Handler for ws::Message message
+impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for MyWs {
+    fn handle(&mut self, msg: Result<ws::Message, ws::ProtocolError>, ctx: &mut Self::Context) {
+        match msg {
+            Ok(ws::Message::Ping(msg)) => ctx.pong(&msg),
+            Ok(ws::Message::Text(text)) => ctx.text(text),
+            Ok(ws::Message::Binary(bin)) => ctx.binary(bin),
+            _ => (),
+        }
+    }
+}
+
+#[get("/ws")]
+pub async fn ws_stats_index(req: HttpRequest, stream: web::Payload) -> Result<HttpResponse, Error> {
+    let resp = ws::start(MyWs {}, &req, stream);
+    println!("{:?}", resp);
+    resp
+}
+
diff --git a/src/main.rs b/src/main.rs
index d1ef730..8cc1c9a 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -9,9 +9,11 @@ use std::env;
 use ::config::Config;
 use actix_governor::{Governor, GovernorConfigBuilder};
 use actix_web::{middleware::Logger, web, App, HttpServer};
+use actix_web::web::ServiceConfig;
 use actix_web_lab::web as web_lab;
 use dotenv::dotenv;
 use handlers::index_page;
+use handlers::ws_stats_index;
 use tokio_postgres::NoTls;
 
 use crate::config::ServerConfig;
@@ -20,6 +22,7 @@ use stats::Stats;
 
 fn scoped_config(cfg: &mut web::ServiceConfig) {
     cfg.service(index_page);
+    cfg.service(ws_stats_index);
 }
 
 #[actix_web::main]