| 1 | #![allow(dead_code)] |
| 2 | |
| 3 | use std::collections::{HashMap, VecDeque}; |
| 4 | use std::sync::{Arc, Mutex, MutexGuard}; |
| 5 | |
| 6 | use fastrand::Rng; |
| 7 | |
| 8 | use crate::block::ClientID; |
| 9 | use crate::encoding::read::{Cursor, Read}; |
| 10 | use crate::transaction::ReadTxn; |
| 11 | use crate::updates::decoder::{Decode, Decoder, DecoderV1}; |
| 12 | use crate::updates::encoder::{Encode, Encoder, EncoderV1}; |
| 13 | use crate::{Doc, StateVector, Transact, Update}; |
| 14 | |
| 15 | pub const EXCHANGE_UPDATES_ORIGIN: &str = "exchange_updates"; |
| 16 | |
| 17 | pub 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 | |
| 35 | const MSG_SYNC_STEP_1: usize = 0; |
| 36 | const MSG_SYNC_STEP_2: usize = 1; |
| 37 | const MSG_SYNC_UPDATE: usize = 2; |
| 38 | |
| 39 | pub fn run_scenario<F>(mut seed: u64, mods: &[F], users: usize, iterations: usize) |
| 40 | where |
| 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 | |
| 80 | pub struct TestConnector(Arc<Mutex<Inner>>); |
| 81 | |
| 82 | struct 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 | |
| 91 | impl 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 | |
| 438 | pub struct Peers<'a> { |
| 439 | inner: MutexGuard<'a, Inner>, |
| 440 | i: usize, |
| 441 | } |
| 442 | |
| 443 | impl<'a> Peers<'a> { |
| 444 | fn new(inner: MutexGuard<'a, Inner>) -> Self { |
| 445 | Self { inner, i: 0 } |
| 446 | } |
| 447 | } |
| 448 | |
| 449 | impl<'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)] |
| 465 | pub struct TestPeer { |
| 466 | state: Arc<Mutex<TestPeerState>>, |
| 467 | } |
| 468 | |
| 469 | #[derive(Debug)] |
| 470 | struct TestPeerState { |
| 471 | doc: Doc, |
| 472 | receiving: HashMap<ClientID, VecDeque<Vec<u8>>>, |
| 473 | updates: VecDeque<Vec<u8>>, |
| 474 | } |
| 475 | |
| 476 | impl 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 | |
| 525 | pub(crate) trait RngExt { |
| 526 | fn between(&mut self, x: u32, y: u32) -> u32; |
| 527 | |
| 528 | fn random_string(&mut self) -> String; |
| 529 | } |
| 530 | |
| 531 | impl 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 | |