|
| 1 | +import gzip |
| 2 | +import http |
| 3 | +import json |
| 4 | +import os |
| 5 | +import shutil |
| 6 | +import time |
| 7 | +from dataclasses import asdict, dataclass |
| 8 | +from typing import Any, Callable, Dict, List |
| 9 | + |
| 10 | +import loguru |
| 11 | +import requests |
| 12 | +import speedtest |
| 13 | +from loguru import logger |
| 14 | + |
| 15 | + |
| 16 | +def formatter(record: "loguru.Record") -> str: |
| 17 | + # Note this function returns the string to be formatted, not the actual message to be logged |
| 18 | + record["extra"]["serialized"] = json.dumps(record["message"]) |
| 19 | + return "{extra[serialized]}\n" |
| 20 | + |
| 21 | + |
| 22 | +def is_online() -> bool: |
| 23 | + return get_latency() > 0 |
| 24 | + |
| 25 | + |
| 26 | +def get_latency() -> float: |
| 27 | + try: |
| 28 | + servers: List[str] = [] |
| 29 | + st = speedtest.Speedtest() |
| 30 | + st.get_servers(servers) |
| 31 | + best_server = st.get_best_server() |
| 32 | + ping = best_server["latency"] |
| 33 | + return float(ping) |
| 34 | + except Exception: |
| 35 | + return -1.0 |
| 36 | + |
| 37 | + |
| 38 | +@dataclass |
| 39 | +class DefaultPayload: |
| 40 | + log_session_uuid: str |
| 41 | + order: int |
| 42 | + timestamp: str |
| 43 | + hardware_id: str |
| 44 | + blueos_id: str |
| 45 | + data: dict[str, Any] |
| 46 | + |
| 47 | + def json(self) -> dict[str, Any]: |
| 48 | + return asdict(self) |
| 49 | + |
| 50 | + |
| 51 | +class TelemetryEngine: |
| 52 | + # pylint: disable=too-many-arguments |
| 53 | + def __init__( |
| 54 | + self, |
| 55 | + label: str, |
| 56 | + endpoint: str, |
| 57 | + s3_endpoint: str, |
| 58 | + create_record: Callable[[Any], Any], |
| 59 | + interval: float, |
| 60 | + max_file_size: int, |
| 61 | + max_file_retention: int, |
| 62 | + buffer_folder: str, |
| 63 | + log_buffer: loguru._logger.Logger, # type: ignore |
| 64 | + ): |
| 65 | + self.buffer_file = f"{buffer_folder}/{label}_usage.log" |
| 66 | + self.buffer_folder = buffer_folder |
| 67 | + |
| 68 | + self.telemetry_endpoint = endpoint |
| 69 | + self.telemetry_s3_endpoint = s3_endpoint |
| 70 | + self.create_record = create_record |
| 71 | + self.interval = interval |
| 72 | + |
| 73 | + self.log_buffer = log_buffer |
| 74 | + self.log_buffer.add( |
| 75 | + self.buffer_file, |
| 76 | + rotation=max_file_size, |
| 77 | + retention=max_file_retention, |
| 78 | + format=formatter, |
| 79 | + compression="gz", |
| 80 | + ) |
| 81 | + |
| 82 | + def __call__(self) -> None: |
| 83 | + order = 0 |
| 84 | + while True: |
| 85 | + order += 1 |
| 86 | + record = self.create_record(order) |
| 87 | + if self.save(record) == "online": |
| 88 | + self.process_buffered_records() |
| 89 | + time.sleep(self.interval) |
| 90 | + |
| 91 | + def upload_file(self, file: str) -> bool: |
| 92 | + """ |
| 93 | + This method requests to telemetry API a presigned url and upload the local archived files. |
| 94 | + """ |
| 95 | + logger.info(f"uploading file... {file}") |
| 96 | + try: |
| 97 | + response = requests.get(self.telemetry_s3_endpoint, timeout=5).json() |
| 98 | + with open(file, "rb") as fh: |
| 99 | + files = {"file": (file, fh)} |
| 100 | + r = requests.post(response["url"], data=response["fields"], files=files, timeout=300) |
| 101 | + if r.status_code == http.client.NO_CONTENT: |
| 102 | + logger.info("[Success!]") |
| 103 | + return True |
| 104 | + except Exception as error: |
| 105 | + logger.info("Ground Control to Major Tom. Your circuit's dead, there's something wrong.") |
| 106 | + logger.error(f"error upload log file: {error}") |
| 107 | + |
| 108 | + return False |
| 109 | + |
| 110 | + def process_buffered_records(self) -> None: |
| 111 | + """ |
| 112 | + Check in the buffered folder if there are archived logs to upload. If the agent connects before an archive |
| 113 | + is created it will also archive the current buffer file and upload it. |
| 114 | + """ |
| 115 | + for file in os.listdir(self.buffer_folder): |
| 116 | + file_path = os.path.join(self.buffer_folder, file) |
| 117 | + |
| 118 | + # Upload regular archive |
| 119 | + if file_path.endswith(".log.gz") and self.upload_file(file_path): |
| 120 | + os.remove(file_path) |
| 121 | + continue |
| 122 | + |
| 123 | + # Archive current buffer and upload it |
| 124 | + if file_path == self.buffer_file and os.path.getsize(file_path): |
| 125 | + timestamp = int(time.time()) |
| 126 | + tmp_name = file_path.replace(".log", f".{timestamp}.log.gz") |
| 127 | + with open(file_path, "rb") as f_in, gzip.open(tmp_name, "wb") as f_out: |
| 128 | + shutil.copyfileobj(f_in, f_out) |
| 129 | + os.remove(file_path) |
| 130 | + if self.upload_file(tmp_name): |
| 131 | + os.remove(tmp_name) |
| 132 | + with open(self.buffer_file, "w", encoding="utf-8"): |
| 133 | + # create new empty file if not there |
| 134 | + pass |
| 135 | + |
| 136 | + def save(self, record: Dict[str, Any]) -> str: |
| 137 | + """ |
| 138 | + Try to POST the telemetry payload, if it fails for any reason, we buffer it locally. |
| 139 | + """ |
| 140 | + try: |
| 141 | + r = requests.post(self.telemetry_endpoint, json=record, timeout=5) |
| 142 | + if r.status_code == http.client.CREATED: |
| 143 | + return "online" |
| 144 | + except Exception as error: |
| 145 | + logger.info("Ground Control to Major Tom. Your circuit's dead, there's something wrong.") |
| 146 | + logger.error(f"error posting telemetry to Ground Control: {error}") |
| 147 | + |
| 148 | + self.log_buffer.info(record) |
| 149 | + return "offline" |
0 commit comments