Re: [PATCH bpf-next v3 4/7] bpf: bpf task work plumbing

[Eduard Zingerman <eddyz87@gmail.com>]

On Fri, 2025-09-05 at 17:45 +0100, Mykyta Yatsenko wrote:

[...]

> diff --git a/kernel/bpf/arraymap.c b/kernel/bpf/arraymap.c
> index 3d080916faf9..4130d8e76dff 100644
> --- a/kernel/bpf/arraymap.c
> +++ b/kernel/bpf/arraymap.c

[...]

> @@ -439,12 +439,14 @@ static void array_map_free_timers_wq(struct bpf_map *map)
>  	/* We don't reset or free fields other than timer and workqueue
>  	 * on uref dropping to zero.
>  	 */
> -	if (btf_record_has_field(map->record, BPF_TIMER | BPF_WORKQUEUE)) {
> +	if (btf_record_has_field(map->record, BPF_TIMER | BPF_WORKQUEUE | BPF_TASK_WORK)) {

I think that hashtab.c:htab_free_internal_structs needs to be renamed
and called here, thus avoiding code duplication.

>  		for (i = 0; i < array->map.max_entries; i++) {
>  			if (btf_record_has_field(map->record, BPF_TIMER))
>  				bpf_obj_free_timer(map->record, array_map_elem_ptr(array, i));
>  			if (btf_record_has_field(map->record, BPF_WORKQUEUE))
>  				bpf_obj_free_workqueue(map->record, array_map_elem_ptr(array, i));
> +			if (btf_record_has_field(map->record, BPF_TASK_WORK))
> +				bpf_obj_free_task_work(map->record, array_map_elem_ptr(array, i));
>  		}
>  	}
>  }

[...]

> diff --git a/kernel/bpf/btf.c b/kernel/bpf/btf.c
> index a1a9bc589518..73ca21911b30 100644
> --- a/kernel/bpf/btf.c
> +++ b/kernel/bpf/btf.c

[...]

> @@ -4034,6 +4037,10 @@ struct btf_record *btf_parse_fields(const struct btf *btf, const struct btf_type
>  		case BPF_LIST_NODE:
>  		case BPF_RB_NODE:
>  			break;
> +		case BPF_TASK_WORK:
> +			WARN_ON_ONCE(rec->task_work_off >= 0);
> +			rec->task_work_off = rec->fields[i].offset;
> +			break;

Nit: let's move this case up to BPF_WORKQUEUE or BPF_REFCOUNT,
     so that similar cases are grouped together.

>  		default:
>  			ret = -EFAULT;
>  			goto end;

[...]

> diff --git a/kernel/bpf/syscall.c b/kernel/bpf/syscall.c
> index 0fbfa8532c39..7da1ca893dfe 100644
> --- a/kernel/bpf/syscall.c
> +++ b/kernel/bpf/syscall.c

[...]

> @@ -840,6 +849,9 @@ void bpf_obj_free_fields(const struct btf_record *rec, void *obj)
>  				continue;
>  			bpf_rb_root_free(field, field_ptr, obj + rec->spin_lock_off);
>  			break;
> +		case BPF_TASK_WORK:
> +			bpf_task_work_cancel_and_free(field_ptr);
> +			break;

Nit: same here, let's keep similar cases together.

>  		case BPF_LIST_NODE:
>  		case BPF_RB_NODE:
>  		case BPF_REFCOUNT:

[...]

> diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
> index a5d19a01d488..6152536a834f 100644
> --- a/kernel/bpf/verifier.c
> +++ b/kernel/bpf/verifier.c
> @@ -2240,6 +2240,8 @@ static void mark_ptr_not_null_reg(struct bpf_reg_state *reg)
>  				reg->map_uid = reg->id;
>  			if (btf_record_has_field(map->inner_map_meta->record, BPF_WORKQUEUE))
>  				reg->map_uid = reg->id;
> +			if (btf_record_has_field(map->inner_map_meta->record, BPF_TASK_WORK))
> +				reg->map_uid = reg->id;

Nit: this can be shorter:

			if (btf_record_has_field(map->inner_map_meta->record,
						 BPF_TIMER | BPF_WORKQUEUE | BPF_TASK_WORK))
				reg->map_uid = reg->id;


>  		} else if (map->map_type == BPF_MAP_TYPE_XSKMAP) {
>  			reg->type = PTR_TO_XDP_SOCK;
>  		} else if (map->map_type == BPF_MAP_TYPE_SOCKMAP ||

[...]

> @@ -10943,6 +10956,35 @@ static int set_rbtree_add_callback_state(struct bpf_verifier_env *env,
>  	return 0;
>  }
>  
> +static int set_task_work_schedule_callback_state(struct bpf_verifier_env *env,
> +						 struct bpf_func_state *caller,
> +						 struct bpf_func_state *callee,
> +						 int insn_idx)
> +{
> +	struct bpf_map *map_ptr = caller->regs[BPF_REG_3].map_ptr;
> +
> +	/*
> +	 * callback_fn(struct bpf_map *map, void *key, void *value);
> +	 */
> +	callee->regs[BPF_REG_1].type = CONST_PTR_TO_MAP;
> +	__mark_reg_known_zero(&callee->regs[BPF_REG_1]);
> +	callee->regs[BPF_REG_1].map_ptr = map_ptr;
> +
> +	callee->regs[BPF_REG_2].type = PTR_TO_MAP_KEY;
> +	__mark_reg_known_zero(&callee->regs[BPF_REG_2]);
> +	callee->regs[BPF_REG_2].map_ptr = map_ptr;
> +
> +	callee->regs[BPF_REG_3].type = PTR_TO_MAP_VALUE;
> +	__mark_reg_known_zero(&callee->regs[BPF_REG_3]);
> +	callee->regs[BPF_REG_3].map_ptr = map_ptr;
> +
> +	/* unused */
> +	__mark_reg_not_init(env, &callee->regs[BPF_REG_4]);
> +	__mark_reg_not_init(env, &callee->regs[BPF_REG_5]);
> +	callee->in_callback_fn = true;

This should be `callee->in_async_callback_fn = true;` to avoid an
infinite loop check in the is_state_visisted() in some cases.

> +	return 0;
> +}
> +
>  static bool is_rbtree_lock_required_kfunc(u32 btf_id);
>  
>  /* Are we currently verifying the callback for a rbtree helper that must

[...]

> @@ -13171,6 +13235,15 @@ static int check_kfunc_args(struct bpf_verifier_env *env, struct bpf_kfunc_call_
>  					return -EINVAL;
>  				}
>  			}
> +			if (meta->map.ptr && reg->map_ptr->record->task_work_off >= 0) {
> +				if (meta->map.ptr != reg->map_ptr ||
> +				    meta->map.uid != reg->map_uid) {
> +					verbose(env,
> +						"bpf_task_work pointer in R2 map_uid=%d doesn't match map pointer in R3 map_uid=%d\n",
> +						meta->map.uid, reg->map_uid);
> +					return -EINVAL;
> +				}
> +			}

Please merge this with the case for wq_off above.

>  			meta->map.ptr = reg->map_ptr;
>  			meta->map.uid = reg->map_uid;
>  			fallthrough;

[...]