Use condvar instead of spinloop at init

Author: syvb

pgvectorscale/src/access_method/build.rs

@@ -3,7 +3,9 @@ use std::time::Instant;
 
 use pg_sys::{FunctionCall0Coll, InvalidOid};
 use pgrx::ffi::c_char;
-use pgrx::pg_sys::{index_getprocinfo, pgstat_progress_update_param, AsPgCStr, Oid};
+use pgrx::pg_sys::{
+    index_getprocinfo, pgstat_progress_update_param, AsPgCStr, ConditionVariable, Oid,
+};
 use pgrx::*;
 
 use crate::access_method::distance::DistanceType;
@@ -112,10 +114,22 @@ impl<'a> BuildStateParallel<'a> {
         // Only update shared counter for the initializing worker until threshold is reached
         if self.is_initializing_worker && self.local_ntuples <= parallel::INITIAL_START_NODES_COUNT
         {
-            self.shared_state
+            let new_count = self
+                .shared_state
                 .build_state
                 .ntuples
-                .fetch_add(1, Ordering::Relaxed);
+                .fetch_add(1, Ordering::Relaxed)
+                + 1;
+
+            // Signal waiting workers when threshold is reached
+            if new_count >= parallel::INITIAL_START_NODES_COUNT {
+                unsafe {
+                    let cv_ptr = self.shared_state as *const ParallelShared as *mut ParallelShared;
+                    pg_sys::ConditionVariableBroadcast(
+                        &raw mut (*cv_ptr).build_state.initialization_cv,
+                    );
+                }
+            }
         }
     }
 
@@ -126,6 +140,14 @@ impl<'a> BuildStateParallel<'a> {
                 .build_state
                 .initializing_worker_done
                 .store(true, Ordering::Relaxed);
+
+            // Signal waiting workers that initialization is done
+            unsafe {
+                let cv_ptr = self.shared_state as *const ParallelShared as *mut ParallelShared;
+                pg_sys::ConditionVariableBroadcast(
+                    &raw mut (*cv_ptr).build_state.initialization_cv,
+                );
+            }
         }
 
         self.update_shared_ntuples();
@@ -191,6 +213,7 @@ struct ParallelBuildState {
     ntuples: AtomicUsize,
     start_nodes_initialized: AtomicBool,
     initializing_worker_done: AtomicBool,
+    initialization_cv: ConditionVariable,
 }
 
 /// Status data for parallel index builds, shared among all parallel workers.
@@ -330,7 +353,7 @@ pub extern "C" fn ambuild(
                 let parallel_shared =
                     pg_sys::shm_toc_allocate((*pcxt).toc, size_of::<ParallelShared>())
                         .cast::<ParallelShared>();
-                parallel_shared.write(ParallelShared {
+                let shared_state = ParallelShared {
                     params: ParallelSharedParams {
                         heaprelid: heap_relation.rd_id,
                         indexrelid: index_relation.rd_id,
@@ -341,8 +364,15 @@ pub extern "C" fn ambuild(
                         ntuples: AtomicUsize::new(0),
                         start_nodes_initialized: AtomicBool::new(false),
                         initializing_worker_done: AtomicBool::new(false),
+                        initialization_cv: std::mem::zeroed(), // Will be initialized below
                     },
-                });
+                };
+                parallel_shared.write(shared_state);
+
+                // Initialize the condition variable
+                pg_sys::ConditionVariableInit(
+                    &raw mut (*parallel_shared).build_state.initialization_cv,
+                );
                 let tablescandesc =
                     pg_sys::shm_toc_allocate((*pcxt).toc, tablescandesc_size_estimate)
                         .cast::<pg_sys::ParallelTableScanDescData>();
@@ -621,23 +651,27 @@ pub extern "C-unwind" fn _vectorscale_build_main(
     };
 
     if !should_initialize {
-        loop {
-            let ntuples = unsafe {
-                (*parallel_shared)
+        unsafe {
+            loop {
+                let ntuples = (*parallel_shared)
                     .build_state
                     .ntuples
-                    .load(Ordering::Relaxed)
-            };
-            let init_done = unsafe {
-                (*parallel_shared)
+                    .load(Ordering::Relaxed);
+                let init_done = (*parallel_shared)
                     .build_state
                     .initializing_worker_done
-                    .load(Ordering::Relaxed)
-            };
-            if ntuples >= parallel::INITIAL_START_NODES_COUNT || init_done {
-                break;
+                    .load(Ordering::Relaxed);
+
+                if ntuples >= parallel::INITIAL_START_NODES_COUNT || init_done {
+                    break;
+                }
+
+                // Wait on condition variable
+                pg_sys::ConditionVariableSleep(
+                    &raw mut (*parallel_shared).build_state.initialization_cv,
+                    pg_sys::PG_WAIT_EXTENSION,
+                );
             }
-            std::thread::sleep(std::time::Duration::from_millis(10));
         }
     }
     debug1!("Worker should initialize: {}", should_initialize);