use same timeout on all calls, refactor bgw wait latch interaction (#63)

* use same timeout on all calls

* only pause in wait latch

Author: ChuckHend

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "vectorize"
-version = "0.11.0"
+version = "0.11.1"
 edition = "2021"
 publish = false
 

Cargo.toml

@@ -6,7 +6,7 @@ description = "The simplest way to orchestrate vector search on Postgres."
 homepage = "https://github.com/tembo-io/pg_vectorize"
 documentation = "https://github.com/tembo-io/pg_vectorize"
 categories = ["orchestration", "machine_learning"]
-version = "0.11.0"
+version = "0.11.1"
 
 [build]
 postgres_version = "15"

Trunk.toml

@@ -94,11 +94,12 @@ pub async fn get_model_info(
 ) -> Result<TransformerMetadata> {
     let svc_url = get_generic_svc_url()?;
     let info_url = svc_url.replace("/embeddings", "/info");
+    let timeout = EMBEDDING_REQ_TIMEOUT_SEC.get();
     let client = reqwest::Client::new();
     let mut req = client
         .get(info_url)
         .query(&[("model_name", model_name)])
-        .timeout(std::time::Duration::from_secs(5)); // model info must always be fast
+        .timeout(std::time::Duration::from_secs(timeout as u64));
     if let Some(key) = api_key {
         req = req.header("Authorization", format!("Bearer {}", key));
     }

sql/vectorize--0.11.0--0.11.1.sql

@@ -4,7 +4,7 @@ use anyhow::Result;
 
 use crate::{
     executor::VectorizeMeta,
-    guc::OPENAI_KEY,
+    guc::{EMBEDDING_REQ_TIMEOUT_SEC, OPENAI_KEY},
     transformers::{
         http_handler::handle_response,
         types::{EmbeddingPayload, EmbeddingRequest, Inputs},
@@ -72,6 +72,7 @@ pub fn trim_inputs(inputs: &[Inputs]) -> Vec<String> {
 
 pub fn validate_api_key(key: &str) -> Result<()> {
     let client = reqwest::Client::new();
+    let timeout = EMBEDDING_REQ_TIMEOUT_SEC.get();
     let runtime = tokio::runtime::Builder::new_current_thread()
         .enable_io()
         .enable_time()
@@ -82,7 +83,7 @@ pub fn validate_api_key(key: &str) -> Result<()> {
             .get("https://api.openai.com/v1/models")
             .header("Content-Type", "application/json")
             .header("Authorization", format!("Bearer {}", key))
-            .timeout(std::time::Duration::from_secs(5)) // api validation should be fast
+            .timeout(std::time::Duration::from_secs(timeout as u64))
             .send()
             .await
             .unwrap_or_else(|e| error!("failed to make Open AI key validation call: {}", e));

src/transformers/http_handler.rs

@@ -1,10 +1,9 @@
 use crate::guc::{init_guc, NUM_BGW_PROC};
 use crate::init::VECTORIZE_QUEUE;
 use crate::util::{get_pg_conn, ready};
-use anyhow::Result;
 use pgrx::bgworkers::*;
 use pgrx::*;
-use std::time::{Duration, Instant};
+use std::time::Duration;
 
 use crate::workers::run_worker;
 
@@ -45,30 +44,30 @@ pub extern "C" fn background_worker_main(_arg: pg_sys::Datum) {
 
     log!("Starting BG Workers {}", BackgroundWorker::get_name(),);
 
-    while BackgroundWorker::wait_latch(Some(Duration::from_millis(10))) {
-        runtime.block_on(async {
-            while !ready(&conn).await {
-                log!("pg-vectorize: waiting for CREATE EXTENSION vectorize");
-                tokio::time::sleep(Duration::from_secs(5)).await;
-            }
-        });
+    let mut ext_ready: bool = false;
+    let mut wait_duration: Duration = Duration::from_secs(6);
+    while BackgroundWorker::wait_latch(Some(wait_duration)) {
+        if !ext_ready {
+            warning!("pg-vectorize: waiting for CREATE EXTENSION vectorize CASCADE;");
+            runtime.block_on(async {
+                ext_ready = ready(&conn).await;
+            });
+            // return to wait_latch if extension is not ready
+            continue;
+        }
 
-        let _worker_ran: Result<()> = runtime.block_on(async {
-            // continue to poll without pauses
-            let start = Instant::now();
-            let duration = Duration::from_secs(6);
-            // return control to wait_latch() after `duration` has elapsed
-            while start.elapsed() < duration {
-                match run_worker(queue.clone(), &conn, VECTORIZE_QUEUE).await {
-                    // sleep for 2 seconds when no messages in the queue
-                    Ok(None) => tokio::time::sleep(Duration::from_secs(2)).await,
-                    // sleep for 6 seconds when there is an error reading messages
-                    Err(_) => tokio::time::sleep(Duration::from_secs(6)).await,
-                    // continue to poll where there was a message consumed
-                    Ok(Some(_)) => continue,
-                }
-            }
-            Ok(())
+        wait_duration = runtime.block_on(async {
+            let wait_dur = match run_worker(queue.clone(), &conn, VECTORIZE_QUEUE).await {
+                // no messages in queue, so wait 2 seconds
+                Ok(None) => 2000,
+                // wait 10 seconds between polls when there is a failure
+                Err(_) => 10000,
+                // when there was a successfully processed message from queue,
+                // only wait 10ms before checking for more messages
+                // this allows postgres to kill or restart the bgw in between messages
+                Ok(Some(_)) => 10,
+            };
+            Duration::from_millis(wait_dur)
         });
     }
     log!("pg-vectorize: shutting down");