rust: allocator: Vmalloc: Support alignments larger than PAGE_SIZE

From: Hui Zhu <[email protected]>

This commit add new function vrealloc_align. vrealloc_align support allocation of aligned vmap pages with __vmalloc_node_noprof.
And vrealloc_align will check the old address. If this address does not meet the current alignment requirements, it will also release the old vmap pages and reallocate new vmap pages that satisfy the alignment requirements.

Co-developed-by Geliang Tang <[email protected]>

Signed-off-by Geliang Tang <[email protected]>

Signed-off-by Hui Zhu <[email protected]>

---
 include/linux/vmalloc.h |  5 +++
 mm/vmalloc.c            | 80 ++++++++++++++++++++++++++---------------
 2 files changed, 57 insertions(+), 28 deletions(-)

include/linux/vmalloc.h +5 -0

index fdc9aeb74a44..0ce0c1ea2427 100644
@@ -201,6 +201,11 @@ void * __must_check vrealloc_noprof(const void *p, size_t size, gfp_t flags)
 		__realloc_size(2);
 #define vrealloc(...)		alloc_hooks(vrealloc_noprof(__VA_ARGS__))
 
+void * __must_check vrealloc_align_noprof(const void *p, size_t size,
+					  size_t align, gfp_t flags)
+		__realloc_size(2);
+#define vrealloc_align(...)	alloc_hooks(vrealloc_align_noprof(__VA_ARGS__))
+
 extern void vfree(const void *addr);
 extern void vfree_atomic(const void *addr);
 

mm/vmalloc.c +52 -28

index ab986dd09b6a..41cb3603b3cc 100644
@@ -4081,9 +4081,11 @@ void *vzalloc_node_noprof(unsigned long size, int node)
 EXPORT_SYMBOL(vzalloc_node_noprof);
 
 /**
- * vrealloc - reallocate virtually contiguous memory; contents remain unchanged
+ * vrealloc_align - reallocate virtually contiguous memory;
+ *                  contents remain unchanged
  * @p: object to reallocate memory for
  * @size: the size to reallocate
+ * @align: requested alignment
  * @flags: the flags for the page level allocator
  *
  * If @p is %NULL, vrealloc() behaves exactly like vmalloc(). If @size is 0 and

@@ -4103,7 +4105,8 @@ EXPORT_SYMBOL(vzalloc_node_noprof);
  * Return: pointer to the allocated memory; %NULL if @size is zero or in case of
  *         failure
  */
-void *vrealloc_noprof(const void *p, size_t size, gfp_t flags)
+void *vrealloc_align_noprof(const void *p, size_t size, size_t align,
+			    gfp_t flags)
 {
 	struct vm_struct *vm = NULL;
 	size_t alloced_size = 0;

@@ -4116,49 +4119,65 @@ void *vrealloc_noprof(const void *p, size_t size, gfp_t flags)
 	}
 
 	if (p) {
+		if (!is_power_of_2(align)) {
+			WARN(1, "Trying to vrealloc_align() align is not power of 2 (%ld)\n",
+			     align);
+			return NULL;
+		}
+
 		vm = find_vm_area(p);
 		if (unlikely(!vm)) {
-			WARN(1, "Trying to vrealloc() nonexistent vm area (%p)\n", p);
+			WARN(1, "Trying to vrealloc_align() nonexistent vm area (%p)\n", p);
 			return NULL;
 		}
 
 		alloced_size = get_vm_area_size(vm);
 		old_size = vm->requested_size;
 		if (WARN(alloced_size < old_size,
-			 "vrealloc() has mismatched area vs requested sizes (%p)\n", p))
+			 "vrealloc_align() has mismatched area vs requested sizes (%p)\n", p))
 			return NULL;
 	}
 
-	/*
-	 * TODO: Shrink the vm_area, i.e. unmap and free unused pages. What
-	 * would be a good heuristic for when to shrink the vm_area?
-	 */
-	if (size <= old_size) {
-		/* Zero out "freed" memory, potentially for future realloc. */
-		if (want_init_on_free() || want_init_on_alloc(flags))
-			memset((void *)p + size, 0, old_size - size);
-		vm->requested_size = size;
-		kasan_poison_vmalloc(p + size, old_size - size);
-		return (void *)p;
-	}
+	if (IS_ALIGNED((unsigned long)p, align)) {
+		/*
+		 * TODO: Shrink the vm_area, i.e. unmap and free unused pages. What
+		 * would be a good heuristic for when to shrink the vm_area?
+		 */
+		if (size <= old_size) {
+			/* Zero out "freed" memory, potentially for future realloc. */
+			if (want_init_on_free() || want_init_on_alloc(flags))
+				memset((void *)p + size, 0, old_size - size);
+			vm->requested_size = size;
+			kasan_poison_vmalloc(p + size, old_size - size);
+			return (void *)p;
+		}
 
-	/*
-	 * We already have the bytes available in the allocation; use them.
-	 */
-	if (size <= alloced_size) {
-		kasan_unpoison_vmalloc(p + old_size, size - old_size,
-				       KASAN_VMALLOC_PROT_NORMAL);
 		/*
-		 * No need to zero memory here, as unused memory will have
-		 * already been zeroed at initial allocation time or during
-		 * realloc shrink time.
+		 * We already have the bytes available in the allocation; use them.
+		 */
+		if (size <= alloced_size) {
+			kasan_unpoison_vmalloc(p + old_size, size - old_size,
+					KASAN_VMALLOC_PROT_NORMAL);
+			/*
+			 * No need to zero memory here, as unused memory will have
+			 * already been zeroed at initial allocation time or during
+			 * realloc shrink time.
+			 */
+			vm->requested_size = size;
+			return (void *)p;
+		}
+	} else {
+		/*
+		 * p is not aligned with align.
+		 * Allocate a new address to handle it.
 		 */
-		vm->requested_size = size;
-		return (void *)p;
+		if (size < old_size)
+			old_size = size;
 	}
 
 	/* TODO: Grow the vm_area, i.e. allocate and map additional pages. */
-	n = __vmalloc_noprof(size, flags);
+	n = __vmalloc_node_noprof(size, align, flags, NUMA_NO_NODE,
+				  __builtin_return_address(0));
 	if (!n)
 		return NULL;
 

@@ -4170,6 +4189,11 @@ void *vrealloc_noprof(const void *p, size_t size, gfp_t flags)
 	return n;
 }
 
+void *vrealloc_noprof(const void *p, size_t size, gfp_t flags)
+{
+	return vrealloc_align_noprof(p, size, 1, flags);
+}
+
 #if defined(CONFIG_64BIT) && defined(CONFIG_ZONE_DMA32)
 #define GFP_VMALLOC32 (GFP_DMA32 | GFP_KERNEL)
 #elif defined(CONFIG_64BIT) && defined(CONFIG_ZONE_DMA)

Signature · 1 line

2.43.0

samples/rust/rust_allocator.rs new file +13 -0

+
+            let (addr, is_ok) = check_ptr(ptr, size, align);
+            if !is_ok {
+                clear_vmalloc_vec(&vmalloc_vec);
+                return Err(EINVAL);
+            }
+
+            vmalloc_vec.push((addr, layout), GFP_KERNEL)?;
+        }
+
+        Ok(RustAllocator { vmalloc_vec })
+    }
+}

samples/rust/rust_allocator.rs line 83

Quoted 7 lines · Hui Zhu

samples/rust/rust_allocator.rs new file +7 -0

+impl kernel::Module for RustAllocator {
+    fn init(_module: &'static ThisModule) -> Result<Self> {
+        pr_info!("Rust allocator sample (init)\n");
+
+        let mut vmalloc_vec = KVec::new();
+        for (size, align) in VMALLOC_ARG {
+            let (ptr, layout) = vmalloc_align(size, align)?;

Ok, I think I get the idea, you want to demonstrate how to use the Allocator trait for raw memory allocations.

However, doing so is discouraged unless there's really no other way. One obvious example are Rust's own memory allocation primitives, such as Box and Vec.

So, instead of this raw allocation, you can just use VBox::new() or VBox::new_uninit() in the following way.

	[repr(align(ALIGN))]
	struct Blob([u8; SIZE]);

	// Creates a vmalloc allocation of size `SIZE` with an alignment of
	// `ALIGN`. The allocation is freed once `b` is dropped.
	let b = VBox::<Blob>::new_uninit(GFP_KERNEL)?;

This way you don't have to handle the layout and the Allocator type yourself and you also don't have to care about explicitly calling vfree(), VBox does all this for you.

samples/rust/rust_allocator.rs line 96

Quoted 13 lines

samples/rust/rust_allocator.rs new file +13 -0

+
+            let (addr, is_ok) = check_ptr(ptr, size, align);
+            if !is_ok {
+                clear_vmalloc_vec(&vmalloc_vec);
+                return Err(EINVAL);
+            }
+
+            vmalloc_vec.push((addr, layout), GFP_KERNEL)?;
+        }
+
+        Ok(RustAllocator { vmalloc_vec })
+    }
+}

From: Hui Zhu <[email protected]>

There is a TODO in Vmalloc::realloc "Support alignments larger than PAGE_SIZE."

These commits make allocator vmalloc support alignments larger than PAGE_SIZE.
The function vrealloc_align is added to vmalloc.c to support reallocating aligned vmap pages.
When Vmalloc::realloc intends to reallocate memory aligned beyond PAGE_SIZE, vrealloc_align should be used instead of vrealloc, thus enabling support for alignments larger than PAGE_SIZE.

And add a sample to the samples memory allocator usage.