lib/kunit/resource.c

   1 // SPDX-License-Identifier: GPL-2.0
   2 /*
   3  * KUnit resource API for test managed resources (allocations, etc.).
   4  *
   5  * Copyright (C) 2022, Google LLC.
   6  * Author: Daniel Latypov <dlatypov@google.com>
   7  */
   8
   9 #include <kunit/resource.h>
  10 #include <kunit/test.h>
  11 #include <linux/kref.h>
  12
  13 /*
  14  * Used for static resources and when a kunit_resource * has been created by
  15  * kunit_alloc_resource().  When an init function is supplied, @data is passed
  16  * into the init function; otherwise, we simply set the resource data field to
  17  * the data value passed in. Doesn't initialize res->should_kfree.
  18  */
  19 int __kunit_add_resource(struct kunit *test,
  20                          kunit_resource_init_t init,
  21                          kunit_resource_free_t free,
  22                          struct kunit_resource *res,
  23                          void *data)
  24 {
  25         int ret = 0;
  26         unsigned long flags;
  27
  28         res->free = free;
  29         kref_init(&res->refcount);
  30
  31         if (init) {
  32                 ret = init(res, data);
  33                 if (ret)
  34                         return ret;
  35         } else {
  36                 res->data = data;
  37         }
  38
  39         spin_lock_irqsave(&test->lock, flags);
  40         list_add_tail(&res->node, &test->resources);
  41         /* refcount for list is established by kref_init() */
  42         spin_unlock_irqrestore(&test->lock, flags);
  43
  44         return ret;
  45 }
  46 EXPORT_SYMBOL_GPL(__kunit_add_resource);
  47
  48 void kunit_remove_resource(struct kunit *test, struct kunit_resource *res)
  49 {
  50         unsigned long flags;
  51         bool was_linked;
  52
  53         spin_lock_irqsave(&test->lock, flags);
  54         was_linked = !list_empty(&res->node);
  55         list_del_init(&res->node);
  56         spin_unlock_irqrestore(&test->lock, flags);
  57
  58         if (was_linked)
  59                 kunit_put_resource(res);
  60 }
  61 EXPORT_SYMBOL_GPL(kunit_remove_resource);
  62
  63 int kunit_destroy_resource(struct kunit *test, kunit_resource_match_t match,
  64                            void *match_data)
  65 {
  66         struct kunit_resource *res = kunit_find_resource(test, match,
  67                                                          match_data);
  68
  69         if (!res)
  70                 return -ENOENT;
  71
  72         kunit_remove_resource(test, res);
  73
  74         /* We have a reference also via _find(); drop it. */
  75         kunit_put_resource(res);
  76
  77         return 0;
  78 }
  79 EXPORT_SYMBOL_GPL(kunit_destroy_resource);
  80
  81 struct kunit_action_ctx {
  82         struct kunit_resource res;
  83         kunit_action_t *func;
  84         void *ctx;
  85 };
  86
  87 static void __kunit_action_free(struct kunit_resource *res)
  88 {
  89         struct kunit_action_ctx *action_ctx = container_of(res, struct kunit_action_ctx, res);
  90
  91         action_ctx->func(action_ctx->ctx);
  92 }
  93
  94
  95 int kunit_add_action(struct kunit *test, void (*action)(void *), void *ctx)
  96 {
  97         struct kunit_action_ctx *action_ctx;
  98
  99         KUNIT_ASSERT_NOT_NULL_MSG(test, action, "Tried to action a NULL function!");
 100
 101         action_ctx = kzalloc(sizeof(*action_ctx), GFP_KERNEL);
 102         if (!action_ctx)
 103                 return -ENOMEM;
 104
 105         action_ctx->func = action;
 106         action_ctx->ctx = ctx;
 107
 108         action_ctx->res.should_kfree = true;
 109         /* As init is NULL, this cannot fail. */
 110         __kunit_add_resource(test, NULL, __kunit_action_free, &action_ctx->res, action_ctx);
 111
 112         return 0;
 113 }
 114 EXPORT_SYMBOL_GPL(kunit_add_action);
 115
 116 int kunit_add_action_or_reset(struct kunit *test, void (*action)(void *),
 117                               void *ctx)
 118 {
 119         int res = kunit_add_action(test, action, ctx);
 120
 121         if (res)
 122                 action(ctx);
 123         return res;
 124 }
 125 EXPORT_SYMBOL_GPL(kunit_add_action_or_reset);
 126
 127 static bool __kunit_action_match(struct kunit *test,
 128                                 struct kunit_resource *res, void *match_data)
 129 {
 130         struct kunit_action_ctx *match_ctx = (struct kunit_action_ctx *)match_data;
 131         struct kunit_action_ctx *res_ctx = container_of(res, struct kunit_action_ctx, res);
 132
 133         /* Make sure this is a free function. */
 134         if (res->free != __kunit_action_free)
 135                 return false;
 136
 137         /* Both the function and context data should match. */
 138         return (match_ctx->func == res_ctx->func) && (match_ctx->ctx == res_ctx->ctx);
 139 }
 140
 141 void kunit_remove_action(struct kunit *test,
 142                         kunit_action_t *action,
 143                         void *ctx)
 144 {
 145         struct kunit_action_ctx match_ctx;
 146         struct kunit_resource *res;
 147
 148         match_ctx.func = action;
 149         match_ctx.ctx = ctx;
 150
 151         res = kunit_find_resource(test, __kunit_action_match, &match_ctx);
 152         if (res) {
 153                 /* Remove the free function so we don't run the action. */
 154                 res->free = NULL;
 155                 kunit_remove_resource(test, res);
 156                 kunit_put_resource(res);
 157         }
 158 }
 159 EXPORT_SYMBOL_GPL(kunit_remove_action);
 160
 161 void kunit_release_action(struct kunit *test,
 162                          kunit_action_t *action,
 163                          void *ctx)
 164 {
 165         struct kunit_action_ctx match_ctx;
 166         struct kunit_resource *res;
 167
 168         match_ctx.func = action;
 169         match_ctx.ctx = ctx;
 170
 171         res = kunit_find_resource(test, __kunit_action_match, &match_ctx);
 172         if (res) {
 173                 kunit_remove_resource(test, res);
 174                 /* We have to put() this here, else free won't be called. */
 175                 kunit_put_resource(res);
 176         }
 177 }
 178 EXPORT_SYMBOL_GPL(kunit_release_action);