Skip to content

Commit

Permalink
fix (#9)
Browse files Browse the repository at this point in the history
  • Loading branch information
kostekIV authored Jan 12, 2024
1 parent 7218798 commit 41edf62
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 6 deletions.
9 changes: 6 additions & 3 deletions src/extensions/server/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,9 @@ use jsonrpsee::Methods;
use prometheus_endpoint::Registry;
use serde::ser::StdError;
use serde::Deserialize;
use std::collections::HashMap;
use std::str::FromStr;
use std::sync::Arc;
use std::sync::{Arc, Mutex};
use std::{future::Future, net::SocketAddr};
use tower::layer::layer_fn;
use tower::ServiceBuilder;
Expand All @@ -23,7 +24,7 @@ use crate::extensions::rate_limit::{MethodWeights, RateLimitBuilder, XFF};
mod prometheus;
mod proxy_get_request;
mod ready_get_request;
use crate::extensions::server::prometheus::PrometheusService;
use crate::extensions::server::prometheus::{MetricPair, PrometheusService};
use proxy_get_request::{ProxyGetRequestLayer, ProxyGetRequestMethod};
use ready_get_request::ReadyProxyLayer;

Expand Down Expand Up @@ -112,6 +113,7 @@ impl SubwayServerBuilder {
let (stop_handle, server_handle) = stop_channel();
let handle = stop_handle.clone();
let rpc_module = rpc_module_builder().await?;
let metrics: Arc<Mutex<HashMap<String, MetricPair>>> = Default::default();

// make_service handle each connection
let make_service = make_service_fn(move |socket: &AddrStream| {
Expand Down Expand Up @@ -139,6 +141,7 @@ impl SubwayServerBuilder {
let rate_limit_builder = rate_limit_builder.clone();
let rpc_method_weights = rpc_method_weights.clone();
let prometheus_registry = prometheus_registry.clone();
let metrics = metrics.clone();

async move {
// service_fn handle each request
Expand Down Expand Up @@ -166,7 +169,7 @@ impl SubwayServerBuilder {
.option_layer(
prometheus_registry
.as_ref()
.map(|r| layer_fn(|s| PrometheusService::new(s, r))),
.map(|r| layer_fn(|s| PrometheusService::new(s, r, metrics.clone()))),
);

let service_builder = ServerBuilder::default()
Expand Down
6 changes: 3 additions & 3 deletions src/extensions/server/prometheus.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ use prometheus_endpoint::{register, Counter, Histogram, HistogramOpts, Registry,
use std::collections::HashMap;
use std::sync::{Arc, Mutex};

type MetricPair = (Counter<U64>, Histogram);
pub type MetricPair = (Counter<U64>, Histogram);

#[derive(Clone)]
pub struct PrometheusService<S> {
Expand All @@ -16,11 +16,11 @@ pub struct PrometheusService<S> {
}

impl<S> PrometheusService<S> {
pub fn new(inner: S, registry: &Registry) -> Self {
pub fn new(inner: S, registry: &Registry, call_metrics: Arc<Mutex<HashMap<String, MetricPair>>>) -> Self {
Self {
inner,
call_metrics,
registry: registry.clone(),
call_metrics: Arc::new(Mutex::new(HashMap::new())),
}
}

Expand Down

0 comments on commit 41edf62

Please sign in to comment.