mirror of
https://github.com/varun-r-mallya/Python-BPF.git
synced 2025-12-31 21:06:25 +00:00
add file io and network stats in container monitor example
This commit is contained in:
@ -1,6 +1,8 @@
|
||||
import logging
|
||||
|
||||
from pythonbpf import bpf, map, section, bpfglobal, struct, compile
|
||||
import time
|
||||
import os
|
||||
from pathlib import Path
|
||||
from pythonbpf import bpf, map, section, bpfglobal, struct, compile, BPF
|
||||
from pythonbpf.maps import HashMap
|
||||
from pythonbpf.helper import get_current_cgroup_id
|
||||
from ctypes import c_int32, c_uint64
|
||||
@ -33,23 +35,18 @@ def write_map() -> HashMap:
|
||||
return HashMap(key=c_uint64, value=write_stats, max_entries=1024)
|
||||
|
||||
|
||||
#
|
||||
# READ PROBE
|
||||
#
|
||||
@bpf
|
||||
@section("kprobe/vfs_read")
|
||||
def trace_read(ctx: struct_pt_regs) -> c_int32:
|
||||
cg = get_current_cgroup_id()
|
||||
count = c_uint64(ctx.dx)
|
||||
ptr = read_map.lookup(cg)
|
||||
|
||||
if ptr:
|
||||
s = read_stats()
|
||||
s.bytes = ptr.bytes + count
|
||||
s.ops = ptr.ops + 1
|
||||
read_map.update(cg, ptr)
|
||||
read_map.update(cg, s)
|
||||
else:
|
||||
print("read init")
|
||||
s = read_stats()
|
||||
s.bytes = count
|
||||
s.ops = c_uint64(1)
|
||||
@ -58,9 +55,6 @@ def trace_read(ctx: struct_pt_regs) -> c_int32:
|
||||
return c_int32(0)
|
||||
|
||||
|
||||
#
|
||||
# WRITE PROBE
|
||||
#
|
||||
@bpf
|
||||
@section("kprobe/vfs_write")
|
||||
def trace_write(ctx1: struct_pt_regs) -> c_int32:
|
||||
@ -74,7 +68,6 @@ def trace_write(ctx1: struct_pt_regs) -> c_int32:
|
||||
s.ops = ptr.ops + 1
|
||||
write_map.update(cg, s)
|
||||
else:
|
||||
print("write init")
|
||||
s = write_stats()
|
||||
s.bytes = count
|
||||
s.ops = c_uint64(1)
|
||||
@ -89,4 +82,110 @@ def LICENSE() -> str:
|
||||
return "GPL"
|
||||
|
||||
|
||||
compile(loglevel=logging.INFO)
|
||||
# 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"]
|
||||
read_map_ref.set_value_struct("read_stats")
|
||||
write_map_ref.set_value_struct("write_stats")
|
||||
|
||||
|
||||
def get_cgroup_ids():
|
||||
"""Get all cgroup IDs from the system"""
|
||||
cgroup_ids = set()
|
||||
|
||||
# Get cgroup IDs from running processes
|
||||
for proc_dir in Path("/proc").glob("[0-9]*"):
|
||||
try:
|
||||
cgroup_file = proc_dir / "cgroup"
|
||||
if cgroup_file.exists():
|
||||
with open(cgroup_file) as f:
|
||||
for line in f:
|
||||
# Parse cgroup path and get inode
|
||||
parts = line.strip().split(":")
|
||||
if len(parts) >= 3:
|
||||
cgroup_path = parts[2]
|
||||
# Try to get the cgroup inode which is used as ID
|
||||
cgroup_mount = f"/sys/fs/cgroup{cgroup_path}"
|
||||
if os.path.exists(cgroup_mount):
|
||||
stat_info = os.stat(cgroup_mount)
|
||||
cgroup_ids.add(stat_info.st_ino)
|
||||
except (PermissionError, FileNotFoundError, OSError):
|
||||
continue
|
||||
|
||||
return cgroup_ids
|
||||
|
||||
|
||||
# Display function
|
||||
def display_stats():
|
||||
"""Read and display file I/O statistics from BPF maps"""
|
||||
print("\n" + "=" * 80)
|
||||
print(f"{'CGROUP ID':<20} {'READ OPS':<15} {'READ BYTES':<20} {'WRITE OPS':<15} {'WRITE BYTES':<20}")
|
||||
print("=" * 80)
|
||||
|
||||
# Get cgroup IDs from the system
|
||||
cgroup_ids = get_cgroup_ids()
|
||||
|
||||
if not cgroup_ids:
|
||||
print("No cgroups found...")
|
||||
print("=" * 80)
|
||||
return
|
||||
|
||||
# Initialize totals
|
||||
total_read_ops = 0
|
||||
total_read_bytes = 0
|
||||
total_write_ops = 0
|
||||
total_write_bytes = 0
|
||||
|
||||
# Track which cgroups have data
|
||||
cgroups_with_data = []
|
||||
|
||||
# Display stats for each cgroup
|
||||
for cgroup_id in sorted(cgroup_ids):
|
||||
# Get read stats using lookup
|
||||
read_ops = 0
|
||||
read_bytes = 0
|
||||
read_stat = read_map_ref.lookup(cgroup_id)
|
||||
if read_stat:
|
||||
read_ops = int(read_stat.ops)
|
||||
read_bytes = int(read_stat.bytes)
|
||||
total_read_ops += read_ops
|
||||
total_read_bytes += read_bytes
|
||||
|
||||
# Get write stats using lookup
|
||||
write_ops = 0
|
||||
write_bytes = 0
|
||||
write_stat = write_map_ref.lookup(cgroup_id)
|
||||
if write_stat:
|
||||
write_ops = int(write_stat.ops)
|
||||
write_bytes = int(write_stat.bytes)
|
||||
total_write_ops += write_ops
|
||||
total_write_bytes += write_bytes
|
||||
|
||||
# Only print if there's data for this cgroup
|
||||
if read_stat or write_stat:
|
||||
print(f"{cgroup_id:<20} {read_ops:<15} {read_bytes:<20} {write_ops:<15} {write_bytes:<20}")
|
||||
cgroups_with_data.append(cgroup_id)
|
||||
|
||||
if not cgroups_with_data:
|
||||
print("No data collected yet...")
|
||||
|
||||
print("=" * 80)
|
||||
print(f"{'TOTAL':<20} {total_read_ops:<15} {total_read_bytes:<20} {total_write_ops:<15} {total_write_bytes:<20}")
|
||||
print()
|
||||
|
||||
|
||||
# Main loop
|
||||
if __name__ == "__main__":
|
||||
print("Tracing file I/O operations... Press Ctrl+C to exit\n")
|
||||
|
||||
try:
|
||||
while True:
|
||||
time.sleep(5) # Update every 5 seconds
|
||||
display_stats()
|
||||
except KeyboardInterrupt:
|
||||
print("\nStopped")
|
||||
|
||||
192
BCC-Examples/container-monitor/net_stats.bpf.py
Normal file
192
BCC-Examples/container-monitor/net_stats.bpf.py
Normal file
@ -0,0 +1,192 @@
|
||||
import logging
|
||||
import time
|
||||
import os
|
||||
from pathlib import Path
|
||||
from pythonbpf import bpf, map, section, bpfglobal, struct, compile, BPF
|
||||
from pythonbpf.maps import HashMap
|
||||
from pythonbpf.helper import get_current_cgroup_id
|
||||
from ctypes import c_int32, c_uint64, c_void_p, c_uint32
|
||||
from vmlinux import struct_sk_buff, struct_pt_regs
|
||||
|
||||
|
||||
@bpf
|
||||
@struct
|
||||
class net_stats:
|
||||
rx_packets: c_uint64
|
||||
tx_packets: c_uint64
|
||||
rx_bytes: c_uint64
|
||||
tx_bytes: c_uint64
|
||||
|
||||
|
||||
@bpf
|
||||
@map
|
||||
def net_stats_map() -> HashMap:
|
||||
return HashMap(key=c_uint64, value=net_stats, max_entries=1024)
|
||||
|
||||
|
||||
@bpf
|
||||
@section("kprobe/__netif_receive_skb")
|
||||
def trace_netif_rx(ctx: struct_pt_regs) -> c_int32:
|
||||
cgroup_id = get_current_cgroup_id()
|
||||
|
||||
# Read skb pointer from first argument (PT_REGS_PARM1)
|
||||
skb = struct_sk_buff(ctx.di) # x86_64: first arg in rdi
|
||||
|
||||
# Read skb->len
|
||||
pkt_len = c_uint64(skb.len)
|
||||
|
||||
|
||||
stats_ptr = net_stats_map.lookup(cgroup_id)
|
||||
|
||||
if stats_ptr:
|
||||
stats = net_stats()
|
||||
stats.rx_packets = stats_ptr.rx_packets + 1
|
||||
stats.tx_packets = stats_ptr.tx_packets
|
||||
stats.rx_bytes = stats_ptr.rx_bytes + pkt_len
|
||||
stats.tx_bytes = stats_ptr.tx_bytes
|
||||
net_stats_map.update(cgroup_id, stats)
|
||||
else:
|
||||
stats = net_stats()
|
||||
stats.rx_packets = c_uint64(1)
|
||||
stats.tx_packets = c_uint64(0)
|
||||
stats.rx_bytes = pkt_len
|
||||
stats.tx_bytes = c_uint64(0)
|
||||
net_stats_map.update(cgroup_id, stats)
|
||||
|
||||
return c_int32(0)
|
||||
|
||||
@bpf
|
||||
@section("kprobe/__dev_queue_xmit")
|
||||
def trace_dev_xmit(ctx1: struct_pt_regs) -> c_int32:
|
||||
cgroup_id = get_current_cgroup_id()
|
||||
|
||||
# Read skb pointer from first argument
|
||||
skb = struct_sk_buff(ctx1.di)
|
||||
pkt_len = c_uint64(skb.len)
|
||||
|
||||
stats_ptr = net_stats_map.lookup(cgroup_id)
|
||||
|
||||
if stats_ptr:
|
||||
stats = net_stats()
|
||||
stats.rx_packets = stats_ptr.rx_packets
|
||||
stats.tx_packets = stats_ptr.tx_packets + 1
|
||||
stats.rx_bytes = stats_ptr.rx_bytes
|
||||
stats.tx_bytes = stats_ptr.tx_bytes + pkt_len
|
||||
net_stats_map.update(cgroup_id, stats)
|
||||
else:
|
||||
stats = net_stats()
|
||||
stats.rx_packets = c_uint64(0)
|
||||
stats.tx_packets = c_uint64(1)
|
||||
stats.rx_bytes = c_uint64(0)
|
||||
stats.tx_bytes = pkt_len
|
||||
net_stats_map.update(cgroup_id, stats)
|
||||
|
||||
return c_int32(0)
|
||||
|
||||
|
||||
@bpf
|
||||
@bpfglobal
|
||||
def LICENSE() -> str:
|
||||
return "GPL"
|
||||
|
||||
# Load and attach BPF program
|
||||
b = BPF()
|
||||
b.load()
|
||||
b.attach_all()
|
||||
|
||||
# Get map reference and enable struct deserialization
|
||||
net_stats_map_ref = b["net_stats_map"]
|
||||
net_stats_map_ref.set_value_struct("net_stats")
|
||||
|
||||
|
||||
def get_cgroup_ids():
|
||||
"""Get all cgroup IDs from the system"""
|
||||
cgroup_ids = set()
|
||||
|
||||
# Get cgroup IDs from running processes
|
||||
for proc_dir in Path("/proc").glob("[0-9]*"):
|
||||
try:
|
||||
cgroup_file = proc_dir / "cgroup"
|
||||
if cgroup_file.exists():
|
||||
with open(cgroup_file) as f:
|
||||
for line in f:
|
||||
# Parse cgroup path and get inode
|
||||
parts = line.strip().split(":")
|
||||
if len(parts) >= 3:
|
||||
cgroup_path = parts[2]
|
||||
# Try to get the cgroup inode which is used as ID
|
||||
cgroup_mount = f"/sys/fs/cgroup{cgroup_path}"
|
||||
if os.path.exists(cgroup_mount):
|
||||
stat_info = os.stat(cgroup_mount)
|
||||
cgroup_ids.add(stat_info.st_ino)
|
||||
except (PermissionError, FileNotFoundError, OSError):
|
||||
continue
|
||||
|
||||
return cgroup_ids
|
||||
|
||||
|
||||
# Display function
|
||||
def display_stats():
|
||||
"""Read and display network I/O statistics from BPF maps"""
|
||||
print("\n" + "=" * 100)
|
||||
print(f"{'CGROUP ID':<20} {'RX PACKETS':<15} {'RX BYTES':<20} {'TX PACKETS':<15} {'TX BYTES':<20}")
|
||||
print("=" * 100)
|
||||
|
||||
# Get cgroup IDs from the system
|
||||
cgroup_ids = get_cgroup_ids()
|
||||
|
||||
if not cgroup_ids:
|
||||
print("No cgroups found...")
|
||||
print("=" * 100)
|
||||
return
|
||||
|
||||
# Initialize totals
|
||||
total_rx_packets = 0
|
||||
total_rx_bytes = 0
|
||||
total_tx_packets = 0
|
||||
total_tx_bytes = 0
|
||||
|
||||
# Track which cgroups have data
|
||||
cgroups_with_data = []
|
||||
|
||||
# Display stats for each cgroup
|
||||
for cgroup_id in sorted(cgroup_ids):
|
||||
# Get network stats using lookup
|
||||
rx_packets = 0
|
||||
rx_bytes = 0
|
||||
tx_packets = 0
|
||||
tx_bytes = 0
|
||||
|
||||
net_stat = net_stats_map_ref.lookup(cgroup_id)
|
||||
if net_stat:
|
||||
rx_packets = int(net_stat.rx_packets)
|
||||
rx_bytes = int(net_stat.rx_bytes)
|
||||
tx_packets = int(net_stat.tx_packets)
|
||||
tx_bytes = int(net_stat.tx_bytes)
|
||||
|
||||
total_rx_packets += rx_packets
|
||||
total_rx_bytes += rx_bytes
|
||||
total_tx_packets += tx_packets
|
||||
total_tx_bytes += tx_bytes
|
||||
|
||||
print(f"{cgroup_id:<20} {rx_packets:<15} {rx_bytes:<20} {tx_packets:<15} {tx_bytes:<20}")
|
||||
cgroups_with_data.append(cgroup_id)
|
||||
|
||||
if not cgroups_with_data:
|
||||
print("No data collected yet...")
|
||||
|
||||
print("=" * 100)
|
||||
print(f"{'TOTAL':<20} {total_rx_packets:<15} {total_rx_bytes:<20} {total_tx_packets:<15} {total_tx_bytes:<20}")
|
||||
print()
|
||||
|
||||
|
||||
# Main loop
|
||||
if __name__ == "__main__":
|
||||
print("Tracing network I/O operations... Press Ctrl+C to exit\n")
|
||||
|
||||
try:
|
||||
while True:
|
||||
time.sleep(5) # Update every 5 seconds
|
||||
display_stats()
|
||||
except KeyboardInterrupt:
|
||||
print("\nStopped")
|
||||
Reference in New Issue
Block a user