ARROW-15067: [C++] Add tracing spans to the scanner

cpp/build-support/lsan-suppressions.txt

@@ -19,3 +19,10 @@
 leak:*__new_exitfn*
 # Leak at shutdown in OpenSSL
 leak:CRYPTO_zalloc
+
+# OpenTelemetry. These seem like false positives and go away if the
+# CPU thread pool is manually shut down before exit.
+# Note that ASan has trouble backtracing these and may not be able to
+# without LSAN_OPTIONS=fast_unwind_on_malloc=0:malloc_context_size=100
+leak:opentelemetry::v1::context::ThreadLocalContextStorage::GetStack
+leak:opentelemetry::v1::context::ThreadLocalContextStorage::Stack::Resize

cpp/src/arrow/dataset/file_csv.cc

@@ -39,6 +39,7 @@
 #include "arrow/util/async_generator.h"
 #include "arrow/util/iterator.h"
 #include "arrow/util/logging.h"
+#include "arrow/util/tracing_internal.h"
 #include "arrow/util/utf8.h"
 
 namespace arrow {
@@ -167,9 +168,14 @@ static inline Result<csv::ReadOptions> GetReadOptions(
 static inline Future<std::shared_ptr<csv::StreamingReader>> OpenReaderAsync(
     const FileSource& source, const CsvFileFormat& format,
     const std::shared_ptr<ScanOptions>& scan_options, Executor* cpu_executor) {
+#ifdef ARROW_WITH_OPENTELEMETRY
+  auto tracer = arrow::internal::tracing::GetTracer();
+  auto span = tracer->StartSpan("arrow::dataset::CsvFileFormat::OpenReaderAsync");
+#endif
   ARROW_ASSIGN_OR_RAISE(auto reader_options, GetReadOptions(format, scan_options));
 
   ARROW_ASSIGN_OR_RAISE(auto input, source.OpenCompressed());
+  const auto& path = source.path();
   ARROW_ASSIGN_OR_RAISE(
       input, io::BufferedInputStream::Create(reader_options.block_size,
                                              default_memory_pool(), std::move(input)));
@@ -190,11 +196,20 @@ static inline Future<std::shared_ptr<csv::StreamingReader>> OpenReaderAsync(
       }));
   return reader_fut.Then(
       // Adds the filename to the error
-      [](const std::shared_ptr<csv::StreamingReader>& reader)
-          -> Result<std::shared_ptr<csv::StreamingReader>> { return reader; },
-      [source](const Status& err) -> Result<std::shared_ptr<csv::StreamingReader>> {
-        return err.WithMessage("Could not open CSV input source '", source.path(),
-                               "': ", err);
+      [=](const std::shared_ptr<csv::StreamingReader>& reader)
+          -> Result<std::shared_ptr<csv::StreamingReader>> {
+#ifdef ARROW_WITH_OPENTELEMETRY
+        span->SetStatus(opentelemetry::trace::StatusCode::kOk);
+        span->End();
+#endif
+        return reader;
+      },
+      [=](const Status& err) -> Result<std::shared_ptr<csv::StreamingReader>> {
+#ifdef ARROW_WITH_OPENTELEMETRY
+        arrow::internal::tracing::MarkSpan(err, span.get());
+        span->End();
+#endif
+        return err.WithMessage("Could not open CSV input source '", path, "': ", err);
       });
 }
 
@@ -250,7 +265,10 @@ Result<RecordBatchGenerator> CsvFileFormat::ScanBatchesAsync(
   auto source = file->source();
   auto reader_fut =
       OpenReaderAsync(source, *this, scan_options, ::arrow::internal::GetCpuThreadPool());
-  return GeneratorFromReader(std::move(reader_fut), scan_options->batch_size);
+  auto generator = GeneratorFromReader(std::move(reader_fut), scan_options->batch_size);
+  WRAP_ASYNC_GENERATOR_WITH_CHILD_SPAN(
+      generator, "arrow::dataset::CsvFileFormat::ScanBatchesAsync::Next");
+  return generator;
 }
 
 Future<util::optional<int64_t>> CsvFileFormat::CountRows(

cpp/src/arrow/dataset/file_ipc.cc

@@ -30,6 +30,7 @@
 #include "arrow/util/checked_cast.h"
 #include "arrow/util/iterator.h"
 #include "arrow/util/logging.h"
+#include "arrow/util/tracing_internal.h"
 
 namespace arrow {
 
@@ -62,16 +63,31 @@ static inline Result<std::shared_ptr<ipc::RecordBatchFileReader>> OpenReader(
 static inline Future<std::shared_ptr<ipc::RecordBatchFileReader>> OpenReaderAsync(
     const FileSource& source,
     const ipc::IpcReadOptions& options = default_read_options()) {
+#ifdef ARROW_WITH_OPENTELEMETRY
+  auto tracer = arrow::internal::tracing::GetTracer();
+  auto span = tracer->StartSpan("arrow::dataset::IpcFileFormat::OpenReaderAsync");
+#endif
   ARROW_ASSIGN_OR_RAISE(auto input, source.Open());
   auto path = source.path();
   return ipc::RecordBatchFileReader::OpenAsync(std::move(input), options)
-      .Then([](const std::shared_ptr<ipc::RecordBatchFileReader>& reader)
-                -> Result<std::shared_ptr<ipc::RecordBatchFileReader>> { return reader; },
-            [path](const Status& status)
-                -> Result<std::shared_ptr<ipc::RecordBatchFileReader>> {
-              return status.WithMessage("Could not open IPC input source '", path,
-                                        "': ", status.message());
-            });
+      .Then(
+          [=](const std::shared_ptr<ipc::RecordBatchFileReader>& reader)
+              -> Result<std::shared_ptr<ipc::RecordBatchFileReader>> {
+#ifdef ARROW_WITH_OPENTELEMETRY
+            span->SetStatus(opentelemetry::trace::StatusCode::kOk);
+            span->End();
+#endif
+            return reader;
+          },
+          [=](const Status& status)
+              -> Result<std::shared_ptr<ipc::RecordBatchFileReader>> {
+#ifdef ARROW_WITH_OPENTELEMETRY
+            arrow::internal::tracing::MarkSpan(status, span.get());
+            span->End();
+#endif
+            return status.WithMessage("Could not open IPC input source '", path,
+                                      "': ", status.message());
+          });
 }
 
 static inline Result<std::vector<int>> GetIncludedFields(
@@ -151,6 +167,8 @@ Result<RecordBatchGenerator> IpcFileFormat::ScanBatchesAsync(
       ARROW_ASSIGN_OR_RAISE(generator, reader->GetRecordBatchGenerator(
                                            /*coalesce=*/false, options->io_context));
     }
+    WRAP_ASYNC_GENERATOR_WITH_CHILD_SPAN(
+        generator, "arrow::dataset::IpcFileFormat::ScanBatchesAsync::Next");
     auto batch_generator = MakeReadaheadGenerator(std::move(generator), readahead_level);
     return MakeChunkedBatchGenerator(std::move(batch_generator), options->batch_size);
   };

cpp/src/arrow/dataset/file_parquet.cc

@@ -34,6 +34,7 @@
 #include "arrow/util/iterator.h"
 #include "arrow/util/logging.h"
 #include "arrow/util/range.h"
+#include "arrow/util/tracing_internal.h"
 #include "parquet/arrow/reader.h"
 #include "parquet/arrow/schema.h"
 #include "parquet/arrow/writer.h"
@@ -415,8 +416,11 @@ Result<RecordBatchGenerator> ParquetFileFormat::ScanBatchesAsync(
                             ::arrow::internal::GetCpuThreadPool(), row_group_readahead));
     return generator;
   };
-  return MakeFromFuture(GetReaderAsync(parquet_fragment->source(), options)
-                            .Then(std::move(make_generator)));
+  auto generator = MakeFromFuture(GetReaderAsync(parquet_fragment->source(), options)
+                                      .Then(std::move(make_generator)));
+  WRAP_ASYNC_GENERATOR_WITH_CHILD_SPAN(
+      generator, "arrow::dataset::ParquetFileFormat::ScanBatchesAsync::Next");
+  return generator;
 }
 
 Future<util::optional<int64_t>> ParquetFileFormat::CountRows(

cpp/src/arrow/dataset/scanner.cc

@@ -35,10 +35,12 @@
 #include "arrow/dataset/plan.h"
 #include "arrow/table.h"
 #include "arrow/util/async_generator.h"
+#include "arrow/util/config.h"
 #include "arrow/util/iterator.h"
 #include "arrow/util/logging.h"
 #include "arrow/util/task_group.h"
 #include "arrow/util/thread_pool.h"
+#include "arrow/util/tracing_internal.h"
 
 namespace arrow {
 
@@ -206,6 +208,18 @@ class AsyncScanner : public Scanner, public std::enable_shared_from_this<AsyncSc
 Result<EnumeratedRecordBatchGenerator> FragmentToBatches(
     const Enumerated<std::shared_ptr<Fragment>>& fragment,
     const std::shared_ptr<ScanOptions>& options) {
+#ifdef ARROW_WITH_OPENTELEMETRY
+  auto tracer = arrow::internal::tracing::GetTracer();
+  auto span = tracer->StartSpan(
+      "arrow::dataset::FragmentToBatches",
+      {
+          {"arrow.dataset.fragment", fragment.value->ToString()},
+          {"arrow.dataset.fragment.index", fragment.index},
+          {"arrow.dataset.fragment.last", fragment.last},
+          {"arrow.dataset.fragment.type_name", fragment.value->type_name()},
+      });
+  auto scope = tracer->WithActiveSpan(span);
+#endif
   ARROW_ASSIGN_OR_RAISE(auto batch_gen, fragment.value->ScanBatchesAsync(options));
   ArrayVector columns;
   for (const auto& field : options->dataset_schema->fields()) {
@@ -214,6 +228,7 @@ Result<EnumeratedRecordBatchGenerator> FragmentToBatches(
                           MakeArrayOfNull(field->type(), /*length=*/0, options->pool));
     columns.push_back(std::move(array));
   }
+  WRAP_ASYNC_GENERATOR(batch_gen);
   batch_gen = MakeDefaultIfEmptyGenerator(
       std::move(batch_gen),
       RecordBatch::Make(options->dataset_schema, /*num_rows=*/0, std::move(columns)));
@@ -230,10 +245,13 @@ Result<EnumeratedRecordBatchGenerator> FragmentToBatches(
 Result<AsyncGenerator<EnumeratedRecordBatchGenerator>> FragmentsToBatches(
     FragmentGenerator fragment_gen, const std::shared_ptr<ScanOptions>& options) {
   auto enumerated_fragment_gen = MakeEnumeratedGenerator(std::move(fragment_gen));
-  return MakeMappedGenerator(std::move(enumerated_fragment_gen),
-                             [=](const Enumerated<std::shared_ptr<Fragment>>& fragment) {
-                               return FragmentToBatches(fragment, options);
-                             });
+  auto batch_gen_gen =
+      MakeMappedGenerator(std::move(enumerated_fragment_gen),
+                          [=](const Enumerated<std::shared_ptr<Fragment>>& fragment) {
+                            return FragmentToBatches(fragment, options);
+                          });
+  PROPAGATE_SPAN_TO_GENERATOR(std::move(batch_gen_gen));
+  return batch_gen_gen;
 }
 
 class OneShotFragment : public Fragment {

cpp/src/arrow/util/future.cc

@@ -27,6 +27,7 @@
 #include "arrow/util/checked_cast.h"
 #include "arrow/util/logging.h"
 #include "arrow/util/thread_pool.h"
+#include "arrow/util/tracing_internal.h"
 
 namespace arrow {
 
@@ -242,6 +243,19 @@ class ConcreteFutureImpl : public FutureImpl {
   void AddCallback(Callback callback, CallbackOptions opts) {
     CheckOptions(opts);
     std::unique_lock<std::mutex> lock(mutex_);
+#ifdef ARROW_WITH_OPENTELEMETRY
+    struct SpanWrapper {
+      void operator()(const FutureImpl& impl) {
+        auto scope = ::arrow::internal::tracing::GetTracer()->WithActiveSpan(active_span);
+        std::move(func)(impl);
+      }
+      Callback func;
+      opentelemetry::nostd::shared_ptr<opentelemetry::trace::Span> active_span;
+    };
+    SpanWrapper wrapper{std::move(callback),
+                        ::arrow::internal::tracing::GetTracer()->GetCurrentSpan()};
+    callback = std::move(wrapper);
+#endif
     CallbackRecord callback_record{std::move(callback), opts};
     if (IsFutureFinished(state_)) {
       lock.unlock();

cpp/src/arrow/util/thread_pool.cc

@@ -30,6 +30,8 @@
 #include "arrow/util/logging.h"
 #include "arrow/util/mutex.h"
 
+#include "arrow/util/tracing_internal.h"
+
 namespace arrow {
 namespace internal {
 
@@ -58,6 +60,20 @@ SerialExecutor::~SerialExecutor() = default;
 
 Status SerialExecutor::SpawnReal(TaskHints hints, FnOnce<void()> task,
                                  StopToken stop_token, StopCallback&& stop_callback) {
+#ifdef ARROW_WITH_OPENTELEMETRY
+  // Wrap the task to propagate a parent tracing span to it
+  struct SpanWrapper {
+    void operator()() {
+      auto scope = ::arrow::internal::tracing::GetTracer()->WithActiveSpan(active_span);
+      std::move(func)();
+    }
+    FnOnce<void()> func;
+    opentelemetry::nostd::shared_ptr<opentelemetry::trace::Span> active_span;
+  };
+  SpanWrapper wrapper{std::move(task),
+                      ::arrow::internal::tracing::GetTracer()->GetCurrentSpan()};
+  task = std::move(wrapper);
+#endif
   // While the SerialExecutor runs tasks synchronously on its main thread,
   // SpawnReal may be called from external threads (e.g. when transferring back
   // from blocking I/O threads), so we need to keep the state alive *and* to
@@ -366,6 +382,19 @@ Status ThreadPool::SpawnReal(TaskHints hints, FnOnce<void()> task, StopToken sto
       // We can still spin up more workers so spin up a new worker
       LaunchWorkersUnlocked(/*threads=*/1);
     }
+#ifdef ARROW_WITH_OPENTELEMETRY
+    // Wrap the task to propagate a parent tracing span to it
+    struct {
+      void operator()() {
+        auto scope = ::arrow::internal::tracing::GetTracer()->WithActiveSpan(activeSpan);
+        std::move(func)();
+      }
+      FnOnce<void()> func;
+      opentelemetry::nostd::shared_ptr<opentelemetry::trace::Span> activeSpan;
+    } wrapper{std::forward<FnOnce<void()>>(task),
+              ::arrow::internal::tracing::GetTracer()->GetCurrentSpan()};
+    task = std::move(wrapper);
+#endif
     state_->pending_tasks_.push_back(
         {std::move(task), std::move(stop_token), std::move(stop_callback)});
   }