Skip to main content

featherbit/plugins/native/
loki_logger.rs

1//! The `loki-logger` node — ships access logs to a Grafana Loki push API in
2//! batches.
3//!
4//! Ports the APISIX `loki-logger` plugin onto featherbit's shared
5//! [`BatchSink`](crate::batch::BatchSink). Each request builds a log entry and
6//! hands it to the sink with a non-blocking `push`; a background task groups a
7//! batch into a single Loki stream (the configured labels) and POSTs the
8//! `{streams:[{stream, values:[[ns_ts, line], ...]}]}` payload to
9//! `<endpoint>/loki/api/v1/push`. The node never mutates the request/response
10//! and never fails, so it belongs in the response pipeline, **after the
11//! upstream node**.
12
13use std::collections::HashMap;
14use std::sync::Arc;
15use std::time::{Duration, SystemTime, UNIX_EPOCH};
16
17use async_trait::async_trait;
18use bytes::Bytes;
19use serde_json::{json, Map, Value};
20
21use crate::batch::{BatchConfig, BatchFlusher, BatchSink, FlushError};
22use crate::context::Context;
23use crate::outbound::{OutboundClient, OutboundRequest};
24use crate::plugins::resources::PluginResources;
25use crate::plugins::util::log_entry::{build_entry, parse_log_format};
26use crate::plugins::{Plugin, PluginOutput, PluginResult};
27
28/// Builds the Loki push payload: one stream carrying `labels`, with one
29/// `[ns_timestamp, json_line]` value per entry. Timestamps are the current
30/// wall-clock time in nanoseconds (as a decimal string, per the Loki API).
31/// Factored out of the network path for unit testing.
32fn build_loki_payload(entries: &[Value], labels: &Map<String, Value>) -> Value {
33    let ts = SystemTime::now()
34        .duration_since(UNIX_EPOCH)
35        .map(|d| d.as_nanos())
36        .unwrap_or(0)
37        .to_string();
38    let values: Vec<Value> = entries
39        .iter()
40        .map(|e| json!([ts, serde_json::to_string(e).unwrap_or_default()]))
41        .collect();
42    json!({
43        "streams": [ { "stream": Value::Object(labels.clone()), "values": values } ]
44    })
45}
46
47/// Delivers batches by POSTing the Loki streams payload to a push endpoint.
48struct LokiLoggerFlusher {
49    client: Arc<OutboundClient>,
50    /// Fully-qualified push URLs (`endpoint_addr + endpoint_uri`).
51    urls: Vec<String>,
52    tenant_id: String,
53    labels: Map<String, Value>,
54    extra_headers: Vec<(String, String)>,
55    ssl_verify: bool,
56    timeout: Duration,
57}
58
59#[async_trait]
60impl BatchFlusher for LokiLoggerFlusher {
61    async fn flush(&self, entries: &[Value]) -> Result<(), FlushError> {
62        let payload = build_loki_payload(entries, &self.labels);
63        let body = Bytes::from(serde_json::to_vec(&payload).unwrap_or_default());
64
65        // Pick an endpoint pseudo-randomly to spread load across replicas.
66        let idx = (SystemTime::now()
67            .duration_since(UNIX_EPOCH)
68            .map(|d| d.subsec_nanos())
69            .unwrap_or(0) as usize)
70            % self.urls.len().max(1);
71        let url = self.urls.get(idx).cloned().unwrap_or_default();
72
73        let mut headers = self.extra_headers.clone();
74        headers.push(("X-Scope-OrgID".to_string(), self.tenant_id.clone()));
75        headers.push(("Content-Type".to_string(), "application/json".to_string()));
76
77        let req = OutboundRequest {
78            method: http::Method::POST,
79            url,
80            headers,
81            body,
82            timeout: self.timeout,
83            ssl_verify: self.ssl_verify,
84            tls: None,
85        };
86        match self.client.request(req).await {
87            Ok(resp) if resp.status < 300 => Ok(()),
88            Ok(resp) => Err(FlushError {
89                message: format!("loki server returned status {}", resp.status),
90                first_fail: None,
91            }),
92            Err(e) => Err(FlushError {
93                message: e.to_string(),
94                first_fail: None,
95            }),
96        }
97    }
98}
99
100/// Batches access-log entries and POSTs them to a Grafana Loki push endpoint.
101pub struct LokiLoggerPlugin {
102    sink: BatchSink,
103    log_format: Option<HashMap<String, Value>>,
104    include_req_body: bool,
105    include_resp_body: bool,
106}
107
108impl LokiLoggerPlugin {
109    /// Builds the plugin from node config.
110    ///
111    /// Accepted keys:
112    /// - `endpoint_addrs` (array of strings, **required**): Loki base
113    ///   addresses (e.g. `http://loki:3100`); one is chosen per flush. A
114    ///   single-string `endpoint` is also accepted. Empty/absent is an error.
115    /// - `endpoint_uri` (string, default `/loki/api/v1/push`): push path
116    ///   appended to each base address.
117    /// - `tenant_id` (string, default `fake`): sent as the `X-Scope-OrgID`
118    ///   header.
119    /// - `log_labels` (object of string→string, default `{job: featherbit}`):
120    ///   Loki stream labels. `labels` is accepted as an alias.
121    /// - `headers` (object of string→string, optional): extra request headers.
122    /// - `ssl_verify` (bool, default `false`): verify TLS certificates.
123    /// - `timeout` (integer ms, default `3000`): whole-call deadline per flush.
124    /// - `log_format`, `include_req_body`, `include_resp_body` — see the
125    ///   shared log-entry builder.
126    /// - Batch keys — see [`BatchConfig`].
127    ///
128    /// ```yaml
129    /// type: loki-logger
130    /// config:
131    ///   endpoint_addrs: [http://loki:3100]
132    ///   tenant_id: my-org
133    ///   log_labels:
134    ///     job: featherbit
135    ///     env: prod
136    /// ```
137    pub fn from_config(
138        config: &HashMap<String, Value>,
139        resources: &Arc<PluginResources>,
140    ) -> Result<Self, String> {
141        let mut addrs: Vec<String> = Vec::new();
142        if let Some(arr) = config.get("endpoint_addrs").and_then(|v| v.as_array()) {
143            for v in arr {
144                if let Some(s) = v.as_str().filter(|s| !s.is_empty()) {
145                    addrs.push(s.trim_end_matches('/').to_string());
146                }
147            }
148        }
149        if let Some(s) = config
150            .get("endpoint")
151            .and_then(|v| v.as_str())
152            .filter(|s| !s.is_empty())
153        {
154            addrs.push(s.trim_end_matches('/').to_string());
155        }
156        if addrs.is_empty() {
157            return Err("loki-logger plugin requires 'endpoint_addrs'".to_string());
158        }
159
160        let endpoint_uri = config
161            .get("endpoint_uri")
162            .and_then(|v| v.as_str())
163            .filter(|s| !s.is_empty())
164            .unwrap_or("/loki/api/v1/push");
165        let urls: Vec<String> = addrs
166            .iter()
167            .map(|a| format!("{}{}", a, endpoint_uri))
168            .collect();
169
170        let tenant_id = config
171            .get("tenant_id")
172            .and_then(|v| v.as_str())
173            .unwrap_or("fake")
174            .to_string();
175
176        let labels = parse_labels(config);
177
178        let mut extra_headers: Vec<(String, String)> = Vec::new();
179        if let Some(obj) = config.get("headers").and_then(|v| v.as_object()) {
180            for (k, v) in obj {
181                if let Some(s) = v.as_str() {
182                    extra_headers.push((k.clone(), s.to_string()));
183                }
184            }
185        }
186
187        let ssl_verify = config
188            .get("ssl_verify")
189            .and_then(|v| v.as_bool())
190            .unwrap_or(false);
191        let timeout = Duration::from_millis(
192            config
193                .get("timeout")
194                .and_then(|v| v.as_u64())
195                .unwrap_or(3000)
196                .max(1),
197        );
198
199        let log_format = parse_log_format(config)?;
200        let include_req_body = config
201            .get("include_req_body")
202            .and_then(|v| v.as_bool())
203            .unwrap_or(false);
204        let include_resp_body = config
205            .get("include_resp_body")
206            .and_then(|v| v.as_bool())
207            .unwrap_or(false);
208
209        let batch_cfg = BatchConfig::from_config(config)?;
210        let flusher = Arc::new(LokiLoggerFlusher {
211            client: resources.outbound.clone(),
212            urls,
213            tenant_id,
214            labels,
215            extra_headers,
216            ssl_verify,
217            timeout,
218        });
219        let sink = BatchSink::spawn("loki-logger", batch_cfg, flusher);
220
221        Ok(Self {
222            sink,
223            log_format,
224            include_req_body,
225            include_resp_body,
226        })
227    }
228}
229
230/// Reads `log_labels` (or `labels`) as a string map, defaulting to
231/// `{job: featherbit}`.
232fn parse_labels(config: &HashMap<String, Value>) -> Map<String, Value> {
233    let src = config
234        .get("log_labels")
235        .or_else(|| config.get("labels"))
236        .and_then(|v| v.as_object());
237    match src {
238        Some(obj) => {
239            let mut m = Map::new();
240            for (k, v) in obj {
241                let val = match v {
242                    Value::String(s) => s.clone(),
243                    Value::Number(n) => n.to_string(),
244                    Value::Bool(b) => b.to_string(),
245                    _ => continue,
246                };
247                m.insert(k.clone(), Value::String(val));
248            }
249            if m.is_empty() {
250                default_labels()
251            } else {
252                m
253            }
254        }
255        None => default_labels(),
256    }
257}
258
259fn default_labels() -> Map<String, Value> {
260    let mut m = Map::new();
261    m.insert("job".to_string(), Value::String("featherbit".to_string()));
262    m
263}
264
265#[async_trait]
266impl Plugin for LokiLoggerPlugin {
267    fn plugin_type(&self) -> &str {
268        "loki-logger"
269    }
270
271    async fn execute(&self, ctx: Context, _named_inputs: &HashMap<String, Value>) -> PluginResult {
272        let entry = build_entry(
273            &ctx,
274            self.log_format.as_ref(),
275            self.include_req_body,
276            self.include_resp_body,
277        );
278        self.sink.push(entry);
279        Ok(PluginOutput {
280            context: ctx,
281            named_outputs: HashMap::new(),
282        })
283    }
284}
285
286#[cfg(test)]
287mod tests {
288    use super::*;
289    use serde_json::json;
290
291    fn cfg(v: Value) -> HashMap<String, Value> {
292        serde_json::from_value(v).unwrap()
293    }
294
295    #[test]
296    fn rejects_missing_endpoint() {
297        assert!(LokiLoggerPlugin::from_config(&HashMap::new(), &PluginResources::empty()).is_err());
298    }
299
300    #[tokio::test]
301    async fn accepts_endpoint_addrs() {
302        let c = cfg(json!({ "endpoint_addrs": ["http://loki:3100"] }));
303        assert!(LokiLoggerPlugin::from_config(&c, &PluginResources::empty()).is_ok());
304    }
305
306    #[test]
307    fn default_labels_are_featherbit() {
308        let labels = parse_labels(&HashMap::new());
309        assert_eq!(
310            labels.get("job"),
311            Some(&Value::String("featherbit".to_string()))
312        );
313    }
314
315    #[test]
316    fn parses_custom_labels() {
317        let c = cfg(json!({ "log_labels": { "job": "gw", "env": "prod" } }));
318        let labels = parse_labels(&c);
319        assert_eq!(labels.get("job"), Some(&json!("gw")));
320        assert_eq!(labels.get("env"), Some(&json!("prod")));
321    }
322
323    #[test]
324    fn payload_has_one_stream_and_value_per_entry() {
325        let mut labels = Map::new();
326        labels.insert("job".to_string(), json!("featherbit"));
327        let entries = vec![json!({"n": 1}), json!({"n": 2})];
328        let payload = build_loki_payload(&entries, &labels);
329
330        let streams = payload["streams"].as_array().unwrap();
331        assert_eq!(streams.len(), 1);
332        assert_eq!(streams[0]["stream"]["job"], "featherbit");
333
334        let values = streams[0]["values"].as_array().unwrap();
335        assert_eq!(values.len(), 2);
336        // each value is [ts_string, json_line_string]
337        assert!(values[0][0]
338            .as_str()
339            .unwrap()
340            .chars()
341            .all(|c| c.is_ascii_digit()));
342        let line: Value = serde_json::from_str(values[1][1].as_str().unwrap()).unwrap();
343        assert_eq!(line["n"], 2);
344    }
345}