mirror of
https://github.com/openai/codex.git
synced 2026-10-18 17:18:52 +00:00
Reload app-server telemetry after account changes (#37339)
## Why An account change can select different OpenTelemetry collector settings. The long-running app server needs to stop exporting telemetry with the previous account's configuration. ## What changed - Rebuild and install the log, trace, and metrics exporters when authentication changes, and shut down the previous provider without blocking the app server. - Make cached global metrics clients and the tracing layer follow the replacement provider so existing callers continue exporting through the current account's configuration. - Clear the cached cloud configuration loader after login and logout so telemetry reloads from the new account context. ## Testing - Add an app-server integration test that switches accounts and verifies logs, traces, and metrics reach the new collector while trace context is preserved. - Add a metrics test covering cached clients after provider replacement. GitOrigin-RevId: a74359777a35edbee5bbe0b2baccc155074cb172
This commit is contained in:
@@ -113,6 +113,14 @@ impl ConfigManager {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn clear_cloud_config_bundle_loader(&self) {
|
||||
if let Ok(mut guard) = self.cloud_config_bundle.write() {
|
||||
*guard = CloudConfigBundleLoader::default();
|
||||
} else {
|
||||
warn!("failed to clear cloud config bundle loader");
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn replace_thread_config_loader(
|
||||
&self,
|
||||
thread_config_loader: Arc<dyn ThreadConfigLoader>,
|
||||
|
||||
@@ -111,6 +111,7 @@ mod mcp_refresh;
|
||||
mod message_processor;
|
||||
mod models;
|
||||
mod models_refresh_worker;
|
||||
mod otel_reloader;
|
||||
mod outgoing_message;
|
||||
mod request_processors;
|
||||
mod request_serialization;
|
||||
@@ -667,15 +668,13 @@ pub async fn run_main_with_transport_options(
|
||||
let log_db_layer = log_db
|
||||
.clone()
|
||||
.map(|layer| layer.with_filter(log_db::default_filter()));
|
||||
let otel_logger_layer = otel.as_ref().and_then(|o| o.logger_layer());
|
||||
let otel_tracing_layer = otel.as_ref().and_then(|o| o.tracing_layer());
|
||||
let (otel_layers, otel_logger_reload_handle) = otel_reloader::layers(otel.as_ref());
|
||||
let _ = tracing_subscriber::registry()
|
||||
.with(stderr_fmt)
|
||||
.with(feedback_layer)
|
||||
.with(feedback_metadata_layer)
|
||||
.with(log_db_layer)
|
||||
.with(otel_logger_layer)
|
||||
.with(otel_tracing_layer)
|
||||
.with(otel_layers)
|
||||
.try_init();
|
||||
for warning in &config_warnings {
|
||||
match &warning.details {
|
||||
@@ -816,6 +815,15 @@ pub async fn run_main_with_transport_options(
|
||||
}
|
||||
transport_accept_handles.push(remote_control_accept_handle);
|
||||
|
||||
let otel_reloader_handle = otel_reloader::spawn(
|
||||
otel,
|
||||
otel_logger_reload_handle,
|
||||
config_manager.clone(),
|
||||
Arc::clone(&auth_manager),
|
||||
default_analytics_enabled,
|
||||
transport_shutdown_token.clone(),
|
||||
);
|
||||
|
||||
let outbound_handle = tokio::spawn(async move {
|
||||
let mut outbound_connections = HashMap::<ConnectionId, OutboundConnectionState>::new();
|
||||
loop {
|
||||
@@ -1176,14 +1184,11 @@ pub async fn run_main_with_transport_options(
|
||||
let _ = outbound_handle.await;
|
||||
|
||||
transport_shutdown_token.cancel();
|
||||
let _ = otel_reloader_handle.await;
|
||||
for handle in transport_accept_handles {
|
||||
let _ = handle.await;
|
||||
}
|
||||
|
||||
if let Some(otel) = otel {
|
||||
otel.shutdown();
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,112 @@
|
||||
use crate::OTEL_SERVICE_NAME;
|
||||
use crate::config_manager::ConfigManager;
|
||||
use codex_login::AuthManager;
|
||||
use codex_otel::OtelProvider;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio::task::JoinHandle;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::Subscriber;
|
||||
use tracing::info;
|
||||
use tracing::warn;
|
||||
use tracing_subscriber::Layer;
|
||||
use tracing_subscriber::registry::LookupSpan;
|
||||
use tracing_subscriber::reload;
|
||||
|
||||
type OtelExportLayer<S> = Option<Box<dyn Layer<S> + Send + Sync + 'static>>;
|
||||
type OtelReloadLayers<S> = (
|
||||
Vec<Box<dyn Layer<S> + Send + Sync + 'static>>,
|
||||
reload::Handle<OtelExportLayer<S>, S>,
|
||||
);
|
||||
|
||||
pub(crate) fn layers<S>(provider: Option<&OtelProvider>) -> OtelReloadLayers<S>
|
||||
where
|
||||
S: Subscriber + for<'span> LookupSpan<'span> + Send + Sync + 'static,
|
||||
{
|
||||
let logger_export_layer: OtelExportLayer<S> = provider
|
||||
.and_then(OtelProvider::logger_export_layer)
|
||||
.map(Layer::boxed);
|
||||
let (logger_layer, logger_handle) = reload::Layer::new(logger_export_layer);
|
||||
|
||||
let mut layers: Vec<Box<dyn Layer<S> + Send + Sync + 'static>> = vec![
|
||||
logger_layer
|
||||
.with_filter(tracing_subscriber::filter::filter_fn(
|
||||
OtelProvider::log_export_filter,
|
||||
))
|
||||
.boxed(),
|
||||
];
|
||||
if provider.is_some_and(|provider| provider.tracer.is_some()) {
|
||||
layers.push(OtelProvider::reloadable_tracing_layer(OTEL_SERVICE_NAME).boxed());
|
||||
}
|
||||
(layers, logger_handle)
|
||||
}
|
||||
|
||||
pub(crate) fn spawn<S>(
|
||||
mut provider: Option<OtelProvider>,
|
||||
logger_reload_handle: reload::Handle<OtelExportLayer<S>, S>,
|
||||
config_manager: ConfigManager,
|
||||
auth_manager: Arc<AuthManager>,
|
||||
default_analytics_enabled: bool,
|
||||
shutdown_token: CancellationToken,
|
||||
) -> JoinHandle<()>
|
||||
where
|
||||
S: Subscriber + for<'span> LookupSpan<'span> + Send + Sync + 'static,
|
||||
{
|
||||
let mut auth_changes = auth_manager.auth_change_receiver();
|
||||
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = shutdown_token.cancelled() => break,
|
||||
changed = auth_changes.changed() => {
|
||||
if changed.is_err() {
|
||||
break;
|
||||
}
|
||||
|
||||
// Account handlers install the new cloud loader after publishing auth changes.
|
||||
tokio::time::sleep(Duration::from_millis(/*millis*/ 50)).await;
|
||||
|
||||
let config = match config_manager.load_latest_config(/*fallback_cwd*/ None).await {
|
||||
Ok(config) => config,
|
||||
Err(error) => {
|
||||
warn!(%error, "failed to reload telemetry config after account change");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let next_provider = match codex_core::otel_init::build_provider(
|
||||
&config,
|
||||
env!("CARGO_PKG_VERSION"),
|
||||
Some(OTEL_SERVICE_NAME),
|
||||
default_analytics_enabled,
|
||||
) {
|
||||
Ok(provider) => provider,
|
||||
Err(error) => {
|
||||
warn!(%error, "failed to rebuild telemetry exporters after account change");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
if let Err(error) = logger_reload_handle.reload(
|
||||
next_provider
|
||||
.as_ref()
|
||||
.and_then(OtelProvider::logger_export_layer)
|
||||
.map(Layer::boxed),
|
||||
) {
|
||||
warn!(%error, "failed to install telemetry exporters after account change");
|
||||
continue;
|
||||
}
|
||||
if let Some(previous_provider) = std::mem::replace(&mut provider, next_provider) {
|
||||
drop(tokio::task::spawn_blocking(move || previous_provider.shutdown()));
|
||||
}
|
||||
info!(
|
||||
event.name = "codex.app_server.otel_reloaded",
|
||||
"reloaded telemetry exporters after account change"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(provider) = provider {
|
||||
let _ = tokio::task::spawn_blocking(move || provider.shutdown()).await;
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -372,6 +372,7 @@ impl AccountRequestProcessor {
|
||||
) {
|
||||
Ok(()) => {
|
||||
self.auth_manager.reload().await;
|
||||
self.config_manager.clear_cloud_config_bundle_loader();
|
||||
Ok(())
|
||||
}
|
||||
Err(err) => Err(internal_error(format!("failed to save api key: {err}"))),
|
||||
@@ -439,6 +440,7 @@ impl AccountRequestProcessor {
|
||||
)
|
||||
.map_err(|err| internal_error(format!("failed to save Amazon Bedrock auth: {err}")))?;
|
||||
self.auth_manager.reload().await;
|
||||
self.config_manager.clear_cloud_config_bundle_loader();
|
||||
Ok(LoginAccountResponse::AmazonBedrock {})
|
||||
}
|
||||
.await;
|
||||
@@ -890,6 +892,8 @@ impl AccountRequestProcessor {
|
||||
}
|
||||
}
|
||||
|
||||
self.config_manager.clear_cloud_config_bundle_loader();
|
||||
|
||||
if managed_bedrock_auth {
|
||||
clear_user_model_provider_if_bedrock(&self.config_manager).await?;
|
||||
}
|
||||
|
||||
@@ -47,6 +47,7 @@ mod memory_reset;
|
||||
mod model_list;
|
||||
mod model_provider_capabilities_read;
|
||||
mod multi_agent_v2_developer_instructions;
|
||||
mod otel;
|
||||
mod output_schema;
|
||||
mod permission_profile_list;
|
||||
mod plan_item;
|
||||
|
||||
@@ -0,0 +1,160 @@
|
||||
use anyhow::Context;
|
||||
use anyhow::Result;
|
||||
use app_test_support::ChatGptAuthFixture;
|
||||
use app_test_support::ChatGptIdTokenClaims;
|
||||
use app_test_support::TestAppServer;
|
||||
use app_test_support::encode_id_token;
|
||||
use app_test_support::write_chatgpt_auth;
|
||||
use app_test_support::write_models_cache;
|
||||
use codex_app_server_protocol::LoginAccountResponse;
|
||||
use codex_app_server_protocol::ThreadStartParams;
|
||||
use codex_config::types::AuthCredentialsStoreMode;
|
||||
use pretty_assertions::assert_eq;
|
||||
use std::path::Path;
|
||||
use std::time::Duration;
|
||||
use tempfile::TempDir;
|
||||
use tokio::time::timeout;
|
||||
use wiremock::Mock;
|
||||
use wiremock::MockServer;
|
||||
use wiremock::ResponseTemplate;
|
||||
use wiremock::matchers::method;
|
||||
|
||||
const TEST_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 30);
|
||||
const INITIAL_ACCOUNT_ID: &str = "123e4567-e89b-42d3-a456-426614174000";
|
||||
const NEXT_ACCOUNT_ID: &str = "123e4567-e89b-42d3-a456-426614174001";
|
||||
const INITIAL_EMAIL: &str = "initial-workspace@example.com";
|
||||
const NEXT_EMAIL: &str = "next-account@example.com";
|
||||
const PARENT_TRACE_ID: &str = "4bf92f3577b34da6a3ce929d0e0e4736";
|
||||
const PARENT_TRACEPARENT: &str = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01";
|
||||
|
||||
#[tokio::test]
|
||||
async fn account_switch_reloads_telemetry_collectors_and_preserves_trace_context() -> Result<()> {
|
||||
let collector = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.respond_with(ResponseTemplate::new(200))
|
||||
.mount(&collector)
|
||||
.await;
|
||||
|
||||
let codex_home = TempDir::new()?;
|
||||
let initial_endpoint = format!("{}/initial", collector.uri());
|
||||
write_otel_config(codex_home.path(), &initial_endpoint)?;
|
||||
write_chatgpt_auth(
|
||||
codex_home.path(),
|
||||
ChatGptAuthFixture::new("initial-access-token")
|
||||
.account_id(INITIAL_ACCOUNT_ID)
|
||||
.chatgpt_account_id(INITIAL_ACCOUNT_ID)
|
||||
.chatgpt_user_id("initial-user")
|
||||
.plan_type("pro")
|
||||
.email(INITIAL_EMAIL),
|
||||
AuthCredentialsStoreMode::File,
|
||||
)?;
|
||||
write_models_cache(codex_home.path())?;
|
||||
|
||||
let mut app_server = TestAppServer::builder()
|
||||
.with_codex_home(codex_home.path())
|
||||
.with_env_overrides(&[("TRACEPARENT", Some(PARENT_TRACEPARENT))])
|
||||
.with_json_logging("codex_app_server::otel_reloader=info")
|
||||
.build_initialized_with_timeout(TEST_TIMEOUT)
|
||||
.await?;
|
||||
app_server
|
||||
.start_thread(ThreadStartParams::default())
|
||||
.await?;
|
||||
|
||||
let next_endpoint = format!("{}/next", collector.uri());
|
||||
write_otel_config(codex_home.path(), &next_endpoint)?;
|
||||
let access_token = encode_id_token(
|
||||
&ChatGptIdTokenClaims::new()
|
||||
.email(NEXT_EMAIL)
|
||||
.plan_type("pro")
|
||||
.chatgpt_account_id(NEXT_ACCOUNT_ID)
|
||||
.chatgpt_user_id("next-user"),
|
||||
)?;
|
||||
let request_id = app_server
|
||||
.send_chatgpt_auth_tokens_login_request(
|
||||
access_token,
|
||||
NEXT_ACCOUNT_ID.to_string(),
|
||||
Some("pro".to_string()),
|
||||
)
|
||||
.await?;
|
||||
let response: LoginAccountResponse =
|
||||
timeout(TEST_TIMEOUT, app_server.read_response(request_id)).await??;
|
||||
assert_eq!(response, LoginAccountResponse::ChatgptAuthTokens {});
|
||||
timeout(
|
||||
TEST_TIMEOUT,
|
||||
app_server.wait_for_json_log_event("codex.app_server.otel_reloaded"),
|
||||
)
|
||||
.await??;
|
||||
app_server
|
||||
.start_thread(ThreadStartParams::default())
|
||||
.await?;
|
||||
|
||||
let status = timeout(TEST_TIMEOUT, app_server.shutdown_gracefully()).await??;
|
||||
assert!(status.success(), "app-server did not shut down cleanly");
|
||||
|
||||
let requests = collector
|
||||
.received_requests()
|
||||
.await
|
||||
.context("collector did not record requests")?;
|
||||
let exported = |path: &str| {
|
||||
requests
|
||||
.iter()
|
||||
.filter(|request| request.url.path() == path)
|
||||
.map(|request| String::from_utf8_lossy(&request.body).into_owned())
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n")
|
||||
};
|
||||
let initial_logs = exported("/initial/logs");
|
||||
let initial_traces = exported("/initial/traces");
|
||||
let next_logs = exported("/next/logs");
|
||||
let next_traces = exported("/next/traces");
|
||||
let next_metrics = exported("/next/metrics");
|
||||
|
||||
assert!(
|
||||
initial_logs.contains(INITIAL_EMAIL),
|
||||
"the initial account's logs were not exported: {initial_logs}"
|
||||
);
|
||||
assert!(
|
||||
initial_traces.contains(PARENT_TRACE_ID),
|
||||
"the initial account's trace context was not propagated: {initial_traces}"
|
||||
);
|
||||
assert!(
|
||||
next_logs.contains(NEXT_EMAIL),
|
||||
"the next account's logs did not reach its collector: {next_logs}"
|
||||
);
|
||||
assert!(
|
||||
next_traces.contains(PARENT_TRACE_ID),
|
||||
"the next account's trace context was not propagated: {next_traces}"
|
||||
);
|
||||
assert!(
|
||||
next_metrics.contains("codex.thread.started"),
|
||||
"the next account's metrics did not reach its collector: {next_metrics}"
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn write_otel_config(codex_home: &Path, collector_endpoint: &str) -> Result<()> {
|
||||
std::fs::write(
|
||||
codex_home.join("config.toml"),
|
||||
format!(
|
||||
r#"model = "mock-model"
|
||||
model_provider = "mock_provider"
|
||||
|
||||
[model_providers.mock_provider]
|
||||
name = "Mock provider"
|
||||
base_url = "http://127.0.0.1:1/v1"
|
||||
wire_api = "responses"
|
||||
|
||||
[analytics]
|
||||
enabled = true
|
||||
|
||||
[otel]
|
||||
environment = "test"
|
||||
exporter = {{ otlp-http = {{ endpoint = "{collector_endpoint}/logs", protocol = "json" }} }}
|
||||
trace_exporter = {{ otlp-http = {{ endpoint = "{collector_endpoint}/traces", protocol = "json" }} }}
|
||||
metrics_exporter = {{ otlp-http = {{ endpoint = "{collector_endpoint}/metrics", protocol = "json" }} }}
|
||||
"#
|
||||
),
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -37,6 +37,7 @@ use std::collections::BTreeMap;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex;
|
||||
use std::sync::RwLock;
|
||||
use std::sync::Weak;
|
||||
use std::time::Duration;
|
||||
use tracing::debug;
|
||||
@@ -98,7 +99,7 @@ impl MetricReader for SharedManualReader {
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct MetricsClientInner {
|
||||
pub(super) struct MetricsClientInner {
|
||||
meter_provider: SdkMeterProvider,
|
||||
meter: Meter,
|
||||
counters: Mutex<HashMap<InstrumentKey, Counter<u64>>>,
|
||||
@@ -287,7 +288,12 @@ impl MetricsClientInner {
|
||||
|
||||
/// OpenTelemetry metrics client used by Codex.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct MetricsClient(std::sync::Arc<MetricsClientInner>);
|
||||
pub struct MetricsClient {
|
||||
// Keep the original provider so its owner only shuts down its own exporter.
|
||||
pub(super) inner: Arc<MetricsClientInner>,
|
||||
// Installed clients share this slot, so existing clones follow account changes.
|
||||
pub(super) active: Option<Arc<RwLock<Arc<MetricsClientInner>>>>,
|
||||
}
|
||||
|
||||
impl MetricsClient {
|
||||
/// Build a metrics client from configuration and validate defaults.
|
||||
@@ -336,22 +342,36 @@ impl MetricsClient {
|
||||
}
|
||||
};
|
||||
|
||||
Ok(Self(std::sync::Arc::new(MetricsClientInner {
|
||||
meter_provider,
|
||||
meter,
|
||||
counters: Mutex::new(HashMap::new()),
|
||||
gauges: Mutex::new(HashMap::new()),
|
||||
histograms: Mutex::new(HashMap::new()),
|
||||
duration_histograms: Mutex::new(HashMap::new()),
|
||||
runtime_reader,
|
||||
runtime_only_metrics,
|
||||
default_tags,
|
||||
})))
|
||||
Ok(Self {
|
||||
inner: Arc::new(MetricsClientInner {
|
||||
meter_provider,
|
||||
meter,
|
||||
counters: Mutex::new(HashMap::new()),
|
||||
gauges: Mutex::new(HashMap::new()),
|
||||
histograms: Mutex::new(HashMap::new()),
|
||||
duration_histograms: Mutex::new(HashMap::new()),
|
||||
runtime_reader,
|
||||
runtime_only_metrics,
|
||||
default_tags,
|
||||
}),
|
||||
active: None,
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn active_inner(&self) -> Arc<MetricsClientInner> {
|
||||
match &self.active {
|
||||
Some(active) => active
|
||||
.read()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.clone(),
|
||||
None => Arc::clone(&self.inner),
|
||||
}
|
||||
}
|
||||
|
||||
/// Send a single counter increment.
|
||||
pub fn counter(&self, name: &str, inc: i64, tags: &[(&str, &str)]) -> Result<()> {
|
||||
self.0.counter(name, /*description*/ None, inc, tags)
|
||||
self.active_inner()
|
||||
.counter(name, /*description*/ None, inc, tags)
|
||||
}
|
||||
|
||||
/// Send a single counter increment with an instrument description.
|
||||
@@ -362,17 +382,19 @@ impl MetricsClient {
|
||||
inc: i64,
|
||||
tags: &[(&str, &str)],
|
||||
) -> Result<()> {
|
||||
self.0.counter(name, Some(description), inc, tags)
|
||||
self.active_inner()
|
||||
.counter(name, Some(description), inc, tags)
|
||||
}
|
||||
|
||||
/// Send a single histogram sample.
|
||||
pub fn histogram(&self, name: &str, value: i64, tags: &[(&str, &str)]) -> Result<()> {
|
||||
self.0.histogram(name, value, tags)
|
||||
self.active_inner().histogram(name, value, tags)
|
||||
}
|
||||
|
||||
/// Send a single gauge measurement.
|
||||
pub fn gauge(&self, name: &str, value: i64, tags: &[(&str, &str)]) -> Result<()> {
|
||||
self.0.gauge(name, /*description*/ None, value, tags)
|
||||
self.active_inner()
|
||||
.gauge(name, /*description*/ None, value, tags)
|
||||
}
|
||||
|
||||
/// Send a single gauge measurement with an instrument description.
|
||||
@@ -383,7 +405,8 @@ impl MetricsClient {
|
||||
value: i64,
|
||||
tags: &[(&str, &str)],
|
||||
) -> Result<()> {
|
||||
self.0.gauge(name, Some(description), value, tags)
|
||||
self.active_inner()
|
||||
.gauge(name, Some(description), value, tags)
|
||||
}
|
||||
|
||||
/// Register a gauge callback that reports the current value on every collection.
|
||||
@@ -394,7 +417,7 @@ impl MetricsClient {
|
||||
observe: impl Fn() -> i64 + Send + Sync + 'static,
|
||||
tags: &[(&str, &str)],
|
||||
) -> Result<()> {
|
||||
self.0
|
||||
self.active_inner()
|
||||
.register_observable_gauge(name, description, observe, tags)
|
||||
}
|
||||
|
||||
@@ -405,7 +428,7 @@ impl MetricsClient {
|
||||
duration: Duration,
|
||||
tags: &[(&str, &str)],
|
||||
) -> Result<()> {
|
||||
self.0.duration_histogram(
|
||||
self.active_inner().duration_histogram(
|
||||
name,
|
||||
duration.as_millis().min(i64::MAX as u128) as f64,
|
||||
MILLISECOND_DURATION_UNIT,
|
||||
@@ -422,7 +445,7 @@ impl MetricsClient {
|
||||
duration_ms: f64,
|
||||
tags: &[(&str, &str)],
|
||||
) -> Result<()> {
|
||||
self.0.duration_histogram(
|
||||
self.active_inner().duration_histogram(
|
||||
name,
|
||||
duration_ms,
|
||||
MILLISECOND_DURATION_UNIT,
|
||||
@@ -440,7 +463,7 @@ impl MetricsClient {
|
||||
duration: Duration,
|
||||
tags: &[(&str, &str)],
|
||||
) -> Result<()> {
|
||||
self.0.duration_histogram(
|
||||
self.active_inner().duration_histogram(
|
||||
name,
|
||||
duration.as_secs_f64(),
|
||||
SECOND_DURATION_UNIT,
|
||||
@@ -460,7 +483,8 @@ impl MetricsClient {
|
||||
|
||||
/// Collect a runtime metrics snapshot without shutting down the provider.
|
||||
pub fn snapshot(&self) -> Result<ResourceMetrics> {
|
||||
let Some(reader) = &self.0.runtime_reader else {
|
||||
let inner = self.active_inner();
|
||||
let Some(reader) = &inner.runtime_reader else {
|
||||
return Err(MetricsError::RuntimeSnapshotUnavailable);
|
||||
};
|
||||
let mut snapshot = ResourceMetrics::default();
|
||||
@@ -472,7 +496,7 @@ impl MetricsClient {
|
||||
|
||||
/// Flush metrics and stop the underlying OTEL meter provider.
|
||||
pub fn shutdown(&self) -> Result<()> {
|
||||
self.0.shutdown()
|
||||
self.inner.shutdown()
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -16,7 +16,9 @@ pub use crate::metrics::error::MetricsError;
|
||||
pub use crate::metrics::error::Result;
|
||||
pub use crate::metrics::process::record_process_start_once;
|
||||
pub use names::*;
|
||||
use std::sync::Arc;
|
||||
use std::sync::OnceLock;
|
||||
use std::sync::RwLock;
|
||||
pub use tags::ORIGINATOR_TAG;
|
||||
pub use tags::SessionMetricTagValues;
|
||||
pub use tags::bounded_originator_tag_value;
|
||||
@@ -24,8 +26,17 @@ pub use tags::bounded_originator_tag_value;
|
||||
static GLOBAL_METRICS: OnceLock<MetricsClient> = OnceLock::new();
|
||||
static GLOBAL_STATSIG_METRICS_SETTINGS: OnceLock<StatsigMetricsSettings> = OnceLock::new();
|
||||
|
||||
pub(crate) fn install_global(metrics: MetricsClient) {
|
||||
let _ = GLOBAL_METRICS.set(metrics);
|
||||
pub(crate) fn install_global(mut metrics: MetricsClient) -> MetricsClient {
|
||||
let active = GLOBAL_METRICS
|
||||
.get()
|
||||
.and_then(|current| current.active.clone())
|
||||
.unwrap_or_else(|| Arc::new(RwLock::new(Arc::clone(&metrics.inner))));
|
||||
*active
|
||||
.write()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner) = Arc::clone(&metrics.inner);
|
||||
metrics.active = Some(active);
|
||||
let _ = GLOBAL_METRICS.set(metrics.clone());
|
||||
metrics
|
||||
}
|
||||
|
||||
pub fn global() -> Option<MetricsClient> {
|
||||
|
||||
@@ -11,6 +11,7 @@ use opentelemetry::Context;
|
||||
use opentelemetry::KeyValue;
|
||||
use opentelemetry::global;
|
||||
use opentelemetry::trace::Span as _;
|
||||
use opentelemetry::trace::SpanBuilder;
|
||||
use opentelemetry::trace::TracerProvider as _;
|
||||
use opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge;
|
||||
use opentelemetry_otlp::LogExporter;
|
||||
@@ -70,6 +71,19 @@ struct ShutdownWorker {
|
||||
completed_tx: tokio::sync::oneshot::Sender<()>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct GlobalTracer {
|
||||
service_name: &'static str,
|
||||
}
|
||||
|
||||
impl opentelemetry::trace::Tracer for GlobalTracer {
|
||||
type Span = global::BoxedSpan;
|
||||
|
||||
fn build_with_context(&self, builder: SpanBuilder, parent: &Context) -> Self::Span {
|
||||
global::tracer(self.service_name).build_with_context(builder, parent)
|
||||
}
|
||||
}
|
||||
|
||||
impl OtelProvider {
|
||||
/// Flushes and shuts down configured exporters at most once.
|
||||
pub fn shutdown(&self) {
|
||||
@@ -156,7 +170,7 @@ impl OtelProvider {
|
||||
}
|
||||
crate::trace_context::validate_tracestate_entries(&settings.tracestate)?;
|
||||
|
||||
let metrics = if matches!(metric_exporter, OtelExporter::None) {
|
||||
let mut metrics = if matches!(metric_exporter, OtelExporter::None) {
|
||||
None
|
||||
} else {
|
||||
let mut config = MetricsConfig::otlp(
|
||||
@@ -196,8 +210,8 @@ impl OtelProvider {
|
||||
global::set_tracer_provider(provider);
|
||||
global::set_text_map_propagator(TraceContextPropagator::new());
|
||||
}
|
||||
if let Some(metrics) = metrics.as_ref() {
|
||||
crate::metrics::install_global(metrics.clone());
|
||||
if let Some(metrics) = metrics.as_mut() {
|
||||
*metrics = crate::metrics::install_global(metrics.clone());
|
||||
if matches!(settings.metrics_exporter, OtelExporter::Statsig) {
|
||||
crate::metrics::install_global_statsig_settings(StatsigMetricsSettings {
|
||||
environment: settings.environment.clone(),
|
||||
@@ -217,13 +231,21 @@ impl OtelProvider {
|
||||
where
|
||||
S: tracing::Subscriber + for<'span> LookupSpan<'span> + Send + Sync,
|
||||
{
|
||||
self.logger.as_ref().map(|logger| {
|
||||
OpenTelemetryTracingBridge::new(logger).with_filter(
|
||||
tracing_subscriber::filter::filter_fn(OtelProvider::log_export_filter),
|
||||
)
|
||||
self.logger_export_layer().map(|layer| {
|
||||
layer.with_filter(tracing_subscriber::filter::filter_fn(
|
||||
OtelProvider::log_export_filter,
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
/// Returns a log-export bridge that must be installed beneath the log export filter.
|
||||
pub fn logger_export_layer<S>(&self) -> Option<impl Layer<S> + Send + Sync>
|
||||
where
|
||||
S: tracing::Subscriber + for<'span> LookupSpan<'span> + Send + Sync,
|
||||
{
|
||||
self.logger.as_ref().map(OpenTelemetryTracingBridge::new)
|
||||
}
|
||||
|
||||
pub fn tracing_layer<S>(&self) -> Option<impl Layer<S> + Send + Sync>
|
||||
where
|
||||
S: tracing::Subscriber + for<'span> LookupSpan<'span> + Send + Sync,
|
||||
@@ -237,6 +259,18 @@ impl OtelProvider {
|
||||
})
|
||||
}
|
||||
|
||||
/// Returns a permanent trace layer that follows the process-global tracer provider.
|
||||
pub fn reloadable_tracing_layer<S>(service_name: &'static str) -> impl Layer<S> + Send + Sync
|
||||
where
|
||||
S: tracing::Subscriber + for<'span> LookupSpan<'span> + Send + Sync,
|
||||
{
|
||||
tracing_opentelemetry::layer()
|
||||
.with_tracer(GlobalTracer { service_name })
|
||||
.with_filter(tracing_subscriber::filter::filter_fn(
|
||||
Self::trace_export_filter,
|
||||
))
|
||||
}
|
||||
|
||||
pub fn codex_export_filter(meta: &tracing::Metadata<'_>) -> bool {
|
||||
Self::log_export_filter(meta)
|
||||
}
|
||||
@@ -586,6 +620,43 @@ mod tests {
|
||||
assert!(!is_trace_safe_target("codex_otel.network_proxy"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cached_global_metrics_follow_reinstalled_provider() -> Result<(), Box<dyn Error>> {
|
||||
let initial =
|
||||
crate::metrics::install_global(MetricsClient::new(MetricsConfig::in_memory(
|
||||
"test",
|
||||
"codex-test",
|
||||
env!("CARGO_PKG_VERSION"),
|
||||
InMemoryMetricExporter::default(),
|
||||
))?);
|
||||
let cached = crate::metrics::global().expect("initial global metrics client");
|
||||
|
||||
let exporter = InMemoryMetricExporter::default();
|
||||
let replacement =
|
||||
crate::metrics::install_global(MetricsClient::new(MetricsConfig::in_memory(
|
||||
"test",
|
||||
"codex-test",
|
||||
env!("CARGO_PKG_VERSION"),
|
||||
exporter.clone(),
|
||||
))?);
|
||||
cached.counter("codex.after_transition", /*inc*/ 1, &[])?;
|
||||
initial.shutdown()?;
|
||||
replacement.shutdown()?;
|
||||
|
||||
let exported_metrics = exporter.get_finished_metrics()?;
|
||||
let mut names: Vec<_> = exported_metrics
|
||||
.iter()
|
||||
.flat_map(opentelemetry_sdk::metrics::data::ResourceMetrics::scope_metrics)
|
||||
.flat_map(opentelemetry_sdk::metrics::data::ScopeMetrics::metrics)
|
||||
.map(opentelemetry_sdk::metrics::data::Metric::name)
|
||||
.collect();
|
||||
names.sort_unstable();
|
||||
names.dedup();
|
||||
assert_eq!(names, vec!["codex.after_transition"]);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn statsig_runtime_only_metrics_are_not_exported() -> Result<(), Box<dyn Error>> {
|
||||
let exporter = InMemoryMetricExporter::default();
|
||||
|
||||
Reference in New Issue
Block a user