Merge pull request #309 from haardikk21/brotli-compression-on-proxy

ws proxy: brotli compress each message from upstream before sending downstream

Author: avalonche

Cargo.lock

@@ -32,6 +32,7 @@ redis = "0.30.0"
 redis-test = { version = "0.10.0", optional = true }
 uuid = { version = "1.16.0", features = ["v4"] }
 tokio-util = "0.7.12"
+brotli = "8.0.1"
 
 [dependencies.ring]
 version = "0.17.12"

crates/websocket-proxy/Cargo.toml

@@ -62,3 +62,8 @@ When Redis is enabled, the following features are available:
 
 If the Redis connection fails, the proxy will automatically fall back to in-memory rate limiting.
 
+### Brotli Compression
+
+The proxy supports compressing messages to downstream clients using Brotli.
+
+To enable this, pass the parameter `--enable-compression`

crates/websocket-proxy/README.md

@@ -18,8 +18,8 @@ impl ClientConnection {
         }
     }
 
-    pub async fn send(&mut self, data: String) -> Result<(), Error> {
-        self.websocket.send(data.into_bytes().into()).await
+    pub async fn send(&mut self, data: Vec<u8>) -> Result<(), Error> {
+        self.websocket.send(data.into()).await
     }
 
     pub fn id(&self) -> String {

crates/websocket-proxy/src/client.rs

@@ -27,7 +27,7 @@ mod test {
         server: Server,
         server_addr: SocketAddr,
         client_id_to_handle: HashMap<usize, JoinHandle<()>>,
-        sender: Sender<String>,
+        sender: Sender<Vec<u8>>,
     }
 
     impl TestHarness {
@@ -153,7 +153,10 @@ mod test {
         }
 
         fn send_messages(&mut self, messages: Vec<&str>) {
-            let messages: Vec<String> = messages.into_iter().map(String::from).collect();
+            let messages: Vec<Vec<u8>> = messages
+                .into_iter()
+                .map(|m| m.as_bytes().to_vec())
+                .collect();
 
             for message in messages.iter() {
                 match self.sender.send(message.clone()) {

crates/websocket-proxy/src/integration.rs

@@ -18,6 +18,7 @@ use clap::Parser;
 use dotenvy::dotenv;
 use metrics_exporter_prometheus::PrometheusBuilder;
 use rate_limit::RedisRateLimit;
+use std::io::Write;
 use std::net::SocketAddr;
 use std::sync::Arc;
 use tokio::signal::unix::{signal, SignalKind};
@@ -68,6 +69,13 @@ struct Args {
         help = "Maximum number of concurrently connected clients"
     )]
     per_ip_connections_limit: usize,
+    #[arg(
+        long,
+        env,
+        default_value = "false",
+        help = "Enable brotli compression on messages to downstream clients"
+    )]
+    enable_compression: bool,
 
     #[arg(
         long,
@@ -208,7 +216,20 @@ async fn main() {
             .active_connections
             .set((send.receiver_count() - 1) as f64);
 
-        match send.send(data) {
+        let message_data = if args.enable_compression {
+            let data_bytes = data.as_bytes();
+            let mut compressed_data_bytes = Vec::new();
+            {
+                let mut compressor =
+                    brotli::CompressorWriter::new(&mut compressed_data_bytes, 4096, 5, 22);
+                compressor.write_all(data_bytes).unwrap();
+            }
+            compressed_data_bytes
+        } else {
+            data.into_bytes()
+        };
+
+        match send.send(message_data) {
             Ok(_) => (),
             Err(e) => error!(message = "failed to send data", error = e.to_string()),
         }

crates/websocket-proxy/src/main.rs

@@ -7,12 +7,12 @@ use tracing::{info, trace, warn};
 
 #[derive(Clone)]
 pub struct Registry {
-    sender: Sender<String>,
+    sender: Sender<Vec<u8>>,
     metrics: Arc<Metrics>,
 }
 
 impl Registry {
-    pub fn new(sender: Sender<String>, metrics: Arc<Metrics>) -> Self {
+    pub fn new(sender: Sender<Vec<u8>>, metrics: Arc<Metrics>) -> Self {
         Self { sender, metrics }
     }