From 4f65030e0a46b025894ad77bd827f7b75fb3740e Mon Sep 17 00:00:00 2001 From: Joaquin Date: Wed, 4 Mar 2026 10:57:11 -0300 Subject: [PATCH] feat: implement LRU cache for JSON schema validation to prevent memory leaks --- .../jackson2/DefaultJsonSchemaValidator.java | 42 +++++++++++++++---- .../jackson3/DefaultJsonSchemaValidator.java | 42 +++++++++++++++---- 2 files changed, 68 insertions(+), 16 deletions(-) diff --git a/mcp-json-jackson2/src/main/java/io/modelcontextprotocol/json/schema/jackson2/DefaultJsonSchemaValidator.java b/mcp-json-jackson2/src/main/java/io/modelcontextprotocol/json/schema/jackson2/DefaultJsonSchemaValidator.java index e07bf1759..2d998057c 100644 --- a/mcp-json-jackson2/src/main/java/io/modelcontextprotocol/json/schema/jackson2/DefaultJsonSchemaValidator.java +++ b/mcp-json-jackson2/src/main/java/io/modelcontextprotocol/json/schema/jackson2/DefaultJsonSchemaValidator.java @@ -3,9 +3,10 @@ */ package io.modelcontextprotocol.json.schema.jackson2; +import java.util.Collections; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; -import java.util.concurrent.ConcurrentHashMap; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -31,21 +32,40 @@ public class DefaultJsonSchemaValidator implements JsonSchemaValidator { private static final Logger logger = LoggerFactory.getLogger(DefaultJsonSchemaValidator.class); + /** + * Default maximum number of cached JSON schemas. + */ + public static final int DEFAULT_MAX_CACHE_SIZE = 1024; + private final ObjectMapper objectMapper; private final SchemaRegistry schemaFactory; - // TODO: Implement a strategy to purge the cache (TTL, size limit, etc.) - private final ConcurrentHashMap schemaCache; + private final Map schemaCache; public DefaultJsonSchemaValidator() { - this(new ObjectMapper()); + this(new ObjectMapper(), DEFAULT_MAX_CACHE_SIZE); } public DefaultJsonSchemaValidator(ObjectMapper objectMapper) { + this(objectMapper, DEFAULT_MAX_CACHE_SIZE); + } + + /** + * Creates a new {@link DefaultJsonSchemaValidator} with the given {@link ObjectMapper} + * and maximum cache size. + * @param objectMapper the object mapper to use for JSON processing + * @param maxCacheSize the maximum number of schemas to cache (LRU) + */ + public DefaultJsonSchemaValidator(ObjectMapper objectMapper, int maxCacheSize) { this.objectMapper = objectMapper; this.schemaFactory = SchemaRegistry.withDialect(Dialects.getDraft202012()); - this.schemaCache = new ConcurrentHashMap<>(); + this.schemaCache = Collections.synchronizedMap(new LinkedHashMap(maxCacheSize, 0.75f, true) { + @Override + protected boolean removeEldestEntry(Map.Entry eldest) { + return size() > maxCacheSize; + } + }); } @Override @@ -140,9 +160,15 @@ protected String generateCacheKey(Map schema) { // Use the (optional) "$id" field as the cache key if present return "" + schema.get("$id"); } - // Fall back to schema's hash code as a simple cache key - // For more sophisticated caching, could use content-based hashing - return String.valueOf(schema.hashCode()); + try { + // Use the stable JSON representation as the cache key to avoid hash + // collisions and map order issues + return this.objectMapper.writeValueAsString(schema); + } + catch (JsonProcessingException e) { + // Fall back to schema's hash code if serialization fails + return String.valueOf(schema.hashCode()); + } } /** diff --git a/mcp-json-jackson3/src/main/java/io/modelcontextprotocol/json/schema/jackson3/DefaultJsonSchemaValidator.java b/mcp-json-jackson3/src/main/java/io/modelcontextprotocol/json/schema/jackson3/DefaultJsonSchemaValidator.java index 8c9b7ccdb..a1d686f86 100644 --- a/mcp-json-jackson3/src/main/java/io/modelcontextprotocol/json/schema/jackson3/DefaultJsonSchemaValidator.java +++ b/mcp-json-jackson3/src/main/java/io/modelcontextprotocol/json/schema/jackson3/DefaultJsonSchemaValidator.java @@ -3,9 +3,10 @@ */ package io.modelcontextprotocol.json.schema.jackson3; +import java.util.Collections; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; -import java.util.concurrent.ConcurrentHashMap; import com.networknt.schema.Schema; import com.networknt.schema.SchemaRegistry; @@ -30,21 +31,40 @@ public class DefaultJsonSchemaValidator implements JsonSchemaValidator { private static final Logger logger = LoggerFactory.getLogger(DefaultJsonSchemaValidator.class); + /** + * Default maximum number of cached JSON schemas. + */ + public static final int DEFAULT_MAX_CACHE_SIZE = 1024; + private final JsonMapper jsonMapper; private final SchemaRegistry schemaFactory; - // TODO: Implement a strategy to purge the cache (TTL, size limit, etc.) - private final ConcurrentHashMap schemaCache; + private final Map schemaCache; public DefaultJsonSchemaValidator() { - this(JsonMapper.shared()); + this(JsonMapper.shared(), DEFAULT_MAX_CACHE_SIZE); } public DefaultJsonSchemaValidator(JsonMapper jsonMapper) { + this(jsonMapper, DEFAULT_MAX_CACHE_SIZE); + } + + /** + * Creates a new {@link DefaultJsonSchemaValidator} with the given {@link JsonMapper} + * and maximum cache size. + * @param jsonMapper the JSON mapper to use for JSON processing + * @param maxCacheSize the maximum number of schemas to cache (LRU) + */ + public DefaultJsonSchemaValidator(JsonMapper jsonMapper, int maxCacheSize) { this.jsonMapper = jsonMapper; this.schemaFactory = SchemaRegistry.withDialect(Dialects.getDraft202012()); - this.schemaCache = new ConcurrentHashMap<>(); + this.schemaCache = Collections.synchronizedMap(new LinkedHashMap(maxCacheSize, 0.75f, true) { + @Override + protected boolean removeEldestEntry(Map.Entry eldest) { + return size() > maxCacheSize; + } + }); } @Override @@ -139,9 +159,15 @@ protected String generateCacheKey(Map schema) { // Use the (optional) "$id" field as the cache key if present return "" + schema.get("$id"); } - // Fall back to schema's hash code as a simple cache key - // For more sophisticated caching, could use content-based hashing - return String.valueOf(schema.hashCode()); + try { + // Use the stable JSON representation as the cache key to avoid hash + // collisions and map order issues + return this.jsonMapper.writeValueAsString(schema); + } + catch (JacksonException e) { + // Fall back to schema's hash code if serialization fails + return String.valueOf(schema.hashCode()); + } } /**