66import datadog .trace .api .sampling .SamplingMechanism ;
77import datadog .trace .common .writer .RemoteResponseListener ;
88import datadog .trace .core .CoreSpan ;
9- import java .util .HashMap ;
9+ import java .util .Collections ;
1010import java .util .Map ;
11+ import java .util .TreeMap ;
1112import java .util .function .Function ;
1213import org .slf4j .Logger ;
1314import org .slf4j .LoggerFactory ;
@@ -71,8 +72,8 @@ public void onResponse(
7172 }
7273
7374 log .debug ("Update service sampler rates: {} -> {}" , endpoint , responseJson );
74- final Map <String , Map <String , RateSampler >> updatedEnvServiceRates =
75- new HashMap <>(newServiceRates . size () * 2 );
75+ final TreeMap <String , TreeMap <String , RateSampler >> updatedEnvServiceRates =
76+ new TreeMap <>(String :: compareToIgnoreCase );
7677
7778 RateSampler fallbackSampler = RateSamplersByEnvAndService .DEFAULT_SAMPLER ;
7879 for (final Map .Entry <String , Number > entry : newServiceRates .entrySet ()) {
@@ -87,7 +88,7 @@ public void onResponse(
8788 } else {
8889 Map <String , RateSampler > serviceRates =
8990 updatedEnvServiceRates .computeIfAbsent (
90- envAndService .lowerEnv , env -> new HashMap <>(newServiceRates . size () * 2 ));
91+ envAndService .lowerEnv , env -> new TreeMap <>(String :: compareToIgnoreCase ));
9192
9293 serviceRates .computeIfAbsent (
9394 envAndService .lowerService ,
@@ -114,38 +115,34 @@ private static RateSampler createRateSampler(final double sampleRate) {
114115 private static final class RateSamplersByEnvAndService {
115116 private static final RateSampler DEFAULT_SAMPLER = createRateSampler (DEFAULT_RATE );
116117
117- private final Map <String , Map <String , RateSampler >> envServiceRates ;
118+ private final Map <String , TreeMap <String , RateSampler >> envServiceRates ;
118119 private final RateSampler fallbackSampler ;
119120
120121 RateSamplersByEnvAndService () {
121- this (new HashMap <>( 0 ), DEFAULT_SAMPLER );
122+ this (Collections . emptyMap ( ), DEFAULT_SAMPLER );
122123 }
123124
124125 RateSamplersByEnvAndService (
125- Map <String , Map <String , RateSampler >> envServiceRates , RateSampler fallbackSampler ) {
126+ Map <String , TreeMap <String , RateSampler >> envServiceRates , RateSampler fallbackSampler ) {
126127 this .envServiceRates = envServiceRates ;
127128 this .fallbackSampler = fallbackSampler ;
128129 }
129130
130131 // used in tests only
131132 RateSampler getSampler (EnvAndService envAndService ) {
132- return getSamplerImpl (envAndService .lowerEnv , envAndService .lowerService );
133+ return getSampler (envAndService .lowerEnv , envAndService .lowerService );
133134 }
134135
135136 public RateSampler getSampler (String env , String service ) {
136- return getSamplerImpl (env .toLowerCase (), service .toLowerCase ());
137- }
138-
139- private RateSampler getSamplerImpl (String lowerEnv , String lowerService ) {
140- if (EnvAndService .isFallback (lowerEnv , lowerService )) {
137+ if (EnvAndService .isFallback (env , service )) {
141138 return fallbackSampler ;
142139 }
143140
144- Map <String , RateSampler > serviceRates = envServiceRates .get (lowerEnv );
141+ Map <String , RateSampler > serviceRates = envServiceRates .get (env );
145142 if (serviceRates == null ) {
146143 return fallbackSampler ;
147144 }
148- RateSampler sampler = serviceRates .get (lowerService );
145+ RateSampler sampler = serviceRates .get (service );
149146 return null == sampler ? fallbackSampler : sampler ;
150147 }
151148 }
0 commit comments