feat: update sticky proto (#282)

* fixup! update sticky proto

Author: vahidlazio

confidence-proto/src/main/proto/confidence/wasm/wasm_api.proto

@@ -20,14 +20,15 @@ message ResolveWithStickyRequest {
   ResolveFlagsRequest resolve_request = 1;
 
   // Context about the materialization required for the resolve
-  MaterializationContext materialization_context = 7;
+  map<string, MaterializationMap> materializations_per_unit = 2;
 
   // if a materialization info is missing, we want to return to the caller immediately
-  bool fail_fast_on_sticky = 8;
+  bool fail_fast_on_sticky = 3;
 }
 
-message MaterializationContext {
-  map<string, MaterializationInfo> unit_materialization_info = 1;
+message MaterializationMap {
+  // materialization name to info
+  map<string, MaterializationInfo> info_map = 1;
 }
 
 message MaterializationInfo {

openfeature-provider-local/pom.xml

@@ -19,7 +19,7 @@
     <maven.compiler.source>17</maven.compiler.source>
     <maven.compiler.target>17</maven.compiler.target>
     <!-- WASM file version. When updated, make sure to mvn clean to force a new download -->
-    <confidence.resolver.wasm.version>v0.3.1</confidence.resolver.wasm.version>
+    <confidence.resolver.wasm.version>v0.4.0</confidence.resolver.wasm.version>
   </properties>
 
 

openfeature-provider-local/src/main/java/com/spotify/confidence/MaterializationRepository.java

@@ -1,12 +1,11 @@
 package com.spotify.confidence;
 
-import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 
 public non-sealed interface MaterializationRepository extends StickyResolveStrategy {
   CompletableFuture<Map<String, MaterializationInfo>> loadMaterializedAssignmentsForUnit(
-      String unit, Map<String, List<String>> materializationsToRules);
+      String unit, String materialization);
 
   CompletableFuture<Void> storeAssignment(
       String unit, Map<String, MaterializationInfo> assignments);

openfeature-provider-local/src/main/java/com/spotify/confidence/SwapWasmResolverApi.java

@@ -1,5 +1,6 @@
 package com.spotify.confidence;
 
+import com.spotify.confidence.flags.resolver.v1.MaterializationMap;
 import com.spotify.confidence.flags.resolver.v1.ResolveWithStickyRequest;
 import com.spotify.confidence.flags.resolver.v1.ResolveWithStickyResponse;
 import com.spotify.confidence.shaded.flags.resolver.v1.ResolveFlagsRequest;
@@ -8,6 +9,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.stream.Collectors;
@@ -147,36 +149,46 @@ private ResolveWithStickyRequest handleMissingMaterializations(
         missingItems.stream()
             .collect(
                 Collectors.groupingBy(
-                    ResolveWithStickyResponse.MissingMaterializationItem::getUnit,
-                    Collectors.toMap(
-                        ResolveWithStickyResponse.MissingMaterializationItem
-                            ::getReadMaterialization,
-                        item -> List.of(item.getRule()),
-                        (existing, replacement) -> {
-                          final var combined = new java.util.ArrayList<>(existing);
-                          combined.addAll(replacement);
-                          return combined;
-                        })));
+                    ResolveWithStickyResponse.MissingMaterializationItem::getUnit));
 
-    // Load materialized assignments for all missing units
-    final var materializationContext = request.getMaterializationContext().toBuilder();
+    final HashMap<String, MaterializationMap> materializationPerUnitMap = new HashMap<>();
 
+    // Load materialized assignments for all missing units
     missingByUnit.forEach(
-        (unit, materializationsToRules) -> {
-          try {
-            final var loadedAssignments =
-                repository.loadMaterializedAssignmentsForUnit(unit, materializationsToRules).get();
-            loadedAssignments.forEach(
-                (materialization, info) ->
-                    materializationContext.putUnitMaterializationInfo(unit, info.toProto()));
-          } catch (Exception e) {
-            throw new RuntimeException(
-                "Failed to load materialized assignments for unit: " + unit, e);
-          }
+        (unit, materializationInfoItem) -> {
+          materializationInfoItem.forEach(
+              item -> {
+                final Map<String, MaterializationInfo> loadedAssignments;
+                try {
+                  loadedAssignments =
+                      repository
+                          .loadMaterializedAssignmentsForUnit(unit, item.getReadMaterialization())
+                          .get();
+                } catch (InterruptedException | ExecutionException e) {
+                  throw new RuntimeException(e);
+                }
+                materializationPerUnitMap.put(
+                    unit,
+                    MaterializationMap.newBuilder()
+                        .putAllInfoMap(
+                            loadedAssignments.entrySet().stream()
+                                .collect(
+                                    Collectors.toMap(
+                                        Map.Entry::getKey,
+                                        e -> {
+                                          final var info = e.getValue();
+                                          return com.spotify.confidence.flags.resolver.v1
+                                              .MaterializationInfo.newBuilder()
+                                              .setUnitInInfo(info.isUnitInMaterialization())
+                                              .putAllRuleToVariant(info.ruleToVariant())
+                                              .build();
+                                        })))
+                        .build());
+              });
         });
 
     // Return new request with updated materialization context
-    return request.toBuilder().setMaterializationContext(materializationContext.build()).build();
+    return request.toBuilder().putAllMaterializationsPerUnit(materializationPerUnitMap).build();
   }
 
   public ResolveFlagsResponse resolve(ResolveFlagsRequest request) {

openfeature-provider-local/src/test/java/com/spotify/confidence/WasmResolveTest.java

@@ -4,7 +4,6 @@
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -19,7 +18,6 @@
 import dev.openfeature.sdk.ImmutableContext;
 import dev.openfeature.sdk.ProviderEvaluation;
 import dev.openfeature.sdk.Value;
-import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import org.junit.jupiter.api.Test;
@@ -145,8 +143,6 @@ public void testMaterializationRepositoryWhenMaterializationsMissing() {
     assertEquals("on", structure.getValue("data").asString());
 
     // Assert that the materialization repository was called with correct input
-    verify(mockRepository)
-        .loadMaterializedAssignmentsForUnit(
-            eq("test-user"), eq(Map.of("read-mat", List.of("MyRule"))));
+    verify(mockRepository).loadMaterializedAssignmentsForUnit("test-user", "read-mat");
   }
 }