Re: [PATCH v4 30/39] unwind_user/deferred: Make unwind deferral requests NMI-safe

[Josh Poimboeuf]

On Thu, Jan 23, 2025 at 11:17:34PM +0100, Peter Zijlstra wrote:
> On Thu, Jan 23, 2025 at 11:48:07AM -0800, Josh Poimboeuf wrote:
> > > > >       cookie = READ_ONCE(current->unwind_info.cookie);
> > > > >       do {
> > > > >               if (cookie)
> > > > >                       return cookie;
> > > > >               cookie = ctx_to_cookie(cpu, ctr+1);
> > > > >       } while (!try_cmpxchg64(&current->unwind_info.cookie, &cookie, 
> > > > > cookie));
> > 
> > Should not the 2nd argument be &zero?
> 
> This I suppose
> 
>       cookie = READ_ONCE(current->unwind_info.cookie);
>       do {
>               if (cookie)
>                       return cookie;
>       } while (!try_cmpxchg64(&current->unwind_info.cookie, &cookie,
>                               ctx_to_cookie(cpu, ctr+1)));

Here's what I have at the moment.  unwind_deferred_request() is vastly
simplified thanks to your suggestions, no more NMI-specific muck.
work->pending is replaced with work->cookie so the task work can use it
in case it crossed the cookie-clearing boundary (return-to-user with
IRQs disabled).

static inline bool nmi_supported(void)
{
        return IS_ENABLED(CONFIG_ARCH_HAVE_NMI_SAFE_CMPXCHG);
}

/*
 * Read the task context cookie, initializing it if this is the first request
 * since the most recent entry from user.
 */
static u64 get_cookie(void)
{
        struct unwind_task_info *info = &current->unwind_info;
        u64 ctr, cookie, old_cookie = 0;
        unsigned int cpu;

        cookie = READ_ONCE(info->cookie);
        if (cookie)
                return cookie;

        guard(irqsave)();

        cpu = raw_smp_processor_id();
        ctr = __this_cpu_read(unwind_ctx_ctr);

        cookie = ctx_to_cookie(cpu, ++ctr);

        /* Make sure the cookie is non-zero. */
        if (!cookie)
                cookie = ctx_to_cookie(cpu, ++ctr);

        /* Write the task cookie unless an NMI just now swooped in to do so. */
        if (!try_cmpxchg64(&info->cookie, &old_cookie, cookie))
                return old_cookie;

        /* Success, update the context counter. */
        __this_cpu_write(unwind_ctx_ctr, ctr);

        return cookie;
}

/*
 * Schedule a user space unwind to be done in task work before exiting the
 * kernel.
 *
 * The returned cookie output is a unique identifer for the current task entry
 * context.  Its value will also be passed to the callback function.  It can be
 * used to stitch kernel and user stack traces together in post-processing.
 *
 * It's valid to call this function multiple times for the same @work within
 * the same task entry context.  Each call will return the same cookie.  If the
 * callback is already pending, an error will be returned along with the
 * cookie.  If the callback is not pending because it has already been
 * previously called for the same entry context, it will be called again with
 * the same stack trace and cookie.
 *
 * Thus are three possible return scenarios:
 *
 *   * return != 0, *cookie == 0: the operation failed, no pending callback.
 *
 *   * return != 0, *cookie != 0: the callback is already pending. The cookie
 *     can still be used to correlate with the pending callback.
 *
 *   * return == 0, *cookie != 0: the callback queued successfully.  The
 *     callback is guaranteed to be called with the given cookie.
 */
int unwind_deferred_request(struct unwind_work *work, u64 *cookie)
{
        enum task_work_notify_mode mode = in_nmi() ? TWA_NMI_CURRENT : 
TWA_RESUME;
        u64 zero = 0;
        int ret;

        *cookie = 0;

        if (!current->mm || !user_mode(task_pt_regs(current)))
                return -EINVAL;

        if (!nmi_supported() && in_nmi())
                return -EINVAL;

        /* Return a valid @cookie even if the callback is already pending. */
        *cookie = READ_ONCE(work->cookie);
        if (*cookie)
                return -EEXIST;

        *cookie = get_cookie();

        /* Claim the work unless an NMI just now swooped in to do so. */
        if (!try_cmpxchg64(&work->cookie, &zero, *cookie))
                return -EEXIST;

        /* The work has been claimed, now schedule it. */
        ret = task_work_add(current, &work->work, mode);
        if (WARN_ON_ONCE(ret)) {
                WRITE_ONCE(work->cookie, 0);
                return ret;
        }

        return 0;
}

-- 
Josh