RESOLVED - issue BATCH-537: Bad ItemKeyGenerator strategy can lead to infinite loop in retry

Some changes to account for retry exceptions and force a failure if strange hashCode/equals are detected.

spring-batch-infrastructure/src/main/java/org/springframework/batch/retry/policy/ItemWriterRetryPolicy.java

@@ -24,6 +24,7 @@ import org.springframework.batch.item.ItemRecoverer;
 import org.springframework.batch.repeat.support.RepeatSynchronizationManager;
 import org.springframework.batch.retry.RetryCallback;
 import org.springframework.batch.retry.RetryContext;
+import org.springframework.batch.retry.RetryException;
 import org.springframework.batch.retry.RetryPolicy;
 import org.springframework.batch.retry.TerminatedRetryException;
 import org.springframework.batch.retry.callback.ItemWriterRetryCallback;
@@ -136,6 +137,10 @@ public class ItemWriterRetryPolicy extends AbstractStatefulRetryPolicy {
 
 		final private Object item;
 
+		final private Object key;
+
+		final private int initialHashCode;
+
 		// The delegate context...
 		private RetryContext delegateContext;
 
@@ -147,10 +152,12 @@ public class ItemWriterRetryPolicy extends AbstractStatefulRetryPolicy {
 
 		public ItemWriterRetryContext(ItemWriterRetryCallback callback, RetryContext parent) {
 			super(parent);
-			this.item = callback.getItem();
 			this.recoverer = callback.getRecoverer();
 			this.keyGenerator = callback.getKeyGenerator();
+			this.item = callback.getItem();
+			this.key = keyGenerator.getKey(item);
 			this.failedItemIdentifier = callback.getFailedItemIdentifier();
+			this.initialHashCode = key.hashCode();
 		}
 
 		public boolean canRetry(RetryContext context) {
@@ -162,10 +169,15 @@ public class ItemWriterRetryPolicy extends AbstractStatefulRetryPolicy {
 		}
 
 		public RetryContext open(RetryCallback callback, RetryContext parent) {
-			if (hasFailed(failedItemIdentifier, keyGenerator, item)) {
-				this.delegateContext = retryContextCache.get(keyGenerator.getKey(item));
+			if (hasFailed(failedItemIdentifier, key)) {
+				this.delegateContext = retryContextCache.get(key);
+				if (this.delegateContext == null) {
+					throw new RetryException("Inconsistent state for failed item: no history found. "
+							+ "Consider whether equals() or hashCode() for the item might be inconsistent, "
+							+ "or if you need to supply a better ItemKeyGenerator");
+				}
 			}
-			if (this.delegateContext == null) {
+			else {
 				// Only create a new context if we don't know the history of
 				// this item:
 				this.delegateContext = delegate.open(callback, null);
@@ -175,7 +187,14 @@ public class ItemWriterRetryPolicy extends AbstractStatefulRetryPolicy {
 		}
 
 		public void registerThrowable(RetryContext context, Throwable throwable) throws TerminatedRetryException {
-			retryContextCache.put(keyGenerator.getKey(item), this.delegateContext);
+			// TODO: this comparison assumes that hashCode is the limiting
+			// factor. Actually the cache should be able to decide for us.
+			if (this.initialHashCode != key.hashCode()) {
+				throw new RetryException("Inconsistent state for failed item key: hashCode has changed. "
+						+ "Consider whether equals() or hashCode() for the item might be inconsistent, "
+						+ "or if you need to supply a better ItemKeyGenerator");
+			}
+			retryContextCache.put(key, this.delegateContext);
 			delegate.registerThrowable(this.delegateContext, throwable);
 		}
 
@@ -191,7 +210,7 @@ public class ItemWriterRetryPolicy extends AbstractStatefulRetryPolicy {
 
 		public Object handleRetryExhausted(RetryContext context) throws Exception {
 			// If there is no going back, then we can remove the history
-			retryContextCache.remove(keyGenerator.getKey(item));
+			retryContextCache.remove(key);
 			RepeatSynchronizationManager.setCompleteOnly();
 			if (recoverer != null) {
 				boolean success = recoverer.recover(item, context.getLastThrowable());
@@ -225,15 +244,13 @@ public class ItemWriterRetryPolicy extends AbstractStatefulRetryPolicy {
 	 * item key.
 	 * 
 	 * @param failedItemIdentifier
-	 * @param keyGenerator
-	 * @param item
+	 * @param key
 	 * @return
 	 */
-	protected boolean hasFailed(FailedItemIdentifier failedItemIdentifier, ItemKeyGenerator keyGenerator, Object item) {
+	protected boolean hasFailed(FailedItemIdentifier failedItemIdentifier, Object key) {
 		if (failedItemIdentifier != null) {
-			return failedItemIdentifier.hasFailed(item);
+			return failedItemIdentifier.hasFailed(key);
 		}
-		return retryContextCache.containsKey(keyGenerator.getKey(item));
+		return retryContextCache.containsKey(key);
 	}
-
 }

spring-batch-infrastructure/src/main/java/org/springframework/batch/retry/policy/MapRetryContextCache.java

@@ -24,15 +24,52 @@ import org.springframework.batch.retry.RetryContext;
 
 /**
  * Map-based implementation of {@link RetryContextCache}. The map backing the
- * cache of contexts is sytchronized.
+ * cache of contexts is synchronized.
  * 
  * @author Dave Syer
  * 
  */
 public class MapRetryContextCache implements RetryContextCache {
 
+	/**
+	 * Default value for maximum capacity of the cache. This is set to a
+	 * reasonably low value (4096) to avoid users inadvertently filling the
+	 * cache with item keys that are inconsistent.
+	 */
+	public static final int DEFAULT_CAPACITY = 4096;
+
 	private Map map = Collections.synchronizedMap(new HashMap());
 
+	private int capacity;
+
+	/**
+	 * Create a {@link MapRetryContextCache} with default capacity.
+	 */
+	public MapRetryContextCache() {
+		this(DEFAULT_CAPACITY);
+	}
+
+	/**
+	 * @param defaultCapacity
+	 */
+	public MapRetryContextCache(int defaultCapacity) {
+		super();
+		this.capacity = defaultCapacity;
+	}
+
+	/**
+	 * Public setter for the capacity. Prevents the cache from growing
+	 * unboundedly if items that fail are misidentified and two references to an
+	 * identical item actually do not have the same key. This can happen when
+	 * users implement equals and hashCode based on mutable fields, for
+	 * instance.
+	 * 
+	 * @param capacity the capacity to set
+	 */
+	public void setCapacity(int capacity) {
+		this.capacity = capacity;
+	}
+
 	public boolean containsKey(Object key) {
 		return map.containsKey(key);
 	}
@@ -42,6 +79,11 @@ public class MapRetryContextCache implements RetryContextCache {
 	}
 
 	public void put(Object key, RetryContext context) {
+		if (map.size() >= capacity) {
+			throw new RetryCacheCapacityExceededException("Retry cache capacity limit breached. " +
+					"Do you need to re-consider the implementation of the key generator, " +
+					"or the equals and hashCode of the items that failed?");
+		}
 		map.put(key, context);
 	}
 

spring-batch-infrastructure/src/main/java/org/springframework/batch/retry/policy/RetryCacheCapacityExceededException.java

@@ -0,0 +1,52 @@
+/*
+ * Copyright 2006-2007 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.batch.retry.policy;
+
+import org.springframework.batch.retry.RetryException;
+
+/**
+ * Exception that indicates that a cache limit was exceeded. This is often a
+ * sign of badly or inconsistently implemented hashCode, equals in failed items.
+ * Items can then fail repeatedly and appear different to the cache, so they get
+ * added over and over again until a limit is reached and this exception is
+ * thrown. Consult the documentation of the {@link RetryContextCache} in use to
+ * determine how to increase the limit if appropriate.
+ * 
+ * @author Dave Syer
+ * 
+ */
+public class RetryCacheCapacityExceededException extends RetryException {
+
+	/**
+	 * Constructs a new instance with a message.
+	 * 
+	 * @param message
+	 */
+	public RetryCacheCapacityExceededException(String message) {
+		super(message);
+	}
+
+	/**
+	 * Constructs a new instance with a message and nested exception.
+	 * 
+	 * @param msg the exception message.
+	 * 
+	 */
+	public RetryCacheCapacityExceededException(String msg, Throwable nested) {
+		super(msg, nested);
+	}
+
+}

spring-batch-infrastructure/src/main/java/org/springframework/batch/retry/policy/RetryContextCache.java

@@ -31,7 +31,7 @@ public interface RetryContextCache {
 
 	RetryContext get(Object key);
 
-	void put(Object key, RetryContext context);
+	void put(Object key, RetryContext context) throws RetryCacheCapacityExceededException;
 
 	void remove(Object key);
 

spring-batch-infrastructure/src/test/java/org/springframework/batch/retry/policy/ItemWriterRetryPolicyTests.java

@@ -31,6 +31,7 @@ import org.springframework.batch.repeat.context.RepeatContextSupport;
 import org.springframework.batch.repeat.support.RepeatSynchronizationManager;
 import org.springframework.batch.retry.RetryCallback;
 import org.springframework.batch.retry.RetryContext;
+import org.springframework.batch.retry.RetryException;
 import org.springframework.batch.retry.StubItemKeyGeneratorRecoverer;
 import org.springframework.batch.retry.callback.ItemWriterRetryCallback;
 import org.springframework.batch.retry.context.RetryContextSupport;
@@ -46,12 +47,6 @@ public class ItemWriterRetryPolicyTests extends TestCase {
 
 	private List list = new ArrayList();
 
-	private ItemKeyGenerator keyGenerator = new ItemKeyGenerator() {
-		public Object getKey(Object item) {
-			return item;
-		}
-	};
-
 	protected void setUp() throws Exception {
 		super.setUp();
 		// The list simulates a failed delivery, redelivery of the same message,
@@ -294,10 +289,96 @@ public class ItemWriterRetryPolicyTests extends TestCase {
 		MapRetryContextCache cache = new MapRetryContextCache();
 		policy.setRetryContextCache(cache);
 		cache.put("foo", new RetryContextSupport(null));
-		assertTrue(policy.hasFailed(null, keyGenerator , "foo"));
+		assertTrue(policy.hasFailed(null, "foo"));
 	}
 
-	private static class MockFailedItemProvider extends ListItemReader implements ItemKeyGenerator, FailedItemIdentifier {
+	public void testKeyGeneratorNotConsistentAfterFailure() throws Throwable {
+
+		AbstractItemWriter writer = new AbstractItemWriter() {
+			public void write(Object data) {
+				// This simulates what happens if someone uses a primary key
+				// for hasCode and equals and then relies on default key
+				// generator
+				((StringHolder) data).string = ((StringHolder) data).string + (count++);
+				throw new RuntimeException("Barf!");
+			}
+		};
+
+		policy = new ItemWriterRetryPolicy();
+		policy.setDelegate(new SimpleRetryPolicy(3));
+		StringHolder item = new StringHolder("bar");
+		ItemWriterRetryCallback callback = new ItemWriterRetryCallback(item, writer);
+		RetryContext context = policy.open(callback, null);
+		assertNotNull(context);
+		try {
+			callback.doWithRetry(context);
+			fail("Expected RuntimeException");
+		}
+		catch (RuntimeException e) {
+			assertEquals("Barf!", e.getMessage());
+			try {
+				policy.registerThrowable(context, e);
+				fail("Expected RetryException");
+			}
+			catch (RetryException ex) {
+				String message = ex.getMessage();
+				assertTrue("Message doesn't contain 'inconsistent': " + message, message.indexOf("inconsistent") >= 0);
+			}
+			assertEquals(0, context.getRetryCount());
+		}
+
+	}
+
+	public void testCacheCapacity() throws Exception {
+		policy = new ItemWriterRetryPolicy();
+		policy.setDelegate(new SimpleRetryPolicy(1));
+		policy.setRetryContextCache(new MapRetryContextCache(1));
+		AbstractItemWriter writer = new AbstractItemWriter() {
+			public void write(Object data) {
+				count++;
+				list.add(data);
+			}
+		};
+		RetryContext context;
+		context = policy.open(new ItemWriterRetryCallback("foo", writer), null);
+		policy.registerThrowable(context, null);
+		assertEquals(0, context.getRetryCount());
+		context = policy.open(new ItemWriterRetryCallback("bar", writer), null);
+		try {
+			policy.registerThrowable(context, new RuntimeException("foo"));
+			fail("Expected RetryException");
+		}
+		catch (RetryException e) {
+			String message = e.getMessage();
+			assertTrue("Message does not contain 'capacity': " + message, message.indexOf("capacity") >= 0);
+		}
+	}
+
+	public void testCacheCapacityNotReachedIfRecovered() throws Exception {
+		policy = new ItemWriterRetryPolicy();
+		policy.setDelegate(new SimpleRetryPolicy(1));
+		policy.setRetryContextCache(new MapRetryContextCache(2));
+		AbstractItemWriter writer = new AbstractItemWriter() {
+			public void write(Object data) {
+				count++;
+				list.add(data);
+			}
+		};
+		RetryContext context;
+		context = policy.open(new ItemWriterRetryCallback("foo", writer), null);
+		policy.registerThrowable(context, null);
+		assertEquals(0, context.getRetryCount());
+		policy.registerThrowable(context, new RuntimeException("foo"));
+		context = policy.open(new ItemWriterRetryCallback("bar", writer), null);
+		policy.registerThrowable(context, null);
+		policy.handleRetryExhausted(context);
+		context = policy.open(new ItemWriterRetryCallback("spam", writer), null);
+		policy.registerThrowable(context, null);
+		assertEquals(0, context.getRetryCount());
+	}
+
+	private static class MockFailedItemProvider extends ListItemReader implements ItemKeyGenerator,
+			FailedItemIdentifier {
 
 		private int hasFailedCount = 0;
 
@@ -316,4 +397,41 @@ public class ItemWriterRetryPolicyTests extends TestCase {
 
 	}
 
+	private static class StringHolder {
+
+		private String string;
+
+		/**
+		 * @param string
+		 */
+		public StringHolder(String string) {
+			this.string = string;
+		}
+
+		/*
+		 * (non-Javadoc)
+		 * @see java.lang.Object#equals(java.lang.Object)
+		 */
+		public boolean equals(Object obj) {
+			return string.equals(((StringHolder) obj).string);
+		}
+
+		/*
+		 * (non-Javadoc)
+		 * @see java.lang.Object#hashCode()
+		 */
+		public int hashCode() {
+			return string.hashCode();
+		}
+
+		/*
+		 * (non-Javadoc)
+		 * @see java.lang.Object#toString()
+		 */
+		public String toString() {
+			return "String: " + string + " (hash = " + hashCode() + ")";
+		}
+
+	}
+
 }