[PATCH bpf-next v7 0/4] Task local data

[PATCH bpf-next v7 0/4] Task local data

[Amery Hung]

v6 -> v7
  - Fix typos and improve the clarity of the cover letter (Emil)
  - Some variable naming changes to make it less confusing (Emil)
  - Add retry in tld_obj_init() as bpf_task_storage_get_recur() with
    BPF_LOCAL_STORAGE_GET_F_CREATE can fail if the task local storage
    is also being modified by other threads on the same CPU. This can
    be especially easy to trigger in CI VM that only has two CPUs.

    Adding bpf_preempt_{disable,enable} around bpf_task_storage_get()
    does not solve the problem as other threads competing for the percpu
    counter lock in task local storage may not limit to bpf helpers.
    Some may be calling bpf_task_storage_free() when threads exit.

    There is no good max retry under heavy contention. For the 1st
    selftest in this set, the max retry to guarantee success grows
    linearly with the thread count. A proposal is to remove the percpu
    counter by changing locks in bpf_local_storage to rqspinlock.

    An alternative is to reduce the probability of failure by allowing
    bpf syscall and struct_ops programs to use bpf_task_storage_get()
    instead of the _recur version by modifying bpf_prog_check_recur().
    This does not solve the problem in tracing programs though.

    v6: https://lore.kernel.org/bpf/20250717164842.1848817-1-ameryhung@gmail.com/

* Overview *

This patchset introduces task local data, an abstract storage type
shared between user space and bpf programs for storing data specific to
each task, and implements a library to access it. The goal is to provide
an abstraction + implementation that is efficient in data sharing and
easy to use. The motivating use case is user space hinting with
sched_ext.


* Motivating use case *

CPU schedulers can potentially make a better decision with hints from
user space process. To support experimenting user space hinting with
sched_ext, there needs a mechanism to pass a "per-task hint" from user
space to the bpf scheduler "efficiently".

A bpf feature, UPTR [0], supported by task local storage is introduced
to serve the needs. It allows pinning a user space page to the kernel
through a special field in task local storage map. This patchset further
builds an abstraction on top of it to allow user space and bpf programs
to easily share per-task data.


* Design *

Task local data is built on top of task local storage map and UPTR to
achieve fast per-task data sharing. UPTR is a type of special field
supported in task local storage map value. A user page assigned to a UPTR
will be pinned to the kernel when the map is updated. Therefore, user
space programs can update data that will be seen by bpf programs without
syscalls.

Additionally, unlike most bpf maps, task local data does not require a
static map value definition. This design is driven by sched_ext, which
would like to allow multiple developers to share a storage without the
need to explicitly agree on the layout of it. While a centralized layout
definition would have worked, the friction of synchronizing it across
different repos is not desirable. This simplify code base management and
makes experimenting easier.


* Introduction to task local data library *

Task local data library provides simple APIs for user space and bpf
through two header files, task_local_data.h and task_local_data.bpf.h,
respectively. The diagram below illustrates the basic usage.

First, an entry of data in task local data, we call it a TLD, needs to
be defined in the user space through TLD_DEFINE_KEY() with information
including the size and the name. The macro will define a global variable
key of opaque type tld_key_t associated with the TLD and initialize it.
Then, the user space program can locate the TLD by passing the key to
tld_get_data() in different thread, where fd is the file descriptor of
the underlying task local storage map. The API returns a pointer to the
TLD specific to the calling thread and will remain valid until the
thread exits. Finally, user space programs can directly read/write the
TLD without bpf syscalls.

To use task local storage on the bpf side, struct tld_keys, needs to
be defined first. The structure that will be used to create tld_key_map
should contain the keys to the TLDs used in a bpf program. In the bpf
program, tld_init_object() first needs to be called to initialize a
struct tld_object variable on the stack. Then, tld_get_data() can be
called to get a pointer to the TLD similar to the user space. The API
will use the name to lookup task local data and cache the key in task
local storage map, tld_key_map, to speed up subsequent access. The size
of the TLD is also needed to prevent out-of-bound access in bpf.


 ┌─ Application ───────────────────────────────────────────────────────┐
 │ TLD_DEFINE_KEY(kx, "X", 4);      ┌─ library A ─────────────────────┐│
 │                                  │ void func(...)                  ││
 │ int main(...)                    │ {                               ││
 │ {                                │     tld_key_t ky;               ││
 │      int *x;                     │     bool *y;                    ││
 │                                  │                                 ││
 │      x = tld_get_data(fd, kx);   │     ky = tld_create_key("Y", 1);││
 │      if (x) *x = 123;            │     y = tld_get_data(fd, ky);   ││
 │                         ┌────────┤     if (y) *y = true;           ││
 │                         │        └─────────────────────────────────┘│
 └───────┬─────────────────│───────────────────────────────────────────┘
         V                 V
 + ─ Task local data ─ ─ ─ ─ ─ +  ┌─ BPF program ──────────────────────┐
 | ┌─ tld_data_map ──────────┐ |  │ struct tld_object obj;             │
 | │ BPF Task local storage  │ |  │ bool *y;                           │
 | │                         │ |  │ int *x;                            │
 | │ tld_data_u __uptr *data │ |  │                                    │
 | │ tld_meta_u __uptr *meta │ |  │ if (tld_init_object(task, &obj))   │
 | └─────────────────────────┘ |  │     return 0;                      │
 | ┌─ tld_key_map ───────────┐ |  │                                    │
 | │ BPF Task local storage  │ |  │ x = tld_get_data(&obj, kx, "X", 4);│
 | │                         │ |<─┤ if (x) /* do something */          │
 | │ tld_key_t kx;           │ |  │                                    │
 | │ tld_key_t ky;           │ |  │ y = tld_get_data(&obj, ky, "Y", 1);│
 | └─────────────────────────┘ |  │ if (y) /* do something */          │
 + ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ +  └────────────────────────────────────┘


* Implementation *

Task local data defines the storage to be a task local storage map with
two UPTRs pointing to tld_data_u and tld_meta_u. tld_data_u, individual
to each thread, contains TLD data and the starting offset of data in a
page. tld_meta_u, shared by threads in a process, consists of the TLD
counts, total size of TLDs and tld_metadata for each TLD. tld_metadata
contains the name and the size of a TLD.

  struct tld_data_u {
          u64 start;
          char data[PAGE_SIZE - sizeof(u64)];
  };

  struct tld_meta_u {
          u8 cnt;
          u16 size;
          struct metadata metadata[TLD_DATA_CNT]; 
  };

Both user space and bpf API follow the same protocol when accessing
task local data. A pointer to a TLD is located using a key. The key is
effectively the offset of a TLD in tld_data_u->data. To define a new
TLD, the user space API TLD_DEFINE_KEY() iterates tld_meta_u->metadata
until an empty slot is found and then update it. It also adds up sizes
of prior TLDs to derive the offset (i.e., key). Then, to locate a TLD,
tld_get_data() can simply return tld_data_u->data + offset.

To locate a TLD in bpf programs, an API with the same name as the user
space, tld_get_data() can be called. It takes more effort in the first
invocation as it fetches the key by name. Internal helper,
__tld_fetch_key() will iterate tld_meta_u->metadata until the name is
found. Then, the offset is cached as key in another task local storage
map, tld_key_map. When the search fails, the current TLD count is
cached instead to skip searching metadata entries that has been searched
in future invocation. The detail of task local data operations can be
found in patch 1.


* Misc *

The metadata can potentially use run-length encoding for names to reduce
memory wastage and support save more TLDs. I have a version that works,
but the selftest takes a bit longer to finish. More investigation needed
to find the root cause. I will save this for the future when there is a
need to store more than 63 TLDs.


[0] https://lore.kernel.org/bpf/20241023234759.860539-1-martin.lau@linux.dev/

---

v5 -> v6
  - Address Andrii's comment
  - Fix verification failure in no_alu32
  - Some cleanup
  v5: https://lore.kernel.org/bpf/20250627233958.2602271-1-ameryhung@gmail.com/

v4 -> v5
  - Add an option to free memory on thread exit to prevent memory leak
  - Add an option to reduce memory waste if the allocator can
    use just enough memory to fullfill aligned_alloc() (e.g., glibc)
  - Tweak bpf API
      - Remove tld_fetch_key() as it does not work in init_tasl
      - tld_get_data() now tries to fetch key if it is not cached yet
  - Optimize bpf side tld_get_data()
      - Faster fast path
      - Less code
  - Use stdatomic.h in user space library with seq_cst order
  - Introduce TLD_DEFINE_KEY() as the default TLD creation API for
    easier memory management.
      - TLD_DEFINE_KEY() can consume memory up to a page and no memory
        is wasted since their size is known before per-thread data
        allocation.
      - tld_create_key() can only use up to TLD_DYN_DATA_SIZE. Since
        tld_create_key can run any time even after per-thread data
        allocation, it is impossible to predict the total size. A
        configurable size of memory is allocated on top of the total
        size of TLD_DEFINE_KEY() to accommodate dynamic key creation.
  - Add tld prefix to all macros
  - Replace map_update(NO_EXIST) in __tld_init_data() with cmpxchg()
  - No more +1,-1 dance on the bpf side
  - Reduce printf from ASSERT in race test
  - Try implementing run-length encoding for name and decide to
    save it for the future
  v4: https://lore.kernel.org/bpf/20250515211606.2697271-1-ameryhung@gmail.com/

v3 -> v4
  - API improvements
      - Simplify API
      - Drop string obfuscation
      - Use opaque type for key
      - Better documentation
  - Implementation
      - Switch to dynamic allocation for per-task data
      - Now offer as header-only libraries
      - No TLS map pinning; leave it to users
  - Drop pthread dependency
  - Add more invalid tld_create_key() test
  - Add a race test for tld_create_key()
  v3: https://lore.kernel.org/bpf/20250425214039.2919818-1-ameryhung@gmail.com/

---

Amery Hung (4):
  bpf: Allow syscall bpf programs to call non-recur helpers
  selftests/bpf: Introduce task local data
  selftests/bpf: Test basic task local data operations
  selftests/bpf: Test concurrent task local data key creation

 include/linux/bpf_verifier.h                  |   1 +
 .../bpf/prog_tests/task_local_data.h          | 386 ++++++++++++++++++
 .../bpf/prog_tests/test_task_local_data.c     | 297 ++++++++++++++
 .../selftests/bpf/progs/task_local_data.bpf.h | 237 +++++++++++
 .../bpf/progs/test_task_local_data.c          |  65 +++
 5 files changed, 986 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/prog_tests/task_local_data.h
 create mode 100644 tools/testing/selftests/bpf/prog_tests/test_task_local_data.c
 create mode 100644 tools/testing/selftests/bpf/progs/task_local_data.bpf.h
 create mode 100644 tools/testing/selftests/bpf/progs/test_task_local_data.c

-- 
2.47.3

[PATCH bpf-next v7 1/4] bpf: Allow syscall bpf programs to call non-recur helpers

[Amery Hung]

Allow syscall programs to call non-recur helpers too since syscall bpf
programs runs in process context through bpf syscall, BPF_PROG_TEST_RUN,
and cannot run recursively.

bpf_task_storage_{get,set} have "_recur" versions that call trylock
instead of taking the lock directly to avoid deadlock when called by
bpf programs that run recursively. Currently, only bpf_lsm, bpf_iter,
struct_ops without private stack are allow to call the non-recur helpers
since they cannot be recursively called in another bpf program.

Signed-off-by: Amery Hung <ameryhung@gmail.com>
---
 include/linux/bpf_verifier.h | 1 +
 1 file changed, 1 insertion(+)

diff --git a/include/linux/bpf_verifier.h b/include/linux/bpf_verifier.h
index 94defa405c85..c823f8efe3ed 100644
--- a/include/linux/bpf_verifier.h
+++ b/include/linux/bpf_verifier.h
@@ -962,6 +962,7 @@ static inline bool bpf_prog_check_recur(const struct bpf_prog *prog)
 	case BPF_PROG_TYPE_STRUCT_OPS:
 		return prog->aux->jits_use_priv_stack;
 	case BPF_PROG_TYPE_LSM:
+	case BPF_PROG_TYPE_SYSCALL:
 		return false;
 	default:
 		return true;
-- 
2.47.3

[PATCH bpf-next v7 2/4] selftests/bpf: Introduce task local data

[Amery Hung]

Task local data defines an abstract storage type for storing task-
specific data (TLD). This patch provides user space and bpf
implementation as header-only libraries for accessing task local data.

Task local data is a bpf task local storage map with two UPTRs:

- tld_meta_u, shared by all tasks of a process, consists of the total
  count and size of TLDs and an array of metadata of TLDs. A TLD
  metadata contains the size and name. The name is used to identify a
  specific TLD in bpf programs.

- u_tld_data points to a task-specific memory. It stores TLD data and
  the starting offset of data in a page.

  Task local design decouple user space and bpf programs. Since bpf
  program does not know the size of TLDs in compile time, u_tld_data
  is declared as a page to accommodate TLDs up to a page. As a result,
  while user space will likely allocate memory smaller than a page for
  actual TLDs, it needs to pin a page to kernel. It will pin the page
  that contains enough memory if the allocated memory spans across the
  page boundary.

The library also creates another task local storage map, tld_key_map,
to cache keys for bpf programs to speed up the access.

Below are the core task local data API:

                   User space                          BPF
  Define TLD       TLD_DEFINE_KEY(), tld_create_key()  -
  Init TLD object  -                                   tld_object_init()
  Get TLD data     tld_get_data()                      tld_get_data()

- TLD_DEFINE_KEY(), tld_create_key()

  A TLD is first defined by the user space with TLD_DEFINE_KEY() or
  tld_create_key(). TLD_DEFINE_KEY() defines a TLD statically and
  allocates just enough memory during initialization. tld_create_key()
  allows creating TLDs on the fly, but has a fix memory budget,
  TLD_DYN_DATA_SIZE.

  Internally, they all call __tld_create_key(), which iterates
  tld_meta_u->metadata to check if a TLD can be added. The total TLD
  size needs to fit into a page (limit of UPTR), and no two TLDs can
  have the same name. If a TLD can be added, u_tld_meta->cnt is
  increased using cmpxchg as there may be other concurrent
  __tld_create_key(). After a successful cmpxchg, the last available
  tld_meta_u->metadata now belongs to the calling thread. To prevent
  other threads from reading incomplete metadata while it is being
  updated, tld_meta_u->metadata->size is used to signal the completion.

  Finally, the offset, derived from adding up prior TLD sizes is then
  encapsulated as an opaque object key to prevent user misuse. The
  offset is guaranteed to be 8-byte aligned to prevent load/store
  tearing and allow atomic operations on it.

- tld_get_data()

  User space programs can pass the key to tld_get_data() to get a
  pointer to the associated TLD. The pointer will remain valid for the
  lifetime of the thread.

  tld_data_u is lazily allocated on the first call to tld_get_data().
  Trying to read task local data from bpf will result in -ENODATA
  during tld_object_init(). The task-specific memory need to be freed
  manually by calling tld_free() on thread exit to prevent memory leak
  or use TLD_FREE_DATA_ON_THREAD_EXIT.

- tld_object_init() (BPF)

  BPF programs need to call tld_object_init() before calling
  tld_get_data(). This is to avoid redundant map lookup in
  tld_get_data() by storing pointers to the map values on stack.
  The pointers are encapsulated as tld_object.

  tld_key_map is also created on the first time tld_object_init()
  is called to cache TLD keys successfully fetched by tld_get_data().

  bpf_task_storage_get(.., F_CREATE) needs to be retried since it may
  fail when another thread has already taken the percpu counter lock
  for the task local storage.

- tld_get_data() (BPF)

  BPF programs can also get a pointer to a TLD with tld_get_data().
  It uses the cached key in tld_key_map to locate the data in
  tld_data_u->data. If the cached key is not set yet (<= 0),
  __tld_fetch_key() will be called to iterate tld_meta_u->metadata
  and find the TLD by name. To prevent redundant string comparison
  in the future when the search fail, the tld_meta_u->cnt is stored
  in the non-positive range of the key. Next time, __tld_fetch_key()
  will be called only if there are new TLDs and the search will start
  from the newly added tld_meta_u->metadata using the old
  tld_meta_u-cnt.

Signed-off-by: Amery Hung <ameryhung@gmail.com>
---
 .../bpf/prog_tests/task_local_data.h          | 386 ++++++++++++++++++
 .../selftests/bpf/progs/task_local_data.bpf.h | 237 +++++++++++
 2 files changed, 623 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/prog_tests/task_local_data.h
 create mode 100644 tools/testing/selftests/bpf/progs/task_local_data.bpf.h

diff --git a/tools/testing/selftests/bpf/prog_tests/task_local_data.h b/tools/testing/selftests/bpf/prog_tests/task_local_data.h
new file mode 100644
index 000000000000..a408d10c3688
--- /dev/null
+++ b/tools/testing/selftests/bpf/prog_tests/task_local_data.h
@@ -0,0 +1,386 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+#ifndef __TASK_LOCAL_DATA_H
+#define __TASK_LOCAL_DATA_H
+
+#include <errno.h>
+#include <fcntl.h>
+#include <sched.h>
+#include <stdatomic.h>
+#include <stddef.h>
+#include <stdlib.h>
+#include <string.h>
+#include <unistd.h>
+#include <sys/syscall.h>
+#include <sys/types.h>
+
+#ifdef TLD_FREE_DATA_ON_THREAD_EXIT
+#include <pthread.h>
+#endif
+
+#include <bpf/bpf.h>
+
+/*
+ * OPTIONS
+ *
+ *   Define the option before including the header
+ *
+ *   TLD_FREE_DATA_ON_THREAD_EXIT - Frees memory on thread exit automatically
+ *
+ *   Thread-specific memory for storing TLD is allocated lazily on the first call to
+ *   tld_get_data(). The thread that calls it must also call tld_free() on thread exit
+ *   to prevent memory leak. Pthread will be included if the option is defined. A pthread
+ *   key will be registered with a destructor that calls tld_free().
+ *
+ *
+ *   TLD_DYN_DATA_SIZE - The maximum size of memory allocated for TLDs created dynamically
+ *   (default: 64 bytes)
+ *
+ *   A TLD can be defined statically using TLD_DEFINE_KEY() or created on the fly using
+ *   tld_create_key(). As the total size of TLDs created with tld_create_key() cannot be
+ *   possibly known statically, a memory area of size TLD_DYN_DATA_SIZE will be allocated
+ *   for these TLDs. This additional memory is allocated for every thread that calls
+ *   tld_get_data() even if no tld_create_key are actually called, so be mindful of
+ *   potential memory wastage. Use TLD_DEFINE_KEY() whenever possible as just enough memory
+ *   will be allocated for TLDs created with it.
+ *
+ *
+ *   TLD_NAME_LEN - The maximum length of the name of a TLD (default: 62)
+ *
+ *   Setting TLD_NAME_LEN will affect the maximum number of TLDs a process can store,
+ *   TLD_MAX_DATA_CNT.
+ *
+ *
+ *   TLD_DATA_USE_ALIGNED_ALLOC - Always use aligned_alloc() instead of malloc()
+ *
+ *   When allocating the memory for storing TLDs, we need to make sure there is a memory
+ *   region of the X bytes within a page. This is due to the limit posed by UPTR: memory
+ *   pinned to the kernel cannot exceed a page nor can it cross the page boundary. The
+ *   library normally calls malloc(2*X) given X bytes of total TLDs, and only uses
+ *   aligned_alloc(PAGE_SIZE, X) when X >= PAGE_SIZE / 2. This is to reduce memory wastage
+ *   as not all memory allocator can use the exact amount of memory requested to fulfill
+ *   aligned_alloc(). For example, some may round the size up to the alignment. Enable the
+ *   option to always use aligned_alloc() if the implementation has low memory overhead.
+ */
+
+#define TLD_PAGE_SIZE getpagesize()
+#define TLD_PAGE_MASK (~(TLD_PAGE_SIZE - 1))
+
+#define TLD_ROUND_MASK(x, y) ((__typeof__(x))((y) - 1))
+#define TLD_ROUND_UP(x, y) ((((x) - 1) | TLD_ROUND_MASK(x, y)) + 1)
+
+#define TLD_READ_ONCE(x) (*(volatile typeof(x) *)&(x))
+
+#ifndef TLD_DYN_DATA_SIZE
+#define TLD_DYN_DATA_SIZE 64
+#endif
+
+#define TLD_MAX_DATA_CNT (TLD_PAGE_SIZE / sizeof(struct tld_metadata) - 1)
+
+#ifndef TLD_NAME_LEN
+#define TLD_NAME_LEN 62
+#endif
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+typedef struct {
+	__s16 off;
+} tld_key_t;
+
+struct tld_metadata {
+	char name[TLD_NAME_LEN];
+	_Atomic __u16 size;
+};
+
+struct tld_meta_u {
+	_Atomic __u8 cnt;
+	__u16 size;
+	struct tld_metadata metadata[];
+};
+
+struct tld_data_u {
+	__u64 start; /* offset of tld_data_u->data in a page */
+	char data[];
+};
+
+struct tld_map_value {
+	void *data;
+	struct tld_meta_u *meta;
+};
+
+struct tld_meta_u * _Atomic tld_meta_p __attribute__((weak));
+__thread struct tld_data_u *tld_data_p __attribute__((weak));
+__thread void *tld_data_alloc_p __attribute__((weak));
+
+#ifdef TLD_FREE_DATA_ON_THREAD_EXIT
+pthread_key_t tld_pthread_key __attribute__((weak));
+
+static void tld_free(void);
+
+static void __tld_thread_exit_handler(void *unused)
+{
+	tld_free();
+}
+#endif
+
+static int __tld_init_meta_p(void)
+{
+	struct tld_meta_u *meta, *uninit = NULL;
+	int err = 0;
+
+	meta = (struct tld_meta_u *)aligned_alloc(TLD_PAGE_SIZE, TLD_PAGE_SIZE);
+	if (!meta) {
+		err = -ENOMEM;
+		goto out;
+	}
+
+	memset(meta, 0, TLD_PAGE_SIZE);
+	meta->size = TLD_DYN_DATA_SIZE;
+
+	if (!atomic_compare_exchange_strong(&tld_meta_p, &uninit, meta)) {
+		free(meta);
+		goto out;
+	}
+
+#ifdef TLD_FREE_DATA_ON_THREAD_EXIT
+	pthread_key_create(&tld_pthread_key, __tld_thread_exit_handler);
+#endif
+out:
+	return err;
+}
+
+static int __tld_init_data_p(int map_fd)
+{
+	bool use_aligned_alloc = false;
+	struct tld_map_value map_val;
+	struct tld_data_u *data;
+	void *data_alloc = NULL;
+	int err, tid_fd = -1;
+
+	tid_fd = syscall(SYS_pidfd_open, gettid(), O_EXCL);
+	if (tid_fd < 0) {
+		err = -errno;
+		goto out;
+	}
+
+#ifdef TLD_DATA_USE_ALIGNED_ALLOC
+	use_aligned_alloc = true;
+#endif
+
+	/*
+	 * tld_meta_p->size = TLD_DYN_DATA_SIZE +
+	 *          total size of TLDs defined via TLD_DEFINE_KEY()
+	 */
+	data_alloc = (use_aligned_alloc || tld_meta_p->size * 2 >= TLD_PAGE_SIZE) ?
+		aligned_alloc(TLD_PAGE_SIZE, tld_meta_p->size) :
+		malloc(tld_meta_p->size * 2);
+	if (!data_alloc) {
+		err = -ENOMEM;
+		goto out;
+	}
+
+	/*
+	 * Always pass a page-aligned address to UPTR since the size of tld_map_value::data
+	 * is a page in BTF. If data_alloc spans across two pages, use the page that contains large
+	 * enough memory.
+	 */
+	if (TLD_PAGE_SIZE - (~TLD_PAGE_MASK & (intptr_t)data_alloc) >= tld_meta_p->size) {
+		map_val.data = (void *)(TLD_PAGE_MASK & (intptr_t)data_alloc);
+		data = data_alloc;
+		data->start = (~TLD_PAGE_MASK & (intptr_t)data_alloc) +
+			      offsetof(struct tld_data_u, data);
+	} else {
+		map_val.data = (void *)(TLD_ROUND_UP((intptr_t)data_alloc, TLD_PAGE_SIZE));
+		data = (void *)(TLD_ROUND_UP((intptr_t)data_alloc, TLD_PAGE_SIZE));
+		data->start = offsetof(struct tld_data_u, data);
+	}
+	map_val.meta = TLD_READ_ONCE(tld_meta_p);
+
+	err = bpf_map_update_elem(map_fd, &tid_fd, &map_val, 0);
+	if (err) {
+		free(data_alloc);
+		goto out;
+	}
+
+	tld_data_p = data;
+	tld_data_alloc_p = data_alloc;
+#ifdef TLD_FREE_DATA_ON_THREAD_EXIT
+	pthread_setspecific(tld_pthread_key, (void *)1);
+#endif
+out:
+	if (tid_fd >= 0)
+		close(tid_fd);
+	return err;
+}
+
+static tld_key_t __tld_create_key(const char *name, size_t size, bool dyn_data)
+{
+	int err, i, sz, off = 0;
+	__u8 cnt;
+
+	if (!TLD_READ_ONCE(tld_meta_p)) {
+		err = __tld_init_meta_p();
+		if (err)
+			return (tld_key_t){err};
+	}
+
+	for (i = 0; i < TLD_MAX_DATA_CNT; i++) {
+retry:
+		cnt = atomic_load(&tld_meta_p->cnt);
+		if (i < cnt) {
+			/* A metadata is not ready until size is updated with a non-zero value */
+			while (!(sz = atomic_load(&tld_meta_p->metadata[i].size)))
+				sched_yield();
+
+			if (!strncmp(tld_meta_p->metadata[i].name, name, TLD_NAME_LEN))
+				return (tld_key_t){-EEXIST};
+
+			off += TLD_ROUND_UP(sz, 8);
+			continue;
+		}
+
+		/*
+		 * TLD_DEFINE_KEY() is given memory upto a page while at most
+		 * TLD_DYN_DATA_SIZE is allocated for tld_create_key()
+		 */
+		if (dyn_data) {
+			if (off + TLD_ROUND_UP(size, 8) > tld_meta_p->size)
+				return (tld_key_t){-E2BIG};
+		} else {
+			if (off + TLD_ROUND_UP(size, 8) > TLD_PAGE_SIZE - sizeof(struct tld_data_u))
+				return (tld_key_t){-E2BIG};
+			tld_meta_p->size += TLD_ROUND_UP(size, 8);
+		}
+
+		/*
+		 * Only one tld_create_key() can increase the current cnt by one and
+		 * takes the latest available slot. Other threads will check again if a new
+		 * TLD can still be added, and then compete for the new slot after the
+		 * succeeding thread update the size.
+		 */
+		if (!atomic_compare_exchange_strong(&tld_meta_p->cnt, &cnt, cnt + 1))
+			goto retry;
+
+		strncpy(tld_meta_p->metadata[i].name, name, TLD_NAME_LEN);
+		atomic_store(&tld_meta_p->metadata[i].size, size);
+		return (tld_key_t){(__s16)off};
+	}
+
+	return (tld_key_t){-ENOSPC};
+}
+
+/**
+ * TLD_DEFINE_KEY() - Define a TLD and a global variable key associated with the TLD.
+ *
+ * @name: The name of the TLD
+ * @size: The size of the TLD
+ * @key: The variable name of the key. Cannot exceed TLD_NAME_LEN
+ *
+ * The macro can only be used in file scope.
+ *
+ * A global variable key of opaque type, tld_key_t, will be declared and initialized before
+ * main() starts. Use tld_key_is_err() or tld_key_err_or_zero() later to check if the key
+ * creation succeeded. Pass the key to tld_get_data() to get a pointer to the TLD.
+ * bpf programs can also fetch the same key by name.
+ *
+ * The total size of TLDs created using TLD_DEFINE_KEY() cannot exceed a page. Just
+ * enough memory will be allocated for each thread on the first call to tld_get_data().
+ */
+#define TLD_DEFINE_KEY(key, name, size)			\
+tld_key_t key;						\
+							\
+__attribute__((constructor))				\
+void __tld_define_key_##key(void)			\
+{							\
+	key = __tld_create_key(name, size, false);	\
+}
+
+/**
+ * tld_create_key() - Create a TLD and return a key associated with the TLD.
+ *
+ * @name: The name the TLD
+ * @size: The size of the TLD
+ *
+ * Return an opaque object key. Use tld_key_is_err() or tld_key_err_or_zero() to check
+ * if the key creation succeeded. Pass the key to tld_get_data() to get a pointer to
+ * locate the TLD. bpf programs can also fetch the same key by name.
+ *
+ * Use tld_create_key() only when a TLD needs to be created dynamically (e.g., @name is
+ * not known statically or a TLD needs to be created conditionally)
+ *
+ * An additional TLD_DYN_DATA_SIZE bytes are allocated per-thread to accommodate TLDs
+ * created dynamically with tld_create_key(). Since only a user page is pinned to the
+ * kernel, when TLDs created with TLD_DEFINE_KEY() uses more than TLD_PAGE_SIZE -
+ * TLD_DYN_DATA_SIZE, the buffer size will be limited to the rest of the page.
+ */
+__attribute__((unused))
+static tld_key_t tld_create_key(const char *name, size_t size)
+{
+	return __tld_create_key(name, size, true);
+}
+
+__attribute__((unused))
+static inline bool tld_key_is_err(tld_key_t key)
+{
+	return key.off < 0;
+}
+
+__attribute__((unused))
+static inline int tld_key_err_or_zero(tld_key_t key)
+{
+	return tld_key_is_err(key) ? key.off : 0;
+}
+
+/**
+ * tld_get_data() - Get a pointer to the TLD associated with the given key of the
+ * calling thread.
+ *
+ * @map_fd: A file descriptor of tld_data_map, the underlying BPF task local storage map
+ * of task local data.
+ * @key: A key object created by TLD_DEFINE_KEY() or tld_create_key().
+ *
+ * Return a pointer to the TLD if the key is valid; NULL if not enough memory for TLD
+ * for this thread, or the key is invalid. The returned pointer is guaranteed to be 8-byte
+ * aligned.
+ *
+ * Threads that call tld_get_data() must call tld_free() on exit to prevent
+ * memory leak if TLD_FREE_DATA_ON_THREAD_EXIT is not defined.
+ */
+__attribute__((unused))
+static void *tld_get_data(int map_fd, tld_key_t key)
+{
+	if (!TLD_READ_ONCE(tld_meta_p))
+		return NULL;
+
+	/* tld_data_p is allocated on the first invocation of tld_get_data() */
+	if (!tld_data_p && __tld_init_data_p(map_fd))
+		return NULL;
+
+	return tld_data_p->data + key.off;
+}
+
+/**
+ * tld_free() - Free task local data memory of the calling thread
+ *
+ * For the calling thread, all pointers to TLDs acquired before will become invalid.
+ *
+ * Users must call tld_free() on thread exit to prevent memory leak. Alternatively,
+ * define TLD_FREE_DATA_ON_THREAD_EXIT and a thread exit handler will be registered
+ * to free the memory automatically.
+ */
+__attribute__((unused))
+static void tld_free(void)
+{
+	if (tld_data_alloc_p) {
+		free(tld_data_alloc_p);
+		tld_data_alloc_p = NULL;
+		tld_data_p = NULL;
+	}
+}
+
+#ifdef __cplusplus
+} /* extern "C" */
+#endif
+
+#endif /* __TASK_LOCAL_DATA_H */
diff --git a/tools/testing/selftests/bpf/progs/task_local_data.bpf.h b/tools/testing/selftests/bpf/progs/task_local_data.bpf.h
new file mode 100644
index 000000000000..432fff2af844
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/task_local_data.bpf.h
@@ -0,0 +1,237 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+#ifndef __TASK_LOCAL_DATA_BPF_H
+#define __TASK_LOCAL_DATA_BPF_H
+
+/*
+ * Task local data is a library that facilitates sharing per-task data
+ * between user space and bpf programs.
+ *
+ *
+ * USAGE
+ *
+ * A TLD, an entry of data in task local data, first needs to be created by the
+ * user space. This is done by calling user space API, TLD_DEFINE_KEY() or
+ * tld_create_key(), with the name of the TLD and the size.
+ *
+ * TLD_DEFINE_KEY(prio, "priority", sizeof(int));
+ *
+ * or
+ *
+ * void func_call(...) {
+ *     tld_key_t prio, in_cs;
+ *
+ *     prio = tld_create_key("priority", sizeof(int));
+ *     in_cs = tld_create_key("in_critical_section", sizeof(bool));
+ *     ...
+ *
+ * A key associated with the TLD, which has an opaque type tld_key_t, will be
+ * initialized or returned. It can be used to get a pointer to the TLD in the
+ * user space by calling tld_get_data().
+ *
+ * In a bpf program, tld_object_init() first needs to be called to initialized a
+ * tld_object on the stack. Then, TLDs can be accessed by calling tld_get_data().
+ * The API will try to fetch the key by the name and use it to locate the data.
+ * A pointer to the TLD will be returned. It also caches the key in a task local
+ * storage map, tld_key_map, whose value type, struct tld_keys, must be defined
+ * by the developer.
+ *
+ * struct tld_keys {
+ *     tld_key_t prio;
+ *     tld_key_t in_cs;
+ * };
+ *
+ * SEC("struct_ops")
+ * void prog(struct task_struct task, ...)
+ * {
+ *     struct tld_object tld_obj;
+ *     int err, *p;
+ *
+ *     err = tld_object_init(task, &tld_obj);
+ *     if (err)
+ *         return;
+ *
+ *     p = tld_get_data(&tld_obj, prio, "priority", sizeof(int));
+ *     if (p)
+ *         // do something depending on *p
+ */
+#include <errno.h>
+#include <bpf/bpf_helpers.h>
+
+#define TLD_ROUND_MASK(x, y) ((__typeof__(x))((y) - 1))
+#define TLD_ROUND_UP(x, y) ((((x) - 1) | TLD_ROUND_MASK(x, y)) + 1)
+
+#define TLD_MAX_DATA_CNT (__PAGE_SIZE / sizeof(struct tld_metadata) - 1)
+
+#ifndef TLD_NAME_LEN
+#define TLD_NAME_LEN 62
+#endif
+
+#ifndef TLD_KEY_MAP_CREATE_RETRY
+#define TLD_KEY_MAP_CREATE_RETRY 10
+#endif
+
+typedef struct {
+	__s16 off;
+} tld_key_t;
+
+struct tld_metadata {
+	char name[TLD_NAME_LEN];
+	__u16 size;
+};
+
+struct tld_meta_u {
+	__u8 cnt;
+	__u16 size;
+	struct tld_metadata metadata[TLD_MAX_DATA_CNT];
+};
+
+struct tld_data_u {
+	__u64 start; /* offset of tld_data_u->data in a page */
+	char data[__PAGE_SIZE - sizeof(__u64)];
+};
+
+struct tld_map_value {
+	struct tld_data_u __uptr *data;
+	struct tld_meta_u __uptr *meta;
+};
+
+typedef struct tld_uptr_dummy {
+	struct tld_data_u data[0];
+	struct tld_meta_u meta[0];
+} *tld_uptr_dummy_t;
+
+struct tld_object {
+	struct tld_map_value *data_map;
+	struct tld_keys *key_map;
+	/*
+	 * Force the compiler to generate the actual definition of tld_meta_u
+	 * and tld_data_u in BTF. Without it, tld_meta_u and u_tld_data will
+	 * be BTF_KIND_FWD.
+	 */
+	tld_uptr_dummy_t dummy[0];
+};
+
+/*
+ * Map value of tld_key_map for caching keys. Must be defined by the developer.
+ * Members should be tld_key_t and passed to the 3rd argument of tld_fetch_key().
+ */
+struct tld_keys;
+
+struct {
+	__uint(type, BPF_MAP_TYPE_TASK_STORAGE);
+	__uint(map_flags, BPF_F_NO_PREALLOC);
+	__type(key, int);
+	__type(value, struct tld_map_value);
+} tld_data_map SEC(".maps");
+
+struct {
+	__uint(type, BPF_MAP_TYPE_TASK_STORAGE);
+	__uint(map_flags, BPF_F_NO_PREALLOC);
+	__type(key, int);
+	__type(value, struct tld_keys);
+} tld_key_map SEC(".maps");
+
+/**
+ * tld_object_init() - Initialize a tld_object.
+ *
+ * @task: The task_struct of the target task
+ * @tld_obj: A pointer to a tld_object to be initialized
+ *
+ * Return 0 on success; -ENODATA if the user space did not initialize task local data
+ * for the current task through tld_get_data(); -ENOMEM if the creation of tld_key_map
+ * fails
+ */
+__attribute__((unused))
+static int tld_object_init(struct task_struct *task, struct tld_object *tld_obj)
+{
+	int i;
+
+	tld_obj->data_map = bpf_task_storage_get(&tld_data_map, task, 0, 0);
+	if (!tld_obj->data_map)
+		return -ENODATA;
+
+	bpf_for(i, 0, TLD_KEY_MAP_CREATE_RETRY) {
+		tld_obj->key_map = bpf_task_storage_get(&tld_key_map, task, 0,
+							BPF_LOCAL_STORAGE_GET_F_CREATE);
+		if (likely(tld_obj->key_map))
+			break;
+	}
+	if (!tld_obj->key_map)
+		return -ENOMEM;
+
+	return 0;
+}
+
+/*
+ * Return the offset of TLD if @name is found. Otherwise, return the current TLD count
+ * using the nonpositive range so that the next tld_get_data() can skip fetching key if
+ * no new TLD is added or start comparing name from the first newly added TLD.
+ */
+__attribute__((unused))
+static int __tld_fetch_key(struct tld_object *tld_obj, const char *name, int i_start)
+{
+	struct tld_metadata *metadata;
+	int i, cnt, start, off = 0;
+
+	if (!tld_obj->data_map || !tld_obj->data_map->data || !tld_obj->data_map->meta)
+		return 0;
+
+	start = tld_obj->data_map->data->start;
+	cnt = tld_obj->data_map->meta->cnt;
+	metadata = tld_obj->data_map->meta->metadata;
+
+	bpf_for(i, 0, cnt) {
+		if (i >= TLD_MAX_DATA_CNT)
+			break;
+
+		if (i >= i_start && !bpf_strncmp(metadata[i].name, TLD_NAME_LEN, name))
+			return start + off;
+
+		off += TLD_ROUND_UP(metadata[i].size, 8);
+	}
+
+	return -cnt;
+}
+
+/**
+ * tld_get_data() - Retrieve a pointer to the TLD associated with the name.
+ *
+ * @tld_obj: A pointer to a valid tld_object initialized by tld_object_init()
+ * @key: The cached key of the TLD in tld_key_map
+ * @name: The name of the key associated with a TLD
+ * @size: The size of the TLD. Must be a known constant value
+ *
+ * Return a pointer to the TLD associated with @name; NULL if not found or @size is too
+ * big. @key is used to cache the key if the TLD is found to speed up subsequent calls.
+ * It should be defined as an member of tld_keys of tld_key_t type by the developer.
+ */
+#define tld_get_data(tld_obj, key, name, size)						\
+	({										\
+		void *data = NULL, *_data = (tld_obj)->data_map->data;			\
+		long off = (tld_obj)->key_map->key.off;					\
+		int cnt;								\
+											\
+		if (likely(_data)) {							\
+			if (likely(off > 0)) {						\
+				barrier_var(off);					\
+				if (likely(off < __PAGE_SIZE - size))			\
+					data = _data + off;				\
+			} else {							\
+				cnt = -(off);						\
+				if (likely((tld_obj)->data_map->meta) &&		\
+				    cnt < (tld_obj)->data_map->meta->cnt) {		\
+					off = __tld_fetch_key(tld_obj, name, cnt);	\
+					(tld_obj)->key_map->key.off = off;		\
+											\
+					if (likely(off < __PAGE_SIZE - size)) {		\
+						barrier_var(off);			\
+						if (off > 0)				\
+							data = _data + off;		\
+					}						\
+				}							\
+			}								\
+		}									\
+		data;									\
+	})
+
+#endif
-- 
2.47.3

[PATCH bpf-next v7 3/4] selftests/bpf: Test basic task local data operations

[Amery Hung]

Test basic operations of task local data with valid and invalid
tld_create_key().

For invalid calls, make sure they return the right error code and check
that the TLDs are not inserted by running tld_get_data("
value_not_exists") on the bpf side. The call should a null pointer.

For valid calls, first make sure the TLDs are created by calling
tld_get_data() on the bpf side. The call should return a valid pointer.

Finally, verify that the TLDs are indeed task-specific (i.e., their
addresses do not overlap) with multiple user threads. This done by
writing values unique to each thread, reading them from both user space
and bpf, and checking if the value read back matches the value written.

Signed-off-by: Amery Hung <ameryhung@gmail.com>
---
 .../bpf/prog_tests/test_task_local_data.c     | 192 ++++++++++++++++++
 .../bpf/progs/test_task_local_data.c          |  65 ++++++
 2 files changed, 257 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/prog_tests/test_task_local_data.c
 create mode 100644 tools/testing/selftests/bpf/progs/test_task_local_data.c

diff --git a/tools/testing/selftests/bpf/prog_tests/test_task_local_data.c b/tools/testing/selftests/bpf/prog_tests/test_task_local_data.c
new file mode 100644
index 000000000000..2e77d3fa2534
--- /dev/null
+++ b/tools/testing/selftests/bpf/prog_tests/test_task_local_data.c
@@ -0,0 +1,192 @@
+// SPDX-License-Identifier: GPL-2.0
+#include <pthread.h>
+#include <bpf/btf.h>
+#include <test_progs.h>
+
+#define TLD_FREE_DATA_ON_THREAD_EXIT
+#define TLD_DYN_DATA_SIZE 4096
+#include "task_local_data.h"
+
+struct test_tld_struct {
+	__u64 a;
+	__u64 b;
+	__u64 c;
+	__u64 d;
+};
+
+#include "test_task_local_data.skel.h"
+
+TLD_DEFINE_KEY(value0_key, "value0", sizeof(int));
+
+/*
+ * Reset task local data between subtests by clearing metadata other
+ * than the statically defined value0. This is safe as subtests run
+ * sequentially. Users of task local data library should not touch
+ * library internal.
+ */
+static void reset_tld(void)
+{
+	if (TLD_READ_ONCE(tld_meta_p)) {
+		/* Remove TLDs created by tld_create_key() */
+		tld_meta_p->cnt = 1;
+		tld_meta_p->size = TLD_DYN_DATA_SIZE;
+		memset(&tld_meta_p->metadata[1], 0,
+		       (TLD_MAX_DATA_CNT - 1) * sizeof(struct tld_metadata));
+	}
+}
+
+/* Serialize access to bpf program's global variables */
+static pthread_mutex_t global_mutex;
+
+static tld_key_t *tld_keys;
+
+#define TEST_BASIC_THREAD_NUM 32
+
+void *test_task_local_data_basic_thread(void *arg)
+{
+	LIBBPF_OPTS(bpf_test_run_opts, opts);
+	struct test_task_local_data *skel = (struct test_task_local_data *)arg;
+	int fd, err, tid, *value0, *value1;
+	struct test_tld_struct *value2;
+
+	fd = bpf_map__fd(skel->maps.tld_data_map);
+
+	value0 = tld_get_data(fd, value0_key);
+	if (!ASSERT_OK_PTR(value0, "tld_get_data"))
+		goto out;
+
+	value1 = tld_get_data(fd, tld_keys[1]);
+	if (!ASSERT_OK_PTR(value1, "tld_get_data"))
+		goto out;
+
+	value2 = tld_get_data(fd, tld_keys[2]);
+	if (!ASSERT_OK_PTR(value2, "tld_get_data"))
+		goto out;
+
+	tid = gettid();
+
+	*value0 = tid + 0;
+	*value1 = tid + 1;
+	value2->a = tid + 2;
+	value2->b = tid + 3;
+	value2->c = tid + 4;
+	value2->d = tid + 5;
+
+	pthread_mutex_lock(&global_mutex);
+	/* Run task_main that read task local data and save to global variables */
+	err = bpf_prog_test_run_opts(bpf_program__fd(skel->progs.task_main), &opts);
+	ASSERT_OK(err, "run task_main");
+	ASSERT_OK(opts.retval, "task_main retval");
+
+	ASSERT_EQ(skel->bss->test_value0, tid + 0, "tld_get_data value0");
+	ASSERT_EQ(skel->bss->test_value1, tid + 1, "tld_get_data value1");
+	ASSERT_EQ(skel->bss->test_value2.a, tid + 2, "tld_get_data value2.a");
+	ASSERT_EQ(skel->bss->test_value2.b, tid + 3, "tld_get_data value2.b");
+	ASSERT_EQ(skel->bss->test_value2.c, tid + 4, "tld_get_data value2.c");
+	ASSERT_EQ(skel->bss->test_value2.d, tid + 5, "tld_get_data value2.d");
+	pthread_mutex_unlock(&global_mutex);
+
+	/* Make sure valueX are indeed local to threads */
+	ASSERT_EQ(*value0, tid + 0, "value0");
+	ASSERT_EQ(*value1, tid + 1, "value1");
+	ASSERT_EQ(value2->a, tid + 2, "value2.a");
+	ASSERT_EQ(value2->b, tid + 3, "value2.b");
+	ASSERT_EQ(value2->c, tid + 4, "value2.c");
+	ASSERT_EQ(value2->d, tid + 5, "value2.d");
+
+	*value0 = tid + 5;
+	*value1 = tid + 4;
+	value2->a = tid + 3;
+	value2->b = tid + 2;
+	value2->c = tid + 1;
+	value2->d = tid + 0;
+
+	/* Run task_main again */
+	pthread_mutex_lock(&global_mutex);
+	err = bpf_prog_test_run_opts(bpf_program__fd(skel->progs.task_main), &opts);
+	ASSERT_OK(err, "run task_main");
+	ASSERT_OK(opts.retval, "task_main retval");
+
+	ASSERT_EQ(skel->bss->test_value0, tid + 5, "tld_get_data value0");
+	ASSERT_EQ(skel->bss->test_value1, tid + 4, "tld_get_data value1");
+	ASSERT_EQ(skel->bss->test_value2.a, tid + 3, "tld_get_data value2.a");
+	ASSERT_EQ(skel->bss->test_value2.b, tid + 2, "tld_get_data value2.b");
+	ASSERT_EQ(skel->bss->test_value2.c, tid + 1, "tld_get_data value2.c");
+	ASSERT_EQ(skel->bss->test_value2.d, tid + 0, "tld_get_data value2.d");
+	pthread_mutex_unlock(&global_mutex);
+
+out:
+	pthread_exit(NULL);
+}
+
+static void test_task_local_data_basic(void)
+{
+	struct test_task_local_data *skel;
+	pthread_t thread[TEST_BASIC_THREAD_NUM];
+	char dummy_key_name[TLD_NAME_LEN];
+	tld_key_t key;
+	int i, err;
+
+	reset_tld();
+
+	ASSERT_OK(pthread_mutex_init(&global_mutex, NULL), "pthread_mutex_init");
+
+	skel = test_task_local_data__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "skel_open_and_load"))
+		return;
+
+	tld_keys = calloc(TLD_MAX_DATA_CNT, sizeof(tld_key_t));
+	if (!ASSERT_OK_PTR(tld_keys, "calloc tld_keys"))
+		goto out;
+
+	ASSERT_FALSE(tld_key_is_err(value0_key), "TLD_DEFINE_KEY");
+	tld_keys[1] = tld_create_key("value1", sizeof(int));
+	ASSERT_FALSE(tld_key_is_err(tld_keys[1]), "tld_create_key");
+	tld_keys[2] = tld_create_key("value2", sizeof(struct test_tld_struct));
+	ASSERT_FALSE(tld_key_is_err(tld_keys[2]), "tld_create_key");
+
+	/*
+	 * Shouldn't be able to store data exceed a page. Create a TLD just big
+	 * enough to exceed a page. TLDs already created are int value0, int
+	 * value1, and struct test_tld_struct value2.
+	 */
+	key = tld_create_key("value_not_exist",
+			     TLD_PAGE_SIZE - 2 * sizeof(int) - sizeof(struct test_tld_struct) + 1);
+	ASSERT_EQ(tld_key_err_or_zero(key), -E2BIG, "tld_create_key");
+
+	key = tld_create_key("value2", sizeof(struct test_tld_struct));
+	ASSERT_EQ(tld_key_err_or_zero(key), -EEXIST, "tld_create_key");
+
+	/* Shouldn't be able to create the (TLD_MAX_DATA_CNT+1)-th TLD */
+	for (i = 3; i < TLD_MAX_DATA_CNT; i++) {
+		snprintf(dummy_key_name, TLD_NAME_LEN, "dummy_value%d", i);
+		tld_keys[i] = tld_create_key(dummy_key_name, sizeof(int));
+		ASSERT_FALSE(tld_key_is_err(tld_keys[i]), "tld_create_key");
+	}
+	key = tld_create_key("value_not_exist", sizeof(struct test_tld_struct));
+	ASSERT_EQ(tld_key_err_or_zero(key), -ENOSPC, "tld_create_key");
+
+	/* Access TLDs from multiple threads and check if they are thread-specific */
+	for (i = 0; i < TEST_BASIC_THREAD_NUM; i++) {
+		err = pthread_create(&thread[i], NULL, test_task_local_data_basic_thread, skel);
+		if (!ASSERT_OK(err, "pthread_create"))
+			goto out;
+	}
+
+out:
+	for (i = 0; i < TEST_BASIC_THREAD_NUM; i++)
+		pthread_join(thread[i], NULL);
+
+	if (tld_keys) {
+		free(tld_keys);
+		tld_keys = NULL;
+	}
+	tld_free();
+	test_task_local_data__destroy(skel);
+}
+
+void test_task_local_data(void)
+{
+	if (test__start_subtest("task_local_data_basic"))
+		test_task_local_data_basic();
+}
diff --git a/tools/testing/selftests/bpf/progs/test_task_local_data.c b/tools/testing/selftests/bpf/progs/test_task_local_data.c
new file mode 100644
index 000000000000..fffafc013044
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/test_task_local_data.c
@@ -0,0 +1,65 @@
+// SPDX-License-Identifier: GPL-2.0
+
+#include <vmlinux.h>
+#include <errno.h>
+#include <bpf/bpf_helpers.h>
+
+#include "task_local_data.bpf.h"
+
+struct tld_keys {
+	tld_key_t value0;
+	tld_key_t value1;
+	tld_key_t value2;
+	tld_key_t value_not_exist;
+};
+
+struct test_tld_struct {
+	__u64 a;
+	__u64 b;
+	__u64 c;
+	__u64 d;
+};
+
+int test_value0;
+int test_value1;
+struct test_tld_struct test_value2;
+
+SEC("syscall")
+int task_main(void *ctx)
+{
+	struct tld_object tld_obj;
+	struct test_tld_struct *struct_p;
+	struct task_struct *task;
+	int err, *int_p;
+
+	task = bpf_get_current_task_btf();
+	err = tld_object_init(task, &tld_obj);
+	if (err)
+		return 1;
+
+	int_p = tld_get_data(&tld_obj, value0, "value0", sizeof(int));
+	if (int_p)
+		test_value0 = *int_p;
+	else
+		return 2;
+
+	int_p = tld_get_data(&tld_obj, value1, "value1", sizeof(int));
+	if (int_p)
+		test_value1 = *int_p;
+	else
+		return 3;
+
+	struct_p = tld_get_data(&tld_obj, value2, "value2", sizeof(struct test_tld_struct));
+	if (struct_p)
+		test_value2 = *struct_p;
+	else
+		return 4;
+
+	int_p = tld_get_data(&tld_obj, value_not_exist, "value_not_exist", sizeof(int));
+	if (int_p)
+		return 5;
+
+	return 0;
+}
+
+char _license[] SEC("license") = "GPL";
-- 
2.47.3

[PATCH bpf-next v7 4/4] selftests/bpf: Test concurrent task local data key creation

[Amery Hung]

Test thread-safety of tld_create_key(). Since tld_create_key() does
not rely on locks but memory barriers and atomic operations to protect
the shared metadata, the thread-safety of the function is non-trivial.
Make sure concurrent tld_key_create(), both valid and invalid, can not
race and corrupt metatada, which may leads to TLDs not being thread-
specific or duplicate TLDs with the same name.

Signed-off-by: Amery Hung <ameryhung@gmail.com>
---
 .../bpf/prog_tests/test_task_local_data.c     | 105 ++++++++++++++++++
 1 file changed, 105 insertions(+)

diff --git a/tools/testing/selftests/bpf/prog_tests/test_task_local_data.c b/tools/testing/selftests/bpf/prog_tests/test_task_local_data.c
index 2e77d3fa2534..3b5cd2cd89c7 100644
--- a/tools/testing/selftests/bpf/prog_tests/test_task_local_data.c
+++ b/tools/testing/selftests/bpf/prog_tests/test_task_local_data.c
@@ -185,8 +185,113 @@ static void test_task_local_data_basic(void)
 	test_task_local_data__destroy(skel);
 }
 
+#define TEST_RACE_THREAD_NUM (TLD_MAX_DATA_CNT - 3)
+
+void *test_task_local_data_race_thread(void *arg)
+{
+	int err = 0, id = (intptr_t)arg;
+	char key_name[32];
+	tld_key_t key;
+
+	key = tld_create_key("value_not_exist", TLD_PAGE_SIZE + 1);
+	if (tld_key_err_or_zero(key) != -E2BIG) {
+		err = 1;
+		goto out;
+	}
+
+	/* Only one thread will succeed in creating value1 */
+	key = tld_create_key("value1", sizeof(int));
+	if (!tld_key_is_err(key))
+		tld_keys[1] = key;
+
+	/* Only one thread will succeed in creating value2 */
+	key = tld_create_key("value2", sizeof(struct test_tld_struct));
+	if (!tld_key_is_err(key))
+		tld_keys[2] = key;
+
+	snprintf(key_name, 32, "thread_%d", id);
+	tld_keys[id] = tld_create_key(key_name, sizeof(int));
+	if (tld_key_is_err(tld_keys[id]))
+		err = 2;
+out:
+	return (void *)(intptr_t)err;
+}
+
+static void test_task_local_data_race(void)
+{
+	LIBBPF_OPTS(bpf_test_run_opts, opts);
+	pthread_t thread[TEST_RACE_THREAD_NUM];
+	struct test_task_local_data *skel;
+	int fd, i, j, err, *data;
+	void *ret = NULL;
+
+	skel = test_task_local_data__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "skel_open_and_load"))
+		return;
+
+	tld_keys = calloc(TLD_MAX_DATA_CNT, sizeof(tld_key_t));
+	if (!ASSERT_OK_PTR(tld_keys, "calloc tld_keys"))
+		goto out;
+
+	fd = bpf_map__fd(skel->maps.tld_data_map);
+
+	ASSERT_FALSE(tld_key_is_err(value0_key), "TLD_DEFINE_KEY");
+	tld_keys[0] = value0_key;
+
+	for (j = 0; j < 100; j++) {
+		reset_tld();
+
+		for (i = 0; i < TEST_RACE_THREAD_NUM; i++) {
+			/*
+			 * Try to make tld_create_key() race with each other. Call
+			 * tld_create_key(), both valid and invalid, from different threads.
+			 */
+			err = pthread_create(&thread[i], NULL, test_task_local_data_race_thread,
+					     (void *)(intptr_t)(i + 3));
+			if (CHECK_FAIL(err))
+				break;
+		}
+
+		/* Wait for all tld_create_key() to return */
+		for (i = 0; i < TEST_RACE_THREAD_NUM; i++) {
+			pthread_join(thread[i], &ret);
+			if (CHECK_FAIL(ret))
+				break;
+		}
+
+		/* Write a unique number to each TLD */
+		for (i = 0; i < TLD_MAX_DATA_CNT; i++) {
+			data = tld_get_data(fd, tld_keys[i]);
+			if (CHECK_FAIL(!data))
+				break;
+			*data = i;
+		}
+
+		/* Read TLDs and check the value to see if any address collides with another */
+		for (i = 0; i < TLD_MAX_DATA_CNT; i++) {
+			data = tld_get_data(fd, tld_keys[i]);
+			if (CHECK_FAIL(*data != i))
+				break;
+		}
+
+		/* Run task_main to make sure no invalid TLDs are added */
+		err = bpf_prog_test_run_opts(bpf_program__fd(skel->progs.task_main), &opts);
+		ASSERT_OK(err, "run task_main");
+		ASSERT_OK(opts.retval, "task_main retval");
+	}
+out:
+	if (tld_keys) {
+		free(tld_keys);
+		tld_keys = NULL;
+	}
+	tld_free();
+	test_task_local_data__destroy(skel);
+}
+
 void test_task_local_data(void)
 {
 	if (test__start_subtest("task_local_data_basic"))
 		test_task_local_data_basic();
+	if (test__start_subtest("task_local_data_race"))
+		test_task_local_data_race();
 }
-- 
2.47.3

Re: [PATCH bpf-next v7 1/4] bpf: Allow syscall bpf programs to call non-recur helpers

[Emil Tsalapatis]

On Wed, Jul 30, 2025 at 2:59 PM Amery Hung <ameryhung@gmail.com> wrote:
>
> Allow syscall programs to call non-recur helpers too since syscall bpf
> programs runs in process context through bpf syscall, BPF_PROG_TEST_RUN,
> and cannot run recursively.
>
> bpf_task_storage_{get,set} have "_recur" versions that call trylock
> instead of taking the lock directly to avoid deadlock when called by
> bpf programs that run recursively. Currently, only bpf_lsm, bpf_iter,
> struct_ops without private stack are allow to call the non-recur helpers
> since they cannot be recursively called in another bpf program.
>
> Signed-off-by: Amery Hung <ameryhung@gmail.com>

Re-adding the tags in the case this is the final version, as we
discussed off-list.

Reviewed-by: Emil Tsalapatis <emil@etsalapatis.com>

> ---
>  include/linux/bpf_verifier.h | 1 +
>  1 file changed, 1 insertion(+)
>
> diff --git a/include/linux/bpf_verifier.h b/include/linux/bpf_verifier.h
> index 94defa405c85..c823f8efe3ed 100644
> --- a/include/linux/bpf_verifier.h
> +++ b/include/linux/bpf_verifier.h
> @@ -962,6 +962,7 @@ static inline bool bpf_prog_check_recur(const struct bpf_prog *prog)
>         case BPF_PROG_TYPE_STRUCT_OPS:
>                 return prog->aux->jits_use_priv_stack;
>         case BPF_PROG_TYPE_LSM:
> +       case BPF_PROG_TYPE_SYSCALL:
>                 return false;
>         default:
>                 return true;
> --
> 2.47.3
>

Re: [PATCH bpf-next v7 2/4] selftests/bpf: Introduce task local data

[Emil Tsalapatis]

On Wed, Jul 30, 2025 at 2:59 PM Amery Hung <ameryhung@gmail.com> wrote:
>
> Task local data defines an abstract storage type for storing task-
> specific data (TLD). This patch provides user space and bpf
> implementation as header-only libraries for accessing task local data.
>
> Task local data is a bpf task local storage map with two UPTRs:
>
> - tld_meta_u, shared by all tasks of a process, consists of the total
>   count and size of TLDs and an array of metadata of TLDs. A TLD
>   metadata contains the size and name. The name is used to identify a
>   specific TLD in bpf programs.
>
> - u_tld_data points to a task-specific memory. It stores TLD data and
>   the starting offset of data in a page.
>
>   Task local design decouple user space and bpf programs. Since bpf
>   program does not know the size of TLDs in compile time, u_tld_data
>   is declared as a page to accommodate TLDs up to a page. As a result,
>   while user space will likely allocate memory smaller than a page for
>   actual TLDs, it needs to pin a page to kernel. It will pin the page
>   that contains enough memory if the allocated memory spans across the
>   page boundary.
>
> The library also creates another task local storage map, tld_key_map,
> to cache keys for bpf programs to speed up the access.
>
> Below are the core task local data API:
>
>                    User space                          BPF
>   Define TLD       TLD_DEFINE_KEY(), tld_create_key()  -
>   Init TLD object  -                                   tld_object_init()
>   Get TLD data     tld_get_data()                      tld_get_data()
>
> - TLD_DEFINE_KEY(), tld_create_key()
>
>   A TLD is first defined by the user space with TLD_DEFINE_KEY() or
>   tld_create_key(). TLD_DEFINE_KEY() defines a TLD statically and
>   allocates just enough memory during initialization. tld_create_key()
>   allows creating TLDs on the fly, but has a fix memory budget,
>   TLD_DYN_DATA_SIZE.
>
>   Internally, they all call __tld_create_key(), which iterates
>   tld_meta_u->metadata to check if a TLD can be added. The total TLD
>   size needs to fit into a page (limit of UPTR), and no two TLDs can
>   have the same name. If a TLD can be added, u_tld_meta->cnt is
>   increased using cmpxchg as there may be other concurrent
>   __tld_create_key(). After a successful cmpxchg, the last available
>   tld_meta_u->metadata now belongs to the calling thread. To prevent
>   other threads from reading incomplete metadata while it is being
>   updated, tld_meta_u->metadata->size is used to signal the completion.
>
>   Finally, the offset, derived from adding up prior TLD sizes is then
>   encapsulated as an opaque object key to prevent user misuse. The
>   offset is guaranteed to be 8-byte aligned to prevent load/store
>   tearing and allow atomic operations on it.
>
> - tld_get_data()
>
>   User space programs can pass the key to tld_get_data() to get a
>   pointer to the associated TLD. The pointer will remain valid for the
>   lifetime of the thread.
>
>   tld_data_u is lazily allocated on the first call to tld_get_data().
>   Trying to read task local data from bpf will result in -ENODATA
>   during tld_object_init(). The task-specific memory need to be freed
>   manually by calling tld_free() on thread exit to prevent memory leak
>   or use TLD_FREE_DATA_ON_THREAD_EXIT.
>
> - tld_object_init() (BPF)
>
>   BPF programs need to call tld_object_init() before calling
>   tld_get_data(). This is to avoid redundant map lookup in
>   tld_get_data() by storing pointers to the map values on stack.
>   The pointers are encapsulated as tld_object.
>
>   tld_key_map is also created on the first time tld_object_init()
>   is called to cache TLD keys successfully fetched by tld_get_data().
>
>   bpf_task_storage_get(.., F_CREATE) needs to be retried since it may
>   fail when another thread has already taken the percpu counter lock
>   for the task local storage.
>
> - tld_get_data() (BPF)
>
>   BPF programs can also get a pointer to a TLD with tld_get_data().
>   It uses the cached key in tld_key_map to locate the data in
>   tld_data_u->data. If the cached key is not set yet (<= 0),
>   __tld_fetch_key() will be called to iterate tld_meta_u->metadata
>   and find the TLD by name. To prevent redundant string comparison
>   in the future when the search fail, the tld_meta_u->cnt is stored
>   in the non-positive range of the key. Next time, __tld_fetch_key()
>   will be called only if there are new TLDs and the search will start
>   from the newly added tld_meta_u->metadata using the old
>   tld_meta_u-cnt.
>
> Signed-off-by: Amery Hung <ameryhung@gmail.com>

Reviewed-by: Emil Tsalapatis <emil@etsalapatis.com>

> ---
>  .../bpf/prog_tests/task_local_data.h          | 386 ++++++++++++++++++
>  .../selftests/bpf/progs/task_local_data.bpf.h | 237 +++++++++++
>  2 files changed, 623 insertions(+)
>  create mode 100644 tools/testing/selftests/bpf/prog_tests/task_local_data.h
>  create mode 100644 tools/testing/selftests/bpf/progs/task_local_data.bpf.h
>
> diff --git a/tools/testing/selftests/bpf/prog_tests/task_local_data.h b/tools/testing/selftests/bpf/prog_tests/task_local_data.h
> new file mode 100644
> index 000000000000..a408d10c3688
> --- /dev/null
> +++ b/tools/testing/selftests/bpf/prog_tests/task_local_data.h
> @@ -0,0 +1,386 @@
> +/* SPDX-License-Identifier: GPL-2.0 */
> +#ifndef __TASK_LOCAL_DATA_H
> +#define __TASK_LOCAL_DATA_H
> +
> +#include <errno.h>
> +#include <fcntl.h>
> +#include <sched.h>
> +#include <stdatomic.h>
> +#include <stddef.h>
> +#include <stdlib.h>
> +#include <string.h>
> +#include <unistd.h>
> +#include <sys/syscall.h>
> +#include <sys/types.h>
> +
> +#ifdef TLD_FREE_DATA_ON_THREAD_EXIT
> +#include <pthread.h>
> +#endif
> +
> +#include <bpf/bpf.h>
> +
> +/*
> + * OPTIONS
> + *
> + *   Define the option before including the header
> + *
> + *   TLD_FREE_DATA_ON_THREAD_EXIT - Frees memory on thread exit automatically
> + *
> + *   Thread-specific memory for storing TLD is allocated lazily on the first call to
> + *   tld_get_data(). The thread that calls it must also call tld_free() on thread exit
> + *   to prevent memory leak. Pthread will be included if the option is defined. A pthread
> + *   key will be registered with a destructor that calls tld_free().
> + *
> + *
> + *   TLD_DYN_DATA_SIZE - The maximum size of memory allocated for TLDs created dynamically
> + *   (default: 64 bytes)
> + *
> + *   A TLD can be defined statically using TLD_DEFINE_KEY() or created on the fly using
> + *   tld_create_key(). As the total size of TLDs created with tld_create_key() cannot be
> + *   possibly known statically, a memory area of size TLD_DYN_DATA_SIZE will be allocated
> + *   for these TLDs. This additional memory is allocated for every thread that calls
> + *   tld_get_data() even if no tld_create_key are actually called, so be mindful of
> + *   potential memory wastage. Use TLD_DEFINE_KEY() whenever possible as just enough memory
> + *   will be allocated for TLDs created with it.
> + *
> + *
> + *   TLD_NAME_LEN - The maximum length of the name of a TLD (default: 62)
> + *
> + *   Setting TLD_NAME_LEN will affect the maximum number of TLDs a process can store,
> + *   TLD_MAX_DATA_CNT.
> + *
> + *
> + *   TLD_DATA_USE_ALIGNED_ALLOC - Always use aligned_alloc() instead of malloc()
> + *
> + *   When allocating the memory for storing TLDs, we need to make sure there is a memory
> + *   region of the X bytes within a page. This is due to the limit posed by UPTR: memory
> + *   pinned to the kernel cannot exceed a page nor can it cross the page boundary. The
> + *   library normally calls malloc(2*X) given X bytes of total TLDs, and only uses
> + *   aligned_alloc(PAGE_SIZE, X) when X >= PAGE_SIZE / 2. This is to reduce memory wastage
> + *   as not all memory allocator can use the exact amount of memory requested to fulfill
> + *   aligned_alloc(). For example, some may round the size up to the alignment. Enable the
> + *   option to always use aligned_alloc() if the implementation has low memory overhead.
> + */
> +
> +#define TLD_PAGE_SIZE getpagesize()
> +#define TLD_PAGE_MASK (~(TLD_PAGE_SIZE - 1))
> +
> +#define TLD_ROUND_MASK(x, y) ((__typeof__(x))((y) - 1))
> +#define TLD_ROUND_UP(x, y) ((((x) - 1) | TLD_ROUND_MASK(x, y)) + 1)
> +
> +#define TLD_READ_ONCE(x) (*(volatile typeof(x) *)&(x))
> +
> +#ifndef TLD_DYN_DATA_SIZE
> +#define TLD_DYN_DATA_SIZE 64
> +#endif
> +
> +#define TLD_MAX_DATA_CNT (TLD_PAGE_SIZE / sizeof(struct tld_metadata) - 1)
> +
> +#ifndef TLD_NAME_LEN
> +#define TLD_NAME_LEN 62
> +#endif
> +
> +#ifdef __cplusplus
> +extern "C" {
> +#endif
> +
> +typedef struct {
> +       __s16 off;
> +} tld_key_t;
> +
> +struct tld_metadata {
> +       char name[TLD_NAME_LEN];
> +       _Atomic __u16 size;
> +};
> +
> +struct tld_meta_u {
> +       _Atomic __u8 cnt;
> +       __u16 size;
> +       struct tld_metadata metadata[];
> +};
> +
> +struct tld_data_u {
> +       __u64 start; /* offset of tld_data_u->data in a page */
> +       char data[];
> +};
> +
> +struct tld_map_value {
> +       void *data;
> +       struct tld_meta_u *meta;
> +};
> +
> +struct tld_meta_u * _Atomic tld_meta_p __attribute__((weak));
> +__thread struct tld_data_u *tld_data_p __attribute__((weak));
> +__thread void *tld_data_alloc_p __attribute__((weak));
> +
> +#ifdef TLD_FREE_DATA_ON_THREAD_EXIT
> +pthread_key_t tld_pthread_key __attribute__((weak));
> +
> +static void tld_free(void);
> +
> +static void __tld_thread_exit_handler(void *unused)
> +{
> +       tld_free();
> +}
> +#endif
> +
> +static int __tld_init_meta_p(void)
> +{
> +       struct tld_meta_u *meta, *uninit = NULL;
> +       int err = 0;
> +
> +       meta = (struct tld_meta_u *)aligned_alloc(TLD_PAGE_SIZE, TLD_PAGE_SIZE);
> +       if (!meta) {
> +               err = -ENOMEM;
> +               goto out;
> +       }
> +
> +       memset(meta, 0, TLD_PAGE_SIZE);
> +       meta->size = TLD_DYN_DATA_SIZE;
> +
> +       if (!atomic_compare_exchange_strong(&tld_meta_p, &uninit, meta)) {
> +               free(meta);
> +               goto out;
> +       }
> +
> +#ifdef TLD_FREE_DATA_ON_THREAD_EXIT
> +       pthread_key_create(&tld_pthread_key, __tld_thread_exit_handler);
> +#endif
> +out:
> +       return err;
> +}
> +
> +static int __tld_init_data_p(int map_fd)
> +{
> +       bool use_aligned_alloc = false;
> +       struct tld_map_value map_val;
> +       struct tld_data_u *data;
> +       void *data_alloc = NULL;
> +       int err, tid_fd = -1;
> +
> +       tid_fd = syscall(SYS_pidfd_open, gettid(), O_EXCL);
> +       if (tid_fd < 0) {
> +               err = -errno;
> +               goto out;
> +       }
> +
> +#ifdef TLD_DATA_USE_ALIGNED_ALLOC
> +       use_aligned_alloc = true;
> +#endif
> +
> +       /*
> +        * tld_meta_p->size = TLD_DYN_DATA_SIZE +
> +        *          total size of TLDs defined via TLD_DEFINE_KEY()
> +        */
> +       data_alloc = (use_aligned_alloc || tld_meta_p->size * 2 >= TLD_PAGE_SIZE) ?
> +               aligned_alloc(TLD_PAGE_SIZE, tld_meta_p->size) :
> +               malloc(tld_meta_p->size * 2);
> +       if (!data_alloc) {
> +               err = -ENOMEM;
> +               goto out;
> +       }
> +
> +       /*
> +        * Always pass a page-aligned address to UPTR since the size of tld_map_value::data
> +        * is a page in BTF. If data_alloc spans across two pages, use the page that contains large
> +        * enough memory.
> +        */
> +       if (TLD_PAGE_SIZE - (~TLD_PAGE_MASK & (intptr_t)data_alloc) >= tld_meta_p->size) {
> +               map_val.data = (void *)(TLD_PAGE_MASK & (intptr_t)data_alloc);
> +               data = data_alloc;
> +               data->start = (~TLD_PAGE_MASK & (intptr_t)data_alloc) +
> +                             offsetof(struct tld_data_u, data);
> +       } else {
> +               map_val.data = (void *)(TLD_ROUND_UP((intptr_t)data_alloc, TLD_PAGE_SIZE));
> +               data = (void *)(TLD_ROUND_UP((intptr_t)data_alloc, TLD_PAGE_SIZE));
> +               data->start = offsetof(struct tld_data_u, data);
> +       }
> +       map_val.meta = TLD_READ_ONCE(tld_meta_p);
> +
> +       err = bpf_map_update_elem(map_fd, &tid_fd, &map_val, 0);
> +       if (err) {
> +               free(data_alloc);
> +               goto out;
> +       }
> +
> +       tld_data_p = data;
> +       tld_data_alloc_p = data_alloc;
> +#ifdef TLD_FREE_DATA_ON_THREAD_EXIT
> +       pthread_setspecific(tld_pthread_key, (void *)1);
> +#endif
> +out:
> +       if (tid_fd >= 0)
> +               close(tid_fd);
> +       return err;
> +}
> +
> +static tld_key_t __tld_create_key(const char *name, size_t size, bool dyn_data)
> +{
> +       int err, i, sz, off = 0;
> +       __u8 cnt;
> +
> +       if (!TLD_READ_ONCE(tld_meta_p)) {
> +               err = __tld_init_meta_p();
> +               if (err)
> +                       return (tld_key_t){err};
> +       }
> +
> +       for (i = 0; i < TLD_MAX_DATA_CNT; i++) {
> +retry:
> +               cnt = atomic_load(&tld_meta_p->cnt);
> +               if (i < cnt) {
> +                       /* A metadata is not ready until size is updated with a non-zero value */
> +                       while (!(sz = atomic_load(&tld_meta_p->metadata[i].size)))
> +                               sched_yield();
> +
> +                       if (!strncmp(tld_meta_p->metadata[i].name, name, TLD_NAME_LEN))
> +                               return (tld_key_t){-EEXIST};
> +
> +                       off += TLD_ROUND_UP(sz, 8);
> +                       continue;
> +               }
> +
> +               /*
> +                * TLD_DEFINE_KEY() is given memory upto a page while at most
> +                * TLD_DYN_DATA_SIZE is allocated for tld_create_key()
> +                */
> +               if (dyn_data) {
> +                       if (off + TLD_ROUND_UP(size, 8) > tld_meta_p->size)
> +                               return (tld_key_t){-E2BIG};
> +               } else {
> +                       if (off + TLD_ROUND_UP(size, 8) > TLD_PAGE_SIZE - sizeof(struct tld_data_u))
> +                               return (tld_key_t){-E2BIG};
> +                       tld_meta_p->size += TLD_ROUND_UP(size, 8);
> +               }
> +
> +               /*
> +                * Only one tld_create_key() can increase the current cnt by one and
> +                * takes the latest available slot. Other threads will check again if a new
> +                * TLD can still be added, and then compete for the new slot after the
> +                * succeeding thread update the size.
> +                */
> +               if (!atomic_compare_exchange_strong(&tld_meta_p->cnt, &cnt, cnt + 1))
> +                       goto retry;
> +
> +               strncpy(tld_meta_p->metadata[i].name, name, TLD_NAME_LEN);
> +               atomic_store(&tld_meta_p->metadata[i].size, size);
> +               return (tld_key_t){(__s16)off};
> +       }
> +
> +       return (tld_key_t){-ENOSPC};
> +}
> +
> +/**
> + * TLD_DEFINE_KEY() - Define a TLD and a global variable key associated with the TLD.
> + *
> + * @name: The name of the TLD
> + * @size: The size of the TLD
> + * @key: The variable name of the key. Cannot exceed TLD_NAME_LEN
> + *
> + * The macro can only be used in file scope.
> + *
> + * A global variable key of opaque type, tld_key_t, will be declared and initialized before
> + * main() starts. Use tld_key_is_err() or tld_key_err_or_zero() later to check if the key
> + * creation succeeded. Pass the key to tld_get_data() to get a pointer to the TLD.
> + * bpf programs can also fetch the same key by name.
> + *
> + * The total size of TLDs created using TLD_DEFINE_KEY() cannot exceed a page. Just
> + * enough memory will be allocated for each thread on the first call to tld_get_data().
> + */
> +#define TLD_DEFINE_KEY(key, name, size)                        \
> +tld_key_t key;                                         \
> +                                                       \
> +__attribute__((constructor))                           \
> +void __tld_define_key_##key(void)                      \
> +{                                                      \
> +       key = __tld_create_key(name, size, false);      \
> +}
> +
> +/**
> + * tld_create_key() - Create a TLD and return a key associated with the TLD.
> + *
> + * @name: The name the TLD
> + * @size: The size of the TLD
> + *
> + * Return an opaque object key. Use tld_key_is_err() or tld_key_err_or_zero() to check
> + * if the key creation succeeded. Pass the key to tld_get_data() to get a pointer to
> + * locate the TLD. bpf programs can also fetch the same key by name.
> + *
> + * Use tld_create_key() only when a TLD needs to be created dynamically (e.g., @name is
> + * not known statically or a TLD needs to be created conditionally)
> + *
> + * An additional TLD_DYN_DATA_SIZE bytes are allocated per-thread to accommodate TLDs
> + * created dynamically with tld_create_key(). Since only a user page is pinned to the
> + * kernel, when TLDs created with TLD_DEFINE_KEY() uses more than TLD_PAGE_SIZE -
> + * TLD_DYN_DATA_SIZE, the buffer size will be limited to the rest of the page.
> + */
> +__attribute__((unused))
> +static tld_key_t tld_create_key(const char *name, size_t size)
> +{
> +       return __tld_create_key(name, size, true);
> +}
> +
> +__attribute__((unused))
> +static inline bool tld_key_is_err(tld_key_t key)
> +{
> +       return key.off < 0;
> +}
> +
> +__attribute__((unused))
> +static inline int tld_key_err_or_zero(tld_key_t key)
> +{
> +       return tld_key_is_err(key) ? key.off : 0;
> +}
> +
> +/**
> + * tld_get_data() - Get a pointer to the TLD associated with the given key of the
> + * calling thread.
> + *
> + * @map_fd: A file descriptor of tld_data_map, the underlying BPF task local storage map
> + * of task local data.
> + * @key: A key object created by TLD_DEFINE_KEY() or tld_create_key().
> + *
> + * Return a pointer to the TLD if the key is valid; NULL if not enough memory for TLD
> + * for this thread, or the key is invalid. The returned pointer is guaranteed to be 8-byte
> + * aligned.
> + *
> + * Threads that call tld_get_data() must call tld_free() on exit to prevent
> + * memory leak if TLD_FREE_DATA_ON_THREAD_EXIT is not defined.
> + */
> +__attribute__((unused))
> +static void *tld_get_data(int map_fd, tld_key_t key)
> +{
> +       if (!TLD_READ_ONCE(tld_meta_p))
> +               return NULL;
> +
> +       /* tld_data_p is allocated on the first invocation of tld_get_data() */
> +       if (!tld_data_p && __tld_init_data_p(map_fd))
> +               return NULL;
> +
> +       return tld_data_p->data + key.off;
> +}
> +
> +/**
> + * tld_free() - Free task local data memory of the calling thread
> + *
> + * For the calling thread, all pointers to TLDs acquired before will become invalid.
> + *
> + * Users must call tld_free() on thread exit to prevent memory leak. Alternatively,
> + * define TLD_FREE_DATA_ON_THREAD_EXIT and a thread exit handler will be registered
> + * to free the memory automatically.
> + */
> +__attribute__((unused))
> +static void tld_free(void)
> +{
> +       if (tld_data_alloc_p) {
> +               free(tld_data_alloc_p);
> +               tld_data_alloc_p = NULL;
> +               tld_data_p = NULL;
> +       }
> +}
> +
> +#ifdef __cplusplus
> +} /* extern "C" */
> +#endif
> +
> +#endif /* __TASK_LOCAL_DATA_H */
> diff --git a/tools/testing/selftests/bpf/progs/task_local_data.bpf.h b/tools/testing/selftests/bpf/progs/task_local_data.bpf.h
> new file mode 100644
> index 000000000000..432fff2af844
> --- /dev/null
> +++ b/tools/testing/selftests/bpf/progs/task_local_data.bpf.h
> @@ -0,0 +1,237 @@
> +/* SPDX-License-Identifier: GPL-2.0 */
> +#ifndef __TASK_LOCAL_DATA_BPF_H
> +#define __TASK_LOCAL_DATA_BPF_H
> +
> +/*
> + * Task local data is a library that facilitates sharing per-task data
> + * between user space and bpf programs.
> + *
> + *
> + * USAGE
> + *
> + * A TLD, an entry of data in task local data, first needs to be created by the
> + * user space. This is done by calling user space API, TLD_DEFINE_KEY() or
> + * tld_create_key(), with the name of the TLD and the size.
> + *
> + * TLD_DEFINE_KEY(prio, "priority", sizeof(int));
> + *
> + * or
> + *
> + * void func_call(...) {
> + *     tld_key_t prio, in_cs;
> + *
> + *     prio = tld_create_key("priority", sizeof(int));
> + *     in_cs = tld_create_key("in_critical_section", sizeof(bool));
> + *     ...
> + *
> + * A key associated with the TLD, which has an opaque type tld_key_t, will be
> + * initialized or returned. It can be used to get a pointer to the TLD in the
> + * user space by calling tld_get_data().
> + *
> + * In a bpf program, tld_object_init() first needs to be called to initialized a
> + * tld_object on the stack. Then, TLDs can be accessed by calling tld_get_data().
> + * The API will try to fetch the key by the name and use it to locate the data.
> + * A pointer to the TLD will be returned. It also caches the key in a task local
> + * storage map, tld_key_map, whose value type, struct tld_keys, must be defined
> + * by the developer.
> + *
> + * struct tld_keys {
> + *     tld_key_t prio;
> + *     tld_key_t in_cs;
> + * };
> + *
> + * SEC("struct_ops")
> + * void prog(struct task_struct task, ...)
> + * {
> + *     struct tld_object tld_obj;
> + *     int err, *p;
> + *
> + *     err = tld_object_init(task, &tld_obj);
> + *     if (err)
> + *         return;
> + *
> + *     p = tld_get_data(&tld_obj, prio, "priority", sizeof(int));
> + *     if (p)
> + *         // do something depending on *p
> + */
> +#include <errno.h>
> +#include <bpf/bpf_helpers.h>
> +
> +#define TLD_ROUND_MASK(x, y) ((__typeof__(x))((y) - 1))
> +#define TLD_ROUND_UP(x, y) ((((x) - 1) | TLD_ROUND_MASK(x, y)) + 1)
> +
> +#define TLD_MAX_DATA_CNT (__PAGE_SIZE / sizeof(struct tld_metadata) - 1)
> +
> +#ifndef TLD_NAME_LEN
> +#define TLD_NAME_LEN 62
> +#endif
> +
> +#ifndef TLD_KEY_MAP_CREATE_RETRY
> +#define TLD_KEY_MAP_CREATE_RETRY 10
> +#endif
> +
> +typedef struct {
> +       __s16 off;
> +} tld_key_t;
> +
> +struct tld_metadata {
> +       char name[TLD_NAME_LEN];
> +       __u16 size;
> +};
> +
> +struct tld_meta_u {
> +       __u8 cnt;
> +       __u16 size;
> +       struct tld_metadata metadata[TLD_MAX_DATA_CNT];
> +};
> +
> +struct tld_data_u {
> +       __u64 start; /* offset of tld_data_u->data in a page */
> +       char data[__PAGE_SIZE - sizeof(__u64)];
> +};
> +
> +struct tld_map_value {
> +       struct tld_data_u __uptr *data;
> +       struct tld_meta_u __uptr *meta;
> +};
> +
> +typedef struct tld_uptr_dummy {
> +       struct tld_data_u data[0];
> +       struct tld_meta_u meta[0];
> +} *tld_uptr_dummy_t;
> +
> +struct tld_object {
> +       struct tld_map_value *data_map;
> +       struct tld_keys *key_map;
> +       /*
> +        * Force the compiler to generate the actual definition of tld_meta_u
> +        * and tld_data_u in BTF. Without it, tld_meta_u and u_tld_data will
> +        * be BTF_KIND_FWD.
> +        */
> +       tld_uptr_dummy_t dummy[0];
> +};
> +
> +/*
> + * Map value of tld_key_map for caching keys. Must be defined by the developer.
> + * Members should be tld_key_t and passed to the 3rd argument of tld_fetch_key().
> + */
> +struct tld_keys;
> +
> +struct {
> +       __uint(type, BPF_MAP_TYPE_TASK_STORAGE);
> +       __uint(map_flags, BPF_F_NO_PREALLOC);
> +       __type(key, int);
> +       __type(value, struct tld_map_value);
> +} tld_data_map SEC(".maps");
> +
> +struct {
> +       __uint(type, BPF_MAP_TYPE_TASK_STORAGE);
> +       __uint(map_flags, BPF_F_NO_PREALLOC);
> +       __type(key, int);
> +       __type(value, struct tld_keys);
> +} tld_key_map SEC(".maps");
> +
> +/**
> + * tld_object_init() - Initialize a tld_object.
> + *
> + * @task: The task_struct of the target task
> + * @tld_obj: A pointer to a tld_object to be initialized
> + *
> + * Return 0 on success; -ENODATA if the user space did not initialize task local data
> + * for the current task through tld_get_data(); -ENOMEM if the creation of tld_key_map
> + * fails
> + */
> +__attribute__((unused))
> +static int tld_object_init(struct task_struct *task, struct tld_object *tld_obj)
> +{
> +       int i;
> +
> +       tld_obj->data_map = bpf_task_storage_get(&tld_data_map, task, 0, 0);
> +       if (!tld_obj->data_map)
> +               return -ENODATA;
> +
> +       bpf_for(i, 0, TLD_KEY_MAP_CREATE_RETRY) {
> +               tld_obj->key_map = bpf_task_storage_get(&tld_key_map, task, 0,
> +                                                       BPF_LOCAL_STORAGE_GET_F_CREATE);
> +               if (likely(tld_obj->key_map))
> +                       break;
> +       }
> +       if (!tld_obj->key_map)
> +               return -ENOMEM;
> +
> +       return 0;
> +}
> +
> +/*
> + * Return the offset of TLD if @name is found. Otherwise, return the current TLD count
> + * using the nonpositive range so that the next tld_get_data() can skip fetching key if
> + * no new TLD is added or start comparing name from the first newly added TLD.
> + */
> +__attribute__((unused))
> +static int __tld_fetch_key(struct tld_object *tld_obj, const char *name, int i_start)
> +{
> +       struct tld_metadata *metadata;
> +       int i, cnt, start, off = 0;
> +
> +       if (!tld_obj->data_map || !tld_obj->data_map->data || !tld_obj->data_map->meta)
> +               return 0;
> +
> +       start = tld_obj->data_map->data->start;
> +       cnt = tld_obj->data_map->meta->cnt;
> +       metadata = tld_obj->data_map->meta->metadata;
> +
> +       bpf_for(i, 0, cnt) {
> +               if (i >= TLD_MAX_DATA_CNT)
> +                       break;
> +
> +               if (i >= i_start && !bpf_strncmp(metadata[i].name, TLD_NAME_LEN, name))
> +                       return start + off;
> +
> +               off += TLD_ROUND_UP(metadata[i].size, 8);
> +       }
> +
> +       return -cnt;
> +}
> +
> +/**
> + * tld_get_data() - Retrieve a pointer to the TLD associated with the name.
> + *
> + * @tld_obj: A pointer to a valid tld_object initialized by tld_object_init()
> + * @key: The cached key of the TLD in tld_key_map
> + * @name: The name of the key associated with a TLD
> + * @size: The size of the TLD. Must be a known constant value
> + *
> + * Return a pointer to the TLD associated with @name; NULL if not found or @size is too
> + * big. @key is used to cache the key if the TLD is found to speed up subsequent calls.
> + * It should be defined as an member of tld_keys of tld_key_t type by the developer.
> + */
> +#define tld_get_data(tld_obj, key, name, size)                                         \
> +       ({                                                                              \
> +               void *data = NULL, *_data = (tld_obj)->data_map->data;                  \
> +               long off = (tld_obj)->key_map->key.off;                                 \
> +               int cnt;                                                                \
> +                                                                                       \
> +               if (likely(_data)) {                                                    \
> +                       if (likely(off > 0)) {                                          \
> +                               barrier_var(off);                                       \
> +                               if (likely(off < __PAGE_SIZE - size))                   \
> +                                       data = _data + off;                             \
> +                       } else {                                                        \
> +                               cnt = -(off);                                           \
> +                               if (likely((tld_obj)->data_map->meta) &&                \
> +                                   cnt < (tld_obj)->data_map->meta->cnt) {             \
> +                                       off = __tld_fetch_key(tld_obj, name, cnt);      \
> +                                       (tld_obj)->key_map->key.off = off;              \
> +                                                                                       \
> +                                       if (likely(off < __PAGE_SIZE - size)) {         \
> +                                               barrier_var(off);                       \
> +                                               if (off > 0)                            \
> +                                                       data = _data + off;             \
> +                                       }                                               \
> +                               }                                                       \
> +                       }                                                               \
> +               }                                                                       \
> +               data;                                                                   \
> +       })
> +
> +#endif
> --
> 2.47.3
>

Re: [PATCH bpf-next v7 3/4] selftests/bpf: Test basic task local data operations

[Emil Tsalapatis]

On Wed, Jul 30, 2025 at 2:59 PM Amery Hung <ameryhung@gmail.com> wrote:
>
> Test basic operations of task local data with valid and invalid
> tld_create_key().
>
> For invalid calls, make sure they return the right error code and check
> that the TLDs are not inserted by running tld_get_data("
> value_not_exists") on the bpf side. The call should a null pointer.
>
> For valid calls, first make sure the TLDs are created by calling
> tld_get_data() on the bpf side. The call should return a valid pointer.
>
> Finally, verify that the TLDs are indeed task-specific (i.e., their
> addresses do not overlap) with multiple user threads. This done by
> writing values unique to each thread, reading them from both user space
> and bpf, and checking if the value read back matches the value written.
>
> Signed-off-by: Amery Hung <ameryhung@gmail.com>

Reviewed-by: Emil Tsalapatis <emil@etsalapatis.com>


> ---
>  .../bpf/prog_tests/test_task_local_data.c     | 192 ++++++++++++++++++
>  .../bpf/progs/test_task_local_data.c          |  65 ++++++
>  2 files changed, 257 insertions(+)
>  create mode 100644 tools/testing/selftests/bpf/prog_tests/test_task_local_data.c
>  create mode 100644 tools/testing/selftests/bpf/progs/test_task_local_data.c
>
> diff --git a/tools/testing/selftests/bpf/prog_tests/test_task_local_data.c b/tools/testing/selftests/bpf/prog_tests/test_task_local_data.c
> new file mode 100644
> index 000000000000..2e77d3fa2534
> --- /dev/null
> +++ b/tools/testing/selftests/bpf/prog_tests/test_task_local_data.c
> @@ -0,0 +1,192 @@
> +// SPDX-License-Identifier: GPL-2.0
> +#include <pthread.h>
> +#include <bpf/btf.h>
> +#include <test_progs.h>
> +
> +#define TLD_FREE_DATA_ON_THREAD_EXIT
> +#define TLD_DYN_DATA_SIZE 4096
> +#include "task_local_data.h"
> +
> +struct test_tld_struct {
> +       __u64 a;
> +       __u64 b;
> +       __u64 c;
> +       __u64 d;
> +};
> +
> +#include "test_task_local_data.skel.h"
> +
> +TLD_DEFINE_KEY(value0_key, "value0", sizeof(int));
> +
> +/*
> + * Reset task local data between subtests by clearing metadata other
> + * than the statically defined value0. This is safe as subtests run
> + * sequentially. Users of task local data library should not touch
> + * library internal.
> + */
> +static void reset_tld(void)
> +{
> +       if (TLD_READ_ONCE(tld_meta_p)) {
> +               /* Remove TLDs created by tld_create_key() */
> +               tld_meta_p->cnt = 1;
> +               tld_meta_p->size = TLD_DYN_DATA_SIZE;
> +               memset(&tld_meta_p->metadata[1], 0,
> +                      (TLD_MAX_DATA_CNT - 1) * sizeof(struct tld_metadata));
> +       }
> +}
> +
> +/* Serialize access to bpf program's global variables */
> +static pthread_mutex_t global_mutex;
> +
> +static tld_key_t *tld_keys;
> +
> +#define TEST_BASIC_THREAD_NUM 32
> +
> +void *test_task_local_data_basic_thread(void *arg)
> +{
> +       LIBBPF_OPTS(bpf_test_run_opts, opts);
> +       struct test_task_local_data *skel = (struct test_task_local_data *)arg;
> +       int fd, err, tid, *value0, *value1;
> +       struct test_tld_struct *value2;
> +
> +       fd = bpf_map__fd(skel->maps.tld_data_map);
> +
> +       value0 = tld_get_data(fd, value0_key);
> +       if (!ASSERT_OK_PTR(value0, "tld_get_data"))
> +               goto out;
> +
> +       value1 = tld_get_data(fd, tld_keys[1]);
> +       if (!ASSERT_OK_PTR(value1, "tld_get_data"))
> +               goto out;
> +
> +       value2 = tld_get_data(fd, tld_keys[2]);
> +       if (!ASSERT_OK_PTR(value2, "tld_get_data"))
> +               goto out;
> +
> +       tid = gettid();
> +
> +       *value0 = tid + 0;
> +       *value1 = tid + 1;
> +       value2->a = tid + 2;
> +       value2->b = tid + 3;
> +       value2->c = tid + 4;
> +       value2->d = tid + 5;
> +
> +       pthread_mutex_lock(&global_mutex);
> +       /* Run task_main that read task local data and save to global variables */
> +       err = bpf_prog_test_run_opts(bpf_program__fd(skel->progs.task_main), &opts);
> +       ASSERT_OK(err, "run task_main");
> +       ASSERT_OK(opts.retval, "task_main retval");
> +
> +       ASSERT_EQ(skel->bss->test_value0, tid + 0, "tld_get_data value0");
> +       ASSERT_EQ(skel->bss->test_value1, tid + 1, "tld_get_data value1");
> +       ASSERT_EQ(skel->bss->test_value2.a, tid + 2, "tld_get_data value2.a");
> +       ASSERT_EQ(skel->bss->test_value2.b, tid + 3, "tld_get_data value2.b");
> +       ASSERT_EQ(skel->bss->test_value2.c, tid + 4, "tld_get_data value2.c");
> +       ASSERT_EQ(skel->bss->test_value2.d, tid + 5, "tld_get_data value2.d");
> +       pthread_mutex_unlock(&global_mutex);
> +
> +       /* Make sure valueX are indeed local to threads */
> +       ASSERT_EQ(*value0, tid + 0, "value0");
> +       ASSERT_EQ(*value1, tid + 1, "value1");
> +       ASSERT_EQ(value2->a, tid + 2, "value2.a");
> +       ASSERT_EQ(value2->b, tid + 3, "value2.b");
> +       ASSERT_EQ(value2->c, tid + 4, "value2.c");
> +       ASSERT_EQ(value2->d, tid + 5, "value2.d");
> +
> +       *value0 = tid + 5;
> +       *value1 = tid + 4;
> +       value2->a = tid + 3;
> +       value2->b = tid + 2;
> +       value2->c = tid + 1;
> +       value2->d = tid + 0;
> +
> +       /* Run task_main again */
> +       pthread_mutex_lock(&global_mutex);
> +       err = bpf_prog_test_run_opts(bpf_program__fd(skel->progs.task_main), &opts);
> +       ASSERT_OK(err, "run task_main");
> +       ASSERT_OK(opts.retval, "task_main retval");
> +
> +       ASSERT_EQ(skel->bss->test_value0, tid + 5, "tld_get_data value0");
> +       ASSERT_EQ(skel->bss->test_value1, tid + 4, "tld_get_data value1");
> +       ASSERT_EQ(skel->bss->test_value2.a, tid + 3, "tld_get_data value2.a");
> +       ASSERT_EQ(skel->bss->test_value2.b, tid + 2, "tld_get_data value2.b");
> +       ASSERT_EQ(skel->bss->test_value2.c, tid + 1, "tld_get_data value2.c");
> +       ASSERT_EQ(skel->bss->test_value2.d, tid + 0, "tld_get_data value2.d");
> +       pthread_mutex_unlock(&global_mutex);
> +
> +out:
> +       pthread_exit(NULL);
> +}
> +
> +static void test_task_local_data_basic(void)
> +{
> +       struct test_task_local_data *skel;
> +       pthread_t thread[TEST_BASIC_THREAD_NUM];
> +       char dummy_key_name[TLD_NAME_LEN];
> +       tld_key_t key;
> +       int i, err;
> +
> +       reset_tld();
> +
> +       ASSERT_OK(pthread_mutex_init(&global_mutex, NULL), "pthread_mutex_init");
> +
> +       skel = test_task_local_data__open_and_load();
> +       if (!ASSERT_OK_PTR(skel, "skel_open_and_load"))
> +               return;
> +
> +       tld_keys = calloc(TLD_MAX_DATA_CNT, sizeof(tld_key_t));
> +       if (!ASSERT_OK_PTR(tld_keys, "calloc tld_keys"))
> +               goto out;
> +
> +       ASSERT_FALSE(tld_key_is_err(value0_key), "TLD_DEFINE_KEY");
> +       tld_keys[1] = tld_create_key("value1", sizeof(int));
> +       ASSERT_FALSE(tld_key_is_err(tld_keys[1]), "tld_create_key");
> +       tld_keys[2] = tld_create_key("value2", sizeof(struct test_tld_struct));
> +       ASSERT_FALSE(tld_key_is_err(tld_keys[2]), "tld_create_key");
> +
> +       /*
> +        * Shouldn't be able to store data exceed a page. Create a TLD just big
> +        * enough to exceed a page. TLDs already created are int value0, int
> +        * value1, and struct test_tld_struct value2.
> +        */
> +       key = tld_create_key("value_not_exist",
> +                            TLD_PAGE_SIZE - 2 * sizeof(int) - sizeof(struct test_tld_struct) + 1);
> +       ASSERT_EQ(tld_key_err_or_zero(key), -E2BIG, "tld_create_key");
> +
> +       key = tld_create_key("value2", sizeof(struct test_tld_struct));
> +       ASSERT_EQ(tld_key_err_or_zero(key), -EEXIST, "tld_create_key");
> +
> +       /* Shouldn't be able to create the (TLD_MAX_DATA_CNT+1)-th TLD */
> +       for (i = 3; i < TLD_MAX_DATA_CNT; i++) {
> +               snprintf(dummy_key_name, TLD_NAME_LEN, "dummy_value%d", i);
> +               tld_keys[i] = tld_create_key(dummy_key_name, sizeof(int));
> +               ASSERT_FALSE(tld_key_is_err(tld_keys[i]), "tld_create_key");
> +       }
> +       key = tld_create_key("value_not_exist", sizeof(struct test_tld_struct));
> +       ASSERT_EQ(tld_key_err_or_zero(key), -ENOSPC, "tld_create_key");
> +
> +       /* Access TLDs from multiple threads and check if they are thread-specific */
> +       for (i = 0; i < TEST_BASIC_THREAD_NUM; i++) {
> +               err = pthread_create(&thread[i], NULL, test_task_local_data_basic_thread, skel);
> +               if (!ASSERT_OK(err, "pthread_create"))
> +                       goto out;
> +       }
> +
> +out:
> +       for (i = 0; i < TEST_BASIC_THREAD_NUM; i++)
> +               pthread_join(thread[i], NULL);
> +
> +       if (tld_keys) {
> +               free(tld_keys);
> +               tld_keys = NULL;
> +       }
> +       tld_free();
> +       test_task_local_data__destroy(skel);
> +}
> +
> +void test_task_local_data(void)
> +{
> +       if (test__start_subtest("task_local_data_basic"))
> +               test_task_local_data_basic();
> +}
> diff --git a/tools/testing/selftests/bpf/progs/test_task_local_data.c b/tools/testing/selftests/bpf/progs/test_task_local_data.c
> new file mode 100644
> index 000000000000..fffafc013044
> --- /dev/null
> +++ b/tools/testing/selftests/bpf/progs/test_task_local_data.c
> @@ -0,0 +1,65 @@
> +// SPDX-License-Identifier: GPL-2.0
> +
> +#include <vmlinux.h>
> +#include <errno.h>
> +#include <bpf/bpf_helpers.h>
> +
> +#include "task_local_data.bpf.h"
> +
> +struct tld_keys {
> +       tld_key_t value0;
> +       tld_key_t value1;
> +       tld_key_t value2;
> +       tld_key_t value_not_exist;
> +};
> +
> +struct test_tld_struct {
> +       __u64 a;
> +       __u64 b;
> +       __u64 c;
> +       __u64 d;
> +};
> +
> +int test_value0;
> +int test_value1;
> +struct test_tld_struct test_value2;
> +
> +SEC("syscall")
> +int task_main(void *ctx)
> +{
> +       struct tld_object tld_obj;
> +       struct test_tld_struct *struct_p;
> +       struct task_struct *task;
> +       int err, *int_p;
> +
> +       task = bpf_get_current_task_btf();
> +       err = tld_object_init(task, &tld_obj);
> +       if (err)
> +               return 1;
> +
> +       int_p = tld_get_data(&tld_obj, value0, "value0", sizeof(int));
> +       if (int_p)
> +               test_value0 = *int_p;
> +       else
> +               return 2;
> +
> +       int_p = tld_get_data(&tld_obj, value1, "value1", sizeof(int));
> +       if (int_p)
> +               test_value1 = *int_p;
> +       else
> +               return 3;
> +
> +       struct_p = tld_get_data(&tld_obj, value2, "value2", sizeof(struct test_tld_struct));
> +       if (struct_p)
> +               test_value2 = *struct_p;
> +       else
> +               return 4;
> +
> +       int_p = tld_get_data(&tld_obj, value_not_exist, "value_not_exist", sizeof(int));
> +       if (int_p)
> +               return 5;
> +
> +       return 0;
> +}
> +
> +char _license[] SEC("license") = "GPL";
> --
> 2.47.3
>

Re: [PATCH bpf-next v7 4/4] selftests/bpf: Test concurrent task local data key creation

[Emil Tsalapatis]

On Wed, Jul 30, 2025 at 2:59 PM Amery Hung <ameryhung@gmail.com> wrote:
>
> Test thread-safety of tld_create_key(). Since tld_create_key() does
> not rely on locks but memory barriers and atomic operations to protect
> the shared metadata, the thread-safety of the function is non-trivial.
> Make sure concurrent tld_key_create(), both valid and invalid, can not
> race and corrupt metatada, which may leads to TLDs not being thread-
> specific or duplicate TLDs with the same name.
>
> Signed-off-by: Amery Hung <ameryhung@gmail.com>
Reviewed-by: Emil Tsalapatis <emil@etsalapatis.com>


> ---
>  .../bpf/prog_tests/test_task_local_data.c     | 105 ++++++++++++++++++
>  1 file changed, 105 insertions(+)
>
> diff --git a/tools/testing/selftests/bpf/prog_tests/test_task_local_data.c b/tools/testing/selftests/bpf/prog_tests/test_task_local_data.c
> index 2e77d3fa2534..3b5cd2cd89c7 100644
> --- a/tools/testing/selftests/bpf/prog_tests/test_task_local_data.c
> +++ b/tools/testing/selftests/bpf/prog_tests/test_task_local_data.c
> @@ -185,8 +185,113 @@ static void test_task_local_data_basic(void)
>         test_task_local_data__destroy(skel);
>  }
>
> +#define TEST_RACE_THREAD_NUM (TLD_MAX_DATA_CNT - 3)
> +
> +void *test_task_local_data_race_thread(void *arg)
> +{
> +       int err = 0, id = (intptr_t)arg;
> +       char key_name[32];
> +       tld_key_t key;
> +
> +       key = tld_create_key("value_not_exist", TLD_PAGE_SIZE + 1);
> +       if (tld_key_err_or_zero(key) != -E2BIG) {
> +               err = 1;
> +               goto out;
> +       }
> +
> +       /* Only one thread will succeed in creating value1 */
> +       key = tld_create_key("value1", sizeof(int));
> +       if (!tld_key_is_err(key))
> +               tld_keys[1] = key;
> +
> +       /* Only one thread will succeed in creating value2 */
> +       key = tld_create_key("value2", sizeof(struct test_tld_struct));
> +       if (!tld_key_is_err(key))
> +               tld_keys[2] = key;
> +
> +       snprintf(key_name, 32, "thread_%d", id);
> +       tld_keys[id] = tld_create_key(key_name, sizeof(int));
> +       if (tld_key_is_err(tld_keys[id]))
> +               err = 2;
> +out:
> +       return (void *)(intptr_t)err;
> +}
> +
> +static void test_task_local_data_race(void)
> +{
> +       LIBBPF_OPTS(bpf_test_run_opts, opts);
> +       pthread_t thread[TEST_RACE_THREAD_NUM];
> +       struct test_task_local_data *skel;
> +       int fd, i, j, err, *data;
> +       void *ret = NULL;
> +
> +       skel = test_task_local_data__open_and_load();
> +       if (!ASSERT_OK_PTR(skel, "skel_open_and_load"))
> +               return;
> +
> +       tld_keys = calloc(TLD_MAX_DATA_CNT, sizeof(tld_key_t));
> +       if (!ASSERT_OK_PTR(tld_keys, "calloc tld_keys"))
> +               goto out;
> +
> +       fd = bpf_map__fd(skel->maps.tld_data_map);
> +
> +       ASSERT_FALSE(tld_key_is_err(value0_key), "TLD_DEFINE_KEY");
> +       tld_keys[0] = value0_key;
> +
> +       for (j = 0; j < 100; j++) {
> +               reset_tld();
> +
> +               for (i = 0; i < TEST_RACE_THREAD_NUM; i++) {
> +                       /*
> +                        * Try to make tld_create_key() race with each other. Call
> +                        * tld_create_key(), both valid and invalid, from different threads.
> +                        */
> +                       err = pthread_create(&thread[i], NULL, test_task_local_data_race_thread,
> +                                            (void *)(intptr_t)(i + 3));
> +                       if (CHECK_FAIL(err))
> +                               break;
> +               }
> +
> +               /* Wait for all tld_create_key() to return */
> +               for (i = 0; i < TEST_RACE_THREAD_NUM; i++) {
> +                       pthread_join(thread[i], &ret);
> +                       if (CHECK_FAIL(ret))
> +                               break;
> +               }
> +
> +               /* Write a unique number to each TLD */
> +               for (i = 0; i < TLD_MAX_DATA_CNT; i++) {
> +                       data = tld_get_data(fd, tld_keys[i]);
> +                       if (CHECK_FAIL(!data))
> +                               break;
> +                       *data = i;
> +               }
> +
> +               /* Read TLDs and check the value to see if any address collides with another */
> +               for (i = 0; i < TLD_MAX_DATA_CNT; i++) {
> +                       data = tld_get_data(fd, tld_keys[i]);
> +                       if (CHECK_FAIL(*data != i))
> +                               break;
> +               }
> +
> +               /* Run task_main to make sure no invalid TLDs are added */
> +               err = bpf_prog_test_run_opts(bpf_program__fd(skel->progs.task_main), &opts);
> +               ASSERT_OK(err, "run task_main");
> +               ASSERT_OK(opts.retval, "task_main retval");
> +       }
> +out:
> +       if (tld_keys) {
> +               free(tld_keys);
> +               tld_keys = NULL;
> +       }
> +       tld_free();
> +       test_task_local_data__destroy(skel);
> +}
> +
>  void test_task_local_data(void)
>  {
>         if (test__start_subtest("task_local_data_basic"))
>                 test_task_local_data_basic();
> +       if (test__start_subtest("task_local_data_race"))
> +               test_task_local_data_race();
>  }
> --
> 2.47.3
>

Re: [PATCH bpf-next v7 0/4] Task local data

[patchwork-bot+netdevbpf]

Hello:

This series was applied to bpf/bpf-next.git (master)
by Alexei Starovoitov <ast@kernel.org>:

On Wed, 30 Jul 2025 11:58:51 -0700 you wrote:
> v6 -> v7
>   - Fix typos and improve the clarity of the cover letter (Emil)
>   - Some variable naming changes to make it less confusing (Emil)
>   - Add retry in tld_obj_init() as bpf_task_storage_get_recur() with
>     BPF_LOCAL_STORAGE_GET_F_CREATE can fail if the task local storage
>     is also being modified by other threads on the same CPU. This can
>     be especially easy to trigger in CI VM that only has two CPUs.
> 
> [...]

Here is the summary with links:
  - [bpf-next,v7,1/4] bpf: Allow syscall bpf programs to call non-recur helpers
    https://git.kernel.org/bpf/bpf-next/c/86de56487e5f
  - [bpf-next,v7,2/4] selftests/bpf: Introduce task local data
    https://git.kernel.org/bpf/bpf-next/c/31e838e1cdf4
  - [bpf-next,v7,3/4] selftests/bpf: Test basic task local data operations
    https://git.kernel.org/bpf/bpf-next/c/120f1a950e49
  - [bpf-next,v7,4/4] selftests/bpf: Test concurrent task local data key creation
    https://git.kernel.org/bpf/bpf-next/c/784181141782

You are awesome, thank you!
-- 
Deet-doot-dot, I am a bot.
https://korg.docs.kernel.org/patchwork/pwbot.html