Split out EntitlementsCache

Author: prdoyle

libs/entitlement/src/main/java/org/elasticsearch/entitlement/initialization/EntitlementInitialization.java

@@ -86,7 +86,8 @@ private static PolicyManager createPolicyManager() {
             EntitlementBootstrap.bootstrapArgs().sourcePaths(),
             ENTITLEMENTS_MODULE,
             pathLookup,
-            bootstrapArgs.suppressFailureLogClasses()
+            bootstrapArgs.suppressFailureLogClasses(),
+            new EntitlementsCacheImpl()
         );
     }
 

libs/entitlement/src/main/java/org/elasticsearch/entitlement/initialization/EntitlementsCacheImpl.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.entitlement.initialization;
+
+import org.elasticsearch.entitlement.runtime.policy.EntitlementsCache;
+
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
+
+import static java.util.Objects.requireNonNull;
+
+/**
+ * The production {@link EntitlementsCache}. (Tests use {@code EntitlementCacheForTesting}.)
+ */
+final class EntitlementsCacheImpl extends ConcurrentHashMap<Module, EntitlementsCache.ModuleEntitlements> implements EntitlementsCache {
+    @Override
+    public ModuleEntitlements computeIfAbsent(Class<?> key, Function<? super Class<?>, ? extends ModuleEntitlements> mappingFunction) {
+        // We cache per module rather than per class to make the cache smaller and increase the hit ratio
+        return computeIfAbsent(key.getModule(), m -> requireNonNull(mappingFunction.apply(key)));
+    }
+}

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

@@ -0,0 +1,76 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.entitlement.runtime.policy;
+
+import org.elasticsearch.entitlement.runtime.policy.entitlements.Entitlement;
+import org.elasticsearch.logging.Logger;
+
+import java.util.List;
+import java.util.Map;
+import java.util.function.Function;
+import java.util.stream.Stream;
+
+/**
+ * Quickly provides entitlement info for a given class. Performance-critical.
+ */
+public interface EntitlementsCache {
+    /**
+     * @return true if entitlement checking should be bypassed for the given class;
+     *         false if the normal built-in "trivially allowed" rules apply.
+     */
+    default boolean isAlwaysAllowed(Class<?> key) {
+        return false;
+    }
+
+    ModuleEntitlements computeIfAbsent(Class<?> key, Function<? super Class<?>, ? extends ModuleEntitlements> mappingFunction);
+
+    /**
+     * This class contains all the entitlements by type, plus the {@link FileAccessTree} for the special case of filesystem entitlements.
+     * <p>
+     * We use layers when computing {@link ModuleEntitlements}; first, we check whether the module we are building it for is in the
+     * server layer ({@link PolicyManager#SERVER_LAYER_MODULES}) (*).
+     * If it is, we use the server policy, using the same caller class module name as the scope, and read the entitlements for that scope.
+     * Otherwise, we use the {@code PluginResolver} to identify the correct plugin layer and find the policy for it (if any).
+     * If the plugin is modular, we again use the same caller class module name as the scope, and read the entitlements for that scope.
+     * If it's not, we use the single {@code ALL-UNNAMED} scope – in this case there is one scope and all entitlements apply
+     * to all the plugin code.
+     * </p>
+     * <p>
+     * (*) implementation detail: this is currently done in an indirect way: we know the module is not in the system layer
+     * (otherwise the check would have been already trivially allowed), so we just check that the module is named, and it belongs to the
+     * boot {@link ModuleLayer}. We might want to change this in the future to make it more consistent/easier to maintain.
+     * </p>
+     *
+     * @param componentName the plugin name or else one of the special component names like "(server)".
+     */
+    record ModuleEntitlements(
+        String componentName,
+        Map<Class<? extends Entitlement>, List<Entitlement>> entitlementsByType,
+        FileAccessTree fileAccess,
+        Logger logger
+    ) {
+
+        public ModuleEntitlements {
+            entitlementsByType = Map.copyOf(entitlementsByType);
+        }
+
+        public boolean hasEntitlement(Class<? extends Entitlement> entitlementClass) {
+            return entitlementsByType.containsKey(entitlementClass);
+        }
+
+        public <E extends Entitlement> Stream<E> getEntitlements(Class<E> entitlementClass) {
+            var entitlements = entitlementsByType.get(entitlementClass);
+            if (entitlements == null) {
+                return Stream.empty();
+            }
+            return entitlements.stream().map(entitlementClass::cast);
+        }
+    }
+}

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

@@ -14,6 +14,7 @@
 import org.elasticsearch.core.SuppressForbidden;
 import org.elasticsearch.entitlement.instrumentation.InstrumentationService;
 import org.elasticsearch.entitlement.runtime.api.NotEntitledException;
+import org.elasticsearch.entitlement.runtime.policy.EntitlementsCache.ModuleEntitlements;
 import org.elasticsearch.entitlement.runtime.policy.FileAccessTree.ExclusiveFileEntitlement;
 import org.elasticsearch.entitlement.runtime.policy.FileAccessTree.ExclusivePath;
 import org.elasticsearch.entitlement.runtime.policy.entitlements.CreateClassLoaderEntitlement;
@@ -121,7 +122,7 @@
  * All these methods start in the same way: the components identified in the previous section are used to establish if and how to check:
  * If the caller class belongs to {@link PolicyManager#SYSTEM_LAYER_MODULES}, no check is performed (the call is trivially allowed, see
  * {@link PolicyManager#isTriviallyAllowed}).
- * Otherwise, we lazily compute and create a {@link PolicyManager.ModuleEntitlements} record (see
+ * Otherwise, we lazily compute and create a {@link ModuleEntitlements} record (see
  * {@link PolicyManager#computeEntitlements}). The record is cached so it can be used in following checks, stored in a
  * {@code Module -> ModuleEntitlement} map.
  * </p>
@@ -181,49 +182,6 @@ public enum ComponentKind {
         }
     }
 
-    /**
-     * This class contains all the entitlements by type, plus the {@link FileAccessTree} for the special case of filesystem entitlements.
-     * <p>
-     * We use layers when computing {@link ModuleEntitlements}; first, we check whether the module we are building it for is in the
-     * server layer ({@link PolicyManager#SERVER_LAYER_MODULES}) (*).
-     * If it is, we use the server policy, using the same caller class module name as the scope, and read the entitlements for that scope.
-     * Otherwise, we use the {@code PluginResolver} to identify the correct plugin layer and find the policy for it (if any).
-     * If the plugin is modular, we again use the same caller class module name as the scope, and read the entitlements for that scope.
-     * If it's not, we use the single {@code ALL-UNNAMED} scope – in this case there is one scope and all entitlements apply
-     * to all the plugin code.
-     * </p>
-     * <p>
-     * (*) implementation detail: this is currently done in an indirect way: we know the module is not in the system layer
-     * (otherwise the check would have been already trivially allowed), so we just check that the module is named, and it belongs to the
-     * boot {@link ModuleLayer}. We might want to change this in the future to make it more consistent/easier to maintain.
-     * </p>
-     *
-     * @param componentName the plugin name or else one of the special component names like "(server)".
-     */
-    record ModuleEntitlements(
-        String componentName,
-        Map<Class<? extends Entitlement>, List<Entitlement>> entitlementsByType,
-        FileAccessTree fileAccess,
-        Logger logger
-    ) {
-
-        ModuleEntitlements {
-            entitlementsByType = Map.copyOf(entitlementsByType);
-        }
-
-        public boolean hasEntitlement(Class<? extends Entitlement> entitlementClass) {
-            return entitlementsByType.containsKey(entitlementClass);
-        }
-
-        public <E extends Entitlement> Stream<E> getEntitlements(Class<E> entitlementClass) {
-            var entitlements = entitlementsByType.get(entitlementClass);
-            if (entitlements == null) {
-                return Stream.empty();
-            }
-            return entitlements.stream().map(entitlementClass::cast);
-        }
-    }
-
     private FileAccessTree getDefaultFileAccess(Path componentPath) {
         return FileAccessTree.withoutExclusivePaths(FilesEntitlement.EMPTY, pathLookup, componentPath);
     }
@@ -249,8 +207,6 @@ ModuleEntitlements policyEntitlements(String componentName, Path componentPath,
         );
     }
 
-    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;
@@ -303,6 +259,8 @@ private static Set<Module> findSystemLayerModules() {
      */
     private final List<ExclusivePath> exclusivePaths;
 
+    final EntitlementsCache moduleEntitlementsCache;
+
     public PolicyManager(
         Policy serverPolicy,
         List<Entitlement> apmAgentEntitlements,
@@ -311,7 +269,8 @@ public PolicyManager(
         Map<String, Path> sourcePaths,
         Module entitlementsModule,
         PathLookup pathLookup,
-        Set<Class<?>> suppressFailureLogClasses
+        Set<Class<?>> suppressFailureLogClasses,
+        EntitlementsCache moduleEntitlementsCache
     ) {
         this.serverEntitlements = buildScopeEntitlementsMap(requireNonNull(serverPolicy));
         this.apmAgentEntitlements = apmAgentEntitlements;
@@ -323,6 +282,7 @@ public PolicyManager(
         this.entitlementsModule = entitlementsModule;
         this.pathLookup = requireNonNull(pathLookup);
         this.mutedClasses = suppressFailureLogClasses;
+        this.moduleEntitlementsCache = moduleEntitlementsCache;
 
         List<ExclusiveFileEntitlement> exclusiveFileEntitlements = new ArrayList<>();
         for (var e : serverEntitlements.entrySet()) {
@@ -723,7 +683,7 @@ private void checkEntitlementPresent(Class<?> callerClass, Class<? extends Entit
     }
 
     ModuleEntitlements getEntitlements(Class<?> requestingClass) {
-        return moduleEntitlementsMap.computeIfAbsent(requestingClass.getModule(), m -> computeEntitlements(requestingClass));
+        return moduleEntitlementsCache.computeIfAbsent(requestingClass, this::computeEntitlements);
     }
 
     private ModuleEntitlements computeEntitlements(Class<?> requestingClass) {
@@ -827,7 +787,7 @@ Optional<StackFrame> findRequestingFrame(Stream<StackFrame> frames) {
     /**
      * @return true if permission is granted regardless of the entitlement
      */
-    private static boolean isTriviallyAllowed(Class<?> requestingClass) {
+    private boolean isTriviallyAllowed(Class<?> requestingClass) {
         if (generalLogger.isTraceEnabled()) {
             generalLogger.trace("Stack trace for upcoming trivially-allowed check", new Exception());
         }
@@ -843,6 +803,10 @@ private static boolean isTriviallyAllowed(Class<?> requestingClass) {
             generalLogger.debug("Entitlement trivially allowed from system module [{}]", requestingClass.getModule().getName());
             return true;
         }
+        if (moduleEntitlementsCache.isAlwaysAllowed(requestingClass)) {
+            generalLogger.debug("Entitlement trivially allowed by the EntitlementsCache");
+            return true;
+        }
         generalLogger.trace("Entitlement not trivially allowed");
         return false;
     }

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

@@ -0,0 +1,24 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.entitlement.runtime.policy;
+
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * When testing, we don't use modules, so we cache per-class.
+ */
+public class EntitlementsCacheForTesting extends ConcurrentHashMap<Class<?>, EntitlementsCache.ModuleEntitlements>
+    implements
+        EntitlementsCache {
+    @Override
+    public boolean isAlwaysAllowed(Class<?> key) {
+        return key.getPackageName().startsWith("org.gradle") || key.getPackageName().startsWith("org.junit");
+    }
+}