Introduce static key in optimizations section

This patch introduces static key technique and adds a kernel module
to demostrate how to use static key to optimize an almost unlikely
branch, the module can interact with userspace thru a character device.

Signed-off-by: Peter Lin <peterlin@qilai.dev>

examples/Makefile

@@ -31,6 +31,7 @@ obj-m += bottomhalf.o
 obj-m += ioctl.o
 obj-m += vinput.o
 obj-m += vkbd.o
+obj-m += static_key.o
 
 PWD := $(CURDIR)
 

examples/static_key.c

@@ -0,0 +1,187 @@
+/*
+ * static_key.c
+ */
+
+#include <linux/atomic.h>
+#include <linux/device.h>
+#include <linux/fs.h>
+#include <linux/kernel.h> /* for sprintf() */
+#include <linux/module.h>
+#include <linux/printk.h>
+#include <linux/types.h>
+#include <linux/uaccess.h> /* for get_user and put_user */
+
+#include <asm/errno.h>
+
+static int device_open(struct inode *inode, struct file *file);
+static int device_release(struct inode *inode, struct file *file);
+static ssize_t device_read(struct file *file, char __user *buf, size_t count,
+                           loff_t *ppos);
+static ssize_t device_write(struct file *file, const char __user *buf,
+                            size_t count, loff_t *ppos);
+
+#define SUCCESS 0
+#define DEVICE_NAME "key_state"
+#define BUF_LEN 10
+
+static int major;
+
+enum {
+    CDEV_NOT_USED = 0,
+    CDEV_EXCLUSIVE_OPEN = 1,
+};
+
+static atomic_t already_open = ATOMIC_INIT(CDEV_NOT_USED);
+
+static char msg[BUF_LEN + 1];
+
+static struct class *cls;
+
+static DEFINE_STATIC_KEY_FALSE(fkey);
+
+static struct file_operations chardev_fops = {
+    .owner = THIS_MODULE,
+    .open = device_open,
+    .release = device_release,
+    .read = device_read,
+    .write = device_write,
+};
+
+static int __init chardev_init(void)
+{
+    major = register_chrdev(0, DEVICE_NAME, &chardev_fops);
+    if (major < 0) {
+        pr_alert("Registering char device failed with %d\n", major);
+        return major;
+    }
+
+    pr_info("I was assigned major number %d\n", major);
+
+    cls = class_create(THIS_MODULE, DEVICE_NAME);
+
+    device_create(cls, NULL, MKDEV(major, 0), NULL, DEVICE_NAME);
+
+    pr_info("Device created on /dev/%s\n", DEVICE_NAME);
+
+    return SUCCESS;
+}
+
+static void __exit chardev_exit(void)
+{
+    device_destroy(cls, MKDEV(major, 0));
+    class_destroy(cls);
+
+    /* Unregister the device */
+    unregister_chrdev(major, DEVICE_NAME);
+}
+
+/* Methods */
+
+/**
+ * Called when a process tried to open the device file, like
+ * cat /dev/key_state
+ */
+static int device_open(struct inode *inode, struct file *file)
+{
+    if (atomic_cmpxchg(&already_open, CDEV_NOT_USED, CDEV_EXCLUSIVE_OPEN))
+        return -EBUSY;
+
+    sprintf(msg, static_key_enabled(&fkey) ? "enabled\n" : "disabled\n");
+
+    pr_info("fastpath 1\n");
+    if (static_branch_unlikely(&fkey))
+        pr_alert("do unlikely thing\n");
+    pr_info("fastpath 2\n");
+
+    try_module_get(THIS_MODULE);
+
+    return SUCCESS;
+}
+
+/**
+ * Called when a process closes the device file
+ */
+static int device_release(struct inode *inode, struct file *file)
+{
+    /* We are now ready for our next caller. */
+    atomic_set(&already_open, CDEV_NOT_USED);
+
+    /**
+     * Decrement the usage count, or else once you opened the file, you will
+     * never get rid of the module.
+     */
+    module_put(THIS_MODULE);
+
+    return SUCCESS;
+}
+
+/**
+ * Called when a process, which already opened the dev file, attempts to
+ * read from it.
+ */
+static ssize_t device_read(struct file *filp, /* see include/linux/fs.h */
+                           char __user *buffer, /* buffer to fill with data */
+                           size_t length, /* length of the buffer */
+                           loff_t *offset)
+{
+    /* Number of the bytes actually written to the buffer */
+    int bytes_read = 0;
+    const char *msg_ptr = msg;
+
+    if (!*(msg_ptr + *offset)) { /* We are at the end of the message */
+        *offset = 0; /* reset the offset */
+        return 0; /* signify end of file */
+    }
+
+    msg_ptr += *offset;
+
+    /* Actually put the date into the buffer */
+    while (length && *msg_ptr) {
+        /**
+         * The buffer is in the user data segment, not the kernel
+         * segment so "*" assignment won't work. We have to use
+         * put_user which copies data from the kernel data segment to
+         * the user data segment.
+         */
+        put_user(*(msg_ptr++), buffer++);
+        length--;
+        bytes_read++;
+    }
+
+    *offset += bytes_read;
+
+    /* Most read functions return the number of bytes put into the buffer. */
+    return bytes_read;
+}
+
+/* Called when a process writes to dev file; echo "enable" > /dev/key_state */
+static ssize_t device_write(struct file *filp, const char __user *buffer,
+                            size_t length, loff_t *offset)
+{
+    char command[10];
+
+    if (length > 10) {
+        pr_err("command exceeded 10 char\n");
+        return -EINVAL;
+    }
+
+    if (copy_from_user(command, buffer, length))
+        return -EFAULT;
+
+    if (strncmp(command, "enable", strlen("enable")) == 0)
+        static_branch_enable(&fkey);
+    else if (strncmp(command, "disable", strlen("disable")) == 0)
+        static_branch_disable(&fkey);
+    else {
+        pr_err("Invalid command: %s\n", command);
+        return -EINVAL;
+    }
+
+    /* Again, return the number of input characters used. */
+    return length;
+}
+
+module_init(chardev_init);
+module_exit(chardev_exit);
+
+MODULE_LICENSE("GPL");

lkmpg.tex

@@ -2044,6 +2044,60 @@ When the \cpp|unlikely| macro is used, the compiler alters its machine instructi
 That avoids flushing the processor pipeline.
 The opposite happens if you use the \cpp|likely| macro.
 
+\subsection{Static keys}
+\label{sec:static_keys}
+Static keys allow us to enable or disable kernel code paths based on the runtime state of key. Its APIs have been available since 2010 (most architectures are already supported), use self-modifying code to eliminate the overhead of cache and branch prediction.
+The most typical use case of static keys is for performance-sensitive kernel code, such as tracepoints, context switching, networking, etc. These hot paths of the kernel often contain branches and can be optimized easily using this technique.
+Before we can use static keys in the kernel, we need to make sure that gcc supports \cpp|asm goto| inline assembly, and the following kernel configurations are set:
+
+\begin{code}
+CONFIG_JUMP_LABEL=y
+CONFIG_HAVE_ARCH_JUMP_LABEL=y
+CONFIG_HAVE_ARCH_JUMP_LABEL_RELATIVE=y
+\end{code}
+
+To declare a static key, we need to define a global variable using the \cpp|DEFINE_STATIC_KEY_FALSE| or \cpp|DEFINE_STATIC_KEY_TRUE| macro defined in \src{include/linux/jump\_label.h}.
+This macro initializes the key with the given initial value, which is either false or true, respectively. For example, to declare a static key with an initial value of false, we can use the following code:
+
+\begin{code}
+DEFINE_STATIC_KEY_FALSE(fkey);
+\end{code}
+
+Once the static key has been declared, we need to add branching code to the module that uses the static key.
+For example, the code includes a fastpath, where a no-op instruction will be generated at compile time as the key is initialized to false and the branch is unlikely to be taken.
+
+\begin{code}
+pr_info("fastpath 1\n");
+if (static_branch_unlikely(&fkey))
+    pr_alert("do unlikely thing\n");
+pr_info("fastpath 2\n");
+\end{code}
+
+If the key is enabled at runtime by calling \cpp|static_branch_enable(&fkey)|, the fastpath will be patched with an unconditional jump instruction to the slowpath code \cpp|pr_alert|, so the branch will always be taken until the key is disabled again.
+
+The following kernel module derived from \verb|chardev.c|, demostrates how the static key works.
+
+\samplec{examples/static_key.c}
+
+To check the state of the static key, we can use the \verb|/dev/key_state| interface.
+
+\begin{codebash}
+cat /dev/key_state
+\end{codebash}
+
+This will display the current state of the key, which is disabled by default.
+
+To change the state of the static key, we can perform a write operation on the file:
+
+\begin{codebash}
+echo enable > /dev/key_state
+\end{codebash}
+
+This will enable the static key, causing the code path to switch from the fastpath to the slowpath.
+
+In some cases, the key is enabled or disabled at initialization and never changed, we can declare a static key as read-only, which means that it can only be toggled in the module init function. To declare a read-only static key, we can use the \cpp|DEFINE_STATIC_KEY_FALSE_RO| or \cpp|DEFINE_STATIC_KEY_TRUE_RO| macro instead. Attempts to change the key at runtime will result in a page fault.
+For more information, see \href{https://www.kernel.org/doc/Documentation/static-keys.txt}{Static keys}
+
 \section{Common Pitfalls}
 \label{sec:opitfall}