varun-r-mallya/python-bpf
1
0
Fork 0
mirror of https://github.com/varun-r-mallya/Python-BPF.git synced 2025-12-31 21:06:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

format chore

Browse Source
This commit is contained in:
varun-r-mallya
2025-11-28 20:59:46 +05:30
parent bbe4990878
commit d73c793989
3 changed files with 98 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
BCC-Examples/container-monitor/container_monitor.py
View File

@ -1,20 +1,18 @@
"""Container Monitor - TUI-based cgroup monitoring combining syscall, file I/O, and network tracking."""
import time
import os
from pathlib import Path
from pythonbpf import bpf, map, section, bpfglobal, struct, BPF
from pythonbpf.maps import HashMap
from pythonbpf.helper import get_current_cgroup_id
from ctypes import c_int32, c_uint64, c_void_p
from vmlinux import struct_pt_regs, struct_sk_buff
from data_collector import ContainerDataCollector
from data_collection import ContainerDataCollector
from tui import ContainerMonitorTUI
# ==================== BPF Structs ====================
@bpf
@struct
class read_stats:
@ -40,6 +38,7 @@ class net_stats:
# ==================== BPF Maps ====================
@bpf
@map
def read_map() -> HashMap:
@ -66,6 +65,7 @@ def syscall_count() -> HashMap:
# ==================== File I/O Tracing ====================
@bpf
@section("kprobe/vfs_read")
def trace_read(ctx: struct_pt_regs) -> c_int32:
@ -109,6 +109,7 @@ def trace_write(ctx1: struct_pt_regs) -> c_int32:
# ==================== Network I/O Tracing ====================
@bpf
@section("kprobe/__netif_receive_skb")
def trace_netif_rx(ctx2: struct_pt_regs) -> c_int32:
@ -165,6 +166,7 @@ def trace_dev_xmit(ctx3: struct_pt_regs) -> c_int32:
# ==================== Syscall Tracing ====================
@bpf
@section("tracepoint/raw_syscalls/sys_enter")
def count_syscalls(ctx: c_void_p) -> c_int32:
@ -190,32 +192,29 @@ def LICENSE() -> str:
if __name__ == "__main__":
print("🔥 Loading BPF programs...")
# Load and attach BPF program
b = BPF()
b.load()
b.attach_all()
# Get map references and enable struct deserialization
read_map_ref = b["read_map"]
write_map_ref = b["write_map"]
net_stats_map_ref = b["net_stats_map"]
syscall_count_ref = b["syscall_count"]
read_map_ref.set_value_struct("read_stats")
write_map_ref.set_value_struct("write_stats")
net_stats_map_ref.set_value_struct("net_stats")
print("✅ BPF programs loaded and attached")
# Setup data collector
collector = ContainerDataCollector(
read_map_ref,
write_map_ref,
net_stats_map_ref,
syscall_count_ref
read_map_ref, write_map_ref, net_stats_map_ref, syscall_count_ref
)
# Create and run TUI
tui = ContainerMonitorTUI(collector)
tui.run()

28
BCC-Examples/container-monitor/data_collector.py → BCC-Examples/container-monitor/data_collection.py
View File

@ -3,7 +3,7 @@
import os
import time
from pathlib import Path
from typing import Dict, List, Set, Optional, Tuple
from typing import Dict, List, Set, Optional
from dataclasses import dataclass
from collections import deque, defaultdict
@ -11,6 +11,7 @@ from collections import deque, defaultdict
@dataclass
class CgroupInfo:
"""Information about a cgroup."""
id: int
name: str
path: str
@ -19,6 +20,7 @@ class CgroupInfo:
@dataclass
class ContainerStats:
"""Statistics for a container/cgroup."""
cgroup_id: int
cgroup_name: str
@ -44,7 +46,9 @@ class ContainerStats:
class ContainerDataCollector:
"""Collects and manages container monitoring data from BPF."""
def __init__(self, read_map, write_map, net_stats_map, syscall_map, history_size: int = 100):
def __init__(
self, read_map, write_map, net_stats_map, syscall_map, history_size: int = 100
):
self.read_map = read_map
self.write_map = write_map
self.net_stats_map = net_stats_map
@ -53,12 +57,14 @@ class ContainerDataCollector:
# Caching
self._cgroup_cache: Dict[int, CgroupInfo] = {}
self._cgroup_cache_time = 0
self._cache_ttl = 5.
self._cache_ttl = 5.0
0 # Refresh cache every 5 seconds
# Historical data for graphing
self._history_size = history_size
self._history: Dict[int, deque] = defaultdict(lambda: deque(maxlen=history_size))
self._history: Dict[int, deque] = defaultdict(
lambda: deque(maxlen=history_size)
)
def get_all_cgroups(self) -> List[CgroupInfo]:
"""Get all cgroups with caching."""
@ -105,11 +111,7 @@ class ContainerDataCollector:
best_path = self._get_best_cgroup_path(paths)
name = self._get_cgroup_name(best_path)
new_cache[cgroup_id] = CgroupInfo(
id=cgroup_id,
name=name,
path=best_path
)
new_cache[cgroup_id] = CgroupInfo(id=cgroup_id, name=name, path=best_path)
self._cgroup_cache = new_cache
self._cgroup_cache_time = time.time()
@ -120,13 +122,13 @@ class ContainerDataCollector:
# Prefer paths with more components (more specific)
# Prefer paths containing docker, podman, etc.
for keyword in ['docker', 'podman', 'kubernetes', 'k8s', 'systemd']:
for keyword in ["docker", "podman", "kubernetes", "k8s", "systemd"]:
for path in path_list:
if keyword in path.lower():
return path
# Return longest path (most specific)
return max(path_list, key=lambda p: (len(p.split('/')), len(p)))
return max(path_list, key=lambda p: (len(p.split("/")), len(p)))
def _get_cgroup_name(self, path: str) -> str:
"""Extract a friendly name from cgroup path."""
@ -165,9 +167,7 @@ class ContainerDataCollector:
cgroup_name = cgroup_info.name if cgroup_info else f"cgroup-{cgroup_id}"
stats = ContainerStats(
cgroup_id=cgroup_id,
cgroup_name=cgroup_name,
timestamp=time.time()
cgroup_id=cgroup_id, cgroup_name=cgroup_name, timestamp=time.time()
)
# Get file I/O stats

107
BCC-Examples/container-monitor/tui.py
View File

@ -1,10 +1,9 @@
"""Terminal User Interface for container monitoring."""
import sys
import time
import curses
from typing import Optional, List
from data_collector import ContainerDataCollector, CgroupInfo, ContainerStats
from data_collection import ContainerDataCollector
class ContainerMonitorTUI:
@ -122,11 +121,11 @@ class ContainerMonitorTUI:
# Highlight selected
stdscr.attron(curses.color_pair(2) | curses.A_BOLD | curses.A_REVERSE)
line = f"► {cgroup.name:<40} ID: {cgroup.id}"
stdscr.addstr(y, 2, line[:width - 4])
stdscr.addstr(y, 2, line[: width - 4])
stdscr.attroff(curses.color_pair(2) | curses.A_BOLD | curses.A_REVERSE)
else:
line = f" {cgroup.name:<40} ID: {cgroup.id}"
stdscr.addstr(y, 2, line[:width - 4])
stdscr.addstr(y, 2, line[: width - 4])
# Footer with count
footer = f"Total cgroups: {len(cgroups)}"
@ -174,21 +173,37 @@ class ContainerMonitorTUI:
# RX graph
y += 2
rx_label = f"RX: {self._format_bytes(stats.rx_bytes)} ({stats.rx_packets:,} packets)"
rx_label = (
f"RX: {self._format_bytes(stats.rx_bytes)} ({stats.rx_packets:,} packets)"
)
stdscr.addstr(y, 2, rx_label)
if len(history) > 1:
self._draw_bar_graph(stdscr, y + 1, 2, width - 4, 3,
[s.rx_bytes for s in history],
curses.color_pair(2))
self._draw_bar_graph(
stdscr,
y + 1,
2,
width - 4,
3,
[s.rx_bytes for s in history],
curses.color_pair(2),
)
# TX graph
y += 5
tx_label = f"TX: {self._format_bytes(stats.tx_bytes)} ({stats.tx_packets:,} packets)"
tx_label = (
f"TX: {self._format_bytes(stats.tx_bytes)} ({stats.tx_packets:,} packets)"
)
stdscr.addstr(y, 2, tx_label)
if len(history) > 1:
self._draw_bar_graph(stdscr, y + 1, 2, width - 4, 3,
[s.tx_bytes for s in history],
curses.color_pair(3))
self._draw_bar_graph(
stdscr,
y + 1,
2,
width - 4,
3,
[s.tx_bytes for s in history],
curses.color_pair(3),
)
# File I/O graphs
y += 5
@ -196,21 +211,37 @@ class ContainerMonitorTUI:
# Read graph
y += 2
read_label = f"READ: {self._format_bytes(stats.read_bytes)} ({stats.read_ops:,} ops)"
read_label = (
f"READ: {self._format_bytes(stats.read_bytes)} ({stats.read_ops:,} ops)"
)
stdscr.addstr(y, 2, read_label)
if len(history) > 1:
self._draw_bar_graph(stdscr, y + 1, 2, width - 4, 3,
[s.read_bytes for s in history],
curses.color_pair(4))
self._draw_bar_graph(
stdscr,
y + 1,
2,
width - 4,
3,
[s.read_bytes for s in history],
curses.color_pair(4),
)
# Write graph
y += 5
write_label = f"WRITE: {self._format_bytes(stats.write_bytes)} ({stats.write_ops:,} ops)"
write_label = (
f"WRITE: {self._format_bytes(stats.write_bytes)} ({stats.write_ops:,} ops)"
)
stdscr.addstr(y, 2, write_label)
if len(history) > 1:
self._draw_bar_graph(stdscr, y + 1, 2, width - 4, 3,
[s.write_bytes for s in history],
curses.color_pair(5))
self._draw_bar_graph(
stdscr,
y + 1,
2,
width - 4,
3,
[s.write_bytes for s in history],
curses.color_pair(5),
)
def _draw_section_header(self, stdscr, y: int, title: str, color_pair: int):
"""Draw a section header."""
@ -220,8 +251,16 @@ class ContainerMonitorTUI:
stdscr.addstr(y, len(title) + 3, "─" * (width - len(title) - 5))
stdscr.attroff(curses.color_pair(color_pair) | curses.A_BOLD)
def _draw_bar_graph(self, stdscr, y: int, x: int, width: int, height: int,
data: List[float], color_pair: int):
def _draw_bar_graph(
self,
stdscr,
y: int,
x: int,
width: int,
height: int,
data: List[float],
color_pair: int,
):
"""Draw a simple bar graph."""
if not data or width < 2:
return
@ -250,25 +289,21 @@ class ContainerMonitorTUI:
else:
bar_line += " "
try:
stdscr.attron(color_pair)
stdscr.addstr(y + row, x, bar_line[:width])
stdscr.attroff(color_pair)
except:
pass # Ignore errors at screen edges
stdscr.attron(color_pair)
stdscr.addstr(y + row, x, bar_line[:width])
stdscr.attroff(color_pair)
def _format_bytes(self, bytes_val: int) -> str:
def _format_bytes(self, bytes_val: float) -> str:
"""Format bytes into human-readable string."""
for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
for unit in ["B", "KB", "MB", "GB", "TB"]:
if bytes_val < 1024.0:
return f"{bytes_val:.2f} {unit}"
bytes_val /= 1024.
0
bytes_val /= 1024.0
return f"{bytes_val:.2f} PB"
def _handle_input(self, key: int) -> bool:
"""Handle keyboard input. Returns False to exit."""
if key == ord('q') or key == ord('Q'):
if key == ord("q") or key == ord("Q"):
return False # Exit
if self.current_screen == "selection":
@ -277,19 +312,19 @@ class ContainerMonitorTUI:
elif key == curses.KEY_DOWN:
cgroups = self.collector.get_all_cgroups()
self.selected_index = min(len(cgroups) - 1, self.selected_index + 1)
elif key == ord('\n') or key == curses.KEY_ENTER or key == 10:
elif key == ord("\n") or key == curses.KEY_ENTER or key == 10:
# Select cgroup
cgroups = self.collector.get_all_cgroups()
if cgroups and 0 <= self.selected_index < len(cgroups):
cgroups.sort(key=lambda c: c.name)
self.selected_cgroup = cgroups[self.selected_index].id
self.current_screen = "monitoring"
elif key == ord('r') or key == ord('R'):
elif key == ord("r") or key == ord("R"):
# Force refresh cache
self.collector._cgroup_cache_time = 0
elif self.current_screen == "monitoring":
if key == 27 or key == ord('b') or key == ord('B'): # ESC or 'b'
if key == 27 or key == ord("b") or key == ord("B"): # ESC or 'b'
self.current_screen = "selection"
self.selected_cgroup = None