Add memory limit (close #267, AI assisted)

Change-Id: Iff996ab4b2de8de9084b996b229c7db6931fe8a7
diff --git a/Changes b/Changes
index a5366f3..39d19c0 100644
--- a/Changes
+++ b/Changes
@@ -1,3 +1,7 @@
+0.65.2 2026-07-08
+    - Add memory limit (#267, margaretha, AI assisted) 
+
+
 0.65.1 2026-06-16
     - [feature] Prevent indexation of documents without
       a token stream (diewald; AI-assisted Claude Opus 4.6)
@@ -37,6 +41,8 @@
     - [bugfix] Fix negation in segment queries following De Morgan's law
       (diewald; fixes #93; diewald; AI-assisted Claude Opus 4.6)
 
+
+
 0.64.6 2026-03-09
     - [performance] Add leaf cache. (diewald)
     - [bugfix] Fix fingerprinter (wasn't threadsafe; diewald)
diff --git a/pom.xml b/pom.xml
index 6e95b7d..97c7aa6 100644
--- a/pom.xml
+++ b/pom.xml
@@ -50,7 +50,7 @@
 
   <groupId>de.ids-mannheim.korap.krill</groupId>
   <artifactId>Krill</artifactId>
-  <version>0.65.1</version>
+  <version>0.65.2</version>
   <packaging>jar</packaging>
 
   <name>Krill</name>
diff --git a/src/main/java/de/ids_mannheim/korap/KrillIndex.java b/src/main/java/de/ids_mannheim/korap/KrillIndex.java
index 892bc05..ec45e38 100644
--- a/src/main/java/de/ids_mannheim/korap/KrillIndex.java
+++ b/src/main/java/de/ids_mannheim/korap/KrillIndex.java
@@ -347,6 +347,22 @@
     };
     
 
+    /** AI generated
+     * 
+     * Check whether JVM heap usage exceeds the configured memory limit.
+     * Returns <tt>false</tt> when the limit is 0 (disabled).
+     *
+     * @param limitMB Maximum allowed heap usage in megabytes (0 = disabled)
+     * @return <tt>true</tt> if heap usage exceeds the limit
+     */
+    private static boolean isMemoryExceeded (long limitMB) {
+        if (limitMB <= 0)
+            return false;
+        Runtime rt = Runtime.getRuntime();
+        long usedMB = (rt.totalMemory() - rt.freeMemory()) / (1024L * 1024L);
+        return usedMB > limitMB;
+    };
+
     // Open index reader
     private void openWriter () {
         if (writerOpen) {
@@ -1527,7 +1543,7 @@
         final TimeOutThread tthread = new TimeOutThread();
         tthread.start();
         final long timeout = meta.getTimeOut();
-        boolean isTimeout = false;
+        boolean stopSearching = false;
        
         // See: http://www.ibm.com/developerworks/java/library/j-benchmark1/index.html
         long t1 = System.nanoTime();
@@ -1563,7 +1579,7 @@
 
                 int oldLocalDocID = -1;
 
-                if (isTimeout)
+                if (stopSearching)
                     break;
 
                 SearchCacheKey finalCacheKey = new SearchCacheKey(prelim, atomic.reader().getCombinedCoreAndDeletesKey().toString());
@@ -1662,7 +1678,14 @@
                     // Timeout!
                     if (tthread.getTime() > timeout) {
                         kr.setTimeExceeded(true);
-                        isTimeout=true;
+                        stopSearching=true;
+                        break;
+                    };
+
+                    // Memory limit!
+                    if (isMemoryExceeded(KrillProperties.maxMemoryMB)) {
+                        kr.setMemoryExceeded(true);
+                        stopSearching = true;
                         break;
                     };
 
@@ -1771,7 +1794,7 @@
                 };
 
                 // Can be disabled TEMPORARILY
-                while (!cutoff && !isTimeout && spans.next()) {
+                while (!cutoff && !stopSearching && spans.next()) {
 
                     // TODO: Deprecated
                     if (limit > 0 && i >= limit)
@@ -1780,7 +1803,14 @@
                     // Timeout!
                     if (tthread.getTime() > timeout) {
                         kr.setTimeExceeded(true);
-                        isTimeout=true;
+                        stopSearching=true;
+                        break;
+                    };
+
+                    // Memory limit!
+                    if (isMemoryExceeded(KrillProperties.maxMemoryMB)) {
+                        kr.setMemoryExceeded(true);
+                        stopSearching = true;
                         break;
                     };
 
@@ -1823,7 +1853,7 @@
                     i++;
                 };
 
-                if (!isTimeout && !cutoff) {
+                if (!stopSearching && !cutoff) {
                     if (foundCache == null) {
                         searchCache.put(
                             finalCacheKey,
diff --git a/src/main/java/de/ids_mannheim/korap/response/Response.java b/src/main/java/de/ids_mannheim/korap/response/Response.java
index fa580d0..67d6ff0 100644
--- a/src/main/java/de/ids_mannheim/korap/response/Response.java
+++ b/src/main/java/de/ids_mannheim/korap/response/Response.java
@@ -52,6 +52,7 @@
             totalResults = -2; // Not set
     private String benchmark;
     private boolean timeExceeded = false;
+    private boolean memoryExceeded = false;
 
     private HashMap<String, ObjectNode> jsonFields;
 
@@ -182,6 +183,26 @@
     };
 
 
+    /** AI generated
+     * 
+     * Set to <tt>true</tt> if the memory limit is exceeded.
+     * 
+     * <p>
+     * Will add a warning (683) to the output.
+     * 
+     * @param exceeded
+     *            Either <tt>true</tt> or <tt>false</tt>,
+     *            in case the memory limit was exceeded
+     * @return Response object for chaining
+     */
+    public Response setMemoryExceeded (boolean exceeded) {
+        if (exceeded)
+            this.addWarning(684, "Memory limit exceeded");
+        this.memoryExceeded = exceeded;
+        return this;
+    };
+
+
     /**
      * Get the benchmark time as a string.
      * 
@@ -524,6 +545,9 @@
         if (this.timeExceeded)
             meta.put("timeExceeded", true);
 
+        if (this.memoryExceeded)
+            meta.put("memoryExceeded", true);
+
         if (this.getNode() != null)
             meta.put("node", this.getNode());
 
diff --git a/src/main/java/de/ids_mannheim/korap/util/KrillProperties.java b/src/main/java/de/ids_mannheim/korap/util/KrillProperties.java
index d88e26b..ba61567 100644
--- a/src/main/java/de/ids_mannheim/korap/util/KrillProperties.java
+++ b/src/main/java/de/ids_mannheim/korap/util/KrillProperties.java
@@ -29,6 +29,13 @@
     public static int kwicMaxToken = -1;
     public static int defaultSearchContextLength = 6;
     public static int maxTextSize = DEFAULT_MAX_STRING_LEN; // Default max text size
+
+    /** AI generated
+     * 
+     * Maximum JVM heap memory used (in MB) before a search is aborted.
+     * 0 means the check is disabled.
+     */
+    public static long maxMemoryMB = 0;
     
     public static boolean matchExpansionIncludeContextSize = false;
     
@@ -108,6 +115,7 @@
         String maxCharContextSize = prop.getProperty("krill.context.max.char");
         String defaultSearchContextLength = prop.getProperty("krill.search.context.default");
         String maxTextSizeValue = prop.getProperty("krill.index.textSize.max");
+        String maxMemoryMBValue = prop.getProperty("krill.search.memory.max");
 
         try {
             if (maxTokenMatchSize != null) {
@@ -136,7 +144,15 @@
                 } else {
                     KrillProperties.maxTextSize = userMaxTextLength;
                 }
-
+            }
+            if (maxMemoryMBValue != null) {
+                long parsedMemoryMB = Long.parseLong(maxMemoryMBValue.trim());
+                if (parsedMemoryMB < 0) {
+                    log.warn("krill.search.memory.max must be >= 0. Memory limit check disabled.");
+                    KrillProperties.maxMemoryMB = 0;
+                } else {
+                    KrillProperties.maxMemoryMB = parsedMemoryMB;
+                }
             }
             if (leftContextMaxShrink != null) {
                 if (leftContextMaxShrink.equals("max")) {
diff --git a/src/main/java/de/ids_mannheim/korap/util/StatusCodes.java b/src/main/java/de/ids_mannheim/korap/util/StatusCodes.java
index 33b330a..28e0108 100644
--- a/src/main/java/de/ids_mannheim/korap/util/StatusCodes.java
+++ b/src/main/java/de/ids_mannheim/korap/util/StatusCodes.java
@@ -18,6 +18,7 @@
     public static final int DOC_ADDED = 681;
     public static final int RESPONSE_TIME_EXCEEDED = 682;
     public static final int STAGED_DATA_COMMITTED = 683;
+    public static final int MEMORY_LIMIT_EXCEEDED = 684;
 
     // 700 - 799 - KoralQuery Deserialization errors
     public static final int NO_QUERY_GIVEN = 700;
diff --git a/src/main/resources/krill.properties.info b/src/main/resources/krill.properties.info
index bbe38d5..b471f19 100644
--- a/src/main/resources/krill.properties.info
+++ b/src/main/resources/krill.properties.info
@@ -16,6 +16,9 @@
 krill.index.relations.max = 100
 krill.index.textSize.max = 20000000
 
+# Maximum JVM heap memory used (in MB) before a search is aborted (0 = disabled)
+krill.search.memory.max = 0
+
 # Token retrieval settings:
 #
 # krill.match.max.token = 5
@@ -77,3 +80,4 @@
 ## the client chooses one or the other per request. When character context
 ## is used, token-based maxShrink adjustment does NOT apply.
 ## Defaults to 500
+
diff --git a/src/test/java/de/ids_mannheim/korap/index/TestMemoryLimit.java b/src/test/java/de/ids_mannheim/korap/index/TestMemoryLimit.java
new file mode 100644
index 0000000..c1cfd3b
--- /dev/null
+++ b/src/test/java/de/ids_mannheim/korap/index/TestMemoryLimit.java
@@ -0,0 +1,92 @@
+package de.ids_mannheim.korap.index;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import java.io.IOException;
+
+import org.apache.lucene.index.Term;
+import org.apache.lucene.search.spans.SpanQuery;
+import org.apache.lucene.search.spans.SpanTermQuery;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import de.ids_mannheim.korap.Krill;
+import de.ids_mannheim.korap.KrillIndex;
+import de.ids_mannheim.korap.response.Result;
+import de.ids_mannheim.korap.util.KrillProperties;
+import de.ids_mannheim.korap.util.StatusCodes;
+
+/** AI generated
+ * 
+ */
+public class TestMemoryLimit {
+
+    private long savedMemoryLimit;
+
+    @Before
+    public void saveMemoryLimit () {
+        savedMemoryLimit = KrillProperties.maxMemoryMB;
+    }
+
+    @After
+    public void restoreMemoryLimit () {
+        KrillProperties.maxMemoryMB = savedMemoryLimit;
+    }
+
+    @Test
+    public void testMemoryLimitAbortsSingleSegment () throws IOException {
+        KrillIndex ki = new KrillIndex();
+        ki.addDoc(getClass().getResourceAsStream("/wiki/00001.json.gz"), true);
+        ki.commit();
+
+        // Set the limit to 1 MB so it is exceeded immediately
+        KrillProperties.maxMemoryMB = 1;
+
+        SpanQuery q = new SpanTermQuery(new Term("tokens", "s:der"));
+        Result kr = new Krill(q).apply(ki);
+        assertTrue("Expected memory exceeded warning", kr.hasWarnings());
+        assertEquals(1, kr.getWarnings().size());
+        assertEquals(StatusCodes.MEMORY_LIMIT_EXCEEDED,
+                kr.getWarning(0).getCode());
+    }
+
+    @Test
+    public void testMemoryLimitAbortsMultipleSegments () throws IOException {
+        KrillIndex ki = new KrillIndex();
+        ki.addDoc(getClass().getResourceAsStream("/wiki/00001.json.gz"), true);
+        ki.commit();
+        ki.addDoc(getClass().getResourceAsStream("/wiki/00002.json.gz"), true);
+        ki.commit();
+
+        // Set the limit to 1 MB so it is exceeded immediately
+        KrillProperties.maxMemoryMB = 1;
+
+        SpanQuery q = new SpanTermQuery(new Term("tokens", "s:der"));
+        Result kr = new Krill(q).apply(ki);
+
+        assertTrue("Expected memory exceeded warning", kr.hasWarnings());
+        assertEquals(1, kr.getWarnings().size());
+        assertEquals(StatusCodes.MEMORY_LIMIT_EXCEEDED,
+                kr.getWarning(0).getCode());
+    }
+
+    @Test
+    public void testNoMemoryLimitWhenDisabled () throws IOException {
+        KrillIndex ki = new KrillIndex();
+        ki.addDoc(getClass().getResourceAsStream("/wiki/00001.json.gz"), true);
+        ki.commit();
+
+        // 0 = disabled
+        KrillProperties.maxMemoryMB = 0;
+
+        SpanQuery q = new SpanTermQuery(new Term("tokens", "s:der"));
+        Result kr = new Krill(q).apply(ki);
+       
+        assertTrue("Expected results when memory limit is disabled",
+                kr.getTotalResults() > 0);
+        assertTrue("Expected no warnings when memory limit is disabled",
+                !kr.hasWarnings());
+    }
+}