feat(request): allow to use custom id generator on client

Author: PanGan21

client/http-client/src/client.rs

@@ -276,7 +276,7 @@ where
 
 		Ok(HttpClient {
 			transport,
-			id_manager: Arc::new(RequestIdManager::new(id_kind)),
+			id_manager: Arc::new(RequestIdManager::new(id_kind, None)),
 			request_timeout,
 			request_guard,
 		})

core/src/client/async_client/mod.rs

@@ -30,7 +30,8 @@ mod helpers;
 mod manager;
 mod utils;
 
-use crate::client::async_client::helpers::{process_subscription_close_response, InnerBatchResponse};
+use crate::JsonRawValue;
+use crate::client::async_client::helpers::{InnerBatchResponse, process_subscription_close_response};
 use crate::client::async_client::utils::MaybePendingFutures;
 use crate::client::{
 	BatchMessage, BatchResponse, ClientT, Error, ReceivedMessage, RegisterNotificationMessage, RequestMessage,
@@ -40,31 +41,30 @@ use crate::error::RegisterMethodError;
 use crate::params::{BatchRequestBuilder, EmptyBatchRequest};
 use crate::tracing::client::{rx_log_from_json, tx_log_from_str};
 use crate::traits::ToRpcParams;
-use crate::JsonRawValue;
 use std::borrow::Cow as StdCow;
 
 use core::time::Duration;
 use helpers::{
 	build_unsubscribe_message, call_with_timeout, process_batch_response, process_notification,
 	process_single_response, process_subscription_response, stop_subscription,
 };
-use jsonrpsee_types::{InvalidRequestId, ResponseSuccess, TwoPointZero};
+use jsonrpsee_types::{IdGeneratorFn, InvalidRequestId, ResponseSuccess, TwoPointZero};
 use manager::RequestManager;
 use std::sync::Arc;
 
 use async_trait::async_trait;
 use futures_timer::Delay;
+use futures_util::Stream;
 use futures_util::future::{self, Either};
 use futures_util::stream::StreamExt;
-use futures_util::Stream;
 use jsonrpsee_types::response::{ResponsePayload, SubscriptionError};
 use jsonrpsee_types::{NotificationSer, RequestSer, Response, SubscriptionResponse};
 use serde::de::DeserializeOwned;
 use tokio::sync::{mpsc, oneshot};
 use tracing::instrument;
 
 use self::utils::{InactivityCheck, IntervalStream};
-use super::{generate_batch_id_range, subscription_channel, FrontToBack, IdKind, RequestIdManager};
+use super::{FrontToBack, IdKind, RequestIdManager, generate_batch_id_range, subscription_channel};
 
 pub(crate) type Notification<'a> = jsonrpsee_types::Notification<'a, Option<serde_json::Value>>;
 
@@ -346,7 +346,7 @@ impl ClientBuilder {
 			to_back: to_back.clone(),
 			request_timeout: self.request_timeout,
 			error: ErrorFromBack::new(to_back, disconnect_reason),
-			id_manager: RequestIdManager::new(self.id_kind),
+			id_manager: RequestIdManager::new(self.id_kind, None),
 			max_log_length: self.max_log_length,
 			on_exit: Some(client_dropped_tx),
 		}
@@ -454,6 +454,11 @@ impl Client {
 	pub fn request_timeout(&self) -> Duration {
 		self.request_timeout
 	}
+
+	/// Registers a custom request id generator
+	pub fn with_request_id_generator(&mut self, request_id_generator: IdGeneratorFn) {
+		self.id_manager.custom_id_generator = Some(request_id_generator);
+	}
 }
 
 impl Drop for Client {

core/src/client/mod.rs

@@ -47,7 +47,7 @@ use crate::traits::ToRpcParams;
 use async_trait::async_trait;
 use core::marker::PhantomData;
 use futures_util::stream::{Stream, StreamExt};
-use jsonrpsee_types::{ErrorObject, Id, SubscriptionId};
+use jsonrpsee_types::{ErrorObject, Id, IdGeneratorFn, SubscriptionId};
 use serde::de::DeserializeOwned;
 use serde_json::Value as JsonValue;
 use tokio::sync::{mpsc, oneshot};
@@ -320,11 +320,7 @@ impl<Notif> Subscription<Notif> {
 			return None;
 		}
 
-		if lagged {
-			Some(SubscriptionCloseReason::Lagged)
-		} else {
-			Some(SubscriptionCloseReason::ConnectionClosed)
-		}
+		if lagged { Some(SubscriptionCloseReason::Lagged) } else { Some(SubscriptionCloseReason::ConnectionClosed) }
 	}
 }
 
@@ -459,17 +455,23 @@ pub struct RequestIdManager {
 	current_id: CurrentId,
 	/// Request ID type.
 	id_kind: IdKind,
+	/// Custom ID generator.
+	custom_id_generator: Option<IdGeneratorFn>,
 }
 
 impl RequestIdManager {
 	/// Create a new `RequestIdGuard` with the provided concurrency limit.
-	pub fn new(id_kind: IdKind) -> Self {
-		Self { current_id: CurrentId::new(), id_kind }
+	pub fn new(id_kind: IdKind, custom_id_generator: Option<IdGeneratorFn>) -> Self {
+		Self { current_id: CurrentId::new(), id_kind, custom_id_generator }
 	}
 
 	/// Attempts to get the next request ID.
 	pub fn next_request_id(&self) -> Id<'static> {
-		self.id_kind.into_id(self.current_id.next())
+		if let Some(ref generator) = self.custom_id_generator {
+			generator.call()
+		} else {
+			self.id_kind.into_id(self.current_id.next()) // Default behavior
+		}
 	}
 
 	/// Get a handle to the `IdKind`.

examples/examples/core_client_with_request_id.rs

@@ -0,0 +1,75 @@
+// Copyright 2019-2021 Parity Technologies (UK) Ltd.
+//
+// Permission is hereby granted, free of charge, to any
+// person obtaining a copy of this software and associated
+// documentation files (the "Software"), to deal in the
+// Software without restriction, including without
+// limitation the rights to use, copy, modify, merge,
+// publish, distribute, sublicense, and/or sell copies of
+// the Software, and to permit persons to whom the Software
+// is furnished to do so, subject to the following
+// conditions:
+//
+// The above copyright notice and this permission notice
+// shall be included in all copies or substantial portions
+// of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
+// ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
+// TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
+// PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
+// SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
+// CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
+// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
+// IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+// DEALINGS IN THE SOFTWARE.
+
+use std::net::SocketAddr;
+use std::time::{SystemTime, UNIX_EPOCH};
+
+use jsonrpsee::client_transport::ws::{Url, WsTransportClientBuilder};
+use jsonrpsee::core::client::{Client, ClientBuilder, ClientT};
+use jsonrpsee::rpc_params;
+use jsonrpsee::server::{RpcModule, Server};
+use jsonrpsee::types::{Id, IdGeneratorFn};
+
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+	tracing_subscriber::FmtSubscriber::builder()
+		.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
+		.try_init()
+		.expect("setting default subscriber failed");
+
+	let addr = run_server().await?;
+	let uri = Url::parse(&format!("ws://{}", addr))?;
+
+	let (tx, rx) = WsTransportClientBuilder::default().build(uri).await?;
+	let mut client: Client = ClientBuilder::default().build_with_tokio(tx, rx);
+
+	let custom_generator = Box::new(|| {
+		let timestamp_in_seconds = SystemTime::now().duration_since(UNIX_EPOCH).expect("Time went backwards").as_secs();
+		Id::Number(timestamp_in_seconds)
+	});
+
+	client.with_request_id_generator(IdGeneratorFn::new(custom_generator));
+
+	let response: String = client.request("say_hello", rpc_params![]).await?;
+	tracing::info!("response: {:?}", response);
+
+	Ok(())
+}
+
+async fn run_server() -> anyhow::Result<SocketAddr> {
+	let server = Server::builder().build("127.0.0.1:0").await?;
+	let mut module = RpcModule::new(());
+	module.register_method("say_hello", |_, _, _| "lo")?;
+	let addr = server.local_addr()?;
+
+	let handle = server.start(module);
+
+	// In this example we don't care about doing shutdown so let's it run forever.
+	// You may use the `ServerHandle` to shut it down or manage it yourself.
+	tokio::spawn(handle.stopped());
+
+	Ok(addr)
+}

types/src/lib.rs

@@ -45,5 +45,5 @@ pub mod error;
 
 pub use error::{ErrorCode, ErrorObject, ErrorObjectOwned};
 pub use params::{Id, InvalidRequestId, Params, ParamsSequence, SubscriptionId, TwoPointZero};
-pub use request::{InvalidRequest, Notification, NotificationSer, Request, RequestSer};
+pub use request::{IdGeneratorFn, InvalidRequest, Notification, NotificationSer, Request, RequestSer};
 pub use response::{Response, ResponsePayload, SubscriptionPayload, SubscriptionResponse, Success as ResponseSuccess};

types/src/request.rs

@@ -27,7 +27,10 @@
 //! Types to handle JSON-RPC requests according to the [spec](https://www.jsonrpc.org/specification#request-object).
 //! Some types come with a "*Ser" variant that implements [`serde::Serialize`]; these are used in the client.
 
-use std::borrow::Cow;
+use std::{
+	borrow::Cow,
+	fmt::{Debug, Formatter, Result},
+};
 
 use crate::{
 	Params,
@@ -173,6 +176,27 @@ impl<'a> NotificationSer<'a> {
 	}
 }
 
+/// Custom id generator function
+pub struct IdGeneratorFn(Box<dyn Fn() -> Id<'static> + Send + Sync>);
+
+impl IdGeneratorFn {
+	/// Creates a new `IdGeneratorFn`
+	pub fn new(generator: Box<dyn Fn() -> Id<'static> + Send + Sync>) -> Self {
+		IdGeneratorFn(generator)
+	}
+
+	/// Calls the id generator function
+	pub fn call(&self) -> Id<'static> {
+		(self.0)()
+	}
+}
+
+impl Debug for IdGeneratorFn {
+	fn fmt(&self, f: &mut Formatter<'_>) -> Result {
+		f.write_str("<CustomIdGenerator>")
+	}
+}
+
 #[cfg(test)]
 mod test {
 	use super::{Cow, Id, InvalidRequest, Notification, NotificationSer, Request, RequestSer, TwoPointZero};