tasks: Add extra MAIN task id

Add TASK_ID_MAIN as a non-shimmed extra task. This task is only present
when CONFIG_TASK_HOSTCMD_THREAD_DEDICATED is enabled.

Add get_hostcmd_thread and get_main_thread helper methods. The hostcmd
thread will be the same as the main thread when
CONFIG_TASK_HOSTCMD_THREAD_MAIN is enabled.

BUG=b:267470086
BRANCH=None
TEST=Unit tests

Change-Id: If61dca427686a9c3cf6fb05dc5ca4e5a87b47127
Signed-off-by: Rob Barnes <robbarnes@google.com>
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/ec/+/4214560
Reviewed-by: Fabio Baltieri <fabiobaltieri@google.com>

7 files changed, 58 insertions, 25 deletions

diff --git a/zephyr/shim/include/shimmed_task_id.h b/zephyr/shim/include/shimmed_task_id.h
index f6300ed150..77c4d6f462 100644
--- a/zephyr/shim/include/shimmed_task_id.h
+++ b/zephyr/shim/include/shimmed_task_id.h
@@ -229,7 +229,8 @@ enum {
  */
 /* clang-format off */
 #define CROS_EC_EXTRA_TASKS(fn)                                         \
-	COND_CODE_1(CONFIG_TASK_HOSTCMD_THREAD_MAIN, (fn(HOSTCMD)), ()) \
+	COND_CODE_1(CONFIG_TASK_HOSTCMD_THREAD_MAIN, (fn(HOSTCMD)),     \
+		(fn(MAIN)))                                             \
 	fn(SYSWORKQ)                                                    \
 	fn(IDLE)
 /* clang-format on */
diff --git a/zephyr/shim/include/shimmed_tasks.h b/zephyr/shim/include/shimmed_tasks.h
index ea44d2a113..8f549f7bb7 100644
--- a/zephyr/shim/include/shimmed_tasks.h
+++ b/zephyr/shim/include/shimmed_tasks.h
@@ -47,6 +47,10 @@
 #define HAS_TASK_USB_MUX 1
 #endif /* CONFIG_PLATFORM_EC_USB_MUX_TASK */
 
+#ifdef CONFIG_TASK_HOSTCMD_THREAD_DEDICATED
+#define HAS_TASK_MAIN 1
+#endif /* CONFIG_TASK_HOSTCMD_THREAD_DEDICATED */
+
 /* These non-shimmed (extra) tasks are always present */
 #define HAS_TASK_IDLE 1
 #define HAS_TASK_SYSWORKQ 1
diff --git a/zephyr/shim/include/zephyr_host_command.h b/zephyr/shim/include/zephyr_host_command.h
index 8619508c42..844f9ae395 100644
--- a/zephyr/shim/include/zephyr_host_command.h
+++ b/zephyr/shim/include/zephyr_host_command.h
@@ -13,6 +13,7 @@
 #include <stdbool.h>
 
 #include <zephyr/init.h>
+#include <zephyr/kernel.h>
 
 /* Initializes and runs the host command handler loop.  */
 void host_command_task(void *u);
@@ -20,8 +21,17 @@ void host_command_task(void *u);
 /* Takes over the main thread and runs the host command loop. */
 void host_command_main(void);
 
-/* True if running in the main thread. */
-bool in_host_command_main(void);
+/*
+ * Returns the main thread id. Will be the same as the HOSTCMD thread
+ * when CONFIG_TASK_HOSTCMD_THREAD_MAIN is enabled.
+ */
+k_tid_t get_main_thread(void);
+
+/*
+ * Returns the HOSTCMD thread id. Will be different than the main thread
+ * when CONFIG_TASK_HOSTCMD_THREAD_DEDICATED is enabled.
+ */
+k_tid_t get_hostcmd_thread(void);
 
 #ifdef CONFIG_PLATFORM_EC_HOSTCMD
 
diff --git a/zephyr/shim/src/host_command.c b/zephyr/shim/src/host_command.c
index 2e55099de2..5ca84c944a 100644
--- a/zephyr/shim/src/host_command.c
+++ b/zephyr/shim/src/host_command.c
@@ -19,18 +19,10 @@ struct host_command *zephyr_find_host_command(int command)
 	return NULL;
 }
 
-/* Pointer to the main thread, defined in kernel/init.c */
-extern struct k_thread z_main_thread;
-
 void host_command_main(void)
 {
-	k_thread_priority_set(&z_main_thread,
+	k_thread_priority_set(get_main_thread(),
 			      EC_TASK_PRIORITY(EC_TASK_HOSTCMD_PRIO));
-	k_thread_name_set(&z_main_thread, "HOSTCMD");
+	k_thread_name_set(get_main_thread(), "HOSTCMD");
 	host_command_task(NULL);
 }
-
-test_mockable bool in_host_command_main(void)
-{
-	return (k_current_get() == &z_main_thread);
-}
diff --git a/zephyr/shim/src/tasks.c b/zephyr/shim/src/tasks.c
index 62bca4ed60..129f7b3073 100644
--- a/zephyr/shim/src/tasks.c
+++ b/zephyr/shim/src/tasks.c
@@ -94,17 +94,43 @@ test_export_static k_tid_t get_sysworkq_thread(void)
 	return &k_sys_work_q.thread;
 }
 
+k_tid_t get_main_thread(void)
+{
+	/* Pointer to the main thread, defined in kernel/init.c */
+	extern struct k_thread z_main_thread;
+
+	return &z_main_thread;
+}
+
+test_mockable k_tid_t get_hostcmd_thread(void)
+{
+#if IS_ENABLED(HAS_TASK_HOSTCMD)
+	if (IS_ENABLED(CONFIG_TASK_HOSTCMD_THREAD_MAIN)) {
+		return get_main_thread();
+	}
+	return task_to_k_tid[TASK_ID_HOSTCMD];
+#endif /* HAS_TASK_HOSTCMD */
+	__ASSERT(false, "HOSTCMD task is not enabled");
+	return NULL;
+}
+
 task_id_t task_get_current(void)
 {
 	if (get_sysworkq_thread() == k_current_get()) {
 		return TASK_ID_SYSWORKQ;
 	}
 
-#ifdef CONFIG_TASK_HOSTCMD_THREAD_MAIN
-	if (in_host_command_main()) {
+#if IS_ENABLED(HAS_TASK_HOSTCMD)
+	if (get_hostcmd_thread() == k_current_get()) {
 		return TASK_ID_HOSTCMD;
 	}
-#endif
+#endif /* HAS_TASK_HOSTCMD */
+
+#if IS_ENABLED(HAS_TASK_MAIN)
+	if (get_main_thread() == k_current_get()) {
+		return TASK_ID_MAIN;
+	}
+#endif /* HAS_TASK_MAIN */
 
 	if (get_idle_thread() == k_current_get()) {
 		return TASK_ID_IDLE;
diff --git a/zephyr/test/drivers/host_command_thread/src/main.c b/zephyr/test/drivers/host_command_thread/src/main.c
index 72d12ee8c0..bcd1b97074 100644
--- a/zephyr/test/drivers/host_command_thread/src/main.c
+++ b/zephyr/test/drivers/host_command_thread/src/main.c
@@ -21,9 +21,6 @@
 
 #define CUSTOM_COMMAND_ID 0x0088
 
-/* Pointer to the main thread, defined in kernel/init.c */
-extern struct k_thread z_main_thread;
-
 /* Thread id of fake main thread */
 static k_tid_t fake_main_tid;
 
@@ -33,7 +30,7 @@ static int last_check_main_thread_result;
 static enum ec_status check_main_thread(struct host_cmd_handler_args *args)
 {
 	last_check_main_thread_result =
-		(k_current_get() == &z_main_thread ? 1 : -1);
+		(k_current_get() == get_main_thread() ? 1 : -1);
 	return EC_RES_SUCCESS;
 }
 
@@ -46,13 +43,13 @@ static void fake_main_thread(void *a, void *b, void *c)
 
 K_THREAD_STACK_DEFINE(fake_main_thread_stack, 4000);
 
-/* Override in_host_command_main() from shim/src/host_command.c so
+/* Override get_hostcmd_thread() from shim/src/tasks.c so
  * task_get_current() returns TASK_ID_HOSTCMD when fake main thread
  * is running.
  */
-bool in_host_command_main(void)
+k_tid_t get_hostcmd_thread(void)
 {
-	return (k_current_get() == fake_main_tid);
+	return fake_main_tid;
 }
 
 ZTEST_SUITE(host_cmd_thread, drivers_predicate_post_main, NULL, NULL, NULL,
@@ -74,11 +71,11 @@ ZTEST(host_cmd_thread, test_takeover)
 	k_msleep(500);
 
 	/* Get the name of the thread (must be done after the sleep) */
-	const char *main_thread_name = k_thread_name_get(&z_main_thread);
+	const char *main_thread_name = k_thread_name_get(get_main_thread());
 
 	/* Verify that the thread is not the hostcmd thread */
 	zassert_equal(EC_TASK_PRIORITY(EC_TASK_HOSTCMD_PRIO),
-		      k_thread_priority_get(&z_main_thread));
+		      k_thread_priority_get(get_main_thread()));
 	zassert_equal(strlen(expected_thread_name), strlen(main_thread_name));
 	zassert_mem_equal(expected_thread_name, main_thread_name,
 			  strlen(expected_thread_name));
diff --git a/zephyr/test/tasks/extra_tasks.c b/zephyr/test/tasks/extra_tasks.c
index 6069741001..ac11c20940 100644
--- a/zephyr/test/tasks/extra_tasks.c
+++ b/zephyr/test/tasks/extra_tasks.c
@@ -3,6 +3,7 @@
  * found in the LICENSE file.
  */
 
+#include "host_command.h"
 #include "task.h"
 
 #include <zephyr/kernel.h>
@@ -37,9 +38,11 @@ ZTEST_USER(extra_tasks, test_main_thread_mapping)
 
 	hostcmd_thread = find_thread_by_name("HOSTCMD");
 	zassert_not_null(hostcmd_thread);
+	zassert_equal(hostcmd_thread, get_hostcmd_thread());
 
 	main_thread = find_thread_by_name("main");
 	zassert_not_null(main_thread);
+	zassert_equal(main_thread, get_main_thread());
 	/* Not equal when CONFIG_TASK_HOSTCMD_THREAD_DEDICATED is set */
 	zassert_not_equal(main_thread, hostcmd_thread);
 }