use tetsy_jsonrpc_core::{
Middleware as RequestMiddleware, Metadata,
Request, Response, FutureResponse, FutureOutput
};
use prometheus_endpoint::{
Registry, CounterVec, PrometheusError,
Opts, register, U64
};
use futures::{future::Either, Future};
#[derive(Debug, Clone)]
pub struct RpcMetrics {
rpc_calls: Option<CounterVec<U64>>,
}
impl RpcMetrics {
pub fn new(metrics_registry: Option<&Registry>) -> Result<Self, PrometheusError> {
Ok(Self {
rpc_calls: metrics_registry.map(|r|
register(
CounterVec::new(
Opts::new(
"rpc_calls_total",
"Number of rpc calls received",
),
&["protocol"]
)?,
r,
)
).transpose()?,
})
}
}
pub struct RpcMiddleware {
metrics: RpcMetrics,
transport_label: String,
}
impl RpcMiddleware {
pub fn new(metrics: RpcMetrics, transport_label: &str) -> Self {
RpcMiddleware {
metrics,
transport_label: String::from(transport_label),
}
}
}
impl<M: Metadata> RequestMiddleware<M> for RpcMiddleware {
type Future = FutureResponse;
type CallFuture = FutureOutput;
fn on_request<F, X>(&self, request: Request, meta: M, next: F) -> Either<FutureResponse, X>
where
F: Fn(Request, M) -> X + Send + Sync,
X: Future<Item = Option<Response>, Error = ()> + Send + 'static,
{
if let Some(ref rpc_calls) = self.metrics.rpc_calls {
rpc_calls.with_label_values(&[self.transport_label.as_str()]).inc();
}
Either::B(next(request, meta))
}
}