diff mbox series

[v2,4/6] kunit: Allow function redirection outside of the KUnit thread

Message ID 20240826222015.1484-5-michal.wajdeczko@intel.com (mailing list archive)
State New
Delegated to: Brendan Higgins
Headers show
Series kunit: Add macros to help write more complex tests | expand

Commit Message

Michal Wajdeczko Aug. 26, 2024, 10:20 p.m. UTC
Currently, the 'static stub' API only allows function redirection
for calls made from the kthread of the current test, which prevents
the use of this functionality when the tested code is also used by
other threads, outside of the KUnit test, like from the workqueue.

Add another set of macros to allow redirection to the replacement
functions, which, unlike the KUNIT_STATIC_STUB_REDIRECT, will
affect all calls done during the test execution.

These new stubs, named 'global', must be declared using dedicated
KUNIT_DECLARE_GLOBAL_STUB() macro and then can be placed either as
global static variables or as part of the other structures.

To properly maintain stubs lifecycle, they can be activated only
from the main KUnit context. Some precaution is taken to avoid
changing or deactivating replacement functions if one is still
used by other thread.

Signed-off-by: Michal Wajdeczko <michal.wajdeczko@intel.com>
---
Cc: Rae Moar <rmoar@google.com>
Cc: David Gow <davidgow@google.com>
Cc: Daniel Latypov <dlatypov@google.com>
Cc: Lucas De Marchi <lucas.demarchi@intel.com>
---
v2: s/FIXED_STUB/GLOBAL_STUB (David, Lucas)
    make it little more thread safe (Rae, David)
    wait until stub call finishes before test end (David)
    wait until stub call finishes before changing stub (David)
    allow stub deactivation (Rae)
    prefer kunit log (David)
---
 include/kunit/static_stub.h | 158 ++++++++++++++++++++++++++++++++++++
 lib/kunit/static_stub.c     |  49 +++++++++++
 2 files changed, 207 insertions(+)

Comments

Lucas De Marchi Aug. 27, 2024, 2:46 p.m. UTC | #1
On Tue, Aug 27, 2024 at 12:20:13AM GMT, Michal Wajdeczko wrote:
>Currently, the 'static stub' API only allows function redirection
>for calls made from the kthread of the current test, which prevents
>the use of this functionality when the tested code is also used by
>other threads, outside of the KUnit test, like from the workqueue.
>
>Add another set of macros to allow redirection to the replacement
>functions, which, unlike the KUNIT_STATIC_STUB_REDIRECT, will
>affect all calls done during the test execution.
>
>These new stubs, named 'global', must be declared using dedicated
>KUNIT_DECLARE_GLOBAL_STUB() macro and then can be placed either as
>global static variables or as part of the other structures.
>
>To properly maintain stubs lifecycle, they can be activated only
>from the main KUnit context. Some precaution is taken to avoid
>changing or deactivating replacement functions if one is still
>used by other thread.
>
>Signed-off-by: Michal Wajdeczko <michal.wajdeczko@intel.com>
>---
>Cc: Rae Moar <rmoar@google.com>
>Cc: David Gow <davidgow@google.com>
>Cc: Daniel Latypov <dlatypov@google.com>
>Cc: Lucas De Marchi <lucas.demarchi@intel.com>
>---
>v2: s/FIXED_STUB/GLOBAL_STUB (David, Lucas)
>    make it little more thread safe (Rae, David)
>    wait until stub call finishes before test end (David)
>    wait until stub call finishes before changing stub (David)
>    allow stub deactivation (Rae)
>    prefer kunit log (David)
>---
> include/kunit/static_stub.h | 158 ++++++++++++++++++++++++++++++++++++
> lib/kunit/static_stub.c     |  49 +++++++++++
> 2 files changed, 207 insertions(+)
>
>diff --git a/include/kunit/static_stub.h b/include/kunit/static_stub.h
>index bf940322dfc0..42a70dcefb56 100644
>--- a/include/kunit/static_stub.h
>+++ b/include/kunit/static_stub.h
>@@ -12,12 +12,15 @@
>
> /* If CONFIG_KUNIT is not enabled, these stubs quietly disappear. */
> #define KUNIT_STATIC_STUB_REDIRECT(real_fn_name, args...) do {} while (0)
>+#define KUNIT_GLOBAL_STUB_REDIRECT(stub_name, args...) do {} while (0)
>+#define KUNIT_DECLARE_GLOBAL_STUB(stub_name, stub_type)
>
> #else
>
> #include <kunit/test.h>
> #include <kunit/test-bug.h>
>
>+#include <linux/cleanup.h> /* for CLASS */
> #include <linux/compiler.h> /* for {un,}likely() */
> #include <linux/sched.h> /* for task_struct */
>
>@@ -109,5 +112,160 @@ void __kunit_activate_static_stub(struct kunit *test,
>  */
> void kunit_deactivate_static_stub(struct kunit *test, void *real_fn_addr);
>
>+/**
>+ * struct kunit_global_stub - Represents a context of global function stub.
>+ * @replacement: The address of replacement function.
>+ * @owner: The KUnit test that owns the stub, valid only when @busy > 0.
>+ * @busy: The stub busyness counter incremented on entry to the replacement
>+ *        function, decremented on exit, used to signal if the stub is idle.
>+ * @idle: The completion state to indicate when the stub is idle again.
>+ *
>+ * This structure is for KUnit internal use only.
>+ * See KUNIT_DECLARE_GLOBAL_STUB().
>+ */
>+struct kunit_global_stub {
>+	void *replacement;
>+	struct kunit *owner;
>+	atomic_t busy;
>+	struct completion idle;
>+};
>+
>+/**
>+ * KUNIT_DECLARE_GLOBAL_STUB() - Declare a global function stub.
>+ * @stub_name: The name of the stub, must be a valid identifier
>+ * @stub_type: The type of the function that this stub will replace
>+ *
>+ * This macro will declare new identifier of an anonymous type that will
>+ * represent global stub function that could be used by KUnit. It can be stored
>+ * outside of the KUnit code. If the CONFIG_KUNIT is not enabled this will
>+ * be evaluated to an empty statement.
>+ *
>+ * The anonymous type introduced by this macro is mostly a wrapper to generic
>+ * struct kunit_global_stub but with additional dummy member, that is never
>+ * used directly, but is needed to maintain the type of the stub function.
>+ */
>+#define KUNIT_DECLARE_GLOBAL_STUB(stub_name, stub_type)				\
>+union {										\
>+	struct kunit_global_stub base;						\
>+	typeof(stub_type) dummy;						\
>+} stub_name
>+
>+/* Internal struct to define guard class */
>+struct kunit_global_stub_guard {
>+	struct kunit_global_stub *stub;
>+	void *active_replacement;
>+};
>+
>+/* Internal class used to guard stub calls */
>+DEFINE_CLASS(kunit_global_stub_guard,
>+	     struct kunit_global_stub_guard,
>+	     ({
>+		struct kunit_global_stub *stub = _T.stub;
>+		bool active = !!_T.active_replacement;

I'd call this `bool active_replacement` as it's not the same thing as
the active below.

>+
>+		if (active && !atomic_dec_return(&stub->busy))
>+			complete_all(&stub->idle);
>+	     }),
>+	     ({
>+		class_kunit_global_stub_guard_t guard;
>+		bool active = !!atomic_inc_not_zero(&stub->busy);
>+
>+		guard.stub = stub;
>+		guard.active_replacement = active ? READ_ONCE(stub->replacement) : NULL;
>+
>+		guard;
>+	     }),
>+	     struct kunit_global_stub *stub)
>+
>+/**
>+ * KUNIT_GLOBAL_STUB_REDIRECT() - Call a fixed function stub if activated.
>+ * @stub: The function stub declared using KUNIT_DECLARE_GLOBAL_STUB()
>+ * @args: All of the arguments passed to this stub
>+ *
>+ * This is a function prologue which is used to allow calls to the current
>+ * function to be redirected if a KUnit is running. If the KUnit is not
>+ * running or stub is not yet activated the function will continue execution
>+ * as normal.
>+ *
>+ * The function stub must be declared with KUNIT_DECLARE_GLOBAL_STUB() that is
>+ * stored in a place that is accessible from both the test code, which will
>+ * activate this stub using kunit_activate_global_stub(), and from the function,
>+ * where we will do this redirection using KUNIT_GLOBAL_STUB_REDIRECT().
>+ *
>+ * Unlike the KUNIT_STATIC_STUB_REDIRECT(), this redirection will work
>+ * even if the caller is not in a KUnit context (like a worker thread).
>+ *
>+ * Example:
>+ *
>+ * .. code-block:: c
>+ *
>+ *	KUNIT_DECLARE_GLOBAL_STUB(func_stub, int (*)(int n));
>+ *
>+ *	int real_func(int n)
>+ *	{
>+ *		KUNIT_GLOBAL_STUB_REDIRECT(func_stub, n);
>+ *		return n + 1;
>+ *	}
>+ *
>+ *	int replacement_func(int n)
>+ *	{
>+ *		return n + 100;
>+ *	}
>+ *
>+ *	void example_test(struct kunit *test)
>+ *	{
>+ *		KUNIT_EXPECT_EQ(test, real_func(1), 2);
>+ *		kunit_activate_global_stub(test, func_stub, replacement_func);
>+ *		KUNIT_EXPECT_EQ(test, real_func(1), 101);
>+ *	}
>+ */
>+#define KUNIT_GLOBAL_STUB_REDIRECT(stub, args...) do {					\
>+	if (kunit_is_running()) {							\
>+		typeof(stub) *__stub = &(stub);						\
>+		CLASS(kunit_global_stub_guard, guard)(&__stub->base);			\
>+		typeof(__stub->dummy) replacement = guard.active_replacement;		\
>+		if (unlikely(replacement)) {						\
>+			kunit_info(__stub->base.owner, "%s: redirecting to %ps\n",	\
>+				   __func__, replacement);				\
>+			return replacement(args);					\
>+		}									\
>+	}										\
>+} while (0)
>+
>+void __kunit_activate_global_stub(struct kunit *test, struct kunit_global_stub *stub,
>+				  void *replacement_addr);
>+
>+/**
>+ * kunit_activate_global_stub() - Setup a fixed function stub.

s/fixed/global/ here and every where else below

>+ * @test: Test case that wants to activate a fixed function stub
>+ * @stub: The location of the function stub pointer
>+ * @replacement: The replacement function
>+ *
>+ * This helper setups a function stub with the replacement function.
>+ * It will also automatically deactivate the stub at the test end.
>+ *
>+ * The redirection can be disabled with kunit_deactivate_global_stub().
>+ * The stub must be declared using KUNIT_DECLARE_GLOBAL_STUB().
>+ */
>+#define kunit_activate_global_stub(test, stub, replacement) do {		\
>+	typeof(stub) *__stub = &(stub);						\
>+	typecheck_fn(typeof(__stub->dummy), (replacement));			\
>+	__kunit_activate_global_stub((test), &__stub->base, (replacement));	\
>+} while (0)
>+
>+void __kunit_deactivate_global_stub(struct kunit *test, struct kunit_global_stub *stub);
>+
>+/**
>+ * kunit_deactivate_global_stub() - Disable a fixed function stub.
>+ * @test: Test case that wants to deactivate a fixed function stub
>+ * @stub: The location of the function stub pointer
>+ *
>+ * The stub must be declared using KUNIT_DECLARE_GLOBAL_STUB().
>+ */
>+#define kunit_deactivate_global_stub(test, stub) do {				\
>+	typeof(stub) *__stub = &(stub);						\
>+	__kunit_deactivate_global_stub((test), &__stub->base);			\
>+} while (0)
>+
> #endif
> #endif
>diff --git a/lib/kunit/static_stub.c b/lib/kunit/static_stub.c
>index 92b2cccd5e76..799a7271dc5b 100644
>--- a/lib/kunit/static_stub.c
>+++ b/lib/kunit/static_stub.c
>@@ -121,3 +121,52 @@ void __kunit_activate_static_stub(struct kunit *test,
> 	}
> }
> EXPORT_SYMBOL_GPL(__kunit_activate_static_stub);
>+
>+static void sanitize_global_stub(void *data)
>+{
>+	struct kunit *test = kunit_get_current_test();
>+	struct kunit_global_stub *stub =  data;
>+
>+	KUNIT_EXPECT_NE(test, 0, atomic_read(&stub->busy));

shouldn't sanitize_ be unconditional and do nothing in this case?

>+	KUNIT_EXPECT_PTR_EQ(test, test, READ_ONCE(stub->owner));
>+
>+	reinit_completion(&stub->idle);
>+	if (!atomic_dec_and_test(&stub->busy)) {
>+		kunit_info(test, "waiting for %ps\n", stub->replacement);
>+		KUNIT_EXPECT_EQ(test, 0, wait_for_completion_interruptible(&stub->idle));

what's preventing stub->busy going to 1 again after this?

Lucas De Marchi

>+	}
>+
>+	WRITE_ONCE(stub->owner, NULL);
>+	WRITE_ONCE(stub->replacement, NULL);
>+}
>+
>+/*
>+ * Helper function for kunit_activate_global_stub(). The macro does
>+ * typechecking, so use it instead.
>+ */
>+void __kunit_activate_global_stub(struct kunit *test,
>+				  struct kunit_global_stub *stub,
>+				  void *replacement_addr)
>+{
>+	KUNIT_ASSERT_NOT_ERR_OR_NULL(test, stub);
>+	KUNIT_ASSERT_NOT_ERR_OR_NULL(test, replacement_addr);
>+	if (atomic_read(&stub->busy))
>+		kunit_release_action(test, sanitize_global_stub, stub);
>+	else
>+		init_completion(&stub->idle);
>+	WRITE_ONCE(stub->owner, test);
>+	WRITE_ONCE(stub->replacement, replacement_addr);
>+	KUNIT_ASSERT_EQ(test, 1, atomic_inc_return(&stub->busy));
>+	KUNIT_ASSERT_EQ(test, 0, kunit_add_action_or_reset(test, sanitize_global_stub, stub));
>+}
>+EXPORT_SYMBOL_GPL(__kunit_activate_global_stub);
>+
>+/*
>+ * Helper function for kunit_deactivate_global_stub(). Use it instead.
>+ */
>+void __kunit_deactivate_global_stub(struct kunit *test, struct kunit_global_stub *stub)
>+{
>+	if (atomic_read(&stub->busy))
>+		kunit_release_action(test, sanitize_global_stub, stub);
>+}
>+EXPORT_SYMBOL_GPL(__kunit_deactivate_global_stub);
>-- 
>2.43.0
>
Michal Wajdeczko Aug. 27, 2024, 8:30 p.m. UTC | #2
On 27.08.2024 16:46, Lucas De Marchi wrote:
> On Tue, Aug 27, 2024 at 12:20:13AM GMT, Michal Wajdeczko wrote:
>> Currently, the 'static stub' API only allows function redirection
>> for calls made from the kthread of the current test, which prevents
>> the use of this functionality when the tested code is also used by
>> other threads, outside of the KUnit test, like from the workqueue.
>>
>> Add another set of macros to allow redirection to the replacement
>> functions, which, unlike the KUNIT_STATIC_STUB_REDIRECT, will
>> affect all calls done during the test execution.
>>
>> These new stubs, named 'global', must be declared using dedicated
>> KUNIT_DECLARE_GLOBAL_STUB() macro and then can be placed either as
>> global static variables or as part of the other structures.
>>
>> To properly maintain stubs lifecycle, they can be activated only
>> from the main KUnit context. Some precaution is taken to avoid
>> changing or deactivating replacement functions if one is still
>> used by other thread.
>>
>> Signed-off-by: Michal Wajdeczko <michal.wajdeczko@intel.com>
>> ---
>> Cc: Rae Moar <rmoar@google.com>
>> Cc: David Gow <davidgow@google.com>
>> Cc: Daniel Latypov <dlatypov@google.com>
>> Cc: Lucas De Marchi <lucas.demarchi@intel.com>
>> ---
>> v2: s/FIXED_STUB/GLOBAL_STUB (David, Lucas)
>>    make it little more thread safe (Rae, David)
>>    wait until stub call finishes before test end (David)
>>    wait until stub call finishes before changing stub (David)
>>    allow stub deactivation (Rae)
>>    prefer kunit log (David)
>> ---
>> include/kunit/static_stub.h | 158 ++++++++++++++++++++++++++++++++++++
>> lib/kunit/static_stub.c     |  49 +++++++++++
>> 2 files changed, 207 insertions(+)
>>
>> diff --git a/include/kunit/static_stub.h b/include/kunit/static_stub.h
>> index bf940322dfc0..42a70dcefb56 100644
>> --- a/include/kunit/static_stub.h
>> +++ b/include/kunit/static_stub.h
>> @@ -12,12 +12,15 @@
>>
>> /* If CONFIG_KUNIT is not enabled, these stubs quietly disappear. */
>> #define KUNIT_STATIC_STUB_REDIRECT(real_fn_name, args...) do {} while (0)
>> +#define KUNIT_GLOBAL_STUB_REDIRECT(stub_name, args...) do {} while (0)
>> +#define KUNIT_DECLARE_GLOBAL_STUB(stub_name, stub_type)
>>
>> #else
>>
>> #include <kunit/test.h>
>> #include <kunit/test-bug.h>
>>
>> +#include <linux/cleanup.h> /* for CLASS */
>> #include <linux/compiler.h> /* for {un,}likely() */
>> #include <linux/sched.h> /* for task_struct */
>>
>> @@ -109,5 +112,160 @@ void __kunit_activate_static_stub(struct kunit
>> *test,
>>  */
>> void kunit_deactivate_static_stub(struct kunit *test, void
>> *real_fn_addr);
>>
>> +/**
>> + * struct kunit_global_stub - Represents a context of global function
>> stub.
>> + * @replacement: The address of replacement function.
>> + * @owner: The KUnit test that owns the stub, valid only when @busy > 0.
>> + * @busy: The stub busyness counter incremented on entry to the
>> replacement
>> + *        function, decremented on exit, used to signal if the stub
>> is idle.
>> + * @idle: The completion state to indicate when the stub is idle again.
>> + *
>> + * This structure is for KUnit internal use only.
>> + * See KUNIT_DECLARE_GLOBAL_STUB().
>> + */
>> +struct kunit_global_stub {
>> +    void *replacement;
>> +    struct kunit *owner;
>> +    atomic_t busy;
>> +    struct completion idle;
>> +};
>> +
>> +/**
>> + * KUNIT_DECLARE_GLOBAL_STUB() - Declare a global function stub.
>> + * @stub_name: The name of the stub, must be a valid identifier
>> + * @stub_type: The type of the function that this stub will replace
>> + *
>> + * This macro will declare new identifier of an anonymous type that will
>> + * represent global stub function that could be used by KUnit. It can
>> be stored
>> + * outside of the KUnit code. If the CONFIG_KUNIT is not enabled this
>> will
>> + * be evaluated to an empty statement.
>> + *
>> + * The anonymous type introduced by this macro is mostly a wrapper to
>> generic
>> + * struct kunit_global_stub but with additional dummy member, that is
>> never
>> + * used directly, but is needed to maintain the type of the stub
>> function.
>> + */
>> +#define KUNIT_DECLARE_GLOBAL_STUB(stub_name, stub_type)                \
>> +union {                                        \
>> +    struct kunit_global_stub base;                        \
>> +    typeof(stub_type) dummy;                        \
>> +} stub_name
>> +
>> +/* Internal struct to define guard class */
>> +struct kunit_global_stub_guard {
>> +    struct kunit_global_stub *stub;
>> +    void *active_replacement;
>> +};
>> +
>> +/* Internal class used to guard stub calls */
>> +DEFINE_CLASS(kunit_global_stub_guard,
>> +         struct kunit_global_stub_guard,
>> +         ({
>> +        struct kunit_global_stub *stub = _T.stub;
>> +        bool active = !!_T.active_replacement;
> 
> I'd call this `bool active_replacement` as it's not the same thing as
> the active below.

IMO 'active_replacement' would be even more confusing as by that name we
identify the address and here it's a flag

OTOH the 'active' in both places means more/less the same (in init below
it mean stub was 'activated' and in exit here that we used 'activated'
replacement function)

> 
>> +
>> +        if (active && !atomic_dec_return(&stub->busy))
>> +            complete_all(&stub->idle);
>> +         }),
>> +         ({
>> +        class_kunit_global_stub_guard_t guard;
>> +        bool active = !!atomic_inc_not_zero(&stub->busy);
>> +
>> +        guard.stub = stub;
>> +        guard.active_replacement = active ?
>> READ_ONCE(stub->replacement) : NULL;
>> +
>> +        guard;
>> +         }),
>> +         struct kunit_global_stub *stub)
>> +
>> +/**
>> + * KUNIT_GLOBAL_STUB_REDIRECT() - Call a fixed function stub if
>> activated.
>> + * @stub: The function stub declared using KUNIT_DECLARE_GLOBAL_STUB()
>> + * @args: All of the arguments passed to this stub
>> + *
>> + * This is a function prologue which is used to allow calls to the
>> current
>> + * function to be redirected if a KUnit is running. If the KUnit is not
>> + * running or stub is not yet activated the function will continue
>> execution
>> + * as normal.
>> + *
>> + * The function stub must be declared with
>> KUNIT_DECLARE_GLOBAL_STUB() that is
>> + * stored in a place that is accessible from both the test code,
>> which will
>> + * activate this stub using kunit_activate_global_stub(), and from
>> the function,
>> + * where we will do this redirection using KUNIT_GLOBAL_STUB_REDIRECT().
>> + *
>> + * Unlike the KUNIT_STATIC_STUB_REDIRECT(), this redirection will work
>> + * even if the caller is not in a KUnit context (like a worker thread).
>> + *
>> + * Example:
>> + *
>> + * .. code-block:: c
>> + *
>> + *    KUNIT_DECLARE_GLOBAL_STUB(func_stub, int (*)(int n));
>> + *
>> + *    int real_func(int n)
>> + *    {
>> + *        KUNIT_GLOBAL_STUB_REDIRECT(func_stub, n);
>> + *        return n + 1;
>> + *    }
>> + *
>> + *    int replacement_func(int n)
>> + *    {
>> + *        return n + 100;
>> + *    }
>> + *
>> + *    void example_test(struct kunit *test)
>> + *    {
>> + *        KUNIT_EXPECT_EQ(test, real_func(1), 2);
>> + *        kunit_activate_global_stub(test, func_stub, replacement_func);
>> + *        KUNIT_EXPECT_EQ(test, real_func(1), 101);
>> + *    }
>> + */
>> +#define KUNIT_GLOBAL_STUB_REDIRECT(stub, args...) do
>> {                    \
>> +    if (kunit_is_running()) {                            \
>> +        typeof(stub) *__stub = &(stub);                        \
>> +        CLASS(kunit_global_stub_guard,
>> guard)(&__stub->base);            \
>> +        typeof(__stub->dummy) replacement =
>> guard.active_replacement;        \
>> +        if (unlikely(replacement)) {                        \
>> +            kunit_info(__stub->base.owner, "%s: redirecting to
>> %ps\n",    \
>> +                   __func__, replacement);                \
>> +            return replacement(args);                    \
>> +        }                                    \
>> +    }                                        \
>> +} while (0)
>> +
>> +void __kunit_activate_global_stub(struct kunit *test, struct
>> kunit_global_stub *stub,
>> +                  void *replacement_addr);
>> +
>> +/**
>> + * kunit_activate_global_stub() - Setup a fixed function stub.
> 
> s/fixed/global/ here and every where else below

oops

> 
>> + * @test: Test case that wants to activate a fixed function stub
>> + * @stub: The location of the function stub pointer
>> + * @replacement: The replacement function
>> + *
>> + * This helper setups a function stub with the replacement function.
>> + * It will also automatically deactivate the stub at the test end.
>> + *
>> + * The redirection can be disabled with kunit_deactivate_global_stub().
>> + * The stub must be declared using KUNIT_DECLARE_GLOBAL_STUB().
>> + */
>> +#define kunit_activate_global_stub(test, stub, replacement) do
>> {        \
>> +    typeof(stub) *__stub = &(stub);                        \
>> +    typecheck_fn(typeof(__stub->dummy), (replacement));            \
>> +    __kunit_activate_global_stub((test), &__stub->base,
>> (replacement));    \
>> +} while (0)
>> +
>> +void __kunit_deactivate_global_stub(struct kunit *test, struct
>> kunit_global_stub *stub);
>> +
>> +/**
>> + * kunit_deactivate_global_stub() - Disable a fixed function stub.
>> + * @test: Test case that wants to deactivate a fixed function stub
>> + * @stub: The location of the function stub pointer
>> + *
>> + * The stub must be declared using KUNIT_DECLARE_GLOBAL_STUB().
>> + */
>> +#define kunit_deactivate_global_stub(test, stub) do {                \
>> +    typeof(stub) *__stub = &(stub);                        \
>> +    __kunit_deactivate_global_stub((test), &__stub->base);            \
>> +} while (0)
>> +
>> #endif
>> #endif
>> diff --git a/lib/kunit/static_stub.c b/lib/kunit/static_stub.c
>> index 92b2cccd5e76..799a7271dc5b 100644
>> --- a/lib/kunit/static_stub.c
>> +++ b/lib/kunit/static_stub.c
>> @@ -121,3 +121,52 @@ void __kunit_activate_static_stub(struct kunit
>> *test,
>>     }
>> }
>> EXPORT_SYMBOL_GPL(__kunit_activate_static_stub);
>> +
>> +static void sanitize_global_stub(void *data)
>> +{
>> +    struct kunit *test = kunit_get_current_test();
>> +    struct kunit_global_stub *stub =  data;
>> +
>> +    KUNIT_EXPECT_NE(test, 0, atomic_read(&stub->busy));
> 
> shouldn't sanitize_ be unconditional and do nothing in this case?

I just didn't like early return here, but maybe it's more correct

> 
>> +    KUNIT_EXPECT_PTR_EQ(test, test, READ_ONCE(stub->owner));
>> +
>> +    reinit_completion(&stub->idle);
>> +    if (!atomic_dec_and_test(&stub->busy)) {
>> +        kunit_info(test, "waiting for %ps\n", stub->replacement);
>> +        KUNIT_EXPECT_EQ(test, 0,
>> wait_for_completion_interruptible(&stub->idle));
> 
> what's preventing stub->busy going to 1 again after this?

at the redirection point in kunit_global_stub_guard we have

	atomic_inc_not_zero(&stub->busy);

and the activation/deactivation can only be done from the main KUnit
thread (which is here)

> 
> Lucas De Marchi
> 
>> +    }
>> +
>> +    WRITE_ONCE(stub->owner, NULL);
>> +    WRITE_ONCE(stub->replacement, NULL);
>> +}
>> +
>> +/*
>> + * Helper function for kunit_activate_global_stub(). The macro does
>> + * typechecking, so use it instead.
>> + */
>> +void __kunit_activate_global_stub(struct kunit *test,
>> +                  struct kunit_global_stub *stub,
>> +                  void *replacement_addr)
>> +{
>> +    KUNIT_ASSERT_NOT_ERR_OR_NULL(test, stub);
>> +    KUNIT_ASSERT_NOT_ERR_OR_NULL(test, replacement_addr);
>> +    if (atomic_read(&stub->busy))
>> +        kunit_release_action(test, sanitize_global_stub, stub);
>> +    else
>> +        init_completion(&stub->idle);
>> +    WRITE_ONCE(stub->owner, test);
>> +    WRITE_ONCE(stub->replacement, replacement_addr);
>> +    KUNIT_ASSERT_EQ(test, 1, atomic_inc_return(&stub->busy));
>> +    KUNIT_ASSERT_EQ(test, 0, kunit_add_action_or_reset(test,
>> sanitize_global_stub, stub));
>> +}
>> +EXPORT_SYMBOL_GPL(__kunit_activate_global_stub);
>> +
>> +/*
>> + * Helper function for kunit_deactivate_global_stub(). Use it instead.
>> + */
>> +void __kunit_deactivate_global_stub(struct kunit *test, struct
>> kunit_global_stub *stub)
>> +{
>> +    if (atomic_read(&stub->busy))
>> +        kunit_release_action(test, sanitize_global_stub, stub);
>> +}
>> +EXPORT_SYMBOL_GPL(__kunit_deactivate_global_stub);
>> -- 
>> 2.43.0
>>
Michal Wajdeczko Aug. 29, 2024, 3:56 p.m. UTC | #3
On 27.08.2024 22:30, Michal Wajdeczko wrote:
> 
> 
> On 27.08.2024 16:46, Lucas De Marchi wrote:
>> On Tue, Aug 27, 2024 at 12:20:13AM GMT, Michal Wajdeczko wrote:
>>> Currently, the 'static stub' API only allows function redirection
>>> for calls made from the kthread of the current test, which prevents
>>> the use of this functionality when the tested code is also used by
>>> other threads, outside of the KUnit test, like from the workqueue.
>>>
>>> Add another set of macros to allow redirection to the replacement
>>> functions, which, unlike the KUNIT_STATIC_STUB_REDIRECT, will
>>> affect all calls done during the test execution.
>>>
>>> These new stubs, named 'global', must be declared using dedicated
>>> KUNIT_DECLARE_GLOBAL_STUB() macro and then can be placed either as
>>> global static variables or as part of the other structures.
>>>
>>> To properly maintain stubs lifecycle, they can be activated only
>>> from the main KUnit context. Some precaution is taken to avoid
>>> changing or deactivating replacement functions if one is still
>>> used by other thread.
>>>
>>> Signed-off-by: Michal Wajdeczko <michal.wajdeczko@intel.com>
>>> ---
>>> Cc: Rae Moar <rmoar@google.com>
>>> Cc: David Gow <davidgow@google.com>
>>> Cc: Daniel Latypov <dlatypov@google.com>
>>> Cc: Lucas De Marchi <lucas.demarchi@intel.com>
>>> ---
>>> v2: s/FIXED_STUB/GLOBAL_STUB (David, Lucas)
>>>    make it little more thread safe (Rae, David)
>>>    wait until stub call finishes before test end (David)
>>>    wait until stub call finishes before changing stub (David)
>>>    allow stub deactivation (Rae)
>>>    prefer kunit log (David)
>>> ---
>>> include/kunit/static_stub.h | 158 ++++++++++++++++++++++++++++++++++++
>>> lib/kunit/static_stub.c     |  49 +++++++++++
>>> 2 files changed, 207 insertions(+)
>>>
>>> diff --git a/include/kunit/static_stub.h b/include/kunit/static_stub.h
>>> index bf940322dfc0..42a70dcefb56 100644
>>> --- a/include/kunit/static_stub.h
>>> +++ b/include/kunit/static_stub.h
>>> @@ -12,12 +12,15 @@
>>>
>>> /* If CONFIG_KUNIT is not enabled, these stubs quietly disappear. */
>>> #define KUNIT_STATIC_STUB_REDIRECT(real_fn_name, args...) do {} while (0)
>>> +#define KUNIT_GLOBAL_STUB_REDIRECT(stub_name, args...) do {} while (0)
>>> +#define KUNIT_DECLARE_GLOBAL_STUB(stub_name, stub_type)
>>>
>>> #else
>>>
>>> #include <kunit/test.h>
>>> #include <kunit/test-bug.h>
>>>
>>> +#include <linux/cleanup.h> /* for CLASS */
>>> #include <linux/compiler.h> /* for {un,}likely() */
>>> #include <linux/sched.h> /* for task_struct */
>>>
>>> @@ -109,5 +112,160 @@ void __kunit_activate_static_stub(struct kunit
>>> *test,
>>>  */
>>> void kunit_deactivate_static_stub(struct kunit *test, void
>>> *real_fn_addr);
>>>
>>> +/**
>>> + * struct kunit_global_stub - Represents a context of global function
>>> stub.
>>> + * @replacement: The address of replacement function.
>>> + * @owner: The KUnit test that owns the stub, valid only when @busy > 0.
>>> + * @busy: The stub busyness counter incremented on entry to the
>>> replacement
>>> + *        function, decremented on exit, used to signal if the stub
>>> is idle.
>>> + * @idle: The completion state to indicate when the stub is idle again.
>>> + *
>>> + * This structure is for KUnit internal use only.
>>> + * See KUNIT_DECLARE_GLOBAL_STUB().
>>> + */
>>> +struct kunit_global_stub {
>>> +    void *replacement;
>>> +    struct kunit *owner;
>>> +    atomic_t busy;
>>> +    struct completion idle;
>>> +};
>>> +
>>> +/**
>>> + * KUNIT_DECLARE_GLOBAL_STUB() - Declare a global function stub.
>>> + * @stub_name: The name of the stub, must be a valid identifier
>>> + * @stub_type: The type of the function that this stub will replace
>>> + *
>>> + * This macro will declare new identifier of an anonymous type that will
>>> + * represent global stub function that could be used by KUnit. It can
>>> be stored
>>> + * outside of the KUnit code. If the CONFIG_KUNIT is not enabled this
>>> will
>>> + * be evaluated to an empty statement.
>>> + *
>>> + * The anonymous type introduced by this macro is mostly a wrapper to
>>> generic
>>> + * struct kunit_global_stub but with additional dummy member, that is
>>> never
>>> + * used directly, but is needed to maintain the type of the stub
>>> function.
>>> + */
>>> +#define KUNIT_DECLARE_GLOBAL_STUB(stub_name, stub_type)                \
>>> +union {                                        \
>>> +    struct kunit_global_stub base;                        \
>>> +    typeof(stub_type) dummy;                        \
>>> +} stub_name
>>> +
>>> +/* Internal struct to define guard class */
>>> +struct kunit_global_stub_guard {
>>> +    struct kunit_global_stub *stub;
>>> +    void *active_replacement;
>>> +};
>>> +
>>> +/* Internal class used to guard stub calls */
>>> +DEFINE_CLASS(kunit_global_stub_guard,
>>> +         struct kunit_global_stub_guard,
>>> +         ({
>>> +        struct kunit_global_stub *stub = _T.stub;
>>> +        bool active = !!_T.active_replacement;
>>
>> I'd call this `bool active_replacement` as it's not the same thing as
>> the active below.
> 
> IMO 'active_replacement' would be even more confusing as by that name we
> identify the address and here it's a flag
> 
> OTOH the 'active' in both places means more/less the same (in init below
> it mean stub was 'activated' and in exit here that we used 'activated'
> replacement function)
> 
>>
>>> +
>>> +        if (active && !atomic_dec_return(&stub->busy))
>>> +            complete_all(&stub->idle);
>>> +         }),
>>> +         ({
>>> +        class_kunit_global_stub_guard_t guard;
>>> +        bool active = !!atomic_inc_not_zero(&stub->busy);
>>> +
>>> +        guard.stub = stub;
>>> +        guard.active_replacement = active ?
>>> READ_ONCE(stub->replacement) : NULL;
>>> +
>>> +        guard;
>>> +         }),
>>> +         struct kunit_global_stub *stub)
>>> +
>>> +/**
>>> + * KUNIT_GLOBAL_STUB_REDIRECT() - Call a fixed function stub if
>>> activated.
>>> + * @stub: The function stub declared using KUNIT_DECLARE_GLOBAL_STUB()
>>> + * @args: All of the arguments passed to this stub
>>> + *
>>> + * This is a function prologue which is used to allow calls to the
>>> current
>>> + * function to be redirected if a KUnit is running. If the KUnit is not
>>> + * running or stub is not yet activated the function will continue
>>> execution
>>> + * as normal.
>>> + *
>>> + * The function stub must be declared with
>>> KUNIT_DECLARE_GLOBAL_STUB() that is
>>> + * stored in a place that is accessible from both the test code,
>>> which will
>>> + * activate this stub using kunit_activate_global_stub(), and from
>>> the function,
>>> + * where we will do this redirection using KUNIT_GLOBAL_STUB_REDIRECT().
>>> + *
>>> + * Unlike the KUNIT_STATIC_STUB_REDIRECT(), this redirection will work
>>> + * even if the caller is not in a KUnit context (like a worker thread).
>>> + *
>>> + * Example:
>>> + *
>>> + * .. code-block:: c
>>> + *
>>> + *    KUNIT_DECLARE_GLOBAL_STUB(func_stub, int (*)(int n));
>>> + *
>>> + *    int real_func(int n)
>>> + *    {
>>> + *        KUNIT_GLOBAL_STUB_REDIRECT(func_stub, n);
>>> + *        return n + 1;
>>> + *    }
>>> + *
>>> + *    int replacement_func(int n)
>>> + *    {
>>> + *        return n + 100;
>>> + *    }
>>> + *
>>> + *    void example_test(struct kunit *test)
>>> + *    {
>>> + *        KUNIT_EXPECT_EQ(test, real_func(1), 2);
>>> + *        kunit_activate_global_stub(test, func_stub, replacement_func);
>>> + *        KUNIT_EXPECT_EQ(test, real_func(1), 101);
>>> + *    }
>>> + */
>>> +#define KUNIT_GLOBAL_STUB_REDIRECT(stub, args...) do
>>> {                    \
>>> +    if (kunit_is_running()) {                            \
>>> +        typeof(stub) *__stub = &(stub);                        \
>>> +        CLASS(kunit_global_stub_guard,
>>> guard)(&__stub->base);            \
>>> +        typeof(__stub->dummy) replacement =
>>> guard.active_replacement;        \
>>> +        if (unlikely(replacement)) {                        \
>>> +            kunit_info(__stub->base.owner, "%s: redirecting to
>>> %ps\n",    \
>>> +                   __func__, replacement);                \
>>> +            return replacement(args);                    \
>>> +        }                                    \
>>> +    }                                        \
>>> +} while (0)
>>> +
>>> +void __kunit_activate_global_stub(struct kunit *test, struct
>>> kunit_global_stub *stub,
>>> +                  void *replacement_addr);
>>> +
>>> +/**
>>> + * kunit_activate_global_stub() - Setup a fixed function stub.
>>
>> s/fixed/global/ here and every where else below
> 
> oops
> 
>>
>>> + * @test: Test case that wants to activate a fixed function stub
>>> + * @stub: The location of the function stub pointer
>>> + * @replacement: The replacement function
>>> + *
>>> + * This helper setups a function stub with the replacement function.
>>> + * It will also automatically deactivate the stub at the test end.
>>> + *
>>> + * The redirection can be disabled with kunit_deactivate_global_stub().
>>> + * The stub must be declared using KUNIT_DECLARE_GLOBAL_STUB().
>>> + */
>>> +#define kunit_activate_global_stub(test, stub, replacement) do
>>> {        \
>>> +    typeof(stub) *__stub = &(stub);                        \
>>> +    typecheck_fn(typeof(__stub->dummy), (replacement));            \
>>> +    __kunit_activate_global_stub((test), &__stub->base,
>>> (replacement));    \
>>> +} while (0)
>>> +
>>> +void __kunit_deactivate_global_stub(struct kunit *test, struct
>>> kunit_global_stub *stub);
>>> +
>>> +/**
>>> + * kunit_deactivate_global_stub() - Disable a fixed function stub.
>>> + * @test: Test case that wants to deactivate a fixed function stub
>>> + * @stub: The location of the function stub pointer
>>> + *
>>> + * The stub must be declared using KUNIT_DECLARE_GLOBAL_STUB().
>>> + */
>>> +#define kunit_deactivate_global_stub(test, stub) do {                \
>>> +    typeof(stub) *__stub = &(stub);                        \
>>> +    __kunit_deactivate_global_stub((test), &__stub->base);            \
>>> +} while (0)
>>> +
>>> #endif
>>> #endif
>>> diff --git a/lib/kunit/static_stub.c b/lib/kunit/static_stub.c
>>> index 92b2cccd5e76..799a7271dc5b 100644
>>> --- a/lib/kunit/static_stub.c
>>> +++ b/lib/kunit/static_stub.c
>>> @@ -121,3 +121,52 @@ void __kunit_activate_static_stub(struct kunit
>>> *test,
>>>     }
>>> }
>>> EXPORT_SYMBOL_GPL(__kunit_activate_static_stub);
>>> +
>>> +static void sanitize_global_stub(void *data)
>>> +{
>>> +    struct kunit *test = kunit_get_current_test();
>>> +    struct kunit_global_stub *stub =  data;
>>> +
>>> +    KUNIT_EXPECT_NE(test, 0, atomic_read(&stub->busy));
>>
>> shouldn't sanitize_ be unconditional and do nothing in this case?
> 
> I just didn't like early return here, but maybe it's more correct

hmm, in fact we don't need to check stub->busy prior to calling
kunit_release_action() since our goal is to detect whether stub was
activated, but this action will be released/called only if we have added
this action after the stub activation, so we can just rely on the action
management code and just keep the EXPECT_NE here as a guard

> 
>>
>>> +    KUNIT_EXPECT_PTR_EQ(test, test, READ_ONCE(stub->owner));
>>> +
>>> +    reinit_completion(&stub->idle);
>>> +    if (!atomic_dec_and_test(&stub->busy)) {
>>> +        kunit_info(test, "waiting for %ps\n", stub->replacement);
>>> +        KUNIT_EXPECT_EQ(test, 0,
>>> wait_for_completion_interruptible(&stub->idle));
>>
>> what's preventing stub->busy going to 1 again after this?
> 
> at the redirection point in kunit_global_stub_guard we have
> 
> 	atomic_inc_not_zero(&stub->busy);
> 
> and the activation/deactivation can only be done from the main KUnit
> thread (which is here)
> 
>>
>> Lucas De Marchi
>>
>>> +    }
>>> +
>>> +    WRITE_ONCE(stub->owner, NULL);
>>> +    WRITE_ONCE(stub->replacement, NULL);
>>> +}
>>> +
>>> +/*
>>> + * Helper function for kunit_activate_global_stub(). The macro does
>>> + * typechecking, so use it instead.
>>> + */
>>> +void __kunit_activate_global_stub(struct kunit *test,
>>> +                  struct kunit_global_stub *stub,
>>> +                  void *replacement_addr)
>>> +{
>>> +    KUNIT_ASSERT_NOT_ERR_OR_NULL(test, stub);
>>> +    KUNIT_ASSERT_NOT_ERR_OR_NULL(test, replacement_addr);
>>> +    if (atomic_read(&stub->busy))
>>> +        kunit_release_action(test, sanitize_global_stub, stub);
>>> +    else
>>> +        init_completion(&stub->idle);
>>> +    WRITE_ONCE(stub->owner, test);
>>> +    WRITE_ONCE(stub->replacement, replacement_addr);
>>> +    KUNIT_ASSERT_EQ(test, 1, atomic_inc_return(&stub->busy));
>>> +    KUNIT_ASSERT_EQ(test, 0, kunit_add_action_or_reset(test,
>>> sanitize_global_stub, stub));
>>> +}
>>> +EXPORT_SYMBOL_GPL(__kunit_activate_global_stub);
>>> +
>>> +/*
>>> + * Helper function for kunit_deactivate_global_stub(). Use it instead.
>>> + */
>>> +void __kunit_deactivate_global_stub(struct kunit *test, struct
>>> kunit_global_stub *stub)
>>> +{
>>> +    if (atomic_read(&stub->busy))
>>> +        kunit_release_action(test, sanitize_global_stub, stub);
>>> +}
>>> +EXPORT_SYMBOL_GPL(__kunit_deactivate_global_stub);
>>> -- 
>>> 2.43.0
>>>
diff mbox series

Patch

diff --git a/include/kunit/static_stub.h b/include/kunit/static_stub.h
index bf940322dfc0..42a70dcefb56 100644
--- a/include/kunit/static_stub.h
+++ b/include/kunit/static_stub.h
@@ -12,12 +12,15 @@ 
 
 /* If CONFIG_KUNIT is not enabled, these stubs quietly disappear. */
 #define KUNIT_STATIC_STUB_REDIRECT(real_fn_name, args...) do {} while (0)
+#define KUNIT_GLOBAL_STUB_REDIRECT(stub_name, args...) do {} while (0)
+#define KUNIT_DECLARE_GLOBAL_STUB(stub_name, stub_type)
 
 #else
 
 #include <kunit/test.h>
 #include <kunit/test-bug.h>
 
+#include <linux/cleanup.h> /* for CLASS */
 #include <linux/compiler.h> /* for {un,}likely() */
 #include <linux/sched.h> /* for task_struct */
 
@@ -109,5 +112,160 @@  void __kunit_activate_static_stub(struct kunit *test,
  */
 void kunit_deactivate_static_stub(struct kunit *test, void *real_fn_addr);
 
+/**
+ * struct kunit_global_stub - Represents a context of global function stub.
+ * @replacement: The address of replacement function.
+ * @owner: The KUnit test that owns the stub, valid only when @busy > 0.
+ * @busy: The stub busyness counter incremented on entry to the replacement
+ *        function, decremented on exit, used to signal if the stub is idle.
+ * @idle: The completion state to indicate when the stub is idle again.
+ *
+ * This structure is for KUnit internal use only.
+ * See KUNIT_DECLARE_GLOBAL_STUB().
+ */
+struct kunit_global_stub {
+	void *replacement;
+	struct kunit *owner;
+	atomic_t busy;
+	struct completion idle;
+};
+
+/**
+ * KUNIT_DECLARE_GLOBAL_STUB() - Declare a global function stub.
+ * @stub_name: The name of the stub, must be a valid identifier
+ * @stub_type: The type of the function that this stub will replace
+ *
+ * This macro will declare new identifier of an anonymous type that will
+ * represent global stub function that could be used by KUnit. It can be stored
+ * outside of the KUnit code. If the CONFIG_KUNIT is not enabled this will
+ * be evaluated to an empty statement.
+ *
+ * The anonymous type introduced by this macro is mostly a wrapper to generic
+ * struct kunit_global_stub but with additional dummy member, that is never
+ * used directly, but is needed to maintain the type of the stub function.
+ */
+#define KUNIT_DECLARE_GLOBAL_STUB(stub_name, stub_type)				\
+union {										\
+	struct kunit_global_stub base;						\
+	typeof(stub_type) dummy;						\
+} stub_name
+
+/* Internal struct to define guard class */
+struct kunit_global_stub_guard {
+	struct kunit_global_stub *stub;
+	void *active_replacement;
+};
+
+/* Internal class used to guard stub calls */
+DEFINE_CLASS(kunit_global_stub_guard,
+	     struct kunit_global_stub_guard,
+	     ({
+		struct kunit_global_stub *stub = _T.stub;
+		bool active = !!_T.active_replacement;
+
+		if (active && !atomic_dec_return(&stub->busy))
+			complete_all(&stub->idle);
+	     }),
+	     ({
+		class_kunit_global_stub_guard_t guard;
+		bool active = !!atomic_inc_not_zero(&stub->busy);
+
+		guard.stub = stub;
+		guard.active_replacement = active ? READ_ONCE(stub->replacement) : NULL;
+
+		guard;
+	     }),
+	     struct kunit_global_stub *stub)
+
+/**
+ * KUNIT_GLOBAL_STUB_REDIRECT() - Call a fixed function stub if activated.
+ * @stub: The function stub declared using KUNIT_DECLARE_GLOBAL_STUB()
+ * @args: All of the arguments passed to this stub
+ *
+ * This is a function prologue which is used to allow calls to the current
+ * function to be redirected if a KUnit is running. If the KUnit is not
+ * running or stub is not yet activated the function will continue execution
+ * as normal.
+ *
+ * The function stub must be declared with KUNIT_DECLARE_GLOBAL_STUB() that is
+ * stored in a place that is accessible from both the test code, which will
+ * activate this stub using kunit_activate_global_stub(), and from the function,
+ * where we will do this redirection using KUNIT_GLOBAL_STUB_REDIRECT().
+ *
+ * Unlike the KUNIT_STATIC_STUB_REDIRECT(), this redirection will work
+ * even if the caller is not in a KUnit context (like a worker thread).
+ *
+ * Example:
+ *
+ * .. code-block:: c
+ *
+ *	KUNIT_DECLARE_GLOBAL_STUB(func_stub, int (*)(int n));
+ *
+ *	int real_func(int n)
+ *	{
+ *		KUNIT_GLOBAL_STUB_REDIRECT(func_stub, n);
+ *		return n + 1;
+ *	}
+ *
+ *	int replacement_func(int n)
+ *	{
+ *		return n + 100;
+ *	}
+ *
+ *	void example_test(struct kunit *test)
+ *	{
+ *		KUNIT_EXPECT_EQ(test, real_func(1), 2);
+ *		kunit_activate_global_stub(test, func_stub, replacement_func);
+ *		KUNIT_EXPECT_EQ(test, real_func(1), 101);
+ *	}
+ */
+#define KUNIT_GLOBAL_STUB_REDIRECT(stub, args...) do {					\
+	if (kunit_is_running()) {							\
+		typeof(stub) *__stub = &(stub);						\
+		CLASS(kunit_global_stub_guard, guard)(&__stub->base);			\
+		typeof(__stub->dummy) replacement = guard.active_replacement;		\
+		if (unlikely(replacement)) {						\
+			kunit_info(__stub->base.owner, "%s: redirecting to %ps\n",	\
+				   __func__, replacement);				\
+			return replacement(args);					\
+		}									\
+	}										\
+} while (0)
+
+void __kunit_activate_global_stub(struct kunit *test, struct kunit_global_stub *stub,
+				  void *replacement_addr);
+
+/**
+ * kunit_activate_global_stub() - Setup a fixed function stub.
+ * @test: Test case that wants to activate a fixed function stub
+ * @stub: The location of the function stub pointer
+ * @replacement: The replacement function
+ *
+ * This helper setups a function stub with the replacement function.
+ * It will also automatically deactivate the stub at the test end.
+ *
+ * The redirection can be disabled with kunit_deactivate_global_stub().
+ * The stub must be declared using KUNIT_DECLARE_GLOBAL_STUB().
+ */
+#define kunit_activate_global_stub(test, stub, replacement) do {		\
+	typeof(stub) *__stub = &(stub);						\
+	typecheck_fn(typeof(__stub->dummy), (replacement));			\
+	__kunit_activate_global_stub((test), &__stub->base, (replacement));	\
+} while (0)
+
+void __kunit_deactivate_global_stub(struct kunit *test, struct kunit_global_stub *stub);
+
+/**
+ * kunit_deactivate_global_stub() - Disable a fixed function stub.
+ * @test: Test case that wants to deactivate a fixed function stub
+ * @stub: The location of the function stub pointer
+ *
+ * The stub must be declared using KUNIT_DECLARE_GLOBAL_STUB().
+ */
+#define kunit_deactivate_global_stub(test, stub) do {				\
+	typeof(stub) *__stub = &(stub);						\
+	__kunit_deactivate_global_stub((test), &__stub->base);			\
+} while (0)
+
 #endif
 #endif
diff --git a/lib/kunit/static_stub.c b/lib/kunit/static_stub.c
index 92b2cccd5e76..799a7271dc5b 100644
--- a/lib/kunit/static_stub.c
+++ b/lib/kunit/static_stub.c
@@ -121,3 +121,52 @@  void __kunit_activate_static_stub(struct kunit *test,
 	}
 }
 EXPORT_SYMBOL_GPL(__kunit_activate_static_stub);
+
+static void sanitize_global_stub(void *data)
+{
+	struct kunit *test = kunit_get_current_test();
+	struct kunit_global_stub *stub =  data;
+
+	KUNIT_EXPECT_NE(test, 0, atomic_read(&stub->busy));
+	KUNIT_EXPECT_PTR_EQ(test, test, READ_ONCE(stub->owner));
+
+	reinit_completion(&stub->idle);
+	if (!atomic_dec_and_test(&stub->busy)) {
+		kunit_info(test, "waiting for %ps\n", stub->replacement);
+		KUNIT_EXPECT_EQ(test, 0, wait_for_completion_interruptible(&stub->idle));
+	}
+
+	WRITE_ONCE(stub->owner, NULL);
+	WRITE_ONCE(stub->replacement, NULL);
+}
+
+/*
+ * Helper function for kunit_activate_global_stub(). The macro does
+ * typechecking, so use it instead.
+ */
+void __kunit_activate_global_stub(struct kunit *test,
+				  struct kunit_global_stub *stub,
+				  void *replacement_addr)
+{
+	KUNIT_ASSERT_NOT_ERR_OR_NULL(test, stub);
+	KUNIT_ASSERT_NOT_ERR_OR_NULL(test, replacement_addr);
+	if (atomic_read(&stub->busy))
+		kunit_release_action(test, sanitize_global_stub, stub);
+	else
+		init_completion(&stub->idle);
+	WRITE_ONCE(stub->owner, test);
+	WRITE_ONCE(stub->replacement, replacement_addr);
+	KUNIT_ASSERT_EQ(test, 1, atomic_inc_return(&stub->busy));
+	KUNIT_ASSERT_EQ(test, 0, kunit_add_action_or_reset(test, sanitize_global_stub, stub));
+}
+EXPORT_SYMBOL_GPL(__kunit_activate_global_stub);
+
+/*
+ * Helper function for kunit_deactivate_global_stub(). Use it instead.
+ */
+void __kunit_deactivate_global_stub(struct kunit *test, struct kunit_global_stub *stub)
+{
+	if (atomic_read(&stub->busy))
+		kunit_release_action(test, sanitize_global_stub, stub);
+}
+EXPORT_SYMBOL_GPL(__kunit_deactivate_global_stub);