Skip to content

Add Slot migration testing support - #15

Open
Fniakate8 wants to merge 6 commits into
valkey-io:unstablefrom
Fniakate8:slot-migration
Open

Add Slot migration testing support #15
Fniakate8 wants to merge 6 commits into
valkey-io:unstablefrom
Fniakate8:slot-migration

Conversation

@Fniakate8

@Fniakate8 Fniakate8 commented Aug 13, 2026

Copy link
Copy Markdown

Add slot migration testing support

This PR adds a helper so cluster tests can move a slot (and the keys in it) from one node to another on a live cluster, without hand-writing the actual CLUSTER SETSLOT / GETKEYSINSLOT / MIGRATE protocol each time.

What's Added

  • ClusterTestCase.migrate_slot(source, target, slot, dbs=(0,)) — runs the manual migration protocol end to end: mark the slot IMPORTING on the target and
    MIGRATING on the source, batch-move every key with MIGRATE ... KEYS, then announce the new owner on every primary.
  • ClusterTestCase.get_slot_owner(slot) — the node that currently owns a slot.
  • ClusterTestCase.wait_for_slot_owner(slot, owner) — polls until every node agrees on the new owner (no fixed sleeps).
  • Supporting primitives on ClusterNodeHandle: count_keys_in_slot, start_importing_slot, start_migrating_slot, assign_slot_owner, get_slot_owner_id, is_primary.
  • tests/test_slot_migration.py — 7 tests: single-key move, many-key drain (250 keys), multi-DB, caller-DB isolation, cluster-wide ownership, empty slot, and
    replica shards.
  • README: a slot migration subsection under the existing "Testing Cluster Mode Enabled" docs.

Why this is needed

Slot migration is a core cluster operation (adding/removing nodes, rebalancing), but the framework has no helper for it — modules that want to test it hand-roll ~20 lines of low-level CLUSTER commands today (e.g. Search's test_multidb_slot_migration_CME). This gives them one call and handles the fiddly parts (batch key move, draining slots with >100 keys, per-DB selection, primary-only ownership announcement).

How to use

Once a slot's owner is known, migrating it is two calls:

self.migrate_slot(source_node, target_node, slot) # move slot + its keys
self.wait_for_slot_owner(slot, target_node) # wait until the cluster agrees

Use self.get_slot_owner(slot) to find the current owner, and pass dbs=(0, 1, ...) to migrate keys across multiple databases. Migrating any DB other than 0, requires the cluster to be started withcluster-databases > 1(plain cluster mode only has DB 0).

Design notes

  • Uses the legacy SETSLOT/MIGRATE protocol (what Search's real migration test uses today) rather than native CLUSTER MIGRATESLOTS, which has no current consumer.
  • Batch MIGRATE ... KEYS + a drain loop so all keys move, not just the first 100.
  • Key-moving runs on a dedicated connection so the caller's client keeps its own selected DB.
  • Ownership is announced only on primaries — CLUSTER SETSLOT is rejected on replicas (they learn from their primary).

Test plan

  • Single-key and many-key (250) migrations move all keys to the target
  • Multi-DB migration moves keys across databases (with cluster-databases)
  • migrate_slot does not disturb the caller's selected DB
  • All nodes agree on the new owner after migration
  • Migration works on shards with replicas, and the key reaches the target's replica

detemmienation and others added 6 commits June 10, 2026 22:31
Signed-off-by: Tracy <yuningt@amazon.com>
Signed-off-by: Tracy <yuningt@amazon.com>
Signed-off-by: Tracy <yuningt@amazon.com>
Signed-off-by: Tracy <yuningt@amazon.com>
Signed-off-by: Tracy <yuningt@amazon.com>
Add ClusterTestCase.migrate_slot() so cluster tests can move a slot (and its
keys) between live nodes without hand-writing the CLUSTER SETSLOT /
GETKEYSINSLOT / MIGRATE protocol each time.

- ClusterTestCase: migrate_slot (mark IMPORTING/MIGRATING, batch-move keys with
  MIGRATE KEYS, announce the new owner on every primary), get_slot_owner,
  wait_for_slot_owner (polls until all nodes agree, no fixed sleeps)
- ClusterNodeHandle primitives: count_keys_in_slot, start_importing_slot,
  start_migrating_slot, assign_slot_owner, get_slot_owner_id, is_primary
- Ownership is announced only on primaries (CLUSTER SETSLOT is rejected on
  replicas; they learn from their primary)
- Multi-DB migration supported via dbs=...; requires cluster-databases > 1
- tests/test_slot_migration.py: single-key, many-key drain, multi-DB, caller-DB
  isolation, cluster-wide ownership, empty slot, and replica shards
- README: slot migration subsection under the CME docs

Built on the ClusterTestCase/ClusterNodeHandle helpers from valkey-io#12.

Signed-off-by: Fanta Niakate <niakatf@amazon.com>

@zackcam zackcam left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good to me at a high level overall!

target = next(n for n in self.nodes if n.nodeid != source.nodeid)
return source, target

def test_migrate_slot_moves_key(self):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

could maybe combine this and the test below mostly achieve the same purpose with the seocnd one being just more keys

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants