feat: implement state machine for vCPU state transitions

This patch implements the state machine which tracks the various
legal transitions allowed for a vCPU as defined by the rules in
the FF-A v1.3 ALP2 specification.

Change-Id: I03f2d398b1ffc965e69449277536403605267766
Signed-off-by: Madhukar Pappireddy <madhukar.pappireddy@arm.com>
diff --git a/inc/hf/vcpu.h b/inc/hf/vcpu.h
index 0c12324..c845e20 100644
--- a/inc/hf/vcpu.h
+++ b/inc/hf/vcpu.h
@@ -352,3 +352,4 @@
 			    ffa_id_t sender_vm_id, struct ffa_value args);
 
 void vcpu_dir_req_reset_state(struct vcpu_locked vcpu_locked);
+bool vcpu_state_set(struct vcpu_locked vcpu_locked, enum vcpu_state to);
diff --git a/src/api.c b/src/api.c
index aff2efa..cf3cd18 100644
--- a/src/api.c
+++ b/src/api.c
@@ -143,7 +143,7 @@
 	arch_regs_set_retval(&next->regs, to_ret);
 
 	/* Set the current vCPU state. */
-	current_locked.vcpu->state = vcpu_state;
+	CHECK(vcpu_state_set(current_locked, vcpu_state));
 
 	return next;
 }
@@ -1319,7 +1319,8 @@
 			vcpu->rt_model = RTM_NONE;
 
 			vcpu_was_init_state = true;
-			vcpu->state = VCPU_STATE_STARTING;
+			CHECK(vcpu_state_set(vcpu_next_locked,
+					     VCPU_STATE_STARTING));
 			break;
 		}
 		*run_ret = ffa_error(FFA_BUSY);
@@ -1487,7 +1488,7 @@
 
 	assert(!vm_id_is_current_world(current->vm->id) ||
 	       next_state == VCPU_STATE_BLOCKED);
-	current->state = VCPU_STATE_BLOCKED;
+	CHECK(vcpu_state_set(current_locked, VCPU_STATE_BLOCKED));
 
 	/*
 	 * Set a placeholder return code to the scheduler. This will be
@@ -2970,7 +2971,8 @@
 				"Receiver VM %#x aborted, cannot run vCPU %u\n",
 				receiver_vcpu->vm->id,
 				vcpu_index(receiver_vcpu));
-			receiver_vcpu->state = VCPU_STATE_ABORTED;
+			CHECK(vcpu_state_set(receiver_vcpu_locked,
+					     VCPU_STATE_ABORTED));
 		}
 	}
 
@@ -3016,7 +3018,7 @@
 
 	assert(!vm_id_is_current_world(current->vm->id) ||
 	       next_state == VCPU_STATE_BLOCKED);
-	current->state = VCPU_STATE_BLOCKED;
+	CHECK(vcpu_state_set(current_locked, VCPU_STATE_BLOCKED));
 
 	ffa_direct_msg_wind_call_chain_ffa_direct_req(
 		current_locked, receiver_vcpu_locked, sender_vm_id);
@@ -3252,6 +3254,9 @@
 	ffa_direct_msg_unwind_call_chain_ffa_direct_resp(current_locked,
 							 next_locked);
 
+	/* Schedule the receiver's vCPU now. */
+	CHECK(vcpu_state_set(next_locked, VCPU_STATE_RUNNING));
+
 	/*
 	 * Check if there is a pending interrupt, and if the partition
 	 * is expects to notify the scheduler or resume straight away.
diff --git a/src/arch/aarch64/hypervisor/handler.c b/src/arch/aarch64/hypervisor/handler.c
index df62f4d..ed950dd 100644
--- a/src/arch/aarch64/hypervisor/handler.c
+++ b/src/arch/aarch64/hypervisor/handler.c
@@ -1136,7 +1136,7 @@
 	 * A partition exection is in ABORTED state after it encounters a fatal
 	 * error.
 	 */
-	current->state = VCPU_STATE_ABORTED;
+	CHECK(vcpu_state_set(current_locked, VCPU_STATE_ABORTED));
 
 	/*
 	 * SPMC de-allocates and/or uninitializes all the resources allocated
@@ -1146,6 +1146,12 @@
 	vm_unlock(&vm_locked);
 
 	/*
+	 * Move the vCPU to STOPPED state as all its resources have been
+	 * reclaimed by SPMC by now.
+	 */
+	CHECK(vcpu_state_set(current_locked, VCPU_STATE_STOPPED));
+
+	/*
 	 * SPMC performs necessary operations based on the abort action for
 	 * SP.
 	 */
diff --git a/src/arch/aarch64/plat/psci/spmc.c b/src/arch/aarch64/plat/psci/spmc.c
index f1c4b11..56a583f 100644
--- a/src/arch/aarch64/plat/psci/spmc.c
+++ b/src/arch/aarch64/plat/psci/spmc.c
@@ -148,7 +148,7 @@
 				       0ULL);
 
 	/* Set the vCPU's state to STARTING. */
-	boot_vcpu->state = VCPU_STATE_STARTING;
+	CHECK(vcpu_state_set(vcpu_locked, VCPU_STATE_STARTING));
 	boot_vcpu->regs_available = false;
 
 	/* vCPU restarts in runtime model for SP initialization. */
diff --git a/src/ffa/spmc/cpu_cycles.c b/src/ffa/spmc/cpu_cycles.c
index 8dfd823..96aa1aa 100644
--- a/src/ffa/spmc/cpu_cycles.c
+++ b/src/ffa/spmc/cpu_cycles.c
@@ -209,7 +209,7 @@
 	ffa_cpu_cycles_exit_spmc_schedule_mode(current_locked);
 	assert(current->call_chain.prev_node == NULL);
 
-	current->state = VCPU_STATE_WAITING;
+	CHECK(vcpu_state_set(current_locked, VCPU_STATE_WAITING));
 
 	vcpu_set_running(target_locked, NULL);
 
@@ -307,7 +307,7 @@
 		next_vm = vm_get_next_boot_secondary_core(current->vm);
 	}
 
-	current->state = VCPU_STATE_WAITING;
+	CHECK(vcpu_state_set(current_locked, VCPU_STATE_WAITING));
 	current->rt_model = RTM_NONE;
 	current->scheduling_mode = NONE;
 
@@ -330,10 +330,13 @@
 	 */
 	if (vcpu_next->rt_model == RTM_SP_INIT ||
 	    vcpu_next->state == VCPU_STATE_OFF) {
+		struct vcpu_locked vcpu_next_locked;
+
+		vcpu_next_locked = vcpu_lock(vcpu_next);
 		vcpu_next->rt_model = RTM_SP_INIT;
 		arch_regs_reset(vcpu_next);
 		vcpu_next->cpu = current->cpu;
-		vcpu_next->state = VCPU_STATE_STARTING;
+		CHECK(vcpu_state_set(vcpu_next_locked, VCPU_STATE_STARTING));
 		vcpu_next->regs_available = false;
 		vcpu_set_phys_core_idx(vcpu_next);
 		arch_regs_set_pc_arg(&vcpu_next->regs,
@@ -347,6 +350,7 @@
 			vcpu_set_boot_info_gp_reg(vcpu_next);
 		}
 
+		vcpu_unlock(&vcpu_next_locked);
 		*next = vcpu_next;
 
 		return true;
@@ -487,6 +491,8 @@
 			*next = api_switch_to_other_world(current_locked, ret,
 							  VCPU_STATE_BLOCKED);
 		} else {
+			struct two_vcpu_locked vcpus_locked;
+
 			/*
 			 * Relinquish cycles to the SP that sent direct request
 			 * message to the current SP.
@@ -494,6 +500,14 @@
 			*next = api_switch_to_vm(
 				current_locked, ret, VCPU_STATE_BLOCKED,
 				current->direct_request_origin.vm_id);
+
+			vcpu_unlock(&current_locked);
+
+			/* Lock both vCPUs at once to avoid deadlock. */
+			vcpus_locked = vcpu_lock_both(current, *next);
+
+			vcpu_state_set(vcpus_locked.vcpu2, VCPU_STATE_RUNNING);
+			vcpu_unlock(&vcpus_locked.vcpu2);
 		}
 		break;
 	case RTM_SEC_INTERRUPT: {
diff --git a/src/ffa/spmc/direct_messaging.c b/src/ffa/spmc/direct_messaging.c
index d6b06f3..85e80df 100644
--- a/src/ffa/spmc/direct_messaging.c
+++ b/src/ffa/spmc/direct_messaging.c
@@ -626,7 +626,7 @@
 	vcpu_dir_req_reset_state(current_locked);
 
 	/* Turn off the vCPU.*/
-	current->state = VCPU_STATE_OFF;
+	CHECK(vcpu_state_set(current_locked, VCPU_STATE_OFF));
 	dlog_verbose("SPMC turned off vCPU%zu of SP%#x\n",
 		     cpu_index(current->cpu), current->vm->id);
 
diff --git a/src/ffa/spmc/interrupts.c b/src/ffa/spmc/interrupts.c
index 94594b0..da3b9fc 100644
--- a/src/ffa/spmc/interrupts.c
+++ b/src/ffa/spmc/interrupts.c
@@ -135,7 +135,7 @@
 	assert(preempted_vcpu != NULL);
 
 	target_vcpu->preempted_vcpu = preempted_vcpu;
-	preempted_vcpu->state = VCPU_STATE_PREEMPTED;
+	CHECK(vcpu_state_set(current_locked, VCPU_STATE_PREEMPTED));
 }
 
 /**
@@ -564,7 +564,7 @@
 
 	/* Removing a node from an existing call chain. */
 	current_vcpu->call_chain.prev_node = NULL;
-	current_vcpu->state = VCPU_STATE_PREEMPTED;
+	CHECK(vcpu_state_set(both_vcpu_locked.vcpu1, VCPU_STATE_PREEMPTED));
 
 	/*
 	 * SPMC applies the runtime model till when the vCPU transitions from
diff --git a/src/vcpu.c b/src/vcpu.c
index ba52d3b..efefc9c 100644
--- a/src/vcpu.c
+++ b/src/vcpu.c
@@ -84,7 +84,7 @@
 void vcpu_prepare(struct vcpu_locked vcpu, ipaddr_t entry, uintreg_t arg)
 {
 	arch_regs_set_pc_arg(&vcpu.vcpu->regs, entry, arg);
-	vcpu.vcpu->state = VCPU_STATE_CREATED;
+	CHECK(vcpu_state_set(vcpu, VCPU_STATE_CREATED));
 }
 
 ffa_vcpu_index_t vcpu_index(const struct vcpu *vcpu)
@@ -687,3 +687,146 @@
 	vcpu->scheduling_mode = NONE;
 	vcpu->rt_model = RTM_NONE;
 }
+
+static inline const char *vcpu_state_print_name(enum vcpu_state state)
+{
+	switch (state) {
+	case VCPU_STATE_OFF:
+		return "VCPU_STATE_OFF";
+	case VCPU_STATE_RUNNING:
+		return "VCPU_STATE_RUNNING";
+	case VCPU_STATE_WAITING:
+		return "VCPU_STATE_WAITING";
+	case VCPU_STATE_BLOCKED:
+		return "VCPU_STATE_BLOCKED";
+	case VCPU_STATE_PREEMPTED:
+		return "VCPU_STATE_PREEMPTED";
+	case VCPU_STATE_BLOCKED_INTERRUPT:
+		return "VCPU_STATE_BLOCKED_INTERRUPT";
+	case VCPU_STATE_ABORTED:
+		return "VCPU_STATE_ABORTED";
+	case VCPU_STATE_NULL:
+		return "VCPU_STATE_NULL";
+	case VCPU_STATE_STOPPED:
+		return "VCPU_STATE_STOPPED";
+	case VCPU_STATE_CREATED:
+		return "VCPU_STATE_CREATED";
+	case VCPU_STATE_STARTING:
+		return "VCPU_STATE_STARTING";
+	case VCPU_STATE_STOPPING:
+		return "VCPU_STATE_STOPPING";
+	}
+}
+
+static bool vcpu_is_valid_transition_from_starting_running_state(
+	enum vcpu_state to_state)
+{
+	switch (to_state) {
+	case VCPU_STATE_WAITING:
+	case VCPU_STATE_BLOCKED:
+	case VCPU_STATE_PREEMPTED:
+	case VCPU_STATE_ABORTED:
+	case VCPU_STATE_STOPPING:
+		return true;
+	default:
+		return false;
+	}
+}
+
+static bool vcpu_is_valid_transition_from_stopping_state(
+	enum vcpu_state to_state)
+{
+	switch (to_state) {
+	case VCPU_STATE_WAITING:
+	case VCPU_STATE_BLOCKED:
+	case VCPU_STATE_PREEMPTED:
+	case VCPU_STATE_ABORTED:
+	case VCPU_STATE_STOPPED:
+		return true;
+	default:
+		return false;
+	}
+}
+
+bool vcpu_state_set(struct vcpu_locked vcpu_locked, enum vcpu_state to_state)
+{
+	struct vcpu *vcpu = vcpu_locked.vcpu;
+	enum vcpu_state from_state = vcpu->state;
+	bool ret;
+
+	/*
+	 * SPMC can turn off a vCPU, irrespective of the present state, when it
+	 * receives CPU_OFF power management message from SPMD.
+	 */
+	if (ffa_is_vm_id(vcpu->vm->id) || from_state == to_state ||
+	    to_state == VCPU_STATE_OFF) {
+		ret = true;
+		goto out;
+	}
+
+	switch (from_state) {
+	case VCPU_STATE_OFF:
+		/*
+		 * When a CPU is powered on after previously being powered down,
+		 * the pinned vCPU of an SP moves from OFF to CREATED before
+		 * transitioning to STARTING state.
+		 */
+		ret = (to_state == VCPU_STATE_STARTING ||
+		       to_state == VCPU_STATE_CREATED);
+		break;
+	case VCPU_STATE_NULL:
+		ret = (to_state == VCPU_STATE_CREATED);
+		break;
+	case VCPU_STATE_CREATED:
+		ret = (to_state == VCPU_STATE_STARTING);
+		break;
+	case VCPU_STATE_STARTING:
+	case VCPU_STATE_RUNNING:
+		ret = vcpu_is_valid_transition_from_starting_running_state(
+			to_state);
+		break;
+	case VCPU_STATE_STOPPING:
+		ret = vcpu_is_valid_transition_from_stopping_state(to_state);
+		break;
+	case VCPU_STATE_WAITING:
+		/*
+		 * A transition of vCPU from WAITING to ABORTED is legal. One
+		 * such scenario is when the endpoint was aborted.
+		 */
+		ret = (to_state == VCPU_STATE_RUNNING ||
+		       to_state == VCPU_STATE_OFF ||
+		       to_state == VCPU_STATE_ABORTED);
+		break;
+	case VCPU_STATE_BLOCKED:
+	case VCPU_STATE_BLOCKED_INTERRUPT:
+	case VCPU_STATE_PREEMPTED:
+		ret = (to_state == VCPU_STATE_RUNNING);
+		break;
+	case VCPU_STATE_ABORTED:
+		ret = (to_state == VCPU_STATE_NULL ||
+		       to_state == VCPU_STATE_STOPPED);
+		break;
+	case VCPU_STATE_STOPPED:
+		/*
+		 * A legacy SP will be put in ABORTED state by SPMC if it
+		 * encounters a fatal error in runtime.
+		 */
+		ret = (to_state == VCPU_STATE_CREATED ||
+		       to_state == VCPU_STATE_NULL ||
+		       to_state == VCPU_STATE_ABORTED);
+		break;
+	}
+out:
+	if (!ret) {
+		dlog_error(
+			"Partition %#x vCPU %u: transition from %s to %s "
+			"failed.\n",
+			vcpu->vm->id, (uint32_t)cpu_index(vcpu->cpu),
+			vcpu_state_print_name(from_state),
+			vcpu_state_print_name(to_state));
+	} else {
+		vcpu->state = to_state;
+	}
+
+	return ret;
+}