Ensure retry context is saved to cache after it is updated

If the order is wrong (as it was), then a cache that serializes the
context will only ever contain the new value, and never be
updated.

src/main/java/org/springframework/retry/support/RetryTemplate.java

@@ -402,15 +402,15 @@ public class RetryTemplate implements RetryOperations {
 
 	protected void registerThrowable(RetryPolicy retryPolicy, RetryState state,
 			RetryContext context, Throwable e) {
-		registerContext(context, state);
 		retryPolicy.registerThrowable(context, e);
+		registerContext(context, state);
 	}
 
 	private void registerContext(RetryContext context, RetryState state) {
 		if (state != null) {
 			Object key = state.getKey();
 			if (key != null) {
-				if (context.getRetryCount() > 0
+				if (context.getRetryCount() > 1
 						&& !this.retryContextCache.containsKey(key)) {
 					throw new RetryException(
 							"Inconsistent state for failed item key: cache key has changed. "

src/test/java/org/springframework/retry/policy/RetryContextSerializationTests.java

@@ -39,6 +39,7 @@ import org.springframework.retry.RetryPolicy;
 import org.springframework.util.ClassUtils;
 import org.springframework.util.SerializationUtils;
 
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
 /**
@@ -47,7 +48,7 @@ import static org.junit.Assert.assertTrue;
  */
 @RunWith(Parameterized.class)
 public class RetryContextSerializationTests {
-	
+
 	private static Log logger = LogFactory.getLog(RetryContextSerializationTests.class);
 
 	private RetryPolicy policy;
@@ -62,7 +63,8 @@ public class RetryContextSerializationTests {
 		Set<BeanDefinition> candidates = scanner.findCandidateComponents("org.springframework.retry.policy");
 		for (BeanDefinition beanDefinition : candidates) {
 			try {
-				result.add(new Object[] { BeanUtils.instantiate(ClassUtils.resolveClassName(beanDefinition.getBeanClassName(), null)) });
+				result.add(new Object[] {
+						BeanUtils.instantiate(ClassUtils.resolveClassName(beanDefinition.getBeanClassName(), null)) });
 			} catch (Exception e) {
 				logger.warn("Cannot create instance of " + beanDefinition.getBeanClassName(), e);
 			}
@@ -79,14 +81,17 @@ public class RetryContextSerializationTests {
 
 	@Test
 	public void testSerializationCycleForContext() {
-		assertTrue(SerializationUtils.deserialize(
-				SerializationUtils.serialize(policy.open(null))) instanceof RetryContext);
+		RetryContext context = policy.open(null);
+		assertEquals(0, context.getRetryCount());
+		policy.registerThrowable(context, new RuntimeException());
+		assertEquals(1, context.getRetryCount());
+		assertEquals(1,
+				((RetryContext) SerializationUtils.deserialize(SerializationUtils.serialize(context))).getRetryCount());
 	}
 
 	@Test
 	public void testSerializationCycleForPolicy() {
-		assertTrue(SerializationUtils.deserialize(
-				SerializationUtils.serialize(policy)) instanceof RetryPolicy);
+		assertTrue(SerializationUtils.deserialize(SerializationUtils.serialize(policy)) instanceof RetryPolicy);
 	}
 
 }

src/test/java/org/springframework/retry/policy/SerializedMapRetryContextCache.java

@@ -0,0 +1,56 @@
+/*
+ * Copyright 2016 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.retry.policy;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+import org.springframework.retry.RetryContext;
+import org.springframework.util.SerializationUtils;
+
+public class SerializedMapRetryContextCache implements RetryContextCache {
+	private static final int DEFAULT_CAPACITY = 4096;
+
+	private Map<Object, byte[]> map = Collections.synchronizedMap(new HashMap<Object, byte[]>());
+
+	@Override
+	public boolean containsKey(Object key) {
+		return map.containsKey(key);
+	}
+
+	@Override
+	public RetryContext get(Object key) {
+		byte[] bytes = map.get(key);
+		return (RetryContext) SerializationUtils.deserialize(bytes);
+	}
+
+	@Override
+	public void put(Object key, RetryContext context) {
+		if (map.size() >= DEFAULT_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?");
+		}
+		byte[] serialized = SerializationUtils.serialize(context);
+		map.put(key, serialized);
+	}
+
+	@Override
+	public void remove(Object key) {
+		map.remove(key);
+	}
+}

src/test/java/org/springframework/retry/policy/StatefulRetryIntegrationTests.java

@@ -113,6 +113,41 @@ public class StatefulRetryIntegrationTests {
 		assertFalse(cache.containsKey("foo"));
 
 		assertEquals(2, callback.attempts);
+		assertEquals(1, callback.context.getRetryCount());
+		assertEquals("bar", result);
+	}
+
+	@Test
+	public void testExternalRetryWithSuccessOnRetryAndSerializedContext() throws Throwable {
+		MockRetryCallback callback = new MockRetryCallback();
+
+		RetryState retryState = new DefaultRetryState("foo");
+
+		RetryTemplate retryTemplate = new RetryTemplate();
+		RetryContextCache cache = new SerializedMapRetryContextCache();
+		retryTemplate.setRetryContextCache(cache);
+		retryTemplate.setRetryPolicy(new SimpleRetryPolicy(2));
+
+		assertFalse(cache.containsKey("foo"));
+
+		Object result = "start_foo";
+		try {
+			result = retryTemplate.execute(callback, retryState);
+			// The first failed attempt we expect to retry...
+			fail("Expected RuntimeException");
+		}
+		catch (RuntimeException e) {
+			assertNull(e.getMessage());
+		}
+
+		assertTrue(cache.containsKey("foo"));
+
+		result = retryTemplate.execute(callback, retryState);
+
+		assertFalse(cache.containsKey("foo"));
+
+		assertEquals(2, callback.attempts);
+		assertEquals(1, callback.context.getRetryCount());
 		assertEquals("bar", result);
 	}
 
@@ -182,9 +217,11 @@ public class StatefulRetryIntegrationTests {
 	private static final class MockRetryCallback
 			implements RetryCallback<String, Exception> {
 		int attempts = 0;
-
+		RetryContext context;
+		
 		public String doWithRetry(RetryContext context) throws Exception {
 			attempts++;
+			this.context = context;
 			if (attempts < 2) {
 				throw new RuntimeException();
 			}