Skip to content

Commit a8aaf06

Browse files
cdxkercdxker
committed
feature: add tracing and sentry back
1 parent 0737586 commit a8aaf06

19 files changed

+639
-36
lines changed

‎server/Cargo.lock‎

Lines changed: 358 additions & 28 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎server/Cargo.toml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -196,6 +196,11 @@ pagefind = { version = "1.3.0" }
196196
tl = "0.7.8"
197197
url = "2.5.4"
198198
actix-request-reply-cache = "0.1.6"
199+
sentry = { version = "0.41.0", features = ["tracing"] }
200+
tracing = "0.1.41"
201+
tracing-subscriber = { version = "0.3.19", features = ["env-filter", "registry"] }
202+
sentry-tracing = "0.41.0"
203+
sentry-actix = "0.41.0"
199204

200205
[build-dependencies]
201206
dotenvy = "0.15.7"

‎server/src/bin/ingestion-worker.rs‎

Lines changed: 41 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ use signal_hook::consts::SIGTERM;
1111
use std::collections::HashMap;
1212
use std::error::Error;
1313
use std::sync::{atomic::AtomicBool, Arc};
14+
use tracing_subscriber::{prelude::*, EnvFilter, Layer};
1415
use trieve_server::data::models::{
1516
self, ChunkBoost, ChunkData, ChunkGroup, ChunkMetadata, DatasetConfiguration,
1617
PagefindIndexWorkerMessage, QdrantPayload, WorkerEvent,
@@ -43,10 +44,46 @@ use trieve_server::{establish_connection, get_env};
4344
#[tokio::main]
4445
async fn main() -> Result<(), Box<dyn Error>> {
4546
dotenvy::dotenv().ok();
46-
env_logger::builder()
47-
.target(env_logger::Target::Stdout)
48-
.filter_level(log::LevelFilter::Info)
49-
.init();
47+
// env_logger::builder()
48+
// .target(env_logger::Target::Stdout)
49+
// .filter_level(log::LevelFilter::Info)
50+
// .init();
51+
let sentry_url = std::env::var("SENTRY_URL");
52+
let _guard = if let Ok(sentry_url) = sentry_url {
53+
let guard = sentry::init((
54+
sentry_url,
55+
sentry::ClientOptions {
56+
release: sentry::release_name!(),
57+
traces_sample_rate: 1.0,
58+
send_default_pii: true,
59+
..Default::default()
60+
},
61+
));
62+
63+
tracing_subscriber::Registry::default()
64+
.with(sentry::integrations::tracing::layer())
65+
.with(
66+
tracing_subscriber::fmt::layer().with_filter(
67+
EnvFilter::from_default_env()
68+
.add_directive(tracing_subscriber::filter::LevelFilter::INFO.into()),
69+
),
70+
)
71+
.init();
72+
73+
log::info!("Sentry monitoring enabled");
74+
Some(guard)
75+
} else {
76+
tracing_subscriber::Registry::default()
77+
.with(
78+
tracing_subscriber::fmt::layer().with_filter(
79+
tracing_subscriber::EnvFilter::from_default_env()
80+
.add_directive(tracing_subscriber::filter::LevelFilter::INFO.into()),
81+
),
82+
)
83+
.init();
84+
85+
None
86+
};
5087

5188
let database_url = get_env!("DATABASE_URL", "DATABASE_URL is not set");
5289

‎server/src/handlers/chunk_handler.rs‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -312,6 +312,7 @@ pub enum CreateChunkReqPayloadEnum {
312312
("ApiKey" = ["admin"]),
313313
)
314314
)]
315+
#[tracing::instrument(skip_all)]
315316
pub async fn create_chunk(
316317
create_chunk_data: web::Json<CreateChunkReqPayloadEnum>,
317318
pool: web::Data<Pool>,
@@ -568,6 +569,7 @@ pub async fn create_chunk(
568569
}
569570

570571
/// Delete Chunk
572+
#[tracing::instrument(skip_all)]
571573
///
572574
/// Delete a chunk by its id. Auth'ed user or api key must have an admin or owner role for the specified dataset's organization.
573575
#[utoipa::path(
@@ -639,6 +641,7 @@ pub struct BulkDeleteChunkPayload {
639641
("ApiKey" = ["admin"]),
640642
)
641643
)]
644+
#[tracing::instrument(skip_all)]
642645
pub async fn bulk_delete_chunk(
643646
chunk_filter: web::Json<BulkDeleteChunkPayload>,
644647
redis_pool: web::Data<RedisPool>,
@@ -690,6 +693,7 @@ pub async fn bulk_delete_chunk(
690693
("ApiKey" = ["admin"]),
691694
)
692695
)]
696+
#[tracing::instrument(skip_all)]
693697
pub async fn delete_chunk_by_tracking_id(
694698
tracking_id: web::Path<String>,
695699
pool: web::Data<Pool>,
@@ -796,6 +800,7 @@ pub struct UpdateIngestionMessage {
796800
("ApiKey" = ["admin"]),
797801
)
798802
)]
803+
#[tracing::instrument(skip_all)]
799804
pub async fn update_chunk(
800805
update_chunk_data: web::Json<UpdateChunkReqPayload>,
801806
pool: web::Data<Pool>,
@@ -957,6 +962,7 @@ pub struct UpdateChunkByTrackingIdData {
957962
)
958963
)]
959964
#[deprecated]
965+
#[tracing::instrument(skip_all)]
960966
pub async fn update_chunk_by_tracking_id(
961967
update_chunk_data: web::Json<UpdateChunkByTrackingIdData>,
962968
pool: web::Data<Pool>,
@@ -1242,6 +1248,7 @@ pub fn is_audio(query: QueryTypes) -> bool {
12421248
("ApiKey" = ["readonly"]),
12431249
)
12441250
)]
1251+
#[tracing::instrument(skip_all)]
12451252
pub async fn search_chunks(
12461253
data: web::Json<SearchChunksReqPayload>,
12471254
_user: LoggedUser,
@@ -1526,6 +1533,7 @@ impl From<AutocompleteReqPayload> for SearchChunksReqPayload {
15261533
("ApiKey" = ["readonly"]),
15271534
)
15281535
)]
1536+
#[tracing::instrument(skip_all)]
15291537
pub async fn autocomplete(
15301538
data: web::Json<AutocompleteReqPayload>,
15311539
_user: LoggedUser,
@@ -1671,6 +1679,7 @@ pub struct ScrollChunksReqPayload {
16711679
("ApiKey" = ["readonly"]),
16721680
)
16731681
)]
1682+
#[tracing::instrument(skip_all)]
16741683
pub async fn scroll_dataset_chunks(
16751684
data: web::Json<ScrollChunksReqPayload>,
16761685
_user: LoggedUser,
@@ -1764,6 +1773,7 @@ pub async fn scroll_dataset_chunks(
17641773
("ApiKey" = ["readonly"]),
17651774
)
17661775
)]
1776+
#[tracing::instrument(skip_all)]
17671777
pub async fn get_chunk_by_id(
17681778
chunk_id: web::Path<uuid::Uuid>,
17691779
_user: LoggedUser,
@@ -1862,6 +1872,7 @@ pub struct CountChunkQueryResponseBody {
18621872
("ApiKey" = ["readonly"]),
18631873
)
18641874
)]
1875+
#[tracing::instrument(skip_all)]
18651876
pub async fn count_chunks(
18661877
data: web::Json<CountChunksReqPayload>,
18671878
_user: LoggedUser,
@@ -1960,6 +1971,7 @@ pub async fn count_chunks(
19601971
("ApiKey" = ["readonly"]),
19611972
)
19621973
)]
1974+
#[tracing::instrument(skip_all)]
19631975
pub async fn get_chunk_by_tracking_id(
19641976
tracking_id: web::Path<String>,
19651977
_user: LoggedUser,
@@ -2019,6 +2031,7 @@ pub struct GetChunksData {
20192031
("ApiKey" = ["readonly"]),
20202032
)
20212033
)]
2034+
#[tracing::instrument(skip_all)]
20222035
pub async fn get_chunks_by_ids(
20232036
chunk_payload: web::Json<GetChunksData>,
20242037
_user: LoggedUser,
@@ -2090,6 +2103,7 @@ pub struct GetTrackingChunksData {
20902103
("ApiKey" = ["readonly"]),
20912104
)
20922105
)]
2106+
#[tracing::instrument(skip_all)]
20932107
pub async fn get_chunks_by_tracking_ids(
20942108
chunk_payload: web::Json<GetTrackingChunksData>,
20952109
_user: LoggedUser,
@@ -2202,6 +2216,7 @@ pub enum RecommendResponseTypes {
22022216
("ApiKey" = ["readonly"]),
22032217
)
22042218
)]
2219+
#[tracing::instrument(skip_all)]
22052220
pub async fn get_recommended_chunks(
22062221
data: web::Json<RecommendChunksRequest>,
22072222
pool: web::Data<Pool>,
@@ -2563,6 +2578,7 @@ pub struct GenerateOffChunksReqPayload {
25632578
("ApiKey" = ["readonly"]),
25642579
)
25652580
)]
2581+
#[tracing::instrument(skip_all)]
25662582
pub async fn generate_off_chunks(
25672583
data: web::Json<GenerateOffChunksReqPayload>,
25682584
pool: web::Data<Pool>,
@@ -3108,6 +3124,7 @@ pub struct ChunkedContent {
31083124
),
31093125
),
31103126
)]
3127+
#[tracing::instrument(skip_all)]
31113128
pub async fn split_html_content(
31123129
body: web::Json<ChunkHtmlContentReqPayload>,
31133130
) -> Result<HttpResponse, ServiceError> {

‎server/src/handlers/message_handler.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1091,6 +1091,7 @@ pub struct GetToolFunctionParamsRespBody {
10911091
("ApiKey" = ["admin"]),
10921092
)
10931093
)]
1094+
#[tracing::instrument(name = "get_tool_function_params", skip(_required_user))]
10941095
pub async fn get_tool_function_params(
10951096
data: web::Json<GetToolFunctionParamsReqPayload>,
10961097
dataset_org_plan_sub: DatasetAndOrgWithSubAndPlan,

‎server/src/lib.rs‎

Lines changed: 38 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ use once_cell::sync::Lazy;
3737
use openssl::ssl::SslVerifyMode;
3838
use openssl::ssl::{SslConnector, SslMethod};
3939
use postgres_openssl::MakeTlsConnector;
40+
use tracing_subscriber::{prelude::*, EnvFilter, Layer};
4041
use ureq::json;
4142
use utoipa::{
4243
openapi::security::{ApiKey, ApiKeyValue, SecurityScheme},
@@ -746,10 +747,42 @@ pub struct ApiDoc;
746747
pub fn main() -> std::io::Result<()> {
747748
dotenvy::dotenv().ok();
748749

749-
env_logger::builder()
750-
.target(env_logger::Target::Stdout)
751-
.filter_level(log::LevelFilter::Info)
752-
.init();
750+
let sentry_url = std::env::var("SENTRY_URL");
751+
let _guard = if let Ok(sentry_url) = sentry_url {
752+
let guard = sentry::init((
753+
sentry_url,
754+
sentry::ClientOptions {
755+
release: sentry::release_name!(),
756+
traces_sample_rate: 1.0,
757+
send_default_pii: true,
758+
..Default::default()
759+
},
760+
));
761+
762+
tracing_subscriber::Registry::default()
763+
.with(sentry::integrations::tracing::layer())
764+
.with(
765+
tracing_subscriber::fmt::layer().with_filter(
766+
EnvFilter::from_default_env()
767+
.add_directive(tracing_subscriber::filter::LevelFilter::INFO.into()),
768+
),
769+
)
770+
.init();
771+
772+
log::info!("Sentry monitoring enabled");
773+
Some(guard)
774+
} else {
775+
tracing_subscriber::Registry::default()
776+
.with(
777+
tracing_subscriber::fmt::layer().with_filter(
778+
tracing_subscriber::EnvFilter::from_default_env()
779+
.add_directive(tracing_subscriber::filter::LevelFilter::INFO.into()),
780+
),
781+
)
782+
.init();
783+
784+
None
785+
};
753786

754787
let database_url = get_env!("DATABASE_URL", "DATABASE_URL should be set");
755788
let redis_url = get_env!("REDIS_URL", "REDIS_URL should be set");
@@ -923,6 +956,7 @@ pub fn main() -> std::io::Result<()> {
923956
.app_data(web::Data::new(event_queue.clone()))
924957
.app_data(web::Data::new(clickhouse_client.clone()))
925958
.app_data(web::Data::new(metrics.clone()))
959+
.wrap(sentry_actix::Sentry::new())
926960
.app_data(detector.clone())
927961
.app_data(web::Data::new(broccoli_queue.clone()))
928962
.wrap(from_fn(middleware::timeout_middleware::timeout_15secs))

‎server/src/middleware/auth_middleware.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -382,6 +382,7 @@ fn bytes_to_payload(buf: web::Bytes) -> Payload {
382382
Payload::from(pl)
383383
}
384384

385+
#[tracing::instrument(skip_all)]
385386
pub async fn insert_api_key_payload(
386387
req: &mut ServiceRequest,
387388
api_key_params: ApiKeyRequestParams,

‎server/src/middleware/metrics_middleware.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ use actix_web::{
66
web,
77
};
88

9+
#[tracing::instrument(skip_all)]
910
pub async fn error_logging_middleware(
1011
metrics: web::Data<Metrics>,
1112
req: ServiceRequest,

‎server/src/middleware/timeout_middleware.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ use actix_web::{
77
web,
88
};
99

10+
#[tracing::instrument(skip_all)]
1011
pub async fn timeout_15secs(
1112
service_req: ServiceRequest,
1213
next: Next<impl MessageBody + 'static>,

0 commit comments

Comments
 (0)