llext: add reference counting

Extend the llext_load() / llext_unload() API to let it be called
repeatedly for the same extension to increment or decrement its
reference counter respectively. We use a mutex to protect the counter
and make both llext_load() and llext_unload() return the use-count to
let the caller identify when the first loading and the last unloading
took place.

Signed-off-by: Guennadi Liakhovetski <guennadi.liakhovetski@linux.intel.com>

include/zephyr/llext/llext.h

@@ -62,6 +62,9 @@ struct llext {
 
 	/** Exported symbols from the llext, may be linked against by other llext */
 	struct llext_symtable sym_tab;
+
+	/** Extension use counter, prevents unloading while in use */
+	unsigned int use_count;
 };
 
 /**
@@ -115,7 +118,7 @@ int llext_load(struct llext_loader *loader, const char *name, struct llext **ext
  *
  * @param[in] ext Extension to unload
  */
-void llext_unload(struct llext *ext);
+int llext_unload(struct llext **ext);
 
 /**
  * @brief Find the address for an arbitrary symbol name.

subsys/llext/llext.c

@@ -807,12 +807,25 @@ out:
 	return ret;
 }
 
+static struct k_mutex llext_lock = Z_MUTEX_INITIALIZER(llext_lock);
+
 int llext_load(struct llext_loader *ldr, const char *name, struct llext **ext,
 	       struct llext_load_param *ldr_parm)
 {
 	int ret;
 	elf_ehdr_t ehdr;
 
+	k_mutex_lock(&llext_lock, K_FOREVER);
+
+	if (*ext) {
+		/* The use count is at least 1 */
+		ret = (*ext)->use_count++;
+		k_mutex_unlock(&llext_lock);
+		return ret;
+	}
+
+	k_mutex_unlock(&llext_lock);
+
 	ret = llext_seek(ldr, 0);
 	if (ret != 0) {
 		LOG_ERR("Failed to seek for ELF header");
@@ -848,40 +861,58 @@ int llext_load(struct llext_loader *ldr, const char *name, struct llext **ext,
 
 		ldr->hdr = ehdr;
 		ret = do_llext_load(ldr, *ext, ldr_parm);
+		if (ret < 0)
+			return ret;
+
+		strncpy((*ext)->name, name, sizeof((*ext)->name));
+		(*ext)->name[sizeof((*ext)->name) - 1] = '\0';
+		(*ext)->use_count++;
+
+		sys_slist_append(&_llext_list, &(*ext)->_llext_list);
+		LOG_INF("Loaded extension %s", (*ext)->name);
+
 		break;
 	default:
 		LOG_ERR("Unsupported elf file type %x", ehdr.e_type);
-		*ext = NULL;
 		return -EINVAL;
 	}
 
-	if (ret == 0) {
-		strncpy((*ext)->name, name, sizeof((*ext)->name));
-		(*ext)->name[sizeof((*ext)->name) - 1] = '\0';
-		sys_slist_append(&_llext_list, &(*ext)->_llext_list);
-		LOG_INF("Loaded extension %s", (*ext)->name);
+	return 0;
 }
 
+int llext_unload(struct llext **ext)
+{
+	__ASSERT(*ext, "Expected non-null extension");
+	struct llext *tmp = *ext;
+
+	k_mutex_lock(&llext_lock, K_FOREVER);
+	__ASSERT(tmp->use_count, "A valid LLEXT cannot have a zero use-count!");
+
+	if (tmp->use_count-- != 1) {
+		unsigned int ret = tmp->use_count;
+
+		k_mutex_unlock(&llext_lock);
 		return ret;
 	}
 
-void llext_unload(struct llext *ext)
-{
-	__ASSERT(ext, "Expected non-null extension");
+	/* FIXME: protect the global list */
+	sys_slist_find_and_remove(&_llext_list, &tmp->_llext_list);
 
-	sys_slist_find_and_remove(&_llext_list, &ext->_llext_list);
+	*ext = NULL;
+	k_mutex_unlock(&llext_lock);
 
 	for (int i = 0; i < LLEXT_MEM_COUNT; i++) {
-		if (ext->mem_on_heap[i]) {
+		if (tmp->mem_on_heap[i]) {
 			LOG_DBG("freeing memory region %d", i);
-			k_heap_free(&llext_heap, ext->mem[i]);
-			ext->mem[i] = NULL;
+			k_heap_free(&llext_heap, tmp->mem[i]);
+			tmp->mem[i] = NULL;
 		}
 	}
 
-	k_heap_free(&llext_heap, ext->sym_tab.syms);
+	k_heap_free(&llext_heap, tmp->sym_tab.syms);
+	k_heap_free(&llext_heap, tmp);
 
-	k_heap_free(&llext_heap, ext);
+	return 0;
 }
 
 int llext_call_fn(struct llext *ext, const char *sym_name)

subsys/llext/shell.c

@@ -139,7 +139,7 @@ static int cmd_llext_unload(const struct shell *sh, size_t argc, char *argv[])
 		return -EINVAL;
 	}
 
-	llext_unload(ext);
+	llext_unload(&ext);
 	shell_print(sh, "Unloaded extension %s\n", argv[1]);
 
 	return 0;

tests/subsys/llext/src/test_llext_simple.c

@@ -29,7 +29,7 @@ ZTEST(llext, test_llext_simple)
 		LLEXT_BUF_LOADER(hello_world_elf, ARRAY_SIZE(hello_world_elf));
 	struct llext_loader *loader = &buf_loader.loader;
 	struct llext_load_param ldr_parm = LLEXT_LOAD_PARAM_DEFAULT;
-	struct llext *ext;
+	struct llext *ext = NULL;
 	const void * const printk_fn = llext_find_sym(NULL, "printk");
 
 	zassert_equal(printk_fn, printk, "printk should be an exported symbol");
@@ -46,7 +46,7 @@ ZTEST(llext, test_llext_simple)
 
 	zassert_ok(res, "calling hello world should succeed");
 
-	llext_unload(ext);
+	llext_unload(&ext);
 }
 
 ZTEST_SUITE(llext, NULL, NULL, NULL, NULL, NULL);