#include "scxtest/scx_test.h"
#include <scx/common.bpf.h>
#include <lib/sdt_task.h>
struct scx_task_map_val {
union sdt_id tid;
__u64 tptr;
struct sdt_data __arena *data;
};
struct {
__uint(type, BPF_MAP_TYPE_TASK_STORAGE);
__uint(map_flags, BPF_F_NO_PREALLOC);
__type(key, int);
__type(value, struct scx_task_map_val);
} scx_task_map SEC(".maps");
struct scx_allocator scx_task_allocator;
__hidden
void __arena *scx_task_alloc(struct task_struct *p)
{
struct sdt_data __arena *data = NULL;
struct scx_task_map_val *mval;
mval = bpf_task_storage_get(&scx_task_map, p, 0,
BPF_LOCAL_STORAGE_GET_F_CREATE);
if (!mval)
return NULL;
data = scx_alloc(&scx_task_allocator);
if (unlikely(!data))
return NULL;
mval->tid = data->tid;
mval->tptr = (__u64) p;
mval->data = data;
return (void __arena *)data->payload;
}
__hidden
int scx_task_init(__u64 data_size)
{
return scx_alloc_init(&scx_task_allocator, data_size);
}
__hidden
void __arena *scx_task_data(struct task_struct *p)
{
struct sdt_data __arena *data;
struct scx_task_map_val *mval;
scx_arena_subprog_init();
mval = bpf_task_storage_get(&scx_task_map, p, 0, 0);
if (!mval)
return NULL;
data = mval->data;
return (void __arena *)data->payload;
}
__hidden
void scx_task_free(struct task_struct *p)
{
struct scx_task_map_val *mval;
scx_arena_subprog_init();
mval = bpf_task_storage_get(&scx_task_map, p, 0, 0);
if (!mval)
return;
scx_alloc_free_idx(&scx_task_allocator, mval->tid.idx);
bpf_task_storage_delete(&scx_task_map, p);
}