Merge pull request rabbitmq#638 from janssk1/master

637: Configurable correlatorId generation for RPCClient

Author: acogoluegnes

src/main/java/com/rabbitmq/client/IncrementingCorrelationIdGenerator.java

@@ -0,0 +1,22 @@
+package com.rabbitmq.client;
+
+import java.util.function.Supplier;
+
+public class IncrementingCorrelationIdGenerator implements Supplier<String> {
+
+ private final String _prefix;
+ private int _correlationId;
+
+ public IncrementingCorrelationIdGenerator(String _prefix) {
+ this._prefix = _prefix;
+ }
+
+ @Override
+ public String get() {
+ return _prefix + _correlationId++;
+ }
+
+ public int getCorrelationId() {
+ return _correlationId;
+ }
+}

src/main/java/com/rabbitmq/client/RpcClient.java

@@ -28,6 +28,7 @@
 import java.util.Map.Entry;
 import java.util.concurrent.TimeoutException;
 import java.util.function.Function;
+import java.util.function.Supplier;
 
 import com.rabbitmq.client.impl.MethodArgumentReader;
 import com.rabbitmq.client.impl.MethodArgumentWriter;
@@ -84,7 +85,7 @@ public class RpcClient {
  /** Map from request correlation ID to continuation BlockingCell */
  private final Map<String, BlockingCell<Object>> _continuationMap = new HashMap<String, BlockingCell<Object>>();
  /** Contains the most recently-used request correlation ID */
- private int _correlationId;
+ private final Supplier<String> _correlationIdGenerator;
 
  /** Consumer attached to our reply queue */
  private DefaultConsumer _consumer;
@@ -109,7 +110,7 @@ public RpcClient(RpcClientParams params) throws
  _timeout = params.getTimeout();
  _useMandatory = params.shouldUseMandatory();
  _replyHandler = params.getReplyHandler();
- _correlationId = 0;
+ _correlationIdGenerator = params.getCorrelationIdGenerator();
 
  _consumer = setupConsumer();
  if (_useMandatory) {
@@ -208,8 +209,7 @@ public Response doCall(AMQP.BasicProperties props, byte[] message, int timeout)
  BlockingCell<Object> k = new BlockingCell<Object>();
  String replyId;
  synchronized (_continuationMap) {
- _correlationId++;
- replyId = "" + _correlationId;
+ replyId = _correlationIdGenerator.get();
  props = ((props==null) ? new AMQP.BasicProperties.Builder() : props.builder())
  .correlationId(replyId).replyTo(_replyTo).build();
  _continuationMap.put(replyId, k);
@@ -392,9 +392,14 @@ public Map<String, BlockingCell<Object>> getContinuationMap() {
  /**
  * Retrieve the correlation id.
  * @return the most recently used correlation id
+ * @deprecated Only works for {@link IncrementingCorrelationIdGenerator}
  */
  public int getCorrelationId() {
- return _correlationId;
+ if (_correlationIdGenerator instanceof IncrementingCorrelationIdGenerator) {
+ return ((IncrementingCorrelationIdGenerator) _correlationIdGenerator).getCorrelationId();
+ } else {
+ throw new UnsupportedOperationException();
+ }
  }
 
  /**

src/main/java/com/rabbitmq/client/RpcClientParams.java

@@ -16,6 +16,7 @@
 package com.rabbitmq.client;
 
 import java.util.function.Function;
+import java.util.function.Supplier;
 
 /**
  * Holder class to configure a {@link RpcClient}.
@@ -54,6 +55,8 @@ public class RpcClientParams {
  */
  private Function<Object, RpcClient.Response> replyHandler = RpcClient.DEFAULT_REPLY_HANDLER;
 
+ private Supplier<String> correlationIdGenerator = new IncrementingCorrelationIdGenerator("");
+
  /**
  * Set the channel to use for communication.
  *
@@ -170,6 +173,15 @@ public boolean shouldUseMandatory() {
  return useMandatory;
  }
 
+ public RpcClientParams correlationIdGenerator(Supplier<String> correlationIdGenerator) {
+ this.correlationIdGenerator = correlationIdGenerator;
+ return this;
+ }
+
+ public Supplier<String> getCorrelationIdGenerator() {
+ return correlationIdGenerator;
+ }
+
  public Function<Object, RpcClient.Response> getReplyHandler() {
  return replyHandler;
  }

src/test/java/com/rabbitmq/client/test/RpcTest.java

@@ -24,6 +24,7 @@
 import com.rabbitmq.client.impl.recovery.RecordedQueue;
 import com.rabbitmq.client.impl.recovery.TopologyRecoveryFilter;
 import com.rabbitmq.tools.Host;
+import org.hamcrest.CoreMatchers;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -39,9 +40,7 @@
 import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.awaitility.Awaitility.waitAtMost;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.fail;
+import static org.junit.Assert.*;
 
 public class RpcTest {
 
@@ -138,6 +137,25 @@ public void rpcUnroutableWithMandatoryFlagShouldThrowUnroutableException() throw
  client.close();
  }
 
+ @Test
+ public void rpcCustomCorrelatorId() throws Exception {
+ rpcServer = new TestRpcServer(serverChannel, queue);
+ new Thread(() -> {
+ try {
+ rpcServer.mainloop();
+ } catch (Exception e) {
+ // safe to ignore when loops ends/server is canceled
+ }
+ }).start();
+ RpcClient client = new RpcClient(new RpcClientParams()
+ .channel(clientChannel).exchange("").routingKey(queue).timeout(1000)
+ .correlationIdGenerator(new IncrementingCorrelationIdGenerator("myPrefix-"))
+ );
+ RpcClient.Response response = client.doCall(null, "hello".getBytes());
+ assertThat(response.getProperties().getCorrelationId(), CoreMatchers.equalTo("myPrefix-0"));
+ client.close();
+ }
+
  @Test
  public void rpcCustomReplyHandler() throws Exception {
  rpcServer = new TestRpcServer(serverChannel, queue);
@@ -156,7 +174,6 @@ public void rpcCustomReplyHandler() throws Exception {
  return RpcClient.DEFAULT_REPLY_HANDLER.apply(reply);
  })
  );
- assertEquals(0, replyHandlerCalls.get());
  RpcClient.Response response = client.doCall(null, "hello".getBytes());
  assertEquals(1, replyHandlerCalls.get());
  assertEquals("*** hello ***", new String(response.getBody()));