feat(orchestrator): return subnet assignment also on upgrade loop errors

[pierugo-dfinity]

The upgrade loop in the orchestrator is responsible both for executing upgrades and determining the subnet ID of the node, used to provision SSH keys and rotate IDKG keys. Though there are multiple code flows where the orchestrator determines the subnet ID but there is an error later in the loop, which makes the function return an error and the caller not apply the subnet ID. This prevents SSH keys from being provisioned even though the subnet ID had correctly been identified.

An example of such a code flow is if the local CUP is not deserializable but the NiDkgId is, which allows the subnet ID to be correctly determined (i.e. we hit here). But since the CUP is not deserializable and currently has the highest height compared to a recovery or peers CUP (we imagine it's at the very start of a recovery, before applying SSH keys -> there is no recovery CUP yet), we return an error here and the subnet ID is not updated, and SSH keys are not provisioned. If it does not have the highest height (i.e. there is a recovery CUP), then we can use the latter, which explains why we can still recover.

Note: the existing system test sr_app_no_upgrade_with_chain_keys_test is testing that we can recover a subnet exactly in that case (if the CUP is not deserializable but the NiDkgId is). As explained, nodes can see the Recovery CUP, but we do not apply readonly keys even though we could. In a parallel PR, I distinguished cases where the NiDkgId was corrupted or not. If yes, then there's indeed no way of provisioning SSH keys, but there's also no way of seeing the Recovery CUP -> thus use failover nodes. If not, then we should be able to provision SSH keys. When the second case runs on the current implementation, it fails because we cannot provision SSH keys. When merging this branch to it, the test succeeds, which is a positive sign towards the added value of this change.

Another example is if we detected we need to leave the subnet but removing the state failed (i.e. hit here). Then, we'd return an error again and fail to remove SSH keys of the subnet.

This PR is not supposed to bring any functional change to the upgrade logic but instead moves the responsibility of setting the subnet assignment from the caller of check_for_upgrade to the latter directly.

PS: The PR also uses the same registry version for the entire loop, instead of determining multiple times the latest registry version (in functions prepare_upgrade_if_scheduled, check_for_upgrade_as_unassigned, should_node_become_unassigned), in order to have a more consistent and predictable behaviour.

[pierugo-dfinity]

Note: I would also argue changing both of these false to true. I.e., in case the registry version is unavailable locally or somehow the field is empty or not deserializable, I would prefer not to accidentally remove the state, keep the subnet's SSH keys for a bit too long and try to rotate IDKG keys (in which case the registry should anyways deny the request because we would have left the subnet) rather than the opposite.

As of today, I cannot see a way where the registry version is not available locally, since it is always lower than the latest that we have. But if this function gets re-used somewhere else, I feel like returning true is more fail-safe than false. What do you guys think?

Note that replacing this to true could also mean launching the replica even though we are unassigned. But again, I do not think it hurts a lot if it's a single node doing so since other nodes would ignore it.

[eichhorl]

This is called "UpgradeCheckResult", but it is not actually a Result type, which is confusing. The enum does have different Ok and Err variants, however. Could we just turn it into a "real" Result? For instance, could we have a new error type that wraps both an OrchestratorError and SubnetAssignment, or similar?

[eichhorl]

Should there also be an ErrorAsLeaving variant?

[eichhorl]

Does this make a difference?

[pierugo-dfinity]

Are you talking about the return type being bool? Not much, I just did not want to return some ErrorAsUnassigned here and instead keep a simple Err to be more generic (i.e. the function does not need to “know” about UpgradeCheckResult). This was reverted anyways.

[eichhorl]

If returning the subnet assignment is too much of a hassle, especially since it needs to be updated in both the Ok and Err case, maybe another idea could be to have the Upgrade struct own a copy to this RwLock and mutating it directly there?

[eichhorl]

I think this block is a bit hard to understand. If this refactoring is needed, is there a way to avoid the named blocks and break statements? We could also consider extracting some functions if that makes it easier.

[pierugo-dfinity]

Indeed. I reverted this and kept most of what was already there while improving it to return just the subnet ID out of the block, instead of also the CUP and CUP proto. LMK what you think. (I recommend reviewing with whitespaces disabled)

[pierugo-dfinity]

Thanks for the review @eichhorl, I agree with all your remarks, that function has quite a high coupling so it’s tricky 😅. I’ll cover the ones that I haven’t answered directly here since they are interconnected.

This is called "UpgradeCheckResult", but it is not actually a Result type, which is confusing. The enum does have different Ok and Err variants, however. Could we just turn it into a "real" Result? For instance, could we have a new error type that wraps both an OrchestratorError and SubnetAssignment, or similar?

and

Should there also be an ErrorAsLeaving variant?

I think these translate to keep the current OrchestratorControlFlowenum and returning a Result<OrchestratorControlFlow, (Option<OrchestratorControlFlow>, OrchestratorError)>(obviously with intermediary types) (the Option in case we have an error and can't tell the flow). This could even encode an error during Stop. But I don’t think it helps readability much:

• We probably would still need a get_control_flow function or smth similar to get the flow from either the Ok or Err variant without doing a pattern matching in the caller, but we would need to introduce a new trait for that because Result is a foreign type (or wrap the Result in a 1-tuple).
• Each time we’d like to return an error in check(), we would have to return something like Err(UpgradeCheckError::new(OrchestratorControlFlow::Unassigned, err)), in contrast to the current UpgradeCheckResult::ErrorAsUnassigned(err) which reads more easily imo.

Which comes to

If returning the subnet assignment is too much of a hassle, especially since it needs to be updated in both the Ok and Err case, maybe another idea could be to have the Upgrade struct own a copy to this RwLock and mutating it directly there?

I think this is the cleanest solution. Unfortunately, it is impossible to hold the lock during an .await. We could either use an async RwLock, but it does not feel super idiomatic (slow and would turn the dashboard’s response async), or instead commit the Assigned assignment early and overwrite it later if we discover we are actually not part of the subnet anymore. This means that other tasks could potentially see the unassignment a bit late, but that is not a big deal (since the commitment of being Assigned is sound: we were previously in the subnet). This is what I finally implemented, reverting a lot of the original changes of the PR. I have to agree it looks much nicer.
I'm also open to using an async RwLock to prevent requesting an exclusive lock twice, but that could also introduce some blocking on the readers (including the dashboard's response) while we are fetching a CUP from our peers (in get_latest_cup), so I'm more tempted towards the current solution.