oficina / docs Public

Clone
git clone https://amadan.net/oficina/docs

Plain git — no account needed to clone.

Download

Download this file

1#![allow(dead_code)]
2
3use std::collections::{HashMap, VecDeque};
4use std::sync::{Arc, Mutex, MutexGuard};
5
6use fastrand::Rng;
7
8use crate::block::ClientID;
9use crate::encoding::read::{Cursor, Read};
10use crate::transaction::ReadTxn;
11use crate::updates::decoder::{Decode, Decoder, DecoderV1};
12use crate::updates::encoder::{Encode, Encoder, EncoderV1};
13use crate::{Doc, StateVector, Transact, Update};
14
15pub const EXCHANGE_UPDATES_ORIGIN: &str = "exchange_updates";
16
17pub fn exchange_updates(docs: &[&Doc]) {
18 for i in 0..docs.len() {
19 for j in 0..docs.len() {
20 if i != j {
21 let a = docs[i];
22 let ta = a.transact();
23 let b = docs[j];
24 let mut tb = b.transact_mut_with(EXCHANGE_UPDATES_ORIGIN);
25
26 let sv = tb.state_vector().encode_v1();
27 let update = ta.encode_diff_v1(&StateVector::decode_v1(sv.as_slice()).unwrap());
28 let update = Update::decode_v1(update.as_slice()).unwrap();
29 tb.apply_update(update).unwrap();
30 }
31 }
32 }
33}
34
35const MSG_SYNC_STEP_1: usize = 0;
36const MSG_SYNC_STEP_2: usize = 1;
37const MSG_SYNC_UPDATE: usize = 2;
38
39pub fn run_scenario<F>(mut seed: u64, mods: &[F], users: usize, iterations: usize)
40where
41 F: Fn(&mut Doc, &mut Rng),
42{
43 if seed == 0 {
44 seed = fastrand::get_seed();
45 println!("run scenario with seed: {}", seed);
46 }
47
48 let rng = Rng::with_seed(seed);
49 let tc = TestConnector::with_peer_num(rng, users as u64);
50 for _ in 0..iterations {
51 if tc.0.lock().unwrap().rng.u32(0..100) <= 2 {
52 // 2% chance to disconnect/reconnect a random user
53 if tc.0.lock().unwrap().rng.bool() {
54 tc.disconnect_random();
55 } else {
56 tc.reconnect_random();
57 }
58 } else if tc.0.lock().unwrap().rng.u32(0..100) <= 1 {
59 // 1% chance to flush all
60 tc.flush_all();
61 } else if tc.0.lock().unwrap().rng.u32(0..100) <= 50 {
62 tc.flush_random();
63 }
64
65 {
66 let mut inner = tc.0.lock().unwrap();
67 let inner = &mut *inner;
68 let rng = &mut inner.rng;
69 let idx = rng.usize(0..inner.peers.len());
70 let peer = &mut inner.peers[idx];
71 let test = rng.choice(mods).unwrap();
72 let mut peer_state = peer.state();
73 test(&mut peer_state.doc, rng);
74 };
75 }
76
77 tc.assert_final_state();
78}
79
80pub struct TestConnector(Arc<Mutex<Inner>>);
81
82struct Inner {
83 rng: Rng,
84 peers: Vec<TestPeer>,
85 /// Maps all Client IDs to indexes in the `docs` vector.
86 all: HashMap<ClientID, usize>,
87 /// Maps online Client IDs to indexes in the `docs` vector.
88 online: HashMap<ClientID, usize>,
89}
90
91impl TestConnector {
92 /// Create new [TestConnector] with provided randomizer.
93 pub fn with_rng(rng: Rng) -> Self {
94 TestConnector(Arc::new(Mutex::new(Inner {
95 rng,
96 peers: Vec::new(),
97 all: HashMap::new(),
98 online: HashMap::new(),
99 })))
100 }
101
102 /// Create a new [TestConnector] with pre-initialized number of peers.
103 pub fn with_peer_num(rng: Rng, peer_num: u64) -> Self {
104 let mut tc = Self::with_rng(rng);
105 for client_id in 0..peer_num {
106 let peer = tc.create_peer(client_id as ClientID);
107 let peer_state = peer.state();
108 peer_state.doc.get_or_insert_text("text");
109 peer_state.doc.get_or_insert_map("map");
110 }
111 tc.sync_all();
112 tc
113 }
114
115 /// Create a new [TestPeer] with provided `client_id` or return one, if such `client_id`
116 /// was already created before.
117 pub fn create_peer(&self, client_id: ClientID) -> TestPeer {
118 if let Some(peer) = self.get(&client_id) {
119 peer
120 } else {
121 let rc = self.0.clone();
122 let instance = TestPeer::new(client_id);
123 let _sub = {
124 let rc = rc.clone();
125 let peer_state = instance.state();
126 peer_state
127 .doc
128 .observe_update_v1(move |_, e| {
129 let mut inner = rc.lock().unwrap();
130 Self::broadcast(&mut inner, client_id, &e.update);
131 })
132 .unwrap()
133 };
134 let mut inner = rc.lock().unwrap();
135 let idx = inner.peers.len();
136 inner.peers.push(instance);
137 inner.all.insert(client_id, idx);
138 inner.online.insert(client_id, idx);
139 inner.peers[idx].clone()
140 }
141 }
142
143 fn broadcast(inner: &mut Inner, sender: ClientID, payload: &Vec<u8>) {
144 let online: Vec<_> = inner
145 .online
146 .iter()
147 .filter_map(|(&id, &idx)| if id != sender { Some(idx) } else { None })
148 .collect();
149 for idx in online {
150 let peer = &mut inner.peers[idx];
151 peer.receive(sender, payload.clone());
152 }
153 }
154
155 /// Try to retrieve a reference to [TestPeer] for a given `client_id`, if such node was created.
156 pub fn get(&self, client_id: &ClientID) -> Option<TestPeer> {
157 let inner = self.0.lock().unwrap();
158 let idx = inner.all.get(client_id)?;
159 Some(inner.peers[*idx].clone())
160 }
161
162 /// Disconnects test node with given `client_id` from the rest of known nodes.
163 pub fn disconnect(&self, client_id: ClientID) {
164 if let Some(peer) = self.get(&client_id) {
165 peer.clear();
166 }
167 let mut inner = self.0.lock().unwrap();
168 inner.online.remove(&client_id);
169 }
170
171 /// Append `client_id` to the list of known Y instances in [TestConnector].
172 /// Also initiate sync with all clients.
173 pub fn connect(&self, client_id: ClientID) {
174 let mut inner = self.0.lock().unwrap();
175 Self::connect_inner(&mut *inner, client_id);
176 }
177
178 fn connect_inner(inner: &mut Inner, client_id: ClientID) {
179 if !inner.online.contains_key(&client_id) {
180 let idx = *inner.all.get(&client_id).expect("unknown client_id");
181 inner.online.insert(client_id, idx);
182 }
183
184 let client_idx = *inner.all.get(&client_id).unwrap();
185 let payload = {
186 let sender = &inner.peers[client_idx];
187 let mut sender = sender.state();
188 let mut encoder = EncoderV1::new();
189 Self::write_step1(&mut *sender, &mut encoder);
190 encoder.to_vec()
191 };
192 Self::broadcast(inner, client_id, &payload);
193
194 let online: Vec<_> = inner
195 .online
196 .iter()
197 .filter_map(|(&id, &idx)| {
198 if id != client_id {
199 Some((id, idx))
200 } else {
201 None
202 }
203 })
204 .collect();
205 for (remote_id, idx) in online {
206 let payload = {
207 let peer = &inner.peers[idx];
208 let mut peer = peer.state();
209 let mut encoder = EncoderV1::new();
210 Self::write_step1(&mut *peer, &mut encoder);
211 encoder.to_vec()
212 };
213
214 let sender = &mut inner.peers[client_idx];
215 sender.receive(remote_id, payload);
216 }
217 }
218
219 /// Reconnects back all known peers.
220 pub fn reconnect_all(&mut self) {
221 let mut inner = self.0.lock().unwrap();
222 let all_ids: Vec<_> = inner.all.keys().cloned().collect();
223 for client_id in all_ids {
224 Self::connect_inner(&mut inner, client_id);
225 }
226 }
227
228 /// Disconnects all known peers from each other.
229 pub fn disconnect_all(&mut self) {
230 let mut inner = self.0.lock().unwrap();
231 let all_ids: Vec<_> = inner.all.keys().cloned().collect();
232 for client_id in all_ids {
233 Self::connect_inner(&mut inner, client_id);
234 }
235 }
236
237 /// Reconnects all known peers and processes their pending messages.
238 pub fn sync_all(&mut self) {
239 self.reconnect_all();
240 self.flush_all();
241 }
242
243 /// Processes all pending messages of connected peers in random order.
244 pub fn flush_all(&self) -> bool {
245 let mut did_something = false;
246 while self.flush_random() {
247 did_something = true;
248 }
249 did_something
250 }
251
252 /// Choose random connection and flush a random message from a random sender.
253 /// If this function was unable to flush a message, because there are no more messages to flush,
254 /// it returns false. true otherwise.
255 pub fn flush_random(&self) -> bool {
256 let mut inner = self.0.lock().unwrap();
257 Self::flush_random_inner(&mut *inner)
258 }
259
260 fn flush_random_inner(inner: &mut Inner) -> bool {
261 if let Some((receiver, sender)) = Self::pick_random_pair(inner) {
262 if let Some(m) = receiver.try_recv(&sender.client_id()) {
263 let mut encoder = EncoderV1::new();
264 let mut decoder = DecoderV1::new(Cursor::new(m.as_slice()));
265 Self::read_sync_message(receiver, &mut decoder, &mut encoder);
266 let payload = encoder.to_vec();
267 if !payload.is_empty() {
268 sender.receive(receiver.client_id(), payload); // send reply message
269 }
270
271 // If update message, add the received message to the list of received messages
272 {
273 let mut decoder = DecoderV1::new(Cursor::new(m.as_slice()));
274 let msg_type: usize = decoder.read_var().unwrap();
275 if msg_type == MSG_SYNC_STEP_2 || msg_type == MSG_SYNC_UPDATE {
276 receiver.send(decoder.read_buf().unwrap().to_vec())
277 }
278 }
279 true
280 } else {
281 receiver.remove(&sender.client_id());
282 Self::flush_random_inner(inner)
283 }
284 } else {
285 false
286 }
287 }
288
289 fn pick_random_pair(inner: &mut Inner) -> Option<(&mut TestPeer, &mut TestPeer)> {
290 let pairs: Vec<_> = inner
291 .peers
292 .iter()
293 .enumerate()
294 .flat_map(|(receiver_idx, conn)| {
295 let conn = conn.state();
296 if conn.receiving.is_empty() {
297 vec![]
298 } else {
299 conn.receiving
300 .iter()
301 .map(|(key, _)| (receiver_idx, *inner.all.get(key).unwrap()))
302 .collect()
303 }
304 })
305 .collect();
306 let (receiver_idx, sender_idx) = inner.rng.choice(pairs)?;
307 unsafe {
308 let ptr = inner.peers.as_ptr() as *mut TestPeer;
309 let receiver = ptr.offset(receiver_idx as isize);
310 let sender = ptr.offset(sender_idx as isize);
311 Some((receiver.as_mut().unwrap(), sender.as_mut().unwrap()))
312 }
313 }
314
315 /// Disconnects one peer at random.
316 pub fn disconnect_random(&self) -> bool {
317 let id = {
318 let mut inner = self.0.lock().unwrap();
319 let keys: Vec<_> = inner.online.keys().cloned().collect();
320 let rng = &mut inner.rng;
321 rng.choice(keys).clone()
322 };
323 if let Some(id) = id {
324 self.disconnect(id);
325 true
326 } else {
327 false
328 }
329 }
330
331 /// Reconnects one previously disconnected peer at random.
332 pub fn reconnect_random(&self) -> bool {
333 let mut inner = self.0.lock().unwrap();
334 let reconnectable: Vec<_> = inner
335 .all
336 .keys()
337 .filter(|&id| !inner.online.contains_key(id))
338 .cloned()
339 .collect();
340 if let Some(id) = inner.rng.choice(reconnectable) {
341 Self::connect_inner(&mut inner, id);
342 true
343 } else {
344 false
345 }
346 }
347
348 fn read_sync_message<D: Decoder, E: Encoder>(
349 peer: &TestPeer,
350 decoder: &mut D,
351 encoder: &mut E,
352 ) -> usize {
353 let msg_type = decoder.read_var().unwrap();
354 let mut peer = peer.state();
355 match msg_type {
356 MSG_SYNC_STEP_1 => Self::read_sync_step1(&mut *peer, decoder, encoder),
357 MSG_SYNC_STEP_2 => Self::read_sync_step2(&mut *peer, decoder),
358 MSG_SYNC_UPDATE => Self::read_update(&mut *peer, decoder),
359 other => panic!(
360 "Unknown message type: {} to {}",
361 other,
362 peer.doc.client_id()
363 ),
364 }
365 msg_type
366 }
367
368 fn read_sync_step1<D: Decoder, E: Encoder>(
369 peer: &mut TestPeerState,
370 decoder: &mut D,
371 encoder: &mut E,
372 ) {
373 Self::write_step2(peer, decoder.read_buf().unwrap(), encoder)
374 }
375
376 fn read_sync_step2<D: Decoder>(peer: &mut TestPeerState, decoder: &mut D) {
377 let mut txn = peer.doc.transact_mut();
378
379 let update = Update::decode_v1(decoder.read_buf().unwrap()).unwrap();
380 txn.apply_update(update).unwrap();
381 }
382
383 fn read_update<D: Decoder>(peer: &mut TestPeerState, decoder: &mut D) {
384 Self::read_sync_step2(peer, decoder)
385 }
386
387 /// Create a sync step 1 message based on the state of the current shared document.
388 fn write_step1<E: Encoder>(peer: &mut TestPeerState, encoder: &mut E) {
389 let txn = peer.doc.transact_mut();
390
391 encoder.write_var(MSG_SYNC_STEP_1);
392 encoder.write_buf(txn.state_vector().encode_v1());
393 }
394
395 fn write_step2<E: Encoder>(peer: &mut TestPeerState, sv: &[u8], encoder: &mut E) {
396 let txn = peer.doc.transact_mut();
397 let remote_sv = StateVector::decode_v1(sv).unwrap();
398
399 encoder.write_var(MSG_SYNC_STEP_2);
400 encoder.write_buf(txn.encode_diff_v1(&remote_sv));
401 }
402
403 pub fn assert_final_state(mut self) {
404 self.reconnect_all();
405 while self.flush_all() { /* do nothing */ }
406 // For each document, merge all received document updates with Y.mergeUpdates
407 // and create a new document which will be added to the list of "users"
408 // This ensures that mergeUpdates works correctly
409 /*
410 const mergedDocs = users.map(user => {
411 const ydoc = new Y.Doc()
412 enc.applyUpdate(ydoc, enc.mergeUpdates(user.updates))
413 return ydoc
414 })
415 users.push(.../** @type {any} */(mergedDocs))
416 */
417 let inner = self.0.lock().unwrap();
418 for i in 0..(inner.peers.len() - 1) {
419 let p1 = inner.peers[i].state();
420 let p2 = inner.peers[i + 1].state();
421 let a = p1.doc.transact_mut();
422 let b = p2.doc.transact_mut();
423
424 let astore = a.store();
425 let bstore = b.store();
426 assert_eq!(astore.blocks, bstore.blocks);
427 assert_eq!(astore.pending, bstore.pending);
428 assert_eq!(astore.pending_ds, bstore.pending_ds);
429 }
430 }
431
432 pub fn peers(&self) -> Peers {
433 let inner = self.0.lock().unwrap();
434 Peers::new(inner)
435 }
436}
437
438pub struct Peers<'a> {
439 inner: MutexGuard<'a, Inner>,
440 i: usize,
441}
442
443impl<'a> Peers<'a> {
444 fn new(inner: MutexGuard<'a, Inner>) -> Self {
445 Self { inner, i: 0 }
446 }
447}
448
449impl<'a> Iterator for Peers<'a> {
450 type Item = TestPeer;
451
452 fn next(&mut self) -> Option<Self::Item> {
453 if self.i >= self.inner.peers.len() {
454 None
455 } else {
456 let peer = self.inner.peers[self.i].clone();
457 self.i += 1;
458 Some(peer)
459 }
460 }
461}
462
463#[repr(transparent)]
464#[derive(Debug, Clone)]
465pub struct TestPeer {
466 state: Arc<Mutex<TestPeerState>>,
467}
468
469#[derive(Debug)]
470struct TestPeerState {
471 doc: Doc,
472 receiving: HashMap<ClientID, VecDeque<Vec<u8>>>,
473 updates: VecDeque<Vec<u8>>,
474}
475
476impl TestPeer {
477 pub fn new(client_id: ClientID) -> Self {
478 TestPeer {
479 state: Arc::new(Mutex::new(TestPeerState {
480 doc: Doc::with_client_id(client_id),
481 receiving: HashMap::new(),
482 updates: VecDeque::new(),
483 })),
484 }
485 }
486
487 fn state(&self) -> MutexGuard<TestPeerState> {
488 self.state.lock().unwrap()
489 }
490
491 fn remove(&self, client: &ClientID) {
492 let mut state = self.state();
493 state.receiving.remove(client);
494 }
495
496 fn send(&self, data: Vec<u8>) {
497 let mut state = self.state();
498 state.updates.push_back(data);
499 }
500
501 fn try_recv(&self, sender: &ClientID) -> Option<Vec<u8>> {
502 let mut state = self.state();
503 let client = state.receiving.get_mut(sender)?;
504 client.pop_front()
505 }
506
507 pub fn client_id(&self) -> ClientID {
508 self.state().doc.client_id()
509 }
510
511 fn clear(&self) {
512 let mut state = self.state();
513 state.receiving.clear();
514 }
515
516 /// Receive a message from another client. This message is only appended to the list of
517 /// receiving messages. TestConnector decides when this client actually reads this message.
518 fn receive(&self, from: ClientID, message: Vec<u8>) {
519 let mut state = self.state();
520 let messages = state.receiving.entry(from).or_default();
521 messages.push_back(message);
522 }
523}
524
525pub(crate) trait RngExt {
526 fn between(&mut self, x: u32, y: u32) -> u32;
527
528 fn random_string(&mut self) -> String;
529}
530
531impl RngExt for Rng {
532 fn between(&mut self, x: u32, y: u32) -> u32 {
533 let a = x.min(y);
534 let b = x.max(y);
535 if a == b {
536 a
537 } else {
538 self.u32(a..b)
539 }
540 }
541
542 fn random_string(&mut self) -> String {
543 let mut res = String::new();
544 for _ in 0..self.usize(1..10) {
545 res.push(self.alphanumeric());
546 }
547 res
548 }
549}
550