Skip to content

feat(migration): add support for back migration without restart #3550

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
barakeinav1 wants to merge 3 commits into
base: main
Choose a base branch
from
Draft

feat(migration): add support for back migration without restart #3550

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
b72e64a
fix(migration): make onboarding state machine re-entrant (#3406)
barakeinav1 Jun 14, 2026
9781816
fix(migration): apply Copilot review findings
barakeinav1 Jun 14, 2026
2cd0db0
docs(design): link #3551 from the Risk section
barakeinav1 Jun 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions crates/e2e-tests/tests/common.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ pub const CONTRACT_UPGRADE_COMPATIBILITY_TESTNET_PORT_SEED: u16 = 19;
pub const TIMEOUT_METRIC_PORT_SEED: u16 = 20;
pub const MIGRATION_BACK_PORT_SEED: u16 = 21;
pub const SIGTERM_HANDLER_PORT_SEED: u16 = 22;
pub const MIGRATION_BACK_NO_RESTART_PORT_SEED: u16 = 23;

/// Start a cluster, wait for Running state and presignatures to buffer.
///
Expand Down
193 changes: 95 additions & 98 deletions crates/e2e-tests/tests/migration_service.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -655,46 +655,33 @@ async fn migration_service__should_migrate_nodes_via_backup_cli() {
}
}

/// Back-migration (A → B → A) end-to-end via the backup CLI. Investigation
/// context for the production scenario lives at near/mpc#2121; this test
/// does NOT regress against that bug (see Caveat below).
///
/// Starts a cluster with 2 participating nodes (A0, A1) + 1 migration
/// target (B0). A1 stays a participant throughout so threshold (2) holds
/// during both directions. Per direction:
/// 1. Register backup service
/// 2. GET keyshares from source node
/// 3. Initiate node migration
/// 4. PUT keyshares to target node
/// 5. Verify migration finalizes and sign + ckd requests succeed
/// Between the forward and back directions, A0's CVM is killed and
/// restarted so its keyshares stay on disk through the restart.
/// After the back direction, B0 is killed so the final sign + ckd
/// assertions prove A0 + A1 carry the workload alone.
///
/// Caveat: attestation is mocked as `{"Mock": "Valid"}` throughout, so
/// this test cannot exercise the on-chain stale-attestation failure mode —
/// which is the leading suspect for #2121's production symptom. This is a
/// positive-baseline regression test for the back-migration code path, not
/// a reproducer for #2121.
#[tokio::test]
#[expect(non_snake_case)]
async fn migration_service__should_handle_back_migration_a_to_b_to_a() {
/// Action applied to A0 between the forward and back directions.
enum BetweenDirections {
/// A0 keeps running.
None,
/// A0 is killed and restarted from its on-disk home dir; asserts
/// keyshares survive.
KillAndRestartA,
}

/// Shared scaffolding: cluster setup, forward A → B, `between` hook,
/// back B → A, kill B0, then asserts sign+ckd succeed with A0+A1 alone.
/// Attestation is mocked throughout.
async fn run_back_migration_round_trip(port_seed: u16, between: BetweenDirections) {
let backup_cli = must_get_backup_cli_path();

// Given: 2 participants (A0, A1) + 1 migration target (B0).
// 2 participants (A0, A1) + 1 migration target (B0). A1 stays a
// participant throughout so threshold (2) holds in both directions.
let num_nodes = 2;
let (mut cluster, running) =
common::must_setup_cluster(common::MIGRATION_BACK_PORT_SEED, |c| {
c.num_nodes = num_nodes;
c.threshold = num_nodes;
c.migration_targets = vec![0];
})
.await;
let (mut cluster, running) = common::must_setup_cluster(port_seed, |c| {
c.num_nodes = num_nodes;
c.threshold = num_nodes;
c.migration_targets = vec![0];
})
.await;

// A0 (idx 0) is the source we migrate from; B0 lives at `num_nodes`
// since `migration_targets[0] = 0` puts the first target right after
// the initial participants.
// B0 lives at `num_nodes` because `migration_targets[0] = 0` places
// the first target right after the initial participants.
let (a_idx, b_idx) = (0usize, num_nodes);
assert_eq!(
cluster.nodes[a_idx].account_id().to_string(),
Expand All @@ -707,7 +694,6 @@ async fn migration_service__should_handle_back_migration_a_to_b_to_a() {
"migration target must have a different p2p key"
);

// Sanity: cluster is healthy before any migration.
let mut rng = rand::rngs::StdRng::seed_from_u64(42);
common::send_sign_request(&cluster, &running, &mut rng, cluster.default_user_account())
.await
Expand All @@ -716,87 +702,98 @@ async fn migration_service__should_handle_back_migration_a_to_b_to_a() {
.await
.expect("ckd request failed before forward migration");

// When: forward migration A0 → B0.
run_migration_round(&cluster, a_idx, b_idx, &backup_cli, "forward").await;

// Then: forward worked — B0 is the active participant, sign + ckd still
// succeed (B0 + A1).
common::send_sign_request(&cluster, &running, &mut rng, cluster.default_user_account())
.await
.expect("sign request failed after forward migration");
common::send_ckd_request(&cluster, &running, &mut rng, cluster.default_user_account())
.await
.expect("ckd request failed after forward migration");

// Simulate operator decommissioning the old node and bringing it back
// online later for the back-migration. `kill_nodes` + `start_nodes`
// preserves the node's home dir (keyshares stay on disk) — that disk
// state is the trigger for the suspected P1 early-return.
//
// Snapshot A0's keyshare-file listing before kill so we can verify
// it's still there after start. Keyshares live in
// `home_dir/permanent_keys/` (see `crates/node/src/keyshare/local.rs`).
let a0_keyshares_before = snapshot_permanent_keys(&cluster, a_idx);
assert!(
!a0_keyshares_before.is_empty(),
"A0's permanent_keys/ is empty before kill — test setup did not produce keyshares, \
so the back-migration P1 precondition can't be modeled"
);
// Snapshot A0's indexer height before the kill so the post-restart
// readiness check (below) can prove the indexer has resumed and
// advanced past where it was — not just bound the web port.
let a0_indexer_height_before_kill = common::current_node_indexer_height(&cluster, a_idx)
.await
.expect("failed to read A0's indexer height before kill")
.expect("A0's indexer should be reporting a block height before the kill");
cluster
.kill_nodes(&[a_idx])
.expect("failed to kill A0 before back-migration");
cluster
.start_nodes(&[a_idx])
.expect("failed to restart A0 for back-migration");
cluster
.wait_for_node_healthy(a_idx)
.await
.expect("A0 did not become healthy after restart");
// `wait_for_node_healthy` only checks the web port; wait for the
// indexer to actually resume before back-migration (see #3366).
common::wait_for_node_indexer_height_above(
&cluster,
a_idx,
a0_indexer_height_before_kill,
Duration::from_secs(60),
)
.await
.expect("A0's indexer did not resume + advance within 60s after restart");
// Validate the assumption the test rests on: A0's keyshares are
// still on disk after the restart. If they're missing or different,
// the test no longer models production's back-migration scenario
// (e.g. someone swapped `start_nodes` for `reset_and_start_nodes`,
// which wipes the home dir).
let a0_keyshares_after = snapshot_permanent_keys(&cluster, a_idx);
assert_eq!(
a0_keyshares_before, a0_keyshares_after,
"A0's permanent_keys/ contents changed across kill+start — \
keyshares were NOT preserved on disk"
);
match between {
BetweenDirections::None => {}
BetweenDirections::KillAndRestartA => {
// Keyshares live in `home_dir/permanent_keys/`
// (`crates/node/src/keyshare/local.rs`). Snapshot before/after
// the restart to verify they're preserved on disk.
let a0_keyshares_before = snapshot_permanent_keys(&cluster, a_idx);
assert!(
!a0_keyshares_before.is_empty(),
"A0's permanent_keys/ is empty before kill — test setup did not produce \
keyshares, so the back-migration P1 precondition can't be modeled"
);
let a0_indexer_height_before_kill =
common::current_node_indexer_height(&cluster, a_idx)
.await
.expect("failed to read A0's indexer height before kill")
.expect("A0's indexer should be reporting a block height before the kill");
cluster
.kill_nodes(&[a_idx])
.expect("failed to kill A0 before back-migration");
cluster
.start_nodes(&[a_idx])
.expect("failed to restart A0 for back-migration");
cluster
.wait_for_node_healthy(a_idx)
.await
.expect("A0 did not become healthy after restart");
// `wait_for_node_healthy` only checks the web port; wait for
// the indexer to actually resume before back-migration
// (see #3366).
common::wait_for_node_indexer_height_above(
&cluster,
a_idx,
a0_indexer_height_before_kill,
Duration::from_secs(60),
)
.await
.expect("A0's indexer did not resume + advance within 60s after restart");
let a0_keyshares_after = snapshot_permanent_keys(&cluster, a_idx);
assert_eq!(
a0_keyshares_before, a0_keyshares_after,
"A0's permanent_keys/ contents changed across kill+start — \
keyshares were NOT preserved on disk"
);
}
}

// When: back-migration B0 → A0 via the same helper, indices swapped.
run_migration_round(&cluster, b_idx, a_idx, &backup_cli, "back").await;

// Kill B0 before the final assertions so the sign + ckd requests cannot
// be served by the now-demoted node — that proves A0 + A1 are carrying
// the workload alone, mirroring what the forward test does by killing
// the source after migration.
// Kill B0 so the final assertions can't be served by the demoted node
// — proves A0 + A1 carry the workload alone.
cluster
.kill_nodes(&[b_idx])
.expect("failed to kill B0 after back-migration");

// Then: A0 + A1 are the participants again and sign + ckd still succeed.
common::send_sign_request(&cluster, &running, &mut rng, cluster.default_user_account())
.await
.expect("sign request failed after back-migration");
common::send_ckd_request(&cluster, &running, &mut rng, cluster.default_user_account())
.await
.expect("ckd request failed after back-migration");
}

/// Back-migration A → B → A with A0 staying alive between the two
/// directions (no kill+restart).
#[tokio::test]
#[expect(non_snake_case)]
async fn migration_service__should_handle_back_migration_a_to_b_to_a_no_restart() {
run_back_migration_round_trip(
common::MIGRATION_BACK_NO_RESTART_PORT_SEED,
BetweenDirections::None,
)
.await;
}

/// Back-migration A → B → A with A0 killed + restarted between the two
/// directions.
#[tokio::test]
#[expect(non_snake_case)]
async fn migration_service__should_handle_back_migration_a_to_b_to_a() {
run_back_migration_round_trip(
common::MIGRATION_BACK_PORT_SEED,
BetweenDirections::KillAndRestartA,
)
.await;
}
45 changes: 33 additions & 12 deletions crates/node/src/migration_service.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,10 @@
use std::{net::SocketAddr, sync::Arc};

use anyhow::Context as _;
use ed25519_dalek::SigningKey;
use near_account_id::AccountId;
use onboarding::onboard;
use tokio::sync::{RwLock, watch};
use tokio::sync::{RwLock, oneshot, watch};
use types::MigrationInfo;

use crate::{
Expand All @@ -30,14 +31,19 @@ impl From<&SecretsConfig> for MigrationSecrets {
}
}

/// Spawns the migration web server and the onboarding state machine,
/// then awaits the first `OnboardingJob::Done` before returning. The
/// onboarding loop keeps running in the background past that point so
/// future migrations can re-enter `Onboard(keyset)` without a restart.
/// See `docs/design/migration-onboarding-reentry.md`.
pub async fn spawn_recovery_server_and_run_onboarding(
migration_web_ui: SocketAddr,
migration_secrets: MigrationSecrets,
my_near_account_id: AccountId,
keyshare_storage: Arc<RwLock<KeyshareStorage>>,
my_migration_info_receiver: watch::Receiver<MigrationInfo>,
contract_state_receiver: watch::Receiver<ContractState>,
tx_sender: impl TransactionSender,
tx_sender: impl TransactionSender + 'static,
) -> anyhow::Result<()> {
let (import_keyshares_sender, import_keyshares_receiver) = tokio::sync::watch::channel(vec![]);
let web_server_state = web::types::WebServerState {
Expand All @@ -53,15 +59,30 @@ pub async fn spawn_recovery_server_and_run_onboarding(
&migration_secrets.p2p_private_key,
)
.await?;
onboard(
contract_state_receiver,
my_migration_info_receiver.clone(),
my_near_account_id.clone(),
migration_secrets.p2p_private_key.verifying_key(),
tx_sender,
keyshare_storage.clone(),
import_keyshares_receiver,
)
.await?;

let (first_done_tx, first_done_rx) = oneshot::channel();
let tls_public_key = migration_secrets.p2p_private_key.verifying_key();
tokio::spawn(async move {
if let Err(err) = onboard(
contract_state_receiver,
my_migration_info_receiver,
my_near_account_id,
tls_public_key,
tx_sender,
keyshare_storage,
import_keyshares_receiver,
Some(first_done_tx),
)
.await
{
tracing::error!(?err, "onboarding state machine exited unexpectedly");
}
});

// Propagate so startup fails if onboarding dies before signaling —
// otherwise the keyshare channel has no listener.
first_done_rx
.await
.context("onboarding task died before reaching first Done")?;
Ok(())
}
Loading
Loading