Merge pull request #77 from KaminariOS/log_dedup

gaganso · web-flow · commit e5fe70e96422 · 2025-10-07T00:44:59.000-06:00
Add log dedup logic
diff --git a/aiopslab/orchestrator/actions/base.py b/aiopslab/orchestrator/actions/base.py
@@ -15,6 +15,16 @@
 from aiopslab.observer.metric_api import PrometheusAPI
 from aiopslab.observer.trace_api import TraceAPI
 
+from aiopslab.orchestrator.actions.log_deduplication import greedy_compress_lines 
+
+import re
+
+LOG_COMMAND_PATTERN: str = (
+    r"\b(?:"
+    r"kubectl\s+(?:logs|get\s+events|describe|get\s+\S+\s+-w)"  # logs/events/describe/watch
+    r"|docker\s+(?:logs|events)"                               # docker logs/events
+    r")\b(?:[^\n]*)"
+)
 
 class TaskActions:
     """Base class for task actions."""
@@ -60,6 +70,7 @@ def get_logs(namespace: str, service: str) -> str:
             except Exception as e:
                 return "Error: Your service/namespace does not exist. Use kubectl to check."
 
+        logs = greedy_compress_lines(logs) 
         print(logs)
 
         return logs
@@ -90,7 +101,14 @@ def exec_shell(command: str, timeout: int = 30) -> str:
             if pattern in command:
                 return error
 
-        return Shell.exec(command, timeout=timeout)
+        result = Shell.exec(command) 
+
+        if re.search(LOG_COMMAND_PATTERN, command):
+            result = greedy_compress_lines(result)
+
+        print(result)
+
+        return result
 
     @staticmethod
     @read
diff --git a/aiopslab/orchestrator/actions/log_deduplication.py b/aiopslab/orchestrator/actions/log_deduplication.py
@@ -0,0 +1,173 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+import re
+import sys
+import argparse
+from pathlib import Path
+import os
+
+# -------------------------------
+# Timestamp detection
+# -------------------------------
+DEFAULT_TIMESTAMP_REGEX = (
+    r"(?:"
+    # ISO-like: 2025-09-24 18:41:09 or 2025-09-24T18:41:09Z
+    r"\d{4}-\d{2}-\d{2}(?:[ T]\d{2}:\d{2}:\d{2}(?:\.\d+)?(?:Z|[+-]\d{2}:\d{2})?)"
+    r"|"
+    # Abbreviated month: 2025-Sep-24 18:41:09.830218
+    r"\d{4}-[A-Z][a-z]{2}-\d{2} \d{2}:\d{2}:\d{2}(?:\.\d+)?"
+    r"|"
+    # Time only: 18:41:09
+    r"\b\d{2}:\d{2}:\d{2}\b"
+    r"|"
+    # Durations: 10m5s, 30s
+    r"\b\d+m(?:\d+s)?\b"
+    r"|"
+    r"\b\d+s\b"
+    r"|"
+    # syslog-like: Wed Sep 24 18:41:09 2025
+    r"[A-Z][a-z]{2} [A-Z][a-z]{2} \d{2} \d{2}:\d{2}:\d{2} \d{4}"
+    r")"
+)
+
+DEFAULT_TS_RX = re.compile(DEFAULT_TIMESTAMP_REGEX)
+
+def find_timestamp_spans(line: str, ts_rx: re.Pattern[str]) -> list[tuple[int, int]]:
+    return [m.span() for m in ts_rx.finditer(line)]
+
+def make_blocks(lines: list[str], block_size: int) -> list[str]:
+    if block_size <= 0:
+        raise ValueError("block_size must be positive")
+    return ["\n".join(lines[i:i+block_size]) for i in range(0, len(lines), block_size)]
+
+# -------------------------------
+# Greedy compressor (single pass)
+# -------------------------------
+def greedy_compress_pass(
+    lines: list[str], 
+    ts_rx: re.Pattern[str], 
+    block_size: int
+) -> list[str]:
+    """Run greedy timestamp dedup for a single block size."""
+    if not lines:
+        return []
+
+    blocks = make_blocks(lines, block_size)
+    result: list[str] = [blocks[0]]
+    prev_spans: list[tuple[int, int]] | None = find_timestamp_spans(blocks[0], ts_rx)
+
+    for block in blocks[1:]:
+        spans = find_timestamp_spans(block, ts_rx)
+
+        if not prev_spans or not spans:
+            result.append(block)
+            prev_spans = spans
+            continue
+
+        if len(spans) != len(prev_spans):
+            result.append(block)
+            prev_spans = spans
+            continue
+
+        if [s[0] for s in spans] != [s[0] for s in prev_spans]:
+            result.append(block)
+            prev_spans = spans
+            continue
+
+        def mask_timestamps(text: str, spans: list[tuple[int, int]]) -> str:
+            parts: list[str] = []
+            last_end = 0
+            for start, end in spans:
+                parts.append(text[last_end:start])
+                # Replace timestamp with spaces of equal length
+                parts.append(" " * (end - start))
+                last_end = end
+            parts.append(text[last_end:])
+            return "".join(parts)
+
+        prev_masked = mask_timestamps(result[-1], prev_spans)
+        curr_masked = mask_timestamps(block, spans)
+
+        if prev_masked == curr_masked:
+            # Replace last block if only timestamps differ
+            result[-1] = block
+        else:
+            # print(f"prev: {prev_masked}\ncurr: {curr_masked}")
+            result.append(block)
+
+        prev_spans = spans
+
+    return result
+
+# -------------------------------
+# Multi-pass driver
+# -------------------------------
+def greedy_compress_lines(
+    raw_str: str, 
+    ts_rx: re.Pattern[str] = DEFAULT_TS_RX, 
+) -> str:
+    """
+    Run greedy dedup with passes from block_size=1 up to block_size=window_size.
+    window_size = LOG_TRIM if LOG_TRIM is int or trimming is disabled  
+    """
+    log_trim = None
+    try:
+        value = os.environ.get("LOG_TRIM")
+        log_trim = int(value) if value is not None else None
+    except ValueError:
+        log_trim = None
+    if not log_trim:
+        return raw_str
+    window_size = log_trim
+    lines = raw_str.splitlines()
+    result = lines[:]
+    for size in range(1, window_size + 1):
+        result = greedy_compress_pass(result, ts_rx, size)
+    return "\n".join(result)
+
+# -------------------------------
+# CLI
+# -------------------------------
+def main():
+    ap = argparse.ArgumentParser(
+        description="Greedy log deduplicator with multi-pass window size support."
+    )
+    ap.add_argument("input", help="Path to input text file")
+    ap.add_argument("output", help="Path to output text file")
+    ap.add_argument(
+        "--timestamp-regex",
+        default=DEFAULT_TIMESTAMP_REGEX,
+        help="Regex for timestamps. Default matches ISO, RFC, and k8s durations.",
+    )
+    ap.add_argument(
+        "--window-size",
+        type=int,
+        default=2,
+        help="Maximum block size for multi-pass deduplication. Default: 2",
+    )
+    args = ap.parse_args()
+
+    in_path = Path(args.input)
+    out_path = Path(args.output)
+
+    if not in_path.exists():
+        print(f"Input file not found: {in_path}", file=sys.stderr)
+        sys.exit(1)
+
+    raw_text = in_path.read_text(encoding="utf-8", errors="replace")
+    lines = raw_text
+
+    try:
+        ts_rx = re.compile(args.timestamp_regex)
+    except re.error as e:
+        print(f"Invalid timestamp regex: {e}", file=sys.stderr)
+        sys.exit(1)
+
+    deduped = greedy_compress_lines(lines, ts_rx, args.window_size)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    out_path.write_text(deduped, encoding="utf-8")
+
+if __name__ == "__main__":
+    main()
+
diff --git a/aiopslab/service/shell.py b/aiopslab/service/shell.py
@@ -55,11 +55,10 @@ def local_exec(command: str, input_data=None, cwd=None, timeout=30):
 
             if out.returncode != 0:
                 error_message = out.stderr.decode("utf-8")
-                print(f"[ERROR] Command execution failed: {error_message}")
+                error_message = f"[ERROR] Command execution failed: {error_message}"
                 return error_message
             else:
                 output_message = out.stdout.decode("utf-8") + out.stderr.decode("utf-8")
-                print(output_message)
                 return output_message
 
         except Exception as e:
@@ -89,7 +88,6 @@ def ssh_exec(host: str, user: str, ssh_key_path: str, command: str, timeout=30):
                 return error_message
             else:
                 output_message = stdout.read().decode("utf-8")
-                print(output_message)
                 return output_message
 
         except Exception as e:
@@ -116,12 +114,10 @@ def docker_exec(container_name: str, command: str, timeout=30):
 
             if out.stderr or out.returncode != 0:
                 error_message = out.stderr.decode("utf-8")
-                print(f"[ERROR] Docker command execution failed: {error_message}")
-                return error_message
+                return f"[ERROR] Docker command execution failed: {error_message}"
             else:
                 output_message = out.stdout.decode("utf-8")
-                print(output_message)
                 return output_message
 
         except Exception as e:
-            raise RuntimeError(f"Failed to execute command in Docker container: {container_name}\nError: {str(e)}")
+            raise RuntimeError(f"Failed to execute command in Docker container: {container_name}\nError: {str(e)}")
diff --git a/tests/orchestrator/log_dedup.py b/tests/orchestrator/log_dedup.py
@@ -0,0 +1,140 @@
+# test_log_deduplication.py
+import os
+import unittest
+from aiopslab.orchestrator.actions import log_deduplication as ld
+
+
+class TestTimestampDetection(unittest.TestCase):
+    def test_find_timestamp_spans_iso(self):
+        line = "2025-09-24 18:41:09 some log message"
+        spans = ld.find_timestamp_spans(line, ld.DEFAULT_TS_RX)
+        self.assertEqual(len(spans), 1)
+        start, end = spans[0]
+        self.assertEqual(line[start:end], "2025-09-24 18:41:09")
+
+    def test_find_timestamp_spans_multiple(self):
+        line = "Start 2025-09-24T18:41:09Z end 2025-09-24T19:00:00Z"
+        spans = ld.find_timestamp_spans(line, ld.DEFAULT_TS_RX)
+        self.assertEqual(len(spans), 2)
+
+    def test_find_timestamp_spans_no_match(self):
+        line = "no timestamp here"
+        spans = ld.find_timestamp_spans(line, ld.DEFAULT_TS_RX)
+        self.assertEqual(spans, [])
+
+
+class TestMakeBlocks(unittest.TestCase):
+    def test_make_blocks_basic(self):
+        lines = ["a", "b", "c", "d", "e"]
+        blocks = ld.make_blocks(lines, 2)
+        self.assertEqual(blocks, ["a\nb", "c\nd", "e"])
+
+    def test_make_blocks_invalid(self):
+        with self.assertRaises(ValueError):
+            ld.make_blocks(["x"], 0)
+
+
+class TestGreedyCompressPass(unittest.TestCase):
+    def setUp(self):
+        self.ts_rx = ld.DEFAULT_TS_RX
+
+    def test_no_lines(self):
+        result = ld.greedy_compress_pass([], self.ts_rx, 1)
+        self.assertEqual(result, [])
+
+    def test_dedup_identical_except_timestamps(self):
+        lines = [
+            '{"time":"2025-09-25T04:36:33Z","msg":"foo"}',
+            '{"time":"2025-09-25T04:36:34Z","msg":"foo"}',
+        ]
+        result = ld.greedy_compress_pass(lines, self.ts_rx, 1)
+        self.assertEqual(len(result), 1)
+        self.assertIn('"foo"', result[0])
+
+    def test_different_messages_not_deduped(self):
+        lines = [
+            '{"time":"2025-09-25T04:36:33Z","msg":"foo"}',
+            '{"time":"2025-09-25T04:36:33Z","msg":"bar"}',
+        ]
+        result = ld.greedy_compress_pass(lines, self.ts_rx, 1)
+        self.assertEqual(len(result), 2)
+
+
+class TestGreedyCompressLines(unittest.TestCase):
+    def setUp(self):
+        self.ts_rx = ld.DEFAULT_TS_RX
+
+    def test_env_disabled(self):
+        """Should return input unchanged when LOG_TRIM is unset or invalid."""
+        os.environ.pop("LOG_TRIM", None)
+        text = "some logs\nmore logs"
+        result = ld.greedy_compress_lines(text)
+        self.assertEqual(result, text)
+
+    def test_env_enabled(self):
+        """Should dedup when LOG_TRIM is set."""
+        os.environ["LOG_TRIM"] = "2"
+        text = (
+            '{"time":"2025-09-25T04:36:33Z","message":"TLS disabled."}\n'
+            '{"time":"2025-09-25T04:36:34Z","message":"TLS disabled."}\n'
+        )
+        result = ld.greedy_compress_lines(text, self.ts_rx)
+        # Expect only one entry remains after deduplication
+        self.assertIn("TLS disabled", result)
+        self.assertEqual(result.count("TLS disabled"), 1)
+
+    def test_invalid_env(self):
+        """Invalid LOG_TRIM should skip dedup."""
+        os.environ["LOG_TRIM"] = "invalid"
+        text = "hello world"
+        result = ld.greedy_compress_lines(text)
+        self.assertEqual(result, text)
+
+
+class TestRealisticLogs(unittest.TestCase):
+    """Integration-style tests using the provided log sequences."""
+
+    def setUp(self):
+        self.ts_rx = ld.DEFAULT_TS_RX
+        os.environ["LOG_TRIM"] = "3"
+
+    def test_dedup_kubernetes_events(self):
+        sample = """
+Error: Your service/namespace does not exist. Use kubectl to check.
+{"level":"info","time":"2025-09-25T04:36:33Z","message":"TLS disabled."}
+{"level":"info","time":"2025-09-25T04:36:33Z","message":"TLS disabled."}
+{"level":"info","time":"2025-09-25T04:36:34Z","message":"TLS disabled."}
+"""
+        result = ld.greedy_compress_lines(sample.strip(), self.ts_rx)
+        # Should remove duplicate lines differing only by timestamps
+        self.assertIn("TLS disabled", result)
+        self.assertEqual(result.count("TLS disabled"), 1)
+
+    def test_k8s_container_events(self):
+        sample = """
+[90m2025-09-25T04:36:33Z[0m [32mINF[0m [1mcmd/geo/db.go:29[0m[36m >[0m [1mNew session successfull...[0m
+[90m2025-09-25T04:36:34Z[0m [32mINF[0m [1mcmd/geo/db.go:29[0m[36m >[0m [1mNew session successfull...[0m
+[90m2025-09-25T04:36:35Z[0m [32mINF[0m [1mcmd/geo/db.go:31[0m[36m >[0m [1mGenerating test data...[0m
+"""
+        result = ld.greedy_compress_lines(sample.strip(), self.ts_rx)
+        # Should collapse identical log messages differing only by timestamp
+        self.assertIn("New session successfull...", result)
+        self.assertEqual(result.count("New session successfull..."), 1)
+        # Should preserve distinct lines
+        self.assertIn("Generating test data...", result)
+
+    def test_multiblock_dedup(self):
+        sample = "\n".join(
+            [
+                f'{{"level":"info","time":"2025-09-25T04:37:{i:02d}Z","message":"Tune: setGCPercent to 100"}}'
+                for i in range(3)
+            ]
+        )
+        result = ld.greedy_compress_lines(sample, self.ts_rx)
+        self.assertIn("Tune: setGCPercent to 100", result)
+        self.assertEqual(result.count("Tune: setGCPercent to 100"), 1)
+
+
+if __name__ == "__main__":
+    unittest.main()
+
