Follow up fixes to RuleCache (#458)

Switch from a global to a per-instance descriptor map now that we're
caching the program. Make a few small improvements to the rule cache.
Update benchmark to better represent intended use of a validator
(creating one instance and re-using it across evaluations).

Author: pkwarren

benchmarks/build.gradle.kts

@@ -1,3 +1,5 @@
+import com.diffplug.gradle.spotless.SpotlessExtension
+
 plugins {
     java
     alias(libs.plugins.jmh)
@@ -10,6 +12,12 @@ java {
     targetCompatibility = JavaVersion.VERSION_21
 }
 
+configure<SpotlessExtension> {
+    java {
+        targetExclude("build/generated/**/*.java")
+    }
+}
+
 val buf: Configuration by configurations.creating
 
 tasks.register("configureBuf") {

benchmarks/src/jmh/java/build/buf/protovalidate/benchmarks/ValidationBenchmark.java

@@ -19,10 +19,8 @@
 import build.buf.protovalidate.benchmarks.gen.ManyUnruledFieldsMessage;
 import build.buf.protovalidate.benchmarks.gen.RepeatedRuleMessage;
 import build.buf.protovalidate.benchmarks.gen.SimpleStringMessage;
-import build.buf.protovalidate.exceptions.CompilationException;
 import build.buf.protovalidate.exceptions.ValidationException;
-import com.google.protobuf.Descriptors.Descriptor;
-import java.util.Collections;
+import com.google.protobuf.Descriptors.FieldDescriptor;
 import java.util.concurrent.TimeUnit;
 import org.openjdk.jmh.annotations.Benchmark;
 import org.openjdk.jmh.annotations.BenchmarkMode;
@@ -41,9 +39,7 @@ public class ValidationBenchmark {
   private Validator validator;
   private SimpleStringMessage simple;
   private ManyUnruledFieldsMessage manyUnruled;
-
-  // Descriptor captured once; cheap to reference during benchmark.
-  private static final Descriptor REPEATED_RULE_DESC = RepeatedRuleMessage.getDescriptor();
+  private RepeatedRuleMessage repeatedRule;
 
   @Setup
   public void setup() throws ValidationException {
@@ -65,13 +61,20 @@ public void setup() throws ValidationException {
             .setF9("v9")
             .build();
 
+    RepeatedRuleMessage.Builder repeatedRuleBuilder = RepeatedRuleMessage.newBuilder();
+    for (FieldDescriptor fd : RepeatedRuleMessage.getDescriptor().getFields()) {
+      repeatedRuleBuilder.setField(fd, "v");
+    }
+    repeatedRule = repeatedRuleBuilder.build();
+
     // Warm evaluator cache for steady-state benchmarks.
     validator.validate(simple);
     validator.validate(manyUnruled);
+    validator.validate(repeatedRule);
   }
 
   // Steady-state validate() benchmarks. These exercise the hot path after the
-  // evaluator cache is warm. PR #451 does not affect this path.
+  // evaluator cache is warm.
 
   @Benchmark
   public void validateSimple(Blackhole bh) throws ValidationException {
@@ -83,19 +86,8 @@ public void validateManyUnruled(Blackhole bh) throws ValidationException {
     bh.consume(validator.validate(manyUnruled));
   }
 
-  // Compile-path benchmark. Measures building a fresh validator and warming
-  // its RuleCache for RepeatedRuleMessage (20 string fields, all min_len).
-  // PR #451 affects exactly this path: without the fix, the AST is rebuilt
-  // for every field; with the fix, fields 2..20 hit the cache.
-  //
-  // Time is dominated by Cel environment construction in newCel(); the #451
-  // signal is the delta on top of that baseline.
   @Benchmark
-  @OutputTimeUnit(TimeUnit.MILLISECONDS)
-  public void compileValidatorForRepeated(Blackhole bh) throws CompilationException {
-    Validator v =
-        ValidatorFactory.newBuilder()
-            .buildWithDescriptors(Collections.singletonList(REPEATED_RULE_DESC), false);
-    bh.consume(v);
+  public void validateRepeatedRule(Blackhole bh) throws ValidationException {
+    bh.consume(validator.validate(repeatedRule));
   }
 }

src/main/java/build/buf/protovalidate/RuleCache.java

@@ -62,11 +62,10 @@ private CelRule(
   }
 
   /**
-   * Concurrent map for caching {@link FieldDescriptor} and their associated List of {@link
-   * AstExpression}.
+   * Compiled rules keyed by rule field descriptor (e.g. {@code StringRules.min_len}), shared across
+   * all user fields that reference the same rule. The rule value is bound per call at eval time.
    */
-  private static final Map<FieldDescriptor, List<CelRule>> descriptorMap =
-      new ConcurrentHashMap<>();
+  private final Map<FieldDescriptor, List<CelRule>> descriptorMap = new ConcurrentHashMap<>();
 
   /** The environment to use for evaluation. */
   private final Cel cel;
@@ -123,14 +122,14 @@ List<CompiledProgram> compile(
     }
     List<CompiledProgram> programs = new ArrayList<>();
     for (CelRule rule : completeProgramList) {
+      Object fieldValue = message.getField(rule.field);
       programs.add(
           new CompiledProgram(
               rule.program,
               rule.astExpression.source,
               rule.rulePath,
-              new ObjectValue(rule.field, message.getField(rule.field)),
-              Variable.newRuleVariable(
-                  message, ProtoAdapter.toCel(rule.field, message.getField(rule.field)))));
+              new ObjectValue(rule.field, fieldValue),
+              Variable.newRuleVariable(message, ProtoAdapter.toCel(rule.field, fieldValue))));
     }
     return Collections.unmodifiableList(programs);
   }
@@ -148,8 +147,31 @@ List<CompiledProgram> compile(
     }
     build.buf.validate.PredefinedRules rules = getFieldRules(ruleFieldDesc);
     if (rules == null) return null;
+    try {
+      return descriptorMap.computeIfAbsent(
+          ruleFieldDesc,
+          key -> {
+            try {
+              return buildCelRules(fieldDescriptor, forItems, setOneof, key, message, rules);
+            } catch (CompilationException e) {
+              throw new UncheckedCompilationException(e);
+            }
+          });
+    } catch (UncheckedCompilationException e) {
+      throw e.getCompilationException();
+    }
+  }
+
+  private List<CelRule> buildCelRules(
+      FieldDescriptor fieldDescriptor,
+      boolean forItems,
+      FieldDescriptor setOneof,
+      FieldDescriptor ruleFieldDesc,
+      Message message,
+      build.buf.validate.PredefinedRules rules)
+      throws CompilationException {
     List<Expression> expressions = Expression.fromRules(rules.getCelList());
-    celRules = new ArrayList<>(expressions.size());
+    List<CelRule> celRules = new ArrayList<>(expressions.size());
     Cel ruleCel = getRuleCel(fieldDescriptor, message, ruleFieldDesc, forItems);
     for (Expression expression : expressions) {
       FieldPath rulePath =
@@ -167,10 +189,22 @@ List<CompiledProgram> compile(
       }
       celRules.add(new CelRule(astExpression, program, ruleFieldDesc, rulePath));
     }
-    descriptorMap.put(ruleFieldDesc, celRules);
     return celRules;
   }
 
+  private static final class UncheckedCompilationException extends RuntimeException {
+    private final CompilationException compilationException;
+
+    UncheckedCompilationException(CompilationException cause) {
+      super(cause);
+      this.compilationException = cause;
+    }
+
+    CompilationException getCompilationException() {
+      return compilationException;
+    }
+  }
+
   private build.buf.validate.@Nullable PredefinedRules getFieldRules(FieldDescriptor ruleFieldDesc)
       throws CompilationException {
     DescriptorProtos.FieldOptions options = ruleFieldDesc.getOptions();
@@ -310,7 +344,7 @@ private ResolvedRule resolveRules(
             DynamicMessage.parseFrom(
                 expectedRuleMessageDescriptor, typeRules.toByteString(), extensionRegistry);
       } catch (InvalidProtocolBufferException e) {
-        throw new RuntimeException(e);
+        throw new CompilationException("failed to reparse rules with extension registry", e);
       }
     }
     if (!allowUnknownFields && !typeRules.getUnknownFields().isEmpty()) {