misc: rework compression feature

Author: Miaxos

src/compression.rs

@@ -0,0 +1,25 @@
+#[cfg(feature = "compression")]
+use libflate::gzip;
+
+const TARGET_LOG_COMPRESSION: &str = "apollo-studio-extension-compression";
+
+#[cfg(feature = "compression")]
+pub fn compress(msg: Vec<u8>) -> Result<Vec<u8>, std::io::Error> {
+    let mut encoder = gzip::Encoder::new(Vec::new()).unwrap();
+    let mut msg = std::io::Cursor::new(msg);
+
+    match std::io::copy(&mut msg, &mut encoder) {
+        Ok(_) => {}
+        Err(e) => {
+            error!(target: TARGET_LOG_COMPRESSION, message = "An issue happened while GZIP compression", err = ?e);
+            return Err(e);
+        }
+    };
+
+    encoder.finish().into_result()
+}
+
+#[cfg(not(feature = "compression"))]
+pub fn compress(msg: Vec<u8>) -> Result<Vec<u8>, std::io::Error> {
+    Ok::<Vec<u8>, std::io::Error>(msg)
+}

src/lib.rs

@@ -1,14 +1,11 @@
 //! # Apollo Studio Extension for Performance Tracing for async_graphql crates
+mod compression;
 mod packages;
 mod proto;
 pub mod register;
 
 #[macro_use]
 extern crate tracing;
-/*
- * TODO:
- *   - Gzip compression (libflate)
- */
 use packages::uname::Uname;
 use std::collections::HashMap;
 use std::sync::Arc;
@@ -203,41 +200,22 @@ impl ApolloTracing {
                         }
                     };
 
-                    let result = if cfg!(features = "compression") {
-                        let mut encoder = libflate::gzip::Encoder::new(Vec::new()).unwrap();
-                        let mut msg = std::io::Cursor::new(msg);
-                        match std::io::copy(&mut msg, &mut encoder) {
-                            Ok(_) => {}
-                            Err(e) => {
-                                error!(target: TARGET_LOG, message = "An issue happened while GZIP compression", err = ?e);
-                                continue;
-                            }
-                        };
-                        match encoder.finish().into_result() {
-                            Ok(msg) => {
-                                client
-                                    .post(REPORTING_URL)
-                                    .body(msg)
-                                    .header("content-type", "application/protobuf")
-                                    .header("X-Api-Key", &authorization_token)
-                                    .send()
-                                    .await
-                            }
-                            Err(e) => {
-                                error!(target: TARGET_LOG, message = "An issue happened while GZIP compression", err = ?e);
-                                continue;
-                            }
+                    let msg = match compression::compress(msg) {
+                        Ok(result) => result,
+                        Err(e) => {
+                            error!(target: TARGET_LOG, message = "An issue happened while GZIP compression", err = ?e);
+                            continue;
                         }
-                    } else {
-                        client
-                            .post(REPORTING_URL)
-                            .body(msg)
-                            .header("content-type", "application/protobuf")
-                            .header("X-Api-Key", &authorization_token)
-                            .send()
-                            .await
                     };
 
+                    let result = client
+                        .post(REPORTING_URL)
+                        .body(msg)
+                        .header("content-type", "application/protobuf")
+                        .header("X-Api-Key", &authorization_token)
+                        .send()
+                        .await;
+
                     match result {
                         Ok(data) => {
                             span_batch.record("response", &debug(&data));