Support trace-level sampling

tracing-honeycomb/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "tracing-honeycomb"
-version = "0.2.0"
+version = "0.2.1"
 authors = ["Inanna Malick <[email protected]>"]
 edition = "2018"
 description = "Honeycomb.io tracing layer for multiprocess telemetry"
@@ -13,7 +13,7 @@ readme = "README.md"
 [dependencies]
 tracing = "0.1.12"
 tracing-core = "0.1.9"
-tracing-distributed =  { path = "../tracing-distributed" , version = "0.2.0"}
+tracing-distributed =  { path = "../tracing-distributed", version = "0.2.0" }
 libhoney-rust = "0.1.3"
 rand = "0.7"
 chrono = "0.4.9"

tracing-honeycomb/src/honeycomb.rs

@@ -9,17 +9,18 @@ use tracing_distributed::{Event, Span, Telemetry};
 #[derive(Debug)]
 pub struct HoneycombTelemetry {
     honeycomb_client: Mutex<libhoney::Client<libhoney::transmission::Transmission>>,
+    sample_rate: usize,
 }
 
 impl HoneycombTelemetry {
-    pub(crate) fn new(cfg: libhoney::Config) -> Self {
+    pub(crate) fn new(cfg: libhoney::Config, sample_rate: usize) -> Self {
         let honeycomb_client = libhoney::init(cfg);
 
         // publishing requires &mut so just mutex-wrap it
         // FIXME: may not be performant, investigate options (eg mpsc)
         let honeycomb_client = Mutex::new(honeycomb_client);
 
-        HoneycombTelemetry { honeycomb_client }
+        HoneycombTelemetry { honeycomb_client, sample_rate }
     }
 
     fn report_data(&self, data: HashMap<String, ::libhoney::Value>) {
@@ -34,6 +35,14 @@ impl HoneycombTelemetry {
             eprintln!("error sending event to honeycomb, {:?}", err);
         }
     }
+
+    fn should_report(&self, trace_id: TraceId) -> bool {
+        if self.sample_rate <= 1 {
+            return true;
+        }
+
+        trace_id.0 as usize % self.sample_rate == 0
+    }
 }
 
 impl Telemetry for HoneycombTelemetry {
@@ -46,13 +55,17 @@ impl Telemetry for HoneycombTelemetry {
     }
 
     fn report_span(&self, span: Span<Self::Visitor, Self::SpanId, Self::TraceId>) {
-        let data = span_to_values(span);
-        self.report_data(data);
+        if self.should_report(span.trace_id) {
+            let data = span_to_values(span);
+            self.report_data(data);
+        }
     }
 
     fn report_event(&self, event: Event<Self::Visitor, Self::SpanId, Self::TraceId>) {
-        let data = event_to_values(event);
-        self.report_data(data);
+        if self.should_report(event.trace_id) {
+            let data = event_to_values(event);
+            self.report_data(data);
+        }
     }
 }
 

tracing-honeycomb/src/lib.rs

@@ -66,7 +66,38 @@ pub fn new_honeycomb_telemetry_layer(
     let instance_id: u64 = rand::thread_rng().gen();
     TelemetryLayer::new(
         service_name,
-        HoneycombTelemetry::new(honeycomb_config),
+        HoneycombTelemetry::new(honeycomb_config, 1),
+        move |tracing_id| SpanId {
+            instance_id,
+            tracing_id,
+        },
+    )
+}
+
+/// Construct a TelemetryLayer that publishes telemetry to honeycomb.io using the
+/// provided honeycomb config, and sample rate. This function differs from
+/// `new_honeycomb_telemetry_layer` and the `sample_rate` on the
+/// `libhoney::Config` there in an important way. `libhoney` samples `Event`
+/// data, which is individual spans on each trace. This means that using the
+/// sampling logic in libhoney may result in missing event data or incomplete
+/// traces. Calling this function provides trace-level sampling, meaning sampling
+/// decisions are based on a modulo of the traceID, and events in a single trace
+/// will not be sampled differently. If the trace is sampled, then all spans
+/// under it will be sent to honeycomb. If a trace is not sampled, no spans or
+/// events under it will be sent. When using this trace-level sampling, the
+/// `sample_rate` parameter on the `libhoney::Config` should be set to 1, which
+/// is the default.
+///
+/// Specialized to the honeycomb.io-specific SpanId and TraceId provided by this crate.
+pub fn new_honeycomb_telemetry_layer_with_trace_sampling(
+    service_name: &'static str,
+    honeycomb_config: libhoney::Config,
+    sample_rate: usize,
+) -> TelemetryLayer<HoneycombTelemetry, SpanId, TraceId> {
+    let instance_id: u64 = rand::thread_rng().gen();
+    TelemetryLayer::new(
+        service_name,
+        HoneycombTelemetry::new(honeycomb_config, sample_rate),
         move |tracing_id| SpanId {
             instance_id,
             tracing_id,