From b29ab647d1c22a0c8e4efd05d6b69898d155569d Mon Sep 17 00:00:00 2001
From: Nick Pillitteri <nick@56quarters.xyz>
Date: Sun, 23 Jun 2024 08:13:24 -0400
Subject: [PATCH] Add support for profiling to `dns` binary

Adding profiling support to `dns` binary to justify future performance
improvements.
---
 mtop/src/bin/dns.rs        | 15 ++++++++++++++-
 mtop/src/bin/mc.rs         | 38 +++-----------------------------------
 mtop/src/profile/mod.rs    | 11 +++++++----
 mtop/src/profile/writer.rs | 35 +++++++++++++++++++++++++++++++++++
 4 files changed, 59 insertions(+), 40 deletions(-)
 create mode 100644 mtop/src/profile/writer.rs

diff --git a/mtop/src/bin/dns.rs b/mtop/src/bin/dns.rs
index b4db349..5b92080 100644
--- a/mtop/src/bin/dns.rs
+++ b/mtop/src/bin/dns.rs
@@ -1,4 +1,5 @@
 use clap::{Args, Parser, Subcommand, ValueHint};
+use mtop::profile;
 use mtop_client::dns::{Flags, Message, MessageId, Name, Question, Record, RecordClass, RecordType};
 use std::fmt::Write;
 use std::io::Cursor;
@@ -21,6 +22,11 @@ struct DnsConfig {
     #[arg(long, default_value_t = DEFAULT_LOG_LEVEL)]
     log_level: Level,
 
+    /// Output pprof protobuf profile data to this file if profiling support was enabled
+    /// at build time.
+    #[arg(long, value_hint = ValueHint::FilePath)]
+    profile_output: Option<PathBuf>,
+
     #[command(subcommand)]
     mode: Action,
 }
@@ -91,11 +97,18 @@ async fn main() -> ExitCode {
         mtop::tracing::console_subscriber(opts.log_level).expect("failed to setup console logging");
     tracing::subscriber::set_global_default(console_subscriber).expect("failed to initialize console logging");
 
-    match &opts.mode {
+    let profiling = profile::Writer::default();
+    let code = match &opts.mode {
         Action::Query(cmd) => run_query(cmd).await,
         Action::Read(cmd) => run_read(cmd).await,
         Action::Write(cmd) => run_write(cmd).await,
+    };
+
+    if let Some(p) = opts.profile_output {
+        profiling.finish(p);
     }
+
+    code
 }
 
 async fn run_query(cmd: &QueryCommand) -> ExitCode {
diff --git a/mtop/src/bin/mc.rs b/mtop/src/bin/mc.rs
index 9e60944..af0f5b5 100644
--- a/mtop/src/bin/mc.rs
+++ b/mtop/src/bin/mc.rs
@@ -1,14 +1,12 @@
 use clap::{Args, Parser, Subcommand, ValueHint};
 use mtop::bench::{Bencher, Percent, Summary};
 use mtop::check::{Checker, TimingBundle};
-use mtop::profile::Profiler;
+use mtop::profile;
 use mtop_client::{
     DiscoveryDefault, MemcachedClient, MemcachedPool, Meta, MtopError, PoolConfig, SelectorRendezvous, Server,
     TLSConfig, Timeout, Value,
 };
-use std::fs::File;
-use std::io::Write;
-use std::path::{Path, PathBuf};
+use std::path::PathBuf;
 use std::process::ExitCode;
 use std::time::Duration;
 use std::{env, io};
@@ -322,7 +320,7 @@ async fn main() -> ExitCode {
         return ExitCode::FAILURE;
     };
 
-    let profiling = Profiling::default();
+    let profiling = profile::Writer::default();
     let code = match &opts.mode {
         Action::Add(cmd) => run_add(&opts, cmd, &client).await,
         Action::Bench(cmd) => run_bench(&opts, cmd, client).await,
@@ -344,36 +342,6 @@ async fn main() -> ExitCode {
     code
 }
 
-#[derive(Debug, Default)]
-pub struct Profiling {
-    profiler: Profiler,
-}
-
-impl Profiling {
-    pub fn finish<P>(&self, path: P)
-    where
-        P: AsRef<Path>,
-    {
-        if let Err(e) = self.profiler.proto().and_then(|b| Self::write_file(&path, &b)) {
-            tracing::warn!(message = "unable to collect and write pprof data", path = ?path.as_ref(), err = %e);
-        }
-    }
-
-    fn write_file<P>(path: P, contents: &[u8]) -> Result<(), MtopError>
-    where
-        P: AsRef<Path>,
-    {
-        if contents.is_empty() {
-            return Err(MtopError::configuration("mtop built without profiling support"));
-        }
-
-        let mut file = File::options().create(true).truncate(true).write(true).open(path)?;
-        file.write_all(contents)?;
-        file.flush()?;
-        Ok(file.sync_all()?)
-    }
-}
-
 async fn new_client(opts: &McConfig, servers: &[Server]) -> Result<MemcachedClient, MtopError> {
     let tls = TLSConfig {
         enabled: opts.tls_enabled,
diff --git a/mtop/src/profile/mod.rs b/mtop/src/profile/mod.rs
index bb8fdb0..757bdce 100644
--- a/mtop/src/profile/mod.rs
+++ b/mtop/src/profile/mod.rs
@@ -1,9 +1,12 @@
 #[cfg(not(feature = "profile"))]
+mod nop;
+#[cfg(not(feature = "profile"))]
 pub use nop::Profiler;
-#[cfg(feature = "profile")]
-pub use proto::Profiler;
 
-#[cfg(not(feature = "profile"))]
-mod nop;
 #[cfg(feature = "profile")]
 mod proto;
+#[cfg(feature = "profile")]
+pub use proto::Profiler;
+
+mod writer;
+pub use writer::Writer;
diff --git a/mtop/src/profile/writer.rs b/mtop/src/profile/writer.rs
new file mode 100644
index 0000000..23cec15
--- /dev/null
+++ b/mtop/src/profile/writer.rs
@@ -0,0 +1,35 @@
+use crate::profile::Profiler;
+use mtop_client::MtopError;
+use std::fs::File;
+use std::io::Write;
+use std::path::Path;
+
+#[derive(Debug, Default)]
+pub struct Writer {
+    profiler: Profiler,
+}
+
+impl Writer {
+    pub fn finish<P>(&self, path: P)
+    where
+        P: AsRef<Path>,
+    {
+        if let Err(e) = self.profiler.proto().and_then(|b| Self::write_file(&path, &b)) {
+            tracing::warn!(message = "unable to collect and write pprof data", path = ?path.as_ref(), err = %e);
+        }
+    }
+
+    fn write_file<P>(path: P, contents: &[u8]) -> Result<(), MtopError>
+    where
+        P: AsRef<Path>,
+    {
+        if contents.is_empty() {
+            return Err(MtopError::configuration("mtop built without profiling support"));
+        }
+
+        let mut file = File::options().create(true).truncate(true).write(true).open(path)?;
+        file.write_all(contents)?;
+        file.flush()?;
+        Ok(file.sync_all()?)
+    }
+}