fix: hanging queries

Author: darthunix

executor/src/ipc.rs

@@ -30,11 +30,17 @@ pub async fn signal_listener(state: Arc<SharedState<'_>>) {
     let mut signal = signal(SignalKind::user_defined1()).expect("failed to create signal");
     while signal.recv().await.is_some() {
         // Notify affected sockets that the signal has been received.
+        let mut woke = 0usize;
         for i in 0..state.size {
             if state.flags[i].load(Ordering::Acquire) {
                 state.wakers[i].wake();
+                woke += 1;
+                tracing::trace!(target = "executor::ipc", conn_id = i, "signal_listener: woke socket");
             }
         }
+        if woke == 0 {
+            tracing::trace!(target = "executor::ipc", "signal_listener: signal received but no flags set");
+        }
     }
 }
 

executor/src/pgscan.rs

@@ -51,6 +51,7 @@ pub struct HeapBlock {
 #[derive(Debug, Default)]
 pub struct ScanRegistry {
     inner: Mutex<HashMap<ScanId, Entry>>,
+    conn_id: usize,
 }
 
 #[derive(Debug)]
@@ -60,9 +61,9 @@ struct Entry {
 }
 
 impl ScanRegistry {
-    pub fn new() -> Self {
-        Self::default()
-    }
+    pub fn new() -> Self { Self { inner: Mutex::new(HashMap::new()), conn_id: 0 } }
+
+    pub fn with_conn(conn_id: usize) -> Self { Self { inner: Mutex::new(HashMap::new()), conn_id } }
 
     #[inline]
     fn lock(&self) -> MutexGuard<HashMap<ScanId, Entry>> {
@@ -111,6 +112,9 @@ impl ScanRegistry {
         let mut map = self.lock();
         let _ = map.remove(&scan_id);
     }
+
+    #[inline]
+    pub fn conn_id(&self) -> usize { self.conn_id }
 }
 
 // No global registry; registries are per-connection and owned by the server Storage.
@@ -249,7 +253,9 @@ impl ExecutionPlan for PgScanExec {
         _ctx: Arc<TaskContext>,
     ) -> DFResult<SendableRecordBatchStream> {
         let rx = self.registry.take_receiver(self.scan_id);
-        let stream = PgScanStream::new(Arc::clone(&self.schema), Arc::clone(&self.attrs), rx);
+        // Use connection id stored in registry to address per-connection slot buffers
+        let conn_id = self.registry.conn_id();
+        let stream = PgScanStream::new(Arc::clone(&self.schema), Arc::clone(&self.attrs), rx, conn_id);
         Ok(Box::pin(stream))
     }
 
@@ -262,15 +268,17 @@ pub struct PgScanStream {
     schema: SchemaRef,
     attrs: Arc<Vec<PgAttrMeta>>,
     rx: Option<mpsc::Receiver<HeapBlock>>,
+    conn_id: usize,
 }
 
 impl PgScanStream {
     pub fn new(
         schema: SchemaRef,
         attrs: Arc<Vec<PgAttrMeta>>,
         rx: Option<mpsc::Receiver<HeapBlock>>,
+        conn_id: usize,
     ) -> Self {
-        Self { schema, attrs, rx }
+        Self { schema, attrs, rx, conn_id }
     }
 }
 
@@ -285,9 +293,9 @@ impl Stream for PgScanStream {
         match rx.poll_recv(cx) {
             Poll::Ready(Some(block)) => {
                 // Borrow page and (optional) visibility bitmap from shared memory (no copy)
-                let page = unsafe { shm::block_slice(block.slot_id as usize) };
+                let page = unsafe { shm::block_slice(this.conn_id, block.slot_id as usize) };
                 let _vis = if block.vis_len > 0 {
-                    Some(unsafe { shm::vis_slice(block.slot_id as usize, block.vis_len as usize) })
+                    Some(unsafe { shm::vis_slice(this.conn_id, block.slot_id as usize, block.vis_len as usize) })
                 } else {
                     None
                 };
@@ -309,9 +317,23 @@ impl Stream for PgScanStream {
                     .map_err(|e| datafusion::error::DataFusionError::Execution(format!("{e}")))?;
                 // Use tuples_by_offset to iterate LP_NORMAL tuples in page order
                 let mut pairs: Vec<(u16, u16)> = Vec::new();
+                // Pre-scan to populate pairs and log LP_NORMAL count
+                {
+                    let _ = hp.tuples_by_offset(None, std::ptr::null_mut(), &mut pairs);
+                }
+                let pairs_len = pairs.len();
+                // Create iterator borrowing the filled pairs slice
                 let mut it = hp.tuples_by_offset(None, std::ptr::null_mut(), &mut pairs);
+                tracing::trace!(
+                    target = "executor::server",
+                    blkno = block.blkno,
+                    num_offsets = block.num_offsets,
+                    lp_normal = pairs_len,
+                    "pgscan: tuples_by_offset summary"
+                );
                 let page_hdr = unsafe { &*(page.as_ptr() as *const pg_sys::PageHeaderData) }
                     as *const pg_sys::PageHeaderData;
+                let mut decoded_rows = 0usize;
                 while let Some(tup) = it.next() {
                     // Decode projected columns for tuple using iterator over all columns
                     let iter =
@@ -326,6 +348,7 @@ impl Stream for PgScanStream {
                             None => append_null(&mut builders[col_idx]),
                         }
                     }
+                    decoded_rows += 1;
                 }
 
                 // Build Arrow arrays and RecordBatch
@@ -335,6 +358,12 @@ impl Stream for PgScanStream {
                 }
                 let rb = RecordBatch::try_new(Arc::clone(&this.schema), arrs)
                     .map_err(|e| datafusion::error::DataFusionError::Execution(format!("{e}")))?;
+                tracing::trace!(
+                    target = "executor::server",
+                    rows = decoded_rows,
+                    blkno = block.blkno,
+                    "pgscan: decoded rows"
+                );
                 Poll::Ready(Some(Ok(rb)))
             }
             Poll::Ready(None) => Poll::Ready(None),

executor/src/server.rs

@@ -57,6 +57,10 @@ impl Storage {
         *self = Self::default();
         self.registry = registry;
     }
+
+    pub fn set_registry_conn(&mut self, conn_id: usize) {
+        self.registry = Arc::new(crate::pgscan::ScanRegistry::with_conn(conn_id));
+    }
 }
 
 pub struct Connection<'bytes> {
@@ -87,6 +91,11 @@ impl<'bytes> Connection<'bytes> {
     }
 
     pub async fn poll(&mut self) -> Result<()> {
+        let pending = self.recv_socket.buffer.len();
+        if pending > 0 {
+            trace!("poll: data already available: {} bytes", pending);
+            return Ok(());
+        }
         trace!("poll: waiting for socket signal");
         (&mut self.recv_socket).await?;
         trace!(
@@ -135,6 +144,13 @@ impl<'bytes> Connection<'bytes> {
     }
 
     pub async fn process_message(&mut self, storage: &mut Storage) -> Result<()> {
+        // Guard against spurious wakeups where no bytes are available yet.
+        let pending = self.recv_socket.buffer.len();
+        if pending < 6 {
+            // Header is at least 6 bytes (3 fixints + u16)
+            trace!("process_message: woke but no header bytes available",);
+            return Ok(());
+        }
         let header = consume_header(&mut self.recv_socket.buffer)?;
         debug!(
             direction = ?header.direction,
@@ -152,6 +168,8 @@ impl<'bytes> Connection<'bytes> {
         // Column layout arrives during planning/prepare; store it for later encoding use.
         if header.tag == ControlPacket::ColumnLayout as u8 {
             let attrs = consume_column_layout(&mut self.recv_socket.buffer)?;
+            let len = attrs.len();
+            tracing::trace!(target = "executor::server", pg_attrs = len, "column layout received");
             storage.pg_attrs = Some(attrs);
             return Ok(());
         }
@@ -162,6 +180,16 @@ impl<'bytes> Connection<'bytes> {
                     // Read metadata; we expect the visibility bitmap to be stored in shared memory.
                     let meta =
                         read_heap_block_bitmap_meta(&mut self.recv_socket.buffer, header.length)?;
+                    trace!(
+                        target = "executor::server",
+                        scan_id = meta.scan_id,
+                        slot_id = meta.slot_id,
+                        table_oid = meta.table_oid,
+                        blkno = meta.blkno,
+                        num_offsets = meta.num_offsets,
+                        bitmap_inline_len = meta.bitmap_len,
+                        "heap bitmap meta received"
+                    );
                     // If payload included inline bitmap bytes (legacy), drain them without allocation.
                     let mut remaining = meta.bitmap_len as usize;
                     while remaining > 0 {
@@ -172,6 +200,13 @@ impl<'bytes> Connection<'bytes> {
                     }
                     // Trailing u16 carries the visibility bitmap length in shared memory
                     let vis_len = read_u16_msgpack(&mut self.recv_socket.buffer)?;
+                    trace!(
+                        target = "executor::server",
+                        scan_id = meta.scan_id,
+                        slot_id = meta.slot_id,
+                        vis_len,
+                        "heap bitmap vis_len (shm) received"
+                    );
                     if let Some(tx) = storage.registry.sender(meta.scan_id as u64) {
                         let block = crate::pgscan::HeapBlock {
                             scan_id: meta.scan_id as u64,
@@ -190,6 +225,13 @@ impl<'bytes> Connection<'bytes> {
                             );
                         }
                         // Request the next heap block for this scan using the same slot
+                        trace!(
+                            target = "executor::server",
+                            scan_id = meta.scan_id,
+                            table_oid = meta.table_oid,
+                            slot_id = meta.slot_id,
+                            "requesting next heap block"
+                        );
                         if let Err(e) = request_heap_block(
                             &mut self.send_buffer,
                             meta.scan_id,
@@ -411,6 +453,13 @@ fn encode_and_write_rows(
     let rows = batch.num_rows();
     let cols = std::cmp::min(batch.num_columns(), attrs.len());
     if cols == 0 || rows == 0 {
+        tracing::trace!(
+            target = "executor::server",
+            rows,
+            cols,
+            attrs = attrs.len(),
+            "execution: empty batch or no attrs; skipping write"
+        );
         return 0;
     }
     // Working buffers reused across rows to minimize allocations
@@ -554,9 +603,11 @@ fn open_data_flow(_conn: &mut Connection, storage: &mut Storage) -> Result<TaskR
     let phys = storage.physical_plan.as_ref().expect("checked above");
     // Reserve capacity to avoid HashMap reallocation while registering
     let scan_count = count_scans(phys) as usize;
+    trace!(scan_count, "open_data_flow: reserving scan registry capacity");
     storage.registry.reserve(scan_count);
     // Register channels per scan in the connection-local registry; no response payload
-    let _ = for_each_scan::<_, Error>(phys, |id, _table_oid| {
+    let _ = for_each_scan::<_, Error>(phys, |id, table_oid| {
+        trace!(scan_id = id, table_oid, "open_data_flow: registering scan channel");
         let _ = storage.registry.register(id, 16);
         Ok(())
     });
@@ -578,6 +629,7 @@ fn start_data_flow(conn: &mut Connection, storage: &mut Storage) -> Result<TaskR
 
     let plan = Arc::clone(storage.physical_plan.as_ref().expect("checked above"));
     let pg_attrs = storage.pg_attrs.clone().unwrap_or_default();
+    tracing::trace!(target = "executor::server", pg_attrs = pg_attrs.len(), "start_data_flow: attrs snapshot");
     // Build a fresh TaskContext for execution
     let state = SessionStateBuilder::new().build();
     let ctx = state.task_ctx();
@@ -586,6 +638,7 @@ fn start_data_flow(conn: &mut Connection, storage: &mut Storage) -> Result<TaskR
     let client_pid_val = conn
         .client_pid
         .load(std::sync::atomic::Ordering::Relaxed); // snapshot pid
+    trace!(conn_id, "start_data_flow: spawning execution task");
     let handle = tokio::spawn(async move {
         // Execute a single partition; PgScanExec uses a single stream
         match plan.execute(0, ctx) {
@@ -606,6 +659,12 @@ fn start_data_flow(conn: &mut Connection, storage: &mut Storage) -> Result<TaskR
                                 continue;
                             }
                             let wrote = encode_and_write_rows(&batch, &pg_attrs, &mut ring);
+                            tracing::trace!(
+                                target = "executor::server",
+                                rows = wrote,
+                                cols = batch.num_columns(),
+                                "execution: batch encoded and written"
+                            );
                             if wrote > 0 {
                                 let pid = client_pid_val;
                                 if pid > 0 && pid != i32::MAX {
@@ -618,6 +677,11 @@ fn start_data_flow(conn: &mut Connection, storage: &mut Storage) -> Result<TaskR
                 tracing::trace!(target = "pg_fusion::server", "execution stream completed");
                 // Write EOF sentinel row to result ring to unblock backend
                 let _ = protocol::result::write_eof(&mut ring);
+                // Nudge the backend in case it is waiting on latch
+                let pid = client_pid_val;
+                if pid > 0 && pid != i32::MAX {
+                    unsafe { libc::kill(pid as libc::pid_t, libc::SIGUSR1) };
+                }
             }
             Err(e) => {
                 tracing::error!(
@@ -629,11 +693,13 @@ fn start_data_flow(conn: &mut Connection, storage: &mut Storage) -> Result<TaskR
     });
     storage.exec_task = Some(handle);
     // Send a tiny ack to the backend indicating execution is ready
+    trace!("start_data_flow: sending ExecReady to backend");
     protocol::exec::prepare_exec_ready(&mut conn.send_buffer)?;
     // No demo row: result rows will be written by the execution task
     // Kick off initial heap block requests for each scan using slot 0
     if let Some(phys) = storage.physical_plan.as_ref() {
         let _ = for_each_scan::<_, Error>(phys, |id, table_oid| {
+            trace!(scan_id = id, table_oid, slot_id = 0, "start_data_flow: initial heap request");
             request_heap_block(&mut conn.send_buffer, id, table_oid, 0)?;
             Ok(())
         });
@@ -642,12 +708,9 @@ fn start_data_flow(conn: &mut Connection, storage: &mut Storage) -> Result<TaskR
 }
 
 fn end_data_flow(storage: &mut Storage) -> Result<TaskResult> {
-    // Cancel running execution task if present
-    if let Some(handle) = storage.exec_task.take() {
-        handle.abort();
-    }
-    // Close and clear all registered scan channels
-    storage.registry.close_and_clear();
+    // Use unified reset to ensure state machine returns to Initialized and
+    // any in-flight execution task is aborted, while preserving registry object.
+    storage.flush();
     Ok(TaskResult::Noop)
 }
 

executor/src/shm.rs

@@ -41,20 +41,24 @@ pub fn copy_vis(slot: usize, vis_len: usize) -> Vec<u8> {
     unsafe { std::slice::from_raw_parts(ptr, len) }.to_vec()
 }
 
-/// Borrow the heap page bytes from shared memory for the given `slot`.
+/// Borrow the heap page bytes from shared memory for the given connection and `slot`.
 /// Lifetime is 'static since the region lives for the process lifetime; caller
 /// must ensure the producer won't mutate the slot while reading.
-pub unsafe fn block_slice(slot: usize) -> &'static [u8] {
+pub unsafe fn block_slice(conn_offset: usize, slot: usize) -> &'static [u8] {
     let (base, layout) = get();
-    let ptr = slot_ptr_calc(base, layout, slot, 0);
+    let per = layout.layout.size();
+    let conn_base = base.add(conn_offset * per);
+    let ptr = slot_ptr_calc(conn_base, layout, slot, 0);
     std::slice::from_raw_parts(ptr, layout.block_len)
 }
 
 /// Borrow the visibility bitmap bytes from shared memory for the given `slot`.
 /// Lifetime is 'static; see notes in `block_slice`.
-pub unsafe fn vis_slice(slot: usize, vis_len: usize) -> &'static [u8] {
+pub unsafe fn vis_slice(conn_offset: usize, slot: usize, vis_len: usize) -> &'static [u8] {
     let (base, layout) = get();
-    let ptr = vis_ptr_calc(base, layout, slot, 0);
+    let per = layout.layout.size();
+    let conn_base = base.add(conn_offset * per);
+    let ptr = vis_ptr_calc(conn_base, layout, slot, 0);
     let len = vis_len.min(layout.vis_bytes_per_block);
     std::slice::from_raw_parts(ptr, len)
 }