Add stack querying helpers

Add stack helpers, test get_stackid
Expore BPF_MAP_GET_NEXT_KEY as nextkey()
Add host map keys iteration
Add Base.get for runtime map
Improve bpfcall to handle maps

Author: jpsamaroo

src/BPFnative.jl

@@ -66,6 +66,9 @@ end
 # Runtime API
 module RT
 import ..API
+using ..LLVM
+using ..LLVM.Interop
+include("runtime/maps_core.jl")
 include("runtime/bpfcall.jl")
 include("runtime/maps.jl")
 include("runtime/buffers.jl")

src/common.jl

@@ -255,6 +255,8 @@ end
     sk_assign
 end
 
+const PERF_MAX_STACK_DEPTH = 127
+
 # Kernel structures
 
 if has_vmlinux

src/host/maps.jl

@@ -15,6 +15,7 @@ struct HashMap{K,V} <: AbstractHashMap{K,V}
     fd::Cint
 end
 maptype_to_jltype(::Val{API.BPF_MAP_TYPE_HASH}) = HashMap
+maptype_to_jltype(::Val{API.BPF_MAP_TYPE_STACK_TRACE}) = HashMap
 struct ArrayMap{K,V} <: AbstractArrayMap{K,V}
     fd::Cint
 end
@@ -134,4 +135,47 @@ function Base.haskey(map::HostMap{K,V}, idx) where {K,V}
     end
 end
 
-# TODO: keys(map) using BPF_MAP_GET_NEXT_KEY
+function nextkey(map::HostMap{K,V}, idx) where {K,V}
+    key = Ref{K}(idx)
+    nkey = Ref{K}()
+    key_ptr = Base.unsafe_convert(Ptr{K}, key)
+    nkey_ptr = Base.unsafe_convert(Ptr{K}, nkey)
+    attr = Ref(map_access_elem_attr(map.fd,
+                                    key_ptr,
+                                    nkey_ptr,
+                                    0))
+    ret = GC.@preserve key nkey begin
+        bpf(API.BPF_MAP_GET_NEXT_KEY, attr)
+    end
+    if (ret == -1) && (Libc.errno() == Libc.ENOENT)
+        return nothing
+    end
+    ret >= 0 || Base.systemerror(ret)
+    nkey[]
+end
+
+struct HostMapKeySet{K,V,H<:HostMap{K,V}}
+    map::H
+end
+Base.keys(map::H) where {K,V,H<:HostMap{K,V}} = HostMapKeySet{K,V,H}(map)
+Base.IteratorSize(::Type{<:HostMapKeySet}) = Base.SizeUnknown()
+Base.eltype(::Type{HostMapKeySet{K,V,H}}) where {K,V,H} = K
+function Base.iterate(hmks::HostMapKeySet{K,V,H}) where {K,V,H}
+    fakekey_ref = Ref{K}()
+    iterate(hmks, fakekey_ref[])
+end
+function Base.iterate(hmks::HostMapKeySet{K,V,H}, key) where {K,V,H}
+    nkey = nextkey(hmks.map, key)
+    if nkey === nothing
+        return nothing
+    end
+    return nkey, nkey
+end
+function Base.length(map::HostMap)
+    # TODO: This sucks!
+    ctr = 0
+    for k in keys(map)
+        ctr += 1
+    end
+    ctr
+end

src/runtime/bpfcall.jl

@@ -3,21 +3,52 @@ export bpfcall
 @generated function _bpfcall(::Val{cmd}, ::Type{rettype}, ::Type{argtypes}, args::Vararg{Any}) where {cmd,rettype,argtypes}
     JuliaContext() do ctx
         T_ret = convert(LLVMType, rettype, ctx)
-        T_args = map(x->convert(LLVMType, x, ctx), argtypes.parameters)
+        T_jlptr = convert(LLVMType, Ptr{Cvoid}, ctx)
+        T_ptr_i8 = LLVM.PointerType(LLVM.Int8Type(ctx))
 
-        llvm_f, _ = create_function(T_ret, LLVMType[T_args...])
+        outer_args = filter(arg->!(arg <: RTMap), args)
+        T_outer_args = LLVMType[convert(LLVMType, arg, ctx) for arg in outer_args]
+
+        llvm_f, _ = create_function(T_ret, T_outer_args)
         mod = LLVM.parent(llvm_f)
 
         Builder(ctx) do builder
             entry = BasicBlock(llvm_f, "entry", ctx)
             position!(builder, entry)
-            ft = LLVM.FunctionType(T_ret, LLVMType[T_args...])
+
+            inner_args = LLVM.Value[]
+            T_inner_args = LLVMType[]
+            outer_args_ex = Expr[]
+            jlidx = 1
+            pidx = 1
+            for (idx, arg) in enumerate(argtypes.parameters)
+                if arg <: RTMap
+                    map_gv = _genmap!(mod, arg, ctx)
+                    push!(inner_args, map_gv)
+                    push!(T_inner_args, llvmtype(map_gv))
+                else
+                    llvm_arg = parameters(llvm_f)[pidx]
+                    T_arg = T_outer_args[pidx]
+                    if arg <: Ptr
+                        llvm_arg = inttoptr!(builder, llvm_arg, T_ptr_i8)
+                        T_arg = llvmtype(llvm_arg)
+                    end
+                    push!(inner_args, llvm_arg)
+                    push!(T_inner_args, T_arg)
+                    push!(outer_args_ex, Expr(:ref, :args, jlidx))
+                    pidx += 1
+                end
+                jlidx += 1
+            end
+
+            ft = LLVM.FunctionType(T_ret, T_inner_args)
             ftp = LLVM.PointerType(ft)
             f = inttoptr!(builder, ConstantInt(cmd, ctx), ftp)
-            value = call!(builder, f, LLVM.Value[parameters(llvm_f)...])
+            value = call!(builder, f, inner_args)
             ret!(builder, value)
         end
-        call_function(llvm_f, rettype, Base.to_tuple_type(args), :((args...,)))
+        outer_args_ex = Expr(:tuple, outer_args_ex...)
+        call_function(llvm_f, rettype, Base.to_tuple_type(outer_args), outer_args_ex)
     end
 end
 @inline bpfcall(cmd::API.BPFHelper, RT, AT, args...) =

src/runtime/helpers.jl

@@ -1,6 +1,7 @@
 # BPF helpers
 
 const ptr_sk_buff = API.pointertype(API.sk_buff)
+const ptr_task_struct = Ptr{Cvoid} #API.pointertype(API.task_struct)
 
 bpfconvert(x) = x
 bpfconvert(x::AbstractBuffer) = pointer(x)
@@ -54,4 +55,19 @@ end
 @inline get_current_comm(buf::AbstractSizedBuffer) =
     bpfcall(API.get_current_comm, Clong, Tuple{BufPtr, UInt32}, pointer(buf), length(buf))
 
+@inline get_stackid(ctx::T, map::M, flags::Integer) where {T,M<:RTMap} =
+    bpfcall(API.get_stackid, Clong, Tuple{T, M, UInt64}, ctx, map, unsafe_trunc(UInt64,flags))
+@inline function get_current_task()
+    res = bpfcall(API.get_current_task, UInt64)
+    if res > 0
+        unsafe_load(reinterpret(ptr_task_struct, res))
+    else
+        nothing
+    end
+end
+@inline get_stack(ctx::T, buf::AbstractSizedBuffer, flags::UInt64) where {T} =
+    bpfcall(API.get_stack, Clong, Tuple{T, BufPtr, UInt32, UInt64}, ctx, pointer(buf), length(buf), flags)
+@inline get_task_stack(ctx::ptr_task_struct, buf::AbstractSizedBuffer, flags::UInt64) where {T} =
+    bpfcall(API.get_task_stack, Clong, Tuple{ptr_task_struct, BufPtr, UInt32, UInt64}, ctx, pointer(buf), length(buf), flags)
+
 # TODO: The rest!

src/runtime/maps.jl

@@ -1,20 +1,3 @@
-using ..LLVM
-using ..LLVM.Interop
-
-abstract type RTMap{Name,MT,K,V,ME,F} end
-abstract type AbstractHashMap{Name,MT,K,V,ME,F} <: RTMap{Name,MT,K,V,ME,F} end
-abstract type AbstractArrayMap{Name,MT,K,V,ME,F} <: RTMap{Name,MT,K,V,ME,F} end
-
-struct HashMap{Name,MT,K,V,ME,F} <: AbstractHashMap{Name,MT,K,V,ME,F} end
-maptype_to_jltype(::Val{API.BPF_MAP_TYPE_HASH}) = HashMap
-struct ArrayMap{Name,MT,K,V,ME,F} <: AbstractArrayMap{Name,MT,K,V,ME,F} end
-maptype_to_jltype(::Val{API.BPF_MAP_TYPE_ARRAY}) = ArrayMap
-
-function RTMap(; name, maptype, keytype, valuetype, maxentries=1, flags=0)
-    jltype = maptype_to_jltype(Val(maptype))
-    jltype{Symbol(name), maptype, keytype, valuetype, maxentries, flags}()
-end
-
 function map_lookup_elem(map::RTMap{Name,MT,K,V,ME,F}, key::K) where {Name,MT,K,V,ME,F}
     keyref = Ref{K}(key)
     GC.@preserve keyref begin
@@ -37,19 +20,6 @@ function map_delete_elem(map::RTMap{Name,MT,K,V,ME,F}, key::K) where {Name,MT,K,
         _map_delete_elem(map, Base.unsafe_convert(Ptr{K}, keyref))
     end
 end
-function _genmap!(mod::LLVM.Module, ::Type{<:RTMap{Name,MT,K,V,ME,F}}, ctx) where {Name,MT,K,V,ME,F}
-    T_i32 = LLVM.Int32Type(ctx)
-    T_map = LLVM.StructType([T_i32, T_i32, T_i32, T_i32, T_i32], ctx)
-    name = string(Name)
-    gv = GlobalVariable(mod, T_map, name)
-    section!(gv, "maps")
-    alignment!(gv, 4)
-    vec = Any[Int32(MT),Int32(sizeof(K)),Int32(sizeof(V)),Int32(ME),Int32(F)]
-    init = ConstantStruct([ConstantInt(v, ctx) for v in vec], ctx)
-    initializer!(gv, init)
-    linkage!(gv, LLVM.API.LLVMLinkOnceODRLinkage)
-    return gv
-end
 @generated function _map_lookup_elem(map::RTMap{Name,MT,K,V,ME,F}, key::Ptr{K}) where {Name,MT,K,V,ME,F}
     JuliaContext() do ctx
         T_keyp = LLVM.PointerType(convert(LLVMType, K, ctx))
@@ -217,6 +187,17 @@ function Base.haskey(map::AbstractArrayMap{Name,MT,K,V,ME,F}, idx) where {Name,M
     end
 end
 
+function Base.get(map::RTMap{Name,MT,K,V,ME,F}, k::K, v::V) where {Name,MT,K,V,ME,F}
+    map_v = map[k]
+    if map_v !== nothing
+        return map_v
+    else
+        return v
+    end
+end
+@inline Base.get(map::RTMap{Name,MT,K,V,ME,F}, k, v) where {Name,MT,K,V,ME,F} =
+    get(map, bpfconvert(K, k), bpfconvert(V, v))
+
 ## Perf
 
 #= TODO

src/runtime/maps_core.jl

@@ -0,0 +1,29 @@
+abstract type RTMap{Name,MT,K,V,ME,F} end
+abstract type AbstractHashMap{Name,MT,K,V,ME,F} <: RTMap{Name,MT,K,V,ME,F} end
+abstract type AbstractArrayMap{Name,MT,K,V,ME,F} <: RTMap{Name,MT,K,V,ME,F} end
+
+struct HashMap{Name,MT,K,V,ME,F} <: AbstractHashMap{Name,MT,K,V,ME,F} end
+maptype_to_jltype(::Val{API.BPF_MAP_TYPE_HASH}) = HashMap
+maptype_to_jltype(::Val{API.BPF_MAP_TYPE_STACK_TRACE}) = HashMap
+struct ArrayMap{Name,MT,K,V,ME,F} <: AbstractArrayMap{Name,MT,K,V,ME,F} end
+maptype_to_jltype(::Val{API.BPF_MAP_TYPE_ARRAY}) = ArrayMap
+
+function RTMap(; name, maptype, keytype, valuetype, maxentries=1, flags=0)
+    jltype = maptype_to_jltype(Val(maptype))
+    jltype{Symbol(name), maptype, keytype, valuetype, maxentries, flags}()
+end
+
+function _genmap!(mod::LLVM.Module, ::Type{<:RTMap{Name,MT,K,V,ME,F}}, ctx) where {Name,MT,K,V,ME,F}
+    T_i32 = LLVM.Int32Type(ctx)
+    T_map = LLVM.StructType([T_i32, T_i32, T_i32, T_i32, T_i32], ctx)
+    name = string(Name)
+    gv = GlobalVariable(mod, T_map, name)
+    section!(gv, "maps")
+    alignment!(gv, 4)
+    vec = Any[Int32(MT),Int32(sizeof(K)),Int32(sizeof(V)),Int32(ME),Int32(F)]
+    init = ConstantStruct([ConstantInt(v, ctx) for v in vec], ctx)
+    initializer!(gv, init)
+    linkage!(gv, LLVM.API.LLVMLinkOnceODRLinkage)
+    return gv
+end
+

test/runtests.jl

@@ -452,5 +452,40 @@ if run_root_tests
                 @test haskey(hmap, 2)
             end
         end
+        # TODO: get_current_comm
+        @testset "get_stackid" begin
+            # from BCC's stackcount
+            # FIXME: { i64, i32 } causes verifier errors
+            struct StackKey
+                tgid::UInt64
+                sid::Clong
+            end
+            kp = KProbe("ksys_write"; license="GPL") do x
+                stacks = RT.RTMap(;name="stacks",maptype=API.BPF_MAP_TYPE_STACK_TRACE,keytype=UInt32,valuetype=NTuple{API.PERF_MAX_STACK_DEPTH,UInt64},maxentries=100)
+                counts = RT.RTMap(;name="counts",maptype=API.BPF_MAP_TYPE_HASH,keytype=StackKey,valuetype=UInt32,maxentries=100)
+                pid, tgid = RT.get_current_pid_tgid()
+                sid = RT.get_stackid(x, stacks, 0)
+                key = StackKey(tgid, sid)
+                old_count = get(counts, key, 0)
+                counts[key] = old_count + 1
+                0
+            end
+            API.load(kp) do
+                stacks = Host.hostmap(API.findmap(kp.obj, "stacks"); K=UInt32, V=NTuple{API.PERF_MAX_STACK_DEPTH,UInt64})
+                counts = Host.hostmap(API.findmap(kp.obj, "counts"); K=StackKey, V=UInt32)
+                write(test_io, "1"); flush(test_io)
+                @test length(counts) > 0
+                key = nothing
+                for k in keys(counts)
+                    if k.tgid == getpid()
+                        key = k
+                        break
+                    end
+                end
+                @test key !== nothing
+                @test haskey(counts, key)
+                @test haskey(stacks, key.sid)
+            end
+        end
     end
 end