fix: Fix deadlock in new streaming CSV / NDJSON sinks (#21598)

pola-rs · Mar 5, 2025 · 7943dc1 · 7943dc1
1 parent 241d34b
commit 7943dc1
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/crates/polars-stream/src/nodes/io_sinks/csv.rs b/crates/polars-stream/src/nodes/io_sinks/csv.rs
@@ -122,11 +122,14 @@ impl SinkNode for CsvSinkNode {
                     writer.write_batch(&df)?;
 
                     allocation_size = allocation_size.max(buffer.len());
+
+                    // Must drop before linearizer insert or will deadlock.
+                    drop(consume_token); // Keep the consume_token until here to increase the
+                                         // backpressure.
+
                     if lin_tx.insert(Priority(Reverse(seq), buffer)).await.is_err() {
                         return Ok(());
                     }
-                    drop(consume_token); // Keep the consume_token until here to increase the
-                                         // backpressure.
                 }
 
                 PolarsResult::Ok(())

diff --git a/crates/polars-stream/src/nodes/io_sinks/json.rs b/crates/polars-stream/src/nodes/io_sinks/json.rs
@@ -89,11 +89,14 @@ impl SinkNode for NDJsonSinkNode {
                     writer.write_batch(&df)?;
 
                     allocation_size = allocation_size.max(buffer.len());
+
+                    // Must drop before linearizer insert or will deadlock.
+                    drop(consume_token); // Keep the consume_token until here to increase the
+                                         // backpressure.
+
                     if lin_tx.insert(Priority(Reverse(seq), buffer)).await.is_err() {
                         return Ok(());
                     }
-                    drop(consume_token); // Keep the consume_token until here to increase the
-                                         // backpressure.
                 }
 
                 PolarsResult::Ok(())