1 files changed, 57 insertions, 0 deletions
diff --git a/cli/metrics.rs b/cli/metrics.rs
index 43f11cd36..227c10f60 100644
--- a/cli/metrics.rs
+++ b/cli/metrics.rs
@@ -71,3 +71,60 @@ impl Metrics {
     self.op_completed(bytes_received);
   }
 }
+
+use deno_core::BufVec;
+use deno_core::Op;
+use deno_core::OpFn;
+use deno_core::OpState;
+use std::cell::RefCell;
+use std::rc::Rc;
+
+pub fn metrics_op(op_fn: Box<OpFn>) -> Box<OpFn> {
+  Box::new(move |op_state: Rc<RefCell<OpState>>, bufs: BufVec| -> Op {
+    // TODOs:
+    // * The 'bytes' metrics seem pretty useless, especially now that the
+    //   distinction between 'control' and 'data' buffers has become blurry.
+    // * Tracking completion of async ops currently makes us put the boxed
+    //   future into _another_ box. Keeping some counters may not be expensive
+    //   in itself, but adding a heap allocation for every metric seems bad.
+    let mut buf_len_iter = bufs.iter().map(|buf| buf.len());
+    let bytes_sent_control = buf_len_iter.next().unwrap_or(0);
+    let bytes_sent_data = buf_len_iter.sum();
+
+    let op = (op_fn)(op_state.clone(), bufs);
+
+    let cli_state = crate::ops::cli_state2(&op_state);
+    let cli_state_ = cli_state.clone();
+    let mut metrics = cli_state.metrics.borrow_mut();
+
+    use futures::future::FutureExt;
+
+    match op {
+      Op::Sync(buf) => {
+        metrics.op_sync(bytes_sent_control, bytes_sent_data, buf.len());
+        Op::Sync(buf)
+      }
+      Op::Async(fut) => {
+        metrics.op_dispatched_async(bytes_sent_control, bytes_sent_data);
+        let fut = fut
+          .inspect(move |buf| {
+            let mut metrics = cli_state_.metrics.borrow_mut();
+            metrics.op_completed_async(buf.len());
+          })
+          .boxed_local();
+        Op::Async(fut)
+      }
+      Op::AsyncUnref(fut) => {
+        metrics.op_dispatched_async_unref(bytes_sent_control, bytes_sent_data);
+        let fut = fut
+          .inspect(move |buf| {
+            let mut metrics = cli_state_.metrics.borrow_mut();
+            metrics.op_completed_async_unref(buf.len());
+          })
+          .boxed_local();
+        Op::AsyncUnref(fut)
+      }
+      other => other,
+    }
+  })
+}