perf(core): cache generated JSON schemas to remove lock from hot path

agentscope-core/src/main/java/io/agentscope/core/util/JsonSchemaUtils.java

@@ -27,8 +27,11 @@
 import com.github.victools.jsonschema.module.jackson.JacksonModule;
 import com.github.victools.jsonschema.module.jackson.JacksonOption;
 import io.agentscope.core.tool.ToolSchemaModule;
+import java.lang.reflect.ParameterizedType;
 import java.lang.reflect.Type;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * Utility class for JSON Schema operations.
@@ -53,9 +56,17 @@
  *   <li>{@code @JsonClassDescription(...)} - add class description</li>
  * </ul>
  *
- * <p>All public methods are thread-safe. Schema generation through the shared victools
- * {@code SchemaGenerator} is serialized by an internal lock, because the generator itself
- * is not designed for concurrent use.</p>
+ * <p>All public methods are thread-safe. Generated schemas are cached per {@link Class} /
+ * {@link Type}; the shared victools {@code SchemaGenerator} is not designed for concurrent use,
+ * so the internal lock is taken only on a cache miss (the first time a given class or type is
+ * seen). A cache hit converts a fresh, independently mutable {@code Map} from the cached,
+ * never-mutated {@link JsonNode}, so it needs no lock.</p>
+ *
+ * <p>Cache entries are scoped to the class they describe instead of living in a static map keyed
+ * by {@code Class}, so an entry cannot outlive that class or pin the classloader that defined it.
+ * That matters because the structured-output and tool-parameter classes reaching this utility are
+ * not always compile-time-fixed: extensions can load skills and tools at runtime, and
+ * multi-tenant deployments may load classes per tenant.</p>
  *
  * @hidden
  */
@@ -68,10 +79,48 @@ public class JsonSchemaUtils {
     /**
      * Guards the shared victools {@link SchemaGenerator}, which is not thread-safe: its
      * JacksonModule keeps an unsynchronized introspection cache, so concurrent schema
-     * generation must be serialized.
+     * generation must be serialized. Only cache misses in {@link #CLASS_SCHEMA_SLOT} and
+     * {@link #TYPE_SCHEMA_SLOT} take this lock; cache hits never do.
      */
     private static final Object SCHEMA_LOCK = new Object();
 
+    /**
+     * Schema cache slot of each class. A schema is a deterministic function of the class and the
+     * static, never-changing generator config, so entries never need invalidation. Cached nodes
+     * are never mutated after being stored: every call still converts a fresh, independently
+     * mutable {@link Map} from the cached node, so callers that mutate the returned map (e.g.
+     * {@code ToolSchemaGenerator}) cannot corrupt the cache or interfere with one another.
+     *
+     * <p>Creating the slot through {@link ClassValue} holds it on the class it describes, rather
+     * than in a static map that strongly references the class as a key, so a slot cannot keep that
+     * class — or the classloader which defined it — reachable once the rest of the application has
+     * let go of them.
+     */
+    private static final ClassValue<AtomicReference<JsonNode>> CLASS_SCHEMA_SLOT =
+            new ClassValue<>() {
+                @Override
+                protected AtomicReference<JsonNode> computeValue(Class<?> clazz) {
+                    return new AtomicReference<>();
+                }
+            };
+
+    /**
+     * Same caching strategy as {@link #CLASS_SCHEMA_SLOT}, keyed by generic {@link Type} to support
+     * parameterized structured-output and tool-parameter types. The variants of one raw class (e.g.
+     * {@code List<String>} versus {@code List<Integer>}) share the map held on that raw class, so
+     * these slots are scoped to a classloader in the same way.
+     *
+     * <p>Because unrelated raw classes never share a map, a miss takes {@link #SCHEMA_LOCK} without
+     * grouping unrelated types behind the same lock.
+     */
+    private static final ClassValue<Map<Type, JsonNode>> TYPE_SCHEMA_SLOT =
+            new ClassValue<>() {
+                @Override
+                protected Map<Type, JsonNode> computeValue(Class<?> rawType) {
+                    return new ConcurrentHashMap<>();
+                }
+            };
+
     static {
         // JacksonModule to support @JsonProperty, @JsonPropertyDescription annotations
         JacksonModule jacksonModule =
@@ -106,10 +155,7 @@ public class JsonSchemaUtils {
      */
     public static Map<String, Object> generateSchemaFromClass(Class<?> clazz) {
         try {
-            JsonNode schemaNode;
-            synchronized (SCHEMA_LOCK) {
-                schemaNode = schemaGenerator.generateSchema(clazz);
-            }
+            JsonNode schemaNode = cachedSchemaNode(clazz);
             return JsonUtils.getJsonCodec()
                     .convertValue(schemaNode, new TypeReference<Map<String, Object>>() {});
         } catch (Exception e) {
@@ -144,10 +190,7 @@ public static Map<String, Object> generateSchemaFromJsonNode(JsonNode schema) {
      */
     public static Map<String, Object> generateSchemaFromType(Type type) {
         try {
-            JsonNode schemaNode;
-            synchronized (SCHEMA_LOCK) {
-                schemaNode = schemaGenerator.generateSchema(type);
-            }
+            JsonNode schemaNode = cachedSchemaNode(type);
             return JsonUtils.getJsonCodec()
                     .convertValue(schemaNode, new TypeReference<Map<String, Object>>() {});
         } catch (Exception e) {
@@ -156,6 +199,84 @@ public static Map<String, Object> generateSchemaFromType(Type type) {
         }
     }
 
+    /**
+     * Returns the cached schema node for a class, generating it under {@link #SCHEMA_LOCK} on a
+     * miss. The slot is re-read inside the lock so that a class several threads reach at once is
+     * still generated exactly once.
+     */
+    private static JsonNode cachedSchemaNode(Class<?> clazz) {
+        AtomicReference<JsonNode> slot = CLASS_SCHEMA_SLOT.get(clazz);
+        JsonNode cached = slot.get();
+        if (cached != null) {
+            return cached;
+        }
+
+        synchronized (SCHEMA_LOCK) {
+            cached = slot.get();
+            if (cached == null) {
+                cached = schemaGenerator.generateSchema(clazz);
+                slot.set(cached);
+            }
+        }
+        return cached;
+    }
+
+    /**
+     * Returns the cached schema node for a type, generating it under {@link #SCHEMA_LOCK} on a
+     * miss.
+     *
+     * <p>The miss path re-reads the slot inside the lock rather than calling {@code
+     * ConcurrentHashMap#computeIfAbsent}, which runs its mapping function while holding the map's
+     * bin lock: nesting the schema lock inside it would block unrelated types hashing to the same
+     * bin for the whole generation, and would break the contract that a mapping function must not
+     * modify its own map.
+     */
+    private static JsonNode cachedSchemaNode(Type type) {
+        Class<?> rawType = rawTypeOf(type);
+        if (rawType == null) {
+            // Type variables, wildcards and generic arrays have no raw class to hang a slot on;
+            // generate them without caching.
+            synchronized (SCHEMA_LOCK) {
+                return schemaGenerator.generateSchema(type);
+            }
+        }
+
+        Map<Type, JsonNode> slot = TYPE_SCHEMA_SLOT.get(rawType);
+        JsonNode cached = slot.get(type);
+        if (cached != null) {
+            return cached;
+        }
+
+        synchronized (SCHEMA_LOCK) {
+            cached = slot.get(type);
+            if (cached == null) {
+                cached = schemaGenerator.generateSchema(type);
+                slot.put(type, cached);
+            }
+        }
+        return cached;
+    }
+
+    /**
+     * Returns the raw class whose cache a type's schema belongs to, or {@code null} when the type
+     * has no raw class and therefore cannot be cached.
+     *
+     * @throws NullPointerException if the type is null
+     */
+    private static Class<?> rawTypeOf(Type type) {
+        if (type == null) {
+            throw new NullPointerException("type must not be null");
+        }
+        if (type instanceof Class<?> clazz) {
+            return clazz;
+        }
+        if (type instanceof ParameterizedType parameterizedType
+                && parameterizedType.getRawType() instanceof Class<?> rawClass) {
+            return rawClass;
+        }
+        return null;
+    }
+
     /**
      * Convert Map to typed object.
      *

agentscope-core/src/test/java/io/agentscope/core/util/JsonSchemaUtilsTest.java

@@ -17,6 +17,7 @@
 package io.agentscope.core.util;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -168,6 +169,120 @@ void testGenerateSchemaFromType() {
         assertEquals("object", mapSchema.get("type"));
     }
 
+    @Test
+    void testGenerateSchemaFromClassRepeatedCallsReturnEqualIndependentMaps() {
+        Map<String, Object> first = JsonSchemaUtils.generateSchemaFromClass(SimpleModel.class);
+        Map<String, Object> second = JsonSchemaUtils.generateSchemaFromClass(SimpleModel.class);
+
+        // A repeated class must yield an equal schema, so caching cannot change the result.
+        assertEquals(first, second);
+
+        // Each call must return a fresh, independently mutable map: mutating one must not leak
+        // into another, matching in-place-mutating callers such as ToolSchemaGenerator.
+        first.put("description", "mutated");
+        assertFalse(second.containsKey("description"));
+
+        Map<String, Object> third = JsonSchemaUtils.generateSchemaFromClass(SimpleModel.class);
+        assertFalse(third.containsKey("description"));
+        assertEquals(second, third);
+    }
+
+    @Test
+    void testGenerateSchemaFromTypeRepeatedCallsReturnEqualIndependentMaps() {
+        Type listType = new TypeReference<List<String>>() {}.getType();
+
+        Map<String, Object> first = JsonSchemaUtils.generateSchemaFromType(listType);
+        Map<String, Object> second = JsonSchemaUtils.generateSchemaFromType(listType);
+
+        assertEquals(first, second);
+
+        first.put("description", "mutated");
+        assertFalse(second.containsKey("description"));
+
+        Map<String, Object> third = JsonSchemaUtils.generateSchemaFromType(listType);
+        assertFalse(third.containsKey("description"));
+        assertEquals(second, third);
+    }
+
+    @Test
+    void testGenerateSchemaFromClassNestedMutationDoesNotAffectLaterCalls() {
+        Map<String, Object> first = JsonSchemaUtils.generateSchemaFromClass(NestedModel.class);
+
+        @SuppressWarnings("unchecked")
+        Map<String, Object> firstProperties = (Map<String, Object>) first.get("properties");
+        assertNotNull(firstProperties);
+
+        // Real callers mutate below the top level: ToolSchemaGenerator hoists "$defs" out of
+        // nested schemas and ReActAgent rewrites nested properties in place. A later call must
+        // still observe the pristine schema, which is exactly the deep-copy invariant the cache
+        // relies on.
+        assertNotNull(firstProperties.remove("tags"));
+        @SuppressWarnings("unchecked")
+        Map<String, Object> firstAuthor = (Map<String, Object>) firstProperties.get("author");
+        assertNotNull(firstAuthor);
+        firstAuthor.put("description", "mutated");
+
+        Map<String, Object> second = JsonSchemaUtils.generateSchemaFromClass(NestedModel.class);
+
+        @SuppressWarnings("unchecked")
+        Map<String, Object> secondProperties = (Map<String, Object>) second.get("properties");
+        assertNotNull(secondProperties);
+        assertTrue(secondProperties.containsKey("tags"));
+
+        @SuppressWarnings("unchecked")
+        Map<String, Object> secondAuthor = (Map<String, Object>) secondProperties.get("author");
+        assertNotNull(secondAuthor);
+        assertFalse(secondAuthor.containsKey("description"));
+    }
+
+    @Test
+    void testGenerateSchemaFromTypeNestedMutationDoesNotAffectLaterCalls() {
+        Type listType = new TypeReference<List<SimpleModel>>() {}.getType();
+
+        Map<String, Object> first = JsonSchemaUtils.generateSchemaFromType(listType);
+
+        @SuppressWarnings("unchecked")
+        Map<String, Object> firstItems = (Map<String, Object>) first.get("items");
+        assertNotNull(firstItems);
+        firstItems.put("description", "mutated");
+
+        Map<String, Object> second = JsonSchemaUtils.generateSchemaFromType(listType);
+
+        @SuppressWarnings("unchecked")
+        Map<String, Object> secondItems = (Map<String, Object>) second.get("items");
+        assertNotNull(secondItems);
+        assertFalse(secondItems.containsKey("description"));
+    }
+
+    @Test
+    void testGenerateSchemaFromClassNullThrows() {
+        // Caching routes a null class through ClassValue#get, which rejects null keys; the
+        // resulting NPE must match the pre-cache behavior for a null argument.
+        assertThrows(
+                NullPointerException.class, () -> JsonSchemaUtils.generateSchemaFromClass(null));
+    }
+
+    @Test
+    void testGenerateSchemaFromTypeNullThrows() {
+        // A null type has no raw class to cache under, so the NPE surfaces from the raw-class
+        // lookup, matching the pre-cache behavior for a null argument.
+        assertThrows(
+                NullPointerException.class, () -> JsonSchemaUtils.generateSchemaFromType(null));
+    }
+
+    @Test
+    void testGenerateSchemaFromTypeVariableStillGenerates() {
+        // A type variable has no raw class to hang an entry on, so it is generated without
+        // caching; that fallback must still yield a usable scheme.
+        Type typeVariable = List.class.getTypeParameters()[0];
+
+        Map<String, Object> first = JsonSchemaUtils.generateSchemaFromType(typeVariable);
+        Map<String, Object> second = JsonSchemaUtils.generateSchemaFromType(typeVariable);
+
+        assertNotNull(first);
+        assertEquals(first, second);
+    }
+
     static class ConcurrentClassA {
         public String name;
         public int age;