Adding usage of must_obey_client in wrapper to improve performace

.github/workflows/ci.yml

@@ -24,7 +24,12 @@ jobs:
         cargo fmt --check
         cargo clippy --profile release --all-targets -- -D clippy::all
     - name: Release Build
-      run: RUSTFLAGS="-D warnings" cargo build --all --all-targets  --release
+      run: |
+        if [ "${{ matrix.server_version }}" = "8.0.0" ]; then
+          RUSTFLAGS="-D warnings" cargo build --all --all-targets --release --features valkey_8_0
+        else
+          RUSTFLAGS="-D warnings" cargo build --all --all-targets --release
+        fi
     - name: Run unit tests
       run: cargo test --features enable-system-alloc
     - name: Make valkey-server binary
@@ -77,7 +82,12 @@ jobs:
         cargo fmt --check
         cargo clippy --profile release --all-targets -- -D clippy::all
     - name: Release Build
-      run: RUSTFLAGS="-D warnings" cargo build --all --all-targets  --release
+      run: |
+        if [ "${{ matrix.server_version }}" = "8.0.0" ]; then
+          RUSTFLAGS="-D warnings" cargo build --all --all-targets --release --features valkey_8_0
+        else
+          RUSTFLAGS="-D warnings" cargo build --all --all-targets --release
+        fi
     - name: Run unit tests
       run: cargo test --features enable-system-alloc
     - name: Make Valkey-server binary with asan

Cargo.toml

@@ -38,4 +38,5 @@ debug-assertions = true
 default = ["min-valkey-compatibility-version-8-0"]
 enable-system-alloc = ["valkey-module/enable-system-alloc"]
 min-valkey-compatibility-version-8-0 = []
+valkey_8_0 = []  # Valkey-bloom is intended to be loaded on server versions >= Valkey 8.1 and by default it is built this way (unless this flag is provided). It is however compatible with Valkey version 8.0 if the user explicitly provides this feature flag in their cargo build command.
 use-redismodule-api = [] # We don't support this feature flag which is why it is empty. 

README.md

@@ -2,7 +2,7 @@
 
 Valkey-Bloom (BSD-3-Clause) is a Rust Valkey-Module which brings a native and space efficient probabilistic Module data type to Valkey. With this, users can create filters (space-efficient probabilistic Module data type) to add elements, perform “check” operation to test whether an element exists, auto scale their filters, perform RDB Save and load operations, etc.
 
-Valkey-Bloom is built using bloomfilter::Bloom (https://crates.io/crates/bloomfilter which has a BSD-2-Clause license).
+Valkey-Bloom is built using `bloomfilter::Bloom` (https://crates.io/crates/bloomfilter which has a BSD-2-Clause license).
 
 It is compatible with the BloomFilter (BF.*) command APIs in Redis offerings.
 
@@ -63,4 +63,12 @@ valkey-server --loadmodule /path/to/libvalkey_bloom.so
 1. Connect to a running Valkey instance using valkey-cli
 2. Execute Valkey command:
     MODULE LOAD /path/to/libvalkey_bloom.so
-```
+```
+## Feature Flags
+
+* valkey_8_0: valkey-bloom is intended to be loaded on server versions >= Valkey 8.1 and by default it is built this way (unless this flag is provided). It is however compatible with Valkey version 8.0 if the user explicitly provides this feature flag in their cargo build command.
+```
+cargo build --release --features valkey_8_0
+```
+
+This can also be done by specifiyng SERVER_VERSION=8.0.0 and then running `./build.sh`

build.sh

@@ -12,8 +12,6 @@ echo "Running cargo and clippy format checks..."
 cargo fmt --check
 cargo clippy --profile release --all-targets -- -D clippy::all
 
-echo "Running cargo build release..."
-RUSTFLAGS="-D warnings" cargo build --all --all-targets  --release
 
 echo "Running unit tests..."
 cargo test --features enable-system-alloc
@@ -29,6 +27,14 @@ if [ "$SERVER_VERSION" != "unstable" ] && [ "$SERVER_VERSION" != "8.0.0" ] ; the
   exit 1
 fi
 
+echo "Running cargo build release..."
+if [ "$SERVER_VERSION" == "8.0.0" ] ; then
+    RUSTFLAGS="-D warnings" cargo build --all --all-targets  --release --features valkey_8_0
+else
+    RUSTFLAGS="-D warnings" cargo build --all --all-targets  --release
+fi
+
+
 REPO_URL="https://github.com/valkey-io/valkey.git"
 BINARY_PATH="tests/.build/binaries/$SERVER_VERSION/valkey-server"
 

src/bloom/command_handler.rs

@@ -6,11 +6,10 @@ use crate::configs::{
     BLOOM_CAPACITY_MAX, BLOOM_CAPACITY_MIN, BLOOM_EXPANSION_MAX, BLOOM_EXPANSION_MIN,
     BLOOM_FP_RATE_MAX, BLOOM_FP_RATE_MIN, BLOOM_TIGHTENING_RATIO_MAX, BLOOM_TIGHTENING_RATIO_MIN,
 };
+use crate::wrapper::must_obey_client;
 use std::sync::atomic::Ordering;
-use valkey_module::ContextFlags;
 use valkey_module::NotifyEvent;
 use valkey_module::{Context, ValkeyError, ValkeyResult, ValkeyString, ValkeyValue, VALKEY_OK};
-
 /// Helper function used to add items to a bloom object. It handles both multi item and single item add operations.
 /// It is used by any command that allows adding of items: BF.ADD, BF.MADD, and BF.INSERT.
 /// Returns the result of the item add operation on success as a ValkeyValue and a ValkeyError on failure.
@@ -177,7 +176,7 @@ pub fn bloom_filter_add_value(
         }
     };
     // Skip bloom filter size validation on replicated cmds.
-    let validate_size_limit = !ctx.get_flags().contains(ContextFlags::REPLICATED);
+    let validate_size_limit = !must_obey_client(ctx);
     let mut add_succeeded = false;
     match value {
         Some(bloom) => {
@@ -404,7 +403,7 @@ pub fn bloom_filter_reserve(ctx: &Context, input_args: &[ValkeyString]) -> Valke
                 false => (Some(configs::FIXED_SEED), false),
             };
             // Skip bloom filter size validation on replicated cmds.
-            let validate_size_limit = !ctx.get_flags().contains(ContextFlags::REPLICATED);
+            let validate_size_limit = !must_obey_client(ctx);
             let tightening_ratio = *configs::BLOOM_TIGHTENING_F64
                 .lock()
                 .expect("Unable to get a lock on tightening ratio static");
@@ -615,7 +614,7 @@ pub fn bloom_filter_insert(ctx: &Context, input_args: &[ValkeyString]) -> Valkey
         }
     };
     // Skip bloom filter size validation on replicated cmds.
-    let validate_size_limit = !ctx.get_flags().contains(ContextFlags::REPLICATED);
+    let validate_size_limit = !must_obey_client(ctx);
     let mut add_succeeded = false;
     match value {
         Some(bloom) => {
@@ -811,7 +810,7 @@ pub fn bloom_filter_load(ctx: &Context, input_args: &[ValkeyString]) -> ValkeyRe
         None => {
             // if filter not exists, create it.
             let hex = value.to_vec();
-            let validate_size_limit = !ctx.get_flags().contains(ContextFlags::REPLICATED);
+            let validate_size_limit = !must_obey_client(ctx);
             let bloom = match BloomObject::decode_object(&hex, validate_size_limit) {
                 Ok(v) => v,
                 Err(err) => {

src/wrapper/mod.rs

@@ -1 +1,29 @@
+use valkey_module::Context;
+
 pub mod bloom_callback;
+
+/// Wrapper for the ValkeyModule_MustObeyClient function.
+/// Takes in an Context and returns true if the if commands are arriving
+/// from the primary client or AOF client and should never be rejected.
+/// False otherwise.
+pub fn must_obey_client(ctx: &Context) -> bool {
+    #[cfg(not(feature = "valkey_8_0"))]
+    {
+        // ValkeyModule_MustObeyClient exists since Valkey 8.1, so we use this as it is a more performant check.
+        let ctx_raw = ctx.get_raw() as *mut valkey_module::ValkeyModuleCtx;
+
+        let status = unsafe { valkey_module::raw::ValkeyModule_MustObeyClient.unwrap()(ctx_raw) };
+        match status {
+            1 => true,
+            0 => false,
+            _ => panic!("We do not expect ValkeyModule_MustObeyClient to return anything other than 1 or 0."),
+        }
+    }
+
+    #[cfg(feature = "valkey_8_0")]
+    {
+        // On Valkey 8.0, fallback to checking for replicated flag in the GetContextFlags API as a best effort.
+        ctx.get_flags()
+            .contains(valkey_module::ContextFlags::REPLICATED)
+    }
+}

tests/test_bloom_correctness.py

@@ -135,7 +135,7 @@ def test_scaling_filter(self):
 
     def test_max_and_validate_scale_to_correctness(self):
         validate_scale_to_commands = [
-            ('BF.INSERT key ERROR 0.00000001 VALIDATESCALETO 13107101', "provided VALIDATESCALETO causes bloom object to exceed memory limit" ),
+            ('BF.INSERT MemLimitKey EXPANSION 25 ERROR 0.00000000000000001 VALIDATESCALETO 1627601', "provided VALIDATESCALETO causes bloom object to exceed memory limit" ),
             ('BF.INSERT key EXPANSION 1 VALIDATESCALETO 101601', "provided VALIDATESCALETO causes false positive to degrade to 0" )
         ]
         for cmd in validate_scale_to_commands:
@@ -144,12 +144,12 @@ def test_max_and_validate_scale_to_correctness(self):
                 assert False, "Expect BF.INSERT to fail if the wanted capacity would cause an error"
             except Exception as e:
                 assert cmd[1] == str(e), f"Unexpected error message: {e}" 
-        self.client.execute_command('BF.INSERT MemLimitKey ERROR 0.00000001 VALIDATESCALETO 13107100')
+        self.client.execute_command('BF.INSERT MemLimitKey EXPANSION 25 ERROR 0.00000000000000001 VALIDATESCALETO 1627600')
         self.client.execute_command('BF.INSERT FPKey VALIDATESCALETO 101600 EXPANSION 1')
         FPKey_max_capacity = self.client.execute_command(f'BF.INFO FPKey MAXSCALEDCAPACITY')
         MemLimitKeyMaxCapacity = self.client.execute_command(f'BF.INFO MemLimitKey MAXSCALEDCAPACITY')
         self.add_items_till_capacity(self.client, "FPKey", 101600,  1, "item")
-        self.add_items_till_capacity(self.client, "MemLimitKey", 13107100,  1, "item")
+        self.add_items_till_capacity(self.client, "MemLimitKey", 1627600,  1, "item")
         key_names = [("MemLimitKey", MemLimitKeyMaxCapacity, "operation exceeds bloom object memory limit"), ("FPKey", FPKey_max_capacity, "false positive degrades to 0 on scale out")]
         for key in key_names:
             try: