shuttle_engine/scheduler/
serialization.rs

1//! This module implements a simple serialization scheme for schedules (`Schedule`) that tries to
2//! produce small printable strings. This is useful for roundtripping schedules in test outputs.
3
4use crate::runtime::task::TaskId;
5use crate::scheduler::{Schedule, ScheduleStep};
6use bitvec::prelude::*;
7
8/// A simplified version of the deprecated [`varmint`](https://github.com/mycorrhiza/varmint-rs)
9/// crate, used under the MIT license.
10mod varint {
11    pub fn space_needed(val: u64) -> usize {
12        let used_bits = u64::MIN.leading_zeros() - val.leading_zeros();
13        std::cmp::max((used_bits + 6) as usize / 7, 1)
14    }
15
16    pub trait WriteVarInt {
17        fn write_u64_varint(&mut self, val: u64) -> std::io::Result<()>;
18    }
19
20    impl<R: std::io::Write> WriteVarInt for R {
21        fn write_u64_varint(&mut self, mut val: u64) -> std::io::Result<()> {
22            loop {
23                let current = (val & 0x7F) as u8;
24                val >>= 7;
25                if val == 0 {
26                    self.write_all(&[current])?;
27                    return Ok(());
28                } else {
29                    self.write_all(&[current | 0x80])?;
30                }
31            }
32        }
33    }
34
35    pub trait ReadVarInt {
36        fn read_u64_varint(&mut self) -> std::io::Result<u64>;
37    }
38
39    fn read_u8<R: std::io::Read>(reader: &mut R) -> std::io::Result<u8> {
40        let mut buffer = [0u8];
41        reader.read_exact(&mut buffer)?;
42        Ok(buffer[0])
43    }
44
45    impl<R: std::io::Read> ReadVarInt for R {
46        fn read_u64_varint(&mut self) -> std::io::Result<u64> {
47            let first = read_u8(self)?;
48            if first & 0x80 == 0 {
49                return Ok(u64::from(first));
50            }
51
52            let mut result = u64::from(first & 0x7F);
53            let mut offset = 7;
54
55            loop {
56                let current = read_u8(self)?;
57                result += u64::from(current & 0x7F) << offset;
58                if current & 0x80 == 0 {
59                    return Ok(result);
60                }
61                offset += 7;
62                if offset == 63 {
63                    let last = read_u8(self)?;
64                    if last == 0x01 {
65                        return Ok(result + (1 << offset));
66                    } else {
67                        return Err(std::io::Error::other("varint exceeded 64 bits long"));
68                    }
69                }
70            }
71        }
72    }
73}
74
75// The serialization format is this:
76//   [task id bitwidth] [number of schedule steps] [seed] [step]*
77// The bitwidth, number of steps, and seed are encoded as VarInts, so are at least one byte.
78// The steps are densely packed bitstrings. The leading bit of a step is 0 if it's a task ID or 1
79// if it's a random value. If it's a task ID, the following `bitwidth` bits are the task ID. If it's
80// a random value, there are no following bits.
81//
82// We encode the binary serialization as a hex string for easy copy/pasting.
83
84const SCHEDULE_MAGIC_V2: u8 = 0x91;
85
86const LINE_WIDTH: usize = 76;
87
88pub fn serialize_schedule(schedule: &Schedule) -> String {
89    use self::varint::{space_needed, WriteVarInt};
90
91    let &max_task_id = schedule
92        .steps
93        .iter()
94        .filter_map(|s| match s {
95            ScheduleStep::Task(tid) => Some(tid),
96            _ => None,
97        })
98        .max()
99        .unwrap_or(&TaskId::from(0));
100    let task_id_bits = std::mem::size_of_val(&max_task_id) * 8 - usize::from(max_task_id).leading_zeros() as usize;
101    let task_id_bits = task_id_bits.max(1);
102
103    let mut encoded = bitvec![u8, Lsb0; 0; schedule.steps.len() * (1 + task_id_bits)];
104    let mut offset = 0usize;
105    for step in &schedule.steps {
106        match step {
107            ScheduleStep::Task(tid) => {
108                encoded.set(offset, false);
109                encoded[offset + 1..offset + 1 + task_id_bits].store(usize::from(*tid));
110                offset += 1 + task_id_bits;
111            }
112            ScheduleStep::Random => {
113                encoded.set(offset, true);
114                offset += 1;
115            }
116        }
117    }
118
119    let mut buf = Vec::with_capacity(
120        1 + space_needed(task_id_bits as u64)
121            + space_needed(schedule.len() as u64)
122            + space_needed(schedule.seed)
123            + encoded.len(),
124    );
125    buf.push(SCHEDULE_MAGIC_V2);
126    buf.write_u64_varint(task_id_bits as u64).unwrap();
127    buf.write_u64_varint(schedule.len() as u64).unwrap();
128    buf.write_u64_varint(schedule.seed).unwrap();
129    buf.extend(encoded.as_raw_slice());
130
131    let serialized = hex::encode(buf);
132    let lines = serialized.as_bytes().chunks(LINE_WIDTH).collect::<Vec<_>>();
133    let wrapped = lines.join(&b'\n');
134    String::from_utf8(wrapped).unwrap()
135}
136
137pub fn deserialize_schedule(str: &str) -> Option<Schedule> {
138    use self::varint::ReadVarInt;
139
140    let str: String = str.chars().filter(|c| !c.is_whitespace()).collect();
141    let bytes = hex::decode(str).ok()?;
142
143    let version = bytes[0];
144    if version != SCHEDULE_MAGIC_V2 {
145        return None;
146    }
147    let mut bytes = &bytes[1..];
148
149    let task_id_bits = bytes.read_u64_varint().ok()? as usize;
150    let schedule_len = bytes.read_u64_varint().ok()? as usize;
151    let seed = bytes.read_u64_varint().ok()?;
152
153    let encoded = BitSlice::<_, Lsb0>::from_slice(bytes);
154    let mut offset = 0usize;
155    let mut steps = Vec::with_capacity(schedule_len);
156    while steps.len() < schedule_len {
157        if *encoded.get(offset).unwrap() {
158            steps.push(ScheduleStep::Random);
159            offset += 1;
160        } else {
161            let tid = encoded[offset + 1..offset + 1 + task_id_bits].load::<usize>();
162            steps.push(ScheduleStep::Task(TaskId::from(tid)));
163            offset += 1 + task_id_bits;
164        }
165    }
166
167    Some(Schedule { seed, steps })
168}
169
170#[cfg(test)]
171mod test {
172    use super::*;
173    use proptest::{collection::vec, prelude::*};
174
175    // Schedules of up to 100 steps with TaskIds up to 10000
176    fn schedule_strategy() -> impl Strategy<Value = Schedule> {
177        let step_strategy = prop_oneof![
178            Just(ScheduleStep::Random),
179            (0usize..10000).prop_map(|tid| ScheduleStep::Task(TaskId::from(tid)))
180        ];
181        let steps_strategy = vec(step_strategy, (0, 100));
182        (any::<u64>(), steps_strategy).prop_map(|(seed, steps)| Schedule { seed, steps })
183    }
184
185    fn check_roundtrip(schedule: Schedule) {
186        let encoded = serialize_schedule(&schedule);
187        let decoded = deserialize_schedule(encoded.as_str()).unwrap();
188        assert_eq!(schedule, decoded);
189    }
190
191    #[test]
192    fn serialization_roundtrip_basic() {
193        check_roundtrip(Schedule {
194            seed: 10,
195            steps: vec![ScheduleStep::Random],
196        });
197        check_roundtrip(Schedule {
198            seed: 10,
199            steps: vec![ScheduleStep::Task(TaskId::from(0))],
200        });
201    }
202
203    proptest! {
204        #[test]
205        fn serialization_roundtrip_proptest(schedule in schedule_strategy()) {
206            check_roundtrip(schedule);
207        }
208    }
209}
shuttle_engine/scheduler/serialization.rs

shuttle_engine/scheduler/
serialization.rs
