Re: [PATCH v2 2/4] alloc_tag: Introduce IOCTLs to toggle allocation tracepoints

[Hao Ge <hao.ge@linux.dev>]

Hi Abhishek


On 10/15/26 01:18, Abhishek Bapat wrote:
> On Tue, Oct 13, 2026 at 12:41 AM Hao Ge <hao.ge@linux.dev> wrote:
>> Hi Abhishek
>>
>>
>> On 2026/10/8 06:40, Abhishek Bapat wrote:
>>> Introduce a new IOCTL (`ALLOCINFO_IOC_TOGGLE_TRACE`) to selectively
>>> toggle tracing on 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 the 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 <abhishekbapat@google.com>
>>> ---
>>>   Documentation/mm/allocation-profiling.rst |  20 ++++
>>>   include/linux/alloc_tag.h                 |  47 +++++---
>>>   include/linux/codetag.h                   |   5 +-
>>>   include/uapi/linux/alloc_tag.h            |  13 +++
>>>   mm/alloc_tag.c                            | 128 ++++++++++++++++++++--
>>>   5 files changed, 187 insertions(+), 26 deletions(-)
>>>
>>> diff --git a/Documentation/mm/allocation-profiling.rst b/Documentation/mm/allocation-profiling.rst
>>> index b241450394df..63b668466342 100644
>>> --- a/Documentation/mm/allocation-profiling.rst
>>> +++ b/Documentation/mm/allocation-profiling.rst
>>> @@ -149,6 +149,26 @@ expose the full call stack and the lifetime of individual allocations:
>>>   - slab and percpu: the address of the object's `codetag_ref`
>>>   - page allocator: the `struct page` pointer of the head page
>>>
>>> +Events are only emitted for tags that have tracing turned on and while profiling
>>> +is enabled. When no tag is traced, the hooks are behind a static branch.
>>> +
>>> +Enabling Tracing
>>> +----------------
>>> +
>>> +Tracing is toggled on and off per tag with the `ALLOCINFO_IOC_TOGGLE_TRACE`
>>> +ioctl on `/proc/allocinfo`. It takes the same filter as used by
>>> +`ALLOCINFO_IOC_GET_AT`; an empty mask selects all tags. Size limits in the
>>> +filter are checked once, against the tag's size at the time of the call. Tags
>>> +already in the requested state are skipped and a -ENOENT is returned if no
>>> +tag matched.
>>> +
>>> +Tracing is turned off for a module's tags when the module is unloaded.
>>> +
>>> +Memory allocated before a tag was enabled has no alloc event but will still
>>> +produce a free event. Likewise, memory alive when tag is disabled will not
>>> +produce a free event. Tools should expect unmatched frees after enabling and
>>> +missing frees after disabling.
>>> +
>>>   Correlating the events
>>>   ----------------------
>>>
>>> diff --git a/include/linux/alloc_tag.h b/include/linux/alloc_tag.h
>>> index 49825177d8c8..94697b28053a 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 bool alloc_tag_trace_enabled(void)
>>> +static inline void alloc_tag_set_inaccurate(struct alloc_tag *tag)
>>>   {
>>> -     return static_branch_unlikely(&alloc_tag_trace_key);
>>> +     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)
>>> +{
>>> +     if (static_branch_unlikely(&alloc_tag_trace_key))
>>> +             return tag && alloc_tag_is_traced(tag);
>>> +     return false;
>>>   }
>>>
>>>   void alloc_tag_trace_mem_alloc(const void *ptr, struct alloc_tag *tag,
>>> @@ -151,7 +178,7 @@ void __alloc_tag_trace_hit(struct alloc_tag *tag);
>>>
>>>   static __always_inline void alloc_tag_trace_hit(struct alloc_tag *tag)
>>>   {
>>> -     if (alloc_tag_trace_enabled())
>>> +     if (alloc_tag_trace_enabled(tag))
>>>                __alloc_tag_trace_hit(tag);
>>>   }
>>>
>>> @@ -225,7 +252,7 @@ static inline void alloc_tag_add(union codetag_ref *ref, struct alloc_tag *tag,
>>>        if (likely(alloc_tag_ref_set(ref, tag))) {
>>>                this_cpu_add(tag->counters->bytes, bytes);
>>>
>>> -             if (alloc_tag_trace_enabled())
>>> +             if (alloc_tag_trace_enabled(tag))
>>>                        /* Trace successful allocs with their unique ptr */
>>>                        alloc_tag_trace_mem_alloc(ptr, tag, bytes);
>>>        }
>>> @@ -249,23 +276,13 @@ static inline void alloc_tag_sub(union codetag_ref *ref, size_t bytes, const voi
>>>        this_cpu_sub(tag->counters->bytes, bytes);
>>>        this_cpu_dec(tag->counters->calls);
>>>
>>> -     if (alloc_tag_trace_enabled())
>>> +     if (alloc_tag_trace_enabled(tag))
>>>                /* Trace frees with their unique ptr */
>>>                alloc_tag_trace_mem_free(ptr, tag, 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..dafcba2f7e2c 100644
>>> --- a/include/uapi/linux/alloc_tag.h
>>> +++ b/include/uapi/linux/alloc_tag.h
>>> @@ -85,9 +85,20 @@ struct allocinfo_get_at {
>>>        struct allocinfo_tag_data data;
>>>   };
>>>
>>> +struct allocinfo_toggle_traces {
>>> +     /* inputs */
>>> +     /*
>>> +      * Tags matching the filter are toggled; an empty mask selects all tags.
>>> +      * Size criteria are checked once, against the tag's size at call time.
>>> +      */
>>> +     struct allocinfo_filter filter;
>>> +     __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 +107,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 ca2412a67312..262f00f2bba9 100644
>>> --- a/mm/alloc_tag.c
>>> +++ b/mm/alloc_tag.c
>>> @@ -60,6 +60,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;
>>> @@ -296,7 +305,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;
>>>        }
>>> @@ -317,6 +326,23 @@ static bool matches_filter(struct codetag *ct, struct allocinfo_filter *filter,
>>>        return true;
>>>   }
>>>
>>> +/*
>>> + * Checks that a user supplied filter only uses known fields and that its size
>>> + * range, if any, is not inverted.
>>> + */
>>> +static bool allocinfo_filter_valid(const struct allocinfo_filter *filter)
>>> +{
>>> +     if (filter->mask & ~ALLOCINFO_FILTER_MASKS)
>>> +             return false;
>>> +
>>> +     if ((filter->mask & ALLOCINFO_FILTER_MASK_MIN_SIZE) &&
>>> +         (filter->mask & ALLOCINFO_FILTER_MASK_MAX_SIZE) &&
>>> +         filter->min_size > filter->max_size)
>>> +             return false;
>>> +
>>> +     return true;
>>> +}
>>> +
>>>   /*
>>>    * Seeks the ioctl iterator to the specified 0-indexed tag position, reads its
>>>    * profiling data and returns it to userspace.
>>> @@ -333,12 +359,7 @@ static int allocinfo_ioctl_get_at(struct seq_file *m, void __user *arg)
>>>        if (copy_from_user(&params, arg, sizeof(params)))
>>>                return -EFAULT;
>>>
>>> -     if (params.filter.mask & ~ALLOCINFO_FILTER_MASKS)
>>> -             return -EINVAL;
>>> -
>>> -     if ((params.filter.mask & ALLOCINFO_FILTER_MASK_MIN_SIZE) &&
>>> -         (params.filter.mask & ALLOCINFO_FILTER_MASK_MAX_SIZE) &&
>>> -         params.filter.min_size > params.filter.max_size)
>>> +     if (!allocinfo_filter_valid(&params.filter))
>>>                return -EINVAL;
>>>
>>>        priv = m->private;
>>> @@ -443,6 +464,75 @@ static int allocinfo_ioctl_get_next(struct seq_file *m, void __user *arg)
>>>        return ret;
>>>   }
>>>
>>> +static void alloc_tag_trace_toggle(struct alloc_tag *tag, bool enable)
>>> +{
>>> +     if (enable) {
>>> +             if (alloc_tag_is_traced(tag))
>>> +                     return;
>>> +
>>> +             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;
>>> +
>>> +             alloc_tag_clear_traced(tag);
>>> +             if (atomic_dec_and_test(&alloc_tag_trace_cnt))
>>> +                     static_branch_disable(&alloc_tag_trace_key);
>>> +     }
>>> +}
>>> +
>>> +/*
>>> + * Toggles tracing on every allocation tag that matches the user supplied
>>> + * filter. An empty filter mask selects all tags, same as for
>>> + * ALLOCINFO_IOC_GET_AT. Tags already in the requested state are not an error.
>>> + */
>>> +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;
>>> +     struct alloc_tag_counters counters;
>>> +     bool fetched_counters;
>>> +     int matches = 0, ret;
>>> +
>>> +     if (!capable(CAP_SYS_ADMIN))
>>> +             return -EPERM;
>>> +
>>> +     if (copy_from_user(&params, arg, sizeof(params)))
>>> +             return -EFAULT;
>>> +
>>> +     if (!allocinfo_filter_valid(&params.filter))
>>> +             return -EINVAL;
>>> +
>>> +     codetag_lock_module_list(alloc_tag_cttype);
>>> +
>>> +     iter = codetag_get_ct_iter(alloc_tag_cttype);
>>> +
>>> +     /* Toggle tracing on all codetags that match */
>>> +     while ((ct = codetag_next_ct(&iter))) {
>>> +             fetched_counters = false;
>>> +             if (matches_filter(ct, &params.filter, &counters, &fetched_counters)) {
>>> +                     matches++;
>>> +
>>> +                     mutex_lock(&alloc_tag_trace_mutex);
>> Could we grab and drop alloc_tag_trace_mutex outside the loop?
> Well, that is possible. But I did it this way to minimize the critical section.
>
> Since this mutex is specific to the trace ioctl, I'm comfortable doing
> what you recommended too. It would just mean that one thread would be
> blocked until another thread fully finishes its loop, which is not too
> bad. I'll let Suren comment on this as well.


Hmm, sorry, I misread this part of the code. (sorry for this) . Putting 
|mutex_lock| outside the loop is not a good idea.

It greatly increases the critical section.

My original concern was that we need to mutex_lock and mutex_unlock 
every time we get a match.

As the function comment states, An empty filter mask selects all tags, 
that means we need a

mutex_lock / mutex_unlock for every alloc_tag.

Let's see what Suren thinks.


Thanks

Best Regards

Hao

>>
>>> +                     alloc_tag_trace_toggle(ct_to_alloc_tag(ct), !!params.enable);
>>> +                     mutex_unlock(&alloc_tag_trace_mutex);
>>> +             }
>>> +     }
>>> +
>>> +     if (matches == 0)
>>> +             /* Nothing matched the filter */
>>> +             ret = -ENOENT;
>>> +     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.
>>> @@ -463,6 +553,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;
>>> @@ -492,8 +585,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);
>>> @@ -1031,6 +1122,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);
>>> @@ -1357,6 +1466,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
>>>        };