Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: expose ConnectionGuard as request extension #1443

Merged
merged 3 commits into from
Aug 19, 2024
Merged

feat: expose ConnectionGuard as request extension #1443

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
d92c6eb
feat: expose ConnectionGuard as request extension
dinhani Aug 16, 2024
2331c36
test: assert connection guard information via json-rpc interface
dinhani Aug 16, 2024
80f6c5c
chore: format code
dinhani Aug 16, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion server/src/server.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1047,7 +1047,9 @@ where
let curr_conns = max_conns - conn_guard.available_connections();
tracing::debug!(target: LOG_TARGET, "Accepting new connection {}/{}", curr_conns, max_conns);

request.extensions_mut().insert::<ConnectionId>(conn.conn_id.into());
let req_ext = request.extensions_mut();
req_ext.insert::<ConnectionGuard>(conn_guard.clone());
req_ext.insert::<ConnectionId>(conn.conn_id.into());
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cool, we could probably merge the ConnectionID and ConnectionGuard as some point to another type e.g. ConnectionDetails or something


let is_upgrade_request = is_upgrade_request(&request);

Expand Down
12 changes: 10 additions & 2 deletions tests/tests/helpers.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,8 +39,8 @@ use jsonrpsee::server::middleware::http::ProxyGetRequestLayer;

use jsonrpsee::server::middleware::rpc::RpcServiceT;
use jsonrpsee::server::{
serve_with_graceful_shutdown, stop_channel, PendingSubscriptionSink, RpcModule, RpcServiceBuilder, Server,
ServerBuilder, ServerHandle, SubscriptionMessage, TrySendError,
serve_with_graceful_shutdown, stop_channel, ConnectionGuard, PendingSubscriptionSink, RpcModule, RpcServiceBuilder,
Server, ServerBuilder, ServerHandle, SubscriptionMessage, TrySendError,
};
use jsonrpsee::types::{ErrorObject, ErrorObjectOwned};
use jsonrpsee::{Methods, SubscriptionCloseResponse};
Expand Down Expand Up @@ -165,6 +165,14 @@ pub async fn server() -> SocketAddr {
let mut module = RpcModule::new(());
module.register_method("say_hello", |_, _, _| "hello").unwrap();
module.register_method("get_connection_id", |_, _, ext| *ext.get::<u32>().unwrap()).unwrap();
module
.register_method("get_available_connections", |_, _, ext| {
ext.get::<ConnectionGuard>().unwrap().available_connections()
})
.unwrap();
module
.register_method("get_max_connections", |_, _, ext| ext.get::<ConnectionGuard>().unwrap().max_connections())
.unwrap();

module
.register_async_method("slow_hello", |_, _, _| async {
Expand Down
27 changes: 26 additions & 1 deletion tests/tests/integration_tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -207,7 +207,7 @@ async fn ws_method_call_works_over_proxy_stream() {
}

#[tokio::test]
async fn extensions_with_different_ws_clients() {
async fn connection_id_extension_with_different_ws_clients() {
init_logger();

let server_addr = server().await;
Expand All @@ -225,6 +225,31 @@ async fn extensions_with_different_ws_clients() {
assert_ne!(connection_id, second_connection_id);
}

#[tokio::test]
async fn connection_guard_extension_with_different_ws_clients() {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it's sufficient with one client but this does hurt, thanks :)

init_logger();

let server_addr = server().await;
let server_url = format!("ws://{}", server_addr);

// First connected retrieves initial information from ConnectionGuard.
let first_client = WsClientBuilder::default().build(&server_url).await.unwrap();
let first_max_connections: usize = first_client.request("get_max_connections", rpc_params![]).await.unwrap();
let first_available_connections: usize =
first_client.request("get_available_connections", rpc_params![]).await.unwrap();

assert_eq!(first_available_connections, first_max_connections - 1);

// Second client ensure max connections stays the same, but available connections is decreased.
let second_client = WsClientBuilder::default().build(&server_url).await.unwrap();
let second_max_connections: usize = second_client.request("get_max_connections", rpc_params![]).await.unwrap();
let second_available_connections: usize =
second_client.request("get_available_connections", rpc_params![]).await.unwrap();

assert_eq!(second_max_connections, first_max_connections);
assert_eq!(second_available_connections, second_max_connections - 2);
}

#[tokio::test]
async fn ws_method_call_str_id_works() {
init_logger();
Expand Down
Loading