From 847ec5e60e445ef689ebb548a22409f8772e7570 Mon Sep 17 00:00:00 2001
From: Hubert Badocha <hubert.badocha@phoenix-rtos.com>
Date: Tue, 3 Oct 2023 16:54:08 +0200
Subject: [PATCH] atexit: allow for calling atexit inside of atexit

JIRA: RTOS-636
---
 include/stdlib.h |  2 --
 stdlib/atexit.c  | 76 ++++++++++++++++++++++++++----------------------
 2 files changed, 41 insertions(+), 37 deletions(-)

diff --git a/include/stdlib.h b/include/stdlib.h
index daacdd9c..ae452429 100644
--- a/include/stdlib.h
+++ b/include/stdlib.h
@@ -32,8 +32,6 @@ extern "C" {
 #define EXIT_SUCCESS 0
 #define EXIT_FAILURE 1
 
-#define ATEXIT_MAX 32
-
 #define MB_CUR_MAX 1
 #define MB_LEN_MAX 4
 
diff --git a/stdlib/atexit.c b/stdlib/atexit.c
index 88c52399..db77b2fe 100644
--- a/stdlib/atexit.c
+++ b/stdlib/atexit.c
@@ -21,8 +21,14 @@
 /* The atexit_node structure uses uint32_t for flags that
  indicate whether to call a function with arguments or not, therefore
  why ATEXIT_MAX must be 32 as type "fntype" */
+#define ATEXIT_MAX 32
+
+
+typedef void (*destructor_t)(void);
+
+
 struct atexit_node {
-	void (*destructors[ATEXIT_MAX])(void);
+	destructor_t destructors[ATEXIT_MAX];
 	void *args[ATEXIT_MAX];
 	uint32_t fntype;
 	struct atexit_node *prev;
@@ -33,6 +39,7 @@ struct atexit_node {
 static struct {
 	handle_t lock;
 	struct atexit_node *head;
+	unsigned int idx;
 } atexit_common = { .head = &((struct atexit_node) {}) };
 
 
@@ -41,6 +48,7 @@ void _atexit_init(void)
 {
 	mutexCreate(&atexit_common.lock);
 	memset(atexit_common.head, 0, sizeof(struct atexit_node));
+	atexit_common.idx = 0;
 }
 
 
@@ -48,36 +56,30 @@ void _atexit_init(void)
 static int _atexit_register(int isarg, void (*fn)(void), void *arg)
 {
 	struct atexit_node *node;
-	int i;
 
 	mutexLock(atexit_common.lock);
 	node = atexit_common.head;
 
-	/* Looking for an empty slot */
-	for (i = 0; i < ATEXIT_MAX; i++) {
-		if (node->destructors[i] == NULL) {
-			break;
-		}
-	}
-
 	/* Allocate new node if there are no free slots left */
-	if (i == ATEXIT_MAX) {
+	if (atexit_common.idx == ATEXIT_MAX) {
 		node = (struct atexit_node *)calloc(1, sizeof(struct atexit_node));
 		if (node == NULL) {
 			mutexUnlock(atexit_common.lock);
 			return -1;
 		}
 		node->prev = atexit_common.head;
-		atexit_common.head = node;
 
-		i = 0;
+		atexit_common.head = node;
+		atexit_common.idx = 0;
 	}
 
-	if (isarg) {
-		node->fntype |= 1u << i;
-		node->args[i] = arg;
+	if (isarg != 0) {
+		node->fntype |= (1u << atexit_common.idx);
+		node->args[atexit_common.idx] = arg;
 	}
-	node->destructors[i] = fn;
+	node->destructors[atexit_common.idx] = fn;
+
+	atexit_common.idx++;
 
 	mutexUnlock(atexit_common.lock);
 	return 0;
@@ -87,28 +89,32 @@ static int _atexit_register(int isarg, void (*fn)(void), void *arg)
 /* Call all registered destructors */
 void _atexit_call(void)
 {
-	struct atexit_node *last, *node = atexit_common.head;
-	int i;
-
-	while (node != NULL) {
-		for (i = ATEXIT_MAX - 1; i >= 0; i--) {
-			if (node->destructors[i] != NULL) {
-				if (((node->fntype >> i) & 1) == 1) {
-					((void (*)(void *))node->destructors[i])(node->args[i]);
-				}
-				else {
-					node->destructors[i]();
-				}
-			}
+	mutexLock(atexit_common.lock);
+	while ((atexit_common.head != NULL) && (atexit_common.idx != 0)) {
+		atexit_common.idx--;
+		destructor_t destructor = atexit_common.head->destructors[atexit_common.idx];
+		if (((atexit_common.head->fntype >> atexit_common.idx) & 1) == 1) {
+			void *arg = atexit_common.head->args[atexit_common.idx];
+			mutexUnlock(atexit_common.lock);
+			((void (*)(void *))destructor)(arg);
 		}
-
-		last = node;
-		node = node->prev;
-		if (node != NULL) {
-			/* Free only if it's not the first node (statically allocated) */
-			free(last);
+		else {
+			mutexUnlock(atexit_common.lock);
+			destructor();
+		}
+		mutexLock(atexit_common.lock);
+
+		if (atexit_common.idx == 0) {
+			struct atexit_node *last = atexit_common.head;
+			atexit_common.head = atexit_common.head->prev;
+			if (atexit_common.head != NULL) {
+				/* Free only if it's not the first node (statically allocated) */
+				free(last);
+				atexit_common.idx = ATEXIT_MAX;
+			}
 		}
 	}
+	mutexUnlock(atexit_common.lock);
 }