Merge pull request #890 from adwinwhite/ws_metrics

Author: shepmaster

tests/spec/features/tools_spec.rb

@@ -61,7 +61,7 @@ def code_with_undefined_behavior
 
     within(:output, :stdout) do
       # First-party
-      expect(page).to have_content('core::fmt::Arguments::new_v1')
+      expect(page).to have_content('::std::io::_print')
 
       # Third-party procedural macro
       expect(page).to have_content('block_on(body)')

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

@@ -24,6 +24,11 @@ import PageSwitcher from './PageSwitcher';
 import playgroundApp from './reducers';
 import Router from './Router';
 import configureStore from './configureStore';
+import openWebSocket from './websocket';
+
+const params = new URLSearchParams(window.location.search);
+// openWebSocket() may return null.
+const socket = params.has('websocket') ? openWebSocket(window.location) : null;
 
 const store = configureStore(window);
 
@@ -49,6 +54,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

@@ -0,0 +1,23 @@
+export default function openWebSocket(currentLocation: Location) {
+  try {
+    const wsProtocol = currentLocation.protocol === 'https:' ? 'wss://' : 'ws://';
+    const wsUri = [wsProtocol, currentLocation.host, '/websocket'].join('');
+    return new WebSocket(wsUri);
+  } catch {
+    // WebSocket URL error or WebSocket is not supported by browser.
+    // Assume it's the second case since URL error is easy to notice.
+    (async () => {
+      try {
+        await fetch('/nowebsocket', {
+          method: 'post',
+          headers: {
+            'Content-Length': '0',
+          },
+        });
+      } catch (e) {
+        console.log('Error:', e);
+      }
+    })();
+    return null;
+  }
+}

ui/frontend/websocket.ts

@@ -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_counter, register_int_gauge,
+    Histogram, HistogramVec, IntCounter, IntGauge,
+};
 use regex::Regex;
 use std::{future::Future, time::Instant};
 
@@ -14,6 +17,22 @@ 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();
+    pub(crate) static ref UNAVAILABLE_WS: IntCounter = register_int_counter!(
+        "websocket_unavailability_count",
+        "Number of failed WebSocket connections"
+    )
+    .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, UNAVAILABLE_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,8 @@ 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))
+        .route("/nowebsocket", post(nowebsocket))
         .layer(Extension(Arc::new(SandboxCache::default())))
         .layer(Extension(config.github_token()));
 
@@ -386,6 +392,23 @@ 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());
+}
+
+async fn nowebsocket() {
+    UNAVAILABLE_WS.inc();
+}
+
 #[derive(Debug)]
 struct MetricsAuthorization;