M2 audit: excise the Computer Hub stack — Kigi's last remote-cloud surface
Removed root-and-branch for the zero-egress guarantee (the hub was xAI's remote-workspace/cloud-sandbox service): - Crates deleted: kigi-computer-hub-core, kigi-computer-hub-sdk, kigi-computer-hub-mcp-adapter, kigi-workspace-client (hub-proxied workspace RPC client), and kigi-tracing (its sole network path was the OTLP gRPC exporter; zero consumers remained). kigi-tracing-macros (purely local) stays. - kigi-workspace: every hub surface deleted — hub server/channel/auth, HITL-over-hub permissions, donation/metrics pumps, file upload RPCs, hub tool-snapshot merge (resolve pipeline is MCP-only now), WorkspaceOps::Proxy. Local worktrees, sessions, leader IPC, MCP, and the ACP permission prompt path are untouched; LocalRegistry re-homed into kigi-tool-runtime on the existing ToolDyn types so in-process tool dispatch is unchanged. - kigi-shell: leader workspace-exposure control surface (incl. the wss://computer-hub... URL), [hub] config, ObservabilityBridge, hub WebSocket proxy, dead OTLP config knobs. ClientMode::Headless (never constructed) removed. - kigi-tui/bin: hidden `kigi workspace` command removed (`kigi worktree` stays). - Renames: --xai-api-base-url → --api-base-url / KIGI_API_BASE_URL / [endpoints] api_base_url (serde alias keeps old configs working; the flag feeds BYOK/custom-endpoint routing, not main inference); grok_version → kigi_version in inspect/models-cache/trace metadata (old caches self-heal via version-mismatch refetch). - Dependency tree: dropped fastrace*, opentelemetry-otlp/http/proto, tokio-tungstenite from the workspace; fixed the 4 real useless_format violations the fastrace lint allowance was masking and removed the allowance. - marketplaceAllowlist kept: it gates the LOCAL plugin-marketplace feature, not an xAI service. Known §9 leftover (deliberate, for the M3 sweep): the BYOK default base URL string. Gates: workspace check/clippy 0/0, fmt, deny ok; suites green (workspace 1042, shell 4918, tui 6634, tools 2608, tool-runtime 47, mcp 154).
This commit is contained in:
@@ -3,7 +3,7 @@ license = "Apache-2.0"
|
||||
name = "kigi-tool-protocol"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
description = "Wire-protocol types for the xAI Computer Hub"
|
||||
description = "Tool wire-protocol types"
|
||||
|
||||
[dependencies]
|
||||
kigi-tool-types = { workspace = true }
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//! xAI Computer Hub — wire-protocol types.
|
||||
//! Tool wire-protocol types.
|
||||
//!
|
||||
//! Identifier newtypes, registration payloads, capabilities, hook events,
|
||||
//! handshake messages, the JSON-RPC 2.0 envelope and method catalog, the
|
||||
|
||||
@@ -3,17 +3,18 @@ license = "Apache-2.0"
|
||||
name = "kigi-tool-runtime"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
description = "Unified Tool trait, dispatch trait, error taxonomy, notifications, and search index for the xAI Computer Hub"
|
||||
description = "Unified Tool trait, dispatch trait, error taxonomy, notifications, local tool registry, and search index"
|
||||
|
||||
[dependencies]
|
||||
anyhow = { workspace = true }
|
||||
async-trait = { workspace = true }
|
||||
futures = { workspace = true }
|
||||
indexmap = { workspace = true }
|
||||
parking_lot = { workspace = true }
|
||||
schemars = { workspace = true }
|
||||
serde = { workspace = true, features = ["derive"] }
|
||||
serde_json = { workspace = true }
|
||||
tokio-util = { workspace = true }
|
||||
kigi-tools-api = { workspace = true }
|
||||
kigi-tool-protocol = { workspace = true }
|
||||
kigi-tool-types = { workspace = true }
|
||||
|
||||
|
||||
@@ -152,185 +152,3 @@ pub struct WorkspaceViewerContext {
|
||||
#[serde(default)]
|
||||
pub stream_tool_progress: bool,
|
||||
}
|
||||
|
||||
/// Wire shape of the Computer Hub `session.bind` metadata — one definition
|
||||
/// shared by the emitter (serializes) and the workspace consumer
|
||||
/// (deserializes), so the two can't drift on field names/types.
|
||||
///
|
||||
/// Excludes anything not meant for the workspace (cached tool definitions,
|
||||
/// and terminal-provisioning inputs like image/fuse/isolation) so they can
|
||||
/// never reach the wire. Every field tolerates a missing/malformed value
|
||||
/// (drops to default) to keep valid siblings and mixed-version compatibility.
|
||||
#[derive(Clone, Debug, Default, serde::Serialize, serde::Deserialize)]
|
||||
pub struct WorkspaceBindMetadata {
|
||||
#[serde(
|
||||
default,
|
||||
deserialize_with = "ok_or_default",
|
||||
skip_serializing_if = "Option::is_none"
|
||||
)]
|
||||
pub preset: Option<String>,
|
||||
/// Raw string; the workspace maps it to its own capability enum.
|
||||
#[serde(
|
||||
default,
|
||||
deserialize_with = "ok_or_default",
|
||||
skip_serializing_if = "Option::is_none"
|
||||
)]
|
||||
pub capability_mode: Option<String>,
|
||||
/// Explicit toolset in the grok-tools gRPC wire shape. Empty = unset.
|
||||
#[serde(
|
||||
default,
|
||||
deserialize_with = "ok_or_default",
|
||||
skip_serializing_if = "Vec::is_empty"
|
||||
)]
|
||||
pub tools: Vec<kigi_tools_api::ToolConfigEntry>,
|
||||
#[serde(
|
||||
default,
|
||||
deserialize_with = "ok_or_default",
|
||||
skip_serializing_if = "Option::is_none"
|
||||
)]
|
||||
pub viewer_ctx: Option<WorkspaceViewerContext>,
|
||||
/// Initial auto-approve (YOLO) state for the bound session. Omitted when
|
||||
/// unset (legacy emitters / wire compat with older workspace servers);
|
||||
/// consumers fail closed on `None`.
|
||||
#[serde(
|
||||
default,
|
||||
deserialize_with = "ok_or_default",
|
||||
skip_serializing_if = "Option::is_none"
|
||||
)]
|
||||
pub yolo_mode: Option<bool>,
|
||||
/// Optional/additive: omitted by emitters that don't yet write it.
|
||||
#[serde(
|
||||
default,
|
||||
deserialize_with = "ok_or_default",
|
||||
skip_serializing_if = "Option::is_none"
|
||||
)]
|
||||
pub manifest_version: Option<String>,
|
||||
#[serde(
|
||||
default,
|
||||
deserialize_with = "ok_or_default",
|
||||
skip_serializing_if = "Option::is_none"
|
||||
)]
|
||||
pub manifest_hash: Option<String>,
|
||||
/// Opt-in: forward SystemNotifications produced in this session to the gateway.
|
||||
#[serde(
|
||||
default,
|
||||
deserialize_with = "ok_or_default",
|
||||
skip_serializing_if = "Option::is_none"
|
||||
)]
|
||||
pub system_notifications: Option<bool>,
|
||||
#[serde(
|
||||
default,
|
||||
deserialize_with = "ok_or_default",
|
||||
skip_serializing_if = "std::ops::Not::not"
|
||||
)]
|
||||
pub rpc_only: bool,
|
||||
}
|
||||
|
||||
/// Deserialize a field, falling back to its default on a malformed value
|
||||
/// instead of failing the whole struct.
|
||||
fn ok_or_default<'de, D, T>(deserializer: D) -> Result<T, D::Error>
|
||||
where
|
||||
D: serde::Deserializer<'de>,
|
||||
T: serde::de::DeserializeOwned + Default,
|
||||
{
|
||||
let value = <serde_json::Value as serde::Deserialize>::deserialize(deserializer)?;
|
||||
Ok(serde_json::from_value(value).unwrap_or_default())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod bind_metadata_tests {
|
||||
use super::WorkspaceBindMetadata;
|
||||
|
||||
#[test]
|
||||
fn serialize_omits_empty_fields() {
|
||||
let md = WorkspaceBindMetadata::default();
|
||||
assert_eq!(serde_json::to_value(&md).unwrap(), serde_json::json!({}));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn round_trips_populated() {
|
||||
let md = WorkspaceBindMetadata {
|
||||
preset: Some("explore".to_owned()),
|
||||
capability_mode: Some("read_only".to_owned()),
|
||||
tools: vec![kigi_tools_api::ToolConfigEntry {
|
||||
id: "GrokBuild:grep".to_owned(),
|
||||
..Default::default()
|
||||
}],
|
||||
viewer_ctx: Some(super::WorkspaceViewerContext {
|
||||
stream_tool_progress: true,
|
||||
}),
|
||||
yolo_mode: Some(true),
|
||||
manifest_version: Some("v1".to_owned()),
|
||||
manifest_hash: Some("abc123".to_owned()),
|
||||
system_notifications: Some(true),
|
||||
rpc_only: true,
|
||||
};
|
||||
let value = serde_json::to_value(&md).unwrap();
|
||||
let back: WorkspaceBindMetadata = serde_json::from_value(value).unwrap();
|
||||
assert_eq!(back.preset.as_deref(), Some("explore"));
|
||||
assert_eq!(back.capability_mode.as_deref(), Some("read_only"));
|
||||
assert_eq!(back.tools.len(), 1);
|
||||
assert!(back.viewer_ctx.unwrap().stream_tool_progress);
|
||||
assert_eq!(back.yolo_mode, Some(true));
|
||||
assert_eq!(back.manifest_version.as_deref(), Some("v1"));
|
||||
assert_eq!(back.manifest_hash.as_deref(), Some("abc123"));
|
||||
assert_eq!(back.system_notifications, Some(true));
|
||||
assert!(back.rpc_only);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rpc_only_omitted_when_false_wire_compatible() {
|
||||
let md = WorkspaceBindMetadata::default();
|
||||
let value = serde_json::to_value(&md).unwrap();
|
||||
assert!(value.get("rpc_only").is_none());
|
||||
|
||||
let md: WorkspaceBindMetadata =
|
||||
serde_json::from_value(serde_json::json!({"preset": "explore"})).unwrap();
|
||||
assert!(!md.rpc_only);
|
||||
|
||||
let md: WorkspaceBindMetadata =
|
||||
serde_json::from_value(serde_json::json!({"rpc_only": true})).unwrap();
|
||||
assert!(md.rpc_only);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn system_notifications_is_wire_compatible() {
|
||||
let md = WorkspaceBindMetadata::default();
|
||||
let value = serde_json::to_value(&md).unwrap();
|
||||
assert!(value.get("system_notifications").is_none());
|
||||
|
||||
let md = WorkspaceBindMetadata {
|
||||
system_notifications: Some(true),
|
||||
..Default::default()
|
||||
};
|
||||
let value = serde_json::to_value(&md).unwrap();
|
||||
let back: WorkspaceBindMetadata = serde_json::from_value(value).unwrap();
|
||||
assert_eq!(back.system_notifications, Some(true));
|
||||
|
||||
let md: WorkspaceBindMetadata =
|
||||
serde_json::from_value(serde_json::json!({"preset": "explore"})).unwrap();
|
||||
assert!(md.system_notifications.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_field_falls_back_to_default_keeping_siblings() {
|
||||
// `tools` is the wrong type and `capability_mode` is fine: the bad
|
||||
// field drops to default, the good sibling survives.
|
||||
let value = serde_json::json!({
|
||||
"preset": "explore",
|
||||
"capability_mode": "read_only",
|
||||
"tools": "not-a-list",
|
||||
});
|
||||
let md: WorkspaceBindMetadata = serde_json::from_value(value).unwrap();
|
||||
assert_eq!(md.preset.as_deref(), Some("explore"));
|
||||
assert_eq!(md.capability_mode.as_deref(), Some("read_only"));
|
||||
assert!(md.tools.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn legacy_payload_without_viewer_ctx_parses() {
|
||||
let md: WorkspaceBindMetadata =
|
||||
serde_json::from_value(serde_json::json!({"preset": "explore"})).unwrap();
|
||||
assert!(md.viewer_ctx.is_none());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,16 +1,17 @@
|
||||
//! xAI Computer Hub — unified runtime contract.
|
||||
//! Unified tool runtime contract.
|
||||
//!
|
||||
//! Single home for the `Tool` trait, `ToolDispatch`, `ToolError`,
|
||||
//! `ToolNotification`, `ToolSearchIndex`, `ToolCallContext`, `ToolStream`,
|
||||
//! and the helper constructors that build well-formed streams. Adapters
|
||||
//! for individual tool sources re-export from here so every tool author
|
||||
//! sees the same surface.
|
||||
//! the in-process `LocalRegistry`, and the helper constructors that build
|
||||
//! well-formed streams. Adapters for individual tool sources re-export
|
||||
//! from here so every tool author sees the same surface.
|
||||
|
||||
#![forbid(unsafe_code)]
|
||||
|
||||
pub mod context;
|
||||
pub mod dispatch;
|
||||
pub mod error;
|
||||
pub mod local_registry;
|
||||
pub mod notification;
|
||||
pub mod render;
|
||||
pub mod search;
|
||||
@@ -19,10 +20,11 @@ pub mod tool;
|
||||
|
||||
pub use context::{
|
||||
BehaviorVersion, Cancellation, Cwd, ListToolsContext, SessionContext, ToolCallContext,
|
||||
TraceContext, TypedExtensions, WorkspaceBindMetadata, WorkspaceViewerContext,
|
||||
TraceContext, TypedExtensions, WorkspaceViewerContext,
|
||||
};
|
||||
pub use dispatch::ToolDispatch;
|
||||
pub use error::{ToolError, ToolErrorKind};
|
||||
pub use local_registry::LocalRegistry;
|
||||
pub use notification::{
|
||||
BashExecutionBackgrounded, BashExecutionComplete, BashExecutionFailed, BashExecutionTimeout,
|
||||
BashNotificationBase, BashOutputChunk, FileRead, FileWritten, LspServerCrashed,
|
||||
|
||||
@@ -0,0 +1,117 @@
|
||||
//! In-process tool registry for local dispatch.
|
||||
//!
|
||||
//! [`LocalRegistry`] maps [`ToolId`]s to type-erased
|
||||
//! [`ToolDyn`](crate::tool::ToolDyn) handles.
|
||||
//! Toolset finalization registers every config-enabled tool here and
|
||||
//! dispatch resolves handles via [`LocalRegistry::find`], so a call
|
||||
//! executes in-process without any wire round-trip.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use indexmap::IndexMap;
|
||||
use parking_lot::RwLock;
|
||||
|
||||
use crate::context::ListToolsContext;
|
||||
use crate::tool::{ArcTool, Tool};
|
||||
use kigi_tool_protocol::ToolId;
|
||||
use kigi_tool_types::ToolDescription;
|
||||
|
||||
/// In-process registry of tool handles.
|
||||
///
|
||||
/// Mutations are concurrency-safe (`RwLock` on the entry map), so
|
||||
/// callers MAY hot-add or hot-remove tools while dispatch is in use.
|
||||
///
|
||||
/// Entries use `RwLock<IndexMap>` to preserve insertion order so that
|
||||
/// [`list_tools`](Self::list_tools) returns descriptions in the same
|
||||
/// order tools were registered (matching the config-defined order).
|
||||
#[derive(Clone, Default)]
|
||||
pub struct LocalRegistry {
|
||||
entries: Arc<RwLock<IndexMap<ToolId, ArcTool>>>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for LocalRegistry {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("LocalRegistry")
|
||||
.field("entries", &self.entries.read().len())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl LocalRegistry {
|
||||
/// Construct an empty registry.
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Register a typed [`Tool`] implementation by value. Subsequent
|
||||
/// registrations of the same id replace the previous handle and
|
||||
/// return the displaced handle for inspection / drop ordering.
|
||||
pub fn register<T>(&self, tool: T) -> Option<ArcTool>
|
||||
where
|
||||
T: Tool + 'static,
|
||||
{
|
||||
self.register_arc(Arc::new(tool))
|
||||
}
|
||||
|
||||
/// Register a typed [`Tool`] implementation already wrapped in `Arc`.
|
||||
pub fn register_arc<T>(&self, tool: Arc<T>) -> Option<ArcTool>
|
||||
where
|
||||
T: Tool + 'static,
|
||||
{
|
||||
let id = tool.id();
|
||||
self.entries.write().insert(id, tool as ArcTool)
|
||||
}
|
||||
|
||||
/// Register a type-erased [`ToolDyn`](crate::tool::ToolDyn) directly.
|
||||
///
|
||||
/// Use this for inherently dynamic tools (e.g. MCP tools retrieved
|
||||
/// from a registry as `Arc<dyn ToolDyn>`) where the concrete type
|
||||
/// is not available. For native tools with a concrete type, prefer
|
||||
/// [`register`](Self::register).
|
||||
pub fn register_dyn(&self, tool: ArcTool) -> Option<ArcTool> {
|
||||
let id = tool.id();
|
||||
self.entries.write().insert(id, tool)
|
||||
}
|
||||
|
||||
/// Resolve `tool_id` to its in-process handle, if registered.
|
||||
/// Returns a clone of the handle so the caller can dispatch without
|
||||
/// holding the lock across an await point.
|
||||
pub fn find(&self, tool_id: &ToolId) -> Option<ArcTool> {
|
||||
self.entries.read().get(tool_id).cloned()
|
||||
}
|
||||
|
||||
/// Drop the handle bound to `tool_id`. Returns `true` iff a
|
||||
/// matching entry was removed.
|
||||
pub fn unregister(&self, tool_id: &ToolId) -> bool {
|
||||
self.entries.write().shift_remove(tool_id).is_some()
|
||||
}
|
||||
|
||||
/// Number of tools currently registered.
|
||||
pub fn len(&self) -> usize {
|
||||
self.entries.read().len()
|
||||
}
|
||||
|
||||
/// `true` iff no tools are registered.
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.entries.read().is_empty()
|
||||
}
|
||||
|
||||
/// `true` iff `tool_id` is currently registered.
|
||||
pub fn contains(&self, tool_id: &ToolId) -> bool {
|
||||
self.entries.read().contains_key(tool_id)
|
||||
}
|
||||
|
||||
/// Descriptions of registered tools filtered by `should_list`.
|
||||
///
|
||||
/// Returns descriptions in **insertion order** — the order tools
|
||||
/// were registered — so the caller sees the same ordering as the
|
||||
/// config-defined tool list.
|
||||
pub fn list_tools(&self, ctx: &ListToolsContext) -> Vec<ToolDescription> {
|
||||
self.entries
|
||||
.read()
|
||||
.values()
|
||||
.filter(|handle| handle.should_list(ctx))
|
||||
.map(|handle| handle.description(ctx))
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
@@ -3,7 +3,7 @@ license = "Apache-2.0"
|
||||
name = "kigi-tool-types"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
description = "Canonical tool-description types for the xAI platform"
|
||||
description = "Canonical tool-description types"
|
||||
|
||||
[features]
|
||||
# Enables `BuiltinSubagent::render_prompt` (MiniJinja rendering of the
|
||||
|
||||
@@ -1,37 +0,0 @@
|
||||
[package]
|
||||
license = "Apache-2.0"
|
||||
edition.workspace = true
|
||||
name = "kigi-tracing"
|
||||
version.workspace = true
|
||||
|
||||
[dependencies]
|
||||
async-trait = { workspace = true }
|
||||
http = { workspace = true }
|
||||
log = { workspace = true }
|
||||
fastrace = { workspace = true }
|
||||
fastrace-tonic = { workspace = true }
|
||||
fastrace-opentelemetry = { workspace = true }
|
||||
fastrace-reqwest = { workspace = true }
|
||||
opentelemetry = { workspace = true }
|
||||
opentelemetry-http = { workspace = true }
|
||||
opentelemetry-otlp = { workspace = true }
|
||||
opentelemetry_sdk = { workspace = true, features = ["testing"] }
|
||||
reqwest = { workspace = true }
|
||||
reqwest-middleware = { workspace = true }
|
||||
tokio = { workspace = true }
|
||||
tonic = { workspace = true }
|
||||
tower = { workspace = true, features = ["full"] }
|
||||
tower-http = { workspace = true, features = ["trace"] }
|
||||
tracing = { workspace = true }
|
||||
tracing-opentelemetry = { workspace = true }
|
||||
uuid = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
bytes = { workspace = true }
|
||||
http-body-util = { workspace = true }
|
||||
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
|
||||
tracing-subscriber = { workspace = true }
|
||||
wiremock = { workspace = true }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -1,45 +0,0 @@
|
||||
use tracing::subscriber::NoSubscriber;
|
||||
|
||||
/// Returns `true` when a `tracing` dispatcher (subscriber) is active in the
|
||||
/// current context — either the thread-scoped default
|
||||
/// (`tracing::subscriber::with_default` / `set_default`) or the global one.
|
||||
///
|
||||
/// When this returns `false`, spans have no consumer. Worse than useless: when
|
||||
/// `tracing` is compiled with its `log` compatibility feature, every span
|
||||
/// creation and every later `Span::record(...)` is downgraded to a `log`
|
||||
/// record at the span's level. In processes that only configure a `log`
|
||||
/// logger — e.g. integration tests or fastrace-only binaries — that prints
|
||||
/// noise like:
|
||||
///
|
||||
/// ```text
|
||||
/// I grpc; otel.name="POST_/pkg.Service/Method" ...
|
||||
/// I grpc; status_code=200 OK
|
||||
/// I grpc; trace_id=00000000000000000000000000000000
|
||||
/// ```
|
||||
///
|
||||
/// Request-span factories (gRPC/HTTP server and client middleware) call this
|
||||
/// and return `Span::none()` when no dispatcher is active, so the span is
|
||||
/// neither built nor downgraded to log spam.
|
||||
pub fn dispatcher_active() -> bool {
|
||||
tracing::dispatcher::get_default(|dispatch| !dispatch.is::<NoSubscriber>())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// NOTE: relies on no test in this binary installing a *global* subscriber
|
||||
// (`OtelTestEnv` and friends use thread-scoped `set_default` guards).
|
||||
#[test]
|
||||
fn without_dispatcher_inactive() {
|
||||
assert!(!dispatcher_active());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scoped_dispatcher_active() {
|
||||
tracing::subscriber::with_default(tracing_subscriber::registry(), || {
|
||||
assert!(dispatcher_active());
|
||||
});
|
||||
assert!(!dispatcher_active());
|
||||
}
|
||||
}
|
||||
@@ -1,83 +0,0 @@
|
||||
use fastrace::prelude::*;
|
||||
use fastrace_opentelemetry::OpenTelemetryReporter;
|
||||
use opentelemetry::InstrumentationScope;
|
||||
use opentelemetry::KeyValue;
|
||||
use opentelemetry_otlp::WithExportConfig;
|
||||
use opentelemetry_otlp::{ExporterBuildError, SpanExporter};
|
||||
use opentelemetry_sdk::Resource;
|
||||
use std::borrow::Cow;
|
||||
use std::iter;
|
||||
|
||||
// Fastrace initialization
|
||||
pub fn init_fastrace(
|
||||
endpoint: String,
|
||||
name: String,
|
||||
resource_attributes: impl IntoIterator<Item = (String, String)>,
|
||||
) -> Result<(), ExporterBuildError> {
|
||||
let exporter = SpanExporter::builder()
|
||||
.with_tonic()
|
||||
.with_endpoint(endpoint)
|
||||
.with_protocol(opentelemetry_otlp::Protocol::Grpc)
|
||||
.with_timeout(opentelemetry_otlp::OTEL_EXPORTER_OTLP_TIMEOUT_DEFAULT)
|
||||
.build()?;
|
||||
let attributes = resource_attributes
|
||||
.into_iter()
|
||||
.chain(iter::once(("service.name".into(), name.clone())))
|
||||
.map(|(k, v)| KeyValue::new(k, v));
|
||||
let reporter = OpenTelemetryReporter::new(
|
||||
exporter,
|
||||
Cow::Owned(Resource::builder().with_attributes(attributes).build()),
|
||||
InstrumentationScope::builder(name)
|
||||
.with_version(env!("CARGO_PKG_VERSION"))
|
||||
.build(),
|
||||
);
|
||||
fastrace::set_reporter(reporter, fastrace::collector::Config::default());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn current_trace_id() -> Option<String> {
|
||||
SpanContext::current_local_parent().map(|current| current.encode_w3c_traceparent())
|
||||
}
|
||||
|
||||
pub fn local_or_random_span_ctx() -> SpanContext {
|
||||
SpanContext::current_local_parent().unwrap_or_else(SpanContext::random)
|
||||
}
|
||||
|
||||
pub fn enter_span_with_traceparent(name: impl Into<Cow<'static, str>>, traceparent: &str) -> Span {
|
||||
if let Some(span_ctx) = SpanContext::decode_w3c_traceparent(traceparent) {
|
||||
Span::root(name, span_ctx)
|
||||
} else {
|
||||
Span::enter_with_local_parent(name)
|
||||
}
|
||||
}
|
||||
|
||||
// Tonic channel (TODO: Move into grpc_client when deprecated tracing)
|
||||
#[allow(dead_code)]
|
||||
pub type FastraceChannel = fastrace_tonic::FastraceClientService<tonic::transport::Channel>;
|
||||
|
||||
pub fn fastrace_channel(
|
||||
channel: tonic::transport::Channel,
|
||||
) -> fastrace_tonic::FastraceClientService<tonic::transport::Channel> {
|
||||
tower::ServiceBuilder::new()
|
||||
.layer(fastrace_tonic::FastraceClientLayer)
|
||||
.service(channel)
|
||||
}
|
||||
|
||||
// Request middleware (TODO: Move into http_client when deprecated tracing)
|
||||
#[derive(Clone)]
|
||||
#[allow(dead_code)]
|
||||
pub struct TraceparentMiddleware;
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl reqwest_middleware::Middleware for TraceparentMiddleware {
|
||||
async fn handle(
|
||||
&self,
|
||||
mut req: reqwest::Request,
|
||||
extensions: &mut http::Extensions,
|
||||
next: reqwest_middleware::Next<'_>,
|
||||
) -> reqwest_middleware::Result<reqwest::Response> {
|
||||
req.headers_mut()
|
||||
.extend(fastrace_reqwest::traceparent_headers());
|
||||
next.run(req, extensions).await
|
||||
}
|
||||
}
|
||||
@@ -1,388 +0,0 @@
|
||||
use http::{HeaderMap, Request};
|
||||
use opentelemetry::{global, propagation::Extractor, propagation::Injector};
|
||||
use std::task::{Context, Poll};
|
||||
use tonic::transport::Channel;
|
||||
use tonic::{
|
||||
Status,
|
||||
metadata::{MetadataKey, MetadataMap, MetadataValue},
|
||||
};
|
||||
use tower::{Layer, Service, ServiceBuilder};
|
||||
use tower_http::classify::{GrpcErrorsAsFailures, SharedClassifier};
|
||||
use tower_http::trace::{MakeSpan, Trace, TraceLayer};
|
||||
use tracing::{Span, warn};
|
||||
use tracing_opentelemetry::OpenTelemetrySpanExt;
|
||||
|
||||
pub type TracedChannel = Trace<
|
||||
InjectTraceContextService<Channel>,
|
||||
SharedClassifier<GrpcErrorsAsFailures>,
|
||||
MakeClientSpan,
|
||||
>;
|
||||
|
||||
/// Wraps the input channel with a tracing layer. This function can be used to create a traced gRPC
|
||||
/// client as follows:
|
||||
///
|
||||
/// ```rust
|
||||
/// use tonic::transport::Endpoint;
|
||||
/// use std::str::FromStr;
|
||||
/// use kigi_tracing::traced_channel;
|
||||
///
|
||||
/// let channel = Endpoint::from_str("http://foo").unwrap();
|
||||
/// //let client = SomeClient::new(traced_channel(channel));
|
||||
///```
|
||||
pub fn traced_channel(channel: Channel) -> TracedChannel {
|
||||
ServiceBuilder::new()
|
||||
.layer(TraceLayer::new_for_grpc().make_span_with(MakeClientSpan))
|
||||
.layer(InjectTraceContextLayer)
|
||||
.service(channel)
|
||||
}
|
||||
|
||||
/// Implements the [`MakeSpan`] trait, to trace outgoing gRPC requests.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct MakeClientSpan;
|
||||
|
||||
impl<B> MakeSpan<B> for MakeClientSpan {
|
||||
fn make_span(&mut self, request: &Request<B>) -> Span {
|
||||
// No active dispatcher → the span has no consumer and would only be
|
||||
// downgraded to `log` spam. See `crate::dispatcher_active`.
|
||||
if !crate::dispatcher_active() {
|
||||
return Span::none();
|
||||
}
|
||||
tracing::info_span!(
|
||||
"grpc_request",
|
||||
otel.kind = "client",
|
||||
method = %request.method(),
|
||||
uri = %request.uri(),
|
||||
version = ?request.version(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default)]
|
||||
pub struct InjectTraceContextLayer;
|
||||
|
||||
impl<S> Layer<S> for InjectTraceContextLayer {
|
||||
type Service = InjectTraceContextService<S>;
|
||||
|
||||
fn layer(&self, inner: S) -> Self::Service {
|
||||
InjectTraceContextService { inner }
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct InjectTraceContextService<S> {
|
||||
inner: S,
|
||||
}
|
||||
|
||||
impl<S, B> Service<Request<B>> for InjectTraceContextService<S>
|
||||
where
|
||||
S: Service<Request<B>>,
|
||||
{
|
||||
type Response = S::Response;
|
||||
type Error = S::Error;
|
||||
type Future = S::Future;
|
||||
|
||||
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||
self.inner.poll_ready(cx)
|
||||
}
|
||||
|
||||
fn call(&mut self, mut req: Request<B>) -> Self::Future {
|
||||
crate::http_client::attach_trace_to_http_request(req.headers_mut());
|
||||
self.inner.call(req)
|
||||
}
|
||||
}
|
||||
|
||||
/// Inject W3C `traceparent` / `tracestate` from the active span into gRPC
|
||||
/// metadata. Mutates in place so callers never lose the request body.
|
||||
pub fn attach_trace_to_grpc_request_mut(metadata: &mut MetadataMap) {
|
||||
global::get_text_map_propagator(|propagator| {
|
||||
let context = Span::current().context();
|
||||
propagator.inject_context(&context, &mut MetadataInjector(metadata));
|
||||
});
|
||||
}
|
||||
|
||||
/// Trace context propagation: send the trace context by injecting it into the metadata of the given
|
||||
/// request.
|
||||
pub fn attach_trace_to_grpc_request<T>(
|
||||
mut request: tonic::Request<T>,
|
||||
) -> Result<tonic::Request<T>, Status> {
|
||||
attach_trace_to_grpc_request_mut(request.metadata_mut());
|
||||
Ok(request)
|
||||
}
|
||||
|
||||
// Need a custom Injector to inject OTel headers
|
||||
pub struct MetadataInjector<'a>(&'a mut MetadataMap);
|
||||
|
||||
impl Injector for MetadataInjector<'_> {
|
||||
fn set(&mut self, key: &str, value: String) {
|
||||
match MetadataKey::from_bytes(key.as_bytes()) {
|
||||
Ok(key) => match MetadataValue::try_from(&value) {
|
||||
Ok(value) => {
|
||||
self.0.insert(key, value);
|
||||
}
|
||||
|
||||
Err(error) => warn!(value, error = format!("{error:#}"), "parse metadata value"),
|
||||
},
|
||||
|
||||
Err(error) => warn!(key, error = format!("{error:#}"), "parse metadata key"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct HeaderExtractor<'a>(pub &'a HeaderMap);
|
||||
|
||||
impl Extractor for HeaderExtractor<'_> {
|
||||
fn get(&self, key: &str) -> Option<&str> {
|
||||
self.0.get(key).and_then(|v| {
|
||||
let s = v.to_str();
|
||||
if let Err(ref error) = s {
|
||||
warn!(%error, ?v, "cannot convert header value to ASCII")
|
||||
};
|
||||
s.ok()
|
||||
})
|
||||
}
|
||||
|
||||
fn keys(&self) -> Vec<&str> {
|
||||
self.0.keys().map(|k| k.as_str()).collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::testing::{OtelTestEnv, otel_span_id_hex, otel_trace_id_hex, parse_traceparent};
|
||||
use http_body_util::Empty;
|
||||
use std::convert::Infallible;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
use tower_http::classify::GrpcFailureClass;
|
||||
use tracing::Instrument;
|
||||
|
||||
type EmptyBody = Empty<bytes::Bytes>;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct CaptureService {
|
||||
seen: Arc<Mutex<Option<HeaderMap>>>,
|
||||
response_grpc_status: Option<&'static str>,
|
||||
}
|
||||
|
||||
impl CaptureService {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
seen: Arc::new(Mutex::new(None)),
|
||||
response_grpc_status: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn with_grpc_status(status: &'static str) -> Self {
|
||||
Self {
|
||||
seen: Arc::new(Mutex::new(None)),
|
||||
response_grpc_status: Some(status),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<B> Service<Request<B>> for CaptureService {
|
||||
type Response = http::Response<EmptyBody>;
|
||||
type Error = Infallible;
|
||||
type Future = std::future::Ready<Result<Self::Response, Self::Error>>;
|
||||
|
||||
fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
|
||||
fn call(&mut self, req: Request<B>) -> Self::Future {
|
||||
*self.seen.lock().unwrap() = Some(req.headers().clone());
|
||||
let mut builder = http::Response::builder().status(200);
|
||||
if let Some(status) = self.response_grpc_status {
|
||||
builder = builder.header("grpc-status", status);
|
||||
}
|
||||
std::future::ready(Ok(builder.body(Empty::new()).unwrap()))
|
||||
}
|
||||
}
|
||||
|
||||
fn post_req() -> Request<EmptyBody> {
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("http://svc/package.Service/Method")
|
||||
.body(Empty::new())
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
// With no dispatcher active, the client span must not be created —
|
||||
// `tracing`'s `log` compat would downgrade it into `grpc_request; ...`
|
||||
// log spam in processes that only configure a `log` logger.
|
||||
// See `crate::dispatcher_active`.
|
||||
#[test]
|
||||
fn make_client_span_without_dispatcher_is_none() {
|
||||
assert!(MakeClientSpan.make_span(&post_req()).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn make_client_span_with_scoped_dispatcher_is_enabled() {
|
||||
let _env = OtelTestEnv::install();
|
||||
assert!(!MakeClientSpan.make_span(&post_req()).is_disabled());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn inject_under_trace_layer_uses_client_span_not_parent() {
|
||||
let _env = OtelTestEnv::install();
|
||||
|
||||
let capture = CaptureService::new();
|
||||
let seen = Arc::clone(&capture.seen);
|
||||
let mut svc = ServiceBuilder::new()
|
||||
.layer(TraceLayer::new_for_grpc().make_span_with(MakeClientSpan))
|
||||
.layer(InjectTraceContextLayer)
|
||||
.service(capture);
|
||||
|
||||
let parent = tracing::info_span!("parent_handler");
|
||||
let parent_span_id = otel_span_id_hex(&parent);
|
||||
let parent_trace_id = otel_trace_id_hex(&parent);
|
||||
assert_ne!(parent_span_id, "0000000000000000");
|
||||
|
||||
async {
|
||||
let fut = Service::call(&mut svc, post_req());
|
||||
fut.await.unwrap();
|
||||
}
|
||||
.instrument(parent)
|
||||
.await;
|
||||
|
||||
let headers = seen.lock().unwrap().clone().expect("headers");
|
||||
let tp = headers
|
||||
.get("traceparent")
|
||||
.expect("traceparent")
|
||||
.to_str()
|
||||
.unwrap();
|
||||
let (_ver, injected_trace_id, injected_span_id) = parse_traceparent(tp);
|
||||
|
||||
assert_eq!(injected_trace_id, parent_trace_id);
|
||||
assert_ne!(injected_span_id, parent_span_id);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn inject_without_trace_layer_uses_parent_span() {
|
||||
let _env = OtelTestEnv::install();
|
||||
|
||||
let capture = CaptureService::new();
|
||||
let seen = Arc::clone(&capture.seen);
|
||||
let mut svc = ServiceBuilder::new()
|
||||
.layer(InjectTraceContextLayer)
|
||||
.service(capture);
|
||||
|
||||
let parent = tracing::info_span!("parent_handler");
|
||||
let parent_span_id = otel_span_id_hex(&parent);
|
||||
let parent_trace_id = otel_trace_id_hex(&parent);
|
||||
|
||||
async {
|
||||
let fut = Service::call(&mut svc, post_req());
|
||||
fut.await.unwrap();
|
||||
}
|
||||
.instrument(parent)
|
||||
.await;
|
||||
|
||||
let headers = seen.lock().unwrap().clone().expect("headers captured");
|
||||
let tp = headers.get("traceparent").unwrap().to_str().unwrap();
|
||||
let (_ver, injected_trace_id, injected_span_id) = parse_traceparent(tp);
|
||||
|
||||
assert_eq!(injected_trace_id, parent_trace_id);
|
||||
assert_eq!(injected_span_id, parent_span_id);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn grpc_status_non_ok_invokes_on_failure_classifier() {
|
||||
let _env = OtelTestEnv::install();
|
||||
|
||||
let failures: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
|
||||
let failures_cb = Arc::clone(&failures);
|
||||
|
||||
let capture = CaptureService::with_grpc_status("13");
|
||||
let mut svc = ServiceBuilder::new()
|
||||
.layer(
|
||||
TraceLayer::new_for_grpc()
|
||||
.make_span_with(MakeClientSpan)
|
||||
.on_failure(
|
||||
move |class: GrpcFailureClass,
|
||||
_latency: Duration,
|
||||
_span: &tracing::Span| {
|
||||
failures_cb.lock().unwrap().push(class.to_string());
|
||||
},
|
||||
),
|
||||
)
|
||||
.layer(InjectTraceContextLayer)
|
||||
.service(capture);
|
||||
|
||||
let fut = Service::call(&mut svc, post_req());
|
||||
let _ = fut.await.unwrap();
|
||||
|
||||
let recorded = failures.lock().unwrap().clone();
|
||||
assert_eq!(recorded.len(), 1, "{recorded:?}");
|
||||
assert!(
|
||||
recorded[0].contains("13") || recorded[0].to_lowercase().contains("code"),
|
||||
"{}",
|
||||
recorded[0]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn grpc_status_ok_does_not_invoke_on_failure() {
|
||||
let _env = OtelTestEnv::install();
|
||||
|
||||
let failures: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
|
||||
let failures_cb = Arc::clone(&failures);
|
||||
|
||||
let capture = CaptureService::with_grpc_status("0");
|
||||
let mut svc = ServiceBuilder::new()
|
||||
.layer(
|
||||
TraceLayer::new_for_grpc()
|
||||
.make_span_with(MakeClientSpan)
|
||||
.on_failure(
|
||||
move |class: GrpcFailureClass,
|
||||
_latency: Duration,
|
||||
_span: &tracing::Span| {
|
||||
failures_cb.lock().unwrap().push(class.to_string());
|
||||
},
|
||||
),
|
||||
)
|
||||
.layer(InjectTraceContextLayer)
|
||||
.service(capture);
|
||||
|
||||
let fut = Service::call(&mut svc, post_req());
|
||||
let _ = fut.await.unwrap();
|
||||
|
||||
assert!(failures.lock().unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn attach_trace_to_grpc_request_sets_traceparent_metadata() {
|
||||
let _env = OtelTestEnv::install();
|
||||
let span = tracing::info_span!("handler");
|
||||
let span_id = otel_span_id_hex(&span);
|
||||
let _enter = span.enter();
|
||||
|
||||
let req = attach_trace_to_grpc_request(tonic::Request::new(())).unwrap();
|
||||
let tp = req
|
||||
.metadata()
|
||||
.get("traceparent")
|
||||
.expect("traceparent in metadata")
|
||||
.to_str()
|
||||
.unwrap();
|
||||
let (_ver, _tid, injected_span_id) = parse_traceparent(tp);
|
||||
assert_eq!(injected_span_id, span_id);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn make_client_span_records_otel_kind_client() {
|
||||
let env = OtelTestEnv::install();
|
||||
{
|
||||
let req = post_req();
|
||||
let mut make = MakeClientSpan;
|
||||
let span = make.make_span(&req);
|
||||
let _e = span.enter();
|
||||
}
|
||||
let spans = env.finished_spans();
|
||||
let grpc = spans
|
||||
.iter()
|
||||
.find(|s| s.name == "grpc_request")
|
||||
.expect("grpc_request");
|
||||
assert_eq!(grpc.span_kind, opentelemetry::trace::SpanKind::Client);
|
||||
}
|
||||
}
|
||||
@@ -1,155 +0,0 @@
|
||||
use async_trait::async_trait;
|
||||
use opentelemetry::global;
|
||||
use opentelemetry_http::HeaderInjector;
|
||||
use reqwest::header::HeaderMap;
|
||||
use reqwest_middleware::{ClientBuilder, ClientWithMiddleware, Middleware};
|
||||
use tracing::{Instrument, Span, field};
|
||||
use tracing_opentelemetry::OpenTelemetrySpanExt;
|
||||
|
||||
pub fn attach_trace_to_http_request(headers: &mut HeaderMap) {
|
||||
global::get_text_map_propagator(|propagator| {
|
||||
let context = Span::current().context();
|
||||
propagator.inject_context(&context, &mut HeaderInjector(headers));
|
||||
});
|
||||
}
|
||||
|
||||
pub type TracedHttpClient = ClientWithMiddleware;
|
||||
|
||||
pub fn traced_client(client: reqwest::Client) -> TracedHttpClient {
|
||||
ClientBuilder::new(client).with(TracingMiddleware).build()
|
||||
}
|
||||
|
||||
pub fn traced_client_new() -> TracedHttpClient {
|
||||
traced_client(reqwest::Client::new())
|
||||
}
|
||||
|
||||
pub fn traced_client_from_builder(
|
||||
builder: reqwest::ClientBuilder,
|
||||
) -> Result<TracedHttpClient, reqwest::Error> {
|
||||
Ok(traced_client(builder.build()?))
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
struct TracingMiddleware;
|
||||
|
||||
#[async_trait]
|
||||
impl Middleware for TracingMiddleware {
|
||||
async fn handle(
|
||||
&self,
|
||||
mut req: reqwest::Request,
|
||||
extensions: &mut http::Extensions,
|
||||
next: reqwest_middleware::Next<'_>,
|
||||
) -> reqwest_middleware::Result<reqwest::Response> {
|
||||
let method = req.method().as_str().to_owned();
|
||||
let url = req.url().clone();
|
||||
// No active dispatcher → the span has no consumer and would only be
|
||||
// downgraded to `log` spam. See `crate::dispatcher_active`.
|
||||
let span = if crate::dispatcher_active() {
|
||||
tracing::info_span!(
|
||||
"http_request",
|
||||
otel.kind = "client",
|
||||
"http.request.method" = %method,
|
||||
"url.full" = %url,
|
||||
"http.response.status_code" = field::Empty,
|
||||
)
|
||||
} else {
|
||||
Span::none()
|
||||
};
|
||||
|
||||
let result = async move {
|
||||
attach_trace_to_http_request(req.headers_mut());
|
||||
next.run(req, extensions).await
|
||||
}
|
||||
.instrument(span.clone())
|
||||
.await;
|
||||
|
||||
if let Ok(ref response) = result {
|
||||
span.record("http.response.status_code", response.status().as_u16());
|
||||
}
|
||||
result
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::testing::{OtelTestEnv, otel_span_id_hex, otel_trace_id_hex, parse_traceparent};
|
||||
use tracing::Instrument;
|
||||
use wiremock::matchers::method;
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate};
|
||||
|
||||
#[tokio::test]
|
||||
async fn attach_trace_to_http_request_writes_traceparent() {
|
||||
let _env = OtelTestEnv::install();
|
||||
let span = tracing::info_span!("http_request", otel.kind = "client");
|
||||
let span_id = otel_span_id_hex(&span);
|
||||
let mut headers = HeaderMap::new();
|
||||
let _enter = span.enter();
|
||||
attach_trace_to_http_request(&mut headers);
|
||||
let tp = headers.get("traceparent").unwrap().to_str().unwrap();
|
||||
let (_ver, _tid, injected_span_id) = parse_traceparent(tp);
|
||||
assert_eq!(injected_span_id, span_id);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn traced_client_injects_client_span_not_parent_on_wire() {
|
||||
let _env = OtelTestEnv::install();
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.respond_with(ResponseTemplate::new(200))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let client = traced_client(reqwest::Client::new());
|
||||
let parent = tracing::info_span!("parent_handler");
|
||||
let parent_span_id = otel_span_id_hex(&parent);
|
||||
let parent_trace_id = otel_trace_id_hex(&parent);
|
||||
|
||||
async {
|
||||
client
|
||||
.get(format!("{}/health", server.uri()))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
.instrument(parent)
|
||||
.await;
|
||||
|
||||
let received = server.received_requests().await.unwrap();
|
||||
assert_eq!(received.len(), 1);
|
||||
let tp = received[0]
|
||||
.headers
|
||||
.get("traceparent")
|
||||
.expect("traceparent on wire")
|
||||
.to_str()
|
||||
.unwrap();
|
||||
let (_ver, injected_trace_id, injected_span_id) = parse_traceparent(tp);
|
||||
|
||||
assert_eq!(injected_trace_id, parent_trace_id);
|
||||
assert_ne!(injected_span_id, parent_span_id);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn traced_client_returns_response_status() {
|
||||
let _env = OtelTestEnv::install();
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.respond_with(ResponseTemplate::new(404))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let client = traced_client_new();
|
||||
let resp = client
|
||||
.get(format!("{}/missing", server.uri()))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status().as_u16(), 404);
|
||||
assert!(
|
||||
server.received_requests().await.unwrap()[0]
|
||||
.headers
|
||||
.get("traceparent")
|
||||
.is_some()
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,19 +0,0 @@
|
||||
mod dispatch;
|
||||
mod grpc_client;
|
||||
mod timer;
|
||||
|
||||
pub mod fastrace;
|
||||
pub mod http_client;
|
||||
pub mod tokio;
|
||||
|
||||
#[cfg(test)]
|
||||
mod testing;
|
||||
|
||||
pub use dispatch::*;
|
||||
pub use fastrace::*;
|
||||
pub use grpc_client::*;
|
||||
pub use http_client::{
|
||||
TracedHttpClient, attach_trace_to_http_request, traced_client, traced_client_from_builder,
|
||||
traced_client_new,
|
||||
};
|
||||
pub use timer::*;
|
||||
@@ -1,65 +0,0 @@
|
||||
use opentelemetry::global;
|
||||
use opentelemetry::trace::{SpanContext, TraceContextExt, TracerProvider as _};
|
||||
use opentelemetry_sdk::propagation::TraceContextPropagator;
|
||||
use opentelemetry_sdk::trace::{
|
||||
InMemorySpanExporter, InMemorySpanExporterBuilder, SdkTracerProvider, SimpleSpanProcessor,
|
||||
};
|
||||
use tracing::Span;
|
||||
use tracing_opentelemetry::OpenTelemetrySpanExt;
|
||||
use tracing_subscriber::prelude::*;
|
||||
|
||||
pub struct OtelTestEnv {
|
||||
_guard: tracing::subscriber::DefaultGuard,
|
||||
provider: SdkTracerProvider,
|
||||
exporter: InMemorySpanExporter,
|
||||
}
|
||||
|
||||
impl OtelTestEnv {
|
||||
pub fn install() -> Self {
|
||||
global::set_text_map_propagator(TraceContextPropagator::new());
|
||||
let exporter = InMemorySpanExporterBuilder::new().build();
|
||||
let provider = SdkTracerProvider::builder()
|
||||
.with_span_processor(SimpleSpanProcessor::new(exporter.clone()))
|
||||
.build();
|
||||
let tracer = provider.tracer("kigi-tracing-test");
|
||||
let otel_layer = tracing_opentelemetry::layer()
|
||||
.with_tracer(tracer)
|
||||
.with_context_activation(false)
|
||||
.with_filter(tracing_subscriber::filter::LevelFilter::INFO);
|
||||
let guard = tracing_subscriber::registry()
|
||||
.with(otel_layer)
|
||||
.set_default();
|
||||
Self {
|
||||
_guard: guard,
|
||||
provider,
|
||||
exporter,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn finished_spans(&self) -> Vec<opentelemetry_sdk::trace::SpanData> {
|
||||
let _ = self.provider.force_flush();
|
||||
self.exporter.get_finished_spans().unwrap_or_default()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn parse_traceparent(value: &str) -> (&str, &str, &str) {
|
||||
let mut parts = value.split('-');
|
||||
let version = parts.next().expect("version");
|
||||
let trace_id = parts.next().expect("trace_id");
|
||||
let span_id = parts.next().expect("span_id");
|
||||
(version, trace_id, span_id)
|
||||
}
|
||||
|
||||
pub fn otel_span_id_hex(span: &Span) -> String {
|
||||
let cx = span.context();
|
||||
let span_ref = cx.span();
|
||||
let sc: &SpanContext = span_ref.span_context();
|
||||
format!("{:016x}", u64::from_be_bytes(sc.span_id().to_bytes()))
|
||||
}
|
||||
|
||||
pub fn otel_trace_id_hex(span: &Span) -> String {
|
||||
let cx = span.context();
|
||||
let span_ref = cx.span();
|
||||
let sc: &SpanContext = span_ref.span_context();
|
||||
format!("{:032x}", u128::from_be_bytes(sc.trace_id().to_bytes()))
|
||||
}
|
||||
@@ -1,62 +0,0 @@
|
||||
use log::error;
|
||||
use log::info;
|
||||
use tokio::time::Instant;
|
||||
|
||||
/// A simple timer that logs the runtime of an operation.
|
||||
pub struct Timer {
|
||||
/// Time the operation started.
|
||||
start: Instant,
|
||||
/// An ID shown in the logs to associate the log messages from the timer with each other-
|
||||
id: uuid::Uuid,
|
||||
/// A string that is being logged.
|
||||
message: String,
|
||||
/// True if the timer has been stopped already.
|
||||
stopped: bool,
|
||||
}
|
||||
|
||||
impl Timer {
|
||||
/// Creates a new Timer instance and starts the timer.
|
||||
pub fn new<S: AsRef<str>>(message: S) -> Self {
|
||||
let id = uuid::Uuid::new_v4();
|
||||
info!("[{}] START: {}", id, message.as_ref());
|
||||
Self {
|
||||
start: Instant::now(),
|
||||
id,
|
||||
message: message.as_ref().to_string(),
|
||||
stopped: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Stops the timer and logs the result.
|
||||
pub fn stop<T>(&mut self, result: T) -> T {
|
||||
if !self.stopped {
|
||||
let runtime = self.start.elapsed().as_secs_f32();
|
||||
info!(
|
||||
"[{}] FINISHED in {:.3}s: {}",
|
||||
self.id, runtime, self.message
|
||||
);
|
||||
self.stopped = true;
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
/// Stop the timer prematurely and logs an error.
|
||||
pub fn force_stop(&mut self) {
|
||||
if !self.stopped {
|
||||
let runtime = self.start.elapsed().as_secs_f32();
|
||||
error!(
|
||||
"[{}] FAILED after {:.3}s: {}",
|
||||
self.id, runtime, self.message
|
||||
);
|
||||
self.stopped = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Automatically report the runtime when the object is dropped.
|
||||
impl Drop for Timer {
|
||||
fn drop(&mut self) {
|
||||
self.force_stop();
|
||||
}
|
||||
}
|
||||
@@ -1,22 +0,0 @@
|
||||
use std::future::Future;
|
||||
use tokio::task::JoinHandle;
|
||||
use tracing::{Instrument, Span};
|
||||
|
||||
/// Utility macro for propagating the current tracing context to a newly spawned task.
|
||||
///
|
||||
/// Note: The spawned task will be associated with the currently active span. To create a *new*
|
||||
/// span for the spawned task, manually instrument the future using [tracing::Instrument] instead of
|
||||
/// using this macro. For example:
|
||||
///
|
||||
/// use tracing::{info_span, Instrument};
|
||||
///
|
||||
/// let fut = tokio::spawn(async move {
|
||||
/// print!("do stuff")
|
||||
/// }.instrument(info_span!("spawned task")));
|
||||
pub fn spawn_traced<F>(future: F) -> JoinHandle<F::Output>
|
||||
where
|
||||
F: Future + Send + 'static,
|
||||
F::Output: Send + 'static,
|
||||
{
|
||||
tokio::spawn(future.instrument(Span::current()))
|
||||
}
|
||||
Reference in New Issue
Block a user