Fixing deployment plan

Author: nadav-govari

quickwit/quickwit-metastore/src/metastore/file_backed/file_backed_index/shards.rs

@@ -52,13 +52,20 @@ impl fmt::Debug for Shards {
 }
 
 /// Whether a shard recording `existing_token` can be acquired by a pipeline presenting
-/// `presented_token`. Acquisition is monotonic; a missing or legacy (`'/'`-containing) token ranks
-/// below any ULID.
+/// `presented_token`. Acquisition between ULIDs is monotonic (newer-or-equal wins). A legacy
+/// (`'/'`-containing, pre-ULID) presented token always wins, so a rolling upgrade can still hand a
+/// shard to an old indexer; a missing or legacy recorded token loses to any ULID.
 fn can_acquire_shard(existing_token: &str, presented_token: &str) -> bool {
+    // An old indexer presenting a legacy token keeps the pre-upgrade overwrite behavior.
+    if presented_token.contains('/') {
+        return true;
+    }
+    // A missing or legacy recorded token loses to any ULID.
     if existing_token.is_empty() || existing_token.contains('/') {
         return true;
     }
-    !presented_token.contains('/') && presented_token >= existing_token
+    // Both are ULIDs: acquire only if ours is newer-or-equal.
+    presented_token >= existing_token
 }
 
 impl Shards {
@@ -572,8 +579,9 @@ mod tests {
         assert!(can_acquire_shard(NEWER, NEWER));
         // An older ULID cannot steal a shard owned by a newer one.
         assert!(!can_acquire_shard(NEWER, OLDER));
-        // A legacy presented token cannot supersede a recorded ULID.
-        assert!(!can_acquire_shard(NEWER, LEGACY));
+        // A legacy presented token always wins, so a rolling upgrade can move a shard from a new
+        // indexer back to an old one.
+        assert!(can_acquire_shard(NEWER, LEGACY));
     }
 
     #[test]

quickwit/quickwit-metastore/src/metastore/postgres/metastore.rs

@@ -1436,11 +1436,11 @@ impl MetastoreService for PostgresqlMetastore {
             .bind(&request.publish_token)
             .fetch_all(&self.connection_pool)
             .await?;
-        
+
         if pg_shards.len() != request.shard_ids.len() {
             warn_on_unacquired_shards(&request, &pg_shards);
         }
-        
+
         let acquired_shards = pg_shards
             .into_iter()
             .map(|pg_shard| pg_shard.into())

quickwit/quickwit-metastore/src/metastore/postgres/queries/shards/acquire.sql

@@ -6,13 +6,13 @@ WHERE
     index_uid = $1
     AND source_id = $2
     AND shard_id = ANY ($3)
-    -- Monotonic acquisition: the presented token ($4) must rank >= the recorded one.
+    -- Acquisition is monotonic between ULIDs; a legacy presented token keeps pre-upgrade behavior.
     AND (
-        publish_token IS NULL -- never acquired: free to take
+        $4 LIKE '%/%' -- presented token is legacy (pre-ULID): always takes, for rolling upgrades
+        OR publish_token IS NULL -- never acquired: free to take
         OR publish_token = '' -- empty placeholder: free to take
-        OR publish_token LIKE '%/%' -- legacy pre-ULID token: ranks below any ULID, so superseded
-        -- both are ULIDs: take it only if ours is newer-or-equal (and ours is itself a ULID)
-        OR ($4 NOT LIKE '%/%' AND $4 >= publish_token)
+        OR publish_token LIKE '%/%' -- recorded token is legacy: superseded by any ULID
+        OR $4 >= publish_token -- both are ULIDs: take only if ours is newer-or-equal
     )
 RETURNING
     *

quickwit/quickwit-metastore/src/tests/shard.rs

@@ -229,6 +229,18 @@ pub async fn test_metastore_acquire_shards<
     )
     .await;
 
+    // Publish tokens are the ULID of the indexing plan that minted them; ULIDs are time-ordered, so
+    // lexicographic order is chronological order. A token containing '/' is the legacy pre-ULID
+    // format: it loses to a ULID when recorded, but always wins when presented, so a rolling
+    // upgrade can still hand a shard back to an old indexer.
+    const OLDER_TOKEN: &str = "01000000000000000000000000";
+    const TOKEN: &str = "02000000000000000000000000";
+    const NEWER_TOKEN: &str = "03000000000000000000000000";
+    const LEGACY_TOKEN: &str =
+        "indexer/test-node/test-index:0/test-source/01000000000000000000000000";
+
+    // Shard 1 owned by `TOKEN`, shard 2 unowned, shard 3 owned by a legacy token, shard 4 owned by
+    // `TOKEN`.
     let shards = vec![
         Shard {
             index_uid: Some(test_index.index_uid.clone()),
@@ -239,7 +251,7 @@ pub async fn test_metastore_acquire_shards<
             follower_id: Some("test-ingester-bar".to_string()),
             doc_mapping_uid: Some(DocMappingUid::default()),
             publish_position_inclusive: Some(Position::Beginning),
-            publish_token: Some("test-publish-token-foo".to_string()),
+            publish_token: Some(TOKEN.to_string()),
             update_timestamp: 1724158996,
         },
         Shard {
@@ -251,7 +263,7 @@ pub async fn test_metastore_acquire_shards<
             follower_id: Some("test-ingester-qux".to_string()),
             doc_mapping_uid: Some(DocMappingUid::default()),
             publish_position_inclusive: Some(Position::Beginning),
-            publish_token: Some("test-publish-token-bar".to_string()),
+            publish_token: None,
             update_timestamp: 1724158996,
         },
         Shard {
@@ -263,7 +275,7 @@ pub async fn test_metastore_acquire_shards<
             follower_id: Some("test-ingester-baz".to_string()),
             doc_mapping_uid: Some(DocMappingUid::default()),
             publish_position_inclusive: Some(Position::Beginning),
-            publish_token: None,
+            publish_token: Some(LEGACY_TOKEN.to_string()),
             update_timestamp: 1724158996,
         },
         Shard {
@@ -275,150 +287,113 @@ pub async fn test_metastore_acquire_shards<
             follower_id: Some("test-ingester-tux".to_string()),
             doc_mapping_uid: Some(DocMappingUid::default()),
             publish_position_inclusive: Some(Position::Beginning),
-            publish_token: None,
+            publish_token: Some(TOKEN.to_string()),
             update_timestamp: 1724158996,
         },
     ];
     metastore
         .insert_shards(&test_index.index_uid, &test_index.source_id, shards)
         .await;
 
-    // Test acquire shards.
-    let acquire_shards_request = AcquireShardsRequest {
-        index_uid: Some(test_index.index_uid.clone()),
-        source_id: test_index.source_id.clone(),
-        shard_ids: vec![
-            ShardId::from(1),
-            ShardId::from(2),
-            ShardId::from(3),
-            ShardId::from(666),
-        ], // shard 666 does not exist
-        publish_token: "test-publish-token-foo".to_string(),
-    };
-    let mut acquire_shards_response = metastore
-        .acquire_shards(acquire_shards_request)
-        .await
-        .unwrap();
-
-    acquire_shards_response
-        .acquired_shards
-        .sort_unstable_by(|left, right| left.shard_id().cmp(right.shard_id()));
-
-    let shard = &acquire_shards_response.acquired_shards[0];
-    assert_eq!(shard.index_uid(), &test_index.index_uid);
-    assert_eq!(shard.source_id, test_index.source_id);
-    assert_eq!(shard.shard_id(), ShardId::from(1));
-    assert_eq!(shard.shard_state(), ShardState::Closed);
-    assert_eq!(shard.leader_id, "test-ingester-foo");
-    assert_eq!(shard.follower_id(), "test-ingester-bar");
-    assert_eq!(shard.publish_position_inclusive(), Position::Beginning);
-    assert_eq!(shard.publish_token(), "test-publish-token-foo");
-
-    let shard = &acquire_shards_response.acquired_shards[1];
-    assert_eq!(shard.index_uid(), &test_index.index_uid);
-    assert_eq!(shard.source_id, test_index.source_id);
-    assert_eq!(shard.shard_id(), ShardId::from(2));
-    assert_eq!(shard.shard_state(), ShardState::Open);
-    assert_eq!(shard.leader_id, "test-ingester-bar");
-    assert_eq!(shard.follower_id(), "test-ingester-qux");
-    assert_eq!(shard.publish_position_inclusive(), Position::Beginning);
-    assert_eq!(shard.publish_token(), "test-publish-token-foo");
-
-    let shard = &acquire_shards_response.acquired_shards[2];
-    assert_eq!(shard.index_uid(), &test_index.index_uid);
-    assert_eq!(shard.source_id, test_index.source_id);
-    assert_eq!(shard.shard_id(), ShardId::from(3));
-    assert_eq!(shard.shard_state(), ShardState::Open);
-    assert_eq!(shard.leader_id, "test-ingester-qux");
-    assert_eq!(shard.follower_id(), "test-ingester-baz");
-    assert_eq!(shard.publish_position_inclusive(), Position::Beginning);
-    assert_eq!(shard.publish_token(), "test-publish-token-foo");
-
-    // A token ranking below the recorded one is refused: the shard is left untouched and absent
-    // from the response.
+    // A token ranking below the recorded one — and a non-existent shard — are refused: both are
+    // omitted from the response and the recorded token is left untouched.
     let acquire_shards_response = metastore
         .acquire_shards(AcquireShardsRequest {
             index_uid: Some(test_index.index_uid.clone()),
             source_id: test_index.source_id.clone(),
-            shard_ids: vec![ShardId::from(1)],
-            publish_token: "test-publish-token-aaa".to_string(),
+            shard_ids: vec![ShardId::from(1), ShardId::from(666)],
+            publish_token: OLDER_TOKEN.to_string(),
         })
         .await
         .unwrap();
     assert!(acquire_shards_response.acquired_shards.is_empty());
 
-    // The same token re-acquires successfully (idempotent, e.g. after a local respawn).
+    // The same token re-acquires successfully (idempotent, e.g. after a local respawn); the full
+    // shard is returned.
     let acquire_shards_response = metastore
         .acquire_shards(AcquireShardsRequest {
             index_uid: Some(test_index.index_uid.clone()),
             source_id: test_index.source_id.clone(),
             shard_ids: vec![ShardId::from(1)],
-            publish_token: "test-publish-token-foo".to_string(),
+            publish_token: TOKEN.to_string(),
         })
         .await
         .unwrap();
     assert_eq!(acquire_shards_response.acquired_shards.len(), 1);
-    assert_eq!(
-        acquire_shards_response.acquired_shards[0].publish_token(),
-        "test-publish-token-foo"
-    );
+    let shard = &acquire_shards_response.acquired_shards[0];
+    assert_eq!(shard.index_uid(), &test_index.index_uid);
+    assert_eq!(shard.source_id, test_index.source_id);
+    assert_eq!(shard.shard_id(), ShardId::from(1));
+    assert_eq!(shard.shard_state(), ShardState::Closed);
+    assert_eq!(shard.leader_id, "test-ingester-foo");
+    assert_eq!(shard.follower_id(), "test-ingester-bar");
+    assert_eq!(shard.publish_position_inclusive(), Position::Beginning);
+    assert_eq!(shard.publish_token(), TOKEN);
 
-    // A strictly greater token takes the shard over.
+    // A strictly newer ULID takes the shard over.
     let acquire_shards_response = metastore
         .acquire_shards(AcquireShardsRequest {
             index_uid: Some(test_index.index_uid.clone()),
             source_id: test_index.source_id.clone(),
             shard_ids: vec![ShardId::from(1)],
-            publish_token: "test-publish-token-zzz".to_string(),
+            publish_token: NEWER_TOKEN.to_string(),
         })
         .await
         .unwrap();
     assert_eq!(acquire_shards_response.acquired_shards.len(), 1);
     assert_eq!(
         acquire_shards_response.acquired_shards[0].publish_token(),
-        "test-publish-token-zzz"
+        NEWER_TOKEN
     );
 
-    // Legacy (pre-ULID, '/'-containing) tokens rank below any ULID. Shard 4 starts unowned: a
-    // legacy token can claim it, but is then superseded by a ULID and can no longer take it back.
-    let legacy_token = "indexer/test-node/test-index:0/test-source/01000000000000000000000000";
-    let ulid_token = "02000000000000000000000000";
+    // An unowned shard can be acquired by any ULID.
     let acquire_shards_response = metastore
         .acquire_shards(AcquireShardsRequest {
             index_uid: Some(test_index.index_uid.clone()),
             source_id: test_index.source_id.clone(),
-            shard_ids: vec![ShardId::from(4)],
-            publish_token: legacy_token.to_string(),
+            shard_ids: vec![ShardId::from(2)],
+            publish_token: TOKEN.to_string(),
         })
         .await
         .unwrap();
     assert_eq!(acquire_shards_response.acquired_shards.len(), 1);
+    assert_eq!(
+        acquire_shards_response.acquired_shards[0].publish_token(),
+        TOKEN
+    );
 
+    // A ULID supersedes a recorded legacy token.
     let acquire_shards_response = metastore
         .acquire_shards(AcquireShardsRequest {
             index_uid: Some(test_index.index_uid.clone()),
             source_id: test_index.source_id.clone(),
-            shard_ids: vec![ShardId::from(4)],
-            publish_token: ulid_token.to_string(),
+            shard_ids: vec![ShardId::from(3)],
+            publish_token: TOKEN.to_string(),
         })
         .await
         .unwrap();
     assert_eq!(acquire_shards_response.acquired_shards.len(), 1);
     assert_eq!(
         acquire_shards_response.acquired_shards[0].publish_token(),
-        ulid_token
+        TOKEN
     );
 
+    // Shard 4 is owned by a ULID, yet a legacy token reclaims it: the rolling-upgrade hand-back,
+    // where the control plane moves a shard from a new indexer back to an old one.
     let acquire_shards_response = metastore
         .acquire_shards(AcquireShardsRequest {
             index_uid: Some(test_index.index_uid.clone()),
             source_id: test_index.source_id.clone(),
             shard_ids: vec![ShardId::from(4)],
-            publish_token: legacy_token.to_string(),
+            publish_token: LEGACY_TOKEN.to_string(),
         })
         .await
         .unwrap();
-    assert!(acquire_shards_response.acquired_shards.is_empty());
+    assert_eq!(acquire_shards_response.acquired_shards.len(), 1);
+    assert_eq!(
+        acquire_shards_response.acquired_shards[0].publish_token(),
+        LEGACY_TOKEN
+    );
 
     cleanup_index(&mut metastore, test_index.index_uid).await;
 }