use threadlocal ringhash key provider

pull/970/head
seanyu 2 years ago
parent d0f50a21d7
commit 6a4a3fe94f

@ -17,6 +17,8 @@
package com.tencent.cloud.polaris.loadbalancer; package com.tencent.cloud.polaris.loadbalancer;
import com.tencent.cloud.common.metadata.MetadataContextHolder;
/** /**
* PolarisLoadBalancerRingHashKeyProvider. * PolarisLoadBalancerRingHashKeyProvider.
* *
@ -24,22 +26,17 @@ package com.tencent.cloud.polaris.loadbalancer;
*/ */
public final class PolarisLoadBalancerRingHashKeyProvider { public final class PolarisLoadBalancerRingHashKeyProvider {
private static final ThreadLocal<String> hashKeys = new ThreadLocal<>(); private static final String LOAD_BALANCER_HASH_KEY = "LOAD_BALANCER_HASH_KEY";
private PolarisLoadBalancerRingHashKeyProvider() { private PolarisLoadBalancerRingHashKeyProvider() {
} }
public static void hashKey(String key) { public static void hashKey(String key) {
hashKeys.set(key); MetadataContextHolder.get().setLoadbalancer(LOAD_BALANCER_HASH_KEY, key);
}
static void remove() {
hashKeys.remove();
} }
static String getHashKey() { static String getHashKey() {
return hashKeys.get(); return MetadataContextHolder.get().getLoadbalancerMetadata().get(LOAD_BALANCER_HASH_KEY);
} }
} }

@ -81,10 +81,11 @@ public class PolarisRingHashLoadBalancer implements ReactorServiceInstanceLoadBa
public Mono<Response<ServiceInstance>> choose(Request request) { public Mono<Response<ServiceInstance>> choose(Request request) {
ServiceInstanceListSupplier supplier = supplierObjectProvider ServiceInstanceListSupplier supplier = supplierObjectProvider
.getIfAvailable(NoopServiceInstanceListSupplier::new); .getIfAvailable(NoopServiceInstanceListSupplier::new);
return supplier.get(request).next().map(this::getInstanceResponse); String hashKey = Optional.ofNullable(PolarisLoadBalancerRingHashKeyProvider.getHashKey()).orElse("");
return supplier.get(request).next().map(serviceInstances -> getInstanceResponse(serviceInstances, hashKey));
} }
private Response<ServiceInstance> getInstanceResponse(List<ServiceInstance> serviceInstances) { private Response<ServiceInstance> getInstanceResponse(List<ServiceInstance> serviceInstances, String hashKey) {
if (serviceInstances.isEmpty()) { if (serviceInstances.isEmpty()) {
log.warn("No servers available for service: " + this.serviceId); log.warn("No servers available for service: " + this.serviceId);
return new EmptyResponse(); return new EmptyResponse();
@ -94,12 +95,11 @@ public class PolarisRingHashLoadBalancer implements ReactorServiceInstanceLoadBa
request.setDstInstances(convertToPolarisServiceInstances(serviceInstances)); request.setDstInstances(convertToPolarisServiceInstances(serviceInstances));
request.setLbPolicy(LoadBalanceConfig.LOAD_BALANCE_RING_HASH); request.setLbPolicy(LoadBalanceConfig.LOAD_BALANCE_RING_HASH);
Criteria criteria = new Criteria(); Criteria criteria = new Criteria();
criteria.setHashKey(Optional.ofNullable(PolarisLoadBalancerRingHashKeyProvider.getHashKey()).orElse("")); criteria.setHashKey(hashKey);
request.setCriteria(criteria); request.setCriteria(criteria);
try { try {
ProcessLoadBalanceResponse response = routerAPI.processLoadBalance(request); ProcessLoadBalanceResponse response = routerAPI.processLoadBalance(request);
PolarisLoadBalancerRingHashKeyProvider.remove();
return new DefaultResponse(new PolarisServiceInstance(response.getTargetInstance())); return new DefaultResponse(new PolarisServiceInstance(response.getTargetInstance()));
} }
catch (Exception e) { catch (Exception e) {

Loading…
Cancel
Save