143 lines
4.6 KiB
Python
143 lines
4.6 KiB
Python
"""Logging setup for the MQTT client.
|
|
|
|
Supports:
|
|
- Console logging (always).
|
|
- File logging (optional, controlled by LOG_TO_FILE).
|
|
- Loki push (optional, controlled by LOKI_URL). Loki pushes are batched
|
|
and sent from a background thread so a slow or unavailable Loki never
|
|
blocks ingestion.
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
import queue
|
|
import sys
|
|
import threading
|
|
from logging.handlers import RotatingFileHandler
|
|
from typing import Optional
|
|
|
|
import requests
|
|
from pythonjsonlogger import jsonlogger
|
|
|
|
LOKI_URL = os.getenv("LOKI_URL", "").strip()
|
|
LOKI_USER = os.getenv("LOKI_USER", "")
|
|
LOKI_PASSWORD = os.getenv("LOKI_PASSWORD", "")
|
|
LOKI_FLUSH_SECONDS = float(os.getenv("LOKI_FLUSH_SECONDS", "5"))
|
|
LOKI_TIMEOUT_SECONDS = float(os.getenv("LOKI_TIMEOUT_SECONDS", "5"))
|
|
|
|
|
|
class LokiHandler(logging.Handler):
|
|
"""Buffer log records and push them to Loki in batches.
|
|
|
|
``emit()`` only enqueues. A daemon thread flushes the buffer every
|
|
``flush_interval`` seconds, grouping records by level into one stream
|
|
per level so a single HTTP request carries many entries.
|
|
"""
|
|
|
|
_MAX_BATCH = 500
|
|
|
|
def __init__(
|
|
self,
|
|
url: str,
|
|
user: str = "",
|
|
password: str = "",
|
|
flush_interval: Optional[float] = None,
|
|
timeout: Optional[float] = None,
|
|
) -> None:
|
|
super().__init__()
|
|
self.url = url.rstrip("/") + "/loki/api/v1/push"
|
|
self.auth = (user, password) if user else None
|
|
self._flush_interval = (
|
|
LOKI_FLUSH_SECONDS if flush_interval is None else flush_interval
|
|
)
|
|
self._timeout = LOKI_TIMEOUT_SECONDS if timeout is None else timeout
|
|
self._buffer: queue.SimpleQueue[logging.LogRecord] = queue.SimpleQueue()
|
|
self._stop = threading.Event()
|
|
self._thread = threading.Thread(
|
|
target=self._publish_loop, name="loki-publisher", daemon=True
|
|
)
|
|
self._thread.start()
|
|
|
|
def emit(self, record: logging.LogRecord) -> None:
|
|
self._buffer.put(record)
|
|
|
|
def _publish_loop(self) -> None:
|
|
while not self._stop.wait(self._flush_interval):
|
|
self.flush()
|
|
|
|
def flush(self) -> None:
|
|
records: list[logging.LogRecord] = []
|
|
while len(records) < self._MAX_BATCH:
|
|
try:
|
|
records.append(self._buffer.get_nowait())
|
|
except queue.Empty:
|
|
break
|
|
if not records:
|
|
return
|
|
|
|
streams: dict[str, list[list[str]]] = {}
|
|
for record in records:
|
|
level = record.levelname.lower()
|
|
streams.setdefault(level, []).append(
|
|
[str(int(record.created * 1_000_000_000)), self.format(record)]
|
|
)
|
|
# Loki requires increasing timestamps within a stream.
|
|
for values in streams.values():
|
|
values.sort(key=lambda value: int(value[0]))
|
|
|
|
payload = {
|
|
"streams": [
|
|
{
|
|
"stream": {"service": "mqtt-client", "level": level},
|
|
"values": values,
|
|
}
|
|
for level, values in streams.items()
|
|
]
|
|
}
|
|
try:
|
|
response = requests.post(
|
|
self.url, json=payload, auth=self.auth, timeout=self._timeout
|
|
)
|
|
response.raise_for_status()
|
|
except Exception: # noqa: BLE001 - never let logging break ingestion
|
|
pass
|
|
|
|
def close(self) -> None:
|
|
self._stop.set()
|
|
self.flush()
|
|
super().close()
|
|
|
|
|
|
def setup_logging() -> logging.Logger:
|
|
level = os.getenv("LOG_LEVEL", "INFO").upper()
|
|
logger = logging.getLogger("mqtt-client")
|
|
logger.setLevel(level)
|
|
logger.handlers.clear()
|
|
|
|
formatter = jsonlogger.JsonFormatter(
|
|
"%(asctime)s %(levelname)s %(name)s %(message)s"
|
|
)
|
|
|
|
console = logging.StreamHandler(sys.stdout)
|
|
console.setFormatter(formatter)
|
|
logger.addHandler(console)
|
|
|
|
if os.getenv("LOG_TO_FILE", "true").lower() == "true":
|
|
path = os.getenv("LOG_FILE_PATH", "/var/log/mqtt-client/mqtt-client.log")
|
|
try:
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
file_handler = RotatingFileHandler(
|
|
path, maxBytes=10 * 1024 * 1024, backupCount=5
|
|
)
|
|
file_handler.setFormatter(formatter)
|
|
logger.addHandler(file_handler)
|
|
except OSError as exc:
|
|
# Logging must never prevent startup (e.g. root-owned or
|
|
# read-only log directory).
|
|
logger.warning("File logging disabled: %s", exc)
|
|
|
|
if LOKI_URL:
|
|
logger.addHandler(LokiHandler(LOKI_URL, LOKI_USER, LOKI_PASSWORD))
|
|
|
|
return logger
|