featherbit/plugins/native/
loki_logger.rs1use 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, LogFormat};
26use crate::plugins::{Plugin, PluginOutput, PluginResult};
27
28fn 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
47struct LokiLoggerFlusher {
49 client: Arc<OutboundClient>,
50 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 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
100pub struct LokiLoggerPlugin {
102 sink: BatchSink,
103 log_format: Option<LogFormat>,
104 include_req_body: bool,
105 include_resp_body: bool,
106}
107
108impl LokiLoggerPlugin {
109 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
230fn 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 fn reads_response_body(&self) -> bool {
272 crate::plugins::util::log_entry::reads_response_body(
273 self.log_format.as_ref(),
274 self.include_resp_body,
275 )
276 }
277
278 async fn execute(&self, ctx: Context) -> PluginResult {
279 let entry = build_entry(
280 &ctx,
281 self.log_format.as_ref(),
282 self.include_req_body,
283 self.include_resp_body,
284 );
285 self.sink.push(entry);
286 Ok(PluginOutput::success(ctx))
287 }
288}
289
290#[cfg(test)]
291mod tests {
292 use super::*;
293 use serde_json::json;
294
295 fn cfg(v: Value) -> HashMap<String, Value> {
296 serde_json::from_value(v).unwrap()
297 }
298
299 #[test]
300 fn rejects_missing_endpoint() {
301 assert!(LokiLoggerPlugin::from_config(&HashMap::new(), &PluginResources::empty()).is_err());
302 }
303
304 #[tokio::test]
305 async fn accepts_endpoint_addrs() {
306 let c = cfg(json!({ "endpoint_addrs": ["http://loki:3100"] }));
307 assert!(LokiLoggerPlugin::from_config(&c, &PluginResources::empty()).is_ok());
308 }
309
310 #[test]
311 fn default_labels_are_featherbit() {
312 let labels = parse_labels(&HashMap::new());
313 assert_eq!(
314 labels.get("job"),
315 Some(&Value::String("featherbit".to_string()))
316 );
317 }
318
319 #[test]
320 fn parses_custom_labels() {
321 let c = cfg(json!({ "log_labels": { "job": "gw", "env": "prod" } }));
322 let labels = parse_labels(&c);
323 assert_eq!(labels.get("job"), Some(&json!("gw")));
324 assert_eq!(labels.get("env"), Some(&json!("prod")));
325 }
326
327 #[test]
328 fn payload_has_one_stream_and_value_per_entry() {
329 let mut labels = Map::new();
330 labels.insert("job".to_string(), json!("featherbit"));
331 let entries = vec![json!({"n": 1}), json!({"n": 2})];
332 let payload = build_loki_payload(&entries, &labels);
333
334 let streams = payload["streams"].as_array().unwrap();
335 assert_eq!(streams.len(), 1);
336 assert_eq!(streams[0]["stream"]["job"], "featherbit");
337
338 let values = streams[0]["values"].as_array().unwrap();
339 assert_eq!(values.len(), 2);
340 assert!(values[0][0]
342 .as_str()
343 .unwrap()
344 .chars()
345 .all(|c| c.is_ascii_digit()));
346 let line: Value = serde_json::from_str(values[1][1].as_str().unwrap()).unwrap();
347 assert_eq!(line["n"], 2);
348 }
349}