From 6669e5265fb5c29a574d1e0f18941a2521f1f49c Mon Sep 17 00:00:00 2001
From: Anoop Panicker <apanicker@netflix.com>
Date: Tue, 15 Jan 2019 15:03:03 -0800
Subject: [PATCH] fix for token supplier provider with multiple dyno hosts

---
 .../jedis/TokenMapSupplierProvider.java       | 44 ++++++++++++-------
 1 file changed, 27 insertions(+), 17 deletions(-)

diff --git a/redis-persistence/src/main/java/com/netflix/conductor/jedis/TokenMapSupplierProvider.java b/redis-persistence/src/main/java/com/netflix/conductor/jedis/TokenMapSupplierProvider.java
index 627fed0729..92a7ef207f 100644
--- a/redis-persistence/src/main/java/com/netflix/conductor/jedis/TokenMapSupplierProvider.java
+++ b/redis-persistence/src/main/java/com/netflix/conductor/jedis/TokenMapSupplierProvider.java
@@ -1,43 +1,53 @@
+/*
+ * Copyright 2016 Netflix, Inc.
+ * <p>
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+ * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations under the License.
+ */
 package com.netflix.conductor.jedis;
 
-import com.google.common.collect.Lists;
-
 import com.netflix.dyno.connectionpool.Host;
-import com.netflix.dyno.connectionpool.HostSupplier;
 import com.netflix.dyno.connectionpool.TokenMapSupplier;
 import com.netflix.dyno.connectionpool.impl.lb.HostToken;
-
-import java.util.Arrays;
-import java.util.List;
-import java.util.Set;
+import com.netflix.dyno.connectionpool.impl.utils.CollectionUtils;
 
 import javax.inject.Inject;
 import javax.inject.Provider;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
 
 public class TokenMapSupplierProvider implements Provider<TokenMapSupplier> {
-    private final HostSupplier hostSupplier;
+    private final List<HostToken> hostTokens;
 
     @Inject
-    public TokenMapSupplierProvider(HostSupplier hostSupplier) {
-        this.hostSupplier = hostSupplier;
+    public TokenMapSupplierProvider() {
+        this.hostTokens = new ArrayList<>();
     }
 
     @Override
     public TokenMapSupplier get() {
         return new TokenMapSupplier() {
-
-            // FIXME This isn't particularly safe, but it is equivalent to the existing code.
-            // FIXME It seems like we should be supply tokens for more than one host?
-            HostToken token = new HostToken(1L, Lists.newArrayList(hostSupplier.getHosts()).get(0));
-
             @Override
             public List<HostToken> getTokens(Set<Host> activeHosts) {
-                return Arrays.asList(token);
+                long i = activeHosts.size();
+                for (Host host : activeHosts) {
+                    HostToken hostToken = new HostToken(i, host);
+                    hostTokens.add(hostToken);
+                    i--;
+                }
+                return hostTokens;
             }
 
             @Override
             public HostToken getTokenForHost(Host host, Set<Host> activeHosts) {
-                return token;
+                return CollectionUtils.find(hostTokens, token -> token.getHost().compareTo(host) == 0);
             }
         };
     }