Re: [PATCH v2 09/11] zsmalloc: embed struct zs_handle into zram and zswap entries

[Yosry Ahmed <yosry@kernel.org>]

On Tue, Nov 10, 2026 at 02:48:14PM +0900, Sergey Senozhatsky wrote:
> zsmalloc allocates an object from handle_cachep for every stored
> object, just to have a stable address which compaction and page
> migration can update when the object moves.  Both zram and zswap
> already have per-entry structures that hold the handle, so we basically
> keep a pointer to a pointer, which costs 8 bytes per stored object,
> slab alloc/free on every store and free, and one more memory load on
> every access.
> 
> Make struct zs_handle public and embed it into zram_table_entry and
> zswap_entry (their sizes don't change), so that zsmalloc writes object
> location directly there, and remove handle_cachep.  zs_malloc() now
> takes a handle pointer and returns an error code, zs_free() resets the
> handle.
> 
> Since compaction and migration update the handle in place, it must be
> naturally aligned, must not be on the stack (zs_malloc() warns about
> both), must not be copied, and must stay allocated until zs_free() or
> zs_rebind_handle().  zram recompression allocates a new object before
> it frees the old one, so it uses zram->recomp_handle, protected by
> ->dev_lock, and zs_rebind_handle() moves the new object to the slot
> handle.  zram recompression is the only user of zs_rebind_handle().
> 
> This saves 8 bytes per stored page, e.g. 2MB per 1GB of swapped out
> 4K pages.  Since handle is per-entry the memory savings scale.
> 
> Signed-off-by: Sergey Senozhatsky <senozhatsky@chromium.org>
> ---
>  Documentation/admin-guide/mm/slab.rst |   2 +-
>  Documentation/mm/zsmalloc.rst         |  18 ++--
>  drivers/block/zram/zram_drv.c         |  70 ++++++-------
>  drivers/block/zram/zram_drv.h         |   6 +-
>  include/linux/zsmalloc.h              |  29 +++++-
>  mm/zsmalloc.c                         | 144 ++++++++++++++++----------
>  mm/zswap.c                            |  22 ++--
>  7 files changed, 177 insertions(+), 114 deletions(-)

zsmalloc and zswap bits LGTM with a small nit below:

Reviewed-by: Yosry Ahmed <yosry@kernel.org>

[..]
>  
>  /*
> - * Pairs with READ_ONCE() in handle_to_obj(): zs_free() may read the
> - * handle locklessly, so prevent store tearing here.
> + * Pairs with READ_ONCE() in handle_to_obj() and zs_handle_is_allocated():
> + * zs_free(), zs_rebind_handle(), and zs_handle_is_allocated() may read
> + * the handle locklessly, so prevent store tearing here.
>   */
>  static void record_obj(struct zs_handle *handle, unsigned long obj)
>  {
> @@ -1307,31 +1292,37 @@ static unsigned long obj_malloc(struct zs_pool *pool,
>  /**
>   * zs_malloc - Allocate block of given size from pool.
>   * @pool: pool to allocate from
> + * @handle: handle to store allocated object in
>   * @size: size of block to allocate
>   * @gfp: gfp flags when allocating object
>   * @nid: The preferred node id to allocate new zspage (if needed)
>   *
> - * On success, handle to the allocated object is returned,
> - * otherwise an ERR_PTR().
>   * Allocation requests with size > ZS_MAX_ALLOC_SIZE will fail.
> + * zsmalloc keeps a back-pointer to @handle and updates it asynchronously
> + * when the object moves, so @handle must be unsigned-long aligned, must
> + * not be on the stack or modified/copied by the caller, and must remain
> + * allocated until zs_free() or zs_rebind_handle().
> + *
> + * Return: 0 on success, negative error code otherwise.
>   */
> -struct zs_handle *zs_malloc(struct zs_pool *pool, size_t size, gfp_t gfp,
> -			    const int nid)
> +int zs_malloc(struct zs_pool *pool, struct zs_handle *handle, size_t size,
> +	      gfp_t gfp, const int nid)
>  {
> -	struct zs_handle *handle;
>  	struct size_class *class;
>  	int newfg;
>  	struct zspage *zspage;
>  
> +	if (WARN_ON_ONCE(!IS_ALIGNED((unsigned long)handle,
> +				     __alignof__(unsigned long)) ||
> +			 object_is_on_stack(handle) ||
> +			 zs_handle_is_allocated(handle)))

Nit: Put this in a helper to reuse in zs_rebind_handle().

> +		return -EINVAL;
> +
[..]  
> +/**
> + * zs_rebind_handle() - Rebind an allocated object to a new handle.
> + * @pool: pool the object belongs to
> + * @dst: unallocated destination handle (must not be on the stack)
> + * @src: source handle that currently owns the object
> + *
> + * Rebinds the object from @src to @dst, updating the object header so
> + * compaction and page migration update @dst instead of @src, and resets
> + * @src.
> + */
> +void zs_rebind_handle(struct zs_pool *pool, struct zs_handle *dst,
> +		      struct zs_handle *src)
> +{
> +	struct size_class *class;
> +	struct zspage *zspage;
> +	struct zpdesc *zpdesc;
> +	struct link_free *link;
> +	unsigned int obj_idx;
> +	unsigned long obj;
> +	unsigned long off;
> +	void *vaddr;
> +
> +	/* Should never happen; bail out rather than corrupting handles */
> +	if (WARN_ON_ONCE(!IS_ALIGNED((unsigned long)dst,
> +				     __alignof__(unsigned long)) ||
> +			 object_is_on_stack(dst) ||
> +			 zs_handle_is_allocated(dst) ||
> +			 !zs_handle_is_allocated(src)))
> +		return;
[..]