GH-48741: [C++] Fix deadlock in CSV AsyncThreadedTableReader destructor

Author: pitrou

cpp/src/arrow/csv/column_builder.cc

@@ -36,18 +36,21 @@
 #include "arrow/status.h"
 #include "arrow/type.h"
 #include "arrow/type_fwd.h"
+#include "arrow/util/checked_cast.h"
 #include "arrow/util/logging_internal.h"
 #include "arrow/util/task_group.h"
 
 namespace arrow {
 
+using internal::checked_pointer_cast;
 using internal::TaskGroup;
 
 namespace csv {
 
 class BlockParser;
 
-class ConcreteColumnBuilder : public ColumnBuilder {
+class ConcreteColumnBuilder : public ColumnBuilder,
+                              public std::enable_shared_from_this<ConcreteColumnBuilder> {
  public:
   explicit ConcreteColumnBuilder(MemoryPool* pool, std::shared_ptr<TaskGroup> task_group,
                                  int32_t col_index = -1)
@@ -152,15 +155,17 @@ void NullColumnBuilder::Insert(int64_t block_index,
   const int32_t num_rows = parser->num_rows();
   DCHECK_GE(num_rows, 0);
 
-  task_group_->Append([this, block_index, num_rows]() -> Status {
-    std::unique_ptr<ArrayBuilder> builder;
-    RETURN_NOT_OK(MakeBuilder(pool_, type_, &builder));
-    std::shared_ptr<Array> res;
-    RETURN_NOT_OK(builder->AppendNulls(num_rows));
-    RETURN_NOT_OK(builder->Finish(&res));
-
-    return SetChunk(block_index, res);
-  });
+  // `self` keeps us alive, `this` allows easy access to derived / protected member vars
+  task_group_->Append(
+      [self = shared_from_this(), this, block_index, num_rows]() -> Status {
+        std::unique_ptr<ArrayBuilder> builder;
+        RETURN_NOT_OK(MakeBuilder(pool_, type_, &builder));
+        std::shared_ptr<Array> res;
+        RETURN_NOT_OK(builder->AppendNulls(num_rows));
+        RETURN_NOT_OK(builder->Finish(&res));
+
+        return SetChunk(block_index, res);
+      });
 }
 
 //////////////////////////////////////////////////////////////////////////
@@ -202,7 +207,7 @@ void TypedColumnBuilder::Insert(int64_t block_index,
   ReserveChunks(block_index);
 
   // We're careful that all references in the closure outlive the Append() call
-  task_group_->Append([this, parser, block_index]() -> Status {
+  task_group_->Append([self = shared_from_this(), this, parser, block_index]() -> Status {
     return SetChunk(block_index, converter_->Convert(*parser, col_index_));
   });
 }
@@ -257,7 +262,9 @@ Status InferringColumnBuilder::UpdateType() {
 }
 
 void InferringColumnBuilder::ScheduleConvertChunk(int64_t chunk_index) {
-  task_group_->Append([this, chunk_index]() { return TryConvertChunk(chunk_index); });
+  task_group_->Append([self = shared_from_this(), this, chunk_index]() {
+    return TryConvertChunk(chunk_index);
+  });
 }
 
 Status InferringColumnBuilder::TryConvertChunk(int64_t chunk_index) {

cpp/src/arrow/csv/reader.cc

@@ -1021,13 +1021,9 @@ class AsyncThreadedTableReader
                         convert_options, /*count_rows=*/false),
         cpu_executor_(cpu_executor) {}
 
-  ~AsyncThreadedTableReader() override {
-    if (task_group_) {
-      // In case of error, make sure all pending tasks are finished before
-      // we start destroying BaseTableReader members
-      ARROW_UNUSED(task_group_->Finish());
-    }
-  }
+  // XXX Ideally we can create a child StopToken for the tasks spawned here, and
+  // request cancellation in our destructor, to avoid running superfluous tasks
+  // in case of early error.
 
   Status Init() override {
     ARROW_ASSIGN_OR_RAISE(auto istream_it,
@@ -1048,6 +1044,8 @@ class AsyncThreadedTableReader
   Result<std::shared_ptr<Table>> Read() override { return ReadAsync().result(); }
 
   Future<std::shared_ptr<Table>> ReadAsync() override {
+    // Note that this task group doesn't survive our destruction, so it's essential
+    // that async callbacks own a strong ref to us.
     task_group_ = TaskGroup::MakeThreaded(cpu_executor_, io_context_.stop_token());
 
     auto self = shared_from_this();