randomized_test_helpers.rs

  1use crate::{
  2    db::{self, NewUserParams, UserId},
  3    rpc::{CLEANUP_TIMEOUT, RECONNECT_TIMEOUT},
  4    tests::{TestClient, TestServer},
  5};
  6use async_trait::async_trait;
  7use futures::StreamExt;
  8use gpui::{BackgroundExecutor, Task, TestAppContext};
  9use parking_lot::Mutex;
 10use rand::prelude::*;
 11use rpc::RECEIVE_TIMEOUT;
 12use serde::{Deserialize, Serialize, de::DeserializeOwned};
 13use settings::SettingsStore;
 14use std::sync::OnceLock;
 15use std::{
 16    env,
 17    path::PathBuf,
 18    rc::Rc,
 19    sync::{
 20        Arc,
 21        atomic::{AtomicBool, Ordering::SeqCst},
 22    },
 23};
 24
 25fn plan_load_path() -> &'static Option<PathBuf> {
 26    static PLAN_LOAD_PATH: OnceLock<Option<PathBuf>> = OnceLock::new();
 27    PLAN_LOAD_PATH.get_or_init(|| path_env_var("LOAD_PLAN"))
 28}
 29
 30fn plan_save_path() -> &'static Option<PathBuf> {
 31    static PLAN_SAVE_PATH: OnceLock<Option<PathBuf>> = OnceLock::new();
 32    PLAN_SAVE_PATH.get_or_init(|| path_env_var("SAVE_PLAN"))
 33}
 34
 35fn max_peers() -> usize {
 36    static MAX_PEERS: OnceLock<usize> = OnceLock::new();
 37    *MAX_PEERS.get_or_init(|| {
 38        env::var("MAX_PEERS")
 39            .map(|i| i.parse().expect("invalid `MAX_PEERS` variable"))
 40            .unwrap_or(3)
 41    })
 42}
 43
 44fn max_operations() -> usize {
 45    static MAX_OPERATIONS: OnceLock<usize> = OnceLock::new();
 46    *MAX_OPERATIONS.get_or_init(|| {
 47        env::var("OPERATIONS")
 48            .map(|i| i.parse().expect("invalid `OPERATIONS` variable"))
 49            .unwrap_or(10)
 50    })
 51}
 52
 53static LOADED_PLAN_JSON: Mutex<Option<Vec<u8>>> = Mutex::new(None);
 54static LAST_PLAN: Mutex<Option<Box<dyn Send + FnOnce() -> Vec<u8>>>> = Mutex::new(None);
 55
 56struct TestPlan<T: RandomizedTest> {
 57    rng: StdRng,
 58    replay: bool,
 59    stored_operations: Vec<(StoredOperation<T::Operation>, Arc<AtomicBool>)>,
 60    max_operations: usize,
 61    operation_ix: usize,
 62    users: Vec<UserTestPlan>,
 63    next_batch_id: usize,
 64    allow_server_restarts: bool,
 65    allow_client_reconnection: bool,
 66    allow_client_disconnection: bool,
 67}
 68
 69pub struct UserTestPlan {
 70    pub user_id: UserId,
 71    pub username: String,
 72    pub allow_client_disconnection: bool,
 73    next_root_id: usize,
 74    operation_ix: usize,
 75    online: bool,
 76}
 77
 78#[derive(Clone, Debug, Serialize, Deserialize)]
 79#[serde(untagged)]
 80enum StoredOperation<T> {
 81    Server(ServerOperation),
 82    Client {
 83        user_id: UserId,
 84        batch_id: usize,
 85        operation: T,
 86    },
 87}
 88
 89#[derive(Clone, Debug, Serialize, Deserialize)]
 90enum ServerOperation {
 91    AddConnection {
 92        user_id: UserId,
 93    },
 94    RemoveConnection {
 95        user_id: UserId,
 96    },
 97    BounceConnection {
 98        user_id: UserId,
 99    },
100    RestartServer,
101    MutateClients {
102        batch_id: usize,
103        #[serde(skip_serializing)]
104        #[serde(skip_deserializing)]
105        user_ids: Vec<UserId>,
106        quiesce: bool,
107    },
108}
109
110pub enum TestError {
111    Inapplicable,
112    Other(anyhow::Error),
113}
114
115#[async_trait(?Send)]
116pub trait RandomizedTest: 'static + Sized {
117    type Operation: Send + Clone + Serialize + DeserializeOwned;
118
119    fn generate_operation(
120        client: &TestClient,
121        rng: &mut StdRng,
122        plan: &mut UserTestPlan,
123        cx: &TestAppContext,
124    ) -> Self::Operation;
125
126    async fn apply_operation(
127        client: &TestClient,
128        operation: Self::Operation,
129        cx: &mut TestAppContext,
130    ) -> Result<(), TestError>;
131
132    async fn initialize(server: &mut TestServer, users: &[UserTestPlan]);
133
134    async fn on_client_added(_client: &Rc<TestClient>, _cx: &mut TestAppContext) {}
135
136    async fn on_quiesce(server: &mut TestServer, client: &mut [(Rc<TestClient>, TestAppContext)]);
137}
138
139pub async fn run_randomized_test<T: RandomizedTest>(
140    cx: &mut TestAppContext,
141    executor: BackgroundExecutor,
142    rng: StdRng,
143) {
144    let mut server = TestServer::start(executor.clone()).await;
145    let plan = TestPlan::<T>::new(&mut server, rng).await;
146
147    LAST_PLAN.lock().replace({
148        let plan = plan.clone();
149        Box::new(move || plan.lock().serialize())
150    });
151
152    let mut clients = Vec::new();
153    let mut client_tasks = Vec::new();
154    let mut operation_channels = Vec::new();
155    loop {
156        let Some((next_operation, applied)) = plan.lock().next_server_operation(&clients) else {
157            break;
158        };
159        applied.store(true, SeqCst);
160        let did_apply = TestPlan::apply_server_operation(
161            plan.clone(),
162            executor.clone(),
163            &mut server,
164            &mut clients,
165            &mut client_tasks,
166            &mut operation_channels,
167            next_operation,
168            cx,
169        )
170        .await;
171        if !did_apply {
172            applied.store(false, SeqCst);
173        }
174    }
175
176    drop(operation_channels);
177    futures::future::join_all(client_tasks).await;
178
179    executor.run_until_parked();
180    T::on_quiesce(&mut server, &mut clients).await;
181
182    for (client, cx) in clients {
183        cx.update(|cx| {
184            let settings = cx.remove_global::<SettingsStore>();
185            cx.clear_globals();
186            cx.set_global(settings);
187            theme::init(theme::LoadThemes::JustBase, cx);
188            drop(client);
189        });
190    }
191    executor.run_until_parked();
192
193    if let Some(path) = plan_save_path() {
194        eprintln!("saved test plan to path {:?}", path);
195        std::fs::write(path, plan.lock().serialize()).unwrap();
196    }
197}
198
199pub fn save_randomized_test_plan() {
200    if let Some(serialize_plan) = LAST_PLAN.lock().take()
201        && let Some(path) = plan_save_path()
202    {
203        eprintln!("saved test plan to path {:?}", path);
204        std::fs::write(path, serialize_plan()).unwrap();
205    }
206}
207
208impl<T: RandomizedTest> TestPlan<T> {
209    pub async fn new(server: &mut TestServer, mut rng: StdRng) -> Arc<Mutex<Self>> {
210        let allow_server_restarts = rng.random_bool(0.7);
211        let allow_client_reconnection = rng.random_bool(0.7);
212        let allow_client_disconnection = rng.random_bool(0.1);
213
214        let mut users = Vec::new();
215        for ix in 0..max_peers() {
216            let username = format!("user-{}", ix + 1);
217            let user_id = server
218                .app_state
219                .db
220                .create_user(
221                    &format!("{username}@example.com"),
222                    None,
223                    false,
224                    NewUserParams {
225                        github_login: username.clone(),
226                        github_user_id: ix as i32,
227                    },
228                )
229                .await
230                .unwrap()
231                .user_id;
232            users.push(UserTestPlan {
233                user_id,
234                username,
235                online: false,
236                next_root_id: 0,
237                operation_ix: 0,
238                allow_client_disconnection,
239            });
240        }
241
242        T::initialize(server, &users).await;
243
244        let plan = Arc::new(Mutex::new(Self {
245            replay: false,
246            allow_server_restarts,
247            allow_client_reconnection,
248            allow_client_disconnection,
249            stored_operations: Vec::new(),
250            operation_ix: 0,
251            next_batch_id: 0,
252            max_operations: max_operations(),
253            users,
254            rng,
255        }));
256
257        if let Some(path) = plan_load_path() {
258            let json = LOADED_PLAN_JSON
259                .lock()
260                .get_or_insert_with(|| {
261                    eprintln!("loaded test plan from path {:?}", path);
262                    std::fs::read(path).unwrap()
263                })
264                .clone();
265            plan.lock().deserialize(json);
266        }
267
268        plan
269    }
270
271    fn deserialize(&mut self, json: Vec<u8>) {
272        let stored_operations: Vec<StoredOperation<T::Operation>> =
273            serde_json::from_slice(&json).unwrap();
274        self.replay = true;
275        self.stored_operations = stored_operations
276            .iter()
277            .cloned()
278            .enumerate()
279            .map(|(i, mut operation)| {
280                let did_apply = Arc::new(AtomicBool::new(false));
281                if let StoredOperation::Server(ServerOperation::MutateClients {
282                    batch_id: current_batch_id,
283                    user_ids,
284                    ..
285                }) = &mut operation
286                {
287                    assert!(user_ids.is_empty());
288                    user_ids.extend(stored_operations[i + 1..].iter().filter_map(|operation| {
289                        if let StoredOperation::Client {
290                            user_id, batch_id, ..
291                        } = operation
292                            && batch_id == current_batch_id
293                        {
294                            return Some(user_id);
295                        }
296                        None
297                    }));
298                    user_ids.sort_unstable();
299                }
300                (operation, did_apply)
301            })
302            .collect()
303    }
304
305    fn serialize(&mut self) -> Vec<u8> {
306        // Format each operation as one line
307        let mut json = Vec::new();
308        json.push(b'[');
309        for (operation, applied) in &self.stored_operations {
310            if !applied.load(SeqCst) {
311                continue;
312            }
313            if json.len() > 1 {
314                json.push(b',');
315            }
316            json.extend_from_slice(b"\n  ");
317            serde_json::to_writer(&mut json, operation).unwrap();
318        }
319        json.extend_from_slice(b"\n]\n");
320        json
321    }
322
323    fn next_server_operation(
324        &mut self,
325        clients: &[(Rc<TestClient>, TestAppContext)],
326    ) -> Option<(ServerOperation, Arc<AtomicBool>)> {
327        if self.replay {
328            while let Some(stored_operation) = self.stored_operations.get(self.operation_ix) {
329                self.operation_ix += 1;
330                if let (StoredOperation::Server(operation), applied) = stored_operation {
331                    return Some((operation.clone(), applied.clone()));
332                }
333            }
334            None
335        } else {
336            let operation = self.generate_server_operation(clients)?;
337            let applied = Arc::new(AtomicBool::new(false));
338            self.stored_operations
339                .push((StoredOperation::Server(operation.clone()), applied.clone()));
340            Some((operation, applied))
341        }
342    }
343
344    fn next_client_operation(
345        &mut self,
346        client: &TestClient,
347        current_batch_id: usize,
348        cx: &TestAppContext,
349    ) -> Option<(T::Operation, Arc<AtomicBool>)> {
350        let current_user_id = client.current_user_id(cx);
351        let user_ix = self
352            .users
353            .iter()
354            .position(|user| user.user_id == current_user_id)
355            .unwrap();
356        let user_plan = &mut self.users[user_ix];
357
358        if self.replay {
359            while let Some(stored_operation) = self.stored_operations.get(user_plan.operation_ix) {
360                user_plan.operation_ix += 1;
361                if let (
362                    StoredOperation::Client {
363                        user_id, operation, ..
364                    },
365                    applied,
366                ) = stored_operation
367                    && user_id == &current_user_id
368                {
369                    return Some((operation.clone(), applied.clone()));
370                }
371            }
372            None
373        } else {
374            if self.operation_ix == self.max_operations {
375                return None;
376            }
377            self.operation_ix += 1;
378            let operation = T::generate_operation(
379                client,
380                &mut self.rng,
381                self.users
382                    .iter_mut()
383                    .find(|user| user.user_id == current_user_id)
384                    .unwrap(),
385                cx,
386            );
387            let applied = Arc::new(AtomicBool::new(false));
388            self.stored_operations.push((
389                StoredOperation::Client {
390                    user_id: current_user_id,
391                    batch_id: current_batch_id,
392                    operation: operation.clone(),
393                },
394                applied.clone(),
395            ));
396            Some((operation, applied))
397        }
398    }
399
400    fn generate_server_operation(
401        &mut self,
402        clients: &[(Rc<TestClient>, TestAppContext)],
403    ) -> Option<ServerOperation> {
404        if self.operation_ix == self.max_operations {
405            return None;
406        }
407
408        Some(loop {
409            break match self.rng.random_range(0..100) {
410                0..=29 if clients.len() < self.users.len() => {
411                    let user = self
412                        .users
413                        .iter()
414                        .filter(|u| !u.online)
415                        .choose(&mut self.rng)
416                        .unwrap();
417                    self.operation_ix += 1;
418                    ServerOperation::AddConnection {
419                        user_id: user.user_id,
420                    }
421                }
422                30..=34 if clients.len() > 1 && self.allow_client_disconnection => {
423                    let (client, cx) = &clients[self.rng.random_range(0..clients.len())];
424                    let user_id = client.current_user_id(cx);
425                    self.operation_ix += 1;
426                    ServerOperation::RemoveConnection { user_id }
427                }
428                35..=39 if clients.len() > 1 && self.allow_client_reconnection => {
429                    let (client, cx) = &clients[self.rng.random_range(0..clients.len())];
430                    let user_id = client.current_user_id(cx);
431                    self.operation_ix += 1;
432                    ServerOperation::BounceConnection { user_id }
433                }
434                40..=44 if self.allow_server_restarts && clients.len() > 1 => {
435                    self.operation_ix += 1;
436                    ServerOperation::RestartServer
437                }
438                _ if !clients.is_empty() => {
439                    let count = self
440                        .rng
441                        .random_range(1..10)
442                        .min(self.max_operations - self.operation_ix);
443                    let batch_id = util::post_inc(&mut self.next_batch_id);
444                    let mut user_ids = (0..count)
445                        .map(|_| {
446                            let ix = self.rng.random_range(0..clients.len());
447                            let (client, cx) = &clients[ix];
448                            client.current_user_id(cx)
449                        })
450                        .collect::<Vec<_>>();
451                    user_ids.sort_unstable();
452                    ServerOperation::MutateClients {
453                        user_ids,
454                        batch_id,
455                        quiesce: self.rng.random_bool(0.7),
456                    }
457                }
458                _ => continue,
459            };
460        })
461    }
462
463    async fn apply_server_operation(
464        plan: Arc<Mutex<Self>>,
465        deterministic: BackgroundExecutor,
466        server: &mut TestServer,
467        clients: &mut Vec<(Rc<TestClient>, TestAppContext)>,
468        client_tasks: &mut Vec<Task<()>>,
469        operation_channels: &mut Vec<futures::channel::mpsc::UnboundedSender<usize>>,
470        operation: ServerOperation,
471        cx: &mut TestAppContext,
472    ) -> bool {
473        match operation {
474            ServerOperation::AddConnection { user_id } => {
475                let username;
476                {
477                    let mut plan = plan.lock();
478                    let user = plan.user(user_id);
479                    if user.online {
480                        return false;
481                    }
482                    user.online = true;
483                    username = user.username.clone();
484                };
485                log::info!("adding new connection for {}", username);
486
487                let mut client_cx = cx.new_app();
488
489                let (operation_tx, operation_rx) = futures::channel::mpsc::unbounded();
490                let client = Rc::new(server.create_client(&mut client_cx, &username).await);
491                operation_channels.push(operation_tx);
492                clients.push((client.clone(), client_cx.clone()));
493
494                let foreground_executor = client_cx.foreground_executor().clone();
495                let simulate_client =
496                    Self::simulate_client(plan.clone(), client, operation_rx, client_cx);
497                client_tasks.push(foreground_executor.spawn(simulate_client));
498
499                log::info!("added connection for {}", username);
500            }
501
502            ServerOperation::RemoveConnection {
503                user_id: removed_user_id,
504            } => {
505                log::info!("simulating full disconnection of user {}", removed_user_id);
506                let client_ix = clients
507                    .iter()
508                    .position(|(client, cx)| client.current_user_id(cx) == removed_user_id);
509                let Some(client_ix) = client_ix else {
510                    return false;
511                };
512                let user_connection_ids = server
513                    .connection_pool
514                    .lock()
515                    .user_connection_ids(removed_user_id)
516                    .collect::<Vec<_>>();
517                assert_eq!(user_connection_ids.len(), 1);
518                let removed_peer_id = user_connection_ids[0].into();
519                let (client, client_cx) = clients.remove(client_ix);
520                let client_task = client_tasks.remove(client_ix);
521                operation_channels.remove(client_ix);
522                server.forbid_connections();
523                server.disconnect_client(removed_peer_id);
524                deterministic.advance_clock(RECEIVE_TIMEOUT + RECONNECT_TIMEOUT);
525                log::info!("waiting for user {} to exit...", removed_user_id);
526                client_task.await;
527                server.allow_connections();
528
529                for project in client.dev_server_projects().iter() {
530                    project.read_with(&client_cx, |project, cx| {
531                        assert!(
532                            project.is_disconnected(cx),
533                            "project {:?} should be read only",
534                            project.remote_id()
535                        )
536                    });
537                }
538
539                for (client, cx) in clients {
540                    let contacts = server
541                        .app_state
542                        .db
543                        .get_contacts(client.current_user_id(cx))
544                        .await
545                        .unwrap();
546                    let pool = server.connection_pool.lock();
547                    for contact in contacts {
548                        if let db::Contact::Accepted { user_id, busy, .. } = contact
549                            && user_id == removed_user_id
550                        {
551                            assert!(!pool.is_user_online(user_id));
552                            assert!(!busy);
553                        }
554                    }
555                }
556
557                log::info!("{} removed", client.username);
558                plan.lock().user(removed_user_id).online = false;
559                client_cx.update(|cx| {
560                    cx.clear_globals();
561                    drop(client);
562                });
563            }
564
565            ServerOperation::BounceConnection { user_id } => {
566                log::info!("simulating temporary disconnection of user {}", user_id);
567                let user_connection_ids = server
568                    .connection_pool
569                    .lock()
570                    .user_connection_ids(user_id)
571                    .collect::<Vec<_>>();
572                if user_connection_ids.is_empty() {
573                    return false;
574                }
575                assert_eq!(user_connection_ids.len(), 1);
576                let peer_id = user_connection_ids[0].into();
577                server.disconnect_client(peer_id);
578                deterministic.advance_clock(RECEIVE_TIMEOUT + RECONNECT_TIMEOUT);
579            }
580
581            ServerOperation::RestartServer => {
582                log::info!("simulating server restart");
583                server.reset().await;
584                deterministic.advance_clock(RECEIVE_TIMEOUT);
585                server.start().await.unwrap();
586                deterministic.advance_clock(CLEANUP_TIMEOUT);
587                let environment = &server.app_state.config.zed_environment;
588                let (stale_room_ids, _) = server
589                    .app_state
590                    .db
591                    .stale_server_resource_ids(environment, server.id())
592                    .await
593                    .unwrap();
594                assert_eq!(stale_room_ids, vec![]);
595            }
596
597            ServerOperation::MutateClients {
598                user_ids,
599                batch_id,
600                quiesce,
601            } => {
602                let mut applied = false;
603                for user_id in user_ids {
604                    let client_ix = clients
605                        .iter()
606                        .position(|(client, cx)| client.current_user_id(cx) == user_id);
607                    let Some(client_ix) = client_ix else { continue };
608                    applied = true;
609                    if let Err(err) = operation_channels[client_ix].unbounded_send(batch_id) {
610                        log::error!("error signaling user {user_id}: {err}");
611                    }
612                }
613
614                if quiesce && applied {
615                    deterministic.run_until_parked();
616                    T::on_quiesce(server, clients).await;
617                }
618
619                return applied;
620            }
621        }
622        true
623    }
624
625    async fn simulate_client(
626        plan: Arc<Mutex<Self>>,
627        client: Rc<TestClient>,
628        mut operation_rx: futures::channel::mpsc::UnboundedReceiver<usize>,
629        mut cx: TestAppContext,
630    ) {
631        T::on_client_added(&client, &mut cx).await;
632
633        while let Some(batch_id) = operation_rx.next().await {
634            let Some((operation, applied)) =
635                plan.lock().next_client_operation(&client, batch_id, &cx)
636            else {
637                break;
638            };
639            applied.store(true, SeqCst);
640            match T::apply_operation(&client, operation, &mut cx).await {
641                Ok(()) => {}
642                Err(TestError::Inapplicable) => {
643                    applied.store(false, SeqCst);
644                    log::info!("skipped operation");
645                }
646                Err(TestError::Other(error)) => {
647                    log::error!("{} error: {}", client.username, error);
648                }
649            }
650            cx.executor().simulate_random_delay().await;
651        }
652        log::info!("{}: done", client.username);
653    }
654
655    fn user(&mut self, user_id: UserId) -> &mut UserTestPlan {
656        self.users
657            .iter_mut()
658            .find(|user| user.user_id == user_id)
659            .unwrap()
660    }
661}
662
663impl UserTestPlan {
664    pub fn next_root_dir_name(&mut self) -> String {
665        let user_id = self.user_id;
666        let root_id = util::post_inc(&mut self.next_root_id);
667        format!("dir-{user_id}-{root_id}")
668    }
669}
670
671impl From<anyhow::Error> for TestError {
672    fn from(value: anyhow::Error) -> Self {
673        Self::Other(value)
674    }
675}
676
677fn path_env_var(name: &str) -> Option<PathBuf> {
678    let value = env::var(name).ok()?;
679    let mut path = PathBuf::from(value);
680    if path.is_relative() {
681        let mut abs_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
682        abs_path.pop();
683        abs_path.pop();
684        abs_path.push(path);
685        path = abs_path
686    }
687    Some(path)
688}