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

From: Sergey Senozhatsky

Date: Wed Nov 04 2026 - 23:48:53 EST


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@xxxxxxxxxxxx>
---
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.
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/drivers/block/zram/zram_drv.c b/drivers/block/zram/zram_drv.c
index 7c80db51b29a..223e4151fc9d 100644
--- a/drivers/block/zram/zram_drv.c
+++ b/drivers/block/zram/zram_drv.c
@@ -2262,9 +2262,11 @@ static int write_incompressible_page(struct zram *zram, struct page *page,

slot_lock(zram, index);
slot_free(zram, index);
- ret = zs_malloc(zram->mem_pool, &zram->table[index].handle, PAGE_SIZE,
- GFP_NOIO | __GFP_NOWARN |
- __GFP_HIGHMEM | __GFP_MOVABLE, page_to_nid(page));
+ src = kmap_local_page(page);
+ ret = zs_store(zram->mem_pool, &zram->table[index].handle, src,
+ PAGE_SIZE, GFP_NOIO | __GFP_NOWARN |
+ __GFP_HIGHMEM | __GFP_MOVABLE, page_to_nid(page));
+ kunmap_local(src);
if (ret) {
slot_unlock(zram, index);
return ret;
@@ -2276,11 +2278,6 @@ static int write_incompressible_page(struct zram *zram, struct page *page,
return -ENOMEM;
}

- src = kmap_local_page(page);
- zs_obj_write(zram->mem_pool, &zram->table[index].handle, src,
- PAGE_SIZE);
- kunmap_local(src);
-
set_slot_flag(zram, index, ZRAM_HUGE);
set_slot_size(zram, index, PAGE_SIZE);
mark_slot_accessed(zram, index);
@@ -2329,26 +2326,22 @@ static int zram_write_page(struct zram *zram, struct page *page,

slot_lock(zram, index);
slot_free(zram, index);
- ret = zs_malloc(zram->mem_pool, &zram->table[index].handle, comp_len,
- GFP_NOIO | __GFP_NOWARN |
- __GFP_HIGHMEM | __GFP_MOVABLE, page_to_nid(page));
+ ret = zs_store(zram->mem_pool, &zram->table[index].handle,
+ zstrm->buffer, comp_len,
+ GFP_NOIO | __GFP_NOWARN |
+ __GFP_HIGHMEM | __GFP_MOVABLE, page_to_nid(page));
+ zcomp_cstrm_put(zstrm);
if (ret) {
- zcomp_cstrm_put(zstrm);
slot_unlock(zram, index);
return ret;
}

if (!zram_can_store_page(zram)) {
- zcomp_cstrm_put(zstrm);
zs_free(zram->mem_pool, &zram->table[index].handle);
slot_unlock(zram, index);
return -ENOMEM;
}

- zs_obj_write(zram->mem_pool, &zram->table[index].handle, zstrm->buffer,
- comp_len);
- zcomp_cstrm_put(zstrm);
-
set_slot_size(zram, index, comp_len);
mark_slot_accessed(zram, index);
slot_unlock(zram, index);
@@ -2533,18 +2526,12 @@ static int recompress_slot(struct zram *zram, unsigned long index,
* modify zsmalloc API to return this information. For now, we will
* make do with the node of the page allocated for recompression.
*/
- ret = zs_malloc(zram->mem_pool, &zram->recomp_handle, comp_len_new,
- GFP_NOIO | __GFP_NOWARN |
- __GFP_HIGHMEM | __GFP_MOVABLE,
- page_to_nid(page));
- if (ret) {
- zcomp_cstrm_put(zstrm);
- return ret;
- }
-
- zs_obj_write(zram->mem_pool, &zram->recomp_handle, zstrm->buffer,
- comp_len_new);
+ ret = zs_store(zram->mem_pool, &zram->recomp_handle, zstrm->buffer,
+ comp_len_new, GFP_NOIO | __GFP_NOWARN |
+ __GFP_HIGHMEM | __GFP_MOVABLE, page_to_nid(page));
zcomp_cstrm_put(zstrm);
+ if (ret)
+ return ret;

slot_free(zram, index);
zs_rebind_handle(zram->mem_pool, &zram->table[index].handle,
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)
{
- 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))
@@ -1216,10 +1205,7 @@ void zs_obj_write(struct zs_pool *pool, struct zs_handle *handle,
memcpy_to_page(zpdesc_page(zpdesc), 0,
handle_mem + sizes[0], sizes[1]);
}
-
- zspage_read_unlock(zspage);
}
-EXPORT_SYMBOL_GPL(zs_obj_write);

/**
* zs_huge_class_size() - Returns the size (in bytes) of the first huge
@@ -1284,9 +1270,10 @@ static unsigned long obj_malloc(struct zs_pool *pool,


/**
- * zs_malloc - Allocate block of given size from pool.
+ * zs_store() - Allocate block of given size from pool and copy data into it.
* @pool: pool to allocate from
* @handle: handle to store allocated object in (must not be on the stack)
+ * @src: data to copy into the allocated object
* @size: size of block to allocate
* @gfp: gfp flags when allocating object
* @nid: The preferred node id to allocate new zspage (if needed)
@@ -1295,10 +1282,11 @@ static unsigned long obj_malloc(struct zs_pool *pool,
*
* Return: 0 on success, negative error code otherwise.
*/
-int zs_malloc(struct zs_pool *pool, struct zs_handle *handle, size_t size,
- gfp_t gfp, const int nid)
+int zs_store(struct zs_pool *pool, struct zs_handle *handle, const void *src,
+ size_t size, gfp_t gfp, const int nid)
{
struct size_class *class;
+ unsigned long obj;
int newfg;
struct zspage *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);
+ 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);
newfg = get_fullness_group(class, zspage);
insert_zspage(class, zspage, newfg);
atomic_long_add(class->pages_per_zspage, &pool->pages_allocated);
@@ -1342,12 +1340,11 @@ int zs_malloc(struct zs_pool *pool, struct zs_handle *handle, size_t size,

/* We completely set up zspage so mark them as movable */
SetZsPageMovable(pool, zspage);
-out:
spin_unlock(&class->lock);

return 0;
}
-EXPORT_SYMBOL_GPL(zs_malloc);
+EXPORT_SYMBOL_GPL(zs_store);

static void obj_free(int class_size, unsigned long obj)
{
@@ -2224,7 +2221,7 @@ struct zs_pool *zs_create_pool(const char *name)
huge_class_size = size;
/*
* The object uses ZS_HANDLE_SIZE bytes to store the
- * handle. We need to subtract it, because zs_malloc()
+ * handle. We need to subtract it, because zs_store()
* unconditionally adds handle size before it performs
* size class search - so object may be smaller than
* huge class size, yet it still can end up in the huge
diff --git a/mm/zswap.c b/mm/zswap.c
index f28a46cae1e1..f4cfc552a6fc 100644
--- a/mm/zswap.c
+++ b/mm/zswap.c
@@ -948,12 +948,11 @@ static bool zswap_compress(struct folio *folio, long index,
}

gfp = GFP_NOWAIT | __GFP_NORETRY | __GFP_HIGHMEM | __GFP_MOVABLE;
- alloc_ret = zs_malloc(pool->zs_pool, &entry->handle, dlen, gfp,
- folio_nid(folio));
+ alloc_ret = zs_store(pool->zs_pool, &entry->handle, dst, dlen, gfp,
+ folio_nid(folio));
if (alloc_ret)
goto unlock;

- zs_obj_write(pool->zs_pool, &entry->handle, dst, dlen);
entry->length = dlen;

unlock:
--
2.56.0.398.g882cbad275-goog