Line data Source code
1 : use std::collections::BTreeMap;
2 :
3 : use utils::id::TenantId;
4 : use utils::shard::TenantShardId;
5 :
6 : use crate::scheduler::{ScheduleContext, ScheduleMode};
7 : use crate::tenant_shard::TenantShard;
8 :
9 : /// When making scheduling decisions, it is useful to have the ScheduleContext for a whole
10 : /// tenant while considering the individual shards within it. This iterator is a helper
11 : /// that gathers all the shards in a tenant and then yields them together with a ScheduleContext
12 : /// for the tenant.
13 : pub(super) struct TenantShardContextIterator<'a> {
14 : schedule_mode: ScheduleMode,
15 : inner: std::collections::btree_map::IterMut<'a, TenantShardId, TenantShard>,
16 : }
17 :
18 : impl<'a> TenantShardContextIterator<'a> {
19 1 : pub(super) fn new(
20 1 : tenants: &'a mut BTreeMap<TenantShardId, TenantShard>,
21 1 : schedule_mode: ScheduleMode,
22 1 : ) -> Self {
23 1 : Self {
24 1 : schedule_mode,
25 1 : inner: tenants.iter_mut(),
26 1 : }
27 1 : }
28 : }
29 :
30 : impl<'a> Iterator for TenantShardContextIterator<'a> {
31 : type Item = (TenantId, ScheduleContext, Vec<&'a mut TenantShard>);
32 :
33 3 : fn next(&mut self) -> Option<Self::Item> {
34 3 : let mut tenant_shards = Vec::new();
35 3 : let mut schedule_context = ScheduleContext::new(self.schedule_mode.clone());
36 : loop {
37 6 : let (tenant_shard_id, shard) = self.inner.next()?;
38 :
39 6 : if tenant_shard_id.is_shard_zero() {
40 : // Cleared on last shard of previous tenant
41 3 : assert!(tenant_shards.is_empty());
42 3 : }
43 :
44 : // Accumulate the schedule context for all the shards in a tenant
45 6 : schedule_context.avoid(&shard.intent.all_pageservers());
46 6 : if let Some(attached) = shard.intent.get_attached() {
47 6 : schedule_context.push_attached(*attached);
48 6 : }
49 6 : tenant_shards.push(shard);
50 6 :
51 6 : if tenant_shard_id.shard_number.0 == tenant_shard_id.shard_count.count() - 1 {
52 3 : return Some((tenant_shard_id.tenant_id, schedule_context, tenant_shards));
53 3 : }
54 : }
55 3 : }
56 : }
57 :
58 : #[cfg(test)]
59 : mod tests {
60 : use std::{collections::BTreeMap, str::FromStr};
61 :
62 : use pageserver_api::controller_api::PlacementPolicy;
63 : use utils::shard::{ShardCount, ShardNumber};
64 :
65 : use crate::{
66 : scheduler::test_utils::make_test_nodes, service::Scheduler,
67 : tenant_shard::tests::make_test_tenant_with_id,
68 : };
69 :
70 : use super::*;
71 :
72 : #[test]
73 1 : fn test_context_iterator() {
74 1 : // Hand-crafted tenant IDs to ensure they appear in the expected order when put into
75 1 : // a btreemap & iterated
76 1 : let mut t_1_shards = make_test_tenant_with_id(
77 1 : TenantId::from_str("af0480929707ee75372337efaa5ecf96").unwrap(),
78 1 : PlacementPolicy::Attached(1),
79 1 : ShardCount(1),
80 1 : None,
81 1 : );
82 1 : let t_2_shards = make_test_tenant_with_id(
83 1 : TenantId::from_str("bf0480929707ee75372337efaa5ecf96").unwrap(),
84 1 : PlacementPolicy::Attached(1),
85 1 : ShardCount(4),
86 1 : None,
87 1 : );
88 1 : let mut t_3_shards = make_test_tenant_with_id(
89 1 : TenantId::from_str("cf0480929707ee75372337efaa5ecf96").unwrap(),
90 1 : PlacementPolicy::Attached(1),
91 1 : ShardCount(1),
92 1 : None,
93 1 : );
94 1 :
95 1 : let t1_id = t_1_shards[0].tenant_shard_id.tenant_id;
96 1 : let t2_id = t_2_shards[0].tenant_shard_id.tenant_id;
97 1 : let t3_id = t_3_shards[0].tenant_shard_id.tenant_id;
98 1 :
99 1 : let mut tenants = BTreeMap::new();
100 1 : tenants.insert(t_1_shards[0].tenant_shard_id, t_1_shards.pop().unwrap());
101 5 : for shard in t_2_shards {
102 4 : tenants.insert(shard.tenant_shard_id, shard);
103 4 : }
104 1 : tenants.insert(t_3_shards[0].tenant_shard_id, t_3_shards.pop().unwrap());
105 1 :
106 1 : let nodes = make_test_nodes(3, &[]);
107 1 : let mut scheduler = Scheduler::new(nodes.values());
108 1 : let mut context = ScheduleContext::default();
109 6 : for shard in tenants.values_mut() {
110 6 : shard.schedule(&mut scheduler, &mut context).unwrap();
111 6 : }
112 :
113 1 : let mut iter = TenantShardContextIterator::new(&mut tenants, ScheduleMode::Speculative);
114 1 : let (tenant_id, context, shards) = iter.next().unwrap();
115 1 : assert_eq!(tenant_id, t1_id);
116 1 : assert_eq!(shards[0].tenant_shard_id.shard_number, ShardNumber(0));
117 1 : assert_eq!(shards.len(), 1);
118 1 : assert_eq!(context.attach_count(), 1);
119 :
120 1 : let (tenant_id, context, shards) = iter.next().unwrap();
121 1 : assert_eq!(tenant_id, t2_id);
122 1 : assert_eq!(shards[0].tenant_shard_id.shard_number, ShardNumber(0));
123 1 : assert_eq!(shards[1].tenant_shard_id.shard_number, ShardNumber(1));
124 1 : assert_eq!(shards[2].tenant_shard_id.shard_number, ShardNumber(2));
125 1 : assert_eq!(shards[3].tenant_shard_id.shard_number, ShardNumber(3));
126 1 : assert_eq!(shards.len(), 4);
127 1 : assert_eq!(context.attach_count(), 4);
128 :
129 1 : let (tenant_id, context, shards) = iter.next().unwrap();
130 1 : assert_eq!(tenant_id, t3_id);
131 1 : assert_eq!(shards[0].tenant_shard_id.shard_number, ShardNumber(0));
132 1 : assert_eq!(shards.len(), 1);
133 1 : assert_eq!(context.attach_count(), 1);
134 :
135 6 : for shard in tenants.values_mut() {
136 6 : shard.intent.clear(&mut scheduler);
137 6 : }
138 1 : }
139 : }
|
