use crate::common::ValidationError;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct RolloutConfig {
#[serde(default)]
pub cohorts: Vec<RolloutCohort>,
}
impl RolloutConfig {
pub fn validate(&self) -> Result<(), ValidationError> {
let mut total = 0_u16;
for (index, cohort) in self.cohorts.iter().enumerate() {
if cohort.name.trim().is_empty() {
return Err(ValidationError::new(format!(
"cohorts[{index}].name must not be empty"
)));
}
if cohort.percentage == 0 {
return Err(ValidationError::new(format!(
"cohorts[{index}].percentage must be greater than zero"
)));
}
if self
.cohorts
.iter()
.take(index)
.any(|previous| previous.name == cohort.name)
{
return Err(ValidationError::new(format!(
"duplicate rollout cohort name: {}",
cohort.name
)));
}
total = total.saturating_add(u16::from(cohort.percentage));
if total > 100 {
return Err(ValidationError::new(
"rollout cohort percentages must total at most 100",
));
}
}
Ok(())
}
pub fn assign(&self, task_id: &str) -> RolloutAssignment {
let bucket = bucket_for(task_id);
let mut upper_bound = 0_u8;
let mut selected = None;
for cohort in &self.cohorts {
upper_bound = upper_bound.saturating_add(cohort.percentage);
if bucket < upper_bound {
selected = Some(cohort);
break;
}
}
RolloutAssignment {
task_id: task_id.to_owned(),
bucket,
cohort: selected.map(|cohort| cohort.name.clone()),
eligibility_override: selected.is_some_and(|cohort| cohort.eligibility_override),
kill_switch: selected.is_some_and(|cohort| cohort.kill_switch),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct RolloutCohort {
pub name: String,
pub percentage: u8,
pub eligibility_override: bool,
pub kill_switch: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct RolloutAssignment {
pub task_id: String,
pub bucket: u8,
pub cohort: Option<String>,
pub eligibility_override: bool,
pub kill_switch: bool,
}
impl RolloutAssignment {
pub fn is_unassigned(&self) -> bool {
self.cohort.is_none()
}
}
pub fn assign_rollout(task_id: &str, config: &RolloutConfig) -> RolloutAssignment {
config.assign(task_id)
}
fn bucket_for(task_id: &str) -> u8 {
let digest = Sha256::digest(task_id.as_bytes());
let value = u64::from_be_bytes(digest[..8].try_into().expect("SHA-256 has eight bytes"));
(value % 100) as u8
}
#[cfg(test)]
mod tests {
use super::*;
fn config() -> RolloutConfig {
RolloutConfig {
cohorts: vec![
RolloutCohort {
name: "internal-10".to_owned(),
percentage: 10,
eligibility_override: false,
kill_switch: false,
},
RolloutCohort {
name: "internal-50".to_owned(),
percentage: 50,
eligibility_override: true,
kill_switch: true,
},
],
}
}
#[test]
fn rollout_types_round_trip_and_validate() {
let config = config();
let assignment = config.assign("task-1");
let config_json = serde_json::to_string(&config).expect("config serializes");
let assignment_json = serde_json::to_string(&assignment).expect("assignment serializes");
assert_eq!(
serde_json::from_str::<RolloutConfig>(&config_json).unwrap(),
config
);
assert_eq!(
serde_json::from_str::<RolloutAssignment>(&assignment_json).unwrap(),
assignment
);
config.validate().expect("valid rollout");
}
#[test]
fn assignment_is_sha_based_deterministic_and_metadata_only() {
let config = config();
let first = config.assign("task-1");
let second = config.assign("task-1");
assert_eq!(first, second);
assert!(first.bucket < 100);
assert_eq!(assign_rollout("task-1", &config), first);
assert_eq!(first.task_id, "task-1");
}
#[test]
fn invalid_rollout_percentages_are_rejected() {
let invalid = RolloutConfig {
cohorts: vec![
RolloutCohort {
name: "a".to_owned(),
percentage: 60,
eligibility_override: false,
kill_switch: false,
},
RolloutCohort {
name: "b".to_owned(),
percentage: 60,
eligibility_override: false,
kill_switch: false,
},
],
};
assert!(invalid.validate().is_err());
}
}