Tag DB connections in VM runner

matter-labs · Sep 10, 2024 · b85071a · b85071a
1 parent 0c6dd0b
commit b85071a
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 7 deletions.
diff --git a/core/node/vm_runner/src/impls/bwip.rs b/core/node/vm_runner/src/impls/bwip.rs
@@ -168,7 +168,7 @@ impl OutputHandler for BasicWitnessInputProducerOutputHandler {
     )]
     async fn handle_l1_batch(self: Box<Self>, output: Arc<L1BatchOutput>) -> anyhow::Result<()> {
         let l1_batch_number = self.l1_batch_number;
-        let mut connection = self.pool.connection().await?;
+        let mut connection = self.pool.connection_tagged("bwip").await?;
 
         tracing::info!(%l1_batch_number, "Started saving VM run data");
 

diff --git a/core/node/vm_runner/src/process.rs b/core/node/vm_runner/src/process.rs
@@ -88,7 +88,10 @@ impl VmRunner {
             .create_handler(batch_data.system_env, batch_data.l1_batch_env)
             .await?;
         self.io
-            .mark_l1_batch_as_processing(&mut self.pool.connection().await?, number)
+            .mark_l1_batch_as_processing(
+                &mut self.pool.connection_tagged("vm_runner").await?,
+                number,
+            )
             .await?;
 
         let latency = METRICS.run_vm_time.start();
@@ -145,7 +148,7 @@ impl VmRunner {
         let mut task_handles: Vec<(L1BatchNumber, JoinHandle<anyhow::Result<()>>)> = Vec::new();
         let mut next_batch = self
             .io
-            .latest_processed_batch(&mut self.pool.connection().await?)
+            .latest_processed_batch(&mut self.pool.connection_tagged("vm_runner").await?)
             .await?
             + 1;
         loop {
@@ -174,7 +177,7 @@ impl VmRunner {
 
             let last_ready_batch = self
                 .io
-                .last_ready_to_be_loaded_batch(&mut self.pool.connection().await?)
+                .last_ready_to_be_loaded_batch(&mut self.pool.connection_tagged("vm_runner").await?)
                 .await?;
             METRICS.last_ready_batch.set(last_ready_batch.0.into());
             if next_batch > last_ready_batch {

diff --git a/core/node/vm_runner/src/storage.rs b/core/node/vm_runner/src/storage.rs
@@ -49,7 +49,7 @@ pub(crate) struct PostgresLoader {
 impl PostgresLoader {
     pub async fn new(pool: ConnectionPool<Core>, chain_id: L2ChainId) -> anyhow::Result<Self> {
         let mut l1_batch_params_provider = L1BatchParamsProvider::new();
-        let mut conn = pool.connection().await?;
+        let mut conn = pool.connection_tagged("vm_runner").await?;
         l1_batch_params_provider.initialize(&mut conn).await?;
         Ok(Self {
             pool,
@@ -72,7 +72,7 @@ impl StorageLoader for PostgresLoader {
         &self,
         l1_batch_number: L1BatchNumber,
     ) -> anyhow::Result<Option<(BatchExecuteData, OwnedStorage)>> {
-        let mut conn = self.pool.connection().await?;
+        let mut conn = self.pool.connection_tagged("vm_runner").await?;
         let Some(data) = load_batch_execute_data(
             &mut conn,
             l1_batch_number,
@@ -94,7 +94,7 @@ impl StorageLoader for PostgresLoader {
         tracing::info!(
             "Incomplete data to create storage snapshot for batch; will use sequential storage"
         );
-        let conn = self.pool.connection().await?;
+        let conn = self.pool.connection_tagged("vm_runner").await?;
         let storage = OwnedStorage::postgres(conn, l1_batch_number - 1).await?;
         Ok(Some((data, storage.into())))
     }