[RFC,v2,1/3] kunit: Add kunit_add_action() to defer a call until test exit

Commit Message

David Gow March 31, 2023, 8:04 a.m. UTC

Many uses of the KUnit resource system are intended to simply defer
calling a function until the test exits (be it due to success or
failure). The existing kunit_alloc_resource() function is often used for
this, but was awkward to use (requiring passing NULL init functions, etc),
and returned a resource without incrementing its reference count, which
-- while okay for this use-case -- could cause problems in others.

Instead, introduce a simple kunit_add_action() API: a simple function
(returning nothing, accepting a single void* argument) can be scheduled
to be called when the test exits. Deferred actions are called in the
opposite order to that which they were registered.

This mimics the devres API, devm_add_action(), and also provides
kunit_remove_action(), to cancel a deferred action, and
kunit_release_action() to trigger one early.

This is implemented as a resource under the hood, so the ordering
between resource cleanup and deferred functions is maintained.

Signed-off-by: David Gow <[e-mail]>
---

Changes since RFC v1:
https://lore.kernel.org/linux-kselftest/[e-mail]/
- Rename functions to better match the devm_* APIs. (Thanks Maxime)
- Embed the kunit_resource in struct kunit_action_ctx to avoid an extra
  allocation (Thanks Benjamin)
- Use 'struct kunit_action_ctx' as the type for cancellation tokens
  (Thanks Benjamin)
- Add tests.

---
 include/kunit/resource.h |  89 ++++++++++++++++++++++++++++
 lib/kunit/kunit-test.c   | 123 ++++++++++++++++++++++++++++++++++++++-
 lib/kunit/resource.c     |  99 +++++++++++++++++++++++++++++++
 3 files changed, 310 insertions(+), 1 deletion(-)

Comments

Maxime Ripard April 4, 2023, 1:32 p.m. UTC | #1

Hi David,

Looks great, thanks for sending a second version

On Fri, Mar 31, 2023 at 04:04:09PM +0800, David Gow wrote:
> Many uses of the KUnit resource system are intended to simply defer
> calling a function until the test exits (be it due to success or
> failure). The existing kunit_alloc_resource() function is often used for
> this, but was awkward to use (requiring passing NULL init functions, etc),
> and returned a resource without incrementing its reference count, which
> -- while okay for this use-case -- could cause problems in others.
> 
> Instead, introduce a simple kunit_add_action() API: a simple function
> (returning nothing, accepting a single void* argument) can be scheduled
> to be called when the test exits. Deferred actions are called in the
> opposite order to that which they were registered.
> 
> This mimics the devres API, devm_add_action(), and also provides
> kunit_remove_action(), to cancel a deferred action, and
> kunit_release_action() to trigger one early.
> 
> This is implemented as a resource under the hood, so the ordering
> between resource cleanup and deferred functions is maintained.
> 
> Signed-off-by: David Gow <[e-mail]>
> ---
> 
> Changes since RFC v1:
> https://lore.kernel.org/linux-kselftest/[e-mail]/
> - Rename functions to better match the devm_* APIs. (Thanks Maxime)
> - Embed the kunit_resource in struct kunit_action_ctx to avoid an extra
>   allocation (Thanks Benjamin)
> - Use 'struct kunit_action_ctx' as the type for cancellation tokens
>   (Thanks Benjamin)
> - Add tests.
> 
> ---
>  include/kunit/resource.h |  89 ++++++++++++++++++++++++++++
>  lib/kunit/kunit-test.c   | 123 ++++++++++++++++++++++++++++++++++++++-
>  lib/kunit/resource.c     |  99 +++++++++++++++++++++++++++++++
>  3 files changed, 310 insertions(+), 1 deletion(-)
> 
> diff --git a/include/kunit/resource.h b/include/kunit/resource.h
> index c0d88b318e90..15efd8924666 100644
> --- a/include/kunit/resource.h
> +++ b/include/kunit/resource.h
> @@ -387,4 +387,93 @@ static inline int kunit_destroy_named_resource(struct kunit *test,
>   */
>  void kunit_remove_resource(struct kunit *test, struct kunit_resource *res);
>  
> +typedef void (*kunit_defer_function_t)(void *ctx);
> +
> +/* An opaque token to a deferred action. */
> +struct kunit_action_ctx;
> +
> +/**
> + * kunit_add_action() - Defer an 'action' (function call) until the test ends.
> + * @test: Test case to associate the action with.
> + * @func: The function to run on test exit
> + * @ctx: Data passed into @func
> + * @internal_gfp: gfp to use for internal allocations, if unsure, use GFP_KERNEL
> + *
> + * Defer the execution of a function until the test exits, either normally or
> + * due to a failure.  @ctx is passed as additional context. All functions
> + * registered with kunit_add_action() will execute in the opposite order to that
> + * they were registered in.
> + *
> + * This is useful for cleaning up allocated memory and resources.
> + *
> + * Returns:
> + *   An opaque "cancellation token", or NULL on error. Pass this token to
> + *   kunit_remove_action_token() in order to cancel the deferred execution of
> + *   func().
> + */
> +struct kunit_action_ctx *kunit_add_action(struct kunit *test, kunit_defer_function_t func,
> +		      void *ctx, gfp_t internal_gfp);

Do we expect any other context than GFP_KERNEL?

If so, then maybe we can have kunit_add_action() assume GFP_KERNEL and
add a variant for the odd case where we would actually need a different
GFP flag.

> +/**
> + * kunit_remove_action_token() - Cancel a deferred action.
> + * @test: Test case the action is associated with.
> + * @cancel_token: The cancellation token returned by kunit_add_action()
> + *
> + * Prevent an action deferred using kunit_add_action() from executing when the
> + * test ends.
> + *
> + * You can also use the (test, function, context) triplet to remove an action
> + * with kunit_remove_action().
> + */
> +void kunit_remove_action_token(struct kunit *test, struct kunit_action_ctx *cancel_token);

It's not clear to me why we still need the token. If
kunit_remove_action() works fine, why would we need to store the token?

Can't we just make kunit_add_action() return an int to indicate whether
it failed or not, and that's it?

> +/**
> + * kunit_release_action_token() - Trigger a deferred action immediately.
> + * @test: Test case the action is associated with.
> + * @cancel_token: The cancellation token returned by kunit_add_action()
> + *
> + * Execute an action immediately, instead of waiting for the test to end.
> + *
> + * You can also use the (test, function, context) triplet to trigger an action
> + * with kunit_release_action().
> + */
> +void kunit_release_action_token(struct kunit *test, struct kunit_action_ctx *cancel_token);
> +
> +/**
> + * kunit_remove_action() - Cancel a matching deferred action.
> + * @test: Test case the action is associated with.
> + * @func: The deferred function to cancel.
> + * @ctx: The context passed to the deferred function to trigger.
> + *
> + * Prevent an action deferred via kunit_add_action() from executing when the
> + * test terminates..
> + * Unlike kunit_remove_action_token(), this takes the (func, ctx) pair instead of
> + * the cancellation token. If that function/context pair was deferred multiple
> + * times, only the most recent one will be cancelled.
> + */
> +void kunit_remove_action(struct kunit *test,
> +			 kunit_defer_function_t func,
> +			 void *ctx);
> +
> +/**
> + * kunit_release_action() - Run a matching action call immediately.
> + * @test: Test case the action is associated with.
> + * @func: The deferred function to trigger.
> + * @ctx: The context passed to the deferred function to trigger.
> + *
> + * Execute a function deferred via kunit_add_action()) immediately, rather than
> + * when the test ends.
> + * Unlike kunit_release_action_token(), this takes the (func, ctx) pair instead of
> + * the cancellation token. If that function/context pair was deferred multiple
> + * times, it will only be executed once here. The most recent deferral will
> + * no longer execute when the test ends.
> + *
> + * kunit_release_action(test, func, ctx);
> + * is equivalent to
> + * func(ctx);
> + * kunit_remove_action(test, func, ctx);
> + */
> +void kunit_release_action(struct kunit *test,
> +			  kunit_defer_function_t func,
> +			  void *ctx);
>  #endif /* _KUNIT_RESOURCE_H */
> diff --git a/lib/kunit/kunit-test.c b/lib/kunit/kunit-test.c
> index b63595d3e241..eaca1b133922 100644
> --- a/lib/kunit/kunit-test.c
> +++ b/lib/kunit/kunit-test.c
> @@ -111,7 +111,7 @@ struct kunit_test_resource_context {
>  	struct kunit test;
>  	bool is_resource_initialized;
>  	int allocate_order[2];
> -	int free_order[2];
> +	int free_order[4];
>  };
>  
>  static int fake_resource_init(struct kunit_resource *res, void *context)
> @@ -402,6 +402,123 @@ static void kunit_resource_test_named(struct kunit *test)
>  	KUNIT_EXPECT_TRUE(test, list_empty(&test->resources));
>  }
>  
> +static void increment_int(void *ctx)
> +{
> +	int *i = (int *)ctx;
> +	(*i)++;
> +}
> +
> +static void kunit_resource_test_action(struct kunit *test)
> +{
> +	int num_actions = 0;
> +
> +	kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
> +	KUNIT_EXPECT_EQ(test, num_actions, 0);
> +	kunit_cleanup(test);
> +	KUNIT_EXPECT_EQ(test, num_actions, 1);
> +
> +	/* Once we've cleaned up, the action queue is empty. */
> +	kunit_cleanup(test);
> +	KUNIT_EXPECT_EQ(test, num_actions, 1);
> +
> +	/* Check the same function can be deferred multiple times. */
> +	kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
> +	kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
> +	kunit_cleanup(test);
> +	KUNIT_EXPECT_EQ(test, num_actions, 3);
> +}
> +static void kunit_resource_test_remove_action(struct kunit *test)
> +{
> +	int num_actions = 0;
> +	struct kunit_action_ctx *cancel_token;
> +	struct kunit_action_ctx *cancel_token2;
> +
> +	cancel_token = kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
> +	KUNIT_EXPECT_EQ(test, num_actions, 0);
> +
> +	kunit_remove_action_token(test, cancel_token);
> +	kunit_cleanup(test);
> +	KUNIT_EXPECT_EQ(test, num_actions, 0);
> +
> +	/* Check calls from the same function/context pair can be cancelled independently*/
> +	cancel_token = kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
> +	cancel_token2 = kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
> +	kunit_remove_action_token(test, cancel_token);
> +	kunit_cleanup(test);
> +	KUNIT_EXPECT_EQ(test, num_actions, 1);
> +
> +	/* Also check that we can cancel just one of the identical function/context pairs. */
> +	cancel_token = kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
> +	cancel_token2 = kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
> +	kunit_remove_action(test, increment_int, &num_actions);
> +	kunit_cleanup(test);
> +	KUNIT_EXPECT_EQ(test, num_actions, 2);
> +}
> +static void kunit_resource_test_release_action(struct kunit *test)
> +{
> +	int num_actions = 0;
> +	struct kunit_action_ctx *cancel_token;
> +	struct kunit_action_ctx *cancel_token2;
> +
> +	cancel_token = kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
> +	KUNIT_EXPECT_EQ(test, num_actions, 0);
> +	/* Runs immediately on trigger. */
> +	kunit_release_action_token(test, cancel_token);
> +	KUNIT_EXPECT_EQ(test, num_actions, 1);
> +
> +	/* Doesn't run again on test exit. */
> +	kunit_cleanup(test);
> +	KUNIT_EXPECT_EQ(test, num_actions, 1);
> +
> +	/* Check calls from the same function/context pair can be triggered independently*/
> +	cancel_token = kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
> +	cancel_token2 = kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
> +	kunit_release_action_token(test, cancel_token);
> +	KUNIT_EXPECT_EQ(test, num_actions, 2);
> +	kunit_cleanup(test);
> +	KUNIT_EXPECT_EQ(test, num_actions, 3);
> +
> +	/* Also check that we can trigger just one of the identical function/context pairs. */
> +	kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
> +	kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
> +	kunit_release_action(test, increment_int, &num_actions);
> +	KUNIT_EXPECT_EQ(test, num_actions, 4);
> +	kunit_cleanup(test);
> +	KUNIT_EXPECT_EQ(test, num_actions, 5);
> +}
> +static void action_order_1(void *ctx)
> +{
> +	struct kunit_test_resource_context *res_ctx = (struct kunit_test_resource_context *)ctx;
> +
> +	KUNIT_RESOURCE_TEST_MARK_ORDER(res_ctx, free_order, 1);
> +	kunit_log(KERN_INFO, current->kunit_test, "action_order_1");
> +}
> +static void action_order_2(void *ctx)
> +{
> +	struct kunit_test_resource_context *res_ctx = (struct kunit_test_resource_context *)ctx;
> +
> +	KUNIT_RESOURCE_TEST_MARK_ORDER(res_ctx, free_order, 2);
> +	kunit_log(KERN_INFO, current->kunit_test, "action_order_2");
> +}
> +static void kunit_resource_test_action_ordering(struct kunit *test)
> +{
> +	struct kunit_test_resource_context *ctx = test->priv;
> +	struct kunit_action_ctx *cancel_token;
> +
> +	kunit_add_action(test, action_order_1, ctx, GFP_KERNEL);
> +	cancel_token = kunit_add_action(test, action_order_2, ctx, GFP_KERNEL);
> +	kunit_add_action(test, action_order_1, ctx, GFP_KERNEL);
> +	kunit_add_action(test, action_order_2, ctx, GFP_KERNEL);
> +	kunit_remove_action(test, action_order_1, ctx);
> +	kunit_release_action_token(test, cancel_token);
> +	kunit_cleanup(test);
> +
> +	/* [2 is triggered] [2], [(1 is cancelled)] [1] */
> +	KUNIT_EXPECT_EQ(test, ctx->free_order[0], 2);
> +	KUNIT_EXPECT_EQ(test, ctx->free_order[1], 2);
> +	KUNIT_EXPECT_EQ(test, ctx->free_order[2], 1);
> +}
> +
>  static int kunit_resource_test_init(struct kunit *test)
>  {
>  	struct kunit_test_resource_context *ctx =
> @@ -433,6 +550,10 @@ static struct kunit_case kunit_resource_test_cases[] = {
>  	KUNIT_CASE(kunit_resource_test_proper_free_ordering),
>  	KUNIT_CASE(kunit_resource_test_static),
>  	KUNIT_CASE(kunit_resource_test_named),
> +	KUNIT_CASE(kunit_resource_test_action),
> +	KUNIT_CASE(kunit_resource_test_remove_action),
> +	KUNIT_CASE(kunit_resource_test_release_action),
> +	KUNIT_CASE(kunit_resource_test_action_ordering),
>  	{}
>  };
>  
> diff --git a/lib/kunit/resource.c b/lib/kunit/resource.c
> index c414df922f34..824cf91e306d 100644
> --- a/lib/kunit/resource.c
> +++ b/lib/kunit/resource.c
> @@ -77,3 +77,102 @@ int kunit_destroy_resource(struct kunit *test, kunit_resource_match_t match,
>  	return 0;
>  }
>  EXPORT_SYMBOL_GPL(kunit_destroy_resource);
> +
> +struct kunit_action_ctx {
> +	struct kunit_resource res;
> +	kunit_defer_function_t func;
> +	void *ctx;
> +};
> +
> +static void __kunit_action_free(struct kunit_resource *res)
> +{
> +	struct kunit_action_ctx *action_ctx = container_of(res, struct kunit_action_ctx, res);
> +
> +	action_ctx->func(action_ctx->ctx);
> +}
> +
> +struct kunit_action_ctx *kunit_add_action(struct kunit *test, kunit_defer_function_t func,
> +					 void *ctx, gfp_t internal_gfp)
> +{
> +	struct kunit_action_ctx *action_ctx;
> +
> +	KUNIT_ASSERT_NOT_NULL_MSG(test, func, "Tried to action a NULL function!");
> +
> +	action_ctx = kzalloc(sizeof(*action_ctx), internal_gfp);
> +	if (!action_ctx)
> +		return NULL;
> +
> +	action_ctx->func = func;
> +	action_ctx->ctx = ctx;
> +
> +	action_ctx->res.should_kfree = true;
> +	/* As init is NULL, this cannot fail. */
> +	__kunit_add_resource(test, NULL, __kunit_action_free, &action_ctx->res, action_ctx);
> +
> +	return action_ctx;
> +}

One thing worth pointing is that if kunit_add_action() fails, the
cleanup function passed as an argument won't run.

So, if the kzalloc call ever fails, patch 2 will leak its res->data()
resource for example.

devm (and drmm) handles this using a variant called
devm_add_action_or_reset, we should either provide the same variant or
just go for that behavior by default.

Maxime

Benjamin Berg April 4, 2023, 5:55 p.m. UTC | #2

Hi,

On Tue, 2023-04-04 at 15:32 +0200, Maxime Ripard wrote:
> [SNIP]
> > +/**
> > + * kunit_add_action() - Defer an 'action' (function call) until the test ends.
> > + * @test: Test case to associate the action with.
> > + * @func: The function to run on test exit
> > + * @ctx: Data passed into @func
> > + * @internal_gfp: gfp to use for internal allocations, if unsure, use GFP_KERNEL
> > + *
> > + * Defer the execution of a function until the test exits, either normally or
> > + * due to a failure.  @ctx is passed as additional context. All functions
> > + * registered with kunit_add_action() will execute in the opposite order to that
> > + * they were registered in.
> > + *
> > + * This is useful for cleaning up allocated memory and resources.
> > + *
> > + * Returns:
> > + *   An opaque "cancellation token", or NULL on error. Pass this token to
> > + *   kunit_remove_action_token() in order to cancel the deferred execution of
> > + *   func().
> > + */
> > +struct kunit_action_ctx *kunit_add_action(struct kunit *test, kunit_defer_function_t func,
> > +                     void *ctx, gfp_t internal_gfp);
> 
> Do we expect any other context than GFP_KERNEL?
> 
> If so, then maybe we can have kunit_add_action() assume GFP_KERNEL and
> add a variant for the odd case where we would actually need a different
> GFP flag.

Does anything other than GFP_KERNEL make sense? I would assume these
functions should only ever be called from a kunit context, i.e. the
passed test is guaranteed to be identical to the value returned by
kunit_get_current_test().

That said, I am happy with merging this in this form. I feel the right
thing here is a patch (with corresponding spatch) that changes all of
the related APIs to remove the gfp argument.

> > +/**
> > + * kunit_remove_action_token() - Cancel a deferred action.
> > + * @test: Test case the action is associated with.
> > + * @cancel_token: The cancellation token returned by kunit_add_action()
> > + *
> > + * Prevent an action deferred using kunit_add_action() from executing when the
> > + * test ends.
> > + *
> > + * You can also use the (test, function, context) triplet to remove an action
> > + * with kunit_remove_action().
> > + */
> > +void kunit_remove_action_token(struct kunit *test, struct kunit_action_ctx *cancel_token);
> 
> It's not clear to me why we still need the token. If
> kunit_remove_action() works fine, why would we need to store the token?
> 
> Can't we just make kunit_add_action() return an int to indicate whether
> it failed or not, and that's it?
> 
> > [SNIP]
> 
> One thing worth pointing is that if kunit_add_action() fails, the
> cleanup function passed as an argument won't run.
> 
> So, if the kzalloc call ever fails, patch 2 will leak its res->data()
> resource for example.
> 
> devm (and drmm) handles this using a variant called
> devm_add_action_or_reset, we should either provide the same variant or
> just go for that behavior by default.

Both version of the function would need a return value. An alternative
might be to make assertions part of the API. But as with dropping the
gfp argument, that seems like a more intrusive change that needs to
happen independently.

Anyway, I am fine with action_or_reset as the default and possibly the
only behaviour. I expect that every API user will want an assertion
that checks for failure here anyway.

Benjamin


If kunit_* functions can assert in error conditions, then the example

void test_func(struct kunit *test)
{
  char u8 *buf = kunit_kzalloc(test, 1024, GFP_KERNEL);
  struct sk_buff *skb_a;
  struct sk_buff *skb_b;
  /* Further variables */

  KUNIT_ASSERT_NOT_NULL(test, buf);

  skb_a = skb_alloc(1024, GFP_KERNEL);
  KUNIT_ASSERT_NOT_NULL(test, skb_a);
  if (kunit_add_cleanup(test, (kunit_defer_function_t) kfree_skb, skb_a))
    KUNIT_ASSERT_FAILURE("Failed to add cleanup");

  /* Or, maybe: */
  skb_b = skb_alloc(1024, GFP_KERNEL);
  KUNIT_ASSERT_NOT_NULL(test, skb_b);
  KUNIT_ASSERT_EQ(test, 0,
                  kunit_add_cleanup(test,
                                    (kunit_defer_function_t) kfree_skb,
                                    skb_b));

  /* run code that may assert */
}


could be shortened to (with a trivial kunit_skb_alloc helper)

void test_func(struct kunit *test)
{
  char u8 *buf = kunit_kzalloc(test, 1024, GFP_KERNEL);
  struct sk_buff *skb_a = kunit_skb_alloc(1024, GFP_KERNEL);
  struct sk_buff *skb_b = kunit_skb_alloc(1024, GFP_KERNEL);
  /* Further variables */

  /* run code that may assert */
}

I should just post a patch for the existing API and see what people say
then ...

David Gow April 5, 2023, 7:47 a.m. UTC | #3

On Tue, 4 Apr 2023 at 21:32, Maxime Ripard <[e-mail]> wrote:
>
> Hi David,
>
> Looks great, thanks for sending a second version
>
> On Fri, Mar 31, 2023 at 04:04:09PM +0800, David Gow wrote:
> > Many uses of the KUnit resource system are intended to simply defer
> > calling a function until the test exits (be it due to success or
> > failure). The existing kunit_alloc_resource() function is often used for
> > this, but was awkward to use (requiring passing NULL init functions, etc),
> > and returned a resource without incrementing its reference count, which
> > -- while okay for this use-case -- could cause problems in others.
> >
> > Instead, introduce a simple kunit_add_action() API: a simple function
> > (returning nothing, accepting a single void* argument) can be scheduled
> > to be called when the test exits. Deferred actions are called in the
> > opposite order to that which they were registered.
> >
> > This mimics the devres API, devm_add_action(), and also provides
> > kunit_remove_action(), to cancel a deferred action, and
> > kunit_release_action() to trigger one early.
> >
> > This is implemented as a resource under the hood, so the ordering
> > between resource cleanup and deferred functions is maintained.
> >
> > Signed-off-by: David Gow <[e-mail]>
> > ---
> >
> > Changes since RFC v1:
> > https://lore.kernel.org/linux-kselftest/[e-mail]/
> > - Rename functions to better match the devm_* APIs. (Thanks Maxime)
> > - Embed the kunit_resource in struct kunit_action_ctx to avoid an extra
> >   allocation (Thanks Benjamin)
> > - Use 'struct kunit_action_ctx' as the type for cancellation tokens
> >   (Thanks Benjamin)
> > - Add tests.
> >
> > ---
> >  include/kunit/resource.h |  89 ++++++++++++++++++++++++++++
> >  lib/kunit/kunit-test.c   | 123 ++++++++++++++++++++++++++++++++++++++-
> >  lib/kunit/resource.c     |  99 +++++++++++++++++++++++++++++++
> >  3 files changed, 310 insertions(+), 1 deletion(-)
> >
> > diff --git a/include/kunit/resource.h b/include/kunit/resource.h
> > index c0d88b318e90..15efd8924666 100644
> > --- a/include/kunit/resource.h
> > +++ b/include/kunit/resource.h
> > @@ -387,4 +387,93 @@ static inline int kunit_destroy_named_resource(struct kunit *test,
> >   */
> >  void kunit_remove_resource(struct kunit *test, struct kunit_resource *res);
> >
> > +typedef void (*kunit_defer_function_t)(void *ctx);
> > +
> > +/* An opaque token to a deferred action. */
> > +struct kunit_action_ctx;
> > +
> > +/**
> > + * kunit_add_action() - Defer an 'action' (function call) until the test ends.
> > + * @test: Test case to associate the action with.
> > + * @func: The function to run on test exit
> > + * @ctx: Data passed into @func
> > + * @internal_gfp: gfp to use for internal allocations, if unsure, use GFP_KERNEL
> > + *
> > + * Defer the execution of a function until the test exits, either normally or
> > + * due to a failure.  @ctx is passed as additional context. All functions
> > + * registered with kunit_add_action() will execute in the opposite order to that
> > + * they were registered in.
> > + *
> > + * This is useful for cleaning up allocated memory and resources.
> > + *
> > + * Returns:
> > + *   An opaque "cancellation token", or NULL on error. Pass this token to
> > + *   kunit_remove_action_token() in order to cancel the deferred execution of
> > + *   func().
> > + */
> > +struct kunit_action_ctx *kunit_add_action(struct kunit *test, kunit_defer_function_t func,
> > +                   void *ctx, gfp_t internal_gfp);
>
> Do we expect any other context than GFP_KERNEL?
>
> If so, then maybe we can have kunit_add_action() assume GFP_KERNEL and
> add a variant for the odd case where we would actually need a different
> GFP flag.
>

That'd be fine. The only cases which don't directly pass GFP_KERNEL in
are in the kunit_kmalloc() functions, which themselves accept a gfp to
pass down to kmalloc(). We didn't want to add an extra GFP_KERNEL
allocation there.

This definitely could be relegated to a separate variant of the
function, though (or we could keep using the old implementation of
kunit_kmalloc_array() which creates resources manually). Trying to
match the devm_add_action() API, which assumed GFP_KERNEL probably
makes sense...

> > +/**
> > + * kunit_remove_action_token() - Cancel a deferred action.
> > + * @test: Test case the action is associated with.
> > + * @cancel_token: The cancellation token returned by kunit_add_action()
> > + *
> > + * Prevent an action deferred using kunit_add_action() from executing when the
> > + * test ends.
> > + *
> > + * You can also use the (test, function, context) triplet to remove an action
> > + * with kunit_remove_action().
> > + */
> > +void kunit_remove_action_token(struct kunit *test, struct kunit_action_ctx *cancel_token);
>
> It's not clear to me why we still need the token. If
> kunit_remove_action() works fine, why would we need to store the token?
>
> Can't we just make kunit_add_action() return an int to indicate whether
> it failed or not, and that's it?
>

So the distinction here is that the (function, context) pair doesn't
uniquely identify an action, as you can add the same action multiple
times, with other actions interleaved. A token encodes _which_ of
these actions is being triggered/cancelled: the non-token variants
always cancel the most recent matching function. Without the token,
there's no way of removing an action "further down the stack".

Take, for example, two functions, add_one() and double(), which are
(*ctx)++ and (*ctx) *= 2, respectively.
int var = 0;
tok1 = kunit_add_action(test, add_one, &var);
kunit_add_action(test, double, &var);
tok3 = kunit_add_action(test, add_one, &var);

// The call:
kunit_remove_action(test, add_one, &var);
// is equivalent to
kunit_remove_action_token(test, tok3);
// and gives var = 1 as a final result

// If instead we want to remove the first add_one, we use:
kunit_remove_action_token(test, tok1);
// which cannot be done using kunit_remove_action()
// gives var = 2 instead.


There's also a (minor) performance benefit to using the token
versions, as we don't need to do a (currently O(n)) search through the
list of KUnit resources to find the matching entry. I doubt too many
tests will defer enough to make this a problem.


That being said, given no-one actually needs this behaviour yet, it's
definitely something we could add later if it becomes necessary. I
doubt it'd be useful for most of the normal resource management
use-cases.

> > +/**
> > + * kunit_release_action_token() - Trigger a deferred action immediately.
> > + * @test: Test case the action is associated with.
> > + * @cancel_token: The cancellation token returned by kunit_add_action()
> > + *
> > + * Execute an action immediately, instead of waiting for the test to end.
> > + *
> > + * You can also use the (test, function, context) triplet to trigger an action
> > + * with kunit_release_action().
> > + */
> > +void kunit_release_action_token(struct kunit *test, struct kunit_action_ctx *cancel_token);
> > +
> > +/**
> > + * kunit_remove_action() - Cancel a matching deferred action.
> > + * @test: Test case the action is associated with.
> > + * @func: The deferred function to cancel.
> > + * @ctx: The context passed to the deferred function to trigger.
> > + *
> > + * Prevent an action deferred via kunit_add_action() from executing when the
> > + * test terminates..
> > + * Unlike kunit_remove_action_token(), this takes the (func, ctx) pair instead of
> > + * the cancellation token. If that function/context pair was deferred multiple
> > + * times, only the most recent one will be cancelled.
> > + */
> > +void kunit_remove_action(struct kunit *test,
> > +                      kunit_defer_function_t func,
> > +                      void *ctx);
> > +
> > +/**
> > + * kunit_release_action() - Run a matching action call immediately.
> > + * @test: Test case the action is associated with.
> > + * @func: The deferred function to trigger.
> > + * @ctx: The context passed to the deferred function to trigger.
> > + *
> > + * Execute a function deferred via kunit_add_action()) immediately, rather than
> > + * when the test ends.
> > + * Unlike kunit_release_action_token(), this takes the (func, ctx) pair instead of
> > + * the cancellation token. If that function/context pair was deferred multiple
> > + * times, it will only be executed once here. The most recent deferral will
> > + * no longer execute when the test ends.
> > + *
> > + * kunit_release_action(test, func, ctx);
> > + * is equivalent to
> > + * func(ctx);
> > + * kunit_remove_action(test, func, ctx);
> > + */
> > +void kunit_release_action(struct kunit *test,
> > +                       kunit_defer_function_t func,
> > +                       void *ctx);
> >  #endif /* _KUNIT_RESOURCE_H */
> > diff --git a/lib/kunit/kunit-test.c b/lib/kunit/kunit-test.c
> > index b63595d3e241..eaca1b133922 100644
> > --- a/lib/kunit/kunit-test.c
> > +++ b/lib/kunit/kunit-test.c
> > @@ -111,7 +111,7 @@ struct kunit_test_resource_context {
> >       struct kunit test;
> >       bool is_resource_initialized;
> >       int allocate_order[2];
> > -     int free_order[2];
> > +     int free_order[4];
> >  };
> >
> >  static int fake_resource_init(struct kunit_resource *res, void *context)
> > @@ -402,6 +402,123 @@ static void kunit_resource_test_named(struct kunit *test)
> >       KUNIT_EXPECT_TRUE(test, list_empty(&test->resources));
> >  }
> >
> > +static void increment_int(void *ctx)
> > +{
> > +     int *i = (int *)ctx;
> > +     (*i)++;
> > +}
> > +
> > +static void kunit_resource_test_action(struct kunit *test)
> > +{
> > +     int num_actions = 0;
> > +
> > +     kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
> > +     KUNIT_EXPECT_EQ(test, num_actions, 0);
> > +     kunit_cleanup(test);
> > +     KUNIT_EXPECT_EQ(test, num_actions, 1);
> > +
> > +     /* Once we've cleaned up, the action queue is empty. */
> > +     kunit_cleanup(test);
> > +     KUNIT_EXPECT_EQ(test, num_actions, 1);
> > +
> > +     /* Check the same function can be deferred multiple times. */
> > +     kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
> > +     kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
> > +     kunit_cleanup(test);
> > +     KUNIT_EXPECT_EQ(test, num_actions, 3);
> > +}
> > +static void kunit_resource_test_remove_action(struct kunit *test)
> > +{
> > +     int num_actions = 0;
> > +     struct kunit_action_ctx *cancel_token;
> > +     struct kunit_action_ctx *cancel_token2;
> > +
> > +     cancel_token = kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
> > +     KUNIT_EXPECT_EQ(test, num_actions, 0);
> > +
> > +     kunit_remove_action_token(test, cancel_token);
> > +     kunit_cleanup(test);
> > +     KUNIT_EXPECT_EQ(test, num_actions, 0);
> > +
> > +     /* Check calls from the same function/context pair can be cancelled independently*/
> > +     cancel_token = kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
> > +     cancel_token2 = kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
> > +     kunit_remove_action_token(test, cancel_token);
> > +     kunit_cleanup(test);
> > +     KUNIT_EXPECT_EQ(test, num_actions, 1);
> > +
> > +     /* Also check that we can cancel just one of the identical function/context pairs. */
> > +     cancel_token = kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
> > +     cancel_token2 = kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
> > +     kunit_remove_action(test, increment_int, &num_actions);
> > +     kunit_cleanup(test);
> > +     KUNIT_EXPECT_EQ(test, num_actions, 2);
> > +}
> > +static void kunit_resource_test_release_action(struct kunit *test)
> > +{
> > +     int num_actions = 0;
> > +     struct kunit_action_ctx *cancel_token;
> > +     struct kunit_action_ctx *cancel_token2;
> > +
> > +     cancel_token = kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
> > +     KUNIT_EXPECT_EQ(test, num_actions, 0);
> > +     /* Runs immediately on trigger. */
> > +     kunit_release_action_token(test, cancel_token);
> > +     KUNIT_EXPECT_EQ(test, num_actions, 1);
> > +
> > +     /* Doesn't run again on test exit. */
> > +     kunit_cleanup(test);
> > +     KUNIT_EXPECT_EQ(test, num_actions, 1);
> > +
> > +     /* Check calls from the same function/context pair can be triggered independently*/
> > +     cancel_token = kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
> > +     cancel_token2 = kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
> > +     kunit_release_action_token(test, cancel_token);
> > +     KUNIT_EXPECT_EQ(test, num_actions, 2);
> > +     kunit_cleanup(test);
> > +     KUNIT_EXPECT_EQ(test, num_actions, 3);
> > +
> > +     /* Also check that we can trigger just one of the identical function/context pairs. */
> > +     kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
> > +     kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
> > +     kunit_release_action(test, increment_int, &num_actions);
> > +     KUNIT_EXPECT_EQ(test, num_actions, 4);
> > +     kunit_cleanup(test);
> > +     KUNIT_EXPECT_EQ(test, num_actions, 5);
> > +}
> > +static void action_order_1(void *ctx)
> > +{
> > +     struct kunit_test_resource_context *res_ctx = (struct kunit_test_resource_context *)ctx;
> > +
> > +     KUNIT_RESOURCE_TEST_MARK_ORDER(res_ctx, free_order, 1);
> > +     kunit_log(KERN_INFO, current->kunit_test, "action_order_1");
> > +}
> > +static void action_order_2(void *ctx)
> > +{
> > +     struct kunit_test_resource_context *res_ctx = (struct kunit_test_resource_context *)ctx;
> > +
> > +     KUNIT_RESOURCE_TEST_MARK_ORDER(res_ctx, free_order, 2);
> > +     kunit_log(KERN_INFO, current->kunit_test, "action_order_2");
> > +}
> > +static void kunit_resource_test_action_ordering(struct kunit *test)
> > +{
> > +     struct kunit_test_resource_context *ctx = test->priv;
> > +     struct kunit_action_ctx *cancel_token;
> > +
> > +     kunit_add_action(test, action_order_1, ctx, GFP_KERNEL);
> > +     cancel_token = kunit_add_action(test, action_order_2, ctx, GFP_KERNEL);
> > +     kunit_add_action(test, action_order_1, ctx, GFP_KERNEL);
> > +     kunit_add_action(test, action_order_2, ctx, GFP_KERNEL);
> > +     kunit_remove_action(test, action_order_1, ctx);
> > +     kunit_release_action_token(test, cancel_token);
> > +     kunit_cleanup(test);
> > +
> > +     /* [2 is triggered] [2], [(1 is cancelled)] [1] */
> > +     KUNIT_EXPECT_EQ(test, ctx->free_order[0], 2);
> > +     KUNIT_EXPECT_EQ(test, ctx->free_order[1], 2);
> > +     KUNIT_EXPECT_EQ(test, ctx->free_order[2], 1);
> > +}
> > +
> >  static int kunit_resource_test_init(struct kunit *test)
> >  {
> >       struct kunit_test_resource_context *ctx =
> > @@ -433,6 +550,10 @@ static struct kunit_case kunit_resource_test_cases[] = {
> >       KUNIT_CASE(kunit_resource_test_proper_free_ordering),
> >       KUNIT_CASE(kunit_resource_test_static),
> >       KUNIT_CASE(kunit_resource_test_named),
> > +     KUNIT_CASE(kunit_resource_test_action),
> > +     KUNIT_CASE(kunit_resource_test_remove_action),
> > +     KUNIT_CASE(kunit_resource_test_release_action),
> > +     KUNIT_CASE(kunit_resource_test_action_ordering),
> >       {}
> >  };
> >
> > diff --git a/lib/kunit/resource.c b/lib/kunit/resource.c
> > index c414df922f34..824cf91e306d 100644
> > --- a/lib/kunit/resource.c
> > +++ b/lib/kunit/resource.c
> > @@ -77,3 +77,102 @@ int kunit_destroy_resource(struct kunit *test, kunit_resource_match_t match,
> >       return 0;
> >  }
> >  EXPORT_SYMBOL_GPL(kunit_destroy_resource);
> > +
> > +struct kunit_action_ctx {
> > +     struct kunit_resource res;
> > +     kunit_defer_function_t func;
> > +     void *ctx;
> > +};
> > +
> > +static void __kunit_action_free(struct kunit_resource *res)
> > +{
> > +     struct kunit_action_ctx *action_ctx = container_of(res, struct kunit_action_ctx, res);
> > +
> > +     action_ctx->func(action_ctx->ctx);
> > +}
> > +
> > +struct kunit_action_ctx *kunit_add_action(struct kunit *test, kunit_defer_function_t func,
> > +                                      void *ctx, gfp_t internal_gfp)
> > +{
> > +     struct kunit_action_ctx *action_ctx;
> > +
> > +     KUNIT_ASSERT_NOT_NULL_MSG(test, func, "Tried to action a NULL function!");
> > +
> > +     action_ctx = kzalloc(sizeof(*action_ctx), internal_gfp);
> > +     if (!action_ctx)
> > +             return NULL;
> > +
> > +     action_ctx->func = func;
> > +     action_ctx->ctx = ctx;
> > +
> > +     action_ctx->res.should_kfree = true;
> > +     /* As init is NULL, this cannot fail. */
> > +     __kunit_add_resource(test, NULL, __kunit_action_free, &action_ctx->res, action_ctx);
> > +
> > +     return action_ctx;
> > +}
>
> One thing worth pointing is that if kunit_add_action() fails, the
> cleanup function passed as an argument won't run.
>
> So, if the kzalloc call ever fails, patch 2 will leak its res->data()
> resource for example.
>
> devm (and drmm) handles this using a variant called
> devm_add_action_or_reset, we should either provide the same variant or
> just go for that behavior by default.
>

Excellent point.

I think I'll add a kunit_add_action_or_reset() variant to the next
revision. If we've gone this far to match the devm_ API, continuing to
do so probably is the best way of handling it.

Cheers,
-- David

David Gow April 5, 2023, 8:09 a.m. UTC | #4

On Wed, 5 Apr 2023 at 01:55, Benjamin Berg <[e-mail]> wrote:
>
> Hi,
>
> On Tue, 2023-04-04 at 15:32 +0200, Maxime Ripard wrote:
> > [SNIP]
> > > +/**
> > > + * kunit_add_action() - Defer an 'action' (function call) until the test ends.
> > > + * @test: Test case to associate the action with.
> > > + * @func: The function to run on test exit
> > > + * @ctx: Data passed into @func
> > > + * @internal_gfp: gfp to use for internal allocations, if unsure, use GFP_KERNEL
> > > + *
> > > + * Defer the execution of a function until the test exits, either normally or
> > > + * due to a failure.  @ctx is passed as additional context. All functions
> > > + * registered with kunit_add_action() will execute in the opposite order to that
> > > + * they were registered in.
> > > + *
> > > + * This is useful for cleaning up allocated memory and resources.
> > > + *
> > > + * Returns:
> > > + *   An opaque "cancellation token", or NULL on error. Pass this token to
> > > + *   kunit_remove_action_token() in order to cancel the deferred execution of
> > > + *   func().
> > > + */
> > > +struct kunit_action_ctx *kunit_add_action(struct kunit *test, kunit_defer_function_t func,
> > > +                     void *ctx, gfp_t internal_gfp);
> >
> > Do we expect any other context than GFP_KERNEL?
> >
> > If so, then maybe we can have kunit_add_action() assume GFP_KERNEL and
> > add a variant for the odd case where we would actually need a different
> > GFP flag.
>
> Does anything other than GFP_KERNEL make sense? I would assume these
> functions should only ever be called from a kunit context, i.e. the
> passed test is guaranteed to be identical to the value returned by
> kunit_get_current_test().

That's not strictly-speaking guaranteed. (Indeed, we have some, albeit
contrived, counterexamples in the test.)

The theoretical use-case here is if the kunit context pointer is
passed to another thread or workqueue or something.

There aren't any such users, yet (apart from, possibly,
kunit_kmalloc_array()), though. So we could use GFP_KERNEL by default
for now, and add a variant if such a use-case appears.

>
> That said, I am happy with merging this in this form. I feel the right
> thing here is a patch (with corresponding spatch) that changes all of
> the related APIs to remove the gfp argument.
>
> > > +/**
> > > + * kunit_remove_action_token() - Cancel a deferred action.
> > > + * @test: Test case the action is associated with.
> > > + * @cancel_token: The cancellation token returned by kunit_add_action()
> > > + *
> > > + * Prevent an action deferred using kunit_add_action() from executing when the
> > > + * test ends.
> > > + *
> > > + * You can also use the (test, function, context) triplet to remove an action
> > > + * with kunit_remove_action().
> > > + */
> > > +void kunit_remove_action_token(struct kunit *test, struct kunit_action_ctx *cancel_token);
> >
> > It's not clear to me why we still need the token. If
> > kunit_remove_action() works fine, why would we need to store the token?
> >
> > Can't we just make kunit_add_action() return an int to indicate whether
> > it failed or not, and that's it?
> >
> > > [SNIP]
> >
> > One thing worth pointing is that if kunit_add_action() fails, the
> > cleanup function passed as an argument won't run.
> >
> > So, if the kzalloc call ever fails, patch 2 will leak its res->data()
> > resource for example.
> >
> > devm (and drmm) handles this using a variant called
> > devm_add_action_or_reset, we should either provide the same variant or
> > just go for that behavior by default.
>
> Both version of the function would need a return value. An alternative
> might be to make assertions part of the API. But as with dropping the
> gfp argument, that seems like a more intrusive change that needs to
> happen independently.
>
> Anyway, I am fine with action_or_reset as the default and possibly the
> only behaviour. I expect that every API user will want an assertion
> that checks for failure here anyway.
>

I'm tempted to just have both kunit_add_action() and
kunit_add_action_or_reset(), just to keep things matching the devm_
API to minimise any confusion.

And if we're not too worried about proliferating variants of these
(and, personally, I quite like them), we could have a
kunit_add_action_or_asserrt() version as well.

> Benjamin
>
>
> If kunit_* functions can assert in error conditions, then the example
>
> void test_func(struct kunit *test)
> {
>   char u8 *buf = kunit_kzalloc(test, 1024, GFP_KERNEL);
>   struct sk_buff *skb_a;
>   struct sk_buff *skb_b;
>   /* Further variables */
>
>   KUNIT_ASSERT_NOT_NULL(test, buf);
>
>   skb_a = skb_alloc(1024, GFP_KERNEL);
>   KUNIT_ASSERT_NOT_NULL(test, skb_a);
>   if (kunit_add_cleanup(test, (kunit_defer_function_t) kfree_skb, skb_a))
>     KUNIT_ASSERT_FAILURE("Failed to add cleanup");
>
>   /* Or, maybe: */
>   skb_b = skb_alloc(1024, GFP_KERNEL);
>   KUNIT_ASSERT_NOT_NULL(test, skb_b);
>   KUNIT_ASSERT_EQ(test, 0,
>                   kunit_add_cleanup(test,
>                                     (kunit_defer_function_t) kfree_skb,
>                                     skb_b));
>
>   /* run code that may assert */
> }
>
>
> could be shortened to (with a trivial kunit_skb_alloc helper)
>
> void test_func(struct kunit *test)
> {
>   char u8 *buf = kunit_kzalloc(test, 1024, GFP_KERNEL);
>   struct sk_buff *skb_a = kunit_skb_alloc(1024, GFP_KERNEL);
>   struct sk_buff *skb_b = kunit_skb_alloc(1024, GFP_KERNEL);
>   /* Further variables */
>
>   /* run code that may assert */
> }
>
> I should just post a patch for the existing API and see what people say
> then ...

We definitely already have some functions (e.g.
__kunit_activate_static_stub()) which just assert on failure. In
general, we've avoided doing so where we think there might be a good
reason to handle failures separately (or it makes the API diverge a
lot from a function we're imitating), but I'm open to using them more.
Specialised handling of allocation failures in a test is likely to be
rare enough that making those who need it write their own helpers
wouldn't be a disaster. Or we could have an _or_assert() variant.

In any case, I think your example pretty comfortably speaks for itself.

Cheers,
-- David

Maxime Ripard April 14, 2023, 9:53 a.m. UTC | #5

On Wed, Apr 05, 2023 at 03:47:55PM +0800, David Gow wrote:
> On Tue, 4 Apr 2023 at 21:32, Maxime Ripard <[e-mail]> wrote:
> >
> > Hi David,
> >
> > Looks great, thanks for sending a second version
> >
> > On Fri, Mar 31, 2023 at 04:04:09PM +0800, David Gow wrote:
> > > +/**
> > > + * kunit_remove_action_token() - Cancel a deferred action.
> > > + * @test: Test case the action is associated with.
> > > + * @cancel_token: The cancellation token returned by kunit_add_action()
> > > + *
> > > + * Prevent an action deferred using kunit_add_action() from executing when the
> > > + * test ends.
> > > + *
> > > + * You can also use the (test, function, context) triplet to remove an action
> > > + * with kunit_remove_action().
> > > + */
> > > +void kunit_remove_action_token(struct kunit *test, struct kunit_action_ctx *cancel_token);
> >
> > It's not clear to me why we still need the token. If
> > kunit_remove_action() works fine, why would we need to store the token?
> >
> > Can't we just make kunit_add_action() return an int to indicate whether
> > it failed or not, and that's it?
> >
> 
> So the distinction here is that the (function, context) pair doesn't
> uniquely identify an action, as you can add the same action multiple
> times, with other actions interleaved. A token encodes _which_ of
> these actions is being triggered/cancelled: the non-token variants
> always cancel the most recent matching function. Without the token,
> there's no way of removing an action "further down the stack".

> Take, for example, two functions, add_one() and double(), which are
> (*ctx)++ and (*ctx) *= 2, respectively.
> int var = 0;
> tok1 = kunit_add_action(test, add_one, &var);
> kunit_add_action(test, double, &var);
> tok3 = kunit_add_action(test, add_one, &var);
> 
> // The call:
> kunit_remove_action(test, add_one, &var);
> // is equivalent to
> kunit_remove_action_token(test, tok3);
> // and gives var = 1 as a final result
> 
> // If instead we want to remove the first add_one, we use:
> kunit_remove_action_token(test, tok1);
> // which cannot be done using kunit_remove_action()
> // gives var = 2 instead.

Sure, I still think we're kind of over-engineering this. request_irq,
devm_add_action and drmm_add_action all use that the irq/device, address
of the function and the context value to differentiate between
instances, and we never had the need for any token despite having
thousand of users.

Given that actions are supposed to be unrolled in the opposite order, I
think that removing the last action that match makes the most sense.

> There's also a (minor) performance benefit to using the token
> versions, as we don't need to do a (currently O(n)) search through the
> list of KUnit resources to find the matching entry. I doubt too many
> tests will defer enough to make this a problem.
> 
> 
> That being said, given no-one actually needs this behaviour yet, it's
> definitely something we could add later if it becomes necessary. I
> doubt it'd be useful for most of the normal resource management
> use-cases.

Yeah, I guess it's the best approach for now.

Thanks!
Maxime

Maxime Ripard April 14, 2023, 10:01 a.m. UTC | #6

Hi David,

On Fri, Mar 31, 2023 at 04:04:09PM +0800, David Gow wrote:
> Many uses of the KUnit resource system are intended to simply defer
> calling a function until the test exits (be it due to success or
> failure). The existing kunit_alloc_resource() function is often used for
> this, but was awkward to use (requiring passing NULL init functions, etc),
> and returned a resource without incrementing its reference count, which
> -- while okay for this use-case -- could cause problems in others.
> 
> Instead, introduce a simple kunit_add_action() API: a simple function
> (returning nothing, accepting a single void* argument) can be scheduled
> to be called when the test exits. Deferred actions are called in the
> opposite order to that which they were registered.
> 
> This mimics the devres API, devm_add_action(), and also provides
> kunit_remove_action(), to cancel a deferred action, and
> kunit_release_action() to trigger one early.
> 
> This is implemented as a resource under the hood, so the ordering
> between resource cleanup and deferred functions is maintained.
> 
> Signed-off-by: David Gow <[e-mail]>
> ---
> 
> Changes since RFC v1:
> https://lore.kernel.org/linux-kselftest/[e-mail]/
> - Rename functions to better match the devm_* APIs. (Thanks Maxime)
> - Embed the kunit_resource in struct kunit_action_ctx to avoid an extra
>   allocation (Thanks Benjamin)
> - Use 'struct kunit_action_ctx' as the type for cancellation tokens
>   (Thanks Benjamin)
> - Add tests.
> 
> ---
>  include/kunit/resource.h |  89 ++++++++++++++++++++++++++++
>  lib/kunit/kunit-test.c   | 123 ++++++++++++++++++++++++++++++++++++++-
>  lib/kunit/resource.c     |  99 +++++++++++++++++++++++++++++++
>  3 files changed, 310 insertions(+), 1 deletion(-)
> 
> diff --git a/include/kunit/resource.h b/include/kunit/resource.h
> index c0d88b318e90..15efd8924666 100644
> --- a/include/kunit/resource.h
> +++ b/include/kunit/resource.h
> @@ -387,4 +387,93 @@ static inline int kunit_destroy_named_resource(struct kunit *test,
>   */
>  void kunit_remove_resource(struct kunit *test, struct kunit_resource *res);
>  
> +typedef void (*kunit_defer_function_t)(void *ctx);
> +
> +/* An opaque token to a deferred action. */
> +struct kunit_action_ctx;
> +
> +/**
> + * kunit_add_action() - Defer an 'action' (function call) until the test ends.
> + * @test: Test case to associate the action with.
> + * @func: The function to run on test exit
> + * @ctx: Data passed into @func
> + * @internal_gfp: gfp to use for internal allocations, if unsure, use GFP_KERNEL
> + *
> + * Defer the execution of a function until the test exits, either normally or
> + * due to a failure.  @ctx is passed as additional context. All functions
> + * registered with kunit_add_action() will execute in the opposite order to that
> + * they were registered in.
> + *
> + * This is useful for cleaning up allocated memory and resources.
> + *
> + * Returns:
> + *   An opaque "cancellation token", or NULL on error. Pass this token to
> + *   kunit_remove_action_token() in order to cancel the deferred execution of
> + *   func().
> + */
> +struct kunit_action_ctx *kunit_add_action(struct kunit *test, kunit_defer_function_t func,
> +		      void *ctx, gfp_t internal_gfp);

I've tried to leverage kunit_add_action() today, and I'm wondering if
passing the struct kunit pointer to the deferred function would help.

The code I'm struggling with is something like:

> static int test_init(struct kunit *test)
> {
> 	priv = kunit_kzalloc(sizeof(*priv), GFP_KERNEL);
>	KUNIT_ASSERT_NOT_NULL(test, priv);
>	test->priv = priv;
>
> 	priv->dev = alloc_device();
>
>	return 0;
> }

and then in the test itself:

> static void actual_test(struct kunit *test)
> {
> 	struct test_priv *priv = test->priv;
>
>	id = allocate_buffer(priv->dev);
>
> 	KUNIT_EXPECT_EQ(test, id, 42);
>
> 	free_buffer(priv->dev, id);
> }

I'd like to turn free_buffer an action registered right after allocate
buffer. However, since it takes several arguments and kunit_add_action
expects a single pointer, we would need to create a structure for it,
allocate it, fill it, and then free it when the action has ran.

It creates a lot of boilerplate, while if we were passing the pointer to
struct kunit we could access the context of the test as well, and things
would be much simpler.

Does that make sense?

Maxime

Benjamin Berg April 14, 2023, 11 a.m. UTC | #7

Hi,

On Fri, 2023-04-14 at 12:01 +0200, [e-mail] wrote:
> Hi David,
> 
> On Fri, Mar 31, 2023 at 04:04:09PM +0800, David Gow wrote:
> > Many uses of the KUnit resource system are intended to simply defer
> > calling a function until the test exits (be it due to success or
> > failure). The existing kunit_alloc_resource() function is often used for
> > this, but was awkward to use (requiring passing NULL init functions, etc),
> > and returned a resource without incrementing its reference count, which
> > -- while okay for this use-case -- could cause problems in others.
> > 
> > Instead, introduce a simple kunit_add_action() API: a simple function
> > (returning nothing, accepting a single void* argument) can be scheduled
> > to be called when the test exits. Deferred actions are called in the
> > opposite order to that which they were registered.
> > 
> > This mimics the devres API, devm_add_action(), and also provides
> > kunit_remove_action(), to cancel a deferred action, and
> > kunit_release_action() to trigger one early.
> > 
> > This is implemented as a resource under the hood, so the ordering
> > between resource cleanup and deferred functions is maintained.
> > 
> > Signed-off-by: David Gow <[e-mail]>
> > ---
> > 
> > Changes since RFC v1:
> > https://lore.kernel.org/linux-kselftest/[e-mail]/
> > - Rename functions to better match the devm_* APIs. (Thanks Maxime)
> > - Embed the kunit_resource in struct kunit_action_ctx to avoid an extra
> >   allocation (Thanks Benjamin)
> > - Use 'struct kunit_action_ctx' as the type for cancellation tokens
> >   (Thanks Benjamin)
> > - Add tests.
> > 
> > ---
> >  include/kunit/resource.h |  89 ++++++++++++++++++++++++++++
> >  lib/kunit/kunit-test.c   | 123 ++++++++++++++++++++++++++++++++++++++-
> >  lib/kunit/resource.c     |  99 +++++++++++++++++++++++++++++++
> >  3 files changed, 310 insertions(+), 1 deletion(-)
> > 
> > diff --git a/include/kunit/resource.h b/include/kunit/resource.h
> > index c0d88b318e90..15efd8924666 100644
> > --- a/include/kunit/resource.h
> > +++ b/include/kunit/resource.h
> > @@ -387,4 +387,93 @@ static inline int kunit_destroy_named_resource(struct kunit *test,
> >   */
> >  void kunit_remove_resource(struct kunit *test, struct kunit_resource *res);
> >  
> > +typedef void (*kunit_defer_function_t)(void *ctx);
> > +
> > +/* An opaque token to a deferred action. */
> > +struct kunit_action_ctx;
> > +
> > +/**
> > + * kunit_add_action() - Defer an 'action' (function call) until the test ends.
> > + * @test: Test case to associate the action with.
> > + * @func: The function to run on test exit
> > + * @ctx: Data passed into @func
> > + * @internal_gfp: gfp to use for internal allocations, if unsure, use GFP_KERNEL
> > + *
> > + * Defer the execution of a function until the test exits, either normally or
> > + * due to a failure.  @ctx is passed as additional context. All functions
> > + * registered with kunit_add_action() will execute in the opposite order to that
> > + * they were registered in.
> > + *
> > + * This is useful for cleaning up allocated memory and resources.
> > + *
> > + * Returns:
> > + *   An opaque "cancellation token", or NULL on error. Pass this token to
> > + *   kunit_remove_action_token() in order to cancel the deferred execution of
> > + *   func().
> > + */
> > +struct kunit_action_ctx *kunit_add_action(struct kunit *test, kunit_defer_function_t func,
> > +                     void *ctx, gfp_t internal_gfp);
> 
> I've tried to leverage kunit_add_action() today, and I'm wondering if
> passing the struct kunit pointer to the deferred function would help.
> 
> The code I'm struggling with is something like:
> 
> > static int test_init(struct kunit *test)
> > {
> >         priv = kunit_kzalloc(sizeof(*priv), GFP_KERNEL);
> >         KUNIT_ASSERT_NOT_NULL(test, priv);
> >         test->priv = priv;
> > 
> >         priv->dev = alloc_device();
> > 
> >         return 0;
> > }
> 
> and then in the test itself:
> 
> > static void actual_test(struct kunit *test)
> > {
> >         struct test_priv *priv = test->priv;
> > 
> >         id = allocate_buffer(priv->dev);
> > 
> >         KUNIT_EXPECT_EQ(test, id, 42);
> > 
> >         free_buffer(priv->dev, id);
> > }
> 
> I'd like to turn free_buffer an action registered right after allocate
> buffer. However, since it takes several arguments and kunit_add_action
> expects a single pointer, we would need to create a structure for it,
> allocate it, fill it, and then free it when the action has ran.
> 
> It creates a lot of boilerplate, while if we were passing the pointer to
> struct kunit we could access the context of the test as well, and things
> would be much simpler.

The question seems to be what about the typical use-case. I was always
imagining calling functions like kfree/kfree_skb which often only
require a single argument.

For arbitrary arguments, a struct and custom free function will be
needed. At that point, maybe it is fair to assume that API users will
use the resource API directly, doing the same trick as kunit_add_action
and storing the arguments together with struct kunit_resource.

That said, maybe one could add it as a second argument? It is a little
bit weird API wise, but it would allow simply casting single-argument
functions in order to ignore "struct kunit *" argument. 

Benjamin

Maxime Ripard April 14, 2023, 11:33 a.m. UTC | #8

On Fri, Apr 14, 2023 at 01:00:26PM +0200, Benjamin Berg wrote:
> Hi,
> 
> On Fri, 2023-04-14 at 12:01 +0200, [e-mail] wrote:
> > Hi David,
> > 
> > On Fri, Mar 31, 2023 at 04:04:09PM +0800, David Gow wrote:
> > > Many uses of the KUnit resource system are intended to simply defer
> > > calling a function until the test exits (be it due to success or
> > > failure). The existing kunit_alloc_resource() function is often used for
> > > this, but was awkward to use (requiring passing NULL init functions, etc),
> > > and returned a resource without incrementing its reference count, which
> > > -- while okay for this use-case -- could cause problems in others.
> > > 
> > > Instead, introduce a simple kunit_add_action() API: a simple function
> > > (returning nothing, accepting a single void* argument) can be scheduled
> > > to be called when the test exits. Deferred actions are called in the
> > > opposite order to that which they were registered.
> > > 
> > > This mimics the devres API, devm_add_action(), and also provides
> > > kunit_remove_action(), to cancel a deferred action, and
> > > kunit_release_action() to trigger one early.
> > > 
> > > This is implemented as a resource under the hood, so the ordering
> > > between resource cleanup and deferred functions is maintained.
> > > 
> > > Signed-off-by: David Gow <[e-mail]>
> > > ---
> > > 
> > > Changes since RFC v1:
> > > https://lore.kernel.org/linux-kselftest/[e-mail]/
> > > - Rename functions to better match the devm_* APIs. (Thanks Maxime)
> > > - Embed the kunit_resource in struct kunit_action_ctx to avoid an extra
> > >   allocation (Thanks Benjamin)
> > > - Use 'struct kunit_action_ctx' as the type for cancellation tokens
> > >   (Thanks Benjamin)
> > > - Add tests.
> > > 
> > > ---
> > >  include/kunit/resource.h |  89 ++++++++++++++++++++++++++++
> > >  lib/kunit/kunit-test.c   | 123 ++++++++++++++++++++++++++++++++++++++-
> > >  lib/kunit/resource.c     |  99 +++++++++++++++++++++++++++++++
> > >  3 files changed, 310 insertions(+), 1 deletion(-)
> > > 
> > > diff --git a/include/kunit/resource.h b/include/kunit/resource.h
> > > index c0d88b318e90..15efd8924666 100644
> > > --- a/include/kunit/resource.h
> > > +++ b/include/kunit/resource.h
> > > @@ -387,4 +387,93 @@ static inline int kunit_destroy_named_resource(struct kunit *test,
> > >   */
> > >  void kunit_remove_resource(struct kunit *test, struct kunit_resource *res);
> > >  
> > > +typedef void (*kunit_defer_function_t)(void *ctx);
> > > +
> > > +/* An opaque token to a deferred action. */
> > > +struct kunit_action_ctx;
> > > +
> > > +/**
> > > + * kunit_add_action() - Defer an 'action' (function call) until the test ends.
> > > + * @test: Test case to associate the action with.
> > > + * @func: The function to run on test exit
> > > + * @ctx: Data passed into @func
> > > + * @internal_gfp: gfp to use for internal allocations, if unsure, use GFP_KERNEL
> > > + *
> > > + * Defer the execution of a function until the test exits, either normally or
> > > + * due to a failure.  @ctx is passed as additional context. All functions
> > > + * registered with kunit_add_action() will execute in the opposite order to that
> > > + * they were registered in.
> > > + *
> > > + * This is useful for cleaning up allocated memory and resources.
> > > + *
> > > + * Returns:
> > > + *   An opaque "cancellation token", or NULL on error. Pass this token to
> > > + *   kunit_remove_action_token() in order to cancel the deferred execution of
> > > + *   func().
> > > + */
> > > +struct kunit_action_ctx *kunit_add_action(struct kunit *test, kunit_defer_function_t func,
> > > +                     void *ctx, gfp_t internal_gfp);
> > 
> > I've tried to leverage kunit_add_action() today, and I'm wondering if
> > passing the struct kunit pointer to the deferred function would help.
> > 
> > The code I'm struggling with is something like:
> > 
> > > static int test_init(struct kunit *test)
> > > {
> > >         priv = kunit_kzalloc(sizeof(*priv), GFP_KERNEL);
> > >         KUNIT_ASSERT_NOT_NULL(test, priv);
> > >         test->priv = priv;
> > > 
> > >         priv->dev = alloc_device();
> > > 
> > >         return 0;
> > > }
> > 
> > and then in the test itself:
> > 
> > > static void actual_test(struct kunit *test)
> > > {
> > >         struct test_priv *priv = test->priv;
> > > 
> > >         id = allocate_buffer(priv->dev);
> > > 
> > >         KUNIT_EXPECT_EQ(test, id, 42);
> > > 
> > >         free_buffer(priv->dev, id);
> > > }
> > 
> > I'd like to turn free_buffer an action registered right after allocate
> > buffer. However, since it takes several arguments and kunit_add_action
> > expects a single pointer, we would need to create a structure for it,
> > allocate it, fill it, and then free it when the action has ran.
> > 
> > It creates a lot of boilerplate, while if we were passing the pointer to
> > struct kunit we could access the context of the test as well, and things
> > would be much simpler.
> 
> The question seems to be what about the typical use-case. I was always
> imagining calling functions like kfree/kfree_skb which often only
> require a single argument.

I guess we can have a look at the devm stuff. I'd expect the scope of
things that will eventually tie their resource to kunit would be
similar. "Straight" allocation/deallocation functions are the obvious
first candidates, but there's a lot of other use cases as well.

I guess my main point is that it assumes that most function to clean
things up will take the resource as its only argument, which isn't
always the case. I guess it's reasonable to optimize for the most
trivial case, but we should strive to keep the boilerplate down as much
as we can in the other case too.

> For arbitrary arguments, a struct and custom free function will be
> needed. At that point, maybe it is fair to assume that API users will
> use the resource API directly, doing the same trick as kunit_add_action
> and storing the arguments together with struct kunit_resource.

kunit_add_resource adds tons of boilerplate as well:

struct test_buffer_priv {
	struct device *dev;
}

struct test_alloc_params {
	struct device *dev;
	void *buffer;
}

static int __alloc_buffer(struct kunit_resource *res, void *ptr)
{
	struct test_alloc_params *params = ptr;
	void *buffer;

	params->buffer = allocate_buffer(params->dev, params->size);
	res->data = params;

	return 0;
}

static void __free_buffer(struct kunit_resource *res)
{
	struct test_alloc_params *params = res->data;

	free_buffer(params->dev, params->buffer);
}

void actual_test(struct kunit_test *test)
{
	struct test_alloc_params *params = test->priv;

	kunit_alloc_resource(test, __alloc_buffer, __free_buffer, GFP_KERNEL, params);
	KUNIT_EXPECT_NOT_NULL(params->buffer);
}

int test_init(struct kunit_test *test)
{
	struct test_alloc_params *params =
		kunit_kmalloc(test, sizeof(*params), GFP_KERNEL);

	test->priv = params;

	params->dev = kunit_allocate_device(...);

	return 0;
}


while we could have something like:


struct test_buffer_priv {
	struct device *dev;
}

static void free_buffer_action(struct kunit *test, void *ptr)
{
	struct test_buffer_priv *priv = test->priv;

	free_buffer(params->dev, ptr);
}

void actual_test(struct kunit_test *test)
{
	struct test_buffer_priv *priv = test->priv;
	void *buffer;

	buffer = allocate_buffer(priv->dev, 42);
	kunit_add_action(test, free_buffer_action, buffer);

	KUNIT_EXPECT_NOT_NULL(buffer);
}

int test_init(struct kunit_test *test)
{
	struct test_buffer_priv *priv =
		kunit_kmalloc(test, sizeof(*priv), GFP_KERNEL);

	test->priv = priv;

	priv->dev = kunit_allocate_device(...);

	return 0;
}

Which is much more compact, more readable, and less error prone.

And sure, kfree and free_skb would need some intermediate function, but
it's like 3 more lines, which can even be shared at the framework or
kunit level, so shouldn't really impact the tests themselves.

> That said, maybe one could add it as a second argument? It is a little
> bit weird API wise, but it would allow simply casting single-argument
> functions in order to ignore "struct kunit *" argument. 

I guess. I'd find it a bit weird to have that one function with the
argument order reversed compared to everything else though.

Maxime

David Gow April 15, 2023, 8:42 a.m. UTC | #9

On Fri, 14 Apr 2023 at 18:02, <[e-mail]> wrote:
>
> Hi David,
>
> On Fri, Mar 31, 2023 at 04:04:09PM +0800, David Gow wrote:
> > Many uses of the KUnit resource system are intended to simply defer
> > calling a function until the test exits (be it due to success or
> > failure). The existing kunit_alloc_resource() function is often used for
> > this, but was awkward to use (requiring passing NULL init functions, etc),
> > and returned a resource without incrementing its reference count, which
> > -- while okay for this use-case -- could cause problems in others.
> >
> > Instead, introduce a simple kunit_add_action() API: a simple function
> > (returning nothing, accepting a single void* argument) can be scheduled
> > to be called when the test exits. Deferred actions are called in the
> > opposite order to that which they were registered.
> >
> > This mimics the devres API, devm_add_action(), and also provides
> > kunit_remove_action(), to cancel a deferred action, and
> > kunit_release_action() to trigger one early.
> >
> > This is implemented as a resource under the hood, so the ordering
> > between resource cleanup and deferred functions is maintained.
> >
> > Signed-off-by: David Gow <[e-mail]>
> > ---
> >
> > Changes since RFC v1:
> > https://lore.kernel.org/linux-kselftest/[e-mail]/
> > - Rename functions to better match the devm_* APIs. (Thanks Maxime)
> > - Embed the kunit_resource in struct kunit_action_ctx to avoid an extra
> >   allocation (Thanks Benjamin)
> > - Use 'struct kunit_action_ctx' as the type for cancellation tokens
> >   (Thanks Benjamin)
> > - Add tests.
> >
> > ---
> >  include/kunit/resource.h |  89 ++++++++++++++++++++++++++++
> >  lib/kunit/kunit-test.c   | 123 ++++++++++++++++++++++++++++++++++++++-
> >  lib/kunit/resource.c     |  99 +++++++++++++++++++++++++++++++
> >  3 files changed, 310 insertions(+), 1 deletion(-)
> >
> > diff --git a/include/kunit/resource.h b/include/kunit/resource.h
> > index c0d88b318e90..15efd8924666 100644
> > --- a/include/kunit/resource.h
> > +++ b/include/kunit/resource.h
> > @@ -387,4 +387,93 @@ static inline int kunit_destroy_named_resource(struct kunit *test,
> >   */
> >  void kunit_remove_resource(struct kunit *test, struct kunit_resource *res);
> >
> > +typedef void (*kunit_defer_function_t)(void *ctx);
> > +
> > +/* An opaque token to a deferred action. */
> > +struct kunit_action_ctx;
> > +
> > +/**
> > + * kunit_add_action() - Defer an 'action' (function call) until the test ends.
> > + * @test: Test case to associate the action with.
> > + * @func: The function to run on test exit
> > + * @ctx: Data passed into @func
> > + * @internal_gfp: gfp to use for internal allocations, if unsure, use GFP_KERNEL
> > + *
> > + * Defer the execution of a function until the test exits, either normally or
> > + * due to a failure.  @ctx is passed as additional context. All functions
> > + * registered with kunit_add_action() will execute in the opposite order to that
> > + * they were registered in.
> > + *
> > + * This is useful for cleaning up allocated memory and resources.
> > + *
> > + * Returns:
> > + *   An opaque "cancellation token", or NULL on error. Pass this token to
> > + *   kunit_remove_action_token() in order to cancel the deferred execution of
> > + *   func().
> > + */
> > +struct kunit_action_ctx *kunit_add_action(struct kunit *test, kunit_defer_function_t func,
> > +                   void *ctx, gfp_t internal_gfp);
>
> I've tried to leverage kunit_add_action() today, and I'm wondering if
> passing the struct kunit pointer to the deferred function would help.
>

I'm tempted, but it does make the case where we just want to cast,
e.g., kfree() directly to an action pointer more difficult. Not that
that's a deal-blocker, but it was convenient...

> The code I'm struggling with is something like:
>
> > static int test_init(struct kunit *test)
> > {
> >       priv = kunit_kzalloc(sizeof(*priv), GFP_KERNEL);
> >       KUNIT_ASSERT_NOT_NULL(test, priv);
> >       test->priv = priv;
> >
> >       priv->dev = alloc_device();
> >
> >       return 0;
> > }
>
> and then in the test itself:
>
> > static void actual_test(struct kunit *test)
> > {
> >       struct test_priv *priv = test->priv;
> >
> >       id = allocate_buffer(priv->dev);
> >
> >       KUNIT_EXPECT_EQ(test, id, 42);
> >
> >       free_buffer(priv->dev, id);
> > }
>
> I'd like to turn free_buffer an action registered right after allocate
> buffer. However, since it takes several arguments and kunit_add_action
> expects a single pointer, we would need to create a structure for it,
> allocate it, fill it, and then free it when the action has ran.

The general case of wanting multiple arguments to an action is a bit
complicated. My plan was to initially support just the one argument,
and deal with more complicated cases later. Ideas included:
- using a struct like you suggest, possibly with some macro magic to
make it easier,
- having a bunch of very similar implementations of
kunit_add_action{2,3,4,..}(), which accept 2,3,4,... arguments,
- something horrible and architecture-specific with manually writing
out arguments to the stack (or registers)

None of those sounded particularly pleasant, though. My suspicion is
that the "right" way of doing this is to maybe have one or two helpers
for common cases (e.g., 2 arguments), and just suggest people create a
structure for anything more complicated, but I'd love a nicer
solution.

>
> It creates a lot of boilerplate, while if we were passing the pointer to
> struct kunit we could access the context of the test as well, and things
> would be much simpler.

For the test context specifically, can you just use kunit_get_current_test()?

There might be an issue with using it during the cleanup process after
a failed assertion (as if the test is aborted early, the cleanup runs
in a different thread), but if so, this should fix it:
---
diff --git a/lib/kunit/test.c b/lib/kunit/test.c
index e2910b261112..2d7cad249863 100644
--- a/lib/kunit/test.c
+++ b/lib/kunit/test.c
@@ -392,10 +392,21 @@ static void kunit_case_internal_cleanup(struct
kunit *test)
static void kunit_run_case_cleanup(struct kunit *test,
                                  struct kunit_suite *suite)
{
+       /*
+        * If we're no-longer running from within the test kthread()
because it failed
+        * or timed out, we still need the context to be okay when
running exit and
+        * cleanup functions.
+        */
+       struct kunit *old_current = current->kunit_test;
+
+       current->kunit_test = test;
       if (suite->exit)
               suite->exit(test);

       kunit_case_internal_cleanup(test);
+
+       /* Restore the thread's previous test context (probably NULL
or test). */
+       current->kunit_test = old_current;
}

struct kunit_try_catch_context {
---

I'll look into tidying that up and sending it through next week, anyway.

Cheers,
-- David

David Gow April 15, 2023, 8:48 a.m. UTC | #10

On Fri, 14 Apr 2023 at 19:00, Benjamin Berg <[e-mail]> wrote:
>
> Hi,
>
> On Fri, 2023-04-14 at 12:01 +0200, [e-mail] wrote:
> > Hi David,
> >
> > On Fri, Mar 31, 2023 at 04:04:09PM +0800, David Gow wrote:
> > > Many uses of the KUnit resource system are intended to simply defer
> > > calling a function until the test exits (be it due to success or
> > > failure). The existing kunit_alloc_resource() function is often used for
> > > this, but was awkward to use (requiring passing NULL init functions, etc),
> > > and returned a resource without incrementing its reference count, which
> > > -- while okay for this use-case -- could cause problems in others.
> > >
> > > Instead, introduce a simple kunit_add_action() API: a simple function
> > > (returning nothing, accepting a single void* argument) can be scheduled
> > > to be called when the test exits. Deferred actions are called in the
> > > opposite order to that which they were registered.
> > >
> > > This mimics the devres API, devm_add_action(), and also provides
> > > kunit_remove_action(), to cancel a deferred action, and
> > > kunit_release_action() to trigger one early.
> > >
> > > This is implemented as a resource under the hood, so the ordering
> > > between resource cleanup and deferred functions is maintained.
> > >
> > > Signed-off-by: David Gow <[e-mail]>
> > > ---
> > >
> > > Changes since RFC v1:
> > > https://lore.kernel.org/linux-kselftest/[e-mail]/
> > > - Rename functions to better match the devm_* APIs. (Thanks Maxime)
> > > - Embed the kunit_resource in struct kunit_action_ctx to avoid an extra
> > >   allocation (Thanks Benjamin)
> > > - Use 'struct kunit_action_ctx' as the type for cancellation tokens
> > >   (Thanks Benjamin)
> > > - Add tests.
> > >
> > > ---
> > >  include/kunit/resource.h |  89 ++++++++++++++++++++++++++++
> > >  lib/kunit/kunit-test.c   | 123 ++++++++++++++++++++++++++++++++++++++-
> > >  lib/kunit/resource.c     |  99 +++++++++++++++++++++++++++++++
> > >  3 files changed, 310 insertions(+), 1 deletion(-)
> > >
> > > diff --git a/include/kunit/resource.h b/include/kunit/resource.h
> > > index c0d88b318e90..15efd8924666 100644
> > > --- a/include/kunit/resource.h
> > > +++ b/include/kunit/resource.h
> > > @@ -387,4 +387,93 @@ static inline int kunit_destroy_named_resource(struct kunit *test,
> > >   */
> > >  void kunit_remove_resource(struct kunit *test, struct kunit_resource *res);
> > >
> > > +typedef void (*kunit_defer_function_t)(void *ctx);
> > > +
> > > +/* An opaque token to a deferred action. */
> > > +struct kunit_action_ctx;
> > > +
> > > +/**
> > > + * kunit_add_action() - Defer an 'action' (function call) until the test ends.
> > > + * @test: Test case to associate the action with.
> > > + * @func: The function to run on test exit
> > > + * @ctx: Data passed into @func
> > > + * @internal_gfp: gfp to use for internal allocations, if unsure, use GFP_KERNEL
> > > + *
> > > + * Defer the execution of a function until the test exits, either normally or
> > > + * due to a failure.  @ctx is passed as additional context. All functions
> > > + * registered with kunit_add_action() will execute in the opposite order to that
> > > + * they were registered in.
> > > + *
> > > + * This is useful for cleaning up allocated memory and resources.
> > > + *
> > > + * Returns:
> > > + *   An opaque "cancellation token", or NULL on error. Pass this token to
> > > + *   kunit_remove_action_token() in order to cancel the deferred execution of
> > > + *   func().
> > > + */
> > > +struct kunit_action_ctx *kunit_add_action(struct kunit *test, kunit_defer_function_t func,
> > > +                     void *ctx, gfp_t internal_gfp);
> >
> > I've tried to leverage kunit_add_action() today, and I'm wondering if
> > passing the struct kunit pointer to the deferred function would help.
> >
> > The code I'm struggling with is something like:
> >
> > > static int test_init(struct kunit *test)
> > > {
> > >         priv = kunit_kzalloc(sizeof(*priv), GFP_KERNEL);
> > >         KUNIT_ASSERT_NOT_NULL(test, priv);
> > >         test->priv = priv;
> > >
> > >         priv->dev = alloc_device();
> > >
> > >         return 0;
> > > }
> >
> > and then in the test itself:
> >
> > > static void actual_test(struct kunit *test)
> > > {
> > >         struct test_priv *priv = test->priv;
> > >
> > >         id = allocate_buffer(priv->dev);
> > >
> > >         KUNIT_EXPECT_EQ(test, id, 42);
> > >
> > >         free_buffer(priv->dev, id);
> > > }
> >
> > I'd like to turn free_buffer an action registered right after allocate
> > buffer. However, since it takes several arguments and kunit_add_action
> > expects a single pointer, we would need to create a structure for it,
> > allocate it, fill it, and then free it when the action has ran.
> >
> > It creates a lot of boilerplate, while if we were passing the pointer to
> > struct kunit we could access the context of the test as well, and things
> > would be much simpler.
>
> The question seems to be what about the typical use-case. I was always
> imagining calling functions like kfree/kfree_skb which often only
> require a single argument.

Yeah, my thought was that having just the one argument would be
easiest for re-using existing functions. That being said, implementing
a simple wrapper which just discards the 'test' argument is probably
more ergonomic than having to write all the struct manipulation stuff,
so it depends a bit on what proves the more common case.

> For arbitrary arguments, a struct and custom free function will be
> needed. At that point, maybe it is fair to assume that API users will
> use the resource API directly, doing the same trick as kunit_add_action
> and storing the arguments together with struct kunit_resource.
>
At this point, I'd still probably use the kunit_add_action() API, just
because the resource one adds yet more complication with the 'init'
function and the reference counting. I have some vague plans to
simplify it a bit, but still definitely wouldn't rule out using the
action API here, even if it involves managing structs.


> That said, maybe one could add it as a second argument? It is a little
> bit weird API wise, but it would allow simply casting single-argument
> functions in order to ignore "struct kunit *" argument.
>

Ooh... that's evil in a particularly fun way. :-) It'd definitely be
convenient in both cases (we usually need to cast kfree() anyway for
const reasons), so I'm a bit tempted. Do we know that this would work
with the calling convention on all architectures? I'm not aware of the
kernel using anything like stdcall where the callee pops the stack on
return, but there definitely could be some architecture which does...

Cheers,

-- David

Maxime Ripard April 17, 2023, 11:07 a.m. UTC | #11

Hi David,

On Sat, Apr 15, 2023 at 04:42:27PM +0800, David Gow wrote:
> On Fri, 14 Apr 2023 at 18:02, <[e-mail]> wrote:
> >
> > Hi David,
> >
> > On Fri, Mar 31, 2023 at 04:04:09PM +0800, David Gow wrote:
> > > Many uses of the KUnit resource system are intended to simply defer
> > > calling a function until the test exits (be it due to success or
> > > failure). The existing kunit_alloc_resource() function is often used for
> > > this, but was awkward to use (requiring passing NULL init functions, etc),
> > > and returned a resource without incrementing its reference count, which
> > > -- while okay for this use-case -- could cause problems in others.
> > >
> > > Instead, introduce a simple kunit_add_action() API: a simple function
> > > (returning nothing, accepting a single void* argument) can be scheduled
> > > to be called when the test exits. Deferred actions are called in the
> > > opposite order to that which they were registered.
> > >
> > > This mimics the devres API, devm_add_action(), and also provides
> > > kunit_remove_action(), to cancel a deferred action, and
> > > kunit_release_action() to trigger one early.
> > >
> > > This is implemented as a resource under the hood, so the ordering
> > > between resource cleanup and deferred functions is maintained.
> > >
> > > Signed-off-by: David Gow <[e-mail]>
> > > ---
> > >
> > > Changes since RFC v1:
> > > https://lore.kernel.org/linux-kselftest/[e-mail]/
> > > - Rename functions to better match the devm_* APIs. (Thanks Maxime)
> > > - Embed the kunit_resource in struct kunit_action_ctx to avoid an extra
> > >   allocation (Thanks Benjamin)
> > > - Use 'struct kunit_action_ctx' as the type for cancellation tokens
> > >   (Thanks Benjamin)
> > > - Add tests.
> > >
> > > ---
> > >  include/kunit/resource.h |  89 ++++++++++++++++++++++++++++
> > >  lib/kunit/kunit-test.c   | 123 ++++++++++++++++++++++++++++++++++++++-
> > >  lib/kunit/resource.c     |  99 +++++++++++++++++++++++++++++++
> > >  3 files changed, 310 insertions(+), 1 deletion(-)
> > >
> > > diff --git a/include/kunit/resource.h b/include/kunit/resource.h
> > > index c0d88b318e90..15efd8924666 100644
> > > --- a/include/kunit/resource.h
> > > +++ b/include/kunit/resource.h
> > > @@ -387,4 +387,93 @@ static inline int kunit_destroy_named_resource(struct kunit *test,
> > >   */
> > >  void kunit_remove_resource(struct kunit *test, struct kunit_resource *res);
> > >
> > > +typedef void (*kunit_defer_function_t)(void *ctx);
> > > +
> > > +/* An opaque token to a deferred action. */
> > > +struct kunit_action_ctx;
> > > +
> > > +/**
> > > + * kunit_add_action() - Defer an 'action' (function call) until the test ends.
> > > + * @test: Test case to associate the action with.
> > > + * @func: The function to run on test exit
> > > + * @ctx: Data passed into @func
> > > + * @internal_gfp: gfp to use for internal allocations, if unsure, use GFP_KERNEL
> > > + *
> > > + * Defer the execution of a function until the test exits, either normally or
> > > + * due to a failure.  @ctx is passed as additional context. All functions
> > > + * registered with kunit_add_action() will execute in the opposite order to that
> > > + * they were registered in.
> > > + *
> > > + * This is useful for cleaning up allocated memory and resources.
> > > + *
> > > + * Returns:
> > > + *   An opaque "cancellation token", or NULL on error. Pass this token to
> > > + *   kunit_remove_action_token() in order to cancel the deferred execution of
> > > + *   func().
> > > + */
> > > +struct kunit_action_ctx *kunit_add_action(struct kunit *test, kunit_defer_function_t func,
> > > +                   void *ctx, gfp_t internal_gfp);
> >
> > I've tried to leverage kunit_add_action() today, and I'm wondering if
> > passing the struct kunit pointer to the deferred function would help.
> >
> 
> I'm tempted, but it does make the case where we just want to cast,
> e.g., kfree() directly to an action pointer more difficult. Not that
> that's a deal-blocker, but it was convenient...
> 
> > The code I'm struggling with is something like:
> >
> > > static int test_init(struct kunit *test)
> > > {
> > >       priv = kunit_kzalloc(sizeof(*priv), GFP_KERNEL);
> > >       KUNIT_ASSERT_NOT_NULL(test, priv);
> > >       test->priv = priv;
> > >
> > >       priv->dev = alloc_device();
> > >
> > >       return 0;
> > > }
> >
> > and then in the test itself:
> >
> > > static void actual_test(struct kunit *test)
> > > {
> > >       struct test_priv *priv = test->priv;
> > >
> > >       id = allocate_buffer(priv->dev);
> > >
> > >       KUNIT_EXPECT_EQ(test, id, 42);
> > >
> > >       free_buffer(priv->dev, id);
> > > }
> >
> > I'd like to turn free_buffer an action registered right after allocate
> > buffer. However, since it takes several arguments and kunit_add_action
> > expects a single pointer, we would need to create a structure for it,
> > allocate it, fill it, and then free it when the action has ran.
> 
> The general case of wanting multiple arguments to an action is a bit
> complicated. My plan was to initially support just the one argument,
> and deal with more complicated cases later. Ideas included:
> - using a struct like you suggest, possibly with some macro magic to
> make it easier,
> - having a bunch of very similar implementations of
> kunit_add_action{2,3,4,..}(), which accept 2,3,4,... arguments,
> - something horrible and architecture-specific with manually writing
> out arguments to the stack (or registers)
> 
> None of those sounded particularly pleasant, though. My suspicion is
> that the "right" way of doing this is to maybe have one or two helpers
> for common cases (e.g., 2 arguments), and just suggest people create a
> structure for anything more complicated, but I'd love a nicer
> solution.
> 
> >
> > It creates a lot of boilerplate, while if we were passing the pointer to
> > struct kunit we could access the context of the test as well, and things
> > would be much simpler.
> 
> For the test context specifically, can you just use kunit_get_current_test()?

I wasn't aware that it was a solution, but it looks like a good compromise :)

Maxime

Patch

diff --git a/include/kunit/resource.h b/include/kunit/resource.h
index c0d88b318e90..15efd8924666 100644
--- a/include/kunit/resource.h
+++ b/include/kunit/resource.h
@@ -387,4 +387,93 @@  static inline int kunit_destroy_named_resource(struct kunit *test,
  */
 void kunit_remove_resource(struct kunit *test, struct kunit_resource *res);
 
+typedef void (*kunit_defer_function_t)(void *ctx);
+
+/* An opaque token to a deferred action. */
+struct kunit_action_ctx;
+
+/**
+ * kunit_add_action() - Defer an 'action' (function call) until the test ends.
+ * @test: Test case to associate the action with.
+ * @func: The function to run on test exit
+ * @ctx: Data passed into @func
+ * @internal_gfp: gfp to use for internal allocations, if unsure, use GFP_KERNEL
+ *
+ * Defer the execution of a function until the test exits, either normally or
+ * due to a failure.  @ctx is passed as additional context. All functions
+ * registered with kunit_add_action() will execute in the opposite order to that
+ * they were registered in.
+ *
+ * This is useful for cleaning up allocated memory and resources.
+ *
+ * Returns:
+ *   An opaque "cancellation token", or NULL on error. Pass this token to
+ *   kunit_remove_action_token() in order to cancel the deferred execution of
+ *   func().
+ */
+struct kunit_action_ctx *kunit_add_action(struct kunit *test, kunit_defer_function_t func,
+		      void *ctx, gfp_t internal_gfp);
+
+/**
+ * kunit_remove_action_token() - Cancel a deferred action.
+ * @test: Test case the action is associated with.
+ * @cancel_token: The cancellation token returned by kunit_add_action()
+ *
+ * Prevent an action deferred using kunit_add_action() from executing when the
+ * test ends.
+ *
+ * You can also use the (test, function, context) triplet to remove an action
+ * with kunit_remove_action().
+ */
+void kunit_remove_action_token(struct kunit *test, struct kunit_action_ctx *cancel_token);
+
+/**
+ * kunit_release_action_token() - Trigger a deferred action immediately.
+ * @test: Test case the action is associated with.
+ * @cancel_token: The cancellation token returned by kunit_add_action()
+ *
+ * Execute an action immediately, instead of waiting for the test to end.
+ *
+ * You can also use the (test, function, context) triplet to trigger an action
+ * with kunit_release_action().
+ */
+void kunit_release_action_token(struct kunit *test, struct kunit_action_ctx *cancel_token);
+
+/**
+ * kunit_remove_action() - Cancel a matching deferred action.
+ * @test: Test case the action is associated with.
+ * @func: The deferred function to cancel.
+ * @ctx: The context passed to the deferred function to trigger.
+ *
+ * Prevent an action deferred via kunit_add_action() from executing when the
+ * test terminates..
+ * Unlike kunit_remove_action_token(), this takes the (func, ctx) pair instead of
+ * the cancellation token. If that function/context pair was deferred multiple
+ * times, only the most recent one will be cancelled.
+ */
+void kunit_remove_action(struct kunit *test,
+			 kunit_defer_function_t func,
+			 void *ctx);
+
+/**
+ * kunit_release_action() - Run a matching action call immediately.
+ * @test: Test case the action is associated with.
+ * @func: The deferred function to trigger.
+ * @ctx: The context passed to the deferred function to trigger.
+ *
+ * Execute a function deferred via kunit_add_action()) immediately, rather than
+ * when the test ends.
+ * Unlike kunit_release_action_token(), this takes the (func, ctx) pair instead of
+ * the cancellation token. If that function/context pair was deferred multiple
+ * times, it will only be executed once here. The most recent deferral will
+ * no longer execute when the test ends.
+ *
+ * kunit_release_action(test, func, ctx);
+ * is equivalent to
+ * func(ctx);
+ * kunit_remove_action(test, func, ctx);
+ */
+void kunit_release_action(struct kunit *test,
+			  kunit_defer_function_t func,
+			  void *ctx);
 #endif /* _KUNIT_RESOURCE_H */
diff --git a/lib/kunit/kunit-test.c b/lib/kunit/kunit-test.c
index b63595d3e241..eaca1b133922 100644
--- a/lib/kunit/kunit-test.c
+++ b/lib/kunit/kunit-test.c
@@ -111,7 +111,7 @@  struct kunit_test_resource_context {
 	struct kunit test;
 	bool is_resource_initialized;
 	int allocate_order[2];
-	int free_order[2];
+	int free_order[4];
 };
 
 static int fake_resource_init(struct kunit_resource *res, void *context)
@@ -402,6 +402,123 @@  static void kunit_resource_test_named(struct kunit *test)
 	KUNIT_EXPECT_TRUE(test, list_empty(&test->resources));
 }
 
+static void increment_int(void *ctx)
+{
+	int *i = (int *)ctx;
+	(*i)++;
+}
+
+static void kunit_resource_test_action(struct kunit *test)
+{
+	int num_actions = 0;
+
+	kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
+	KUNIT_EXPECT_EQ(test, num_actions, 0);
+	kunit_cleanup(test);
+	KUNIT_EXPECT_EQ(test, num_actions, 1);
+
+	/* Once we've cleaned up, the action queue is empty. */
+	kunit_cleanup(test);
+	KUNIT_EXPECT_EQ(test, num_actions, 1);
+
+	/* Check the same function can be deferred multiple times. */
+	kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
+	kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
+	kunit_cleanup(test);
+	KUNIT_EXPECT_EQ(test, num_actions, 3);
+}
+static void kunit_resource_test_remove_action(struct kunit *test)
+{
+	int num_actions = 0;
+	struct kunit_action_ctx *cancel_token;
+	struct kunit_action_ctx *cancel_token2;
+
+	cancel_token = kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
+	KUNIT_EXPECT_EQ(test, num_actions, 0);
+
+	kunit_remove_action_token(test, cancel_token);
+	kunit_cleanup(test);
+	KUNIT_EXPECT_EQ(test, num_actions, 0);
+
+	/* Check calls from the same function/context pair can be cancelled independently*/
+	cancel_token = kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
+	cancel_token2 = kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
+	kunit_remove_action_token(test, cancel_token);
+	kunit_cleanup(test);
+	KUNIT_EXPECT_EQ(test, num_actions, 1);
+
+	/* Also check that we can cancel just one of the identical function/context pairs. */
+	cancel_token = kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
+	cancel_token2 = kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
+	kunit_remove_action(test, increment_int, &num_actions);
+	kunit_cleanup(test);
+	KUNIT_EXPECT_EQ(test, num_actions, 2);
+}
+static void kunit_resource_test_release_action(struct kunit *test)
+{
+	int num_actions = 0;
+	struct kunit_action_ctx *cancel_token;
+	struct kunit_action_ctx *cancel_token2;
+
+	cancel_token = kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
+	KUNIT_EXPECT_EQ(test, num_actions, 0);
+	/* Runs immediately on trigger. */
+	kunit_release_action_token(test, cancel_token);
+	KUNIT_EXPECT_EQ(test, num_actions, 1);
+
+	/* Doesn't run again on test exit. */
+	kunit_cleanup(test);
+	KUNIT_EXPECT_EQ(test, num_actions, 1);
+
+	/* Check calls from the same function/context pair can be triggered independently*/
+	cancel_token = kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
+	cancel_token2 = kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
+	kunit_release_action_token(test, cancel_token);
+	KUNIT_EXPECT_EQ(test, num_actions, 2);
+	kunit_cleanup(test);
+	KUNIT_EXPECT_EQ(test, num_actions, 3);
+
+	/* Also check that we can trigger just one of the identical function/context pairs. */
+	kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
+	kunit_add_action(test, increment_int, &num_actions, GFP_KERNEL);
+	kunit_release_action(test, increment_int, &num_actions);
+	KUNIT_EXPECT_EQ(test, num_actions, 4);
+	kunit_cleanup(test);
+	KUNIT_EXPECT_EQ(test, num_actions, 5);
+}
+static void action_order_1(void *ctx)
+{
+	struct kunit_test_resource_context *res_ctx = (struct kunit_test_resource_context *)ctx;
+
+	KUNIT_RESOURCE_TEST_MARK_ORDER(res_ctx, free_order, 1);
+	kunit_log(KERN_INFO, current->kunit_test, "action_order_1");
+}
+static void action_order_2(void *ctx)
+{
+	struct kunit_test_resource_context *res_ctx = (struct kunit_test_resource_context *)ctx;
+
+	KUNIT_RESOURCE_TEST_MARK_ORDER(res_ctx, free_order, 2);
+	kunit_log(KERN_INFO, current->kunit_test, "action_order_2");
+}
+static void kunit_resource_test_action_ordering(struct kunit *test)
+{
+	struct kunit_test_resource_context *ctx = test->priv;
+	struct kunit_action_ctx *cancel_token;
+
+	kunit_add_action(test, action_order_1, ctx, GFP_KERNEL);
+	cancel_token = kunit_add_action(test, action_order_2, ctx, GFP_KERNEL);
+	kunit_add_action(test, action_order_1, ctx, GFP_KERNEL);
+	kunit_add_action(test, action_order_2, ctx, GFP_KERNEL);
+	kunit_remove_action(test, action_order_1, ctx);
+	kunit_release_action_token(test, cancel_token);
+	kunit_cleanup(test);
+
+	/* [2 is triggered] [2], [(1 is cancelled)] [1] */
+	KUNIT_EXPECT_EQ(test, ctx->free_order[0], 2);
+	KUNIT_EXPECT_EQ(test, ctx->free_order[1], 2);
+	KUNIT_EXPECT_EQ(test, ctx->free_order[2], 1);
+}
+
 static int kunit_resource_test_init(struct kunit *test)
 {
 	struct kunit_test_resource_context *ctx =
@@ -433,6 +550,10 @@  static struct kunit_case kunit_resource_test_cases[] = {
 	KUNIT_CASE(kunit_resource_test_proper_free_ordering),
 	KUNIT_CASE(kunit_resource_test_static),
 	KUNIT_CASE(kunit_resource_test_named),
+	KUNIT_CASE(kunit_resource_test_action),
+	KUNIT_CASE(kunit_resource_test_remove_action),
+	KUNIT_CASE(kunit_resource_test_release_action),
+	KUNIT_CASE(kunit_resource_test_action_ordering),
 	{}
 };
 
diff --git a/lib/kunit/resource.c b/lib/kunit/resource.c
index c414df922f34..824cf91e306d 100644
--- a/lib/kunit/resource.c
+++ b/lib/kunit/resource.c
@@ -77,3 +77,102 @@  int kunit_destroy_resource(struct kunit *test, kunit_resource_match_t match,
 	return 0;
 }
 EXPORT_SYMBOL_GPL(kunit_destroy_resource);
+
+struct kunit_action_ctx {
+	struct kunit_resource res;
+	kunit_defer_function_t func;
+	void *ctx;
+};
+
+static void __kunit_action_free(struct kunit_resource *res)
+{
+	struct kunit_action_ctx *action_ctx = container_of(res, struct kunit_action_ctx, res);
+
+	action_ctx->func(action_ctx->ctx);
+}
+
+struct kunit_action_ctx *kunit_add_action(struct kunit *test, kunit_defer_function_t func,
+					 void *ctx, gfp_t internal_gfp)
+{
+	struct kunit_action_ctx *action_ctx;
+
+	KUNIT_ASSERT_NOT_NULL_MSG(test, func, "Tried to action a NULL function!");
+
+	action_ctx = kzalloc(sizeof(*action_ctx), internal_gfp);
+	if (!action_ctx)
+		return NULL;
+
+	action_ctx->func = func;
+	action_ctx->ctx = ctx;
+
+	action_ctx->res.should_kfree = true;
+	/* As init is NULL, this cannot fail. */
+	__kunit_add_resource(test, NULL, __kunit_action_free, &action_ctx->res, action_ctx);
+
+	return action_ctx;
+}
+
+void kunit_remove_action_token(struct kunit *test, struct kunit_action_ctx *cancel_token)
+{
+	/* Remove the free function so we don't run the action. */
+	cancel_token->res.free = NULL;
+
+	kunit_remove_resource(test, &cancel_token->res);
+}
+
+void kunit_release_action_token(struct kunit *test, struct kunit_action_ctx *cancel_token)
+{
+	/* Removing the resource should trigger the res->free function. */
+	kunit_remove_resource(test, &cancel_token->res);
+}
+
+static bool __kunit_action_match(struct kunit *test,
+				struct kunit_resource *res, void *match_data)
+{
+	struct kunit_action_ctx *match_ctx = (struct kunit_action_ctx *)match_data;
+	struct kunit_action_ctx *res_ctx = container_of(res, struct kunit_action_ctx, res);
+
+	/* Make sure this is a free function. */
+	if (res->free != __kunit_action_free)
+		return false;
+
+	/* Both the function and context data should match. */
+	return (match_ctx->func == res_ctx->func) && (match_ctx->ctx == res_ctx->ctx);
+}
+
+void kunit_remove_action(struct kunit *test,
+			kunit_defer_function_t func,
+			void *ctx)
+{
+	struct kunit_action_ctx match_ctx, *action_ctx;
+	struct kunit_resource *res;
+
+	match_ctx.func = func;
+	match_ctx.ctx = ctx;
+
+	res = kunit_find_resource(test, __kunit_action_match, &match_ctx);
+	if (res) {
+		action_ctx = container_of(res, struct kunit_action_ctx, res);
+		kunit_remove_action_token(test, action_ctx);
+		kunit_put_resource(res);
+	}
+}
+
+void kunit_release_action(struct kunit *test,
+			 kunit_defer_function_t func,
+			 void *ctx)
+{
+	struct kunit_action_ctx match_ctx, *action_ctx;
+	struct kunit_resource *res;
+
+	match_ctx.func = func;
+	match_ctx.ctx = ctx;
+
+	res = kunit_find_resource(test, __kunit_action_match, &match_ctx);
+	if (res) {
+		action_ctx = container_of(res, struct kunit_action_ctx, res);
+		kunit_release_action_token(test, action_ctx);
+		/* We have to put() this here, else free won't be called. */
+		kunit_put_resource(res);
+	}
+}