fix: worker loop error

Author: darthunix

src/backend.rs

@@ -104,6 +104,10 @@ unsafe extern "C" fn create_df_scan_state(cscan: *mut CustomScan) -> *mut Node {
             continue;
         }
         match header.packet {
+            Packet::None => {
+                // No data, just continue waiting.
+                continue;
+            }
             Packet::Failure => {
                 let msg = read_error(&mut stream).expect("Failed to read the error message");
                 error!("Failed to compile the query: {}", msg);

src/protocol.rs

@@ -51,10 +51,11 @@ impl TryFrom<u8> for Direction {
 #[derive(Clone, Debug, Default, PartialEq)]
 pub enum Packet {
     #[default]
-    Bind = 0,
-    Failure = 1,
-    Metadata = 2,
-    Parse = 3,
+    None = 0,
+    Bind = 1,
+    Failure = 2,
+    Metadata = 3,
+    Parse = 4,
 }
 
 impl TryFrom<u8> for Packet {
@@ -63,10 +64,11 @@ impl TryFrom<u8> for Packet {
     fn try_from(value: u8) -> Result<Self, Self::Error> {
         assert!(value < 128);
         match value {
-            0 => Ok(Packet::Bind),
-            1 => Ok(Packet::Failure),
-            2 => Ok(Packet::Metadata),
-            3 => Ok(Packet::Parse),
+            0 => Ok(Packet::None),
+            1 => Ok(Packet::Bind),
+            2 => Ok(Packet::Failure),
+            3 => Ok(Packet::Metadata),
+            4 => Ok(Packet::Parse),
             _ => Err(FusionError::Deserialize("packet".to_string(), value.into())),
         }
     }

src/worker.rs

@@ -6,7 +6,7 @@ use crate::ipc::{
 };
 use crate::protocol::{
     consume_header, prepare_metadata, read_params, read_query, request_params, send_error,
-    send_table_refs, Direction, Flag, Header, Packet,
+    send_table_refs, write_header, Direction, Flag, Header, Packet,
 };
 use crate::sql::Catalog;
 use anyhow::Result;
@@ -84,6 +84,21 @@ impl WorkerContext {
     }
 }
 
+fn init_slots() -> Result<()> {
+    for locked_slot in Bus::new().into_iter().flatten() {
+        let mut stream = SlotStream::from(locked_slot);
+        stream.reset();
+        let header = Header {
+            direction: Direction::ToBackend,
+            packet: Packet::None,
+            length: 0,
+            flag: Flag::Last,
+        };
+        write_header(&mut stream, &header)?;
+    }
+    Ok(())
+}
+
 #[pg_guard]
 #[no_mangle]
 pub extern "C" fn worker_main(_arg: pg_sys::Datum) {
@@ -98,10 +113,13 @@ pub extern "C" fn worker_main(_arg: pg_sys::Datum) {
     let mut do_retry = false;
     let mut slots_with_error: Vec<(SlotNumber, SmolStr)> =
         vec![(INVALID_SLOT_NUMBER, "".into()); max_backends() as usize];
+    init_slots().expect("Failed to initialize slots");
 
-    rt.block_on(async {
-        log!("DataFusion worker is running");
-        while do_retry || BackgroundWorker::wait_latch(Some(WORKER_WAIT_TIMEOUT)) {
+    log!("DataFusion worker is running");
+    while do_retry || BackgroundWorker::wait_latch(Some(WORKER_WAIT_TIMEOUT)) {
+        // Do not use any pgrx API in this loop: it is multithreaded while PostgreSQL
+        // functions can work only in the main thread.
+        rt.block_on(async {
             do_retry = false;
             for (id, locked_slot) in Bus::new().into_iter().enumerate() {
                 let Some(slot) = locked_slot else {
@@ -191,16 +209,15 @@ pub extern "C" fn worker_main(_arg: pg_sys::Datum) {
                     response_error(*slot_id, &mut ctx, stream, msg);
                 }
             }
-        }
-    });
+        });
+    }
 }
 
 async fn parse(header: Header, mut stream: SlotStream) -> Result<TaskResult> {
     assert_eq!(header.packet, Packet::Parse);
     // TODO: handle long queries that span multiple packets.
     assert_eq!(header.flag, Flag::Last);
     let (query, _) = read_query(&mut stream)?;
-    log!("Received query: {}", query);
 
     let stmts = DFParser::parse_sql(query)?;
     let Some(stmt) = stmts.into_iter().next() else {
@@ -234,26 +251,7 @@ async fn bind(
     Ok(TaskResult::Bind(plan))
 }
 
-#[inline]
-fn slot_warning() {
-    // The slot should be locked before sending the response.
-    // Normally this operation can not fail because its backend
-    // is waiting for response and should not hold any locks.
-    // But if it does it means that the old backend was terminated
-    // and some other one acquired the same slot. So, fsm should
-    // be reset to initial value.
-    warning!(
-        "{} {} {} {}",
-        "Failed to lock the slot for error response.",
-        "Looks like the old backend was terminated",
-        "and the slot is acquired by another backend.",
-        "The state machine will be reset to the initial state.",
-    );
-}
-
 fn response_error(id: SlotNumber, ctx: &mut WorkerContext, stream: SlotStream, message: &str) {
     ctx.flush(id);
-    if let Err(err) = send_error(id, stream, message) {
-        warning!("Failed to send the error message: {}", err);
-    }
+    send_error(id, stream, message).expect("Failed to send error response");
 }