|
| 1 | +""" |
| 2 | +Test struct values in HashMap. |
| 3 | +
|
| 4 | +This example stores a struct in a HashMap and reads it back, |
| 5 | +testing the new set_value_struct() functionality in pylibbpf. |
| 6 | +""" |
| 7 | + |
| 8 | +from pythonbpf import bpf, map, struct, section, bpfglobal, BPF |
| 9 | +from pythonbpf.helper import ktime, smp_processor_id, pid, comm |
| 10 | +from pythonbpf.maps import HashMap |
| 11 | +from ctypes import c_void_p, c_int64, c_uint32, c_uint64 |
| 12 | +import time |
| 13 | +import os |
| 14 | + |
| 15 | + |
| 16 | +@bpf |
| 17 | +@struct |
| 18 | +class task_info: |
| 19 | + pid: c_uint64 |
| 20 | + timestamp: c_uint64 |
| 21 | + comm: str(16) |
| 22 | + |
| 23 | + |
| 24 | +@bpf |
| 25 | +@map |
| 26 | +def cpu_tasks() -> HashMap: |
| 27 | + return HashMap(key=c_uint32, value=task_info, max_entries=256) |
| 28 | + |
| 29 | + |
| 30 | +@bpf |
| 31 | +@section("tracepoint/sched/sched_switch") |
| 32 | +def trace_sched_switch(ctx: c_void_p) -> c_int64: |
| 33 | + cpu = smp_processor_id() |
| 34 | + |
| 35 | + # Create task info struct |
| 36 | + info = task_info() |
| 37 | + info.pid = pid() |
| 38 | + info.timestamp = ktime() |
| 39 | + comm(info.comm) |
| 40 | + |
| 41 | + # Store in map |
| 42 | + cpu_tasks.update(cpu, info) |
| 43 | + |
| 44 | + return 0 # type: ignore |
| 45 | + |
| 46 | + |
| 47 | +@bpf |
| 48 | +@bpfglobal |
| 49 | +def LICENSE() -> str: |
| 50 | + return "GPL" |
| 51 | + |
| 52 | + |
| 53 | +# Compile and load |
| 54 | +b = BPF() |
| 55 | +b.load() |
| 56 | +b.attach_all() |
| 57 | + |
| 58 | +print("Testing HashMap with Struct Values") |
| 59 | + |
| 60 | +cpu_map = b["cpu_tasks"] |
| 61 | +cpu_map.set_value_struct("task_info") # Enable struct deserialization |
| 62 | + |
| 63 | +print("Listening for context switches.. .\n") |
| 64 | + |
| 65 | +num_cpus = os.cpu_count() or 16 |
| 66 | + |
| 67 | +try: |
| 68 | + while True: |
| 69 | + time.sleep(1) |
| 70 | + |
| 71 | + print(f"--- Snapshot at {time.strftime('%H:%M:%S')} ---") |
| 72 | + |
| 73 | + for cpu in range(num_cpus): |
| 74 | + try: |
| 75 | + info = cpu_map.lookup(cpu) |
| 76 | + |
| 77 | + if info: |
| 78 | + comm_str = ( |
| 79 | + bytes(info.comm).decode("utf-8", errors="ignore").rstrip("\x00") |
| 80 | + ) |
| 81 | + ts_sec = info.timestamp / 1e9 |
| 82 | + |
| 83 | + print( |
| 84 | + f" CPU {cpu}: PID={info.pid}, comm={comm_str}, ts={ts_sec:.3f}s" |
| 85 | + ) |
| 86 | + except KeyError: |
| 87 | + # No data for this CPU yet |
| 88 | + pass |
| 89 | + |
| 90 | + print() |
| 91 | + |
| 92 | +except KeyboardInterrupt: |
| 93 | + print("\nStopped") |
0 commit comments