lightweight overrides for unit-tests

Summary:
Here I'm adding one more implementation of JustKnobs - one that simply stores stuff in thread-local memory and it easy and cheap to use in unit tests.

In tunables in_memory overrides were part of the main implementation but I think it's easier to have them fully separate.

Reviewed By: gustavoavena

Differential Revision: D50497319

fbshipit-source-id: 74463cd79a151af8d3bd4cfc422c5de323ef31b1

shed/justknobs_stub/Cargo.toml

@@ -14,11 +14,13 @@ license = "MIT OR Apache-2.0"
 anyhow = "=1.0.72"
 arc-swap = "1.5"
 cached_config = { version = "0.1.0", path = "../cached_config" }
+futures = { version = "0.3.28", features = ["async-await", "compat"] }
 just_knobs_struct = { version = "0.1.0", path = "cached_config_thrift_struct" }
 serde = { version = "1.0.185", features = ["derive", "rc"] }
 serde_json = { version = "1.0.100", features = ["float_roundtrip", "unbounded_depth"] }
 slog = { version = "2.7", features = ["max_level_trace", "nested-values"] }
 tokio = { version = "1.29.1", features = ["full", "test-util", "tracing"] }
 
 [dev-dependencies]
+maplit = "1.0"
 slog_glog_fmt = { version = "0.1.0", path = "../slog_glog_fmt" }

shed/justknobs_stub/src/lib.rs

@@ -15,6 +15,8 @@
 //!    Can be used for integration tests where the config can be read from on-disk JSON file and
 //!    fully isolated from prod setup.  Used after being initialized with
 //!    init_cached_config_just_knobs/init_cached_config_just_knobs_worker.
+//!  * thread-local-in-memory, which is useful for testing. It allows to override justknobs within a
+//!    test without affecting other tests. Used always when cfg(test) is true.
 use anyhow as _;
 use anyhow::Result;
 use cached_config::CachedConfigJustKnobs;
@@ -24,8 +26,16 @@ use fb_justknobs as prod_implementation;
 use JustKnobsStub as prod_implementation;
 
 pub mod cached_config;
+pub mod thread_local_in_memory;
 pub use cached_config::init_just_knobs as init_cached_config_just_knobs;
 pub use cached_config::init_just_knobs_worker as init_cached_config_just_knobs_worker;
+#[cfg(test)]
+pub use thread_local_in_memory::with_just_knobs;
+#[cfg(test)]
+pub use thread_local_in_memory::with_just_knobs_async;
+#[cfg(test)]
+pub use thread_local_in_memory::with_just_knobs_async_arc;
+use thread_local_in_memory::ThreadLocalInMemoryJustKnobsImpl;
 
 /// Trait that defines the interface for JustKnobs supported by this library and multiple stub
 /// implementations is contains.
@@ -59,15 +69,19 @@ pub struct JustKnobsCombinedImpl;
 
 impl JustKnobs for JustKnobsCombinedImpl {
     fn eval(name: &str, hash_val: Option<&str>, switch_val: Option<&str>) -> Result<bool> {
-        if cached_config::in_use() {
+        if cfg!(test) {
+            ThreadLocalInMemoryJustKnobsImpl::eval(name, hash_val, switch_val)
+        } else if cached_config::in_use() {
             CachedConfigJustKnobs::eval(name, hash_val, switch_val)
         } else {
             prod_implementation::eval(name, hash_val, switch_val)
         }
     }
 
     fn get(name: &str, switch_val: Option<&str>) -> Result<i64> {
-        if cached_config::in_use() {
+        if cfg!(test) {
+            ThreadLocalInMemoryJustKnobsImpl::get(name, switch_val)
+        } else if cached_config::in_use() {
             CachedConfigJustKnobs::get(name, switch_val)
         } else {
             prod_implementation::get(name, switch_val)

shed/justknobs_stub/src/thread_local_in_memory.rs

@@ -0,0 +1,118 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under both the MIT license found in the
+ * LICENSE-MIT file in the root directory of this source tree and the Apache
+ * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
+ * of this source tree.
+ */
+
+use std::cell::RefCell;
+/// JustKnobs implementation that thread-local memory for storage. Meant to be used in unit tests.
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::thread_local;
+
+use anyhow::anyhow;
+use anyhow::Result;
+use futures::future::poll_fn;
+use futures::Future;
+use futures::FutureExt;
+
+use crate::JustKnobs;
+
+thread_local! {
+    static JUST_KNOBS: RefCell<Arc<JustKnobsInMemory>> = Default::default()
+}
+
+#[derive(Default)]
+pub struct JustKnobsInMemory(HashMap<String, KnobVal>);
+#[derive(Copy, Clone)]
+pub enum KnobVal {
+    Bool(bool),
+    Int(i64),
+}
+
+pub(crate) struct ThreadLocalInMemoryJustKnobsImpl;
+impl JustKnobs for ThreadLocalInMemoryJustKnobsImpl {
+    fn eval(name: &str, _hash_val: Option<&str>, _switch_val: Option<&str>) -> Result<bool> {
+        let value = JUST_KNOBS.with(|jk| *jk.borrow().0.get(name).unwrap_or(&KnobVal::Bool(false)));
+
+        match value {
+            KnobVal::Int(_v) => Err(anyhow!(
+                "JustKnobs knob {} has type int while expected bool",
+                name,
+            )),
+            KnobVal::Bool(b) => Ok(b),
+        }
+    }
+
+    fn get(name: &str, _switch_val: Option<&str>) -> Result<i64> {
+        let value = JUST_KNOBS.with(|jk| *jk.borrow().0.get(name).unwrap_or(&KnobVal::Int(0)));
+
+        match value {
+            KnobVal::Bool(_b) => Err(anyhow!(
+                "JustKnobs knob {} has type bool while expected int",
+                name,
+            )),
+            KnobVal::Int(v) => Ok(v),
+        }
+    }
+}
+
+/// A helper function to override jk during a closure's execution.
+/// This is useful for unit tests.
+pub fn with_just_knobs<T>(new_just_knobs: JustKnobsInMemory, f: impl FnOnce() -> T) -> T {
+    JUST_KNOBS.with(move |jk| *jk.borrow_mut() = Arc::new(new_just_knobs));
+    let res = f();
+    JUST_KNOBS.with(|jk| jk.take());
+    res
+}
+
+/// A helper function to override jk during a async closure's execution.  This is
+/// useful for unit tests.
+pub fn with_just_knobs_async<Out, Fut: Future<Output = Out> + Unpin>(
+    new_just_knobs: JustKnobsInMemory,
+    fut: Fut,
+) -> impl Future<Output = Out> {
+    with_just_knobs_async_arc(Arc::new(new_just_knobs), fut)
+}
+
+pub fn with_just_knobs_async_arc<Out, Fut: Future<Output = Out> + Unpin>(
+    new_just_knobs: Arc<JustKnobsInMemory>,
+    mut fut: Fut,
+) -> impl Future<Output = Out> {
+    poll_fn(move |cx| {
+        JUST_KNOBS.with(|jk| *jk.borrow_mut() = new_just_knobs.clone());
+        let res = fut.poll_unpin(cx);
+        JUST_KNOBS.with(|jk| jk.take());
+        res
+    })
+}
+
+#[cfg(test)]
+mod test {
+    use maplit::hashmap;
+
+    use super::*;
+
+    #[test]
+    fn test_jk_override() -> Result<()> {
+        assert!(!ThreadLocalInMemoryJustKnobsImpl::eval("my/config:knob1", None, None).unwrap());
+
+        let res = with_just_knobs(
+            JustKnobsInMemory(hashmap! {
+                "my/config:knob1".to_string() => KnobVal::Bool(true),
+                "my/config:knob2".to_string() => KnobVal::Int(2),
+            }),
+            || {
+                (
+                    ThreadLocalInMemoryJustKnobsImpl::eval("my/config:knob1", None, None).unwrap(),
+                    ThreadLocalInMemoryJustKnobsImpl::get("my/config:knob2", None).unwrap(),
+                )
+            },
+        );
+        assert_eq!(res, (true, 2));
+        Ok(())
+    }
+}