goose 0.19.1

//! Test plan structures and functions.
//!
//! Internally, Goose represents all load tests as a series of Test Plan steps.

use chrono::prelude::*;
use lazy_static::lazy_static;
use regex::Regex;
use serde::{Deserialize, Serialize};
use std::cmp::Ordering;
use std::str::FromStr;
use std::time;

use crate::config::{timespan, GooseConfiguration};
use crate::util;
use crate::{AttackPhase, GooseAttack, GooseAttackRunState, GooseError};

/// Internal data structure representing a test plan.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TestPlan {
    // A test plan is a vector of tuples each indicating a # of users and milliseconds.
    pub steps: Vec<(usize, usize)>,
    // Which step of the test_plan is currently running.
    pub current: usize,
}

impl Default for TestPlan {
    fn default() -> Self {
        Self::new()
    }
}

lazy_static! {
    // Each step of a test plan is "{users},{timespan}"; white space is ignored.
    // The time span is read by `util::parse_timespan`, as every other time
    // span is, so its units take either case and an empty span is 0.
    static ref STEP: Regex =
        Regex::new(r"^\s*([0-9]+)\s*,\s*([0-9hmsHMS]*)\s*$").unwrap();
}

/// Automatically represent all load tests internally as a test plan.
///
/// Load tests launched using `--users`, `--increase-time`, `--increase-rate`, and/or `--run-time` are
/// automatically converted to a `Vec<(usize, usize)>` test plan.
impl TestPlan {
    /// Create a new, empty TestPlan structure.
    pub fn new() -> TestPlan {
        TestPlan {
            steps: Vec::new(),
            current: 0,
        }
    }

    /// Build a test plan from current configuration.
    ///
    /// Fails when a configured time span is too long for a `usize` in
    /// milliseconds.
    pub(crate) fn build(configuration: &GooseConfiguration) -> Result<TestPlan, GooseError> {
        if let Some(test_plan) = configuration.test_plan.as_ref() {
            // Test plan was manually defined, clone and return as is.
            Ok(test_plan.clone())
        } else {
            let mut steps: Vec<(usize, usize)> = Vec::new();

            // The seconds in each time span, 0 when it is not set.
            let span = |option, value: &Option<String>| {
                value
                    .as_deref()
                    .map_or(Ok(0), |value| timespan(option, value))
            };
            let increase_time = span("increase_time", &configuration.increase_time)?;
            let run_time = span("run_time", &configuration.run_time)?;
            let decrease_time = span("decrease_time", &configuration.decrease_time)?;

            // Build a simple test plan from configured options if possible.
            if let Some(users) = configuration.users {
                if increase_time > 0 {
                    // Load test is configured with --increase-time.
                    steps.push((users, increase_time * 1_000));
                } else {
                    // Load test is configured with --increase-rate.
                    let increase_rate =
                        if let Some(increase_rate) = configuration.increase_rate.as_ref() {
                            util::parse_rate(Some(increase_rate.to_string()))
                        } else {
                            util::parse_rate(None)
                        };
                    // Convert increase_rate to milliseconds.
                    let ms_increase_rate = 1.0 / increase_rate * 1_000.0;
                    // Finally, multiply the increase rate by the number of users to start.
                    let total_time = ms_increase_rate * users as f32;
                    steps.push((users, total_time as usize));
                }

                // If a run-time is set, maintain users for the configured duration.
                if run_time > 0 {
                    steps.push((users, run_time * 1_000));
                }

                // Ramp down using --decrease-time or --decrease-rate if configured.
                // This applies whether or not --run-time is set (e.g. ramp up then
                // immediately ramp down).
                if decrease_time > 0 {
                    // Decrease users over the configured decrease-time.
                    steps.push((0, decrease_time * 1_000));
                } else if let Some(decrease_rate) = configuration.decrease_rate.as_ref() {
                    // Calculate decrease time from rate.
                    let rate = util::parse_rate(Some(decrease_rate.to_string()));
                    let ms_decrease_rate = 1.0 / rate * 1_000.0;
                    let total_time = ms_decrease_rate * users as f32;
                    steps.push((0, total_time as usize));
                } else if run_time > 0 {
                    // No decrease option set but run-time is configured: shut down
                    // as quickly as possible once run-time elapses.
                    steps.push((0, 0));
                }
            }

            // Define test plan from options.
            Ok(TestPlan { steps, current: 0 })
        }
    }

    /// Whether the run enters Maintain at some step: a step keeps the users of the
    /// one before it, or the last step has users.
    pub(crate) fn holds(&self) -> bool {
        self.steps.windows(2).any(|pair| pair[0].0 == pair[1].0)
            || self.steps.last().is_some_and(|step| step.0 > 0)
    }

    // Determine the total number of users required by the test plan.
    pub fn total_users(&self) -> usize {
        self.checked_total_users().unwrap_or(usize::MAX)
    }

    /// The total number of users the test plan starts, `None` if it does not
    /// fit in a `usize`.
    fn checked_total_users(&self) -> Option<usize> {
        let mut previous: usize = 0;
        self.steps
            .iter()
            .try_fold(0usize, |total_users, &(users, _)| {
                // Add to total_users every time there is an increase.
                let increase = users.saturating_sub(previous);
                previous = users;
                total_users.checked_add(increase)
            })
    }

    /// Parse a test plan as [`FromStr`] does, without printing anything: the
    /// Controllers parse client text with this.
    pub(crate) fn parse(s: &str) -> Result<TestPlan, GooseError> {
        // Convert string into a TestPlan.
        let mut steps: Vec<(usize, usize)> = Vec::new();
        // A test plan can have multiple lines split by the semicolon ";".
        for line in s.split(';') {
            // Users and milliseconds, unless either is too large for a usize.
            let step = STEP.captures(line).and_then(|cap| {
                Some((
                    cap[1].parse::<usize>().ok()?,
                    util::parse_timespan(&cap[2])? * 1_000,
                ))
            });
            match step {
                Some(step) => steps.push(step),
                None => {
                    return Err(GooseError::InvalidOption {
                        option: "`configuration.test_plan`".to_string(),
                        value: line.to_string(),
                        detail: INVALID_STEP.to_string(),
                    })
                }
            }
        }
        let test_plan = TestPlan { steps, current: 0 };
        // A plan with no users never starts: weighting zero users never ends.
        match test_plan.checked_total_users() {
            Some(total) if total > 0 => Ok(test_plan),
            _ => Err(GooseError::InvalidOption {
                option: "`configuration.test_plan`".to_string(),
                value: s.to_string(),
                detail: "`configuration.test_plan` must start at least one user, and no more than fit in a usize.".to_string(),
            }),
        }
    }
}

/// Implement [`FromStr`] to convert `"users,timespan"` string formatted test plans to Goose's
/// internal representation of Vec<(usize, usize)>.
///
/// Users are represented simply as an integer.
///
/// Time span can be specified as an integer, indicating seconds. Or can use integers together
/// with one or more of "h", "m", and "s", in that order, indicating "hours", "minutes", and
/// "seconds". Valid formats include: 20, 20s, 3m, 2h, 1h20m, 3h30m10s, etc.
impl FromStr for TestPlan {
    type Err = GooseError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        // Logger isn't initialized yet, print why the plan is refused.
        TestPlan::parse(s).inspect_err(|e| eprint!("{}", refusal(e)))
    }
}

/// Write a test plan as `--test-plan` takes it, each step's duration in whole
/// seconds: the canonical form, so `10,1h30m` is written `10,5400`. It parses
/// back to the same steps with [`FromStr`]. `current` is not written.
impl std::fmt::Display for TestPlan {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        for (index, (users, ms)) in self.steps.iter().enumerate() {
            if index > 0 {
                f.write_str(";")?;
            }
            write!(f, "{},{}", users, ms / 1_000)?;
        }
        Ok(())
    }
}

/// What the command line prints for a refused test plan: the format for a
/// step that does not parse, otherwise the reason it was refused.
fn refusal(e: &GooseError) -> String {
    match e {
        GooseError::InvalidOption { value, detail, .. } if detail == INVALID_STEP => format!(
            "ERROR: invalid `configuration.test_plan` value: '{value}'\n  \
             Expected format: --test-plan \"{{users}},{{timespan}};{{users}},{{timespan}}\"\n    \
             {{users}} must be an integer, ie \"100\"\n    \
             {{timespan}} can be integer seconds or \"30s\", \"20m\", \"3h\", \"1h30m\", etc\n"
        ),
        GooseError::InvalidOption { value, detail, .. } => {
            format!("ERROR: invalid `configuration.test_plan` value: '{value}': {detail}\n")
        }
        e => format!("ERROR: {e}\n"),
    }
}

/// The detail of a test plan step that does not parse.
const INVALID_STEP: &str = "invalid `configuration.test_plan` value.";

/// A test plan is a series of steps performing one of the following actions.
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
pub enum TestPlanStepAction {
    /// A test plan step that is increasing the number of GooseUser threads.
    Increasing,
    /// A test plan step that is maintaining the number of GooseUser threads.
    Maintaining,
    /// A test plan step that is decreasing the number of GooseUser threads.
    Decreasing,
    /// A test plan step that is canceling all GooseUser threads.
    Canceling,
    /// The final step indicating that the load test is finished.
    Finished,
}

/// A historical record of a single test plan step, used to generate reports from the metrics.
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
pub struct TestPlanHistory {
    /// What action happend in this step.
    pub action: TestPlanStepAction,
    /// A timestamp of when the step started.
    pub timestamp: DateTime<Utc>,
    /// The number of users when the step started.
    pub users: usize,
}
impl TestPlanHistory {
    /// A helper to record a new test plan step in the historical record.
    pub(crate) fn step(action: TestPlanStepAction, users: usize) -> TestPlanHistory {
        TestPlanHistory {
            action,
            timestamp: Utc::now(),
            users,
        }
    }
}

impl GooseAttack {
    // Advance the active [`GooseAttack`](./struct.GooseAttack.html) to the next TestPlan step.
    pub(crate) fn advance_test_plan(&mut self, goose_attack_run_state: &mut GooseAttackRunState) {
        // Record the instant this new step starts, for use with timers.
        let now = time::Instant::now();
        self.step_started = Some(now);

        // Each step schedules its user changes from its own start, so a ramp
        // left mid way, as by a Controller, never leaves a stale due time.
        goose_attack_run_state.adjust_user_in_ms = 0;

        let action = if self.test_plan.current == self.test_plan.steps.len() - 1 {
            // If this is the last TestPlan step and there are 0 users, stop the run
            // as a ramp down to 0 does: Decrease finds 0 users on the last step, so
            // it runs test_stop, keeps the metrics, judges the thresholds, writes the
            // reports and records the run as finished.
            if self.test_plan.steps[self.test_plan.current].0 == 0 {
                self.set_attack_phase(goose_attack_run_state, AttackPhase::Decrease);
                None
            }
            // Otherwise maintain the number of GooseUser threads until canceled.
            else {
                self.set_attack_phase(goose_attack_run_state, AttackPhase::Maintain);
                Some(TestPlanStepAction::Maintaining)
            }
        // If this is not the last TestPlan step, determine what happens next.
        } else if self.test_plan.current < self.test_plan.steps.len() {
            match self.test_plan.steps[self.test_plan.current]
                .0
                .cmp(&self.test_plan.steps[self.test_plan.current + 1].0)
            {
                Ordering::Less => {
                    self.set_attack_phase(goose_attack_run_state, AttackPhase::Increase);
                    Some(TestPlanStepAction::Increasing)
                }
                Ordering::Greater => {
                    self.set_attack_phase(goose_attack_run_state, AttackPhase::Decrease);
                    Some(TestPlanStepAction::Decreasing)
                }
                Ordering::Equal => {
                    self.set_attack_phase(goose_attack_run_state, AttackPhase::Maintain);
                    Some(TestPlanStepAction::Maintaining)
                }
            }
        } else {
            unreachable!("Advanced 2 steps beyond the end of the TestPlan.")
        };

        // Record details about new new TestPlan step that is starting.
        if let Some(action) = action {
            self.metrics.history.push(TestPlanHistory::step(
                action,
                self.test_plan.steps[self.test_plan.current].0,
            ));
        }

        // Always advance the TestPlan step
        self.test_plan.current += 1;

        // A users ramp starts its hold, or its ramp down, as it enters it.
        if let Some(shape) = self.run_shape.as_mut() {
            shape.entered(self.test_plan.current, &mut self.test_plan.steps, now);
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;

    /// A test plan is written in canonical form, and parses back to the
    /// same steps.
    #[test]
    fn display_round_trips() {
        for (plan, written) in [
            ("10,5s;10,30s;0,5s", "10,5;10,30;0,5"),
            ("10,1h30m;0,0", "10,5400;0,0"),
            ("100000,1s", "100000,1"),
            (" 3 , 2m10s ;0,1h2m3s", "3,130;0,3723"),
            ("1,720h", "1,2592000"),
        ] {
            let parsed = TestPlan::parse(plan).unwrap();
            assert_eq!(parsed.to_string(), written, "{}", plan);
            let again = TestPlan::parse(&parsed.to_string()).unwrap();
            assert_eq!(again.steps, parsed.steps, "{}", plan);
            assert_eq!(again.to_string(), written, "{}", plan);
        }
    }

    /// A step's time units parse in either case, as every other time span
    /// does, from the command line and from `--test-plan` itself.
    #[test]
    fn step_units_in_either_case() {
        for (upper, lower) in [
            ("2,1M", "2,1m"),
            ("10,1H30M", "10,1h30m"),
            ("5,10S", "5,10s"),
            ("10,1h30M;0,5S", "10,1h30m;0,5s"),
        ] {
            let parsed = TestPlan::from_str(upper).unwrap();
            assert_eq!(
                parsed.steps,
                TestPlan::from_str(lower).unwrap().steps,
                "{upper}"
            );
            let configured =
                GooseConfiguration::parse_args_default(&["--test-plan", upper]).unwrap();
            assert_eq!(configured.test_plan.unwrap().steps, parsed.steps, "{upper}");
        }
    }

    /// Plans that parsed before units took either case parse the same, and
    /// steps that were refused still are.
    #[test]
    fn step_grammar_unchanged() {
        for (plan, steps) in [
            ("10,", vec![(10, 0)]),
            ("10,90", vec![(10, 90_000)]),
            ("10,1h30m", vec![(10, 5_400_000)]),
            ("0,5s;10,20s", vec![(0, 5_000), (10, 20_000)]),
            (" 3 , 2m10s ;0,1h2m3s", vec![(3, 130_000), (0, 3_723_000)]),
        ] {
            assert_eq!(TestPlan::parse(plan).unwrap().steps, steps, "{plan}");
        }
        for plan in [
            "2,15mins",
            "2,5 junk",
            "2,1m1h",
            "2,h",
            "2,1d",
            "2,ms",
            "2,s",
            "2,1h1h",
            "2,1h 30m",
            "2,5\u{17f}",
            "2,\u{661}h",
            "x",
            "2",
        ] {
            assert!(TestPlan::parse(plan).is_err(), "{}", plan);
        }
    }

    /// The command line says why it refuses a plan, not only that it did.
    #[test]
    fn refusal_names_reason() {
        let reason = |plan: &str| refusal(&TestPlan::parse(plan).unwrap_err());
        let no_users = "`configuration.test_plan` must start at least one user, and no more than fit in a usize.";
        assert_eq!(
            reason("0,1s"),
            format!("ERROR: invalid `configuration.test_plan` value: '0,1s': {no_users}\n")
        );
        let overflow = "18446744073709551615,1s;0,1s;1,1s";
        assert_eq!(
            reason(overflow),
            format!("ERROR: invalid `configuration.test_plan` value: '{overflow}': {no_users}\n")
        );
        // A step that does not parse still prints the format.
        assert!(reason("10,1s;x").starts_with(
            "ERROR: invalid `configuration.test_plan` value: 'x'\n  Expected format: --test-plan \"{users},{timespan};{users},{timespan}\"\n"
        ));
    }
}