From c7717c85c7a45144b016e5d799125197ba14c001 Mon Sep 17 00:00:00 2001
From: Erik Grinaker <erik@neon.tech>
Date: Wed, 12 Mar 2025 16:16:54 +0100
Subject: [PATCH] storcon,pageserver: use persisted stripe size when loading
 unsharded tenants (#11193)

## Problem

When the storage controller and Pageserver loads tenants from persisted
storage, it uses `ShardIdentity::unsharded()` for unsharded tenants.
However, this replaces the persisted stripe size of unsharded tenants
with the default stripe size.

This doesn't really matter for practical purposes, since the stripe size
is meaningless for unsharded tenants anyway, but can cause consistency
check failures if the persisted stripe size differs from the default.
This was seen in #11168, where we change the default stripe size.

Touches #11168.

## Summary of changes

Carry over the persisted stripe size from `TenantShardPersistence` for
unsharded tenants, and from `LocationConf` on Pageservers.

Also add bounds checks for type casts when loading persisted shard
metadata.
---
 libs/pageserver_api/src/shard.rs      | 10 +++++++
 pageserver/src/tenant/config.rs       |  6 +++-
 storage_controller/src/persistence.rs | 40 ++++++++++++++++++++++-----
 3 files changed, 48 insertions(+), 8 deletions(-)

diff --git a/libs/pageserver_api/src/shard.rs b/libs/pageserver_api/src/shard.rs
index eca04b1f3d..8386d6e586 100644
--- a/libs/pageserver_api/src/shard.rs
+++ b/libs/pageserver_api/src/shard.rs
@@ -112,6 +112,16 @@ impl ShardIdentity {
         }
     }
 
+    /// An unsharded identity with the given stripe size (if non-zero). This is typically used to
+    /// carry over a stripe size for an unsharded tenant from persistent storage.
+    pub fn unsharded_with_stripe_size(stripe_size: ShardStripeSize) -> Self {
+        let mut shard_identity = Self::unsharded();
+        if stripe_size.0 > 0 {
+            shard_identity.stripe_size = stripe_size;
+        }
+        shard_identity
+    }
+
     /// A broken instance of this type is only used for `TenantState::Broken` tenants,
     /// which are constructed in code paths that don't have access to proper configuration.
     ///
diff --git a/pageserver/src/tenant/config.rs b/pageserver/src/tenant/config.rs
index 334fb04604..4308db84e5 100644
--- a/pageserver/src/tenant/config.rs
+++ b/pageserver/src/tenant/config.rs
@@ -219,7 +219,11 @@ impl LocationConf {
         };
 
         let shard = if conf.shard_count == 0 {
-            ShardIdentity::unsharded()
+            // NB: carry over the persisted stripe size instead of using the default. This doesn't
+            // matter for most practical purposes, since unsharded tenants don't use the stripe
+            // size, but can cause inconsistencies between storcon and Pageserver and cause manual
+            // splits without `new_stripe_size` to use an unintended stripe size.
+            ShardIdentity::unsharded_with_stripe_size(ShardStripeSize(conf.shard_stripe_size))
         } else {
             ShardIdentity::new(
                 ShardNumber(conf.shard_number),
diff --git a/storage_controller/src/persistence.rs b/storage_controller/src/persistence.rs
index 5146fe472e..4a97aac125 100644
--- a/storage_controller/src/persistence.rs
+++ b/storage_controller/src/persistence.rs
@@ -1613,23 +1613,49 @@ pub(crate) struct TenantShardPersistence {
 }
 
 impl TenantShardPersistence {
+    fn get_shard_count(&self) -> Result<ShardCount, ShardConfigError> {
+        self.shard_count
+            .try_into()
+            .map(ShardCount)
+            .map_err(|_| ShardConfigError::InvalidCount)
+    }
+
+    fn get_shard_number(&self) -> Result<ShardNumber, ShardConfigError> {
+        self.shard_number
+            .try_into()
+            .map(ShardNumber)
+            .map_err(|_| ShardConfigError::InvalidNumber)
+    }
+
+    fn get_stripe_size(&self) -> Result<ShardStripeSize, ShardConfigError> {
+        self.shard_stripe_size
+            .try_into()
+            .map(ShardStripeSize)
+            .map_err(|_| ShardConfigError::InvalidStripeSize)
+    }
+
     pub(crate) fn get_shard_identity(&self) -> Result<ShardIdentity, ShardConfigError> {
         if self.shard_count == 0 {
-            Ok(ShardIdentity::unsharded())
+            // NB: carry over the stripe size from the persisted record, to avoid consistency check
+            // failures if the persisted value differs from the default stripe size. The stripe size
+            // doesn't really matter for unsharded tenants anyway.
+            Ok(ShardIdentity::unsharded_with_stripe_size(
+                self.get_stripe_size()?,
+            ))
         } else {
             Ok(ShardIdentity::new(
-                ShardNumber(self.shard_number as u8),
-                ShardCount::new(self.shard_count as u8),
-                ShardStripeSize(self.shard_stripe_size as u32),
+                self.get_shard_number()?,
+                self.get_shard_count()?,
+                self.get_stripe_size()?,
             )?)
         }
     }
 
-    pub(crate) fn get_tenant_shard_id(&self) -> Result<TenantShardId, hex::FromHexError> {
+    pub(crate) fn get_tenant_shard_id(&self) -> anyhow::Result<TenantShardId> {
         Ok(TenantShardId {
             tenant_id: TenantId::from_str(self.tenant_id.as_str())?,
-            shard_number: ShardNumber(self.shard_number as u8),
-            shard_count: ShardCount::new(self.shard_count as u8),
+            shard_number: self.get_shard_number()?,
+            shard_count: self.get_shard_count()?,
         })
     }
 }