chore(coprocessor): add pg_auto_explain_with_min_duration config in zkproof-worker

Author: goshawk-3

coprocessor/fhevm-engine/Cargo.lock

@@ -22,6 +22,7 @@ bincode = { workspace = true }
 thiserror = { workspace = true }
 tracing = { workspace = true }
 tracing-subscriber = { workspace = true }
+humantime = { workspace = true }
 
 # local dependencies
 fhevm-engine-common = { path = "../fhevm-engine-common" }

coprocessor/fhevm-engine/zkproof-worker/Cargo.toml

@@ -1,7 +1,8 @@
 use clap::{command, Parser};
 use fhevm_engine_common::healthz_server::HttpServer;
 use fhevm_engine_common::telemetry;
-use std::sync::Arc;
+use humantime::parse_duration;
+use std::{sync::Arc, time::Duration};
 use tokio::{join, task};
 use tokio_util::sync::CancellationToken;
 use tracing::{error, info, Level};
@@ -26,6 +27,15 @@ pub struct Args {
     #[arg(long, default_value_t = 5)]
     pub pg_pool_connections: u32,
 
+    /// Postgres acquire timeout
+    /// A longer timeout could affect the healthz/liveness updates
+    #[arg(long, default_value = "15s", value_parser = parse_duration)]
+    pub pg_timeout: Duration,
+
+    /// Postgres diagnostics: enable auto_explain extension
+    #[arg(long, value_parser = parse_duration)]
+    pub pg_auto_explain_with_min_duration: Option<Duration>,
+
     /// Postgres database url. If unspecified DATABASE_URL environment variable
     /// is used
     #[arg(long)]
@@ -78,6 +88,8 @@ async fn main() {
         pg_pool_connections: args.pg_pool_connections,
         pg_polling_interval: args.pg_polling_interval,
         worker_thread_count: args.worker_thread_count,
+        pg_timeout: args.pg_timeout,
+        pg_auto_explain_with_min_duration: args.pg_auto_explain_with_min_duration,
     };
 
     if let Err(err) = telemetry::setup_otlp(&args.service_name) {

coprocessor/fhevm-engine/zkproof-worker/src/bin/zkproof_worker.rs

@@ -4,7 +4,7 @@ pub mod auxiliary;
 mod tests;
 
 pub mod verifier;
-use std::io;
+use std::{io, time::Duration};
 
 use fhevm_engine_common::{pg_pool::ServiceError, types::FhevmError};
 use thiserror::Error;
@@ -73,6 +73,8 @@ pub struct Config {
     pub notify_database_channel: String,
     pub pg_pool_connections: u32,
     pub pg_polling_interval: u32,
+    pub pg_timeout: Duration,
+    pub pg_auto_explain_with_min_duration: Option<Duration>,
 
     pub worker_thread_count: u32,
 }

coprocessor/fhevm-engine/zkproof-worker/src/lib.rs

@@ -21,15 +21,17 @@ pub async fn setup() -> anyhow::Result<(PostgresPoolManager, DBInstance)> {
         pg_pool_connections: 10,
         pg_polling_interval: 60,
         worker_thread_count: 1,
+        pg_timeout: Duration::from_secs(15),
+        pg_auto_explain_with_min_duration: None,
     };
 
     let pool_mngr = PostgresPoolManager::connect_pool(
         test_instance.parent_token.child_token(),
         conf.database_url.as_str(),
-        Duration::from_secs(15),
-        10,
+        conf.pg_timeout,
+        conf.pg_pool_connections,
         Duration::from_secs(2),
-        None, // TODO: conf.pg_auto_explain_with_min_duration
+        conf.pg_auto_explain_with_min_duration,
     )
     .await
     .unwrap();

coprocessor/fhevm-engine/zkproof-worker/src/tests/utils.rs

@@ -80,18 +80,18 @@ impl HealthCheckService for ZkProofService {
 impl ZkProofService {
     pub async fn create(conf: Config, token: CancellationToken) -> Option<ZkProofService> {
         // Each worker needs at least 3 pg connections
-        let pool_connections =
+        let max_pool_connections =
             std::cmp::max(conf.pg_pool_connections, 3 * conf.worker_thread_count);
         let t = telemetry::tracer("init_service");
         let _s = t.child_span("pg_connect");
 
         let Some(pool_mngr) = PostgresPoolManager::connect_pool(
             token.child_token(),
             conf.database_url.as_str(),
-            Duration::from_secs(15),
-            pool_connections,
+            conf.pg_timeout,
+            max_pool_connections,
             Duration::from_secs(2),
-            None, // TODO: conf.pg_auto_explain_with_min_duration
+            conf.pg_auto_explain_with_min_duration,
         )
         .await
         else {