picodata_driver/

topology.rs

use std::{
    collections::{HashMap, hash_map::Entry},
    sync::{
        Arc,
        atomic::{AtomicUsize, Ordering},
    },
};

use serde::Deserialize;
use uuid::Uuid;

use crate::instance::{Instance, State};

/// Global counter for assigning unique pool slot indices to instances.
static NEXT_IDX: AtomicUsize = AtomicUsize::new(0);

/// Name of a connection type in `_pico_peer_address`, such as `pgproto` or `iproto`.
type ConnectionType = String;

/// A single bucket range entry from `_pico_bucket()` response.
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct BucketRangeEntry {
    pub tier_name: String,
    pub bucket_id_start: u64,
    pub bucket_id_end: u64,
    pub current_replicaset_uuid: Uuid,
}

/// The cluster topology as handed out by [`crate::Client::topology`].
pub mod view {
    use std::{collections::HashMap, ops::RangeInclusive};

    use uuid::Uuid;

    use super::ConnectionType;
    use crate::error::{Error, IncompleteBucketMap};

    /// Cluster topology: tiers with their replicasets and bucket ranges.
    #[derive(Debug, Clone, PartialEq, Eq)]
    pub struct Topology {
        pub(super) tiers: Vec<Tier>,
    }

    /// A tier with its replicasets.
    #[derive(Debug, Clone, PartialEq, Eq)]
    pub struct Tier {
        pub(super) name: String,
        pub(super) bucket_count: u64,
        pub(super) replicasets: Vec<Replicaset>,
    }

    /// A replicaset with its instances and the buckets it owns.
    #[derive(Debug, Clone, PartialEq, Eq)]
    pub struct Replicaset {
        pub(super) uuid: Uuid,
        pub(super) name: String,
        pub(super) master: Option<Instance>,
        pub(super) replicas: Vec<Instance>,
        pub(super) ranges: Vec<RangeInclusive<u64>>,
    }

    /// An instance of a replicaset.
    #[derive(Debug, Clone, PartialEq, Eq)]
    pub struct Instance {
        pub(super) name: String,
        /// The address of each connection type the instance accepts.
        pub(super) addresses: HashMap<ConnectionType, String>,
        pub(super) online: bool,
        pub(super) is_master: bool,
        pub(super) raft_id: u64,
        pub(super) uuid: Uuid,
    }

    /// Bucket ranges of a tier by the UUID of the replicaset owning them.
    pub type BucketMap<'a> = HashMap<Uuid, &'a [RangeInclusive<u64>]>;

    impl Topology {
        /// The cluster's tiers.
        #[must_use]
        pub fn tiers(&self) -> &[Tier] {
            &self.tiers
        }
    }

    impl Tier {
        /// Tier name.
        #[must_use]
        pub fn name(&self) -> &str {
            &self.name
        }

        /// The tier's replicasets.
        #[must_use]
        pub fn replicasets(&self) -> &[Replicaset] {
            &self.replicasets
        }

        /// Number of buckets the tier is configured with.
        #[must_use]
        pub fn bucket_count(&self) -> u64 {
            self.bucket_count
        }

        /// Bucket ranges by the UUID of the replicaset owning them.
        ///
        /// # Errors
        ///
        /// Returns [`Error::IncompleteBucketMap`] if the ranges leave some of the tier's
        /// buckets uncovered.
        pub fn bucket_map(&self) -> Result<BucketMap<'_>, Error> {
            let mut ranges: Vec<_> = self.replicasets.iter().flat_map(|rs| &rs.ranges).collect();
            ranges.sort_unstable_by_key(|r| *r.start());
            // Buckets covered without a gap, counting from bucket 1.
            let mut covered = 0;
            for r in ranges {
                if *r.start() != covered + 1 || r.is_empty() {
                    break;
                }
                covered = *r.end();
            }
            if covered != self.bucket_count {
                return Err(Error::IncompleteBucketMap(Box::new(IncompleteBucketMap {
                    tier: self.name.clone(),
                    covered,
                    expected: self.bucket_count,
                })));
            }
            Ok(self
                .replicasets
                .iter()
                .map(|rs| (rs.uuid(), rs.ranges()))
                .collect())
        }
    }

    impl Replicaset {
        /// UUID of the replicaset.
        #[must_use]
        pub fn uuid(&self) -> Uuid {
            self.uuid
        }

        /// The current master, or `None` while the cluster has not named one.
        #[must_use]
        pub fn master(&self) -> Option<&Instance> {
            self.master.as_ref()
        }

        /// The remaining instances.
        #[must_use]
        pub fn replicas(&self) -> &[Instance] {
            &self.replicas
        }

        /// All the instances, the master first.
        pub fn instances(&self) -> impl Iterator<Item = &Instance> {
            self.master.iter().chain(&self.replicas)
        }

        /// Bucket ranges the replicaset owns, ordered by start.
        #[must_use]
        pub fn ranges(&self) -> &[RangeInclusive<u64>] {
            &self.ranges
        }

        /// Name of the replicaset.
        #[must_use]
        pub fn name(&self) -> &str {
            &self.name
        }
    }

    impl Instance {
        /// Instance name.
        #[must_use]
        pub fn name(&self) -> &str {
            &self.name
        }

        /// The address the instance accepts `connection_type`, such as `iproto`, on;
        /// `None` if it does not.
        #[must_use]
        pub fn address(&self, connection_type: impl AsRef<str>) -> Option<&str> {
            self.addresses
                .get(connection_type.as_ref())
                .map(String::as_str)
        }

        /// Returns `true` if the instance is online.
        #[must_use]
        pub fn is_online(&self) -> bool {
            self.online
        }

        /// Returns `true` if the instance is the master of its replicaset.
        #[must_use]
        pub fn is_master(&self) -> bool {
            self.is_master
        }

        /// Raft id of the instance.
        #[must_use]
        pub fn raft_id(&self) -> u64 {
            self.raft_id
        }

        /// UUID of the instance.
        #[must_use]
        pub fn uuid(&self) -> Uuid {
            self.uuid
        }
    }
}

/// Cluster topology: instances, tiers, replicasets, and bucket ranges.
#[derive(Debug, Default)]
pub(crate) struct Topology {
    pub(crate) instances: HashMap<String, Arc<Instance>>,
    pgproto_tiers: Vec<Tier>,
    /// Tiers by every other connection type, such as `iproto`; never routed to.
    other_tiers: HashMap<ConnectionType, Vec<Tier>>,
}

/// A single tier with its replicasets and bucket-to-replicaset mapping.
#[derive(Debug, Clone)]
struct Tier {
    name: String,
    bucket_count: u64,
    replicasets: Vec<Replicaset>,
    bucket_map: Vec<BucketSlot>,
}

/// Maps a contiguous bucket range [start, end] to a replicaset index.
#[derive(Debug, Clone)]
struct BucketSlot {
    start: u64,
    end: u64,
    replicaset_idx: usize,
}

/// A replicaset with its master and replica instances.
#[derive(Debug, Clone)]
struct Replicaset {
    uuid: Uuid,
    master: Arc<Instance>,
    replicas: Vec<Arc<Instance>>,
}

impl From<(Vec<pg::Row>, Vec<pg::Row>)> for Topology {
    fn from((topology_rows, bucket_rows): (Vec<pg::Row>, Vec<pg::Row>)) -> Self {
        // One row per advertised address: a topology per connection type, `pgproto` is routed to.
        let mut by_type: HashMap<ConnectionType, Vec<pg::Row>> = HashMap::new();
        for row in topology_rows {
            let Ok(connection_type) = row.try_get::<_, &str>("connection_type") else {
                continue;
            };
            if let Some(rows) = by_type.get_mut(connection_type) {
                rows.push(row);
            } else {
                by_type.insert(connection_type.to_string(), vec![row]);
            }
        }
        // Parsed once: every connection type shares the same bucket ranges.
        let mut buckets: Vec<BucketRangeEntry> = Vec::new();
        for row in bucket_rows {
            let entries = if let Ok(json) = row.try_get::<_, serde_json::Value>(0) {
                serde_json::from_value::<Vec<BucketRangeEntry>>(json)
            } else if let Ok(s) = row.try_get::<_, &str>(0) {
                serde_json::from_str(s)
            } else {
                log::error!("failed to read bucket ranges column");
                continue;
            };
            match entries {
                Ok(v) => buckets.extend(v),
                Err(e) => log::error!("failed to parse bucket ranges: {e}"),
            }
        }

        let mut other_tiers: HashMap<ConnectionType, Vec<Tier>> = by_type
            .into_iter()
            .map(|(connection_type, rows)| (connection_type, build_tiers(rows, &buckets)))
            .collect();
        let pgproto_tiers = other_tiers.remove("pgproto").unwrap_or_default();
        // A master placeholder that is not the master also sits in `replicas`.
        let instances = pgproto_tiers
            .iter()
            .flat_map(|t| &t.replicasets)
            .flat_map(|rs| {
                rs.master
                    .is_master
                    .then_some(&rs.master)
                    .into_iter()
                    .chain(&rs.replicas)
            })
            .map(|inst| (inst.name.clone(), inst.clone()))
            .collect();
        Topology {
            instances,
            pgproto_tiers,
            other_tiers,
        }
    }
}

/// Builds the tiers from the rows of one connection type.
fn build_tiers(topology_rows: Vec<pg::Row>, buckets: &[BucketRangeEntry]) -> Vec<Tier> {
    let mut tiers: Vec<Tier> = Vec::new();
    for row in topology_rows {
        let inst = match Instance::try_from(row) {
            Ok(v) => v,
            Err(e) => {
                log::error!("failed to convert instance: {e}");
                continue;
            }
        };

        let inst = Arc::new(inst);

        let tier = if let Some(t) = tiers.iter_mut().find(|t| t.name == inst.tier) {
            t
        } else {
            tiers.push(Tier {
                name: inst.tier.clone(),
                bucket_count: inst.bucket_count,
                replicasets: Vec::new(),
                bucket_map: Vec::new(),
            });
            tiers.last_mut().expect("just pushed")
        };

        let rs = if let Some(r) = tier
            .replicasets
            .iter_mut()
            .find(|r| r.uuid == inst.replicaset_uuid)
        {
            r
        } else {
            tier.replicasets.push(Replicaset {
                uuid: inst.replicaset_uuid,
                master: inst.clone(),
                replicas: Vec::new(),
            });
            tier.replicasets.last_mut().expect("just pushed")
        };

        if inst.is_master {
            rs.master = inst.clone();
        } else {
            rs.replicas.push(inst.clone());
        }
    }

    for entry in buckets {
        let Some(tier) = tiers.iter_mut().find(|t| t.name == entry.tier_name) else {
            continue;
        };
        let Some(rs_idx) = tier
            .replicasets
            .iter()
            .position(|r| r.uuid == entry.current_replicaset_uuid)
        else {
            continue;
        };
        tier.bucket_map.push(BucketSlot {
            start: entry.bucket_id_start,
            end: entry.bucket_id_end,
            replicaset_idx: rs_idx,
        });
    }

    for tier in &mut tiers {
        tier.bucket_map.sort_unstable_by_key(|s| s.start);
    }

    tiers
}

impl Topology {
    /// Creates an empty topology.
    pub(crate) fn new() -> Self {
        Self::default()
    }

    /// Merges incoming topology: updates existing instance states, adds new instances,
    /// marks missing instances as offline, and replaces tier/bucket data.
    pub(crate) fn merge(&mut self, incoming: Topology) {
        for (name, existing) in &self.instances {
            if !incoming.instances.contains_key(name) {
                existing.set_state(State::Offline);
            }
        }

        // The map keeps the same Arc the new tiers hold, so each instance is one object;
        // only its pool slot index carries over from the previous refresh.
        for (name, inc_inst) in incoming.instances {
            match self.instances.entry(name) {
                Entry::Occupied(mut e) => {
                    inc_inst.set_idx(e.get().idx());
                    e.insert(inc_inst);
                }
                Entry::Vacant(e) => {
                    inc_inst.set_idx(NEXT_IDX.fetch_add(1, Ordering::Relaxed));
                    e.insert(inc_inst);
                }
            }
        }

        self.pgproto_tiers = incoming.pgproto_tiers;
        self.other_tiers = incoming.other_tiers;
    }

    /// Resolves the pool slot index for a given tier, `bucket_id`, and master preference.
    /// Returns `None` if the bucket or replicaset cannot be found, or no instance is online.
    pub(crate) fn resolve(&self, tier: &str, bucket_id: u64, on_master: bool) -> Option<usize> {
        let t = self.pgproto_tiers.iter().find(|t| t.name == tier)?;
        let rs = t.lookup(bucket_id)?;

        if on_master || rs.replicas.is_empty() {
            return if rs.master.is_online() {
                Some(rs.master.idx())
            } else {
                None
            };
        }

        let online: Vec<_> = rs.replicas.iter().filter(|r| r.is_online()).collect();
        if online.is_empty() {
            return if rs.master.is_online() {
                Some(rs.master.idx())
            } else {
                None
            };
        }

        let idx =
            usize::try_from(bucket_id % online.len() as u64).expect("modulo result fits in usize");
        Some(online[idx].idx())
    }

    /// Builds the topology handed out to the user: each instance once, with its address
    /// of every connection type.
    pub(crate) fn snapshot(&self) -> view::Topology {
        let mut view_tiers = Vec::with_capacity(self.pgproto_tiers.len());
        add_to_view(&mut view_tiers, "pgproto", &self.pgproto_tiers);
        for (connection_type, tiers) in &self.other_tiers {
            add_to_view(&mut view_tiers, connection_type, tiers);
        }
        view_tiers.sort_unstable_by(|a, b| a.name.cmp(&b.name));
        for tier in &mut view_tiers {
            sort(&mut tier.replicasets);
        }
        view::Topology { tiers: view_tiers }
    }

    /// Returns the total bucket count for a tier.
    pub(crate) fn bucket_count(&self, tier: &str) -> Option<u64> {
        self.pgproto_tiers
            .iter()
            .find(|t| t.name == tier)
            .map(|t| t.bucket_count)
    }
}

/// Adds the instances of `tiers` to `view_tiers` with their `connection_type` address,
/// next to what the other connection types have added.
fn add_to_view(view_tiers: &mut Vec<view::Tier>, connection_type: &str, tiers: &[Tier]) {
    for tier in tiers {
        let view_tier = if let Some(t) = view_tiers.iter_mut().find(|t| t.name == tier.name) {
            t
        } else {
            view_tiers.push(view::Tier {
                name: tier.name.clone(),
                bucket_count: tier.bucket_count,
                replicasets: Vec::new(),
            });
            view_tiers.last_mut().expect("just pushed")
        };

        for (idx, rs) in tier.replicasets.iter().enumerate() {
            let view_rs =
                if let Some(r) = view_tier.replicasets.iter_mut().find(|r| r.uuid == rs.uuid) {
                    r
                } else {
                    view_tier.replicasets.push(view::Replicaset {
                        uuid: rs.uuid,
                        name: rs.master.replicaset_name.clone(),
                        master: None,
                        replicas: Vec::new(),
                        ranges: tier
                            .bucket_map
                            .iter()
                            .filter(|s| s.replicaset_idx == idx)
                            .map(|s| s.start..=s.end)
                            .collect(),
                    });
                    view_tier.replicasets.last_mut().expect("just pushed")
                };

            if rs.master.is_master {
                match &mut view_rs.master {
                    Some(m) => {
                        m.addresses
                            .insert(connection_type.to_string(), rs.master.addr.to_string());
                    }
                    None => view_rs.master = Some(instance(&rs.master, connection_type)),
                }
            }
            for r in &rs.replicas {
                if let Some(i) = view_rs.replicas.iter_mut().find(|i| i.name == r.name) {
                    i.addresses
                        .insert(connection_type.to_string(), r.addr.to_string());
                } else {
                    view_rs.replicas.push(instance(r, connection_type));
                }
            }
        }
    }
}

impl Tier {
    /// Finds the replicaset owning `bucket_id` via binary search on the bucket map.
    fn lookup(&self, bucket_id: u64) -> Option<&Replicaset> {
        let i = self
            .bucket_map
            .partition_point(|s| s.start <= bucket_id)
            .checked_sub(1)?;
        let slot = &self.bucket_map[i];
        if bucket_id <= slot.end {
            self.replicasets.get(slot.replicaset_idx)
        } else {
            None
        }
    }
}

/// Orders replicasets by UUID and their replicas by name, so equal topologies compare equal.
fn sort(replicasets: &mut [view::Replicaset]) {
    replicasets.sort_unstable_by_key(|rs| rs.uuid);
    for rs in replicasets {
        rs.replicas.sort_unstable_by(|a, b| a.name.cmp(&b.name));
    }
}

/// Builds the view of an instance with its `connection_type` address.
fn instance(inst: &Instance, connection_type: &str) -> view::Instance {
    view::Instance {
        name: inst.name.clone(),
        addresses: HashMap::from([(connection_type.to_string(), inst.addr.to_string())]),
        online: inst.is_online(),
        is_master: inst.is_master,
        raft_id: inst.raft_id,
        uuid: inst.uuid,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::error::Error;

    const RS1: Uuid = Uuid::from_u128(1);
    const RS2: Uuid = Uuid::from_u128(2);

    /// Build a single-instance topology directly, bypassing `pg::Row` parsing.
    fn make_topology(name: &str, tier: &str, rs_uuid: Uuid) -> Topology {
        let inst = Arc::new(Instance::new_for_test(
            name,
            "127.0.0.1:3301",
            tier,
            rs_uuid,
            true,
            3000,
        ));

        let mut instances = HashMap::new();
        instances.insert(inst.name.clone(), inst.clone());

        let rs = Replicaset {
            uuid: rs_uuid,
            master: inst.clone(),
            replicas: vec![],
        };
        let tier_obj = Tier {
            name: tier.to_string(),
            bucket_count: 3000,
            replicasets: vec![rs],
            bucket_map: vec![BucketSlot {
                start: 1,
                end: 3000,
                replicaset_idx: 0,
            }],
        };

        Topology {
            instances,
            pgproto_tiers: vec![tier_obj],
            other_tiers: HashMap::new(),
        }
    }

    /// A range of an unknown replicaset leaves a hole, which must be reported.
    #[test]
    fn bucket_map_rejects_an_incomplete_map() {
        let mut topo = Topology::new();
        let mut incoming = make_topology("i1", "default", RS1);
        incoming.pgproto_tiers[0].bucket_map = vec![
            BucketSlot {
                start: 1,
                end: 1500,
                replicaset_idx: 0,
            },
            // Dangling: `replicasets` holds a single entry, so idx 7 is unknown.
            BucketSlot {
                start: 1501,
                end: 3000,
                replicaset_idx: 7,
            },
        ];
        topo.merge(incoming);

        match topo.snapshot().tiers()[0].bucket_map() {
            Err(Error::IncompleteBucketMap(m)) => {
                assert_eq!(m.tier, "default");
                assert_eq!((m.covered, m.expected), (1500, 3000));
            }
            other => panic!("a half-covered tier must not be returned: {other:?}"),
        }
    }

    /// A tier without ranges must be reported, not returned empty.
    #[test]
    fn bucket_map_rejects_a_tier_with_no_ranges() {
        let mut topo = Topology::new();
        let mut incoming = make_topology("i1", "default", RS1);
        incoming.pgproto_tiers[0].bucket_map.clear();
        topo.merge(incoming);

        assert!(matches!(
            topo.snapshot().tiers()[0].bucket_map(),
            Err(Error::IncompleteBucketMap(ref m)) if m.covered == 0
        ));
    }

    /// Every range a replicaset owns must route to that replicaset's master.
    #[test]
    fn snapshot_ranges_agree_with_resolve() {
        let mut topo = Topology::new();
        let mut incoming = make_topology("i1", "default", RS1);
        let i2 = Arc::new(Instance::new_for_test(
            "i2",
            "127.0.0.1:3302",
            "default",
            RS2,
            true,
            3000,
        ));
        incoming.instances.insert(i2.name.clone(), i2.clone());
        let tier = &mut incoming.pgproto_tiers[0];
        tier.replicasets.push(Replicaset {
            uuid: RS2,
            master: i2,
            replicas: vec![],
        });
        tier.bucket_map = vec![
            BucketSlot {
                start: 1,
                end: 1000,
                replicaset_idx: 0,
            },
            BucketSlot {
                start: 1001,
                end: 2000,
                replicaset_idx: 1,
            },
            BucketSlot {
                start: 2001,
                end: 3000,
                replicaset_idx: 0,
            },
        ];
        topo.merge(incoming);

        let snapshot = topo.snapshot();
        let tier = &snapshot.tiers()[0];
        assert_eq!(tier.name(), "default");
        let owned: Vec<_> = tier
            .replicasets()
            .iter()
            .map(|rs| {
                (
                    rs.uuid(),
                    rs.master().map(view::Instance::name),
                    rs.ranges().len(),
                )
            })
            .collect();
        assert_eq!(owned, [(RS1, Some("i1"), 2), (RS2, Some("i2"), 1)]);
        let bucket_map = tier.bucket_map().expect("map covers the tier");
        assert_eq!(bucket_map[&RS1], [1..=1000, 2001..=3000]);

        for rs in tier.replicasets() {
            let expected = topo.instances[rs.master().expect("named master").name()].idx();
            for r in rs.ranges() {
                for bucket in [*r.start(), *r.end()] {
                    assert_eq!(topo.resolve("default", bucket, true), Some(expected));
                }
            }
        }
    }

    /// Without a named master routing keeps the placeholder, but `master()` is `None`.
    #[test]
    fn snapshot_reports_no_master_until_one_is_named() {
        // What `Topology::from` builds when no row carries is_master.
        let inst = Arc::new(Instance::new_for_test(
            "i1",
            "127.0.0.1:3301",
            "default",
            RS1,
            false,
            3000,
        ));
        let mut incoming = make_topology("i1", "default", RS1);
        incoming.instances.insert(inst.name.clone(), inst.clone());
        let rs = &mut incoming.pgproto_tiers[0].replicasets[0];
        rs.master = inst.clone();
        rs.replicas = vec![inst];
        let mut topo = Topology::new();
        topo.merge(incoming);

        assert_eq!(
            topo.resolve("default", 1, true),
            Some(topo.instances["i1"].idx())
        );
        let snapshot = topo.snapshot();
        let rs = &snapshot.tiers()[0].replicasets()[0];
        assert_eq!(rs.master(), None);
        let replicas: Vec<_> = rs.replicas().iter().map(view::Instance::name).collect();
        assert_eq!(replicas, ["i1"]);
    }

    /// `instances` yields the master first, then the replicas by name.
    #[test]
    fn replicaset_lists_the_master_first() {
        let mut topo = make_topology("m", "default", RS1);
        topo.pgproto_tiers[0].replicasets[0].replicas = ["b", "a"]
            .map(|name| {
                Arc::new(Instance::new_for_test(
                    name,
                    "127.0.0.1:3302",
                    "default",
                    RS1,
                    false,
                    3000,
                ))
            })
            .into();

        let snapshot = topo.snapshot();
        let rs = &snapshot.tiers()[0].replicasets()[0];
        let instances: Vec<_> = rs.instances().map(|i| (i.name(), i.is_master())).collect();
        assert_eq!(instances, [("m", true), ("a", false), ("b", false)]);
    }

    /// Regression: without `inc_inst.set_idx(existing.idx())` in `merge`, the
    /// second merge replaces `self.pgproto_tiers` with `incoming.pgproto_tiers` whose
    /// `Arc<Instance>` objects still carry the default `idx == 0`.  Every
    /// subsequent `resolve` call would then return `Some(0)` instead of the
    /// pool slot that was assigned during the first merge.
    #[test]
    fn merge_preserves_pool_idx_in_tiers() {
        // Bump the global counter so the assigned idx is unambiguously non-zero.
        NEXT_IDX.fetch_add(1, Ordering::Relaxed);

        let mut topo = Topology::new();

        // First merge: "i1" is unknown, so it receives a fresh pool-slot index.
        topo.merge(make_topology("i1", "default", RS1));
        let expected_idx = topo.instances["i1"].idx();
        assert!(expected_idx > 0, "precondition: idx must be non-zero");

        // Second merge: same instance name, but a brand-new Arc<Instance> with
        // idx == 0.  Before the fix, `self.pgproto_tiers = incoming.pgproto_tiers` made the
        // tiers point at those fresh objects, so resolve returned Some(0).
        topo.merge(make_topology("i1", "default", RS1));

        assert_eq!(
            topo.resolve("default", 1, true),
            Some(expected_idx),
            "resolve must return the stable pool slot, not the default 0 from the fresh Arc"
        );
    }

    /// Regression: `merge` kept the first-seen Arc in `instances` while the tiers got the
    /// fresh one, so the pools and routing saw two objects with diverging fields.
    #[test]
    fn merge_keeps_one_object_per_instance() {
        let mut topo = Topology::new();
        topo.merge(make_topology("i1", "default", RS1));
        topo.merge(make_topology("i1", "default", RS1));

        let master = &topo.pgproto_tiers[0].replicasets[0].master;
        assert!(Arc::ptr_eq(&topo.instances["i1"], master));
    }

    /// Builds a topology holding only `tier`, filed under the `iproto` connection type.
    fn with_tier(tier: Tier) -> Topology {
        Topology {
            instances: HashMap::new(),
            pgproto_tiers: Vec::new(),
            other_tiers: HashMap::from([("iproto".to_string(), vec![tier])]),
        }
    }

    /// Builds a one-replicaset tier whose master listens on `addr`.
    fn tier_at(addr: &str) -> Tier {
        let inst = Arc::new(Instance::new_for_test(
            "i1", addr, "default", RS1, true, 3000,
        ));
        Tier {
            name: "default".to_string(),
            bucket_count: 3000,
            replicasets: vec![Replicaset {
                uuid: RS1,
                master: inst,
                replicas: vec![],
            }],
            bucket_map: vec![BucketSlot {
                start: 1,
                end: 3000,
                replicaset_idx: 0,
            }],
        }
    }

    /// An instance is listed once, with the address of each connection type it accepts.
    #[test]
    fn snapshot_merges_the_addresses_of_an_instance() {
        let mut topo = with_tier(tier_at("127.0.0.1:3301"));
        topo.pgproto_tiers = vec![tier_at("127.0.0.1:5432")];

        let snapshot = topo.snapshot();
        let rs = &snapshot.tiers()[0].replicasets()[0];
        assert!(rs.replicas().is_empty());
        let i1 = rs.master().expect("named master");
        assert_eq!(i1.address("pgproto"), Some("127.0.0.1:5432"));
        assert_eq!(i1.address("iproto"), Some("127.0.0.1:3301"));
    }

    /// An instance without pgproto is still listed, and the buckets it owns still count.
    #[test]
    fn snapshot_keeps_an_instance_without_pgproto() {
        let topo = with_tier(tier_at("127.0.0.1:3301"));

        let snapshot = topo.snapshot();
        let tier = &snapshot.tiers()[0];
        let i1 = tier.replicasets()[0].master().expect("named master");
        assert_eq!(i1.address("pgproto"), None);
        assert_eq!(i1.address("iproto"), Some("127.0.0.1:3301"));
        assert!(tier.bucket_map().is_ok());
    }

    /// `merge` replaces the non-pgproto tiers wholesale.
    #[test]
    fn merge_replaces_other_tiers() {
        let mut topo = with_tier(tier_at("127.0.0.1:3301"));
        topo.merge(with_tier(tier_at("127.0.0.1:3302")));

        let snapshot = topo.snapshot();
        let i1 = snapshot.tiers()[0].replicasets()[0]
            .master()
            .expect("named master");
        assert_eq!(i1.address("iproto"), Some("127.0.0.1:3302"));
    }

    /// Insertion order of replicasets and replicas does not affect the snapshot.
    #[test]
    fn snapshot_is_independent_of_insertion_order() {
        let inst = |name: &str, rs: Uuid, master: bool| {
            Arc::new(Instance::new_for_test(
                name,
                "127.0.0.1:3301",
                "default",
                rs,
                master,
                3000,
            ))
        };
        let build = |reversed: bool| {
            let mut sets = vec![
                (RS1, 1, 1500, vec!["a", "b"]),
                (RS2, 1501, 3000, vec!["c", "d"]),
            ];
            if reversed {
                sets.reverse();
            }
            let mut tier = Tier {
                name: "default".to_string(),
                bucket_count: 3000,
                replicasets: Vec::new(),
                bucket_map: Vec::new(),
            };
            for (idx, (uuid, start, end, mut names)) in sets.into_iter().enumerate() {
                if reversed {
                    names.reverse();
                }
                tier.replicasets.push(Replicaset {
                    uuid,
                    master: inst(&format!("{uuid}-m"), uuid, true),
                    replicas: names.into_iter().map(|n| inst(n, uuid, false)).collect(),
                });
                tier.bucket_map.push(BucketSlot {
                    start,
                    end,
                    replicaset_idx: idx,
                });
            }
            tier.bucket_map.sort_unstable_by_key(|s| s.start);
            let mut topo = Topology::new();
            topo.pgproto_tiers = vec![tier];
            topo.snapshot()
        };

        assert_eq!(build(false), build(true));
    }
}