From 0d4e8a82316fdce8e08e20a93ff65042f0afb43d Mon Sep 17 00:00:00 2001
From: Andrew Duffy <a10y@users.noreply.github.com>
Date: Thu, 21 Nov 2024 10:56:55 -0500
Subject: [PATCH] feat: add optional instrumentation to readers (#1431)

Add optional instrumentation using the `tracing` crate for `vortex-io`
and `vortex-file`. Disabled by default.

Part of some work to make it easier to collect deep profiles of
performance on object storage. Example output when hooked up to a
subscriber:


<img width="1265" alt="image"
src="https://github.com/user-attachments/assets/352564c6-9c99-408b-bb63-4e6c5bed8efd">
---
 Cargo.lock                     | 2 ++
 Cargo.toml                     | 1 +
 vortex-file/Cargo.toml         | 2 ++
 vortex-file/src/read/stream.rs | 1 +
 vortex-io/Cargo.toml           | 3 +--
 vortex-io/src/object_store.rs  | 2 ++
 vortex-io/src/tokio.rs         | 2 ++
 7 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 14d60239ed..3ab98b1001 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4835,6 +4835,7 @@ dependencies = [
  "rstest",
  "tempfile",
  "tokio",
+ "tracing",
  "vortex-array",
  "vortex-buffer",
  "vortex-dtype",
@@ -4895,6 +4896,7 @@ dependencies = [
  "object_store",
  "tempfile",
  "tokio",
+ "tracing",
  "vortex-buffer",
  "vortex-error",
 ]
diff --git a/Cargo.toml b/Cargo.toml
index 374959b845..4c15f6968d 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -135,6 +135,7 @@ tar = "0.4"
 tempfile = "3"
 thiserror = "2.0.0"
 tokio = "1.37.0"
+tracing = "0.1"
 uninit = "0.6.2"
 url = "2"
 uuid = "1.8.0"
diff --git a/vortex-file/Cargo.toml b/vortex-file/Cargo.toml
index ed9c9ba964..c6096cea9c 100644
--- a/vortex-file/Cargo.toml
+++ b/vortex-file/Cargo.toml
@@ -29,6 +29,7 @@ tokio = { workspace = true, features = [
     "fs",
     "rt-multi-thread",
 ] }
+tracing = { workspace = true, optional = true }
 vortex-array = { workspace = true }
 vortex-buffer = { workspace = true }
 vortex-dtype = { workspace = true, features = ["flatbuffers"] }
@@ -57,3 +58,4 @@ object_store = [
     "vortex-error/object_store",
     "vortex-io/object_store",
 ]
+tracing = ["dep:tracing", "vortex-io/tracing"]
diff --git a/vortex-file/src/read/stream.rs b/vortex-file/src/read/stream.rs
index ea38cdd734..309e7b33b4 100644
--- a/vortex-file/src/read/stream.rs
+++ b/vortex-file/src/read/stream.rs
@@ -298,6 +298,7 @@ impl<R: VortexReadAt + Unpin> VortexFileArrayStream<R> {
     }
 }
 
+#[cfg_attr(feature = "tracing", tracing::instrument(skip_all))]
 async fn read_ranges<R: VortexReadAt>(
     reader: R,
     ranges: Vec<MessageLocator>,
diff --git a/vortex-io/Cargo.toml b/vortex-io/Cargo.toml
index ebc7896920..272332344a 100644
--- a/vortex-io/Cargo.toml
+++ b/vortex-io/Cargo.toml
@@ -17,6 +17,7 @@ categories.workspace = true
 bytes = { workspace = true }
 compio = { workspace = true, features = ["bytes", "macros"], optional = true }
 tokio = { workspace = true, features = ["fs"], optional = true }
+tracing = { workspace = true, optional = true }
 futures = { workspace = true, features = ["std"] }
 futures-util = { workspace = true }
 flume = { workspace = true }
@@ -31,8 +32,6 @@ tokio = { workspace = true, features = ["full"] }
 
 [features]
 futures = ["futures-util/io"]
-compio = ["dep:compio"]
-tokio = ["dep:tokio"]
 object_store = ["dep:object_store", "vortex-error/object_store"]
 
 [lints]
diff --git a/vortex-io/src/object_store.rs b/vortex-io/src/object_store.rs
index 385faf4aba..be8cda99bb 100644
--- a/vortex-io/src/object_store.rs
+++ b/vortex-io/src/object_store.rs
@@ -65,6 +65,7 @@ impl ObjectStoreReadAt {
 }
 
 impl VortexReadAt for ObjectStoreReadAt {
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self)))]
     fn read_byte_range(
         &self,
         pos: u64,
@@ -82,6 +83,7 @@ impl VortexReadAt for ObjectStoreReadAt {
         })
     }
 
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self)))]
     fn size(&self) -> impl Future<Output = u64> + 'static {
         let object_store = self.object_store.clone();
         let location = self.location.clone();
diff --git a/vortex-io/src/tokio.rs b/vortex-io/src/tokio.rs
index 96ae65a646..7682d6499a 100644
--- a/vortex-io/src/tokio.rs
+++ b/vortex-io/src/tokio.rs
@@ -62,6 +62,7 @@ impl Deref for TokioFile {
 }
 
 impl VortexReadAt for TokioFile {
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self)))]
     fn read_byte_range(
         &self,
         pos: u64,
@@ -79,6 +80,7 @@ impl VortexReadAt for TokioFile {
         }
     }
 
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self)))]
     fn size(&self) -> impl Future<Output = u64> + 'static {
         let this = self.clone();