1use nmbrs_workload::model::ParsedOp;
13
14#[derive(Debug, Clone, Copy, PartialEq)]
16pub enum SequencerType {
17 Bucket,
19 Interval,
21 Concat,
23}
24
25impl SequencerType {
26 pub fn parse(s: &str) -> Result<Self, String> {
28 match s.to_lowercase().as_str() {
29 "bucket" => Ok(Self::Bucket),
30 "interval" => Ok(Self::Interval),
31 "concat" => Ok(Self::Concat),
32 other => Err(format!(
33 "unknown sequencer type: '{other}', use bucket|interval|concat"
34 )),
35 }
36 }
37}
38
39pub struct OpSequence {
44 ops: Vec<ParsedOp>,
45 lut: Vec<usize>,
46 sequencer_type: SequencerType,
47}
48
49impl OpSequence {
50 pub fn build(ops: Vec<ParsedOp>, ratios: &[u64], strategy: SequencerType) -> Self {
52 assert_eq!(
53 ops.len(),
54 ratios.len(),
55 "ops and ratios must have equal length"
56 );
57 assert!(!ops.is_empty(), "must have at least one op");
58
59 let lut = match strategy {
60 SequencerType::Bucket => build_bucket_lut(ratios),
61 SequencerType::Interval => build_interval_lut(ratios),
62 SequencerType::Concat => build_concat_lut(ratios),
63 };
64
65 Self {
66 ops,
67 lut,
68 sequencer_type: strategy,
69 }
70 }
71
72 pub fn uniform(ops: Vec<ParsedOp>) -> Self {
74 let ratios = vec![1u64; ops.len()];
75 Self::build(ops, &ratios, SequencerType::Bucket)
76 }
77
78 pub fn from_ops(ops: Vec<ParsedOp>, strategy: SequencerType) -> Self {
85 let ratios: Vec<u64> = ops
86 .iter()
87 .map(|op| {
88 op.params
89 .get("ratio")
90 .and_then(|v| v.as_u64())
91 .or_else(|| op.op.get("ratio").and_then(|v| v.as_u64()))
92 .unwrap_or(1)
93 })
94 .collect();
95 Self::build(ops, &ratios, strategy)
96 }
97
98 #[inline]
100 pub fn get(&self, cycle: u64) -> &ParsedOp {
101 let idx = self.lut[(cycle as usize) % self.lut.len()];
102 &self.ops[idx]
103 }
104
105 #[inline]
107 pub fn get_with_index(&self, cycle: u64) -> (usize, &ParsedOp) {
108 let idx = self.lut[(cycle as usize) % self.lut.len()];
109 (idx, &self.ops[idx])
110 }
111
112 pub fn templates(&self) -> &[ParsedOp] {
114 &self.ops
115 }
116
117 pub fn op_count(&self) -> usize {
119 self.ops.len()
120 }
121
122 pub fn stanza_length(&self) -> usize {
124 self.lut.len()
125 }
126
127 pub fn stanza_ops(&self, base_cycle: u64) -> Vec<(&ParsedOp, u64)> {
133 (0..self.lut.len())
134 .map(|offset| {
135 let cycle = base_cycle + offset as u64;
136 (self.get(cycle), cycle)
137 })
138 .collect()
139 }
140
141 pub fn sequencer_type(&self) -> SequencerType {
143 self.sequencer_type
144 }
145
146 pub fn lut(&self) -> &[usize] {
148 &self.lut
149 }
150}
151
152fn build_bucket_lut(ratios: &[u64]) -> Vec<usize> {
168 let total: u64 = ratios.iter().sum();
169 let mut lut = Vec::with_capacity(total as usize);
170 let mut remaining: Vec<u64> = ratios.to_vec();
171
172 loop {
173 let mut any_drawn = false;
174 for (i, rem) in remaining.iter_mut().enumerate() {
175 if *rem > 0 {
176 lut.push(i);
177 *rem -= 1;
178 any_drawn = true;
179 }
180 }
181 if !any_drawn {
182 break;
183 }
184 }
185
186 lut
187}
188
189fn build_interval_lut(ratios: &[u64]) -> Vec<usize> {
200 let total: u64 = ratios.iter().sum();
201 if total == 0 {
202 return Vec::new();
203 }
204
205 let mut entries: Vec<(f64, usize, usize)> = Vec::new(); for (i, &ratio) in ratios.iter().enumerate() {
208 if ratio == 0 {
209 continue;
210 }
211 let spacing = total as f64 / ratio as f64;
212 for j in 0..ratio as usize {
213 entries.push((j as f64 * spacing, i, j));
214 }
215 }
216
217 entries.sort_by(|a, b| {
219 a.0.partial_cmp(&b.0)
220 .unwrap_or(std::cmp::Ordering::Equal)
221 .then_with(|| a.1.cmp(&b.1))
222 });
223
224 entries.iter().map(|&(_, idx, _)| idx).collect()
225}
226
227fn build_concat_lut(ratios: &[u64]) -> Vec<usize> {
232 let mut lut = Vec::new();
233 for (i, &ratio) in ratios.iter().enumerate() {
234 for _ in 0..ratio {
235 lut.push(i);
236 }
237 }
238 lut
239}
240
241#[cfg(test)]
242mod tests {
243 use super::*;
244
245 fn names(seq: &OpSequence, count: usize) -> Vec<String> {
246 (0..count as u64).map(|c| seq.get(c).name.clone()).collect()
247 }
248
249 #[test]
252 fn bucket_uniform() {
253 let ops = vec![
254 ParsedOp::simple("A", "a"),
255 ParsedOp::simple("B", "b"),
256 ParsedOp::simple("C", "c"),
257 ];
258 let seq = OpSequence::build(ops, &[1, 1, 1], SequencerType::Bucket);
259 assert_eq!(seq.stanza_length(), 3);
260 assert_eq!(names(&seq, 6), vec!["A", "B", "C", "A", "B", "C"]);
261 }
262
263 #[test]
264 fn bucket_weighted() {
265 let ops = vec![
266 ParsedOp::simple("R", "read"),
267 ParsedOp::simple("W", "write"),
268 ParsedOp::simple("D", "delete"),
269 ];
270 let seq = OpSequence::build(ops, &[3, 2, 1], SequencerType::Bucket);
271 assert_eq!(seq.stanza_length(), 6);
272 let n = names(&seq, 6);
273 assert_eq!(n.iter().filter(|s| *s == "R").count(), 3);
274 assert_eq!(n.iter().filter(|s| *s == "W").count(), 2);
275 assert_eq!(n.iter().filter(|s| *s == "D").count(), 1);
276 }
277
278 #[test]
279 fn bucket_interleaves() {
280 let ops = vec![ParsedOp::simple("A", "a"), ParsedOp::simple("B", "b")];
281 let seq = OpSequence::build(ops, &[3, 1], SequencerType::Bucket);
282 assert_eq!(seq.lut(), &[0, 1, 0, 0]);
284 }
285
286 #[test]
289 fn interval_uniform() {
290 let ops = vec![ParsedOp::simple("A", "a"), ParsedOp::simple("B", "b")];
291 let seq = OpSequence::build(ops, &[1, 1], SequencerType::Interval);
292 assert_eq!(seq.stanza_length(), 2);
293 assert_eq!(names(&seq, 4), vec!["A", "B", "A", "B"]);
294 }
295
296 #[test]
297 fn interval_weighted() {
298 let ops = vec![ParsedOp::simple("A", "a"), ParsedOp::simple("B", "b")];
299 let seq = OpSequence::build(ops, &[4, 2], SequencerType::Interval);
300 assert_eq!(seq.stanza_length(), 6);
301 let n = names(&seq, 6);
302 assert_eq!(n.iter().filter(|s| *s == "A").count(), 4);
303 assert_eq!(n.iter().filter(|s| *s == "B").count(), 2);
304 }
305
306 #[test]
307 fn interval_distributes_evenly() {
308 let ops = vec![ParsedOp::simple("A", "a"), ParsedOp::simple("B", "b")];
309 let seq = OpSequence::build(ops, &[3, 3], SequencerType::Interval);
310 assert_eq!(seq.stanza_length(), 6);
313 assert_eq!(names(&seq, 6), vec!["A", "B", "A", "B", "A", "B"]);
314 }
315
316 #[test]
319 fn concat_sequential() {
320 let ops = vec![
321 ParsedOp::simple("A", "a"),
322 ParsedOp::simple("B", "b"),
323 ParsedOp::simple("C", "c"),
324 ];
325 let seq = OpSequence::build(ops, &[3, 2, 1], SequencerType::Concat);
326 assert_eq!(seq.stanza_length(), 6);
327 assert_eq!(names(&seq, 6), vec!["A", "A", "A", "B", "B", "C"]);
328 }
329
330 #[test]
331 fn concat_wraps() {
332 let ops = vec![ParsedOp::simple("X", "x"), ParsedOp::simple("Y", "y")];
333 let seq = OpSequence::build(ops, &[2, 1], SequencerType::Concat);
334 assert_eq!(names(&seq, 6), vec!["X", "X", "Y", "X", "X", "Y"]);
335 }
336
337 #[test]
340 fn from_ops_extracts_ratios() {
341 let mut op1 = ParsedOp::simple("read", "SELECT");
342 op1.params.insert("ratio".into(), serde_json::json!(5));
343 let op2 = ParsedOp::simple("write", "INSERT"); let seq = OpSequence::from_ops(vec![op1, op2], SequencerType::Bucket);
346 assert_eq!(seq.stanza_length(), 6); }
348
349 #[test]
352 fn sequencer_type_parse() {
353 assert_eq!(
354 SequencerType::parse("bucket").unwrap(),
355 SequencerType::Bucket
356 );
357 assert_eq!(
358 SequencerType::parse("INTERVAL").unwrap(),
359 SequencerType::Interval
360 );
361 assert_eq!(
362 SequencerType::parse("Concat").unwrap(),
363 SequencerType::Concat
364 );
365 assert!(SequencerType::parse("bogus").is_err());
366 }
367
368 #[test]
369 fn single_op_any_strategy() {
370 for strategy in [
371 SequencerType::Bucket,
372 SequencerType::Interval,
373 SequencerType::Concat,
374 ] {
375 let ops = vec![ParsedOp::simple("only", "SELECT 1")];
376 let seq = OpSequence::build(ops, &[1], strategy);
377 assert_eq!(seq.get(0).name, "only");
378 assert_eq!(seq.get(999).name, "only");
379 }
380 }
381
382 #[test]
383 fn stanza_repeats_cleanly() {
384 let ops = vec![ParsedOp::simple("A", "a"), ParsedOp::simple("B", "b")];
385 let seq = OpSequence::build(ops, &[2, 1], SequencerType::Bucket);
386 let stanza = seq.stanza_length();
387 let first: Vec<String> = (0..stanza as u64)
389 .map(|c| seq.get(c).name.clone())
390 .collect();
391 let second: Vec<String> = (stanza as u64..2 * stanza as u64)
392 .map(|c| seq.get(c).name.clone())
393 .collect();
394 assert_eq!(first, second);
395 }
396}