Move some production-only functionality to PolicyManager

Author: prdoyle

libs/entitlement/src/main/java/org/elasticsearch/entitlement/runtime/policy/AbstractPolicyManager.java

@@ -180,8 +180,6 @@ ModuleEntitlements policyEntitlements(
         );
     }
 
-    final Map<Module, ModuleEntitlements> moduleEntitlementsMap = new ConcurrentHashMap<>();
-
     private final Map<String, List<Entitlement>> serverEntitlements;
     private final List<Entitlement> apmAgentEntitlements;
     private final Map<String, Map<String, List<Entitlement>>> pluginsEntitlements;
@@ -301,9 +299,7 @@ private static Logger getLogger(String componentName, String moduleName) {
      */
     private static final ConcurrentHashMap<String, Logger> MODULE_LOGGERS = new ConcurrentHashMap<>();
 
-    protected ModuleEntitlements getEntitlements(Class<?> requestingClass) {
-        return moduleEntitlementsMap.computeIfAbsent(requestingClass.getModule(), m -> computeEntitlements(requestingClass));
-    }
+    protected abstract ModuleEntitlements getEntitlements(Class<?> requestingClass);
 
     protected final ModuleEntitlements computeEntitlements(Class<?> requestingClass) {
         var policyScope = scopeResolver.apply(requestingClass);

libs/entitlement/src/main/java/org/elasticsearch/entitlement/runtime/policy/PolicyManager.java

@@ -15,12 +15,15 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Function;
 
 /**
  * The real implementation of {@link AbstractPolicyManager} we use in production.
  */
 public final class PolicyManager extends AbstractPolicyManager {
+    final Map<Module, ModuleEntitlements> moduleEntitlementsMap = new ConcurrentHashMap<>();
+
     public PolicyManager(
         Policy serverPolicy,
         List<Entitlement> apmAgentEntitlements,
@@ -31,4 +34,10 @@ public PolicyManager(
     ) {
         super(serverPolicy, apmAgentEntitlements, pluginPolicies, scopeResolver, pluginSourcePaths, pathLookup);
     }
+
+    @Override
+    protected ModuleEntitlements getEntitlements(Class<?> requestingClass) {
+        return moduleEntitlementsMap.computeIfAbsent(requestingClass.getModule(), m -> computeEntitlements(requestingClass));
+    }
+
 }

libs/entitlement/src/test/java/org/elasticsearch/entitlement/runtime/policy/PolicyManagerTests.java

@@ -12,8 +12,8 @@
 import org.elasticsearch.bootstrap.ScopeResolver;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.core.Strings;
-import org.elasticsearch.entitlement.runtime.policy.PolicyManager.ModuleEntitlements;
-import org.elasticsearch.entitlement.runtime.policy.PolicyManager.PolicyScope;
+import org.elasticsearch.entitlement.runtime.policy.AbstractPolicyManager.ModuleEntitlements;
+import org.elasticsearch.entitlement.runtime.policy.AbstractPolicyManager.PolicyScope;
 import org.elasticsearch.entitlement.runtime.policy.agent.TestAgent;
 import org.elasticsearch.entitlement.runtime.policy.agent.inner.TestInnerAgent;
 import org.elasticsearch.entitlement.runtime.policy.entitlements.CreateClassLoaderEntitlement;
@@ -39,7 +39,7 @@
 import java.util.concurrent.atomic.AtomicReference;
 
 import static java.util.Map.entry;
-import static org.elasticsearch.entitlement.runtime.policy.PolicyManager.ComponentKind.SERVER;
+import static org.elasticsearch.entitlement.runtime.policy.AbstractPolicyManager.ComponentKind.SERVER;
 import static org.hamcrest.Matchers.allOf;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.is;

test/framework/src/main/java/org/elasticsearch/entitlement/runtime/policy/TestPolicyManager.java

@@ -68,7 +68,6 @@ public void setTriviallyAllowingTestCode(boolean newValue) {
      * Called between tests so each test is not affected by prior tests
      */
     public void reset() {
-        assert moduleEntitlementsMap.isEmpty() : "We're not supposed to be using moduleEntitlementsMap in tests";
         classEntitlementsMap.clear();
         isActive = false;
         isTriviallyAllowingTestCode = true;