From 12a4647ecf8c3880bff5268104bb1a003853ba94 Mon Sep 17 00:00:00 2001
From: seanyu <seansyyu@tencent.com>
Date: Thu, 13 Apr 2023 19:44:29 +0800
Subject: [PATCH] use threadlocal ringhash key provider

---
 .../PolarisLoadBalancerRingHashKeyProvider.java     | 13 +++++--------
 .../loadbalancer/PolarisRingHashLoadBalancer.java   |  8 ++++----
 2 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/spring-cloud-starter-tencent-polaris-discovery/src/main/java/com/tencent/cloud/polaris/loadbalancer/PolarisLoadBalancerRingHashKeyProvider.java b/spring-cloud-starter-tencent-polaris-discovery/src/main/java/com/tencent/cloud/polaris/loadbalancer/PolarisLoadBalancerRingHashKeyProvider.java
index 42a203cab..7725523f6 100644
--- a/spring-cloud-starter-tencent-polaris-discovery/src/main/java/com/tencent/cloud/polaris/loadbalancer/PolarisLoadBalancerRingHashKeyProvider.java
+++ b/spring-cloud-starter-tencent-polaris-discovery/src/main/java/com/tencent/cloud/polaris/loadbalancer/PolarisLoadBalancerRingHashKeyProvider.java
@@ -17,6 +17,8 @@
 
 package com.tencent.cloud.polaris.loadbalancer;
 
+import com.tencent.cloud.common.metadata.MetadataContextHolder;
+
 /**
  * PolarisLoadBalancerRingHashKeyProvider.
  *
@@ -24,22 +26,17 @@ package com.tencent.cloud.polaris.loadbalancer;
  */
 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() {
 	}
 
 	public static void hashKey(String key) {
-		hashKeys.set(key);
-	}
-
-	static void remove() {
-		hashKeys.remove();
+		MetadataContextHolder.get().setLoadbalancer(LOAD_BALANCER_HASH_KEY, key);
 	}
 
 	static String getHashKey() {
-		return hashKeys.get();
+		return MetadataContextHolder.get().getLoadbalancerMetadata().get(LOAD_BALANCER_HASH_KEY);
 	}
 
-
 }
\ No newline at end of file
diff --git a/spring-cloud-starter-tencent-polaris-discovery/src/main/java/com/tencent/cloud/polaris/loadbalancer/PolarisRingHashLoadBalancer.java b/spring-cloud-starter-tencent-polaris-discovery/src/main/java/com/tencent/cloud/polaris/loadbalancer/PolarisRingHashLoadBalancer.java
index ba770f072..3f4603c05 100644
--- a/spring-cloud-starter-tencent-polaris-discovery/src/main/java/com/tencent/cloud/polaris/loadbalancer/PolarisRingHashLoadBalancer.java
+++ b/spring-cloud-starter-tencent-polaris-discovery/src/main/java/com/tencent/cloud/polaris/loadbalancer/PolarisRingHashLoadBalancer.java
@@ -81,10 +81,11 @@ public class PolarisRingHashLoadBalancer implements ReactorServiceInstanceLoadBa
 	public Mono<Response<ServiceInstance>> choose(Request request) {
 		ServiceInstanceListSupplier supplier = supplierObjectProvider
 				.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()) {
 			log.warn("No servers available for service: " + this.serviceId);
 			return new EmptyResponse();
@@ -94,12 +95,11 @@ public class PolarisRingHashLoadBalancer implements ReactorServiceInstanceLoadBa
 		request.setDstInstances(convertToPolarisServiceInstances(serviceInstances));
 		request.setLbPolicy(LoadBalanceConfig.LOAD_BALANCE_RING_HASH);
 		Criteria criteria = new Criteria();
-		criteria.setHashKey(Optional.ofNullable(PolarisLoadBalancerRingHashKeyProvider.getHashKey()).orElse(""));
+		criteria.setHashKey(hashKey);
 		request.setCriteria(criteria);
 
 		try {
 			ProcessLoadBalanceResponse response = routerAPI.processLoadBalance(request);
-			PolarisLoadBalancerRingHashKeyProvider.remove();
 			return new DefaultResponse(new PolarisServiceInstance(response.getTargetInstance()));
 		}
 		catch (Exception e) {