From 3726eef3ba695e52c6fee30ababa301060c8b2fe Mon Sep 17 00:00:00 2001
From: Gary Russell <grussell@pivotal.io>
Date: Wed, 12 Oct 2016 12:47:34 -0400
Subject: [PATCH] Support Raw Keys

If the key generator is guaranteed to provide a unique key, it can be used as
the cache key directly.

Add a setter to support this option.
---
 .../StatefulRetryOperationsInterceptor.java   | 16 ++++++++++++++
 ...atefulRetryOperationsInterceptorTests.java | 22 +++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/src/main/java/org/springframework/retry/interceptor/StatefulRetryOperationsInterceptor.java b/src/main/java/org/springframework/retry/interceptor/StatefulRetryOperationsInterceptor.java
index 9014f76..07210ba 100644
--- a/src/main/java/org/springframework/retry/interceptor/StatefulRetryOperationsInterceptor.java
+++ b/src/main/java/org/springframework/retry/interceptor/StatefulRetryOperationsInterceptor.java
@@ -71,6 +71,8 @@ public class StatefulRetryOperationsInterceptor implements MethodInterceptor {
 
 	private Classifier<? super Throwable, Boolean> rollbackClassifier;
 
+	private boolean useRawKey;
+
 	public StatefulRetryOperationsInterceptor() {
 		RetryTemplate retryTemplate = new RetryTemplate();
 		retryTemplate.setRetryPolicy(new NeverRetryPolicy());
@@ -123,6 +125,17 @@ public class StatefulRetryOperationsInterceptor implements MethodInterceptor {
 		this.newMethodArgumentsIdentifier = newMethodArgumentsIdentifier;
 	}
 
+	/**
+	 * Set to true to use the raw key generated by the key generator. Should only be set
+	 * to true for cases where the key is guaranteed to be unique in all cases. When
+	 * false, a compound key is used, including invocation metadata.
+	 * Default: false.
+	 * @param useRawKey the useRawKey to set.
+	 */
+	public void setUseRawKey(boolean useRawKey) {
+		this.useRawKey = useRawKey;
+	}
+
 	/**
 	 * Wrap the method invocation in a stateful retry with the policy and other helpers
 	 * provided. If there is a failure the exception will generally be re-thrown. The only
@@ -181,6 +194,9 @@ public class StatefulRetryOperationsInterceptor implements MethodInterceptor {
 			// really doesn't want to retry.
 			return null;
 		}
+		if (this.useRawKey) {
+			return generatedKey;
+		}
 		String name = StringUtils.hasText(label) ? label
 				: invocation.getMethod().toGenericString();
 		return Arrays.asList(name, generatedKey);
diff --git a/src/test/java/org/springframework/retry/interceptor/StatefulRetryOperationsInterceptorTests.java b/src/test/java/org/springframework/retry/interceptor/StatefulRetryOperationsInterceptorTests.java
index 2d40c47..fdbbf84 100644
--- a/src/test/java/org/springframework/retry/interceptor/StatefulRetryOperationsInterceptorTests.java
+++ b/src/test/java/org/springframework/retry/interceptor/StatefulRetryOperationsInterceptorTests.java
@@ -268,6 +268,28 @@ public class StatefulRetryOperationsInterceptorTests {
 		assertNull(captor.getValue().getKey());
 	}
 
+	@SuppressWarnings("unchecked")
+	@Test
+	public void testKeyGeneratorAndRawKey() throws Throwable {
+		this.interceptor.setKeyGenerator(new MethodArgumentsKeyGenerator() {
+
+			@Override
+			public Object getKey(Object[] item) {
+				return "bar";
+			}
+		});
+		this.interceptor.setLabel("foo");
+		this.interceptor.setUseRawKey(true);
+		RetryOperations template = mock(RetryOperations.class);
+		this.interceptor.setRetryOperations(template);
+		MethodInvocation invocation = mock(MethodInvocation.class);
+		when(invocation.getArguments()).thenReturn(new Object[] { new Object() });
+		this.interceptor.invoke(invocation);
+		ArgumentCaptor<DefaultRetryState> captor = ArgumentCaptor.forClass(DefaultRetryState.class);
+		verify(template).execute(any(RetryCallback.class), any(RecoveryCallback.class), captor.capture());
+		assertEquals("bar", captor.getValue().getKey());
+	}
+
 	@Test
 	public void testTransformerRecoveryAfterTooManyAttempts() throws Exception {
 		((Advised) transformer).addAdvice(interceptor);