fix: make the max number of ws-connections configurable (#397)

Introduce new paramater `max-ws-connections` to the [rpc] section of the
config, by default this value is set to 16384, but setting it to
anything below 200K is also fine, anything above that value will create
significant memory pressure by slab preallocation, and can even trigger
an OOM kill.

 

<!-- greptile_comment -->

## Greptile Summary

Adds configurable WebSocket connection limits to prevent memory
exhaustion, with a new `max-ws-connections` parameter in the RPC config
defaulting to 16384 connections.

- Added `max_ws_connections` field to
`magicblock-pubsub/src/pubsub_service.rs` to control concurrent
WebSocket connections
- Modified `magicblock-config/src/rpc.rs` to expose the new config
parameter with default of 16384
- Added validation to keep connections under 200K to prevent slab
preallocation memory pressure
- Updated config test files to verify proper parsing and defaults
- Integrated limit through PubsubConfig initialization in
`magicblock-api/src/magic_validator.rs`



<!-- /greptile_comment -->

Author: bmuddha

magicblock-api/src/magic_validator.rs

@@ -325,6 +325,7 @@ impl MagicValidator {
         let pubsub_config = PubsubConfig::from_rpc(
             config.validator_config.rpc.addr,
             config.validator_config.rpc.port,
+            config.validator_config.rpc.max_ws_connections,
         );
         validator::init_validator_authority(identity_keypair);
 

magicblock-config/src/rpc.rs

@@ -3,7 +3,7 @@ use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 use serde::{Deserialize, Serialize};
 
 #[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
-#[serde(deny_unknown_fields)]
+#[serde(deny_unknown_fields, rename_all = "kebab-case")]
 pub struct RpcConfig {
     #[serde(
         default = "default_addr",
@@ -13,13 +13,16 @@ pub struct RpcConfig {
     pub addr: IpAddr,
     #[serde(default = "default_port")]
     pub port: u16,
+    #[serde(default = "default_max_ws_connections")]
+    pub max_ws_connections: usize,
 }
 
 impl Default for RpcConfig {
     fn default() -> Self {
         Self {
             addr: default_addr(),
             port: default_port(),
+            max_ws_connections: default_max_ws_connections(),
         }
     }
 }
@@ -57,3 +60,7 @@ fn default_addr() -> IpAddr {
 fn default_port() -> u16 {
     8899
 }
+
+fn default_max_ws_connections() -> usize {
+    16384
+}

magicblock-config/tests/parse_config.rs

@@ -97,7 +97,8 @@ fn test_local_dev_with_programs_toml() {
             }],
             rpc: RpcConfig {
                 addr: IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
-                port: 7799
+                port: 7799,
+                max_ws_connections: 16384
             },
             validator: ValidatorConfig {
                 millis_per_slot: 14,

magicblock-config/tests/read_config.rs

@@ -45,6 +45,7 @@ fn test_load_local_dev_with_programs_toml() {
             rpc: RpcConfig {
                 addr: IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
                 port: 7799,
+                max_ws_connections: 16384
             },
             geyser_grpc: GeyserGrpcConfig {
                 addr: IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
@@ -128,6 +129,7 @@ fn test_load_local_dev_with_programs_toml_envs_override() {
             rpc: RpcConfig {
                 addr: IpAddr::V4(Ipv4Addr::new(0, 1, 0, 1)),
                 port: 123,
+                max_ws_connections: 16384
             },
             geyser_grpc: GeyserGrpcConfig {
                 addr: IpAddr::V4(Ipv4Addr::new(0, 1, 0, 1)),

magicblock-pubsub/src/pubsub_service.rs

@@ -27,12 +27,18 @@ use crate::{
 #[derive(Clone)]
 pub struct PubsubConfig {
     socket: SocketAddr,
+    max_connections: usize,
 }
 
 impl PubsubConfig {
-    pub fn from_rpc(rpc_addr: IpAddr, rpc_port: u16) -> Self {
+    pub fn from_rpc(
+        rpc_addr: IpAddr,
+        rpc_port: u16,
+        max_connections: usize,
+    ) -> Self {
         Self {
             socket: SocketAddr::new(rpc_addr, rpc_port + 1),
+            max_connections,
         }
     }
 }
@@ -41,6 +47,7 @@ impl Default for PubsubConfig {
     fn default() -> Self {
         Self {
             socket: SocketAddr::from(([0, 0, 0, 0], DEFAULT_RPC_PUBSUB_PORT)),
+            max_connections: 16384,
         }
     }
 }
@@ -89,6 +96,10 @@ impl PubsubService {
             |context: &RequestContext| Arc::new(Session::new(context.sender()));
 
         ServerBuilder::with_meta_extractor(self.io, extractor)
+            // NOTE: we just set the max number of allowed connections to a reasonably high value
+            // to satisfy most of the use cases, however this number cannot be arbitrarily large
+            // due to the preallocation involved, and a large value will trigger an OOM Kill
+            .max_connections(self.config.max_connections)
             .start(&self.config.socket)
     }