nimbus/
targeting.rs
1use crate::{NimbusError, Result};
6
7use jexl_eval::Evaluator;
8use serde::Serialize;
9use serde_json::Value;
10
11cfg_if::cfg_if! {
12 if #[cfg(feature = "stateful")] {
13 use anyhow::anyhow;
14 use crate::{TargetingAttributes, stateful::behavior::{EventStore, EventQueryType, query_event_store}};
15 use std::sync::{Arc, Mutex};
16 use firefox_versioning::compare::version_compare;
17 }
18}
19
20#[derive(Clone)]
21pub struct NimbusTargetingHelper {
22 pub(crate) context: Value,
23 #[cfg(feature = "stateful")]
24 pub(crate) event_store: Arc<Mutex<EventStore>>,
25 #[cfg(feature = "stateful")]
26 pub(crate) targeting_attributes: Option<TargetingAttributes>,
27}
28
29impl NimbusTargetingHelper {
30 pub fn new<C: Serialize>(
31 context: C,
32 #[cfg(feature = "stateful")] event_store: Arc<Mutex<EventStore>>,
33 ) -> Self {
34 Self {
35 context: serde_json::to_value(context).unwrap(),
36 #[cfg(feature = "stateful")]
37 event_store,
38 #[cfg(feature = "stateful")]
39 targeting_attributes: None,
40 }
41 }
42
43 pub fn eval_jexl(&self, expr: String) -> Result<bool> {
44 cfg_if::cfg_if! {
45 if #[cfg(feature = "stateful")] {
46 jexl_eval(&expr, &self.context, self.event_store.clone())
47 } else {
48 jexl_eval(&expr, &self.context)
49 }
50 }
51 }
52
53 pub fn evaluate_jexl_raw_value(&self, expr: &str) -> Result<Value> {
54 cfg_if::cfg_if! {
55 if #[cfg(feature = "stateful")] {
56 jexl_eval_raw(expr, &self.context, self.event_store.clone())
57 } else {
58 jexl_eval_raw(expr, &self.context)
59 }
60 }
61 }
62
63 pub(crate) fn put(&self, key: &str, value: bool) -> Self {
64 let context = if let Value::Object(map) = &self.context {
65 let mut map = map.clone();
66 map.insert(key.to_string(), Value::Bool(value));
67 Value::Object(map)
68 } else {
69 self.context.clone()
70 };
71
72 Self {
73 context,
74 #[cfg(feature = "stateful")]
75 event_store: self.event_store.clone(),
76 #[cfg(feature = "stateful")]
77 targeting_attributes: self.targeting_attributes.clone(),
78 }
79 }
80}
81
82pub fn jexl_eval_raw<Context: serde::Serialize>(
83 expression_statement: &str,
84 context: &Context,
85 #[cfg(feature = "stateful")] event_store: Arc<Mutex<EventStore>>,
86) -> Result<Value> {
87 let evaluator = Evaluator::new();
88
89 #[cfg(feature = "stateful")]
90 let evaluator = evaluator
91 .with_transform("versionCompare", |args| Ok(version_compare(args)?))
92 .with_transform("eventSum", |args| {
93 Ok(query_event_store(
94 event_store.clone(),
95 EventQueryType::Sum,
96 args,
97 )?)
98 })
99 .with_transform("eventCountNonZero", |args| {
100 Ok(query_event_store(
101 event_store.clone(),
102 EventQueryType::CountNonZero,
103 args,
104 )?)
105 })
106 .with_transform("eventAveragePerInterval", |args| {
107 Ok(query_event_store(
108 event_store.clone(),
109 EventQueryType::AveragePerInterval,
110 args,
111 )?)
112 })
113 .with_transform("eventAveragePerNonZeroInterval", |args| {
114 Ok(query_event_store(
115 event_store.clone(),
116 EventQueryType::AveragePerNonZeroInterval,
117 args,
118 )?)
119 })
120 .with_transform("eventLastSeen", |args| {
121 Ok(query_event_store(
122 event_store.clone(),
123 EventQueryType::LastSeen,
124 args,
125 )?)
126 })
127 .with_transform("bucketSample", bucket_sample);
128
129 evaluator
130 .eval_in_context(expression_statement, context)
131 .map_err(|err| NimbusError::EvaluationError(err.to_string()))
132}
133
134pub fn jexl_eval<Context: serde::Serialize>(
139 expression_statement: &str,
140 context: &Context,
141 #[cfg(feature = "stateful")] event_store: Arc<Mutex<EventStore>>,
142) -> Result<bool> {
143 let res = jexl_eval_raw(
144 expression_statement,
145 context,
146 #[cfg(feature = "stateful")]
147 event_store,
148 )?;
149 match res.as_bool() {
150 Some(v) => Ok(v),
151 None => Err(NimbusError::InvalidExpression),
152 }
153}
154
155#[cfg(feature = "stateful")]
156fn bucket_sample(args: &[Value]) -> anyhow::Result<Value> {
157 fn get_arg_as_u32(args: &[Value], idx: usize, name: &str) -> anyhow::Result<u32> {
158 match args.get(idx) {
159 None => Err(anyhow!("{} doesn't exist in jexl transform", name)),
160 Some(Value::Number(n)) => {
161 let n: f64 = if let Some(n) = n.as_u64() {
162 n as f64
163 } else if let Some(n) = n.as_i64() {
164 n as f64
165 } else if let Some(n) = n.as_f64() {
166 n
167 } else {
168 unreachable!();
169 };
170
171 debug_assert!(n >= 0.0, "JEXL parser does not support negative values");
172 if n > u32::MAX as f64 {
173 Err(anyhow!("{} is out of range", name))
174 } else {
175 Ok(n as u32)
176 }
177 }
178 Some(_) => Err(anyhow!("{} is not a number", name)),
179 }
180 }
181
182 let input = args
183 .first()
184 .ok_or_else(|| anyhow!("input doesn't exist in jexl transform"))?;
185 let start = get_arg_as_u32(args, 1, "start")?;
186 let count = get_arg_as_u32(args, 2, "count")?;
187 let total = get_arg_as_u32(args, 3, "total")?;
188
189 let result = crate::sampling::bucket_sample(input, start, count, total)?;
190
191 Ok(Value::Bool(result))
192}