Re: [PATCH 6/7] zsmalloc: fold zs_obj_write() into zs_store()

[Yosry Ahmed <yosry@kernel.org>]

On Thu, Nov 05, 2026 at 01:47:36PM +0900, Sergey Senozhatsky wrote:
> zram and zswap always call zs_obj_write() right after zs_malloc(), and
> zs_obj_write() has to take pool->lock and look up object location
> again, which zs_malloc() just did.
> 
> Replace zs_malloc() and zs_obj_write() with one single zs_store()
> call.
> 
> Signed-off-by: Sergey Senozhatsky <senozhatsky@chromium.org>
> ---
>  Documentation/core-api/memory-allocation.rst |  2 +-
>  Documentation/mm/zsmalloc.rst                |  4 +-
>  drivers/block/zram/zram_drv.c                | 43 ++++++-----------
>  include/linux/zsmalloc.h                     |  6 +--
>  mm/zsmalloc.c                                | 51 +++++++++-----------
>  mm/zswap.c                                   |  5 +-
>  6 files changed, 46 insertions(+), 65 deletions(-)
> 
> diff --git a/Documentation/core-api/memory-allocation.rst b/Documentation/core-api/memory-allocation.rst
> index 648222b6c857..5069ac72b857 100644
> --- a/Documentation/core-api/memory-allocation.rst
> +++ b/Documentation/core-api/memory-allocation.rst
> @@ -9,7 +9,7 @@ allocate small chunks using `kmalloc` or `kmem_cache_alloc` families,
>  large virtually contiguous areas using `vmalloc` and its derivatives,
>  or you can directly request pages from the page allocator with
>  `alloc_pages`. It is also possible to use more specialized allocators,
> -for instance `cma_alloc` or `zs_malloc`.
> +for instance `cma_alloc` or `zs_store`.
>  
>  Most of the memory allocation APIs use GFP flags to express how that
>  memory should be allocated. The GFP acronym stands for "get free
> diff --git a/Documentation/mm/zsmalloc.rst b/Documentation/mm/zsmalloc.rst
> index dbb5c1a8dcc3..52eaba8f97e9 100644
> --- a/Documentation/mm/zsmalloc.rst
> +++ b/Documentation/mm/zsmalloc.rst
> @@ -20,9 +20,9 @@ For simplicity, zsmalloc can only allocate objects of size up to PAGE_SIZE
>  since this satisfies the requirements of all its current users (in the
>  worst case, page is incompressible and is thus stored "as-is" i.e. in
>  uncompressed form). For allocation requests larger than this size, failure
> -is returned (see zs_malloc).
> +is returned (see zs_store).
>  
> -Additionally, zs_malloc() does not return a dereferenceable pointer.
> +Additionally, zs_store() does not return a dereferenceable pointer.

Well, it doesn't return a pointer at all now.

>  Instead, it stores the encoded location of the allocated object in the
>  provided ``struct zs_handle``. The reason for this indirection is that
>  zsmalloc does not keep zspages permanently mapped since that would cause
[..]

> diff --git a/include/linux/zsmalloc.h b/include/linux/zsmalloc.h
> index 6c76fd10344e..674358f39191 100644
> --- a/include/linux/zsmalloc.h
> +++ b/include/linux/zsmalloc.h
> @@ -48,8 +48,8 @@ struct scatterlist;
>  struct zs_pool *zs_create_pool(const char *name);
>  void zs_destroy_pool(struct zs_pool *pool);
>  
> -int zs_malloc(struct zs_pool *pool, struct zs_handle *handle, size_t size,
> -	      gfp_t flags, const int nid);
> +int zs_store(struct zs_pool *pool, struct zs_handle *handle, const void *src,
> +	     size_t size, gfp_t flags, const int nid);
>  void zs_free(struct zs_pool *pool, struct zs_handle *handle);
>  #ifdef CONFIG_ZSMALLOC_REBIND_HANDLE
>  void zs_rebind_handle(struct zs_pool *pool, struct zs_handle *dst,
> @@ -68,8 +68,6 @@ void zs_pool_stats(struct zs_pool *pool, struct zs_pool_stats *stats);
>  void zs_obj_read_sg_begin(struct zs_pool *pool, struct zs_handle *handle,
>  			  struct scatterlist *sg, size_t mem_len);
>  void zs_obj_read_sg_end(struct zs_pool *pool, struct zs_handle *handle);
> -void zs_obj_write(struct zs_pool *pool, struct zs_handle *handle,
> -		  void *handle_mem, size_t mem_len);
>  
>  extern const struct movable_operations zsmalloc_mops;
>  
> diff --git a/mm/zsmalloc.c b/mm/zsmalloc.c
> index 71e6dcb1fb3f..ff1cde4c5dc2 100644
> --- a/mm/zsmalloc.c
> +++ b/mm/zsmalloc.c
> @@ -1172,26 +1172,15 @@ void zs_obj_read_sg_end(struct zs_pool *pool, struct zs_handle *handle)
>  }
>  EXPORT_SYMBOL_GPL(zs_obj_read_sg_end);
>  
> -void zs_obj_write(struct zs_pool *pool, struct zs_handle *handle,
> -		  void *handle_mem, size_t mem_len)
> +static void obj_write(struct size_class *class, struct zspage *zspage,
> +		      unsigned long obj, const void *handle_mem,
> +		      size_t mem_len)

Should handle_mem and mem_len be renamed now?

Also, there's a comment in zs_page_migrate() referring to zs_obj_read/write().

>  {
> -	struct zspage *zspage;
>  	struct zpdesc *zpdesc;
> -	unsigned long obj, off;
> +	unsigned long off;
>  	unsigned int obj_idx;
> -	struct size_class *class;
>  
> -	/* Guarantee we can get zspage from handle safely */
> -	read_lock(&pool->lock);
> -	obj = handle_to_obj(handle);
>  	obj_to_location(obj, &zpdesc, &obj_idx);
> -	zspage = get_zspage(zpdesc);
> -
> -	/* Make sure migration doesn't move any pages in this zspage */
> -	zspage_read_lock(zspage);
> -	read_unlock(&pool->lock);
> -
> -	class = zspage_class(pool, zspage);
>  	off = offset_in_page(class->size * obj_idx);
>  
>  	if (!ZsHugePage(zspage))
[..]
  
> @@ -1318,12 +1306,16 @@ int zs_malloc(struct zs_pool *pool, struct zs_handle *handle, size_t size,
>  	spin_lock(&class->lock);
>  	zspage = find_get_zspage(class);
>  	if (likely(zspage)) {
> -		obj_malloc(pool, zspage, handle);
> +		obj = obj_malloc(pool, zspage, handle);
>  		/* Now move the zspage to another fullness group, if required */
>  		fix_fullness_group(class, zspage);
>  		class_stat_add(class, ZS_OBJS_INUSE, 1);
> +		zspage_read_lock(zspage);

The zspage lock is now acquired under the class lock instead of the pool
lock, this is probably fine but worth calling out explicitly in the
changelog.

> +		spin_unlock(&class->lock);
>  
> -		goto out;
> +		obj_write(class, zspage, obj, src, size);
> +		zspage_read_unlock(zspage);
> +		return 0;
>  	}
>  
>  	spin_unlock(&class->lock);
> @@ -1332,8 +1324,14 @@ int zs_malloc(struct zs_pool *pool, struct zs_handle *handle, size_t size,
>  	if (!zspage)
>  		return -ENOMEM;
>  
> +	/*
> +	 * New zspage is not in class->fullness_list or movable yet, so we
> +	 * can allocate the object and copy data before taking class->lock.
> +	 */
> +	obj = obj_malloc(pool, zspage, handle);
> +	obj_write(class, zspage, obj, src, size);
> +
>  	spin_lock(&class->lock);
> -	obj_malloc(pool, zspage, handle);

Should moving obj_malloc() and obj_write() outside class->lock (and
skipping zspage_read_lock()) for new zspages be split into a separate
patch on top? AFAICT, it's a separate change from merging the
interfaces.

>  	newfg = get_fullness_group(class, zspage);
>  	insert_zspage(class, zspage, newfg);
>  	atomic_long_add(class->pages_per_zspage, &pool->pages_allocated);