perf(python): add host-side optimizations for pre-evaluation, context filtering, and index-based eval

.github/workflows/ci.yml

@@ -33,7 +33,10 @@ jobs:
             ${{ runner.os }}-cargo-test-
 
       - name: Run tests
-        run: cargo test --verbose
+        run: |
+          cargo test --lib --verbose -- --test-threads=1
+          cargo test --test integration_tests --test metadata_merging_tests --verbose -- --test-threads=1
+          cargo test --test gherkin_tests --verbose
 
   clippy:
     name: Clippy

java/src/main/java/dev/openfeature/flagd/evaluator/FlagEvaluator.java

@@ -19,7 +19,10 @@
 import java.nio.charset.StandardCharsets;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 /**
  * Thread-safe flag evaluator using the flagd-evaluator WASM module.
@@ -97,6 +100,15 @@ public class FlagEvaluator implements AutoCloseable {
     // Cache of pre-evaluated results for static/disabled flags (replaced atomically on updateState)
     private volatile Map<String, EvaluationResult<Object>> preEvaluatedCache = Collections.emptyMap();
 
+    // Per-flag required context keys for host-side filtering (read/written inside synchronized methods)
+    private Map<String, Set<String>> requiredContextKeysCache = Collections.emptyMap();
+
+    // Flag key to numeric index mapping (read/written inside synchronized methods)
+    private Map<String, Integer> flagIndexCache = Collections.emptyMap();
+
+    // WASM export for index-based evaluation (may be null if WASM module doesn't support it)
+    private final ExportFunction evaluateByIndexFunction;
+
     /**
      * Creates a new flag evaluator with strict validation mode.
      *
@@ -119,6 +131,15 @@ public FlagEvaluator(ValidationMode validationMode) {
         this.deallocFunction = wasmInstance.export("dealloc");
         this.memory = wasmInstance.memory();
 
+        // Bind evaluate_by_index if available (newer WASM modules)
+        ExportFunction evalByIndex = null;
+        try {
+            evalByIndex = wasmInstance.export("evaluate_by_index");
+        } catch (Exception e) {
+            // Older WASM module without evaluate_by_index — fall back to string-based eval
+        }
+        this.evaluateByIndexFunction = evalByIndex;
+
         // Pre-allocate buffers for evaluation (avoids alloc calls per evaluation)
         this.flagKeyBufferPtr = allocFunction.apply(MAX_FLAG_KEY_SIZE)[0];
         this.contextBufferPtr = allocFunction.apply(MAX_CONTEXT_SIZE)[0];
@@ -170,6 +191,22 @@ public synchronized UpdateStateResult updateState(String jsonConfig) throws Eval
             Map<String, EvaluationResult<Object>> preEval = result.getPreEvaluated();
             this.preEvaluatedCache = (preEval != null) ? preEval : Collections.emptyMap();
 
+            // Update required context keys cache
+            Map<String, List<String>> reqKeys = result.getRequiredContextKeys();
+            if (reqKeys != null) {
+                Map<String, Set<String>> keySets = new HashMap<>(reqKeys.size());
+                for (Map.Entry<String, List<String>> entry : reqKeys.entrySet()) {
+                    keySets.put(entry.getKey(), new HashSet<>(entry.getValue()));
+                }
+                this.requiredContextKeysCache = keySets;
+            } else {
+                this.requiredContextKeysCache = Collections.emptyMap();
+            }
+
+            // Update flag index cache
+            Map<String, Integer> indices = result.getFlagIndices();
+            this.flagIndexCache = (indices != null) ? indices : Collections.emptyMap();
+
             return result;
         } catch (Exception e) {
             throw new EvaluatorException("Failed to update state", e);
@@ -217,6 +254,13 @@ public synchronized <T> EvaluationResult<T> evaluateFlag(Class<T> type, String f
             return (EvaluationResult<T>) (EvaluationResult<?>) cached;
         }
 
+        return evaluateFlagInternal(type, flagKey, contextJson);
+    }
+
+    /**
+     * Internal evaluation using flag key string and evaluate_reusable WASM export.
+     */
+    private <T> EvaluationResult<T> evaluateFlagInternal(Class<T> type, String flagKey, String contextJson) throws EvaluatorException {
         byte[] flagBytes = flagKey.getBytes(StandardCharsets.UTF_8);
         if (flagBytes.length > MAX_FLAG_KEY_SIZE) {
             throw new EvaluatorException("Flag key exceeds maximum size of " + MAX_FLAG_KEY_SIZE + " bytes");
@@ -254,6 +298,39 @@ public synchronized <T> EvaluationResult<T> evaluateFlag(Class<T> type, String f
         }
     }
 
+    /**
+     * Evaluates a flag using the numeric index path (evaluate_by_index WASM export).
+     *
+     * <p>This avoids flag key string serialization and uses O(1) Vec lookup on the Rust side.
+     * The context must already be pre-enriched with {@code $flagd.*} and {@code targetingKey}.
+     */
+    private <T> EvaluationResult<T> evaluateByIndex(Class<T> type, int flagIndex, String contextJson) throws EvaluatorException {
+        long contextPtr = 0;
+        int contextLen = 0;
+        if (contextJson != null && !contextJson.isEmpty()) {
+            byte[] contextBytes = contextJson.getBytes(StandardCharsets.UTF_8);
+            if (contextBytes.length > MAX_CONTEXT_SIZE) {
+                throw new EvaluatorException("Context exceeds maximum size of " + MAX_CONTEXT_SIZE + " bytes");
+            }
+            memory.write((int) contextBufferPtr, contextBytes);
+            contextPtr = contextBufferPtr;
+            contextLen = contextBytes.length;
+        }
+
+        try {
+            long packedResult = evaluateByIndexFunction.apply(flagIndex, contextPtr, contextLen)[0];
+            int resultPtr = (int) (packedResult >>> 32);
+            int resultLen = (int) (packedResult & 0xFFFFFFFFL);
+
+            String resultJson = memory.readString(resultPtr, resultLen);
+            deallocFunction.apply(resultPtr, resultLen);
+
+            return OBJECT_MAPPER.readValue(resultJson, JAVA_TYPE_MAP.get(type));
+        } catch (Exception e) {
+            throw new EvaluatorException("Failed to evaluate flag by index: " + flagIndex, e);
+        }
+    }
+
     /**
      * Evaluates a flag with an EvaluationContext.
      *
@@ -266,19 +343,45 @@ public synchronized <T> EvaluationResult<T> evaluateFlag(Class<T> type, String f
      * @return the evaluation result containing value, variant, reason, and metadata
      * @throws EvaluatorException if the evaluation or serialization fails
      */
-    public <T> EvaluationResult<T> evaluateFlag(Class<T> type, String flagKey, EvaluationContext context) throws EvaluatorException {
+    @SuppressWarnings("unchecked")
+    public synchronized <T> EvaluationResult<T> evaluateFlag(Class<T> type, String flagKey, EvaluationContext context) throws EvaluatorException {
         try {
+            // Fast path: return cached result for static/disabled flags
+            EvaluationResult<Object> cached = preEvaluatedCache.get(flagKey);
+            if (cached != null) {
+                return (EvaluationResult<T>) (EvaluationResult<?>) cached;
+            }
+
             // Fast path: empty context
             if (context == null || context.isEmpty()) {
-                return evaluateFlag(type, flagKey, (String) null);
+                return evaluateFlagInternal(type, flagKey, (String) null);
+            }
+
+            // Check if we can use filtered serialization
+            Set<String> requiredKeys = requiredContextKeysCache.get(flagKey);
+            String contextJson;
+            if (requiredKeys != null) {
+                // Filtered path: only serialize keys the targeting rule references
+                contextJson = EvaluationContextSerializer.serializeFiltered(context, requiredKeys, flagKey);
+            } else {
+                // Full serialization path (flag uses {"var": ""} or older WASM module)
+                ByteArrayOutputStream buffer = JSON_BUFFER.get();
+                buffer.reset();
+                try (JsonGenerator generator = JSON_FACTORY.createGenerator(buffer)) {
+                    OBJECT_MAPPER.writeValue(generator, context);
+                }
+                contextJson = buffer.toString(StandardCharsets.UTF_8.name());
             }
-            // Use ThreadLocal buffer for streaming serialization
-            ByteArrayOutputStream buffer = JSON_BUFFER.get();
-            buffer.reset();
-            try (JsonGenerator generator = JSON_FACTORY.createGenerator(buffer)) {
-                OBJECT_MAPPER.writeValue(generator, context);
+
+            // Check if we can use index-based evaluation
+            Integer flagIndex = flagIndexCache.get(flagKey);
+            if (flagIndex != null && evaluateByIndexFunction != null && requiredKeys != null) {
+                // Index-based path: avoids flag key string overhead
+                return evaluateByIndex(type, flagIndex, contextJson);
             }
-            return evaluateFlag(type, flagKey, buffer.toString(StandardCharsets.UTF_8.name()));
+
+            // Fall back to string-based evaluation
+            return evaluateFlagInternal(type, flagKey, contextJson);
         } catch (EvaluatorException e) {
             throw e;
         } catch (Exception e) {

java/src/main/java/dev/openfeature/flagd/evaluator/UpdateStateResult.java

@@ -18,6 +18,10 @@ public class UpdateStateResult {
 
     private Map<String, EvaluationResult<Object>> preEvaluated;
 
+    private Map<String, java.util.List<String>> requiredContextKeys;
+
+    private Map<String, Integer> flagIndices;
+
     public UpdateStateResult() {
     }
 
@@ -78,6 +82,39 @@ public void setPreEvaluated(Map<String, EvaluationResult<Object>> preEvaluated)
         this.preEvaluated = preEvaluated;
     }
 
+    /**
+     * Gets the required context keys per flag for host-side filtering.
+     *
+     * <p>When present for a flag, the host should only serialize these context keys
+     * (plus {@code $flagd.*} enrichment and {@code targetingKey}) before calling evaluate.
+     * If a flag is absent from this map, send the full context.
+     *
+     * @return map of flag key to required context keys, or null if not available
+     */
+    public Map<String, java.util.List<String>> getRequiredContextKeys() {
+        return requiredContextKeys;
+    }
+
+    public void setRequiredContextKeys(Map<String, java.util.List<String>> requiredContextKeys) {
+        this.requiredContextKeys = requiredContextKeys;
+    }
+
+    /**
+     * Gets the flag key to numeric index mapping for {@code evaluate_by_index}.
+     *
+     * <p>Allows calling the WASM {@code evaluate_by_index(index, ...)} export
+     * instead of passing flag key strings.
+     *
+     * @return map of flag key to numeric index, or null if not available
+     */
+    public Map<String, Integer> getFlagIndices() {
+        return flagIndices;
+    }
+
+    public void setFlagIndices(Map<String, Integer> flagIndices) {
+        this.flagIndices = flagIndices;
+    }
+
     @Override
     public String toString() {
         return "UpdateStateResult{" +

java/src/main/java/dev/openfeature/flagd/evaluator/jackson/EvaluationContextSerializer.java

@@ -7,11 +7,13 @@
 import dev.openfeature.sdk.Structure;
 import dev.openfeature.sdk.Value;
 
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 /**
  * Custom serializer for EvaluationContext (including LayeredEvaluationContext and MutableContext).
@@ -39,6 +41,108 @@ public void serialize(EvaluationContext ctx, JsonGenerator gen, SerializerProvid
         gen.writeEndObject();
     }
 
+    // ThreadLocal buffer for filtered serialization
+    private static final ThreadLocal<ByteArrayOutputStream> FILTERED_BUFFER =
+        ThreadLocal.withInitial(() -> new ByteArrayOutputStream(2048));
+
+    private static final com.fasterxml.jackson.core.JsonFactory SHARED_JSON_FACTORY =
+        new com.fasterxml.jackson.core.JsonFactory();
+
+    /**
+     * Serializes a filtered subset of the evaluation context with $flagd enrichment.
+     *
+     * <p>Only includes the specified required keys from the context, plus the
+     * {@code $flagd} enrichment object and {@code targetingKey}. This dramatically
+     * reduces the serialized size for large contexts where the targeting rule only
+     * references a few fields.
+     *
+     * @param ctx          the evaluation context to filter and serialize
+     * @param requiredKeys the set of context keys that the targeting rule references
+     * @param flagKey      the flag key (for $flagd.flagKey enrichment)
+     * @return the filtered, enriched JSON string
+     * @throws IOException if serialization fails
+     */
+    public static String serializeFiltered(
+            EvaluationContext ctx,
+            Set<String> requiredKeys,
+            String flagKey) throws IOException {
+        ByteArrayOutputStream buffer = FILTERED_BUFFER.get();
+        buffer.reset();
+
+        try (JsonGenerator gen = SHARED_JSON_FACTORY.createGenerator(buffer)) {
+            gen.writeStartObject();
+
+            // Write only the required keys from the context
+            for (String key : requiredKeys) {
+                // Skip $flagd — we add it ourselves below
+                if (key.startsWith("$flagd")) {
+                    continue;
+                }
+                Value value = ctx.getValue(key);
+                if (value != null) {
+                    gen.writeFieldName(key);
+                    writeValue(gen, value);
+                } else if ("targetingKey".equals(key)) {
+                    // targetingKey defaults to empty string if not in context
+                    String tk = ctx.getTargetingKey();
+                    gen.writeStringField("targetingKey", tk != null ? tk : "");
+                }
+            }
+
+            // Ensure targetingKey is always present
+            if (!requiredKeys.contains("targetingKey")) {
+                String tk = ctx.getTargetingKey();
+                gen.writeStringField("targetingKey", tk != null ? tk : "");
+            }
+
+            // Add $flagd enrichment
+            gen.writeObjectFieldStart("$flagd");
+            gen.writeStringField("flagKey", flagKey);
+            gen.writeNumberField("timestamp", System.currentTimeMillis() / 1000);
+            gen.writeEndObject();
+
+            gen.writeEndObject();
+        }
+
+        return buffer.toString("UTF-8");
+    }
+
+    /**
+     * Writes an OpenFeature Value to a JsonGenerator.
+     */
+    private static void writeValue(JsonGenerator gen, Value value) throws IOException {
+        if (value == null || value.isNull()) {
+            gen.writeNull();
+        } else if (value.isBoolean()) {
+            gen.writeBoolean(value.asBoolean());
+        } else if (value.isNumber()) {
+            double d = value.asDouble();
+            if (d == Math.floor(d) && !Double.isInfinite(d)) {
+                gen.writeNumber((long) d);
+            } else {
+                gen.writeNumber(d);
+            }
+        } else if (value.isString()) {
+            gen.writeString(value.asString());
+        } else if (value.isList()) {
+            gen.writeStartArray();
+            for (Value item : value.asList()) {
+                writeValue(gen, item);
+            }
+            gen.writeEndArray();
+        } else if (value.isStructure()) {
+            gen.writeStartObject();
+            Structure structure = value.asStructure();
+            for (String key : structure.keySet()) {
+                gen.writeFieldName(key);
+                writeValue(gen, structure.getValue(key));
+            }
+            gen.writeEndObject();
+        } else {
+            gen.writeNull();
+        }
+    }
+
     /**
      * Extracts the raw Java object from an OpenFeature Value wrapper.
      * Recursively handles nested structures and lists.