Skip to main content

relay_knowledge/application/
service_health.rs

1use std::{sync::Arc, time::Duration};
2
3use crate::{
4    api::{ApiError, ApiMetadata, HealthResponse, RequestContext},
5    domain::{CodeRepositoryTotals, GraphVersion},
6    storage::{
7        FileIndexDiagnostics, GraphInspection, HealthStorageSnapshot, IndexRefreshDiagnostics,
8        KnowledgeStore, StorageError,
9    },
10};
11
12use super::{
13    RelayKnowledgeService,
14    index_refresh::{IndexRefreshOutcome, filter_outcome_to_read_models, metadata_for_indexes},
15    service::{current_time_millis, graph_with_repository_code_totals, storage_api_error},
16    status::runtime_status_with_model_profiles,
17};
18
19const HEALTH_STORAGE_BUDGET: Duration = Duration::from_millis(500);
20
21impl RelayKnowledgeService {
22    /// Returns liveness-safe service and data health diagnostics.
23    pub async fn health(&self, context: RequestContext) -> Result<HealthResponse, ApiError> {
24        let store = self.storage.get().await.map_err(storage_api_error)?;
25        match tokio::time::timeout(HEALTH_STORAGE_BUDGET, self.storage_health_snapshot(&store))
26            .await
27        {
28            Ok(Ok(snapshot)) => {
29                let response = self
30                    .health_from_storage_snapshot(context, snapshot, None)
31                    .await;
32                *self.health_cache.write().await = Some(response.clone());
33                Ok(response)
34            }
35            Ok(Err(StorageError::Busy(message))) => Ok(self
36                .degraded_cached_health(context, format!("storage_busy: {message}"))
37                .await),
38            Ok(Err(error)) => Err(storage_api_error(error)),
39            Err(_) => Ok(self
40                .degraded_cached_health(context, "storage_busy: health snapshot timed out")
41                .await),
42        }
43    }
44
45    async fn storage_health_snapshot(
46        &self,
47        store: &Arc<dyn KnowledgeStore>,
48    ) -> Result<HealthStorageSnapshot, StorageError> {
49        match store.health_snapshot(current_time_millis()).await {
50            Ok(snapshot) => Ok(snapshot),
51            Err(StorageError::InvalidInput(message))
52                if message == "health snapshot storage is unavailable" =>
53            {
54                self.legacy_health_snapshot(store).await
55            }
56            Err(error) => Err(error),
57        }
58    }
59
60    async fn health_from_storage_snapshot(
61        &self,
62        context: RequestContext,
63        snapshot: HealthStorageSnapshot,
64        degraded_reason: Option<String>,
65    ) -> HealthResponse {
66        let HealthStorageSnapshot {
67            graph,
68            repository_code_totals,
69            indexes,
70            index_cursors,
71            index_refresh,
72            file_index,
73        } = snapshot;
74        let graph = graph_with_repository_code_totals(graph, &repository_code_totals);
75        let outcome = filter_outcome_to_read_models(
76            IndexRefreshOutcome {
77                indexes,
78                cursors: index_cursors,
79                diagnostics: index_refresh,
80            },
81            &self.runtime.retrieval,
82        );
83        let healthy = degraded_reason.is_none()
84            && outcome
85                .indexes
86                .iter()
87                .all(|status| !status.is_stale_for(graph.graph_version));
88
89        HealthResponse {
90            metadata: metadata_for_indexes(&context, graph.graph_version, &outcome.indexes),
91            healthy,
92            degraded_reason,
93            graph,
94            repository_code_totals,
95            indexes: outcome.indexes,
96            index_cursors: outcome.cursors,
97            index_refresh: outcome.diagnostics,
98            file_index,
99            runtime: runtime_status_with_model_profiles(
100                &self.runtime,
101                self.model_provider_config()
102                    .profile_summary(&self.runtime.retrieval)
103                    .await,
104            ),
105        }
106    }
107
108    async fn legacy_health_snapshot(
109        &self,
110        store: &Arc<dyn KnowledgeStore>,
111    ) -> Result<HealthStorageSnapshot, StorageError> {
112        Ok(HealthStorageSnapshot {
113            graph: store.inspect_graph().await?,
114            repository_code_totals: store.code_repository_totals().await?,
115            indexes: store.index_statuses().await?,
116            index_cursors: store.index_cursors().await?,
117            index_refresh: store
118                .index_refresh_diagnostics(current_time_millis())
119                .await?,
120            file_index: legacy_file_index_diagnostics_or_default(store).await?,
121        })
122    }
123
124    async fn degraded_cached_health(
125        &self,
126        context: RequestContext,
127        degraded_reason: impl Into<String>,
128    ) -> HealthResponse {
129        let degraded_reason = degraded_reason.into();
130        if let Some(cached) = self.health_cache.read().await.clone() {
131            let mut response = cached;
132            response.metadata.trace_id = context.trace_id;
133            response.metadata.request_id = context.request_id;
134            response.metadata.stale = true;
135            response.healthy = false;
136            response.degraded_reason = Some(degraded_reason);
137            return response;
138        }
139
140        HealthResponse {
141            metadata: ApiMetadata::indexed(&context, GraphVersion::ZERO, None, None, true),
142            healthy: false,
143            degraded_reason: Some(degraded_reason),
144            graph: GraphInspection::default(),
145            repository_code_totals: CodeRepositoryTotals::default(),
146            indexes: Vec::new(),
147            index_cursors: Vec::new(),
148            index_refresh: IndexRefreshDiagnostics::default(),
149            file_index: FileIndexDiagnostics::default(),
150            runtime: runtime_status_with_model_profiles(
151                &self.runtime,
152                self.model_provider_config()
153                    .profile_summary(&self.runtime.retrieval)
154                    .await,
155            ),
156        }
157    }
158}
159
160async fn legacy_file_index_diagnostics_or_default(
161    store: &Arc<dyn KnowledgeStore>,
162) -> Result<FileIndexDiagnostics, StorageError> {
163    match store.file_index_diagnostics().await {
164        Ok(diagnostics) => Ok(diagnostics),
165        Err(StorageError::InvalidInput(message))
166            if message == "file index storage is unavailable" =>
167        {
168            Ok(FileIndexDiagnostics::default())
169        }
170        Err(error) => Err(error),
171    }
172}