Skip to content

Commit

Permalink
Introduce AvailableRandomizationUnits
Browse files Browse the repository at this point in the history
  • Loading branch information
@eoger
eoger committed Oct 9, 2020
1 parent 80edb8c commit 13184e3
Show file tree
Hide file tree
Showing 4 changed files with 92 additions and 25 deletions.
24 changes: 18 additions & 6 deletions nimbus/examples/experiment.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@

use clap::{App, Arg, SubCommand};
use env_logger::Env;
use nimbus::{AppContext, Config, NimbusClient};
use nimbus::{AppContext, AvailableRandomizationUnits, Config, NimbusClient};
use std::io::prelude::*;

const DEFAULT_BASE_URL: &str = "https://settings.stage.mozaws.net"; // TODO: Replace this with prod
Expand Down Expand Up @@ -80,9 +80,17 @@ fn main() {
bucket_name: Some(bucket_name.to_string()),
};

let available_randomization_units = AvailableRandomizationUnits { client_id: None };

// Here we initialize our main `NimbusClient` struct
let nimbus_client =
NimbusClient::new(collection_name.to_string(), context, "", Some(config)).unwrap();
let nimbus_client = NimbusClient::new(
collection_name.to_string(),
context,
"",
Some(config),
available_randomization_units,
)
.unwrap();

// We match against the subcommands
match matches.subcommand() {
Expand Down Expand Up @@ -114,11 +122,15 @@ fn main() {
let all_experiments = nimbus_client.get_all_experiments();
let mut num_of_experiments_enrolled = 0;
let mut uuid = uuid::Uuid::new_v4();
let available_randomization_units = AvailableRandomizationUnits {
client_id: Some("bobo".to_string()),
};
while num_of_experiments_enrolled != num {
uuid = uuid::Uuid::new_v4();
num_of_experiments_enrolled = nimbus::filter_enrolled(&uuid, &all_experiments)
.unwrap()
.len()
num_of_experiments_enrolled =
nimbus::filter_enrolled(&uuid, &available_randomization_units, &all_experiments)
.unwrap()
.len()
}
println!("======================================");
println!("Generated UUID is: {}", uuid);
Expand Down
56 changes: 41 additions & 15 deletions nimbus/src/evaluator.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,10 @@
//! TODO: Implement the bucketing logic from the nimbus project
use crate::error::{Error, Result};
use crate::{
error::{Error, Result},
AvailableRandomizationUnits,
};
use crate::{matcher::AppContext, sampling};
use crate::{Branch, EnrolledExperiment, Experiment};
use jexl_eval::Evaluator;
Expand All @@ -29,7 +32,8 @@ impl Bucket {
///
/// # Arguments:
///
/// - `id` The user's id
/// - `nimbus_id` The auto-generated nimbus_id
/// - `available_randomization_units`: The app provded available randomization units
/// - `experiments` A list of experiments, usually retrieved from the network or persisted storage
///
/// # Returns:
Expand All @@ -45,12 +49,29 @@ impl Bucket {
/// - If the bucket sampling failed (i.e we could not find if the user should or should not be enrolled in the experiment based on the bucketing)
/// - If an error occurs while determining the branch the user should be enrolled in any of the experiments
#[allow(dead_code)]
pub fn filter_enrolled(id: &Uuid, experiments: &[Experiment]) -> Result<Vec<EnrolledExperiment>> {
pub fn filter_enrolled(
nimbus_id: &Uuid,
available_randomization_units: &AvailableRandomizationUnits,
experiments: &[Experiment],
) -> Result<Vec<EnrolledExperiment>> {
let nimbus_id = nimbus_id.to_string();
let mut res = Vec::with_capacity(experiments.len());
for exp in experiments {
let bucket_config = exp.bucket_config.clone();
let id = match available_randomization_units
.get_value(&nimbus_id, &bucket_config.randomization_unit)
{
Some(id) => id,
None => {
log::info!(
"Could not find a suitable randomization unit for {}. Skipping experiment.",
&exp.slug
);
continue;
}
};
if sampling::bucket_sample(
vec![id.to_string(), bucket_config.namespace],
vec![id.to_owned(), bucket_config.namespace],
bucket_config.start,
bucket_config.count,
bucket_config.total,
Expand All @@ -59,7 +80,7 @@ pub fn filter_enrolled(id: &Uuid, experiments: &[Experiment]) -> Result<Vec<Enro
slug: exp.slug.clone(),
user_facing_name: exp.user_facing_name.clone(),
user_facing_description: exp.user_facing_description.clone(),
branch_slug: choose_branch(&exp.slug, &exp.branches, id)?.clone().slug,
branch_slug: choose_branch(&exp.slug, &exp.branches, &id)?.clone().slug,
});
}
}
Expand Down Expand Up @@ -90,7 +111,7 @@ pub fn filter_enrolled(id: &Uuid, experiments: &[Experiment]) -> Result<Vec<Enro
pub(crate) fn choose_branch<'a>(
slug: &str,
branches: &'a [Branch],
id: &Uuid,
id: &str,
) -> Result<&'a Branch> {
let ratios = branches.iter().map(|b| b.ratio).collect::<Vec<_>>();
// Note: The "experiment-manager" here comes from https://searchfox.org/mozilla-central/source/toolkit/components/messaging-system/experiments/ExperimentManager.jsm#421
Expand Down Expand Up @@ -239,11 +260,11 @@ mod tests {
];
// 299eed1e-be6d-457d-9e53-da7b1a03f10d maps to the second index
let id = uuid::Uuid::parse_str("299eed1e-be6d-457d-9e53-da7b1a03f10d").unwrap();
let b = choose_branch(slug, &branches, &id).unwrap();
let b = choose_branch(slug, &branches, &id.to_string()).unwrap();
assert_eq!(b.slug, "blue");
// 542213c0-9aef-47eb-bc6b-3b8529736ba2 maps to the first index
let id = uuid::Uuid::parse_str("542213c0-9aef-47eb-bc6b-3b8529736ba2").unwrap();
let b = choose_branch(slug, &branches, &id).unwrap();
let b = choose_branch(slug, &branches, &id.to_string()).unwrap();
assert_eq!(b.slug, "control");
}

Expand All @@ -268,26 +289,31 @@ mod tests {
};
let mut experiment2 = experiment1.clone();
experiment2.bucket_config = BucketConfig {
randomization_unit: RandomizationUnit::NimbusId,
randomization_unit: RandomizationUnit::ClientId,
namespace:
"bug-1637316-message-aboutwelcome-pull-factor-reinforcement-76-rel-release-76-77"
.to_string(),
start: 2000,
count: 3000,
start: 9000,
count: 1000,
total: 10000,
};
experiment2.slug = "TEST_EXP2".to_string();
let experiments = vec![experiment1, experiment2];
let available_randomization_units = AvailableRandomizationUnits {
client_id: None, // We will not match EXP_2 because we don't have the necessary randomization unit.
};
// 299eed1e-be6d-457d-9e53-da7b1a03f10d uuid fits in start: 0, count: 2000, total: 10000 with the example namespace, to the treatment-variation-b branch
// Tested against the desktop implementation
let id = uuid::Uuid::parse_str("299eed1e-be6d-457d-9e53-da7b1a03f10d").unwrap();
let enrolled = filter_enrolled(&id, &experiments).unwrap();
let enrolled = filter_enrolled(&id, &available_randomization_units, &experiments).unwrap();
assert_eq!(enrolled.len(), 1);
assert_eq!(enrolled[0].slug, "TEST_EXP1");
// 542213c0-9aef-47eb-bc6b-3b8529736ba2 uuid fits in start: 2000, count: 3000, total: 10000 with the example namespace, to the control branch
// Tested against the desktop implementation
// Fits because of the client_id.
let available_randomization_units = AvailableRandomizationUnits {
client_id: Some("bobo".to_string()),
};
let id = uuid::Uuid::parse_str("542213c0-9aef-47eb-bc6b-3b8529736ba2").unwrap();
let enrolled = filter_enrolled(&id, &experiments).unwrap();
let enrolled = filter_enrolled(&id, &available_randomization_units, &experiments).unwrap();
assert_eq!(enrolled.len(), 1);
assert_eq!(enrolled[0].slug, "TEST_EXP2");
}
Expand Down
25 changes: 22 additions & 3 deletions nimbus/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,12 +49,14 @@ impl NimbusClient {
app_context: AppContext,
db_path: P,
config: Option<Config>,
available_randomization_units: AvailableRandomizationUnits,
) -> Result<Self> {
let client = Client::new(&collection_name, config.clone())?;
let resp = client.get_experiments()?;
let db = Database::new(db_path)?;
let nimbus_id = Self::get_or_create_nimbus_id(&db)?;
let enrolled_experiments = evaluator::filter_enrolled(&nimbus_id, &resp)?;
let enrolled_experiments =
evaluator::filter_enrolled(&nimbus_id, &available_randomization_units, &resp)?;
Ok(Self {
experiments: resp,
enrolled_experiments,
Expand Down Expand Up @@ -95,8 +97,8 @@ impl NimbusClient {
unimplemented!()
}

pub fn nimbus_id(&self) -> Result<Uuid> {
Self::get_or_create_nimbus_id(&self.db)
pub fn nimbus_id(&self) -> Uuid {
self.nimbus_id
}

fn get_or_create_nimbus_id(db: &Database) -> Result<Uuid> {
Expand Down Expand Up @@ -177,4 +179,21 @@ impl Default for RandomizationUnit {
}
}

pub struct AvailableRandomizationUnits {
pub client_id: Option<String>,
}

impl AvailableRandomizationUnits {
pub fn get_value<'a>(
&'a self,
nimbus_id: &'a str,
wanted: &'a RandomizationUnit,
) -> Option<&'a str> {
match wanted {
RandomizationUnit::NimbusId => Some(nimbus_id),
RandomizationUnit::ClientId => self.client_id.as_deref(),
}
}
}

include!(concat!(env!("OUT_DIR"), "/nimbus.uniffi.rs"));
12 changes: 11 additions & 1 deletion nimbus/src/nimbus.idl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,10 @@ dictionary Config {
string? bucket_name;
};

dictionary AvailableRandomizationUnits {
string? client_id;
};

[Error]
enum Error {
"InvalidPersistedData", "RkvError", "IOError",
Expand All @@ -35,7 +39,13 @@ enum Error {

interface NimbusClient {
[Throws=Error]
constructor(string collection_name, AppContext app_ctx, string dbpath, Config? config);
constructor(
string collection_name,
AppContext app_ctx,
string dbpath,
Config? config,
AvailableRandomizationUnits available_randomization_units
);


string? get_experiment_branch(string experiment_slug);
Expand Down

0 comments on commit 13184e3

Please sign in to comment.