refactor: avoid potential unnecessary store accesses on dataset resol…

…ution (#4513)

* refactor: avoid potential unnecessary store accesses on dataset resolution

* pr remarks

core/control-plane/control-plane-catalog/src/main/java/org/eclipse/edc/connector/controlplane/catalog/CatalogCoreExtension.java

@@ -17,8 +17,10 @@
 import org.eclipse.edc.connector.controlplane.asset.spi.index.AssetIndex;
 import org.eclipse.edc.connector.controlplane.catalog.spi.DatasetResolver;
 import org.eclipse.edc.connector.controlplane.catalog.spi.DistributionResolver;
-import org.eclipse.edc.connector.controlplane.contract.spi.offer.ContractDefinitionResolver;
+import org.eclipse.edc.connector.controlplane.contract.spi.offer.store.ContractDefinitionStore;
 import org.eclipse.edc.connector.controlplane.policy.spi.store.PolicyDefinitionStore;
+import org.eclipse.edc.policy.engine.spi.PolicyEngine;
+import org.eclipse.edc.policy.engine.spi.PolicyScope;
 import org.eclipse.edc.runtime.metamodel.annotation.Extension;
 import org.eclipse.edc.runtime.metamodel.annotation.Inject;
 import org.eclipse.edc.runtime.metamodel.annotation.Provider;
@@ -30,8 +32,8 @@ public class CatalogCoreExtension implements ServiceExtension {
 
     public static final String NAME = "Catalog Core";
 
-    @Inject
-    private ContractDefinitionResolver contractDefinitionResolver;
+    @PolicyScope
+    public static final String CATALOG_SCOPE = "catalog";
 
     @Inject
     private AssetIndex assetIndex;
@@ -45,14 +47,22 @@ public class CatalogCoreExtension implements ServiceExtension {
     @Inject
     private CriterionOperatorRegistry criterionOperatorRegistry;
 
+    @Inject
+    private ContractDefinitionStore contractDefinitionStore;
+
+    @Inject
+    private PolicyEngine policyEngine;
+
     @Override
     public String name() {
         return NAME;
     }
 
     @Provider
     public DatasetResolver datasetResolver() {
+        var contractDefinitionResolver = new ContractDefinitionResolverImpl(contractDefinitionStore, policyEngine, policyDefinitionStore);
         return new DatasetResolverImpl(contractDefinitionResolver, assetIndex, policyDefinitionStore,
                 distributionResolver, criterionOperatorRegistry);
     }
+
 }

core/control-plane/control-plane-catalog/src/main/java/org/eclipse/edc/connector/controlplane/catalog/CatalogDefaultServicesExtension.java

@@ -52,4 +52,5 @@ public DataServiceRegistry dataServiceRegistry() {
     public DistributionResolver distributionResolver() {
         return new DefaultDistributionResolver(dataServiceRegistry, dataFlowManager);
     }
+
 }

core/control-plane/control-plane-catalog/src/main/java/org/eclipse/edc/connector/controlplane/catalog/ContractDefinitionResolverImpl.java

@@ -0,0 +1,74 @@
+/*
+ *  Copyright (c) 2024 Cofinity-X
+ *
+ *  This program and the accompanying materials are made available under the
+ *  terms of the Apache License, Version 2.0 which is available at
+ *  https://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  SPDX-License-Identifier: Apache-2.0
+ *
+ *  Contributors:
+ *       Cofinity-X - initial API and implementation
+ *
+ */
+
+package org.eclipse.edc.connector.controlplane.catalog;
+
+import org.eclipse.edc.connector.controlplane.catalog.spi.ContractDefinitionResolver;
+import org.eclipse.edc.connector.controlplane.catalog.spi.ResolvedContractDefinitions;
+import org.eclipse.edc.connector.controlplane.contract.spi.offer.store.ContractDefinitionStore;
+import org.eclipse.edc.connector.controlplane.contract.spi.types.offer.ContractDefinition;
+import org.eclipse.edc.connector.controlplane.policy.spi.PolicyDefinition;
+import org.eclipse.edc.connector.controlplane.policy.spi.store.PolicyDefinitionStore;
+import org.eclipse.edc.policy.engine.spi.PolicyContextImpl;
+import org.eclipse.edc.policy.engine.spi.PolicyEngine;
+import org.eclipse.edc.policy.model.Policy;
+import org.eclipse.edc.spi.agent.ParticipantAgent;
+import org.eclipse.edc.spi.query.QuerySpec;
+import org.eclipse.edc.spi.result.Result;
+
+import java.util.HashMap;
+import java.util.Optional;
+
+import static java.lang.String.format;
+import static org.eclipse.edc.connector.controlplane.catalog.CatalogCoreExtension.CATALOG_SCOPE;
+
+/**
+ * Determines the contract definitions applicable to a {@link ParticipantAgent} by evaluating the access control and
+ * usage policies associated with a set of assets as defined by {@link ContractDefinition}s. On the distinction between
+ * access control and usage policy, see {@link ContractDefinition}.
+ */
+public class ContractDefinitionResolverImpl implements ContractDefinitionResolver {
+    private final PolicyEngine policyEngine;
+    private final PolicyDefinitionStore policyStore;
+    private final ContractDefinitionStore definitionStore;
+
+    public ContractDefinitionResolverImpl(ContractDefinitionStore contractDefinitionStore, PolicyEngine policyEngine, PolicyDefinitionStore policyStore) {
+        definitionStore = contractDefinitionStore;
+        this.policyEngine = policyEngine;
+        this.policyStore = policyStore;
+    }
+
+    @Override
+    public ResolvedContractDefinitions resolveFor(ParticipantAgent agent) {
+        var policies = new HashMap<String, Policy>();
+        var definitions = definitionStore.findAll(QuerySpec.max())
+                .filter(definition -> {
+                    var policyContext = PolicyContextImpl.Builder.newInstance().additional(ParticipantAgent.class, agent).build();
+                    var accessResult = Optional.of(definition.getAccessPolicyId())
+                            .map(policyId -> policies.computeIfAbsent(policyId,
+                                    key -> Optional.ofNullable(policyStore.findById(key))
+                                            .map(PolicyDefinition::getPolicy)
+                                            .orElse(null))
+                            )
+                            .map(policy -> policyEngine.evaluate(CATALOG_SCOPE, policy, policyContext))
+                            .orElse(Result.failure(format("Policy %s not found", definition.getAccessPolicyId())));
+
+                    return accessResult.succeeded();
+                })
+                .toList();
+
+        return new ResolvedContractDefinitions(definitions, policies);
+    }
+
+}

core/control-plane/control-plane-catalog/src/main/java/org/eclipse/edc/connector/controlplane/catalog/DatasetResolverImpl.java

@@ -17,15 +17,17 @@
 import org.eclipse.edc.connector.controlplane.asset.spi.domain.Asset;
 import org.eclipse.edc.connector.controlplane.asset.spi.index.AssetIndex;
 import org.eclipse.edc.connector.controlplane.catalog.spi.Catalog;
+import org.eclipse.edc.connector.controlplane.catalog.spi.ContractDefinitionResolver;
 import org.eclipse.edc.connector.controlplane.catalog.spi.DataService;
 import org.eclipse.edc.connector.controlplane.catalog.spi.Dataset;
 import org.eclipse.edc.connector.controlplane.catalog.spi.DatasetResolver;
 import org.eclipse.edc.connector.controlplane.catalog.spi.DistributionResolver;
 import org.eclipse.edc.connector.controlplane.contract.spi.ContractOfferId;
-import org.eclipse.edc.connector.controlplane.contract.spi.offer.ContractDefinitionResolver;
 import org.eclipse.edc.connector.controlplane.contract.spi.types.offer.ContractDefinition;
+import org.eclipse.edc.connector.controlplane.policy.spi.PolicyDefinition;
 import org.eclipse.edc.connector.controlplane.policy.spi.store.PolicyDefinitionStore;
 import org.eclipse.edc.dataaddress.httpdata.spi.HttpDataAddressSchema;
+import org.eclipse.edc.policy.model.Policy;
 import org.eclipse.edc.policy.model.PolicyType;
 import org.eclipse.edc.spi.agent.ParticipantAgent;
 import org.eclipse.edc.spi.query.CriterionOperatorRegistry;
@@ -34,6 +36,7 @@
 
 import java.util.Base64;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.function.Predicate;
 import java.util.stream.Stream;
@@ -61,34 +64,36 @@ public DatasetResolverImpl(ContractDefinitionResolver contractDefinitionResolver
     @Override
     @NotNull
     public Stream<Dataset> query(ParticipantAgent agent, QuerySpec querySpec) {
-        var contractDefinitions = contractDefinitionResolver.definitionsFor(agent).toList();
+        var resolved = contractDefinitionResolver.resolveFor(agent);
+        var contractDefinitions = resolved.contractDefinitions();
         if (contractDefinitions.isEmpty()) {
             return Stream.empty();
         }
 
         var assetsQuery = QuerySpec.Builder.newInstance().offset(0).limit(MAX_VALUE).filter(querySpec.getFilterExpression()).build();
         return assetIndex.queryAssets(assetsQuery)
-                .map(asset -> toDataset(contractDefinitions, asset))
+                .map(asset -> toDataset(contractDefinitions, asset, resolved.policies()))
                 .filter(Dataset::hasOffers)
                 .skip(querySpec.getOffset())
                 .limit(querySpec.getLimit());
     }
 
     @Override
     public Dataset getById(ParticipantAgent agent, String id) {
-        var contractDefinitions = contractDefinitionResolver.definitionsFor(agent).toList();
+        var resolved = contractDefinitionResolver.resolveFor(agent);
+        var contractDefinitions = resolved.contractDefinitions();
         if (contractDefinitions.isEmpty()) {
             return null;
         }
 
         return Optional.of(id)
                 .map(assetIndex::findById)
-                .map(asset -> toDataset(contractDefinitions, asset))
+                .map(asset -> toDataset(contractDefinitions, asset, resolved.policies()))
                 .filter(Dataset::hasOffers)
                 .orElse(null);
     }
 
-    private Dataset.Builder buildDataset(Asset asset) {
+    private Dataset.Builder<?, ?> buildDataset(Asset asset) {
         if (!asset.isCatalog()) {
             return Dataset.Builder.newInstance();
         }
@@ -101,7 +106,7 @@ private Dataset.Builder buildDataset(Asset asset) {
                         .build());
     }
 
-    private Dataset toDataset(List<ContractDefinition> contractDefinitions, Asset asset) {
+    private Dataset toDataset(List<ContractDefinition> contractDefinitions, Asset asset, Map<String, Policy> policies) {
 
         var distributions = distributionResolver.getDistributions(asset);
         var datasetBuilder = buildDataset(asset)
@@ -116,10 +121,15 @@ private Dataset toDataset(List<ContractDefinition> contractDefinitions, Asset as
                         .test(asset)
                 )
                 .forEach(contractDefinition -> {
-                    var policyDefinition = policyDefinitionStore.findById(contractDefinition.getContractPolicyId());
-                    if (policyDefinition != null) {
+                    var policy = policies.computeIfAbsent(contractDefinition.getContractPolicyId(), policyId ->
+                            Optional.ofNullable(policyDefinitionStore.findById(policyId))
+                                    .map(PolicyDefinition::getPolicy)
+                                    .orElse(null)
+                    );
+
+                    if (policy != null) {
                         var contractId = ContractOfferId.create(contractDefinition.getId(), asset.getId());
-                        var offerPolicy = policyDefinition.getPolicy().toBuilder().type(PolicyType.OFFER).build();
+                        var offerPolicy = policy.toBuilder().type(PolicyType.OFFER).build();
                         datasetBuilder.offer(contractId.toString(), offerPolicy);
                     }
                 });

core/control-plane/control-plane-contract/src/test/java/org/eclipse/edc/connector/controlplane/contract/offer/ContractDefinitionResolverImplTest.java → core/control-plane/control-plane-catalog/src/test/java/org/eclipse/edc/connector/controlplane/catalog/ContractDefinitionResolverImplTest.java

@@ -13,7 +13,7 @@
  *
  */
 
-package org.eclipse.edc.connector.controlplane.contract.offer;
+package org.eclipse.edc.connector.controlplane.catalog;
 
 import org.eclipse.edc.connector.controlplane.contract.spi.offer.store.ContractDefinitionStore;
 import org.eclipse.edc.connector.controlplane.contract.spi.types.offer.ContractDefinition;
@@ -23,123 +23,113 @@
 import org.eclipse.edc.policy.engine.spi.PolicyEngine;
 import org.eclipse.edc.policy.model.Policy;
 import org.eclipse.edc.spi.agent.ParticipantAgent;
-import org.eclipse.edc.spi.monitor.Monitor;
 import org.eclipse.edc.spi.query.QuerySpec;
 import org.eclipse.edc.spi.result.Result;
-import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
-import java.util.Map;
 import java.util.stream.Stream;
 
+import static java.util.Collections.emptyMap;
 import static org.assertj.core.api.Assertions.assertThat;
-import static org.eclipse.edc.connector.controlplane.contract.spi.offer.ContractDefinitionResolver.CATALOGING_SCOPE;
+import static org.assertj.core.api.Assertions.entry;
+import static org.eclipse.edc.connector.controlplane.catalog.CatalogCoreExtension.CATALOG_SCOPE;
 import static org.mockito.AdditionalMatchers.and;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.ArgumentMatchers.isA;
 import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.only;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoInteractions;
 import static org.mockito.Mockito.when;
 
 class ContractDefinitionResolverImplTest {
 
-    private final PolicyEngine policyEngine = mock(PolicyEngine.class);
-    private final PolicyDefinitionStore policyStore = mock(PolicyDefinitionStore.class);
-    private final ContractDefinitionStore definitionStore = mock(ContractDefinitionStore.class);
+    private final PolicyEngine policyEngine = mock();
+    private final PolicyDefinitionStore policyStore = mock();
+    private final ContractDefinitionStore definitionStore = mock();
 
-    private ContractDefinitionResolverImpl definitionService;
-
-    @BeforeEach
-    void setUp() {
-        definitionService = new ContractDefinitionResolverImpl(mock(Monitor.class), definitionStore, policyEngine, policyStore);
-    }
+    private final ContractDefinitionResolverImpl resolver = new ContractDefinitionResolverImpl(definitionStore,
+            policyEngine, policyStore);
 
     @Test
-    void definitionsFor_verifySatisfiesPolicies() {
-        var agent = new ParticipantAgent(Map.of(), Map.of());
+    void shouldReturnDefinition_whenAccessPolicySatisfied() {
+        var agent = new ParticipantAgent(emptyMap(), emptyMap());
         var def = PolicyDefinition.Builder.newInstance().policy(Policy.Builder.newInstance().build()).build();
         when(policyStore.findById(any())).thenReturn(def);
         when(policyEngine.evaluate(any(), any(), isA(PolicyContext.class))).thenReturn(Result.success());
         when(definitionStore.findAll(any())).thenReturn(Stream.of(createContractDefinition()));
 
-        var definitions = definitionService.definitionsFor(agent);
+        var result = resolver.resolveFor(agent);
 
-        assertThat(definitions).hasSize(1);
+        assertThat(result.contractDefinitions()).hasSize(1);
+        assertThat(result.policies()).hasSize(1);
         verify(policyEngine, atLeastOnce()).evaluate(
-                eq(CATALOGING_SCOPE),
+                eq(CATALOG_SCOPE),
                 eq(def.getPolicy()),
                 and(isA(PolicyContext.class), argThat(c -> c.getContextData(ParticipantAgent.class).equals(agent)))
         );
         verify(definitionStore).findAll(any());
     }
 
     @Test
-    void definitionsFor_verifyDoesNotSatisfyAccessPolicy() {
-        var agent = new ParticipantAgent(Map.of(), Map.of());
+    void shouldNotReturnDefinition_whenAccessPolicyNotSatisfied() {
+        var agent = new ParticipantAgent(emptyMap(), emptyMap());
         var definition = PolicyDefinition.Builder.newInstance().policy(Policy.Builder.newInstance().build()).id("access").build();
         when(policyStore.findById(any())).thenReturn(definition);
         var contractDefinition = createContractDefinition();
         when(policyEngine.evaluate(any(), any(), isA(PolicyContext.class))).thenReturn(Result.failure("invalid"));
         when(definitionStore.findAll(any())).thenReturn(Stream.of(contractDefinition));
 
-        var result = definitionService.definitionsFor(agent);
+        var result = resolver.resolveFor(agent);
 
-        assertThat(result).isEmpty();
+        assertThat(result.contractDefinitions()).isEmpty();
+        assertThat(result.policies()).hasSize(1);
         verify(definitionStore).findAll(any());
     }
 
     @Test
-    void definitionsFor_verifyPoliciesNotFound() {
-        var agent = new ParticipantAgent(Map.of(), Map.of());
+    void shouldNotReturnDefinition_whenAccessPolicyDoesNotExist() {
+        var agent = new ParticipantAgent(emptyMap(), emptyMap());
         when(policyStore.findById(any())).thenReturn(null);
         when(policyEngine.evaluate(any(), any(), isA(PolicyContext.class))).thenReturn(Result.success());
         when(definitionStore.findAll(QuerySpec.max())).thenReturn(Stream.of(createContractDefinition()));
 
-        var definitions = definitionService.definitionsFor(agent);
+        var result = resolver.resolveFor(agent);
 
-        assertThat(definitions).isEmpty();
+        assertThat(result.contractDefinitions()).isEmpty();
+        assertThat(result.policies()).isEmpty();
         verifyNoInteractions(policyEngine);
     }
 
     @Test
-    void definitionFor_found() {
-        var agent = new ParticipantAgent(Map.of(), Map.of());
-        var definition = PolicyDefinition.Builder.newInstance().policy(Policy.Builder.newInstance().build()).build();
-        var contractDefinition = createContractDefinition();
-        when(policyStore.findById(any())).thenReturn(definition);
+    void shouldQueryPolicyOnce_whenDifferentDefinitionsHaveSamePolicy() {
+        var contractDefinition1 = contractDefinitionBuilder().accessPolicyId("accessPolicyId").build();
+        var contractDefinition2 = contractDefinitionBuilder().accessPolicyId("accessPolicyId").build();
+        var policy = Policy.Builder.newInstance().build();
+        var policyDefinition = PolicyDefinition.Builder.newInstance().policy(policy).build();
+        when(policyStore.findById(any())).thenReturn(policyDefinition);
         when(policyEngine.evaluate(any(), any(), isA(PolicyContext.class))).thenReturn(Result.success());
-        when(definitionStore.findById("1")).thenReturn(contractDefinition);
+        when(definitionStore.findAll(any())).thenReturn(Stream.of(contractDefinition1, contractDefinition2));
 
-        var result = definitionService.definitionFor(agent, "1");
+        var result = resolver.resolveFor(new ParticipantAgent(emptyMap(), emptyMap()));
 
-        assertThat(result).isNotNull();
-        verify(policyEngine, atLeastOnce()).evaluate(
-                eq(CATALOGING_SCOPE),
-                eq(definition.getPolicy()),
-                and(isA(PolicyContext.class), argThat(c -> c.getContextData(ParticipantAgent.class).equals(agent)))
-        );
+        assertThat(result.contractDefinitions()).hasSize(2);
+        assertThat(result.policies()).hasSize(1).containsOnly(entry("accessPolicyId", policy));
+        verify(policyStore, only()).findById("accessPolicyId");
     }
 
-    @Test
-    void definitionFor_notFound() {
-        var agent = new ParticipantAgent(Map.of(), Map.of());
-        when(definitionStore.findById(any())).thenReturn(null);
-
-        var result = definitionService.definitionFor(agent, "nodefinition");
-
-        assertThat(result).isNull();
-        verifyNoInteractions(policyEngine);
+    private ContractDefinition createContractDefinition() {
+        return contractDefinitionBuilder()
+                .build();
     }
 
-    private ContractDefinition createContractDefinition() {
+    private ContractDefinition.Builder contractDefinitionBuilder() {
         return ContractDefinition.Builder.newInstance()
                 .id("1")
                 .accessPolicyId("access")
-                .contractPolicyId("contract")
-                .build();
+                .contractPolicyId("contract");
     }
 }