LCOV - code coverage report
Current view: top level - libs/utils/src - shard.rs (source / functions) Coverage Total Hit
Test: 6df3fc19ec669bcfbbf9aba41d1338898d24eaa0.info Lines: 79.1 % 249 197
Test Date: 2025-03-12 18:28:53 Functions: 42.4 % 92 39

            Line data    Source code
       1              : //! See `pageserver_api::shard` for description on sharding.
       2              : 
       3              : use std::ops::RangeInclusive;
       4              : use std::str::FromStr;
       5              : 
       6              : use hex::FromHex;
       7              : use serde::{Deserialize, Serialize};
       8              : 
       9              : use crate::id::TenantId;
      10              : 
      11            0 : #[derive(Ord, PartialOrd, Eq, PartialEq, Clone, Copy, Serialize, Deserialize, Debug, Hash)]
      12              : pub struct ShardNumber(pub u8);
      13              : 
      14            0 : #[derive(Ord, PartialOrd, Eq, PartialEq, Clone, Copy, Serialize, Deserialize, Debug, Hash)]
      15              : pub struct ShardCount(pub u8);
      16              : 
      17              : /// Combination of ShardNumber and ShardCount.
      18              : ///
      19              : /// For use within the context of a particular tenant, when we need to know which shard we're
      20              : /// dealing with, but do not need to know the full ShardIdentity (because we won't be doing
      21              : /// any page->shard mapping), and do not need to know the fully qualified TenantShardId.
      22              : #[derive(Eq, PartialEq, PartialOrd, Ord, Clone, Copy, Hash)]
      23              : pub struct ShardIndex {
      24              :     pub shard_number: ShardNumber,
      25              :     pub shard_count: ShardCount,
      26              : }
      27              : 
      28              : /// Formatting helper, for generating the `shard_id` label in traces.
      29              : pub struct ShardSlug<'a>(&'a TenantShardId);
      30              : 
      31              : /// TenantShardId globally identifies a particular shard in a particular tenant.
      32              : ///
      33              : /// These are written as `<TenantId>-<ShardSlug>`, for example:
      34              : ///   # The second shard in a two-shard tenant
      35              : ///   072f1291a5310026820b2fe4b2968934-0102
      36              : ///
      37              : /// If the `ShardCount` is _unsharded_, the `TenantShardId` is written without
      38              : /// a shard suffix and is equivalent to the encoding of a `TenantId`: this enables
      39              : /// an unsharded [`TenantShardId`] to be used interchangably with a [`TenantId`].
      40              : ///
      41              : /// The human-readable encoding of an unsharded TenantShardId, such as used in API URLs,
      42              : /// is both forward and backward compatible with TenantId: a legacy TenantId can be
      43              : /// decoded as a TenantShardId, and when re-encoded it will be parseable
      44              : /// as a TenantId.
      45              : #[derive(Eq, PartialEq, PartialOrd, Ord, Clone, Copy, Hash)]
      46              : pub struct TenantShardId {
      47              :     pub tenant_id: TenantId,
      48              :     pub shard_number: ShardNumber,
      49              :     pub shard_count: ShardCount,
      50              : }
      51              : 
      52              : impl ShardCount {
      53              :     pub const MAX: Self = Self(u8::MAX);
      54              :     pub const MIN: Self = Self(0);
      55              : 
      56              :     /// The internal value of a ShardCount may be zero, which means "1 shard, but use
      57              :     /// legacy format for TenantShardId that excludes the shard suffix", also known
      58              :     /// as [`TenantShardId::unsharded`].
      59              :     ///
      60              :     /// This method returns the actual number of shards, i.e. if our internal value is
      61              :     /// zero, we return 1 (unsharded tenants have 1 shard).
      62      9613443 :     pub fn count(&self) -> u8 {
      63      9613443 :         if self.0 > 0 { self.0 } else { 1 }
      64      9613443 :     }
      65              : 
      66              :     /// The literal internal value: this is **not** the number of shards in the
      67              :     /// tenant, as we have a special zero value for legacy unsharded tenants.  Use
      68              :     /// [`Self::count`] if you want to know the cardinality of shards.
      69            2 :     pub fn literal(&self) -> u8 {
      70            2 :         self.0
      71            2 :     }
      72              : 
      73              :     /// Whether the `ShardCount` is for an unsharded tenant, so uses one shard but
      74              :     /// uses the legacy format for `TenantShardId`. See also the documentation for
      75              :     /// [`Self::count`].
      76            0 :     pub fn is_unsharded(&self) -> bool {
      77            0 :         self.0 == 0
      78            0 :     }
      79              : 
      80              :     /// `v` may be zero, or the number of shards in the tenant.  `v` is what
      81              :     /// [`Self::literal`] would return.
      82        10015 :     pub const fn new(val: u8) -> Self {
      83        10015 :         Self(val)
      84        10015 :     }
      85              : }
      86              : 
      87              : impl ShardNumber {
      88              :     pub const MAX: Self = Self(u8::MAX);
      89              : }
      90              : 
      91              : impl TenantShardId {
      92          121 :     pub fn unsharded(tenant_id: TenantId) -> Self {
      93          121 :         Self {
      94          121 :             tenant_id,
      95          121 :             shard_number: ShardNumber(0),
      96          121 :             shard_count: ShardCount(0),
      97          121 :         }
      98          121 :     }
      99              : 
     100              :     /// The range of all TenantShardId that belong to a particular TenantId.  This is useful when
     101              :     /// you have a BTreeMap of TenantShardId, and are querying by TenantId.
     102            0 :     pub fn tenant_range(tenant_id: TenantId) -> RangeInclusive<Self> {
     103            0 :         RangeInclusive::new(
     104            0 :             Self {
     105            0 :                 tenant_id,
     106            0 :                 shard_number: ShardNumber(0),
     107            0 :                 shard_count: ShardCount(0),
     108            0 :             },
     109            0 :             Self {
     110            0 :                 tenant_id,
     111            0 :                 shard_number: ShardNumber::MAX,
     112            0 :                 shard_count: ShardCount::MAX,
     113            0 :             },
     114            0 :         )
     115            0 :     }
     116              : 
     117            0 :     pub fn range(&self) -> RangeInclusive<Self> {
     118            0 :         RangeInclusive::new(*self, *self)
     119            0 :     }
     120              : 
     121        31531 :     pub fn shard_slug(&self) -> impl std::fmt::Display + '_ {
     122        31531 :         ShardSlug(self)
     123        31531 :     }
     124              : 
     125              :     /// Convenience for code that has special behavior on the 0th shard.
     126         1146 :     pub fn is_shard_zero(&self) -> bool {
     127         1146 :         self.shard_number == ShardNumber(0)
     128         1146 :     }
     129              : 
     130              :     /// The "unsharded" value is distinct from simply having a single shard: it represents
     131              :     /// a tenant which is not shard-aware at all, and whose storage paths will not include
     132              :     /// a shard suffix.
     133            0 :     pub fn is_unsharded(&self) -> bool {
     134            0 :         self.shard_number == ShardNumber(0) && self.shard_count.is_unsharded()
     135            0 :     }
     136              : 
     137              :     /// Convenience for dropping the tenant_id and just getting the ShardIndex: this
     138              :     /// is useful when logging from code that is already in a span that includes tenant ID, to
     139              :     /// keep messages reasonably terse.
     140            0 :     pub fn to_index(&self) -> ShardIndex {
     141            0 :         ShardIndex {
     142            0 :             shard_number: self.shard_number,
     143            0 :             shard_count: self.shard_count,
     144            0 :         }
     145            0 :     }
     146              : 
     147              :     /// Calculate the children of this TenantShardId when splitting the overall tenant into
     148              :     /// the given number of shards.
     149            8 :     pub fn split(&self, new_shard_count: ShardCount) -> Vec<TenantShardId> {
     150            8 :         let effective_old_shard_count = std::cmp::max(self.shard_count.0, 1);
     151            8 :         let mut child_shards = Vec::new();
     152           48 :         for shard_number in 0..ShardNumber(new_shard_count.0).0 {
     153              :             // Key mapping is based on a round robin mapping of key hash modulo shard count,
     154              :             // so our child shards are the ones which the same keys would map to.
     155           48 :             if shard_number % effective_old_shard_count == self.shard_number.0 {
     156           44 :                 child_shards.push(TenantShardId {
     157           44 :                     tenant_id: self.tenant_id,
     158           44 :                     shard_number: ShardNumber(shard_number),
     159           44 :                     shard_count: new_shard_count,
     160           44 :                 })
     161            4 :             }
     162              :         }
     163              : 
     164            8 :         child_shards
     165            8 :     }
     166              : }
     167              : 
     168              : impl std::fmt::Display for ShardNumber {
     169            0 :     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
     170            0 :         self.0.fmt(f)
     171            0 :     }
     172              : }
     173              : 
     174              : impl std::fmt::Display for ShardSlug<'_> {
     175        18616 :     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
     176        18616 :         write!(
     177        18616 :             f,
     178        18616 :             "{:02x}{:02x}",
     179        18616 :             self.0.shard_number.0, self.0.shard_count.0
     180        18616 :         )
     181        18616 :     }
     182              : }
     183              : 
     184              : impl std::fmt::Display for TenantShardId {
     185        45065 :     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
     186        45065 :         if self.shard_count != ShardCount(0) {
     187          321 :             write!(f, "{}-{}", self.tenant_id, self.shard_slug())
     188              :         } else {
     189              :             // Legacy case (shard_count == 0) -- format as just the tenant id.  Note that this
     190              :             // is distinct from the normal single shard case (shard count == 1).
     191        44744 :             self.tenant_id.fmt(f)
     192              :         }
     193        45065 :     }
     194              : }
     195              : 
     196              : impl std::fmt::Debug for TenantShardId {
     197        22724 :     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
     198        22724 :         // Debug is the same as Display: the compact hex representation
     199        22724 :         write!(f, "{}", self)
     200        22724 :     }
     201              : }
     202              : 
     203              : impl std::str::FromStr for TenantShardId {
     204              :     type Err = hex::FromHexError;
     205              : 
     206           46 :     fn from_str(s: &str) -> Result<Self, Self::Err> {
     207           46 :         // Expect format: 16 byte TenantId, '-', 1 byte shard number, 1 byte shard count
     208           46 :         if s.len() == 32 {
     209              :             // Legacy case: no shard specified
     210              :             Ok(Self {
     211           40 :                 tenant_id: TenantId::from_str(s)?,
     212           40 :                 shard_number: ShardNumber(0),
     213           40 :                 shard_count: ShardCount(0),
     214              :             })
     215            6 :         } else if s.len() == 37 {
     216            6 :             let bytes = s.as_bytes();
     217            6 :             let tenant_id = TenantId::from_hex(&bytes[0..32])?;
     218            6 :             let mut shard_parts: [u8; 2] = [0u8; 2];
     219            6 :             hex::decode_to_slice(&bytes[33..37], &mut shard_parts)?;
     220            6 :             Ok(Self {
     221            6 :                 tenant_id,
     222            6 :                 shard_number: ShardNumber(shard_parts[0]),
     223            6 :                 shard_count: ShardCount(shard_parts[1]),
     224            6 :             })
     225              :         } else {
     226            0 :             Err(hex::FromHexError::InvalidStringLength)
     227              :         }
     228           46 :     }
     229              : }
     230              : 
     231              : impl From<[u8; 18]> for TenantShardId {
     232           94 :     fn from(b: [u8; 18]) -> Self {
     233           94 :         let tenant_id_bytes: [u8; 16] = b[0..16].try_into().unwrap();
     234           94 : 
     235           94 :         Self {
     236           94 :             tenant_id: TenantId::from(tenant_id_bytes),
     237           94 :             shard_number: ShardNumber(b[16]),
     238           94 :             shard_count: ShardCount(b[17]),
     239           94 :         }
     240           94 :     }
     241              : }
     242              : 
     243              : impl ShardIndex {
     244           28 :     pub fn new(number: ShardNumber, count: ShardCount) -> Self {
     245           28 :         Self {
     246           28 :             shard_number: number,
     247           28 :             shard_count: count,
     248           28 :         }
     249           28 :     }
     250          316 :     pub fn unsharded() -> Self {
     251          316 :         Self {
     252          316 :             shard_number: ShardNumber(0),
     253          316 :             shard_count: ShardCount(0),
     254          316 :         }
     255          316 :     }
     256              : 
     257              :     /// The "unsharded" value is distinct from simply having a single shard: it represents
     258              :     /// a tenant which is not shard-aware at all, and whose storage paths will not include
     259              :     /// a shard suffix.
     260       148493 :     pub fn is_unsharded(&self) -> bool {
     261       148493 :         self.shard_number == ShardNumber(0) && self.shard_count == ShardCount(0)
     262       148493 :     }
     263              : 
     264              :     /// For use in constructing remote storage paths: concatenate this with a TenantId
     265              :     /// to get a fully qualified TenantShardId.
     266              :     ///
     267              :     /// Backward compat: this function returns an empty string if Self::is_unsharded, such
     268              :     /// that the legacy pre-sharding remote key format is preserved.
     269         4019 :     pub fn get_suffix(&self) -> String {
     270         4019 :         if self.is_unsharded() {
     271         4003 :             "".to_string()
     272              :         } else {
     273           16 :             format!("-{:02x}{:02x}", self.shard_number.0, self.shard_count.0)
     274              :         }
     275         4019 :     }
     276              : }
     277              : 
     278              : impl std::fmt::Display for ShardIndex {
     279         4541 :     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
     280         4541 :         write!(f, "{:02x}{:02x}", self.shard_number.0, self.shard_count.0)
     281         4541 :     }
     282              : }
     283              : 
     284              : impl std::fmt::Debug for ShardIndex {
     285         3484 :     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
     286         3484 :         // Debug is the same as Display: the compact hex representation
     287         3484 :         write!(f, "{}", self)
     288         3484 :     }
     289              : }
     290              : 
     291              : impl std::str::FromStr for ShardIndex {
     292              :     type Err = hex::FromHexError;
     293              : 
     294         6257 :     fn from_str(s: &str) -> Result<Self, Self::Err> {
     295         6257 :         // Expect format: 1 byte shard number, 1 byte shard count
     296         6257 :         if s.len() == 4 {
     297         6257 :             let bytes = s.as_bytes();
     298         6257 :             let mut shard_parts: [u8; 2] = [0u8; 2];
     299         6257 :             hex::decode_to_slice(bytes, &mut shard_parts)?;
     300         6257 :             Ok(Self {
     301         6257 :                 shard_number: ShardNumber(shard_parts[0]),
     302         6257 :                 shard_count: ShardCount(shard_parts[1]),
     303         6257 :             })
     304              :         } else {
     305            0 :             Err(hex::FromHexError::InvalidStringLength)
     306              :         }
     307         6257 :     }
     308              : }
     309              : 
     310              : impl From<[u8; 2]> for ShardIndex {
     311            1 :     fn from(b: [u8; 2]) -> Self {
     312            1 :         Self {
     313            1 :             shard_number: ShardNumber(b[0]),
     314            1 :             shard_count: ShardCount(b[1]),
     315            1 :         }
     316            1 :     }
     317              : }
     318              : 
     319              : impl Serialize for TenantShardId {
     320           86 :     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
     321           86 :     where
     322           86 :         S: serde::Serializer,
     323           86 :     {
     324           86 :         if serializer.is_human_readable() {
     325           82 :             serializer.collect_str(self)
     326              :         } else {
     327              :             // Note: while human encoding of [`TenantShardId`] is backward and forward
     328              :             // compatible, this binary encoding is not.
     329            4 :             let mut packed: [u8; 18] = [0; 18];
     330            4 :             packed[0..16].clone_from_slice(&self.tenant_id.as_arr());
     331            4 :             packed[16] = self.shard_number.0;
     332            4 :             packed[17] = self.shard_count.0;
     333            4 : 
     334            4 :             packed.serialize(serializer)
     335              :         }
     336            0 :     }
     337              : }
     338              : 
     339              : impl<'de> Deserialize<'de> for TenantShardId {
     340           15 :     fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
     341           15 :     where
     342           15 :         D: serde::Deserializer<'de>,
     343           15 :     {
     344              :         struct IdVisitor {
     345              :             is_human_readable_deserializer: bool,
     346              :         }
     347              : 
     348              :         impl<'de> serde::de::Visitor<'de> for IdVisitor {
     349              :             type Value = TenantShardId;
     350              : 
     351            0 :             fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
     352            0 :                 if self.is_human_readable_deserializer {
     353            0 :                     formatter.write_str("value in form of hex string")
     354              :                 } else {
     355            0 :                     formatter.write_str("value in form of integer array([u8; 18])")
     356              :                 }
     357            0 :             }
     358              : 
     359            2 :             fn visit_seq<A>(self, seq: A) -> Result<Self::Value, A::Error>
     360            2 :             where
     361            2 :                 A: serde::de::SeqAccess<'de>,
     362            2 :             {
     363            2 :                 let s = serde::de::value::SeqAccessDeserializer::new(seq);
     364            2 :                 let id: [u8; 18] = Deserialize::deserialize(s)?;
     365            2 :                 Ok(TenantShardId::from(id))
     366            0 :             }
     367              : 
     368           13 :             fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
     369           13 :             where
     370           13 :                 E: serde::de::Error,
     371           13 :             {
     372           13 :                 TenantShardId::from_str(v).map_err(E::custom)
     373           13 :             }
     374              :         }
     375              : 
     376           15 :         if deserializer.is_human_readable() {
     377           13 :             deserializer.deserialize_str(IdVisitor {
     378           13 :                 is_human_readable_deserializer: true,
     379           13 :             })
     380              :         } else {
     381            2 :             deserializer.deserialize_tuple(
     382            2 :                 18,
     383            2 :                 IdVisitor {
     384            2 :                     is_human_readable_deserializer: false,
     385            2 :                 },
     386            2 :             )
     387              :         }
     388            0 :     }
     389              : }
     390              : 
     391              : impl Serialize for ShardIndex {
     392           34 :     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
     393           34 :     where
     394           34 :         S: serde::Serializer,
     395           34 :     {
     396           34 :         if serializer.is_human_readable() {
     397           32 :             serializer.collect_str(self)
     398              :         } else {
     399              :             // Binary encoding is not used in index_part.json, but is included in anticipation of
     400              :             // switching various structures (e.g. inter-process communication, remote metadata) to more
     401              :             // compact binary encodings in future.
     402            2 :             let mut packed: [u8; 2] = [0; 2];
     403            2 :             packed[0] = self.shard_number.0;
     404            2 :             packed[1] = self.shard_count.0;
     405            2 :             packed.serialize(serializer)
     406              :         }
     407            0 :     }
     408              : }
     409              : 
     410              : impl<'de> Deserialize<'de> for ShardIndex {
     411         6257 :     fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
     412         6257 :     where
     413         6257 :         D: serde::Deserializer<'de>,
     414         6257 :     {
     415              :         struct IdVisitor {
     416              :             is_human_readable_deserializer: bool,
     417              :         }
     418              : 
     419              :         impl<'de> serde::de::Visitor<'de> for IdVisitor {
     420              :             type Value = ShardIndex;
     421              : 
     422            0 :             fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
     423            0 :                 if self.is_human_readable_deserializer {
     424            0 :                     formatter.write_str("value in form of hex string")
     425              :                 } else {
     426            0 :                     formatter.write_str("value in form of integer array([u8; 2])")
     427              :                 }
     428            0 :             }
     429              : 
     430            1 :             fn visit_seq<A>(self, seq: A) -> Result<Self::Value, A::Error>
     431            1 :             where
     432            1 :                 A: serde::de::SeqAccess<'de>,
     433            1 :             {
     434            1 :                 let s = serde::de::value::SeqAccessDeserializer::new(seq);
     435            1 :                 let id: [u8; 2] = Deserialize::deserialize(s)?;
     436            1 :                 Ok(ShardIndex::from(id))
     437            0 :             }
     438              : 
     439         6256 :             fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
     440         6256 :             where
     441         6256 :                 E: serde::de::Error,
     442         6256 :             {
     443         6256 :                 ShardIndex::from_str(v).map_err(E::custom)
     444         6256 :             }
     445              :         }
     446              : 
     447         6257 :         if deserializer.is_human_readable() {
     448         6256 :             deserializer.deserialize_str(IdVisitor {
     449         6256 :                 is_human_readable_deserializer: true,
     450         6256 :             })
     451              :         } else {
     452            1 :             deserializer.deserialize_tuple(
     453            1 :                 2,
     454            1 :                 IdVisitor {
     455            1 :                     is_human_readable_deserializer: false,
     456            1 :                 },
     457            1 :             )
     458              :         }
     459            0 :     }
     460              : }
        

Generated by: LCOV version 2.1-beta