Hi Abhishek
On 2026/9/22 05:26, Abhishek Bapat wrote:
> Introduce a new IOCTL (`ALLOCINFO_IOC_TOGGLE_TRACE`) to selectively
> toggle tracing on exact allocation call sites. Userspace tools can
> use the existing filtering mechanism to specify the set of tags to
> toggle tracing for.
>
> To facilitate low overhead execution for non-targeted call sites, add a
> new `CODETAG_FLAG_TRACE_ON` flag to `struct codetag` to track per-site
> activation. Protect these conditional branch evaluations using a global
> `alloc_tag_trace_key` static branch and an inline static key check
> pattern in the allocator hooks (`alloc_tag_add`, `alloc_tag_sub`, etc).
> This ensures that the trace events are entirely skipped when no
> allocation call sites are being actively traced, leaving only a NOP on
> the allocation fast path.
>
> Signed-off-by: Abhishek Bapat <[email protected]>
> ---
> include/linux/alloc_tag.h | 39 ++++++++----
> include/linux/codetag.h | 5 +-
> include/uapi/linux/alloc_tag.h | 9 +++
> mm/alloc_tag.c | 110 ++++++++++++++++++++++++++++++++-
> 4 files changed, 147 insertions(+), 16 deletions(-)
>
> diff --git a/include/linux/alloc_tag.h b/include/linux/alloc_tag.h
> index 2994934cf44a..dc86f8997476 100644
> --- a/include/linux/alloc_tag.h
> +++ b/include/linux/alloc_tag.h
> @@ -136,9 +136,36 @@ static inline bool mem_alloc_profiling_enabled(void)
> &mem_alloc_profiling_key);
> }
>
> +static inline void alloc_tag_set_inaccurate(struct alloc_tag *tag)
> +{
> + atomic_or(CODETAG_FLAG_INACCURATE, &tag->ct.flags);
> +}
> +
> +static inline bool alloc_tag_is_inaccurate(struct alloc_tag *tag)
> +{
> + return !!(atomic_read(&tag->ct.flags) & CODETAG_FLAG_INACCURATE);
> +}
> +
> +static inline void alloc_tag_set_traced(struct alloc_tag *tag)
> +{
> + atomic_or(CODETAG_FLAG_TRACE_ON, &tag->ct.flags);
> +}
> +
> +static inline void alloc_tag_clear_traced(struct alloc_tag *tag)
> +{
> + atomic_andnot(CODETAG_FLAG_TRACE_ON, &tag->ct.flags);
> +}
> +
> +static inline bool alloc_tag_is_traced(const struct alloc_tag *tag)
> +{
> + return !!(atomic_read(&tag->ct.flags) & CODETAG_FLAG_TRACE_ON);
> +}
> +
> static inline bool alloc_tag_trace_enabled(const struct alloc_tag *tag)
> {
> - return static_branch_unlikely(&alloc_tag_trace_key);
> + if (static_branch_unlikely(&alloc_tag_trace_key))
> + return tag && alloc_tag_is_traced(tag);
> + return false;
> }
>
> void alloc_tag_trace_mem_alloc(union codetag_ref *ref, struct alloc_tag *tag,
> @@ -255,16 +282,6 @@ static inline void alloc_tag_sub(union codetag_ref *ref,
> size_t bytes)
> ref->ct = NULL;
> }
>
> -static inline void alloc_tag_set_inaccurate(struct alloc_tag *tag)
> -{
> - tag->ct.flags |= CODETAG_FLAG_INACCURATE;
> -}
> -
> -static inline bool alloc_tag_is_inaccurate(struct alloc_tag *tag)
> -{
> - return !!(tag->ct.flags & CODETAG_FLAG_INACCURATE);
> -}
> -
> #define alloc_tag_record(p) ((p) = current->alloc_tag)
>
> #else /* CONFIG_MEM_ALLOC_PROFILING */
> diff --git a/include/linux/codetag.h b/include/linux/codetag.h
> index a25a085c2df1..f728295d50c0 100644
> --- a/include/linux/codetag.h
> +++ b/include/linux/codetag.h
> @@ -18,6 +18,7 @@ struct module;
>
> /* codetag flags */
> #define CODETAG_FLAG_INACCURATE (1 << 0)
> +#define CODETAG_FLAG_TRACE_ON (1 << 1)
>
> /*
> * An instance of this structure is created in a special ELF section at every
> @@ -25,7 +26,7 @@ struct module;
> * an array of these.
> */
> struct codetag {
> - unsigned int flags;
> + atomic_t flags;
> unsigned int lineno;
> const char *modname;
> const char *function;
> @@ -71,7 +72,7 @@ struct codetag_iterator {
> .function = __func__, \
> .filename = __FILE__, \
> .lineno = __LINE__, \
> - .flags = 0, \
> + .flags = ATOMIC_INIT(0), \
> }
>
> void codetag_lock_module_list(struct codetag_type *cttype);
> diff --git a/include/uapi/linux/alloc_tag.h b/include/uapi/linux/alloc_tag.h
> index 7d4618bea043..069ab8341e87 100644
> --- a/include/uapi/linux/alloc_tag.h
> +++ b/include/uapi/linux/alloc_tag.h
> @@ -85,9 +85,16 @@ struct allocinfo_get_at {
> struct allocinfo_tag_data data;
> };
>
> +struct allocinfo_toggle_traces {
> + /* inputs */
> + struct allocinfo_tag fields;
> + __u64 enable;
> +};
> +
> #define _ALLOCINFO_IOC_CONTENT_ID 0
> #define _ALLOCINFO_IOC_GET_AT 1
> #define _ALLOCINFO_IOC_GET_NEXT 2
> +#define _ALLOCINFO_IOC_TOGGLE_TRACE 3
>
> #define ALLOCINFO_IOC_BASE 0xA6
> #define ALLOCINFO_IOC_CONTENT_ID _IOR(ALLOCINFO_IOC_BASE,
> _ALLOCINFO_IOC_CONTENT_ID, \
> @@ -96,5 +103,7 @@ struct allocinfo_get_at {
> struct allocinfo_get_at)
> #define ALLOCINFO_IOC_GET_NEXT _IOR(ALLOCINFO_IOC_BASE,
> _ALLOCINFO_IOC_GET_NEXT, \
> struct allocinfo_tag_data)
> +#define ALLOCINFO_IOC_TOGGLE_TRACE _IOW(ALLOCINFO_IOC_BASE,
> _ALLOCINFO_IOC_TOGGLE_TRACE, \
> + struct allocinfo_toggle_traces)
>
> #endif /* _UAPI_ALLOC_TAG_H */
> diff --git a/mm/alloc_tag.c b/mm/alloc_tag.c
> index a5339767efd5..fb179321a2a1 100644
> --- a/mm/alloc_tag.c
> +++ b/mm/alloc_tag.c
> @@ -61,6 +61,15 @@ DEFINE_STATIC_KEY_FALSE(mem_profiling_compressed);
> DEFINE_STATIC_KEY_FALSE(alloc_tag_trace_key);
> EXPORT_SYMBOL(alloc_tag_trace_key);
>
> +static atomic_t alloc_tag_trace_cnt = ATOMIC_INIT(0);
> +
> +/*
> + * As `codetag_lock_module_list` is a read lock, we need an additional mutex
> + * to protect against the race conditions involved in the alloc tag trace
> + * toggle path.
> + */
> +static DEFINE_MUTEX(alloc_tag_trace_mutex);
> +
> struct alloc_tag_kernel_section kernel_tags = { NULL, 0 };
> unsigned long alloc_tag_ref_mask;
> int alloc_tag_ref_offs;
> @@ -297,7 +306,7 @@ static bool matches_filter(struct codetag *ct, struct
> allocinfo_filter *filter,
> return false;
>
> if (filter->mask & ALLOCINFO_FILTER_MASK_INACCURATE) {
> - inaccurate = !!(ct->flags & CODETAG_FLAG_INACCURATE);
> + inaccurate = alloc_tag_is_inaccurate(ct_to_alloc_tag(ct));
> if (inaccurate != !!(filter->inaccurate))
> return false;
> }
> @@ -444,6 +453,81 @@ static int allocinfo_ioctl_get_next(struct seq_file *m,
> void __user *arg)
> return ret;
> }
>
> +static bool alloc_tag_trace_toggle(struct alloc_tag *tag, bool enable)
> +{
> + if (enable) {
> + if (alloc_tag_is_traced(tag))
> + return false;
> +
> + alloc_tag_set_traced(tag);
> + if (atomic_fetch_inc(&alloc_tag_trace_cnt) == 0)
> + static_branch_enable(&alloc_tag_trace_key);
> + } else {
> + if (!alloc_tag_is_traced(tag))
> + return false;
> +
> + alloc_tag_clear_traced(tag);
> + if (atomic_dec_and_test(&alloc_tag_trace_cnt))
> + static_branch_disable(&alloc_tag_trace_key);
> + }
> +
> + return true;
> +}
> +
> +/*
> + * Toggles context capture for a specified allocation.
> + */
> +static int allocinfo_ioctl_toggle_trace(struct seq_file *m, void __user *arg)
> +{
> + struct allocinfo_toggle_traces params;
> + struct codetag_iterator iter;
> + struct codetag *ct;
> + int matches = 0, successes = 0, ret;
> +
> + if (!capable(CAP_SYS_ADMIN))
> + return -EPERM;
> +
> + if (copy_from_user(¶ms, arg, sizeof(params)))
> + return -EFAULT;
> +
> + codetag_lock_module_list(alloc_tag_cttype);
> +
> + struct allocinfo_filter filter = {
> + .mask = ALLOCINFO_FILTER_MASK_MODNAME |
> + ALLOCINFO_FILTER_MASK_FUNCTION |
> + ALLOCINFO_FILTER_MASK_FILENAME |
> + ALLOCINFO_FILTER_MASK_LINENO,
> + .fields = params.fields,
> + };
> +
> + iter = codetag_get_ct_iter(alloc_tag_cttype);
> +
> + /* Toggle tracing on all codetags that match */
> + while ((ct = codetag_next_ct(&iter))) {
> + if (matches_filter(ct, &filter, NULL, NULL)) {
> + matches++;
> +
> + mutex_lock(&alloc_tag_trace_mutex);
> + if (alloc_tag_trace_toggle(ct_to_alloc_tag(ct),
> !!params.enable))
> + successes++;
> + mutex_unlock(&alloc_tag_trace_mutex);
> + }
> + }
> +
> + if (matches == 0)
> + /* Nothing matched the filter */
> + ret = -ENOENT;
> + else if (successes == 0)
> + /* Items matched, but were already in the requested state */
> + ret = -EINVAL;
Why do we return -EINVAL in this path?
-EINVAL can mislead users into thinking their arguments are wrong.
That's not the case; the trace request they want is already active.
Maybe -EEXIST ?
Thanks
Best Regards
Hao
> + else
> + ret = 0;
> +
> + codetag_unlock_module_list(alloc_tag_cttype);
> +
> + return ret;
> +}
> +
> /*
> * Entry point ioctl function for /proc/allocinfo routing requests to fetch
> the
> * layout content ID, seek to a specific tag, or read sequential tags.
> @@ -464,6 +548,9 @@ static long allocinfo_ioctl(struct file *file, unsigned
> int cmd,
> case ALLOCINFO_IOC_GET_NEXT:
> ret = allocinfo_ioctl_get_next(file->private_data, arg);
> break;
> + case ALLOCINFO_IOC_TOGGLE_TRACE:
> + ret = allocinfo_ioctl_toggle_trace(file->private_data, arg);
> + break;
> default:
> ret = -ENOIOCTLCMD;
> break;
> @@ -493,8 +580,6 @@ static const struct proc_ops allocinfo_proc_ops = {
>
> void __alloc_tag_trace_hit(struct alloc_tag *tag)
> {
> - if (unlikely(!tag))
> - return;
> trace_alloc_tag_hit(tag);
> }
> EXPORT_SYMBOL(__alloc_tag_trace_hit);
> @@ -1043,6 +1128,24 @@ static int load_module(struct module *mod, struct
> codetag *start, struct codetag
> return 0;
> }
>
> +static void unload_module(struct module *mod, struct codetag *start, struct
> codetag *stop)
> +{
> + struct alloc_tag *start_tag = ct_to_alloc_tag(start);
> + struct alloc_tag *stop_tag = ct_to_alloc_tag(stop);
> + struct alloc_tag *tag;
> +
> + /*
> + * Turn tracing off for the tags of the module being unloaded. Without
> + * this, `alloc_tag_trace_cnt` would never reach zero and tracing would
> + * stay enabled forever.
> + *
> + * `alloc_tag_trace_mutex` is not needed here as this code path is
> + * protected by a `down_write(&cttype->mod_lock)`.
> + */
> + for (tag = start_tag; tag < stop_tag; tag++)
> + alloc_tag_trace_toggle(tag, false);
> +}
> +
> static void replace_module(struct module *mod, struct module *new_mod)
> {
> MA_STATE(mas, &mod_area_mt, 0, module_tags.size);
> @@ -1369,6 +1472,7 @@ static int __init alloc_tag_init(void)
> .alloc_section_mem = reserve_module_tags,
> .free_section_mem = release_module_tags,
> .module_load = load_module,
> + .module_unload = unload_module,
> .module_replaced = replace_module,
> #endif
> };
