nimbus/
targeting.rs

1// This Source Code Form is subject to the terms of the Mozilla Public
2// License, v. 2.0. If a copy of the MPL was not distributed with this
3// file, You can obtain one at https://mozilla.org/MPL/2.0/.
4
5use 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
134// This is the common entry point to JEXL evaluation.
135// The targeting attributes and additional context should have been merged and calculated before
136// getting here.
137// Any additional transforms should be added here.
138pub 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}