logforth_layout_google_cloud_logging/
lib.rs

1// Copyright 2024 FastLabs Developers
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Layout for Google Cloud Structured Logging.
16
17#![cfg_attr(docsrs, feature(doc_auto_cfg))]
18
19use std::collections::BTreeMap;
20use std::collections::BTreeSet;
21use std::fmt::Arguments;
22
23use logforth_core::Diagnostic;
24use logforth_core::Error;
25use logforth_core::kv::Key;
26use logforth_core::kv::Value;
27use logforth_core::kv::Visitor;
28use logforth_core::layout::Layout;
29use logforth_core::record::Record;
30use serde::Serialize;
31
32/// A layout for Google Cloud Structured Logging.
33///
34/// See the [Google documentation](https://cloud.google.com/logging/docs/structured-logging) for more
35/// information about the structure of the format.
36///
37/// Example format:
38///
39/// ```json
40/// {"severity":"INFO","timestamp":"2025-04-02T10:34:33.225602Z","message":"Hello label value!","logging.googleapis.com/labels":{"label1":"this is a label value"},"logging.googleapis.com/trace":"projects/project-id/traces/612b91406b684ece2c4137ce0f3fd668", "logging.googleapis.com/sourceLocation":{"file":"examples/google_cloud_logging","line":64,"function":"main"}}
41/// ```
42///
43/// If the trace project ID is set, a few keys are treated specially:
44/// - `trace_id`: Combined with trace project ID, set as the `logging.googleapis.com/trace` field.
45/// - `span_id`: Set as the `logging.googleapis.com/spanId` field.
46/// - `trace_sampled`: Set as the `logging.googleapis.com/trace_sampled` field.
47///
48/// Information may be stored either in the payload, or as labels. The payload allows a structured
49/// value to be stored, but is not indexed by default. Labels are indexed by default, but can only
50/// store strings.
51///
52/// # Examples
53///
54/// ```
55/// use logforth_layout_google_cloud_logging::GoogleCloudLoggingLayout;
56///
57/// let layout = GoogleCloudLoggingLayout::default();
58/// ```
59#[derive(Debug, Clone)]
60pub struct GoogleCloudLoggingLayout {
61    trace_project_id: Option<String>,
62    label_keys: BTreeSet<String>,
63
64    // Heuristic keys to extract trace, spanId and traceSampled info from diagnostics.
65    // These are currently hardcoded but may be customizable in the future.
66    trace_keys: BTreeSet<String>,
67    span_id_keys: BTreeSet<String>,
68    trace_sampled_keys: BTreeSet<String>,
69}
70
71impl Default for GoogleCloudLoggingLayout {
72    fn default() -> Self {
73        Self {
74            trace_project_id: None,
75            label_keys: BTreeSet::new(),
76
77            trace_keys: BTreeSet::from(["trace_id".to_string()]),
78            span_id_keys: BTreeSet::from(["span_id".to_string()]),
79            trace_sampled_keys: BTreeSet::from([
80                "sampled".to_string(),
81                "trace_sampled".to_string(),
82            ]),
83        }
84    }
85}
86
87impl GoogleCloudLoggingLayout {
88    /// Set the trace project ID for traces.
89    ///
90    /// If set, the trace_id, span_id, and trace_sampled fields will be set in the log record, in
91    /// such a way that they can be linked to traces in the Google Cloud Trace service.
92    ///
93    /// # Examples
94    ///
95    /// ```
96    /// use logforth_layout_google_cloud_logging::GoogleCloudLoggingLayout;
97    ///
98    /// let layout = GoogleCloudLoggingLayout::default().trace_project_id("project-id");
99    /// ```
100    pub fn trace_project_id(mut self, project_id: impl Into<String>) -> Self {
101        self.trace_project_id = Some(project_id.into());
102        self
103    }
104
105    /// Extends the set of keys that should be treated as labels.
106    ///
107    /// Any key found in a log entry, and referenced here, will be stored in the labels field rather
108    /// than the payload. Labels are indexed by default, but can only store strings.
109    ///
110    /// # Examples
111    ///
112    /// ```
113    /// use logforth_layout_google_cloud_logging::GoogleCloudLoggingLayout;
114    ///
115    /// let layout = GoogleCloudLoggingLayout::default().label_keys(["label1", "label2"]);
116    /// ```
117    pub fn label_keys(mut self, label_keys: impl IntoIterator<Item = impl Into<String>>) -> Self {
118        let label_keys = label_keys.into_iter().map(Into::into);
119        self.label_keys.extend(label_keys);
120        self
121    }
122}
123
124struct KvCollector<'a> {
125    layout: &'a GoogleCloudLoggingLayout,
126
127    payload_fields: BTreeMap<String, serde_json::Value>,
128    labels: BTreeMap<String, serde_json::Value>,
129    trace: Option<String>,
130    span_id: Option<String>,
131    trace_sampled: Option<bool>,
132}
133
134impl Visitor for KvCollector<'_> {
135    fn visit(&mut self, key: Key, value: Value) -> Result<(), Error> {
136        let key = key.into_string();
137
138        if let Some(trace_project_id) = self.layout.trace_project_id.as_ref() {
139            if self.trace.is_none() && self.layout.trace_keys.contains(&key) {
140                self.trace = Some(format!("projects/{trace_project_id}/traces/{value}"));
141                return Ok(());
142            }
143
144            if self.span_id.is_none() && self.layout.span_id_keys.contains(&key) {
145                self.span_id = Some(value.to_string());
146                return Ok(());
147            }
148
149            if self.trace_sampled.is_none() && self.layout.trace_sampled_keys.contains(&key) {
150                self.trace_sampled = value.to_bool();
151                return Ok(());
152            }
153        }
154
155        let value = match serde_json::to_value(&value) {
156            Ok(value) => value,
157            Err(_) => value.to_string().into(),
158        };
159
160        if self.layout.label_keys.contains(&key) {
161            self.labels.insert(key, value);
162        } else {
163            self.payload_fields.insert(key, value);
164        }
165
166        Ok(())
167    }
168}
169
170#[derive(Debug, Clone, Serialize)]
171struct SourceLocation<'a> {
172    #[serde(skip_serializing_if = "Option::is_none")]
173    file: Option<&'a str>,
174    #[serde(skip_serializing_if = "Option::is_none")]
175    line: Option<u32>,
176    #[serde(skip_serializing_if = "Option::is_none")]
177    function: Option<&'a str>,
178}
179
180#[derive(Debug, Clone, Serialize)]
181struct RecordLine<'a> {
182    #[serde(flatten)]
183    extra_fields: BTreeMap<String, serde_json::Value>,
184    severity: &'a str,
185    timestamp: jiff::Timestamp,
186    #[serde(serialize_with = "serialize_args")]
187    message: &'a Arguments<'a>,
188    #[serde(skip_serializing_if = "BTreeMap::is_empty")]
189    #[serde(rename = "logging.googleapis.com/labels")]
190    labels: BTreeMap<String, serde_json::Value>,
191    #[serde(skip_serializing_if = "Option::is_none")]
192    #[serde(rename = "logging.googleapis.com/trace")]
193    trace: Option<String>,
194    #[serde(skip_serializing_if = "Option::is_none")]
195    #[serde(rename = "logging.googleapis.com/spanId")]
196    span_id: Option<String>,
197    #[serde(skip_serializing_if = "Option::is_none")]
198    #[serde(rename = "logging.googleapis.com/trace_sampled")]
199    trace_sampled: Option<bool>,
200    #[serde(skip_serializing_if = "Option::is_none")]
201    #[serde(rename = "logging.googleapis.com/sourceLocation")]
202    source_location: Option<SourceLocation<'a>>,
203}
204
205fn serialize_args<S>(args: &Arguments, serializer: S) -> Result<S::Ok, S::Error>
206where
207    S: serde::Serializer,
208{
209    serializer.collect_str(args)
210}
211
212impl Layout for GoogleCloudLoggingLayout {
213    fn format(&self, record: &Record, diags: &[Box<dyn Diagnostic>]) -> Result<Vec<u8>, Error> {
214        // SAFETY: jiff::Timestamp::try_from only fails if the time is out of range, which is
215        // very unlikely if the system clock is correct.
216        let timestamp = jiff::Timestamp::try_from(record.time()).unwrap();
217
218        let mut visitor = KvCollector {
219            layout: self,
220            payload_fields: BTreeMap::new(),
221            labels: BTreeMap::new(),
222            trace: None,
223            span_id: None,
224            trace_sampled: None,
225        };
226
227        record.key_values().visit(&mut visitor)?;
228        for d in diags {
229            d.visit(&mut visitor)?;
230        }
231
232        let record_line = RecordLine {
233            extra_fields: visitor.payload_fields,
234            timestamp,
235            severity: record.level().as_str(),
236            message: record.args(),
237            labels: visitor.labels,
238            trace: visitor.trace,
239            span_id: visitor.span_id,
240            trace_sampled: visitor.trace_sampled,
241            source_location: Some(SourceLocation {
242                file: record.file(),
243                line: record.line(),
244                function: record.module_path(),
245            }),
246        };
247
248        // SAFETY: RecordLine is serializable.
249        Ok(serde_json::to_vec(&record_line).unwrap())
250    }
251}