Fix stream consumer failure hang in Lance streaming writer

When the Rust streaming task exits early (disk error, write failure),
push() would block indefinitely: the queue stays full, nothing notifies
not_full_cv_, and the C++ producer hangs forever.

Fix:
- Add lance_writer_stream_is_alive() Rust FFI: polls JoinHandle::is_finished()
  to detect whether the background task is still consuming the stream.
- Add StreamState::set_health_check(fn): stores a callback invoked when
  a push() timeout fires without a real queue-drain notification.
- Rewrite the stall wait from bare wait() to a wait_for(500ms) loop:
  on each timeout, call the health check; if the task is gone, call
  set_error() which unblocks both not_full_cv_ and not_empty_cv_.
- Wire up in initialize_lance_dataset(): register a lambda that calls
  lance_writer_stream_is_alive(rust_writer_) after stream start.

The 500 ms poll interval means producers unblock within half a second of
a Rust task crash, rather than hanging until process kill. Normal
operation (no task failure) adds zero overhead: the health check is only
reached on a wait_for timeout, which only fires when the queue stays full
for 500 ms without a pop — i.e., already a stall event.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: tsafin

include/tpch/lance_ffi.h

@@ -46,6 +46,17 @@ LanceWriter* lance_writer_create(const char* uri, const void* arrow_schema_ptr,
  */
 int lance_writer_start_stream(LanceWriter* writer, void* arrow_stream_ptr);
 
+/**
+ * Check whether the background streaming task is still running.
+ *
+ * Used by C++ producers to detect early task exit (e.g., write error) so they
+ * can unblock from a full queue rather than hanging indefinitely.
+ *
+ * @param writer Pointer to LanceWriter from lance_writer_create()
+ * @return 1 if the background task is running, 0 if finished or not started
+ */
+int lance_writer_stream_is_alive(const LanceWriter* writer);
+
 /**
  * Configure write parameters for Lance writes.
  *

src/writers/lance_writer.cpp

@@ -7,6 +7,8 @@
 #include <mutex>
 #include <condition_variable>
 #include <atomic>
+#include <chrono>
+#include <functional>
 #include <cstdlib>
 #include <cstring>
 #include <arrow/type.h>
@@ -23,13 +25,29 @@ struct StreamState {
     explicit StreamState(size_t max_queue_batches)
         : max_queue_batches_(max_queue_batches) {}
 
+    void set_health_check(std::function<bool()> fn) {
+        health_check_fn_ = std::move(fn);
+    }
+
     void push(const std::shared_ptr<arrow::RecordBatch>& batch, int64_t batch_bytes) {
         std::unique_lock<std::mutex> lock(mu_);
         if (!closed_ && status_.ok() && queue_.size() >= max_queue_batches_) {
             auto wait_start = std::chrono::steady_clock::now();
-            not_full_cv_.wait(lock, [&] {
-                return closed_ || queue_.size() < max_queue_batches_ || !status_.ok();
-            });
+            // Poll with a timeout so we can detect early stream-consumer exit.
+            // A bare wait() would hang indefinitely if the Rust task fails before
+            // draining the queue, because nothing would notify not_full_cv_.
+            while (!closed_ && status_.ok() && queue_.size() >= max_queue_batches_) {
+                bool signaled = not_full_cv_.wait_for(lock, std::chrono::milliseconds(500), [&] {
+                    return closed_ || queue_.size() < max_queue_batches_ || !status_.ok();
+                });
+                if (!signaled && health_check_fn_ && !health_check_fn_()) {
+                    status_ = arrow::Status::IOError(
+                        "Lance streaming task exited unexpectedly (write error or crash)");
+                    not_empty_cv_.notify_all();
+                    not_full_cv_.notify_all();
+                    break;
+                }
+            }
             auto wait_end = std::chrono::steady_clock::now();
             auto waited = std::chrono::duration_cast<std::chrono::nanoseconds>(wait_end - wait_start).count();
             stall_ns_.fetch_add(static_cast<uint64_t>(waited), std::memory_order_relaxed);
@@ -119,6 +137,7 @@ struct StreamState {
     size_t max_queue_batches_;
     bool closed_ = false;
     arrow::Status status_ = arrow::Status::OK();
+    std::function<bool()> health_check_fn_;  // returns false when Rust task has exited
     std::atomic<uint64_t> stall_ns_{0};
     std::atomic<uint64_t> stall_count_{0};
     int64_t current_bytes_ = 0;
@@ -268,6 +287,13 @@ void LanceWriter::initialize_lance_dataset(
         streaming_started_ = true;
         stream_state_ = std::move(state);
         stream_reader_ = std::move(reader);
+
+        // Register health check: if the Rust task exits early (write error etc.),
+        // push() will detect it within 500 ms and unblock producers via set_error().
+        auto* raw = reinterpret_cast<::LanceWriter*>(rust_writer_);
+        stream_state_->set_health_check([raw]() -> bool {
+            return lance_writer_stream_is_alive(raw) != 0;
+        });
     }
 
     std::cout << "Lance: Initialized dataset at " << dataset_path_ << "\n";

third_party/lance-ffi/src/lib.rs

@@ -371,6 +371,21 @@ pub extern "C" fn lance_writer_write_batch(writer_ptr: *mut LanceWriterHandle, a
     })).unwrap_or(7)
 }
 
+/// Returns 1 if the background streaming task is still running, 0 if it has finished
+/// (either successfully or with an error) or was never started. C++ uses this to detect
+/// early task exit and unblock producers waiting on a full queue.
+#[no_mangle]
+pub extern "C" fn lance_writer_stream_is_alive(writer_ptr: *const LanceWriterHandle) -> c_int {
+    if writer_ptr.is_null() { return 0; }
+    let writer = unsafe { &*writer_ptr };
+    match &writer.backend {
+        WriterBackend::Streaming { task } => {
+            task.as_ref().map(|t| if t.is_finished() { 0 } else { 1 }).unwrap_or(0)
+        }
+        _ => 0,
+    }
+}
+
 #[no_mangle]
 pub extern "C" fn lance_writer_start_stream(writer_ptr: *mut LanceWriterHandle, arrow_stream_ptr: *const c_void) -> c_int {
     catch_unwind(AssertUnwindSafe(|| {