Skip to content

Commit

Permalink
refactor(test): setup tracing subscriber in before_test_async()
Browse files Browse the repository at this point in the history
  • Loading branch information
rami3l committed Jun 14, 2024
1 parent 5f34b12 commit 6d5f0f6
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 55 deletions.
2 changes: 1 addition & 1 deletion rustup-macros/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -95,7 +95,7 @@ fn test_inner(mod_path: String, mut input: ItemFn) -> syn::Result<TokenStream> {
let name = input.sig.ident.clone();
let new_block: Block = parse_quote! {
{
#before_ident().await;
let _guard = #before_ident().await;
// Define a function with same name we can instrument inside the
// tracing enablement logic.
#[cfg_attr(feature = "otel", tracing::instrument(skip_all))]
Expand Down
71 changes: 17 additions & 54 deletions src/test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -224,69 +224,32 @@ where
f(&rustup_home)
}

#[cfg(feature = "otel")]
use once_cell::sync::Lazy;

/// A tokio runtime for the sync tests, permitting the use of tracing. This is
/// never shutdown, instead it is just dropped at end of process.
#[cfg(feature = "otel")]
static TRACE_RUNTIME: Lazy<tokio::runtime::Runtime> =
Lazy::new(|| tokio::runtime::Runtime::new().unwrap());
/// A tracer for the tests.
#[cfg(feature = "otel")]
static TRACER: Lazy<opentelemetry_sdk::trace::Tracer> = Lazy::new(|| {
use std::time::Duration;

use opentelemetry::{global, KeyValue};
use opentelemetry_otlp::WithExportConfig;
use opentelemetry_sdk::{
propagation::TraceContextPropagator,
trace::{self, Sampler},
Resource,
};
use tokio::runtime::Handle;
use tracing_subscriber::{layer::SubscriberExt, EnvFilter, Registry};

// Use the current runtime, or the sync test runtime otherwise.
let handle = match Handle::try_current() {
Ok(handle) => handle,
Err(_) => TRACE_RUNTIME.handle().clone(),
};
let _guard = handle.enter();
pub async fn before_test_async() -> Option<tracing::dispatcher::DefaultGuard> {
#[cfg(feature = "otel")]
{
use tracing_subscriber::{layer::SubscriberExt, Registry};

let tracer = opentelemetry_otlp::new_pipeline()
.tracing()
.with_exporter(
opentelemetry_otlp::new_exporter()
.tonic()
.with_timeout(Duration::from_secs(3)),
)
.with_trace_config(
trace::config()
.with_sampler(Sampler::AlwaysOn)
.with_resource(Resource::new(vec![KeyValue::new("service.name", "rustup")])),
)
.install_batch(opentelemetry_sdk::runtime::Tokio)
.unwrap();
let telemetry = {
use opentelemetry::global;
use opentelemetry_sdk::propagation::TraceContextPropagator;

global::set_text_map_propagator(TraceContextPropagator::new());
let env_filter = EnvFilter::try_from_default_env().unwrap_or(EnvFilter::new("INFO"));
let telemetry = tracing_opentelemetry::layer().with_tracer(tracer.clone());
let subscriber = Registry::default().with(env_filter).with(telemetry);
tracing::subscriber::set_global_default(subscriber).unwrap();
tracer
});
global::set_text_map_propagator(TraceContextPropagator::new());
crate::cli::log::telemetry()
};

pub async fn before_test_async() {
#[cfg(feature = "otel")]
let subscriber = Registry::default().with(telemetry);
Some(tracing::subscriber::set_default(subscriber))
}
#[cfg(not(feature = "otel"))]
{
Lazy::force(&TRACER);
None
}
}

pub async fn after_test_async() {
#[cfg(feature = "otel")]
{
TRACER.provider().map(|p| p.force_flush());
// We're tracing, so block until all spans are exported.
opentelemetry::global::shutdown_tracer_provider();
}
}

0 comments on commit 6d5f0f6

Please sign in to comment.