staging: ion: make system and contig heaps modular

Add symmetric heap add/remove APIs and export symbols from
ion core that are needed to make both heaps modular.

Bug: 133508579
Test: ion-unit-test, rmmod, modprobe

Change-Id: Ic5f40936531936c9bfab48ea147108c019d28e2c
Signed-off-by: Sandeep Patil <sspatil@google.com>

drivers/staging/android/ion/heaps/Kconfig

@@ -1,13 +1,13 @@
 # SPDX-License-Identifier: GPL-2.0
 config ION_SYSTEM_HEAP
-	bool "Ion system heap"
+	tristate "Ion system heap"
 	depends on ION
 	help
 	  Choose this option to enable the Ion system heap. The system heap
 	  is backed by pages from the buddy allocator. If in doubt, say Y.
 
 config ION_SYSTEM_CONTIG_HEAP
-	bool "Ion system contig heap"
+	tristate "Ion system contig heap"
 	depends on ION
 	help
 	  Choose this option to enable Ion system contig heap. The system contig heap

drivers/staging/android/ion/heaps/Makefile

@@ -1,4 +1,6 @@
 # SPDX-License-Identifier: GPL-2.0
-obj-$(CONFIG_ION_SYSTEM_HEAP) += ion_system_heap.o ion_page_pool.o
+obj-$(CONFIG_ION_SYSTEM_HEAP) += ion_sys_heap.o
+ion_sys_heap-y := ion_system_heap.o ion_page_pool.o
+
 obj-$(CONFIG_ION_SYSTEM_CONTIG_HEAP) += ion_system_contig_heap.o
 obj-$(CONFIG_ION_CMA_HEAP) += ion_cma_heap.o

drivers/staging/android/ion/heaps/ion_system_contig_heap.c

@@ -11,6 +11,7 @@
 #include <linux/highmem.h>
 #include <linux/ion.h>
 #include <linux/mm.h>
+#include <linux/module.h>
 #include <linux/scatterlist.h>
 #include <linux/slab.h>
 
@@ -81,31 +82,22 @@ static struct ion_heap_ops kmalloc_ops = {
 	.map_user = ion_heap_map_user,
 };
 
-static struct ion_heap *__ion_system_contig_heap_create(void)
+static struct ion_heap contig_heap = {
+	.ops = &kmalloc_ops,
+	.type = ION_HEAP_TYPE_SYSTEM_CONTIG,
+	.name = "ion_system_contig_heap",
+};
+
+static int __init ion_system_contig_heap_init(void)
 {
-	struct ion_heap *heap;
-
-	heap = kzalloc(sizeof(*heap), GFP_KERNEL);
-	if (!heap)
-		return ERR_PTR(-ENOMEM);
-	heap->ops = &kmalloc_ops;
-	heap->type = ION_HEAP_TYPE_SYSTEM_CONTIG;
-	heap->name = "ion_system_contig_heap";
-
-	return heap;
+	return ion_device_add_heap(&contig_heap);
 }
 
-static int ion_system_contig_heap_create(void)
+static void __exit ion_system_contig_heap_exit(void)
 {
-	struct ion_heap *heap;
-
-	heap = __ion_system_contig_heap_create();
-	if (IS_ERR(heap))
-		return PTR_ERR(heap);
-
-	ion_device_add_heap(heap);
-
-	return 0;
+	ion_device_remove_heap(&contig_heap);
 }
-device_initcall(ion_system_contig_heap_create);
 
+module_init(ion_system_contig_heap_init);
+module_exit(ion_system_contig_heap_exit);
+MODULE_LICENSE("GPL v2");

drivers/staging/android/ion/heaps/ion_system_heap.c

@@ -11,6 +11,7 @@
 #include <linux/highmem.h>
 #include <linux/ion.h>
 #include <linux/mm.h>
+#include <linux/module.h>
 #include <linux/scatterlist.h>
 #include <linux/slab.h>
 #include <linux/vmalloc.h>
@@ -204,15 +205,6 @@ static int ion_system_heap_shrink(struct ion_heap *heap, gfp_t gfp_mask,
 	return nr_total;
 }
 
-static struct ion_heap_ops system_heap_ops = {
-	.allocate = ion_system_heap_allocate,
-	.free = ion_system_heap_free,
-	.map_kernel = ion_heap_map_kernel,
-	.unmap_kernel = ion_heap_unmap_kernel,
-	.map_user = ion_heap_map_user,
-	.shrink = ion_system_heap_shrink,
-};
-
 static void ion_system_heap_destroy_pools(struct ion_page_pool **pools)
 {
 	int i;
@@ -246,38 +238,39 @@ err_create_pool:
 	return -ENOMEM;
 }
 
-static struct ion_heap *__ion_system_heap_create(void)
+static struct ion_heap_ops system_heap_ops = {
+	.allocate = ion_system_heap_allocate,
+	.free = ion_system_heap_free,
+	.map_kernel = ion_heap_map_kernel,
+	.unmap_kernel = ion_heap_unmap_kernel,
+	.map_user = ion_heap_map_user,
+	.shrink = ion_system_heap_shrink,
+};
+
+struct ion_system_heap system_heap = {
+	.heap = {
+		.ops = &system_heap_ops,
+		.type = ION_HEAP_TYPE_SYSTEM,
+		.flags = ION_HEAP_FLAG_DEFER_FREE,
+		.name = "ion_system_heap",
+	}
+};
+
+static int __init ion_system_heap_init(void)
 {
-	struct ion_system_heap *heap;
+	int ret = ion_system_heap_create_pools(system_heap.pools);
+	if (ret)
+		return ret;
 
-	heap = kzalloc(sizeof(*heap), GFP_KERNEL);
-	if (!heap)
-		return ERR_PTR(-ENOMEM);
-	heap->heap.ops = &system_heap_ops;
-	heap->heap.type = ION_HEAP_TYPE_SYSTEM;
-	heap->heap.flags = ION_HEAP_FLAG_DEFER_FREE;
-
-	if (ion_system_heap_create_pools(heap->pools))
-		goto free_heap;
-
-	return &heap->heap;
-
-free_heap:
-	kfree(heap);
-	return ERR_PTR(-ENOMEM);
+	return ion_device_add_heap(&system_heap.heap);
 }
 
-static int ion_system_heap_create(void)
+static void __exit ion_system_heap_exit(void)
 {
-	struct ion_heap *heap;
-
-	heap = __ion_system_heap_create();
-	if (IS_ERR(heap))
-		return PTR_ERR(heap);
-	heap->name = "ion_system_heap";
-
-	ion_device_add_heap(heap);
-
-	return 0;
+	ion_device_remove_heap(&system_heap.heap);
+	ion_system_heap_destroy_pools(system_heap.pools);
 }
-device_initcall(ion_system_heap_create);
+
+module_init(ion_system_heap_init);
+module_exit(ion_system_heap_exit);
+MODULE_LICENSE("GPL v2");

drivers/staging/android/ion/ion.c

@@ -222,28 +222,36 @@ static int debug_shrink_get(void *data, u64 *val)
 DEFINE_SIMPLE_ATTRIBUTE(debug_shrink_fops, debug_shrink_get,
 			debug_shrink_set, "%llu\n");
 
-void ion_device_add_heap(struct ion_heap *heap)
+int __ion_device_add_heap(struct ion_heap *heap, struct module *owner)
 {
 	struct ion_device *dev = internal_dev;
 	int ret;
 	struct dentry *heap_root;
 	char debug_name[64];
 
-	if (!heap->ops->allocate || !heap->ops->free)
-		pr_err("%s: can not add heap with invalid ops struct.\n",
-		       __func__);
+	if (!heap || !heap->ops || !heap->ops->allocate || !heap->ops->free) {
+		pr_err("%s: invalid heap or heap_ops\n", __func__);
+		ret = -EINVAL;
+		goto out;
+	}
 
+	heap->owner = owner;
 	spin_lock_init(&heap->free_lock);
 	spin_lock_init(&heap->stat_lock);
 	heap->free_list_size = 0;
 
-	if (heap->flags & ION_HEAP_FLAG_DEFER_FREE)
-		ion_heap_init_deferred_free(heap);
+	if (heap->flags & ION_HEAP_FLAG_DEFER_FREE) {
+		ret = ion_heap_init_deferred_free(heap);
+		if (ret)
+			goto out_heap_cleanup;
+	}
 
 	if ((heap->flags & ION_HEAP_FLAG_DEFER_FREE) || heap->ops->shrink) {
 		ret = ion_heap_init_shrinker(heap);
-		if (ret)
+		if (ret) {
 			pr_err("%s: Failed to register shrinker\n", __func__);
+			goto out_heap_cleanup;
+		}
 	}
 
 	heap->num_of_buffers = 0;
@@ -273,6 +281,7 @@ void ion_device_add_heap(struct ion_heap *heap)
 				    &debug_shrink_fops);
 	}
 
+	heap->debugfs_dir = heap_root;
 	down_write(&dev->lock);
 	heap->id = heap_id++;
 	/*
@@ -284,8 +293,37 @@ void ion_device_add_heap(struct ion_heap *heap)
 
 	dev->heap_cnt++;
 	up_write(&dev->lock);
+
+	return 0;
+
+out_heap_cleanup:
+	ion_heap_cleanup(heap);
+out:
+	return ret;
 }
-EXPORT_SYMBOL(ion_device_add_heap);
+EXPORT_SYMBOL_GPL(__ion_device_add_heap);
+
+void ion_device_remove_heap(struct ion_heap *heap)
+{
+	struct ion_device *dev = internal_dev;
+
+	if (!heap) {
+		pr_err("%s: Invalid argument\n", __func__);
+		return;
+	}
+
+	// take semaphore and remove the heap from dev->heap list
+	down_write(&dev->lock);
+	/* So no new allocations can happen from this heap */
+	plist_del(&heap->node, &dev->heaps);
+	if (ion_heap_cleanup(heap) != 0) {
+		pr_warn("%s: failed to cleanup heap (%s)\n",
+			__func__, heap->name);
+	}
+	debugfs_remove_recursive(heap->debugfs_dir);
+	up_write(&dev->lock);
+}
+EXPORT_SYMBOL(ion_device_remove_heap);
 
 static int ion_device_create(void)
 {

drivers/staging/android/ion/ion_buffer.c

@@ -6,6 +6,7 @@
  */
 
 #include <linux/mm.h>
+#include <linux/module.h>
 #include <linux/slab.h>
 
 #include "ion_private.h"
@@ -76,6 +77,14 @@ static struct ion_buffer *ion_buffer_create(struct ion_heap *heap,
 	heap->num_of_alloc_bytes += len;
 	if (heap->num_of_alloc_bytes > heap->alloc_bytes_wm)
 		heap->alloc_bytes_wm = heap->num_of_alloc_bytes;
+	if (heap->num_of_buffers == 1) {
+		/* This module reference lasts as long as at least one
+		 * buffer is allocated from the heap. We are protected
+		 * against ion_device_remove_heap() with dev->lock, so we can
+		 * safely assume the module reference is going to* succeed.
+		 */
+		__module_get(heap->owner);
+	}
 	spin_unlock(&heap->stat_lock);
 
 	INIT_LIST_HEAD(&buffer->attachments);
@@ -184,6 +193,7 @@ int ion_buffer_zero(struct ion_buffer *buffer)
 
 	return ion_sglist_zero(table->sgl, table->nents, pgprot);
 }
+EXPORT_SYMBOL_GPL(ion_buffer_zero);
 
 void ion_buffer_release(struct ion_buffer *buffer)
 {
@@ -196,7 +206,10 @@ void ion_buffer_release(struct ion_buffer *buffer)
 	spin_lock(&buffer->heap->stat_lock);
 	buffer->heap->num_of_buffers--;
 	buffer->heap->num_of_alloc_bytes -= buffer->size;
+	if (buffer->heap->num_of_buffers == 0)
+		module_put(buffer->heap->owner);
 	spin_unlock(&buffer->heap->stat_lock);
+	/* drop reference to the heap module */
 
 	kfree(buffer);
 }

drivers/staging/android/ion/ion_heap.c

@@ -101,12 +101,15 @@ static int ion_heap_deferred_free(void *data)
 		struct ion_buffer *buffer;
 
 		wait_event_freezable(heap->waitqueue,
-				     ion_heap_freelist_size(heap) > 0);
+				     (ion_heap_freelist_size(heap) > 0 ||
+				      kthread_should_stop()));
 
 		spin_lock(&heap->free_lock);
 		if (list_empty(&heap->free_list)) {
 			spin_unlock(&heap->free_lock);
+			if (!kthread_should_stop())
 				continue;
+			break;
 		}
 		buffer = list_first_entry(&heap->free_list, struct ion_buffer,
 					  list);
@@ -156,12 +159,14 @@ void *ion_heap_map_kernel(struct ion_heap *heap,
 
 	return vaddr;
 }
+EXPORT_SYMBOL_GPL(ion_heap_map_kernel);
 
 void ion_heap_unmap_kernel(struct ion_heap *heap,
 			   struct ion_buffer *buffer)
 {
 	vunmap(buffer->vaddr);
 }
+EXPORT_SYMBOL_GPL(ion_heap_unmap_kernel);
 
 int ion_heap_map_user(struct ion_heap *heap, struct ion_buffer *buffer,
 		      struct vm_area_struct *vma)
@@ -198,6 +203,7 @@ int ion_heap_map_user(struct ion_heap *heap, struct ion_buffer *buffer,
 
 	return 0;
 }
+EXPORT_SYMBOL_GPL(ion_heap_map_user);
 
 void ion_heap_freelist_add(struct ion_heap *heap, struct ion_buffer *buffer)
 {
@@ -256,3 +262,32 @@ int ion_heap_init_shrinker(struct ion_heap *heap)
 
 	return register_shrinker(&heap->shrinker);
 }
+
+int ion_heap_cleanup(struct ion_heap *heap)
+{
+	int ret;
+
+	if (heap->flags & ION_HEAP_FLAG_DEFER_FREE &&
+	    !IS_ERR_OR_NULL(heap->task)) {
+		size_t free_list_size = ion_heap_freelist_size(heap);
+		size_t total_drained = ion_heap_freelist_drain(heap, 0);
+
+		if (total_drained != free_list_size) {
+			pr_err("%s: %s heap drained %zu bytes, requested %zu\n",
+			       __func__, heap->name, free_list_size,
+			       total_drained);
+			return -EBUSY;
+		}
+		ret = kthread_stop(heap->task);
+		if (ret < 0) {
+			pr_err("%s: failed to stop heap free thread\n",
+			       __func__);
+			return ret;
+		}
+	}
+
+	if ((heap->flags & ION_HEAP_FLAG_DEFER_FREE) || heap->ops->shrink)
+		unregister_shrinker(&heap->shrinker);
+
+	return 0;
+}

drivers/staging/android/ion/ion_private.h

@@ -51,4 +51,7 @@ extern struct dma_buf *ion_dmabuf_alloc(struct ion_device *dev, size_t len,
 					unsigned int flags);
 extern int ion_free(struct ion_buffer *buffer);
 
+/* ion heap helpers */
+extern int ion_heap_cleanup(struct ion_heap *heap);
+
 #endif /* _ION_PRIVATE_H */

include/linux/ion.h

@@ -12,6 +12,7 @@
 #include <linux/dma-direction.h>
 #include <linux/kref.h>
 #include <linux/mm_types.h>
+#include <linux/module.h>
 #include <linux/mutex.h>
 #include <linux/rbtree.h>
 #include <linux/sched.h>
@@ -105,6 +106,7 @@ struct ion_heap_ops {
  *			allocating.  These are specified by platform data and
  *			MUST be unique
  * @name:		used for debugging
+ * @owner:		kernel module that implements this heap
  * @shrinker:		a shrinker for the heap
  * @free_list:		free list head if deferred free is used
  * @free_list_size	size of the deferred free list in bytes
@@ -127,6 +129,7 @@ struct ion_heap {
 	unsigned long flags;
 	unsigned int id;
 	const char *name;
+	struct module *owner;
 
 	/* deferred free support */
 	struct shrinker shrinker;
@@ -143,16 +146,30 @@ struct ion_heap {
 
 	/* protect heap statistics */
 	spinlock_t stat_lock;
+
+	/* heap's debugfs root */
+	struct dentry *debugfs_dir;
 };
 
+#define ion_device_add_heap(heap) __ion_device_add_heap(heap, THIS_MODULE)
+
 #ifdef CONFIG_ION
 
 /**
- * ion_device_add_heap - adds a heap to the ion device
+ * __ion_device_add_heap - adds a heap to the ion device
  *
  * @heap:               the heap to add
+ *
+ * Returns 0 on success, negative error otherwise.
  */
-void ion_device_add_heap(struct ion_heap *heap);
+int __ion_device_add_heap(struct ion_heap *heap, struct module *owner);
+
+/**
+ * ion_device_remove_heap - removes a heap from ion device
+ *
+ * @heap:		pointer to the heap to be removed
+ */
+void ion_device_remove_heap(struct ion_heap *heap);
 
 /**
  * ion_heap_init_shrinker
@@ -283,7 +300,8 @@ struct dma_buf *ion_alloc(size_t len, unsigned int heap_id_mask,
 
 #else
 
-static inline int ion_device_add_heap(struct ion_heap *heap)
+static inline int __ion_device_add_heap(struct ion_heap *heap,
+				      struct module *owner)
 {
 	return -ENODEV;
 }