Polish correlation ID supplier support in RPC client

References #637

(cherry picked from commit 96b24b3)

Conflicts:
	src/main/java/com/rabbitmq/client/RpcClient.java
	src/main/java/com/rabbitmq/client/RpcClientParams.java
	src/test/java/com/rabbitmq/client/test/RpcTest.java

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
@@ -81,17 +81,20 @@ public Response handle(Object reply) {
         }
     };
 
-<<<<<<< HEAD
     private final RpcClientReplyHandler _replyHandler;
-=======
-    private final Function<Object, Response> _replyHandler;
->>>>>>> 102cbbde1... 637: Default generator should not be static
 
     /** 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;
 
@@ -115,7 +118,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) {
@@ -302,6 +305,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);
@@ -482,16 +486,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);
     }
 
     /**
@@ -559,5 +568,47 @@ public interface RpcClientReplyHandler {
         Response handle(Object reply);
 
     }
+
+    /**
+     * 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
@@ -52,7 +52,10 @@ public class RpcClientParams {
      */
     private RpcClient.RpcClientReplyHandler 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.
@@ -146,7 +149,7 @@ public RpcClientParams timeout(int timeout) {
      *
      * @param useMandatory
      * @return
-     * @see #replyHandler(RpcClient.RpcClientReplyHandler)
+     * @see #replyHandler(Function)
      */
     public RpcClientParams useMandatory(boolean useMandatory) {
         this.useMandatory = useMandatory;
@@ -170,17 +173,13 @@ public boolean shouldUseMandatory() {
         return useMandatory;
     }
 
-    public RpcClient.RpcClientReplyHandler getReplyHandler() {
-        return replyHandler;
-    }
-
-    public RpcClientParams correlationIdGenerator(Supplier<String> correlationIdGenerator) {
-        this.correlationIdGenerator = correlationIdGenerator;
+    public RpcClientParams correlationIdSupplier(Supplier<String> correlationIdGenerator) {
+        this.correlationIdSupplier = correlationIdGenerator;
         return this;
     }
 
-    public Supplier<String> getCorrelationIdGenerator() {
-        return correlationIdGenerator;
+    public Supplier<String> getCorrelationIdSupplier() {
+        return correlationIdSupplier;
     }
 
     /**
@@ -199,4 +198,8 @@ public RpcClientParams replyHandler(RpcClient.RpcClientReplyHandler replyHandler
         this.replyHandler = replyHandler;
         return this;
     }
+
+    public RpcClient.RpcClientReplyHandler getReplyHandler() {
+        return replyHandler;
+    }
 }

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

@@ -1,4 +1,4 @@
-// Copyright (c) 2017 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
@@ -34,7 +34,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;
@@ -100,6 +100,9 @@ public void run() {
         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();
     }
 
@@ -158,7 +161,7 @@ public void run() {
     }
 
     @Test
-    public void rpcCustomCorrelatorId() throws Exception {
+    public void rpcCustomCorrelationId() throws Exception {
         rpcServer = new TestRpcServer(serverChannel, queue);
         new Thread(() -> {
             try {
@@ -169,10 +172,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();
     }