From 43599c2a025ba50cdd4d01d7cefdd9ebfdeb8869 Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Wed, 5 Aug 2020 00:45:56 -0400
Subject: [PATCH] add newrelic attribute for speaker

---
 src/main.rs            | 33 +++++++++++++++++++++------------
 src/newrelic_logger.rs | 25 ++++++++++++++++++++++---
 2 files changed, 43 insertions(+), 15 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index de6faf2..ef22eb1 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -5,12 +5,16 @@
 pub mod newrelic_logger;
 
 use anyhow::Result as AnyhowResult;
+use crate::newrelic_logger::{NewRelicLogger, MaybeNewRelicTransaction};
 use futures::future::{self, Future};
 use hyper::Method;
+use hyper::header::HeaderValue;
+use hyper::http::header::HeaderName;
 use hyper::rt::Stream;
 use hyper::server::conn::AddrStream;
 use hyper::service::{service_fn, make_service_fn};
 use hyper::{Body, Request, Response, Server};
+use rand::seq::IteratorRandom;
 use std::collections::HashMap;
 use std::env;
 use std::fmt::Debug;
@@ -18,11 +22,6 @@ use std::fmt::Display;
 use std::net::SocketAddr;
 use std::str::FromStr;
 use std::sync::Arc;
-use rand::seq::IteratorRandom;
-use hyper::http::header::HeaderName;
-use hyper::header::HeaderValue;
-use newrelic::App;
-use crate::newrelic_logger::NewRelicLogger;
 
 type BoxFut = Box<dyn Future<Item=Response<Body>, Error=hyper::Error> + Send>;
 
@@ -134,7 +133,13 @@ impl Router {
   }
 }
 
-fn speak_proxy(req: Request<Body>, remote_addr: SocketAddr, router: Arc<Router>, endpoint: &'static str) -> BoxFut {
+fn speak_proxy(
+  req: Request<Body>,
+  remote_addr: SocketAddr,
+  router: Arc<Router>,
+  transaction: MaybeNewRelicTransaction,
+  endpoint: &'static str) -> BoxFut
+{
   let mut headers = req.headers().clone();
 
   // This is the IP of the Digital Ocean load balancer.
@@ -173,6 +178,8 @@ fn speak_proxy(req: Request<Body>, remote_addr: SocketAddr, router: Arc<Router>,
 
     info!("Routing {} for {} to {}", endpoint, &request_details.speaker, &proxy_host);
 
+    transaction.add_attribute("speaker", &request_details.speaker);
+
     let mut request_builder = Request::builder();
     request_builder.method(Method::POST)
       .uri(endpoint);
@@ -247,15 +254,17 @@ fn main() -> AnyhowResult<()> {
 
       match (req.method(), req.uri().path()) {
         (&Method::POST, "/speak") => {
-          let _droppable_transaction = newrelic_logger3.web_transaction("proxy /speak");
-          speak_proxy(req, remote_addr.clone(), router3, "/speak")
+          let nr_transaction = newrelic_logger3.web_transaction("/speak");
+          speak_proxy(req, remote_addr.clone(), router3, nr_transaction, "/speak")
         },
         (&Method::POST, "/speak_spectrogram") => {
-          let _droppable_transaction = newrelic_logger3.web_transaction("proxy /speak_spectrogram");
-          speak_proxy(req, remote_addr.clone(), router3, "/speak_spectrogram")
+          let nr_transaction = newrelic_logger3.web_transaction("/speak_spectrogram");
+          speak_proxy(req, remote_addr.clone(), router3, nr_transaction, "/speak_spectrogram")
+        },
+        (&Method::GET, "/proxy_health") => {
+          let _droppable_transaction = newrelic_logger3.web_transaction("/proxy_health");
+          health_check_response(req)
         },
-        (&Method::GET, "/proxy_health") =>
-          health_check_response(req),
         _ => {
           let _droppable_transaction = newrelic_logger3.web_transaction("unmatched endpoint");
           let forward = router3.get_random_host();
diff --git a/src/newrelic_logger.rs b/src/newrelic_logger.rs
index ea6b74b..031ba0c 100644
--- a/src/newrelic_logger.rs
+++ b/src/newrelic_logger.rs
@@ -26,11 +26,30 @@ impl NewRelicLogger {
     }
   }
 
-  pub fn web_transaction(&self, transaction_name: &str) -> Option<Transaction> {
+  pub fn web_transaction(&self, transaction_name: &str) -> MaybeNewRelicTransaction {
     if let Some(ref newrelic_api) = self.newrelic_api {
       info!("Logging request to NewRelic");
-      return newrelic_api.web_transaction(transaction_name).ok();
+      let maybe_transaction = newrelic_api.web_transaction(transaction_name).ok();
+      return MaybeNewRelicTransaction::from_optional(maybe_transaction);
     }
-    None
+    MaybeNewRelicTransaction::from_optional(None)
   }
 }
+
+pub struct MaybeNewRelicTransaction {
+  transaction: Option<Transaction>,
+}
+
+impl MaybeNewRelicTransaction {
+  pub fn from_optional(optional: Option<Transaction>) -> Self {
+    Self {
+      transaction: optional,
+    }
+  }
+
+  pub fn add_attribute(&self, name: &str, value: &str) {
+    if let Some(ref transaction) = self.transaction {
+      let _result = transaction.add_attribute(name, value);
+    }
+  }
+}
\ No newline at end of file