Add per host queues.

palantir · Jul 16, 2021 · 0f2942a · 0f2942a
1 parent 1972ca2
commit 0f2942a
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 6 deletions.
diff --git a/dialogue-core/src/main/java/com/palantir/dialogue/core/DialogueChannel.java b/dialogue-core/src/main/java/com/palantir/dialogue/core/DialogueChannel.java
@@ -220,19 +220,26 @@ public DialogueChannel build() {
 
             ImmutableList<Channel> perHostChannels = IntStream.range(0, channels.size())
                     .mapToObj(index -> new ChannelAndFactory() {
+
+                        private final LimitedChannel stickyLimitedChannel =
+                                StickyConcurrencyLimitedChannel.createForQueueKey(
+                                        nodeSelectionChannel, cf.channelName());
+                        private final Channel queueOverride =
+                                QueuedChannel.createPerHost(cf, stickyLimitedChannel, index);
+
                         @Override
                         public EndpointChannel endpoint(Endpoint endpoint) {
                             EndpointChannel endpointChannel = channelFactory.endpoint(endpoint);
                             return request -> {
+                                request.attachments().put(QueueAttachments.QUEUE_OVERRIDE, queueOverride);
                                 nodeSelectionChannel.routeToHost(index, request);
                                 return endpointChannel.execute(request);
                             };
                         }
 
                         @Override
                         public ListenableFuture<Response> execute(Endpoint endpoint, Request request) {
-                            nodeSelectionChannel.routeToHost(index, request);
-                            return channelFactory.endpoint(endpoint).execute(request);
+                            return endpoint(endpoint).execute(request);
                         }
                     })
                     .collect(ImmutableList.toImmutableList());

diff --git a/dialogue-core/src/main/java/com/palantir/dialogue/core/QueuedChannel.java b/dialogue-core/src/main/java/com/palantir/dialogue/core/QueuedChannel.java
@@ -106,8 +106,8 @@ static QueuedChannel createPerHost(Config cf, LimitedChannel delegate, int hostI
         return new QueuedChannel(
                 delegate,
                 cf.channelName(),
-                stickyInstrumentation(
-                        DialogueClientMetrics.of(cf.clientConf().taggedMetricRegistry()), cf.channelName()),
+                perHostInstrumentation(
+                        DialogueClientMetrics.of(cf.clientConf().taggedMetricRegistry()), cf.channelName(), hostIndex),
                 cf.maxQueueSize());
     }
 
@@ -401,15 +401,22 @@ public Timer requestQueuedTime() {
 
     static QueuedChannelInstrumentation perHostInstrumentation(
             DialogueClientMetrics metrics, String channelName, int hostIndex) {
+        String hostIndexString = Integer.toString(hostIndex);
         return new QueuedChannelInstrumentation() {
             @Override
             public Counter requestsQueued() {
-                return metrics.requestsPerhostQueued().channelName(channelName).hostIndex(hostIndex);
+                return metrics.requestsPerhostQueued()
+                        .channelName(channelName)
+                        .hostIndex(hostIndexString)
+                        .build();
             }
 
             @Override
             public Timer requestQueuedTime() {
-                return metrics.requestStickyQueuedTime(channelName);
+                return metrics.requestPerhostQueuedTime()
+                        .channelName(channelName)
+                        .hostIndex(hostIndexString)
+                        .build();
             }
         };
     }