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};
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<HashMap<String, Value>>,
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 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 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}