Line data    Source code

       1              : use std::{sync::Arc, time::Duration};
       2              : 
       3              : use rand::seq::SliceRandom;
       4              : use rand::thread_rng;
       5              : use tokio_util::sync::CancellationToken;
       6              : 
       7              : use super::Service;
       8              : 
       9              : pub struct ChaosInjector {
      10              :     service: Arc<Service>,
      11              :     interval: Duration,
      12              : }
      13              : 
      14              : impl ChaosInjector {
      15            0 :     pub fn new(service: Arc<Service>, interval: Duration) -> Self {
      16            0 :         Self { service, interval }
      17            0 :     }
      18              : 
      19            0 :     pub async fn run(&mut self, cancel: CancellationToken) {
      20            0 :         let mut interval = tokio::time::interval(self.interval);
      21              : 
      22              :         loop {
      23            0 :             tokio::select! {
      24            0 :                 _ = interval.tick() => {}
      25            0 :                 _ = cancel.cancelled() => {
      26            0 :                     tracing::info!("Shutting down");
      27            0 :                     return;
      28              :                 }
      29              :             }
      30              : 
      31            0 :             self.inject_chaos().await;
      32              : 
      33            0 :             tracing::info!("Chaos iteration...");
      34              :         }
      35            0 :     }
      36              : 
      37            0 :     async fn inject_chaos(&mut self) {
      38            0 :         // Pick some shards to interfere with
      39            0 :         let batch_size = 128;
      40            0 :         let mut inner = self.service.inner.write().unwrap();
      41            0 :         let (nodes, tenants, scheduler) = inner.parts_mut();
      42            0 :         let tenant_ids = tenants.keys().cloned().collect::<Vec<_>>();
      43            0 :         let victims = tenant_ids.choose_multiple(&mut thread_rng(), batch_size);
      44              : 
      45            0 :         for victim in victims {
      46            0 :             let shard = tenants
      47            0 :                 .get_mut(victim)
      48            0 :                 .expect("Held lock between choosing ID and this get");
      49              : 
      50              :             // Pick a secondary to promote
      51            0 :             let Some(new_location) = shard
      52            0 :                 .intent
      53            0 :                 .get_secondary()
      54            0 :                 .choose(&mut thread_rng())
      55            0 :                 .cloned()
      56              :             else {
      57            0 :                 tracing::info!("Skipping shard {victim}: no secondary location, can't migrate");
      58            0 :                 continue;
      59              :             };
      60              : 
      61            0 :             let Some(old_location) = *shard.intent.get_attached() else {
      62            0 :                 tracing::info!("Skipping shard {victim}: currently has no attached location");
      63            0 :                 continue;
      64              :             };
      65              : 
      66            0 :             shard.intent.demote_attached(scheduler, old_location);
      67            0 :             shard.intent.promote_attached(scheduler, new_location);
      68            0 :             self.service.maybe_reconcile_shard(shard, nodes);
      69              :         }
      70            0 :     }
      71              : }