Make RPC mechanics configurable

Author: acogoluegnes

src/main/java/com/rabbitmq/model/Consumer.java

@@ -22,6 +22,7 @@ public interface Consumer extends AutoCloseable {
   @Override
   void close();
 
+  @FunctionalInterface
   interface MessageHandler {
 
     void handle(Context context, Message message);

src/main/java/com/rabbitmq/model/RpcClientBuilder.java

@@ -17,12 +17,22 @@
 // [email protected].
 package com.rabbitmq.model;
 
+import java.util.function.BiFunction;
+import java.util.function.Function;
+import java.util.function.Supplier;
+
 public interface RpcClientBuilder {
 
   RpcClientAddressBuilder requestAddress();
 
   RpcClientBuilder replyToQueue(String replyToQueue);
 
+  RpcClientBuilder correlationIdSupplier(Supplier<Object> correlationIdSupplier);
+
+  RpcClientBuilder requestPostProcessor(BiFunction<Message, Object, Message> requestPostProcessor);
+
+  RpcClientBuilder correlationIdExtractor(Function<Message, Object> correlationIdExtractor);
+
   RpcClient build();
 
   interface RpcClientAddressBuilder extends AddressBuilder<RpcClientAddressBuilder> {

src/main/java/com/rabbitmq/model/RpcServer.java

@@ -19,6 +19,7 @@
 
 public interface RpcServer extends AutoCloseable {
 
+  @FunctionalInterface
   interface Handler {
 
     Message handle(Context ctx, Message request);

src/main/java/com/rabbitmq/model/RpcServerBuilder.java

@@ -17,6 +17,9 @@
 // [email protected].
 package com.rabbitmq.model;
 
+import java.util.function.BiFunction;
+import java.util.function.Function;
+
 public interface RpcServerBuilder {
 
   RpcServerBuilder requestQueue(String requestQueue);
@@ -25,6 +28,10 @@ public interface RpcServerBuilder {
 
   RpcServerAddressBuilder replyToAddress();
 
+  RpcServerBuilder correlationIdExtractor(Function<Message, Object> correlationIdExtractor);
+
+  RpcServerBuilder replyPostProcessor(BiFunction<Message, Object, Message> replyPostProcessor);
+
   RpcServer build();
 
   interface RpcServerAddressBuilder extends AddressBuilder<RpcServerAddressBuilder> {

src/main/java/com/rabbitmq/model/amqp/AmqpRpcClient.java

@@ -24,6 +24,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.BiFunction;
+import java.util.function.Function;
 import java.util.function.Supplier;
 
 class AmqpRpcClient implements RpcClient {
@@ -32,11 +33,11 @@ class AmqpRpcClient implements RpcClient {
 
   private final Publisher publisher;
   private final Consumer consumer;
-  private final String replyToAddress;
   private final Map<Object, CompletableFuture<Message>> outstandingRequests =
       new ConcurrentHashMap<>();
   private final Supplier<Object> correlationIdSupplier;
   private final BiFunction<Message, Object, Message> requestPostProcessor;
+  private final Function<Message, Object> correlationIdExtractor;
 
   AmqpRpcClient(RpcSupport.AmqpRpcClientBuilder builder) {
     AmqpConnection connection = builder.connection();
@@ -52,6 +53,11 @@ class AmqpRpcClient implements RpcClient {
           connection.management().queue().exclusive(true).autoDelete(true).declare();
       replyTo = queueInfo.name();
     }
+    if (builder.correlationIdExtractor() == null) {
+      this.correlationIdExtractor = Message::correlationId;
+    } else {
+      this.correlationIdExtractor = builder.correlationIdExtractor();
+    }
     this.consumer =
         connection
             .consumerBuilder()
@@ -60,22 +66,31 @@ class AmqpRpcClient implements RpcClient {
                 (ctx, msg) -> {
                   ctx.accept();
                   CompletableFuture<Message> result =
-                      this.outstandingRequests.remove(msg.correlationId());
+                      this.outstandingRequests.remove(this.correlationIdExtractor.apply(msg));
                   if (result != null) {
                     result.complete(msg);
                   }
                 })
             .build();
-    DefaultAddressBuilder<?> addressBuilder = new DefaultAddressBuilder<>(null) {};
-    addressBuilder.queue(replyTo);
-    this.replyToAddress = addressBuilder.address();
 
-    String correlationIdPrefix = UUID.randomUUID().toString();
-    AtomicLong correlationIdSequence = new AtomicLong();
-    this.correlationIdSupplier =
-        () -> correlationIdPrefix + "-" + correlationIdSequence.getAndIncrement();
-    this.requestPostProcessor =
-        (request, correlationId) -> request.replyTo(this.replyToAddress).messageId(correlationId);
+    if (builder.correlationIdSupplier() == null) {
+      String correlationIdPrefix = UUID.randomUUID().toString();
+      AtomicLong correlationIdSequence = new AtomicLong();
+      this.correlationIdSupplier =
+          () -> correlationIdPrefix + "-" + correlationIdSequence.getAndIncrement();
+    } else {
+      this.correlationIdSupplier = builder.correlationIdSupplier();
+    }
+
+    if (builder.requestPostProcessor() == null) {
+      DefaultAddressBuilder<?> addressBuilder = new DefaultAddressBuilder<>(null) {};
+      addressBuilder.queue(replyTo);
+      String replyToAddress = addressBuilder.address();
+      this.requestPostProcessor =
+          (request, correlationId) -> request.replyTo(replyToAddress).messageId(correlationId);
+    } else {
+      this.requestPostProcessor = builder.requestPostProcessor();
+    }
   }
 
   @Override
@@ -93,7 +108,7 @@ public CompletableFuture<Message> publish(Message message) {
     Object correlationId = this.correlationIdSupplier.get();
     message = this.requestPostProcessor.apply(message, correlationId);
     CompletableFuture<Message> result = new CompletableFuture<>();
-    this.outstandingRequests.put(message.messageId(), result);
+    this.outstandingRequests.put(correlationId, result);
     this.publisher.publish(message, NO_OP_CALLBACK);
     return result;
   }

src/main/java/com/rabbitmq/model/amqp/AmqpRpcServer.java

@@ -51,8 +51,16 @@ public Message message(byte[] body) {
             return publisher.message(body);
           }
         };
-    this.correlationIdExtractor = Message::messageId;
-    this.replyPostProcessor = Message::correlationId;
+    if (builder.correlationIdExtractor() == null) {
+      this.correlationIdExtractor = Message::messageId;
+    } else {
+      this.correlationIdExtractor = builder.correlationIdExtractor();
+    }
+    if (builder.replyPostProcessor() == null) {
+      this.replyPostProcessor = Message::correlationId;
+    } else {
+      this.replyPostProcessor = builder.replyPostProcessor();
+    }
     this.consumer =
         connection
             .consumerBuilder()
@@ -61,7 +69,9 @@ public Message message(byte[] body) {
                 (ctx, msg) -> {
                   ctx.accept();
                   Message reply = handler.handle(context, msg);
-                  reply.to(msg.replyTo());
+                  if (msg.replyTo() != null) {
+                    reply.to(msg.replyTo());
+                  }
                   Object correlationId = correlationIdExtractor.apply(msg);
                   reply = replyPostProcessor.apply(reply, correlationId);
                   this.publisher.publish(reply, NO_OP_CALLBACK);

src/main/java/com/rabbitmq/model/amqp/RpcSupport.java

@@ -18,6 +18,9 @@
 package com.rabbitmq.model.amqp;
 
 import com.rabbitmq.model.*;
+import java.util.function.BiFunction;
+import java.util.function.Function;
+import java.util.function.Supplier;
 
 abstract class RpcSupport {
 
@@ -30,6 +33,9 @@ static class AmqpRpcClientBuilder implements RpcClientBuilder {
     private final DefaultRpcClientAddressBuilder requestAddressBuilder =
         new DefaultRpcClientAddressBuilder(this);
     private String replyToQueue;
+    private Supplier<Object> correlationIdSupplier;
+    private BiFunction<Message, Object, Message> requestPostProcessor;
+    private Function<Message, Object> correlationIdExtractor;
 
     AmqpRpcClientBuilder(AmqpConnection connection) {
       this.connection = connection;
@@ -46,8 +52,35 @@ public RpcClientBuilder replyToQueue(String replyToQueue) {
       return this;
     }
 
+    @Override
+    public RpcClientBuilder correlationIdSupplier(Supplier<Object> correlationIdSupplier) {
+      this.correlationIdSupplier = correlationIdSupplier;
+      return this;
+    }
+
+    @Override
+    public RpcClientBuilder requestPostProcessor(
+        BiFunction<Message, Object, Message> requestPostProcessor) {
+      this.requestPostProcessor = requestPostProcessor;
+      return this;
+    }
+
+    @Override
+    public RpcClientBuilder correlationIdExtractor(
+        Function<Message, Object> correlationIdExtractor) {
+      this.correlationIdExtractor = correlationIdExtractor;
+      return this;
+    }
+
+    Function<Message, Object> correlationIdExtractor() {
+      return correlationIdExtractor;
+    }
+
     @Override
     public RpcClient build() {
+      if (this.requestAddressBuilder.address() == null) {
+        throw new IllegalArgumentException("Request address cannot be null");
+      }
       return new AmqpRpcClient(this);
     }
 
@@ -58,6 +91,14 @@ AmqpConnection connection() {
     String replyToQueue() {
       return this.replyToQueue;
     }
+
+    Supplier<Object> correlationIdSupplier() {
+      return this.correlationIdSupplier;
+    }
+
+    BiFunction<Message, Object, Message> requestPostProcessor() {
+      return this.requestPostProcessor;
+    }
   }
 
   private static class DefaultRpcClientAddressBuilder
@@ -90,6 +131,8 @@ static class AmqpRpcServerBuilder implements RpcServerBuilder {
     private RpcServer.Handler handler;
     private final DefaultRpcServerAddressBuilder replyToAddressBuilder =
         new DefaultRpcServerAddressBuilder(this);
+    private Function<Message, Object> correlationIdExtractor;
+    private BiFunction<Message, Object, Message> replyPostProcessor;
 
     AmqpRpcServerBuilder(AmqpConnection connection) {
       this.connection = connection;
@@ -112,6 +155,20 @@ public RpcServerAddressBuilder replyToAddress() {
       return this.replyToAddressBuilder;
     }
 
+    @Override
+    public RpcServerBuilder correlationIdExtractor(
+        Function<Message, Object> correlationIdExtractor) {
+      this.correlationIdExtractor = correlationIdExtractor;
+      return this;
+    }
+
+    @Override
+    public RpcServerBuilder replyPostProcessor(
+        BiFunction<Message, Object, Message> replyPostProcessor) {
+      this.replyPostProcessor = replyPostProcessor;
+      return this;
+    }
+
     @Override
     public RpcServer build() {
       return new AmqpRpcServer(this);
@@ -128,6 +185,14 @@ String requestQueue() {
     RpcServer.Handler handler() {
       return this.handler;
     }
+
+    Function<Message, Object> correlationIdExtractor() {
+      return this.correlationIdExtractor;
+    }
+
+    BiFunction<Message, Object, Message> replyPostProcessor() {
+      return this.replyPostProcessor;
+    }
   }
 
   private static class DefaultRpcServerAddressBuilder