Polish correlation ID supplier support in RPC client

References #637

Author: acogoluegnes

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

@@ -1,4 +1,4 @@
-// Copyright (c) 2007-Present Pivotal Software, Inc.  All rights reserved.
+// Copyright (c) 2007-2020 Pivotal Software, Inc.  All rights reserved.
 //
 // This software, the RabbitMQ Java client library, is triple-licensed under the
 // Mozilla Public License 1.1 ("MPL"), the GNU General Public License version 2
@@ -84,9 +84,16 @@ 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 */
+
+    /**
+     * Generates correlation ID for each request.
+     *
+     * @since 5.9.0
+     */
     private final Supplier<String> _correlationIdGenerator;
 
+    private String lastCorrelationId = "0";
+
     /** Consumer attached to our reply queue */
     private DefaultConsumer _consumer;
 
@@ -110,7 +117,7 @@ public RpcClient(RpcClientParams params) throws
         _timeout = params.getTimeout();
         _useMandatory = params.shouldUseMandatory();
         _replyHandler = params.getReplyHandler();
-        _correlationIdGenerator = params.getCorrelationIdGenerator();
+        _correlationIdGenerator = params.getCorrelationIdSupplier();
 
         _consumer = setupConsumer();
         if (_useMandatory) {
@@ -210,6 +217,7 @@ public Response doCall(AMQP.BasicProperties props, byte[] message, int timeout)
         String replyId;
         synchronized (_continuationMap) {
             replyId = _correlationIdGenerator.get();
+            lastCorrelationId = replyId;
             props = ((props==null) ? new AMQP.BasicProperties.Builder() : props.builder())
                 .correlationId(replyId).replyTo(_replyTo).build();
             _continuationMap.put(replyId, k);
@@ -390,16 +398,21 @@ public Map<String, BlockingCell<Object>> getContinuationMap() {
     }
 
     /**
-     * Retrieve the correlation id.
+     * Retrieve the last correlation id used.
+     * <p>
+     * Note as of 5.9.0, correlation IDs may not always be integers
+     * (by default, they are).
+     * This method will try to parse the last correlation ID string
+     * as an integer, so this may result in {@link NumberFormatException}
+     * if the correlation ID supplier provided by
+     * {@link RpcClientParams#correlationIdSupplier(Supplier)}
+     * does not generate appropriate IDs.
+     *
      * @return the most recently used correlation id
-     * @deprecated Only works for {@link IncrementingCorrelationIdGenerator}
+     * @see RpcClientParams#correlationIdSupplier(Supplier)
      */
     public int getCorrelationId() {
-        if (_correlationIdGenerator instanceof IncrementingCorrelationIdGenerator) {
-            return ((IncrementingCorrelationIdGenerator) _correlationIdGenerator).getCorrelationId();
-        } else {
-            throw new UnsupportedOperationException();
-        }
+        return Integer.valueOf(this.lastCorrelationId);
     }
 
     /**
@@ -447,5 +460,47 @@ public byte[] getBody() {
             return body;
         }
     }
+
+    /**
+     * Creates generation IDs as a sequence of integers.
+     *
+     * @return
+     * @see RpcClientParams#correlationIdSupplier(Supplier)
+     * @since 5.9.0
+     */
+    public static Supplier<String> incrementingCorrelationIdSupplier() {
+        return incrementingCorrelationIdSupplier("");
+    }
+
+    /**
+     * Creates generation IDs as a sequence of integers, with the provided prefix.
+     *
+     * @param prefix
+     * @return
+     * @see RpcClientParams#correlationIdSupplier(Supplier)
+     * @since 5.9.0
+     */
+    public static Supplier<String> incrementingCorrelationIdSupplier(String prefix) {
+        return new IncrementingCorrelationIdSupplier(prefix);
+    }
+
+    /**
+     * @since 5.9.0
+     */
+    private static class IncrementingCorrelationIdSupplier implements Supplier<String> {
+
+        private final String prefix;
+        private int correlationId;
+
+        public IncrementingCorrelationIdSupplier(String prefix) {
+            this.prefix = prefix;
+        }
+
+        @Override
+        public String get() {
+            return prefix + ++correlationId;
+        }
+
+    }
 }
 

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

@@ -1,4 +1,4 @@
-// Copyright (c) 2007-Present Pivotal Software, Inc.  All rights reserved.
+// Copyright (c) 2007-2020 Pivotal Software, Inc.  All rights reserved.
 //
 // This software, the RabbitMQ Java client library, is triple-licensed under the
 // Mozilla Public License 1.1 ("MPL"), the GNU General Public License version 2
@@ -55,7 +55,10 @@ public class RpcClientParams {
      */
     private Function<Object, RpcClient.Response> replyHandler = RpcClient.DEFAULT_REPLY_HANDLER;
 
-    private Supplier<String> correlationIdGenerator = new IncrementingCorrelationIdGenerator("");
+    /**
+     * Logic to generate correlation IDs.
+     */
+    private Supplier<String> correlationIdSupplier = RpcClient.incrementingCorrelationIdSupplier();
 
     /**
      * Set the channel to use for communication.
@@ -149,7 +152,7 @@ public RpcClientParams timeout(int timeout) {
      *
      * @param useMandatory
      * @return
-     * @see #replyHandler(RpcClient.RpcClientReplyHandler)
+     * @see #replyHandler(Function)
      */
     public RpcClientParams useMandatory(boolean useMandatory) {
         this.useMandatory = useMandatory;
@@ -173,13 +176,20 @@ public boolean shouldUseMandatory() {
         return useMandatory;
     }
 
-    public RpcClientParams correlationIdGenerator(Supplier<String> correlationIdGenerator) {
-        this.correlationIdGenerator = correlationIdGenerator;
+    /**
+     * Logic to generate correlation IDs.
+     *
+     * @param correlationIdGenerator
+     * @return
+     * @since 5.9.0
+     */
+    public RpcClientParams correlationIdSupplier(Supplier<String> correlationIdGenerator) {
+        this.correlationIdSupplier = correlationIdGenerator;
         return this;
     }
 
-    public Supplier<String> getCorrelationIdGenerator() {
-        return correlationIdGenerator;
+    public Supplier<String> getCorrelationIdSupplier() {
+        return correlationIdSupplier;
     }
 
     public Function<Object, RpcClient.Response> getReplyHandler() {

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

@@ -1,4 +1,4 @@
-// Copyright (c) 2017-2019 Pivotal Software, Inc.  All rights reserved.
+// Copyright (c) 2017-2020 Pivotal Software, Inc.  All rights reserved.
 //
 // This software, the RabbitMQ Java client library, is triple-licensed under the
 // Mozilla Public License 1.1 ("MPL"), the GNU General Public License version 2
@@ -24,7 +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.assertj.core.api.Assertions;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -86,6 +86,9 @@ public void rpc() throws Exception {
         assertEquals("*** hello ***", new String(response.getBody()));
         assertEquals("pre-hello", response.getProperties().getHeaders().get("pre").toString());
         assertEquals("post-hello", response.getProperties().getHeaders().get("post").toString());
+
+        Assertions.assertThat(client.getCorrelationId()).isEqualTo(Integer.valueOf(response.getProperties().getCorrelationId()));
+
         client.close();
     }
 
@@ -138,7 +141,7 @@ public void rpcUnroutableWithMandatoryFlagShouldThrowUnroutableException() throw
     }
 
     @Test
-    public void rpcCustomCorrelatorId() throws Exception {
+    public void rpcCustomCorrelationId() throws Exception {
         rpcServer = new TestRpcServer(serverChannel, queue);
         new Thread(() -> {
             try {
@@ -149,10 +152,10 @@ public void rpcCustomCorrelatorId() throws Exception {
         }).start();
         RpcClient client = new RpcClient(new RpcClientParams()
                 .channel(clientChannel).exchange("").routingKey(queue).timeout(1000)
-                .correlationIdGenerator(new IncrementingCorrelationIdGenerator("myPrefix-"))
+                .correlationIdSupplier(RpcClient.incrementingCorrelationIdSupplier("myPrefix-"))
         );
         RpcClient.Response response = client.doCall(null, "hello".getBytes());
-        assertThat(response.getProperties().getCorrelationId(), CoreMatchers.equalTo("myPrefix-0"));
+        Assertions.assertThat(response.getProperties().getCorrelationId()).isEqualTo("myPrefix-1");
         client.close();
     }