include/kunit/resource.h

   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 #ifndef _KUNIT_RESOURCE_H
  10 #define _KUNIT_RESOURCE_H
  11
  12 #include <kunit/test.h>
  13
  14 #include <linux/kref.h>
  15 #include <linux/list.h>
  16 #include <linux/slab.h>
  17 #include <linux/spinlock.h>
  18
  19 struct kunit_resource;
  20
  21 typedef int (*kunit_resource_init_t)(struct kunit_resource *, void *);
  22 typedef void (*kunit_resource_free_t)(struct kunit_resource *);
  23
  24 /**
  25  * struct kunit_resource - represents a *test managed resource*
  26  * @data: for the user to store arbitrary data.
  27  * @name: optional name
  28  * @free: a user supplied function to free the resource.
  29  *
  30  * Represents a *test managed resource*, a resource which will automatically be
  31  * cleaned up at the end of a test case. This cleanup is performed by the 'free'
  32  * function. The struct kunit_resource itself is freed automatically with
  33  * kfree() if it was allocated by KUnit (e.g., by kunit_alloc_resource()), but
  34  * must be freed by the user otherwise.
  35  *
  36  * Resources are reference counted so if a resource is retrieved via
  37  * kunit_alloc_and_get_resource() or kunit_find_resource(), we need
  38  * to call kunit_put_resource() to reduce the resource reference count
  39  * when finished with it.  Note that kunit_alloc_resource() does not require a
  40  * kunit_resource_put() because it does not retrieve the resource itself.
  41  *
  42  * Example:
  43  *
  44  * .. code-block:: c
  45  *
  46  *      struct kunit_kmalloc_params {
  47  *              size_t size;
  48  *              gfp_t gfp;
  49  *      };
  50  *
  51  *      static int kunit_kmalloc_init(struct kunit_resource *res, void *context)
  52  *      {
  53  *              struct kunit_kmalloc_params *params = context;
  54  *              res->data = kmalloc(params->size, params->gfp);
  55  *
  56  *              if (!res->data)
  57  *                      return -ENOMEM;
  58  *
  59  *              return 0;
  60  *      }
  61  *
  62  *      static void kunit_kmalloc_free(struct kunit_resource *res)
  63  *      {
  64  *              kfree(res->data);
  65  *      }
  66  *
  67  *      void *kunit_kmalloc(struct kunit *test, size_t size, gfp_t gfp)
  68  *      {
  69  *              struct kunit_kmalloc_params params;
  70  *
  71  *              params.size = size;
  72  *              params.gfp = gfp;
  73  *
  74  *              return kunit_alloc_resource(test, kunit_kmalloc_init,
  75  *                      kunit_kmalloc_free, gfp, &params);
  76  *      }
  77  *
  78  * Resources can also be named, with lookup/removal done on a name
  79  * basis also.  kunit_add_named_resource(), kunit_find_named_resource()
  80  * and kunit_destroy_named_resource().  Resource names must be
  81  * unique within the test instance.
  82  */
  83 struct kunit_resource {
  84         void *data;
  85         const char *name;
  86         kunit_resource_free_t free;
  87
  88         /* private: internal use only. */
  89         struct kref refcount;
  90         struct list_head node;
  91         bool should_kfree;
  92 };
  93
  94 /**
  95  * kunit_get_resource() - Hold resource for use.  Should not need to be used
  96  *                        by most users as we automatically get resources
  97  *                        retrieved by kunit_find_resource*().
  98  * @res: resource
  99  */
 100 static inline void kunit_get_resource(struct kunit_resource *res)
 101 {
 102         kref_get(&res->refcount);
 103 }
 104
 105 /*
 106  * Called when refcount reaches zero via kunit_put_resource();
 107  * should not be called directly.
 108  */
 109 static inline void kunit_release_resource(struct kref *kref)
 110 {
 111         struct kunit_resource *res = container_of(kref, struct kunit_resource,
 112                                                   refcount);
 113
 114         if (res->free)
 115                 res->free(res);
 116
 117         /* 'res' is valid here, as if should_kfree is set, res->free may not free
 118          * 'res' itself, just res->data
 119          */
 120         if (res->should_kfree)
 121                 kfree(res);
 122 }
 123
 124 /**
 125  * kunit_put_resource() - When caller is done with retrieved resource,
 126  *                        kunit_put_resource() should be called to drop
 127  *                        reference count.  The resource list maintains
 128  *                        a reference count on resources, so if no users
 129  *                        are utilizing a resource and it is removed from
 130  *                        the resource list, it will be freed via the
 131  *                        associated free function (if any).  Only
 132  *                        needs to be used if we alloc_and_get() or
 133  *                        find() resource.
 134  * @res: resource
 135  */
 136 static inline void kunit_put_resource(struct kunit_resource *res)
 137 {
 138         kref_put(&res->refcount, kunit_release_resource);
 139 }
 140
 141 /**
 142  * __kunit_add_resource() - Internal helper to add a resource.
 143  *
 144  * res->should_kfree is not initialised.
 145  * @test: The test context object.
 146  * @init: a user-supplied function to initialize the result (if needed).  If
 147  *        none is supplied, the resource data value is simply set to @data.
 148  *        If an init function is supplied, @data is passed to it instead.
 149  * @free: a user-supplied function to free the resource (if needed).
 150  * @res: The resource.
 151  * @data: value to pass to init function or set in resource data field.
 152  */
 153 int __kunit_add_resource(struct kunit *test,
 154                          kunit_resource_init_t init,
 155                          kunit_resource_free_t free,
 156                          struct kunit_resource *res,
 157                          void *data);
 158
 159 /**
 160  * kunit_add_resource() - Add a *test managed resource*.
 161  * @test: The test context object.
 162  * @init: a user-supplied function to initialize the result (if needed).  If
 163  *        none is supplied, the resource data value is simply set to @data.
 164  *        If an init function is supplied, @data is passed to it instead.
 165  * @free: a user-supplied function to free the resource (if needed).
 166  * @res: The resource.
 167  * @data: value to pass to init function or set in resource data field.
 168  */
 169 static inline int kunit_add_resource(struct kunit *test,
 170                                      kunit_resource_init_t init,
 171                                      kunit_resource_free_t free,
 172                                      struct kunit_resource *res,
 173                                      void *data)
 174 {
 175         res->should_kfree = false;
 176         return __kunit_add_resource(test, init, free, res, data);
 177 }
 178
 179 static inline struct kunit_resource *
 180 kunit_find_named_resource(struct kunit *test, const char *name);
 181
 182 /**
 183  * kunit_add_named_resource() - Add a named *test managed resource*.
 184  * @test: The test context object.
 185  * @init: a user-supplied function to initialize the resource data, if needed.
 186  * @free: a user-supplied function to free the resource data, if needed.
 187  * @res: The resource.
 188  * @name: name to be set for resource.
 189  * @data: value to pass to init function or set in resource data field.
 190  */
 191 static inline int kunit_add_named_resource(struct kunit *test,
 192                                            kunit_resource_init_t init,
 193                                            kunit_resource_free_t free,
 194                                            struct kunit_resource *res,
 195                                            const char *name,
 196                                            void *data)
 197 {
 198         struct kunit_resource *existing;
 199
 200         if (!name)
 201                 return -EINVAL;
 202
 203         existing = kunit_find_named_resource(test, name);
 204         if (existing) {
 205                 kunit_put_resource(existing);
 206                 return -EEXIST;
 207         }
 208
 209         res->name = name;
 210         res->should_kfree = false;
 211
 212         return __kunit_add_resource(test, init, free, res, data);
 213 }
 214
 215 /**
 216  * kunit_alloc_and_get_resource() - Allocates and returns a *test managed resource*.
 217  * @test: The test context object.
 218  * @init: a user supplied function to initialize the resource.
 219  * @free: a user supplied function to free the resource (if needed).
 220  * @internal_gfp: gfp to use for internal allocations, if unsure, use GFP_KERNEL
 221  * @context: for the user to pass in arbitrary data to the init function.
 222  *
 223  * Allocates a *test managed resource*, a resource which will automatically be
 224  * cleaned up at the end of a test case. See &struct kunit_resource for an
 225  * example.
 226  *
 227  * This is effectively identical to kunit_alloc_resource, but returns the
 228  * struct kunit_resource pointer, not just the 'data' pointer. It therefore
 229  * also increments the resource's refcount, so kunit_put_resource() should be
 230  * called when you've finished with it.
 231  *
 232  * Note: KUnit needs to allocate memory for a kunit_resource object. You must
 233  * specify an @internal_gfp that is compatible with the use context of your
 234  * resource.
 235  */
 236 static inline struct kunit_resource *
 237 kunit_alloc_and_get_resource(struct kunit *test,
 238                              kunit_resource_init_t init,
 239                              kunit_resource_free_t free,
 240                              gfp_t internal_gfp,
 241                              void *context)
 242 {
 243         struct kunit_resource *res;
 244         int ret;
 245
 246         res = kzalloc(sizeof(*res), internal_gfp);
 247         if (!res)
 248                 return NULL;
 249
 250         res->should_kfree = true;
 251
 252         ret = __kunit_add_resource(test, init, free, res, context);
 253         if (!ret) {
 254                 /*
 255                  * bump refcount for get; kunit_resource_put() should be called
 256                  * when done.
 257                  */
 258                 kunit_get_resource(res);
 259                 return res;
 260         }
 261         return NULL;
 262 }
 263
 264 /**
 265  * kunit_alloc_resource() - Allocates a *test managed resource*.
 266  * @test: The test context object.
 267  * @init: a user supplied function to initialize the resource.
 268  * @free: a user supplied function to free the resource (if needed).
 269  * @internal_gfp: gfp to use for internal allocations, if unsure, use GFP_KERNEL
 270  * @context: for the user to pass in arbitrary data to the init function.
 271  *
 272  * Allocates a *test managed resource*, a resource which will automatically be
 273  * cleaned up at the end of a test case. See &struct kunit_resource for an
 274  * example.
 275  *
 276  * Note: KUnit needs to allocate memory for a kunit_resource object. You must
 277  * specify an @internal_gfp that is compatible with the use context of your
 278  * resource.
 279  */
 280 static inline void *kunit_alloc_resource(struct kunit *test,
 281                                          kunit_resource_init_t init,
 282                                          kunit_resource_free_t free,
 283                                          gfp_t internal_gfp,
 284                                          void *context)
 285 {
 286         struct kunit_resource *res;
 287
 288         res = kzalloc(sizeof(*res), internal_gfp);
 289         if (!res)
 290                 return NULL;
 291
 292         res->should_kfree = true;
 293         if (!__kunit_add_resource(test, init, free, res, context))
 294                 return res->data;
 295
 296         return NULL;
 297 }
 298
 299 typedef bool (*kunit_resource_match_t)(struct kunit *test,
 300                                        struct kunit_resource *res,
 301                                        void *match_data);
 302
 303 /**
 304  * kunit_resource_name_match() - Match a resource with the same name.
 305  * @test: Test case to which the resource belongs.
 306  * @res: The resource.
 307  * @match_name: The name to match against.
 308  */
 309 static inline bool kunit_resource_name_match(struct kunit *test,
 310                                              struct kunit_resource *res,
 311                                              void *match_name)
 312 {
 313         return res->name && strcmp(res->name, match_name) == 0;
 314 }
 315
 316 /**
 317  * kunit_find_resource() - Find a resource using match function/data.
 318  * @test: Test case to which the resource belongs.
 319  * @match: match function to be applied to resources/match data.
 320  * @match_data: data to be used in matching.
 321  */
 322 static inline struct kunit_resource *
 323 kunit_find_resource(struct kunit *test,
 324                     kunit_resource_match_t match,
 325                     void *match_data)
 326 {
 327         struct kunit_resource *res, *found = NULL;
 328         unsigned long flags;
 329
 330         spin_lock_irqsave(&test->lock, flags);
 331
 332         list_for_each_entry_reverse(res, &test->resources, node) {
 333                 if (match(test, res, (void *)match_data)) {
 334                         found = res;
 335                         kunit_get_resource(found);
 336                         break;
 337                 }
 338         }
 339
 340         spin_unlock_irqrestore(&test->lock, flags);
 341
 342         return found;
 343 }
 344
 345 /**
 346  * kunit_find_named_resource() - Find a resource using match name.
 347  * @test: Test case to which the resource belongs.
 348  * @name: match name.
 349  */
 350 static inline struct kunit_resource *
 351 kunit_find_named_resource(struct kunit *test,
 352                           const char *name)
 353 {
 354         return kunit_find_resource(test, kunit_resource_name_match,
 355                                    (void *)name);
 356 }
 357
 358 /**
 359  * kunit_destroy_resource() - Find a kunit_resource and destroy it.
 360  * @test: Test case to which the resource belongs.
 361  * @match: Match function. Returns whether a given resource matches @match_data.
 362  * @match_data: Data passed into @match.
 363  *
 364  * RETURNS:
 365  * 0 if kunit_resource is found and freed, -ENOENT if not found.
 366  */
 367 int kunit_destroy_resource(struct kunit *test,
 368                            kunit_resource_match_t match,
 369                            void *match_data);
 370
 371 static inline int kunit_destroy_named_resource(struct kunit *test,
 372                                                const char *name)
 373 {
 374         return kunit_destroy_resource(test, kunit_resource_name_match,
 375                                       (void *)name);
 376 }
 377
 378 /**
 379  * kunit_remove_resource() - remove resource from resource list associated with
 380  *                           test.
 381  * @test: The test context object.
 382  * @res: The resource to be removed.
 383  *
 384  * Note that the resource will not be immediately freed since it is likely
 385  * the caller has a reference to it via alloc_and_get() or find();
 386  * in this case a final call to kunit_put_resource() is required.
 387  */
 388 void kunit_remove_resource(struct kunit *test, struct kunit_resource *res);
 389
 390 #endif /* _KUNIT_RESOURCE_H */