feat: make tests for async loop

Had to remove MyProcNumber symbol from the slot bus iterator to
make it work with tokio.

Author: darthunix

src/backend.rs

@@ -3,7 +3,7 @@ use libc::c_long;
 use pgrx::pg_sys::{
     error, fetch_search_path_array, get_namespace_oid, get_relname_relid, palloc0,
     CustomExecMethods, CustomScan, CustomScanMethods, CustomScanState, EState, ExplainState,
-    InvalidOid, List, ListCell, MyLatch, Node, NodeTag, Oid, ParamListInfo,
+    InvalidOid, List, ListCell, MyLatch, MyProcNumber, Node, NodeTag, Oid, ParamListInfo,
     RegisterCustomScanMethods, ResetLatch, TupleTableSlot, WaitLatch, PG_WAIT_EXTENSION,
     WL_LATCH_SET, WL_POSTMASTER_DEATH, WL_TIMEOUT,
 };
@@ -68,10 +68,11 @@ struct ScanState {
 #[pg_guard]
 #[no_mangle]
 unsafe extern "C" fn create_df_scan_state(cscan: *mut CustomScan) -> *mut Node {
+    let my_proc_number = unsafe { MyProcNumber };
     let wait_stream = || -> SlotStream {
         let stream;
         loop {
-            let Some(slot) = Bus::new().slot_locked(my_slot()) else {
+            let Some(slot) = Bus::new().slot_locked(my_slot(), my_proc_number) else {
                 wait_latch(Some(BACKEND_WAIT_TIMEOUT));
                 continue;
             };
@@ -95,7 +96,7 @@ unsafe extern "C" fn create_df_scan_state(cscan: *mut CustomScan) -> *mut Node {
             wait_latch(Some(BACKEND_WAIT_TIMEOUT));
             skip_wait = false;
         }
-        let Some(slot) = Bus::new().slot_locked(my_slot()) else {
+        let Some(slot) = Bus::new().slot_locked(my_slot(), my_proc_number) else {
             continue;
         };
         let mut stream = SlotStream::from(slot);

src/ipc.rs

@@ -11,7 +11,7 @@ use std::slice::from_raw_parts_mut;
 use std::sync::atomic::{AtomicBool, AtomicI32, Ordering};
 
 static mut SLOT_FREE_LIST_PTR: OnceCell<*mut c_void> = OnceCell::new();
-static mut BUS_PTR: OnceCell<*mut c_void> = OnceCell::new();
+pub(crate) static mut BUS_PTR: OnceCell<*mut c_void> = OnceCell::new();
 static mut WORKER_PID_PTR: OnceCell<*mut c_void> = OnceCell::new();
 pub(crate) const INVALID_PROC_NUMBER: i32 = -1;
 pub(crate) const DATA_SIZE: usize = 8 * 1024;
@@ -268,19 +268,25 @@ impl Slot {
         unsafe { (*self.locked).load(Ordering::Relaxed) }
     }
 
-    pub(crate) fn lock(&self) -> bool {
+    pub(crate) fn lock(&self, holder: i32) -> bool {
         unsafe {
-            if let Ok(_) =
-                (*self.locked).compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
+            if (*self.locked)
+                .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
+                .is_ok()
             {
-                (*self.holder).store(MyProcNumber, Ordering::Relaxed);
+                (*self.holder).store(holder, Ordering::Relaxed);
                 true
             } else {
                 false
             }
         }
     }
 
+    pub(crate) fn pg_lock(&self) -> bool {
+        let holder = unsafe { MyProcNumber };
+        self.lock(holder)
+    }
+
     pub(crate) fn unlock(&self) {
         unsafe {
             (*self.holder).store(INVALID_PROC_NUMBER, Ordering::Release);
@@ -413,9 +419,9 @@ impl Bus {
         Slot::from_bytes(slot_ptr, Slot::estimated_size())
     }
 
-    pub(crate) fn slot_locked(&mut self, id: SlotNumber) -> Option<Slot> {
+    pub(crate) fn slot_locked(&mut self, id: SlotNumber, holder: i32) -> Option<Slot> {
         let slot = self.slot(id);
-        if slot.lock() {
+        if slot.lock(holder) {
             Some(slot)
         } else {
             None
@@ -427,21 +433,23 @@ impl Bus {
         Self::from_bytes(ptr as *mut u8, Self::estimated_size())
     }
 
-    pub(crate) fn into_iter(self) -> BusIter {
-        BusIter::from(self)
+    pub(crate) fn into_iter(self, holder: i32) -> BusIter {
+        BusIter::new(self, holder)
     }
 }
 
 pub(crate) struct BusIter {
     pos: SlotNumber,
     inner: Bus,
+    holder: i32,
 }
 
-impl From<Bus> for BusIter {
-    fn from(value: Bus) -> Self {
+impl BusIter {
+    pub(crate) fn new(bus: Bus, holder: i32) -> Self {
         Self {
             pos: 0,
-            inner: value,
+            inner: bus,
+            holder,
         }
     }
 }
@@ -453,7 +461,7 @@ impl Iterator for BusIter {
         if self.pos >= max_backends() {
             return None;
         }
-        let res = self.inner.slot_locked(self.pos);
+        let res = self.inner.slot_locked(self.pos, self.holder);
         self.pos += 1;
         Some(res)
     }
@@ -562,8 +570,8 @@ pub(crate) mod tests {
         assert_eq!(Slot::estimated_size(), SLOT_SIZE);
         let mut slot_buf: [u8; SLOT_SIZE] = [1; SLOT_SIZE];
         let mut slot = make_slot(&mut slot_buf);
-        assert!(slot.lock());
-        assert!(!slot.lock());
+        assert!(slot.pg_lock());
+        assert!(!slot.pg_lock());
         unsafe {
             assert_eq!(&slot_buf[Slot::range1()], &[1]);
             write(slot.data_mut(), [1; DATA_SIZE]);
@@ -577,7 +585,7 @@ pub(crate) mod tests {
     fn test_slot_stream() {
         let mut buffer: [u8; SLOT_SIZE] = [1; SLOT_SIZE];
         let slot = make_slot(&mut buffer);
-        slot.lock();
+        slot.pg_lock();
         let mut stream = SlotStream::from(slot);
         let data = [42; 10];
         let len = stream.write(&data).unwrap();
@@ -626,12 +634,13 @@ pub(crate) mod tests {
                 assert_eq!(buffer[i * slot_size], 0);
             }
         }
+        let my_proc_number = unsafe { MyProcNumber };
         let mut bus = Bus::from_bytes(ptr, len);
         for i in 0..max_backends() {
-            let slot = bus.slot_locked(i);
+            let slot = bus.slot_locked(i, my_proc_number);
             assert!(slot.is_some());
             let slot = slot.unwrap();
-            assert!(!slot.lock());
+            assert!(!slot.pg_lock());
             slot.unlock();
             assert_eq!(buffer[i as usize * slot_size], 0);
         }

src/worker.rs

@@ -54,8 +54,7 @@ struct WorkerContext {
 }
 
 impl WorkerContext {
-    fn new() -> Self {
-        let capacity = max_backends() as usize;
+    fn with_capacity(capacity: usize) -> Self {
         let tasks = Vec::with_capacity(capacity);
         let mut states = Vec::with_capacity(capacity);
         let mut statements = Vec::with_capacity(capacity);
@@ -88,8 +87,8 @@ impl WorkerContext {
     }
 }
 
-fn init_slots() -> Result<()> {
-    for locked_slot in Bus::new().into_iter().flatten() {
+fn init_slots(holder: i32) -> Result<()> {
+    for locked_slot in Bus::new().into_iter(holder).flatten() {
         let mut stream = SlotStream::from(locked_slot);
         stream.reset();
         let header = Header {
@@ -115,32 +114,41 @@ fn response_error(id: SlotNumber, ctx: &mut WorkerContext, stream: SlotStream, m
 #[no_mangle]
 pub extern "C" fn worker_main(_arg: pg_sys::Datum) {
     BackgroundWorker::attach_signal_handlers(SignalWakeFlags::SIGHUP | SignalWakeFlags::SIGTERM);
-    let mut ctx = WorkerContext::new();
+    let capacity = max_backends() as usize;
+    let mut ctx = WorkerContext::with_capacity(capacity);
     let rt = Builder::new_multi_thread()
         .worker_threads(TOKIO_THREAD_NUMBER)
         .enable_all()
         .build()
         .unwrap();
     let mut do_retry = false;
-    let capacity = max_backends() as usize;
     let mut errors: Vec<Option<SmolStr>> = vec![None; capacity];
     let mut signals: Vec<bool> = vec![false; capacity];
-    init_slots().expect("Failed to initialize slots");
-    unsafe { set_worker_id(MyProcNumber) };
+    let worker_proc_number = unsafe { MyProcNumber };
+    init_slots(worker_proc_number).expect("Failed to initialize slots");
+    set_worker_id(worker_proc_number);
 
     log!("DataFusion worker is running");
     while do_retry || BackgroundWorker::wait_latch(Some(WORKER_WAIT_TIMEOUT)) {
         rt.block_on(async {
             do_retry = false;
-            create_tasks(&mut ctx, &mut errors).await;
-            wait_results(&mut ctx, &mut errors, &mut signals, &mut do_retry).await;
+            create_tasks(&mut ctx, &mut errors, worker_proc_number).await;
+            wait_results(
+                &mut ctx,
+                &mut errors,
+                &mut signals,
+                &mut do_retry,
+                worker_proc_number,
+            )
+            .await;
         });
         // Process errors returned by the tasks.
         for (slot_id, msg) in errors.iter_mut().enumerate() {
             if let Some(msg) = msg {
                 let stream;
                 loop {
-                    let Some(slot) = Bus::new().slot_locked(slot_id as u32) else {
+                    let Some(slot) = Bus::new().slot_locked(slot_id as u32, worker_proc_number)
+                    else {
                         BackgroundWorker::wait_latch(Some(SLOT_WAIT_TIMEOUT));
                         continue;
                     };
@@ -168,8 +176,8 @@ pub extern "C" fn worker_main(_arg: pg_sys::Datum) {
 // runtime, while postgres functions can work only in single thread.
 
 /// Process packets from the slots and create tasks for them.
-async fn create_tasks(ctx: &mut WorkerContext, errors: &mut [Option<SmolStr>]) {
-    for (id, locked_slot) in Bus::new().into_iter().enumerate() {
+async fn create_tasks(ctx: &mut WorkerContext, errors: &mut [Option<SmolStr>], holder: i32) {
+    for (id, locked_slot) in Bus::new().into_iter(holder).enumerate() {
         let Some(slot) = locked_slot else {
             continue;
         };
@@ -226,12 +234,13 @@ async fn wait_results(
     errors: &mut [Option<SmolStr>],
     signals: &mut [bool],
     do_retry: &mut bool,
+    holder: i32,
 ) {
     for (id, task) in &mut ctx.tasks {
         let result = task.await.expect("Failed to await task");
         match result {
             Ok(TaskResult::Parsing((stmt, tables))) => {
-                let mut stream = wait_stream(*id).await;
+                let mut stream = wait_stream(*id, holder).await;
                 if tables.is_empty() {
                     // We don't need any table metadata for this query.
                     // So, write a fake metadata packet to the slot and proceed it
@@ -260,7 +269,7 @@ async fn wait_results(
                 ctx.statements[*id as usize] = Some(stmt);
             }
             Ok(TaskResult::Compilation(plan)) => {
-                let mut stream = wait_stream(*id).await;
+                let mut stream = wait_stream(*id, holder).await;
                 match request_params(&mut stream) {
                     Ok(()) => signals[*id as usize] = true,
                     Err(err) => {
@@ -282,9 +291,9 @@ async fn wait_results(
 }
 
 #[inline(always)]
-async fn wait_stream(slot_id: u32) -> SlotStream {
+async fn wait_stream(slot_id: u32, holder: i32) -> SlotStream {
     loop {
-        let Some(slot) = Bus::new().slot_locked(slot_id) else {
+        let Some(slot) = Bus::new().slot_locked(slot_id, holder) else {
             tokio::time::sleep(SLOT_WAIT_TIMEOUT).await;
             continue;
         };
@@ -339,6 +348,7 @@ mod tests {
     use super::*;
     use crate::data_type::{write_scalar_value, EncodedType};
     use crate::ipc::tests::{make_slot, SLOT_SIZE};
+    use crate::ipc::{Slot, BUS_PTR};
     use crate::protocol::prepare_query;
 
     #[tokio::test]
@@ -503,4 +513,36 @@ mod tests {
             "Projection: Int64(1) [Int64(1):Int64]\n  EmptyRelation []",
         );
     }
+
+    #[tokio::test]
+    async fn test_loop() {
+        let holder = 42;
+        let capacity = 2;
+        let mut ctx = WorkerContext::with_capacity(capacity);
+        let mut errors: Vec<Option<SmolStr>> = vec![None; capacity];
+        let mut signals: Vec<bool> = vec![false; capacity];
+        let mut do_retry = false;
+        let bus_size = Slot::estimated_size() * capacity;
+        let mut buffer = vec![0; bus_size];
+        unsafe { BUS_PTR.set(buffer.as_mut_ptr() as _).unwrap() };
+        init_slots(holder).expect("Failed to initialize slots");
+        // Check processing of the parse message.
+        let sql = "SELECT * FROM foo";
+        {
+            let slot = Bus::new().slot_locked(0, holder).unwrap();
+            let mut stream = SlotStream::from(slot);
+            prepare_query(&mut stream, sql).unwrap();
+        }
+        create_tasks(&mut ctx, &mut errors, holder).await;
+        wait_results(&mut ctx, &mut errors, &mut signals, &mut do_retry, holder).await;
+        let error = errors[0].take();
+        assert!(error.is_none(), "Error: {:?}", error);
+        let stmt = ctx.statements[0].take().unwrap();
+        let expected_stmt = DFParser::parse_sql(sql)
+            .expect("Failed to parse SQL")
+            .into_iter()
+            .next()
+            .expect("Failed to get statement");
+        assert_eq!(stmt, expected_stmt);
+    }
 }