Add metrics for number of active sessions and sessions' durations

Author: adwinwhite

ui/Cargo.lock

@@ -10,7 +10,7 @@ fork-bomb-prevention = []
 
 [dependencies]
 async-trait = "0.1.52"
-axum = { version = "0.5", features = ["headers"] }
+axum = { version = "0.5", features = ["headers", "ws"] }
 dotenv = "0.15.0"
 env_logger = "0.9.0"
 futures = "0.3.21"

ui/Cargo.toml

@@ -20,5 +20,6 @@ interface Window {
   __REDUX_DEVTOOLS_EXTENSION_COMPOSE__: any;
   rustPlayground: {
     setCode(code: string): void;
+    webSocket: WebSocket | null;
   };
 }

ui/frontend/declarations.d.ts

@@ -25,6 +25,9 @@ import playgroundApp from './reducers';
 import Router from './Router';
 import configureStore from './configureStore';
 
+const wsUri = '/websocket';
+const socket = new WebSocket(wsUri);
+
 const store = configureStore(window);
 
 const z = (evt: MediaQueryList | MediaQueryListEvent) => { store.dispatch(browserWidthChanged(evt.matches)); };
@@ -49,6 +52,9 @@ window.rustPlayground = {
   setCode: code => {
     store.dispatch(editCode(code));
   },
+  // Temporarily storing this as a global to prevent it from being
+  // garbage collected (at least by Safari).
+  webSocket: socket,
 };
 
 const container = document.getElementById('playground');

ui/frontend/index.tsx

@@ -1,6 +1,9 @@
 use futures::future::BoxFuture;
 use lazy_static::lazy_static;
-use prometheus::{self, register_histogram_vec, HistogramVec};
+use prometheus::{
+    self, register_histogram, register_histogram_vec, register_int_gauge, Histogram, HistogramVec,
+    IntGauge,
+};
 use regex::Regex;
 use std::{future::Future, time::Instant};
 
@@ -14,6 +17,17 @@ lazy_static! {
         vec![0.1, 1.0, 2.5, 5.0, 10.0, 15.0]
     )
     .unwrap();
+    pub(crate) static ref LIVE_WS: IntGauge = register_int_gauge!(
+        "active_websocket_connections_count",
+        "Number of active WebSocket connections"
+    )
+    .unwrap();
+    pub(crate) static ref DURATION_WS: Histogram = register_histogram!(
+        "websocket_duration_seconds",
+        "WebSocket connection length",
+        vec![15.0, 60.0, 300.0, 600.0, 1800.0, 3600.0, 7200.0]
+    )
+    .unwrap();
 }
 
 #[derive(Debug, Copy, Clone, strum::IntoStaticStr)]

ui/src/metrics.rs

@@ -2,7 +2,7 @@ use crate::{
     gist,
     metrics::{
         track_metric_async, track_metric_force_endpoint_async, track_metric_no_request_async,
-        Endpoint, GenerateLabels, SuccessDetails,
+        Endpoint, GenerateLabels, SuccessDetails, DURATION_WS, LIVE_WS,
     },
     sandbox::{self, Channel, Sandbox},
     CachingSnafu, ClippyRequest, ClippyResponse, CompilationSnafu, CompileRequest, CompileResponse,
@@ -15,7 +15,11 @@ use crate::{
 };
 use async_trait::async_trait;
 use axum::{
-    extract::{self, Extension, Path, TypedHeader},
+    extract::{
+        self,
+        ws::{WebSocket, WebSocketUpgrade},
+        Extension, Path, TypedHeader,
+    },
     handler::Handler,
     headers::{authorization::Bearer, Authorization, CacheControl, ETag, IfNoneMatch},
     http::{header, uri::PathAndQuery, HeaderValue, Method, Request, StatusCode, Uri},
@@ -78,6 +82,7 @@ pub(crate) async fn serve(config: Config) {
         .route("/meta/gist", post(meta_gist_create))
         .route("/meta/gist/:id", get(meta_gist_get))
         .route("/metrics", get(metrics))
+        .route("/websocket", get(websocket))
         .layer(Extension(Arc::new(SandboxCache::default())))
         .layer(Extension(config.github_token()));
 
@@ -386,6 +391,19 @@ async fn metrics(_: MetricsAuthorization) -> Result<Vec<u8>, StatusCode> {
         .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)
 }
 
+async fn websocket(ws: WebSocketUpgrade) -> impl IntoResponse {
+    ws.on_upgrade(handle_socket)
+}
+
+async fn handle_socket(mut socket: WebSocket) {
+    LIVE_WS.inc();
+    let start = Instant::now();
+    while let Some(Ok(_msg)) = socket.recv().await {}
+    LIVE_WS.dec();
+    let elapsed = start.elapsed();
+    DURATION_WS.observe(elapsed.as_secs_f64());
+}
+
 #[derive(Debug)]
 struct MetricsAuthorization;