diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 88c90df..b59e7b0 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -80,7 +80,7 @@ jobs: path: | ~/.cargo/registry ~/.cargo/git - + - run: bash ./ci/prepare_build_environment.sh - name: Bulid release binary run: cargo build --release --target ${{ matrix.target }} diff --git a/Cargo.lock b/Cargo.lock index 149a2d1..19dea27 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -142,6 +142,28 @@ version = "1.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bddcadddf5e9015d310179a59bb28c4d4b9920ad0f11e8e14dbadf654890c9a6" +[[package]] +name = "async-stream" +version = "0.3.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd56dd203fef61ac097dd65721a419ddccb106b2d2b70ba60a6b529f03961a51" +dependencies = [ + "async-stream-impl", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-stream-impl" +version = "0.3.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16e62a023e7c117e27523144c5d2459f4397fcc3cab0085af8e2224f643a0193" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.18", +] + [[package]] name = "async-trait" version = "0.1.68" @@ -208,6 +230,22 @@ dependencies = [ "tower-service", ] +[[package]] +name = "axum-tracing-opentelemetry" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "164b95427e83b79583c7699a72b4a6b485a12bbdef5b5c054ee5ff2296d82f52" +dependencies = [ + "axum", + "futures", + "http", + "opentelemetry", + "tower", + "tower-http 0.3.5", + "tracing", + "tracing-opentelemetry", +] + [[package]] name = "backtrace" version = "0.3.67" @@ -868,6 +906,12 @@ dependencies = [ "windows-sys 0.48.0", ] +[[package]] +name = "fixedbitset" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80" + [[package]] name = "flate2" version = "1.0.26" @@ -918,6 +962,21 @@ dependencies = [ "winapi", ] +[[package]] +name = "futures" +version = "0.3.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23342abe12aba583913b2e62f22225ff9c950774065e4bfb61a19cd9770fec40" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" version = "0.3.28" @@ -925,6 +984,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "955518d47e09b25bbebc7a18df10b81f0c766eaf4c4f1cccef2fca5f2a4fb5f2" dependencies = [ "futures-core", + "futures-sink", ] [[package]] @@ -933,6 +993,17 @@ version = "0.3.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4bca583b7e26f571124fe5b7561d49cb2868d79116cfa0eefce955557c6fee8c" +[[package]] +name = "futures-executor" +version = "0.3.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccecee823288125bd88b4d7f565c9e58e41858e47ab72e8ea2d64e93624386e0" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + [[package]] name = "futures-io" version = "0.3.28" @@ -968,6 +1039,7 @@ version = "0.3.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "26b01e40b772d54cf6c6d721c1d1abd0647a0106a12ecaa1c186273392a69533" dependencies = [ + "futures-channel", "futures-core", "futures-io", "futures-macro", @@ -1153,6 +1225,18 @@ dependencies = [ "want", ] +[[package]] +name = "hyper-timeout" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbb958482e8c7be4bc3cf272a766a2b0bf1a6755e7a6ae777f017a31d11b13b1" +dependencies = [ + "hyper", + "pin-project-lite", + "tokio", + "tokio-io-timeout", +] + [[package]] name = "hyper-tls" version = "0.5.0" @@ -1592,6 +1676,12 @@ dependencies = [ "syn 2.0.18", ] +[[package]] +name = "multimap" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5ce46fe64a9d73be07dcbe690a38ce1b293be448fd8ce1e6c1b8062c9f72c6a" + [[package]] name = "murmurhash32" version = "0.2.0" @@ -1769,6 +1859,86 @@ dependencies = [ "vcpkg", ] +[[package]] +name = "opentelemetry" +version = "0.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69d6c3d7288a106c0a363e4b0e8d308058d56902adefb16f4936f417ffef086e" +dependencies = [ + "opentelemetry_api", + "opentelemetry_sdk", +] + +[[package]] +name = "opentelemetry-otlp" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1c928609d087790fc936a1067bdc310ae702bdf3b090c3f281b713622c8bbde" +dependencies = [ + "async-trait", + "futures", + "futures-util", + "http", + "opentelemetry", + "opentelemetry-proto", + "prost", + "thiserror", + "tokio", + "tonic", +] + +[[package]] +name = "opentelemetry-proto" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d61a2f56df5574508dd86aaca016c917489e589ece4141df1b5e349af8d66c28" +dependencies = [ + "futures", + "futures-util", + "opentelemetry", + "prost", + "tonic", + "tonic-build", +] + +[[package]] +name = "opentelemetry_api" +version = "0.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c24f96e21e7acc813c7a8394ee94978929db2bcc46cf6b5014fc612bf7760c22" +dependencies = [ + "fnv", + "futures-channel", + "futures-util", + "indexmap", + "js-sys", + "once_cell", + "pin-project-lite", + "thiserror", +] + +[[package]] +name = "opentelemetry_sdk" +version = "0.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ca41c4933371b61c2a2f214bf16931499af4ec90543604ec828f7a625c09113" +dependencies = [ + "async-trait", + "crossbeam-channel", + "dashmap", + "fnv", + "futures-channel", + "futures-executor", + "futures-util", + "once_cell", + "opentelemetry_api", + "percent-encoding", + "rand", + "thiserror", + "tokio", + "tokio-stream", +] + [[package]] name = "overload" version = "0.1.1" @@ -1842,6 +2012,16 @@ version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "478c572c3d73181ff3c2539045f6eb99e5491218eae919370993b890cdbdd98e" +[[package]] +name = "petgraph" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4dd7d28ee937e54fe3080c91faa1c3a46c06de6252988a7f4592ba2310ef22a4" +dependencies = [ + "fixedbitset", + "indexmap", +] + [[package]] name = "pin-project" version = "1.1.0" @@ -1901,6 +2081,16 @@ version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5b40af805b3121feab8a3c29f04d8ad262fa8e0561883e7653e024ae4479e6de" +[[package]] +name = "prettyplease" +version = "0.1.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c8646e95016a7a6c4adea95bafa8a16baab64b583356217f2c85db4a39d9a86" +dependencies = [ + "proc-macro2", + "syn 1.0.109", +] + [[package]] name = "proc-macro-error" version = "1.0.4" @@ -1934,6 +2124,60 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "prost" +version = "0.11.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b82eaa1d779e9a4bc1c3217db8ffbeabaae1dca241bf70183242128d48681cd" +dependencies = [ + "bytes", + "prost-derive", +] + +[[package]] +name = "prost-build" +version = "0.11.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "119533552c9a7ffacc21e099c24a0ac8bb19c2a2a3f363de84cd9b844feab270" +dependencies = [ + "bytes", + "heck", + "itertools 0.10.5", + "lazy_static", + "log", + "multimap", + "petgraph", + "prettyplease", + "prost", + "prost-types", + "regex", + "syn 1.0.109", + "tempfile", + "which", +] + +[[package]] +name = "prost-derive" +version = "0.11.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5d2d8d10f3c6ded6da8b05b5fb3b8a5082514344d56c9f871412d29b4e075b4" +dependencies = [ + "anyhow", + "itertools 0.10.5", + "proc-macro2", + "quote", + "syn 1.0.109", +] + +[[package]] +name = "prost-types" +version = "0.11.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "213622a1460818959ac1181aaeb2dc9c7f63df720db7d788b3e24eacd1983e13" +dependencies = [ + "prost", +] + [[package]] name = "quote" version = "1.0.28" @@ -2513,11 +2757,14 @@ name = "tabby" version = "0.1.0" dependencies = [ "axum", + "axum-tracing-opentelemetry", "clap", "ctranslate2-bindings", "hyper", "lazy_static", "mime_guess", + "opentelemetry", + "opentelemetry-otlp", "rust-embed", "serde", "serde_json", @@ -2529,8 +2776,9 @@ dependencies = [ "tabby-scheduler", "tokio", "tower", - "tower-http", + "tower-http 0.4.0", "tracing", + "tracing-opentelemetry", "tracing-subscriber 0.3.17", "utoipa", "utoipa-swagger-ui", @@ -2843,6 +3091,16 @@ dependencies = [ "windows-sys 0.48.0", ] +[[package]] +name = "tokio-io-timeout" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30b74022ada614a1b4834de765f9bb43877f910cc8ce4be40e89042c9223a8bf" +dependencies = [ + "pin-project-lite", + "tokio", +] + [[package]] name = "tokio-macros" version = "2.1.0" @@ -2864,6 +3122,17 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-stream" +version = "0.1.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "397c988d37662c7dda6d2208364a706264bf3d6138b11d436cbac0ad38832842" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", +] + [[package]] name = "tokio-util" version = "0.7.8" @@ -2912,6 +3181,51 @@ dependencies = [ "winnow", ] +[[package]] +name = "tonic" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f219fad3b929bef19b1f86fbc0358d35daed8f2cac972037ac0dc10bbb8d5fb" +dependencies = [ + "async-stream", + "async-trait", + "axum", + "base64 0.13.1", + "bytes", + "futures-core", + "futures-util", + "h2", + "http", + "http-body", + "hyper", + "hyper-timeout", + "percent-encoding", + "pin-project", + "prost", + "prost-derive", + "tokio", + "tokio-stream", + "tokio-util", + "tower", + "tower-layer", + "tower-service", + "tracing", + "tracing-futures", +] + +[[package]] +name = "tonic-build" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5bf5e9b9c0f7e0a7c027dcfaba7b2c60816c7049171f679d99ee2ff65d0de8c4" +dependencies = [ + "prettyplease", + "proc-macro2", + "prost-build", + "quote", + "syn 1.0.109", +] + [[package]] name = "tower" version = "0.4.13" @@ -2920,9 +3234,32 @@ checksum = "b8fa9be0de6cf49e536ce1851f987bd21a43b771b09473c3549a6c853db37c1c" dependencies = [ "futures-core", "futures-util", + "indexmap", "pin-project", "pin-project-lite", + "rand", + "slab", "tokio", + "tokio-util", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tower-http" +version = "0.3.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f873044bf02dd1e8239e9c1293ea39dad76dc594ec16185d0a1bf31d8dc8d858" +dependencies = [ + "bitflags", + "bytes", + "futures-core", + "futures-util", + "http", + "http-body", + "http-range-header", + "pin-project-lite", "tower-layer", "tower-service", "tracing", @@ -2992,6 +3329,16 @@ dependencies = [ "valuable", ] +[[package]] +name = "tracing-futures" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97d095ae15e245a057c8e8451bab9b3ee1e1f68e9ba2b4fbc18d0ac5237835f2" +dependencies = [ + "pin-project", + "tracing", +] + [[package]] name = "tracing-log" version = "0.1.3" @@ -3003,6 +3350,20 @@ dependencies = [ "tracing-core", ] +[[package]] +name = "tracing-opentelemetry" +version = "0.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "21ebb87a95ea13271332df069020513ab70bdb5637ca42d6e492dc3bbbad48de" +dependencies = [ + "once_cell", + "opentelemetry", + "tracing", + "tracing-core", + "tracing-log", + "tracing-subscriber 0.3.17", +] + [[package]] name = "tracing-serde" version = "0.1.3" @@ -3398,6 +3759,17 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "which" +version = "4.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2441c784c52b289a054b7201fc93253e288f094e2f4be9058343127c4226a269" +dependencies = [ + "either", + "libc", + "once_cell", +] + [[package]] name = "winapi" version = "0.3.9" diff --git a/Dockerfile b/Dockerfile index c09d74a..8461764 100644 --- a/Dockerfile +++ b/Dockerfile @@ -9,6 +9,7 @@ RUN apt-get update && \ curl \ pkg-config \ libssl-dev \ + protobuf-compiler \ && \ apt-get clean && \ rm -rf /var/lib/apt/lists/* diff --git a/ci/prepare_build_environment.sh b/ci/prepare_build_environment.sh new file mode 100755 index 0000000..e1b3c09 --- /dev/null +++ b/ci/prepare_build_environment.sh @@ -0,0 +1,5 @@ +#!/bin/bash + +if [[ "$OSTYPE" == "darwin"* ]]; then + brew install protobuf +fi diff --git a/crates/tabby/Cargo.toml b/crates/tabby/Cargo.toml index 9714b15..1d53364 100644 --- a/crates/tabby/Cargo.toml +++ b/crates/tabby/Cargo.toml @@ -26,6 +26,11 @@ strum = { version = "0.24", features = ["derive"] } strfmt = "0.2.4" tracing = { workspace = true } tracing-subscriber = { workspace = true } +opentelemetry = { version = "0.18.0", features = ["rt-tokio"] } +opentelemetry-otlp = "0.11.0" +axum-tracing-opentelemetry = "0.10.0" +tracing-opentelemetry = "0.18.0" + [dependencies.uuid] version = "1.3.3" diff --git a/crates/tabby/src/main.rs b/crates/tabby/src/main.rs index d48a9e1..4891ef9 100644 --- a/crates/tabby/src/main.rs +++ b/crates/tabby/src/main.rs @@ -2,7 +2,13 @@ mod download; mod serve; use clap::{Parser, Subcommand}; -use tracing_subscriber::EnvFilter; +use opentelemetry::{ + global, + sdk::{propagation::TraceContextPropagator, trace, trace::Sampler, Resource}, + KeyValue, +}; +use opentelemetry_otlp::WithExportConfig; +use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter, Layer}; #[derive(Parser)] #[command(author, version, about, long_about = None)] @@ -10,6 +16,10 @@ use tracing_subscriber::EnvFilter; struct Cli { #[command(subcommand)] command: Commands, + + /// Open Telemetry endpoint. + #[clap(long)] + otlp_endpoint: Option, } #[derive(Subcommand)] @@ -33,11 +43,8 @@ pub struct SchedulerArgs { #[tokio::main] async fn main() { - tracing_subscriber::fmt::fmt() - .with_env_filter(EnvFilter::from_default_env().add_directive("tabby=info".parse().unwrap())) - .init(); - let cli = Cli::parse(); + init_logging(cli.otlp_endpoint); match &cli.command { Commands::Serve(args) => serve::main(args).await, @@ -46,6 +53,8 @@ async fn main() { .await .unwrap_or_else(|err| fatal!("Scheduler failed due to '{}'", err)), } + + opentelemetry::global::shutdown_tracer_provider(); } #[macro_export] @@ -64,3 +73,50 @@ macro_rules! fatal { }) }; } + +fn init_logging(otlp_endpoint: Option) { + let mut layers = Vec::new(); + + let fmt_layer = tracing_subscriber::fmt::layer() + .with_file(true) + .with_line_number(true) + .boxed(); + + layers.push(fmt_layer); + + if let Some(otlp_endpoint) = &otlp_endpoint { + global::set_text_map_propagator(TraceContextPropagator::new()); + + let tracer = opentelemetry_otlp::new_pipeline() + .tracing() + .with_exporter( + opentelemetry_otlp::new_exporter() + .tonic() + .with_endpoint(otlp_endpoint), + ) + .with_trace_config( + trace::config() + .with_resource(Resource::new(vec![KeyValue::new( + "service.name", + "tabby.server", + )])) + .with_sampler(Sampler::AlwaysOn), + ) + .install_batch(opentelemetry::runtime::Tokio); + + if let Ok(tracer) = tracer { + layers.push(tracing_opentelemetry::layer().with_tracer(tracer).boxed()); + axum_tracing_opentelemetry::init_propagator().unwrap(); + }; + } + + let env_filter = EnvFilter::from_default_env() + .add_directive("tabby=info".parse().unwrap()) + .add_directive("axum_tracing_opentelemetry=info".parse().unwrap()) + .add_directive("otel=debug".parse().unwrap()); + + tracing_subscriber::registry() + .with(layers) + .with(env_filter) + .init(); +} diff --git a/crates/tabby/src/serve/completions.rs b/crates/tabby/src/serve/completions.rs index 5d67d65..21d7674 100644 --- a/crates/tabby/src/serve/completions.rs +++ b/crates/tabby/src/serve/completions.rs @@ -8,6 +8,7 @@ use hyper::StatusCode; use serde::{Deserialize, Serialize}; use strfmt::{strfmt, strfmt_builder}; use tabby_common::{events, path::ModelDir}; +use tracing::instrument; use utoipa::ToSchema; use self::languages::get_stop_words; @@ -64,6 +65,7 @@ pub struct CompletionResponse { (status = 400, description = "Bad Request") ) )] +#[instrument(skip(state, request))] pub async fn completion( State(state): State>, Json(request): Json, diff --git a/crates/tabby/src/serve/mod.rs b/crates/tabby/src/serve/mod.rs index 1e622e9..df380d0 100644 --- a/crates/tabby/src/serve/mod.rs +++ b/crates/tabby/src/serve/mod.rs @@ -8,6 +8,7 @@ use std::{ }; use axum::{routing, Router, Server}; +use axum_tracing_opentelemetry::opentelemetry_tracing_layer; use clap::Args; use tower_http::cors::CorsLayer; use tracing::info; @@ -92,8 +93,7 @@ pub async fn main(args: &ServeArgs) { let app = Router::new() .merge(SwaggerUi::new("/swagger-ui").url("/api-docs/openapi.json", ApiDoc::openapi())) .nest("/v1", api_router(args)) - .fallback(fallback(args.experimental_admin_panel)) - .layer(CorsLayer::permissive()); + .fallback(fallback(args.experimental_admin_panel)); let address = SocketAddr::from((Ipv4Addr::UNSPECIFIED, args.port)); info!("Listening at {}", address); @@ -112,6 +112,8 @@ fn api_router(args: &ServeArgs) -> Router { routing::post(completions::completion) .with_state(Arc::new(completions::CompletionState::new(args))), ) + .layer(CorsLayer::permissive()) + .layer(opentelemetry_tracing_layer()) } fn fallback(experimental_admin_panel: bool) -> routing::MethodRouter {