tx-support SetTaskSwitchExtension()

3 files changed, 33 insertions, 16 deletions

diff --git a/testsuites/validation/ts-default.h b/testsuites/validation/ts-default.h
index f2130028e4..216e8b1bfa 100644
--- a/testsuites/validation/ts-default.h
+++ b/testsuites/validation/ts-default.h
@@ -186,11 +186,7 @@ static void task_stack_deallocate( void *stack )
 #define CONFIGURE_INIT_TASK_CONSTRUCT_STORAGE_SIZE TASK_STORAGE_SIZE
 
 #if !defined( CONFIGURE_INITIAL_EXTENSIONS )
-#define CONFIGURE_INITIAL_EXTENSIONS \
-  { \
-    .fatal = FatalInitialExtension, \
-    .thread_switch = TaskSwitchInitialExtension \
-  }
+#define CONFIGURE_INITIAL_EXTENSIONS { .fatal = FatalInitialExtension }
 #endif
 
 #if defined( RTEMS_SMP ) && \
diff --git a/testsuites/validation/tx-support.c b/testsuites/validation/tx-support.c
index 1733b32ef1..263d187aa0 100644
--- a/testsuites/validation/tx-support.c
+++ b/testsuites/validation/tx-support.c
@@ -645,22 +645,45 @@ void SetFatalExtension( rtems_fatal_extension fatal )
   fatal_extension = fatal;
 }
 
+static rtems_id task_switch_id;
+
 static rtems_task_switch_extension task_switch_extension;
 
-void TaskSwitchInitialExtension( rtems_tcb *executing, rtems_tcb *heir )
+static void TaskSwitchExtension( rtems_tcb *executing, rtems_tcb *heir )
 {
-  rtems_task_switch_extension task_switch;
-
-  task_switch = task_switch_extension;
-
-  if ( task_switch != NULL ) {
-    ( *task_switch )( executing, heir );
-  }
+  ( *task_switch_extension )( executing, heir );
 }
 
 void SetTaskSwitchExtension( rtems_task_switch_extension task_switch )
 {
-  task_switch_extension = task_switch;
+  rtems_task_switch_extension last;
+  rtems_status_code           sc;
+
+  last = task_switch_extension;
+
+  if ( task_switch == NULL ) {
+    if ( last != NULL ) {
+      sc = rtems_extension_delete( task_switch_id );
+      T_quiet_rsc_success( sc );
+
+      task_switch_extension = NULL;
+    }
+  } else {
+    task_switch_extension = task_switch;
+
+    if ( last == NULL ) {
+      rtems_extensions_table table = {
+        .thread_switch = TaskSwitchExtension
+      };
+
+      sc = rtems_extension_create(
+        rtems_build_name( 'T', 'S', 'W', 'I' ),
+        &table,
+        &task_switch_id
+      );
+      T_quiet_rsc_success( sc );
+    }
+  }
 }
 
 void ClearExtensionCalls( ExtensionCalls *calls )
diff --git a/testsuites/validation/tx-support.h b/testsuites/validation/tx-support.h
index 759174f188..2ac1e8c1b3 100644
--- a/testsuites/validation/tx-support.h
+++ b/testsuites/validation/tx-support.h
@@ -448,8 +448,6 @@ void FatalInitialExtension(
 
 void SetFatalExtension( rtems_fatal_extension fatal );
 
-void TaskSwitchInitialExtension( rtems_tcb *executing, rtems_tcb *heir );
-
 void SetTaskSwitchExtension( rtems_task_switch_extension task_switch );
 
 typedef struct {