hyperactor_telemetry/
trace.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9use std::env;
10
11use crate::env::execution_id;
12
13const MONARCH_CLIENT_TRACE_ID: &str = "MONARCH_CLIENT_TRACE_ID";
14
15/// Returns the current trace ID if it exists, or None if it doesn't.
16/// This function does not create a new trace ID if one doesn't exist.
17/// Todo: Eventually use Message Headers to relay this traceid instead of
18/// env vars
19pub fn get_trace_id() -> Option<String> {
20    if let Ok(env_id) = env::var(MONARCH_CLIENT_TRACE_ID) {
21        if !env_id.is_empty() {
22            return Some(env_id);
23        }
24    }
25
26    // No trace ID found
27    None
28}
29
30/// Returns the current trace ID, if none exists, set the current execution id as the trace id.
31/// This ensures that the client trace id and execution id is the same.
32/// The trace ID remains the same as long as it is in the same process.
33/// Use this method only on the client side.
34pub fn get_or_create_trace_id() -> String {
35    if let Ok(existing_trace_id) = std::env::var(MONARCH_CLIENT_TRACE_ID) {
36        if !existing_trace_id.is_empty() {
37            return existing_trace_id;
38        }
39    }
40
41    let trace_id = execution_id().clone();
42    // Safety: Can be unsound if there are multiple threads
43    // reading and writing the environment.
44    unsafe {
45        std::env::set_var(MONARCH_CLIENT_TRACE_ID, trace_id.clone());
46    }
47
48    trace_id
49}
hyperactor_telemetry/trace.rs

hyperactor_telemetry/
trace.rs
