query tracing

Author: mangas

Cargo.lock

@@ -34,7 +34,15 @@ bs58 = "0.5.1"
 chrono = "0.4.38"
 clap = { version = "4.5.4", features = ["derive", "env"] }
 derivative = "2.2.0"
-diesel = { version = "2.2.4", features = ["postgres", "serde_json", "numeric", "r2d2", "chrono", "uuid", "i-implement-a-third-party-backend-and-opt-into-breaking-changes"] }
+diesel = { version = "2.2.4", features = [
+    "postgres",
+    "serde_json",
+    "numeric",
+    "r2d2",
+    "chrono",
+    "uuid",
+    "i-implement-a-third-party-backend-and-opt-into-breaking-changes",
+] }
 diesel-derive-enum = { version = "2.1.0", features = ["postgres"] }
 diesel-dynamic-schema = { version = "0.2.1", features = ["postgres"] }
 diesel_derives = "2.1.4"
@@ -56,19 +64,24 @@ serde_derive = "1.0.125"
 serde_json = { version = "1.0", features = ["arbitrary_precision"] }
 serde_regex = "1.1.0"
 serde_yaml = "0.9.21"
-slog = { version = "2.7.0", features = ["release_max_level_trace", "max_level_trace"] }
+slog = { version = "2.7.0", features = [
+    "release_max_level_trace",
+    "max_level_trace",
+] }
 sqlparser = "0.46.0"
 strum = { version = "0.26", features = ["derive"] }
 syn = { version = "2.0.66", features = ["full"] }
 test-store = { path = "./store/test-store" }
 thiserror = "1.0.25"
-tokio = { version = "1.38.0", features = ["full"] }
+tokio = { version = "1", features = ["full"] }
 tonic = { version = "0.11.0", features = ["tls-roots", "gzip"] }
 tonic-build = { version = "0.11.0", features = ["prost"] }
 tower-http = { version = "0.5.2", features = ["cors"] }
 wasmparser = "0.118.1"
 wasmtime = "15.0.1"
 
+indexer-watcher = { git = "https://github.com/graphprotocol/indexer-rs", tag = "indexer-service-rs-v1.4.2" }
+
 # Incremental compilation on Rust 1.58 causes an ICE on build. As soon as graph node builds again, these can be removed.
 [profile.test]
 incremental = false

Cargo.toml

@@ -12,7 +12,9 @@ atomic_refcell = "0.1.13"
 # We require this precise version of bigdecimal. Updating to later versions
 # has caused PoI differences; if you update this version, you will need to
 # make sure that it does not cause PoI changes
-old_bigdecimal = { version = "=0.1.2", features = ["serde"], package = "bigdecimal" }
+old_bigdecimal = { version = "=0.1.2", features = [
+    "serde",
+], package = "bigdecimal" }
 bytes = "1.0.1"
 bs58 = { workspace = true }
 cid = "0.11.1"
@@ -64,14 +66,7 @@ slog-envlogger = "2.1.0"
 slog-term = "2.7.0"
 petgraph = "0.6.5"
 tiny-keccak = "1.5.0"
-tokio = { version = "1.38.0", features = [
-    "time",
-    "sync",
-    "macros",
-    "test-util",
-    "rt-multi-thread",
-    "parking_lot",
-] }
+tokio.workspace = true
 tokio-stream = { version = "0.1.15", features = ["sync"] }
 tokio-retry = "0.3.0"
 toml = "0.8.8"
@@ -92,11 +87,13 @@ defer = "0.2"
 # Our fork contains patches to make some fields optional for Celo and Fantom compatibility.
 # Without the "arbitrary_precision" feature, we get the error `data did not match any variant of untagged enum Response`.
 web3 = { git = "https://github.com/graphprotocol/rust-web3", branch = "graph-patches-onto-0.18", features = [
-    "arbitrary_precision", "test"
+    "arbitrary_precision",
+    "test",
 ] }
 serde_plain = "1.0.2"
 csv = "1.3.0"
 object_store = { version = "0.11.0", features = ["gcp"] }
+indexer-watcher.workspace = true
 
 [dev-dependencies]
 clap.workspace = true

graph/Cargo.toml

@@ -25,4 +25,10 @@ fn main() {
         .out_dir("src/substreams_rpc")
         .compile(&["proto/substreams-rpc.proto"], &["proto"])
         .expect("Failed to compile Substreams RPC proto(s)");
+
+    tonic_build::configure()
+        .out_dir("src/grpc/pb")
+        .include_file("mod.rs")
+        .compile(&["proto/tracing.proto"], &["proto"])
+        .expect("Failed to compile Tracing proto(s)");
 }

graph/build.rs

@@ -0,0 +1,17 @@
+syntax = "proto3";
+
+package graph.tracing.v1;
+
+service Stream {
+  rpc QueryTrace(Request) returns (stream Trace);
+}
+
+message Request {
+    int32 deployment_id = 1;
+}
+
+message Trace {
+    int32 deployment_id = 1;
+    string query = 2;
+    uint64 duration_millis = 3;
+}

graph/proto/tracing.proto

@@ -60,6 +60,8 @@ pub mod metrics;
 /// Components dealing with versioning
 pub mod versions;
 
+pub mod tracing;
+
 /// A component that receives events of type `T`.
 pub trait EventConsumer<E> {
     /// Get the event sink.

graph/src/components/mod.rs

@@ -0,0 +1,127 @@
+use std::{collections::HashMap, sync::Arc, time::Duration};
+
+use tokio::sync::{mpsc, watch::Receiver, RwLock};
+
+use super::store::DeploymentId;
+
+const DEFAULT_BUFFER_SIZE: usize = 100;
+
+#[derive(Debug, Clone)]
+pub struct Subscriptions<T> {
+    inner: Arc<RwLock<HashMap<DeploymentId, mpsc::Sender<T>>>>,
+}
+
+impl<T> Default for Subscriptions<T> {
+    fn default() -> Self {
+        Self {
+            inner: Arc::new(RwLock::new(HashMap::new())),
+        }
+    }
+}
+
+pub type SubscriptionsWatcher<T> = Receiver<Subscriptions<T>>;
+
+/// A control structure for managing tracing subscriptions.
+#[derive(Debug)]
+pub struct TracingControl<T> {
+    watcher: Receiver<HashMap<DeploymentId, mpsc::Sender<T>>>,
+    subscriptions: Subscriptions<T>,
+    default_buffer_size: usize,
+}
+
+impl<T: Send + Clone + 'static> Default for TracingControl<T> {
+    fn default() -> Self {
+        let subscriptions = Subscriptions::default();
+        let subs = subscriptions.clone();
+        let watcher = std::thread::spawn(move || {
+            let runtime = tokio::runtime::Builder::new_current_thread()
+                .enable_all()
+                .build()
+                .unwrap();
+            runtime.block_on(indexer_watcher::new_watcher(
+                Duration::from_secs(30),
+                move || {
+                    let subs = subs.clone();
+
+                    async move { Ok(subs.inner.read().await.clone()) }
+                },
+            ))
+        })
+        .join()
+        .unwrap()
+        .unwrap();
+
+        Self {
+            subscriptions,
+            default_buffer_size: DEFAULT_BUFFER_SIZE,
+            watcher,
+        }
+    }
+}
+
+impl<T: Send + Clone + 'static> TracingControl<T> {
+    pub fn new(default_buffer_size: Option<usize>) -> Self {
+        Self {
+            default_buffer_size: default_buffer_size.unwrap_or(DEFAULT_BUFFER_SIZE),
+            ..Default::default()
+        }
+    }
+
+    pub fn producer(&self, key: DeploymentId) -> Option<mpsc::Sender<T>> {
+        self.watcher
+            .borrow()
+            .get(&key)
+            .cloned()
+            .filter(|sender| !sender.is_closed())
+    }
+
+    pub async fn subscribe_with_chan_size(
+        &self,
+        key: DeploymentId,
+        buffer_size: usize,
+    ) -> mpsc::Receiver<T> {
+        let (tx, rx) = mpsc::channel(buffer_size);
+        let mut guard = self.subscriptions.inner.write().await;
+        guard.insert(key, tx);
+
+        rx
+    }
+
+    /// Creates a new subscription for a given deployment ID. If a subscription already
+    /// exists, it will be replaced.
+    pub async fn subscribe(&self, key: DeploymentId) -> mpsc::Receiver<T> {
+        self.subscribe_with_chan_size(key, self.default_buffer_size)
+            .await
+    }
+}
+
+#[cfg(test)]
+mod test {
+
+    use super::*;
+    use std::sync::Arc;
+
+    #[tokio::test]
+    async fn test_tracing_control() {
+        let control: TracingControl<()> = TracingControl::default();
+        let control = Arc::new(control);
+
+        // produce before subscription
+        let tx = control.producer(DeploymentId(123));
+        assert!(tx.is_none());
+
+        // drop the subscription
+        let rx = control.subscribe(DeploymentId(123));
+        drop(rx);
+
+        // check subscription is none because channel is closed
+        let tx = control.producer(DeploymentId(123));
+        assert!(tx.is_none());
+
+        // re-create subscription
+        let _rx = control.subscribe(DeploymentId(123));
+        // check old subscription was replaced
+        let tx = control.producer(DeploymentId(123));
+        assert!(!tx.unwrap().is_closed())
+    }
+}