1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157

//  Copyright 2017 Palantir Technologies, Inc.
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.

//! Tracers.
use crate::trace_context;
use crate::{
    Attached, Endpoint, OpenSpan, Report, Sample, SamplingFlags, Span, SpanId, SpanState,
    TraceContext, TraceId,
};
use lazycell::AtomicLazyCell;
use rand::Rng;
use std::error::Error;
use std::fmt;
use std::time::{Instant, SystemTime};

pub(crate) static TRACER: AtomicLazyCell<Tracer> = AtomicLazyCell::NONE;

pub(crate) struct Tracer {
    pub sampler: Box<dyn Sample + Sync + Send>,
    pub reporter: Box<dyn Report + Sync + Send>,
    pub local_endpoint: Endpoint,
}

/// Initializes the global tracer.
///
/// The tracer can only be initialized once in the lifetime of a program. Spans created before this function is called
/// will be no-ops.
///
/// Returns an error if the tracer is already initialized.
pub fn set_tracer<S, R>(
    sampler: S,
    reporter: R,
    local_endpoint: Endpoint,
) -> Result<(), SetTracerError>
where
    S: Sample + 'static + Sync + Send,
    R: Report + 'static + Sync + Send,
{
    TRACER
        .fill(Tracer {
            sampler: Box::new(sampler),
            reporter: Box::new(reporter),
            local_endpoint,
        })
        .map_err(|_| SetTracerError(()))
}

/// The error returned when attempting to set a tracer when one is already installed.
#[derive(Debug)]
pub struct SetTracerError(());

impl fmt::Display for SetTracerError {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt.write_str("")
    }
}

impl Error for SetTracerError {}

/// Starts a new trace.
pub fn new_trace() -> OpenSpan<Attached> {
    new_trace_from(SamplingFlags::default())
}

/// Stats a new trace with specific sampling flags.
pub fn new_trace_from(flags: SamplingFlags) -> OpenSpan<Attached> {
    let id = next_id();
    let context = TraceContext::builder()
        .trace_id(TraceId::from(id))
        .span_id(SpanId::from(id))
        .sampling_flags(flags)
        .build();
    make_span(context, false)
}

/// Joins an existing trace.
///
/// The context can come from, for example, the headers of an HTTP request.
pub fn join_trace(context: TraceContext) -> OpenSpan<Attached> {
    make_span(context, true)
}

/// Stats a new span with the specified parent.
pub fn new_child(parent: TraceContext) -> OpenSpan<Attached> {
    let id = next_id();
    let context = TraceContext::builder()
        .trace_id(parent.trace_id())
        .parent_id(parent.span_id())
        .span_id(SpanId::from(id))
        .sampling_flags(parent.sampling_flags())
        .build();
    make_span(context, false)
}

/// Creates a new span parented to the current one if it exists, or starting a new trace otherwise.
pub fn next_span() -> OpenSpan<Attached> {
    match crate::current() {
        Some(context) => new_child(context),
        None => new_trace(),
    }
}

fn next_id() -> [u8; 8] {
    let mut id = [0; 8];
    rand::thread_rng().fill(&mut id);
    id
}

fn make_span(mut context: TraceContext, mut shared: bool) -> OpenSpan<Attached> {
    let tracer = match TRACER.borrow() {
        Some(tracer) => tracer,
        None => return OpenSpan::new(context, SpanState::Nop),
    };

    if context.sampled().is_none() {
        context = trace_context::Builder::from(context)
            .sampled(tracer.sampler.sample(context.trace_id()))
            .build();
        // since the thing we got the context from didn't indicate if it should be sampled,
        // we can't assume they're recording the span as well.
        shared = false;
    }

    let state = match context.sampled() {
        Some(false) => SpanState::Nop,
        _ => {
            let mut span = Span::builder();
            span.trace_id(context.trace_id())
                .id(context.span_id())
                .timestamp(SystemTime::now())
                .shared(shared)
                .local_endpoint(tracer.local_endpoint.clone());

            if let Some(parent_id) = context.parent_id() {
                span.parent_id(parent_id);
            }

            SpanState::Real {
                span,
                start_instant: Instant::now(),
            }
        }
    };

    OpenSpan::new(context, state)
}