diff --git a/spring-cloud-gateway-server/src/main/java/org/springframework/cloud/gateway/filter/WeightCalculatorWebFilter.java b/spring-cloud-gateway-server/src/main/java/org/springframework/cloud/gateway/filter/WeightCalculatorWebFilter.java index 977a64fb..a61e2e69 100644 --- a/spring-cloud-gateway-server/src/main/java/org/springframework/cloud/gateway/filter/WeightCalculatorWebFilter.java +++ b/spring-cloud-gateway-server/src/main/java/org/springframework/cloud/gateway/filter/WeightCalculatorWebFilter.java @@ -25,6 +25,7 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ThreadLocalRandom; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Function; import java.util.function.Supplier; import org.apache.commons.logging.Log; @@ -42,6 +43,7 @@ import org.springframework.context.ApplicationEvent; import org.springframework.context.event.SmartApplicationListener; import org.springframework.core.Ordered; import org.springframework.core.style.ToStringCreator; +import org.springframework.util.Assert; import org.springframework.web.server.ServerWebExchange; import org.springframework.web.server.WebFilter; import org.springframework.web.server.WebFilterChain; @@ -65,7 +67,7 @@ public class WeightCalculatorWebFilter implements WebFilter, Ordered, SmartAppli private final ConfigurationService configurationService; - private Supplier randomSupplier = null; + private Function randomFunction = exchange -> ThreadLocalRandom.current().nextDouble(); private int order = WEIGHT_CALC_FILTER_ORDER; @@ -101,11 +103,18 @@ public class WeightCalculatorWebFilter implements WebFilter, Ordered, SmartAppli @Deprecated public void setRandom(Random random) { - this.randomSupplier = random::nextDouble; + Assert.notNull(random, "random may not be null"); + this.randomFunction = exchange -> random.nextDouble(); } public void setRandomSupplier(Supplier randomSupplier) { - this.randomSupplier = randomSupplier; + Assert.notNull(randomSupplier, "randomSupplier may not be null"); + this.randomFunction = exchange -> randomSupplier.get(); + } + + public void setRandomFunction(Function randomFunction) { + Assert.notNull(randomFunction, "randomFunction may not be null"); + this.randomFunction = randomFunction; } @Override @@ -238,17 +247,9 @@ public class WeightCalculatorWebFilter implements WebFilter, Ordered, SmartAppli continue; // nothing we can do, but this is odd } - /* - * Usually, multiple threads accessing the same random object will have some - * performance problems, so we can use ThreadLocalRandom by default - */ - double r; - if (this.randomSupplier != null) { - r = randomSupplier.get(); - } - else { - r = ThreadLocalRandom.current().nextDouble(); - } + // Usually, multiple threads accessing the same random object will have some + // performance problems, so we can use ThreadLocalRandom by default + double r = randomFunction.apply(exchange); List ranges = config.ranges;