Skip to content

Commit 127852e

Browse files
committed
Add passing test struct_pylib.py
1 parent 2fd4fef commit 127852e

File tree

1 file changed

+93
-0
lines changed

1 file changed

+93
-0
lines changed
Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
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

Comments
 (0)