|
@@ -40,6 +40,8 @@ import javax.management.ObjectName;
|
|
|
|
|
|
import com.fasterxml.jackson.databind.ObjectMapper;
|
|
|
import com.fasterxml.jackson.databind.ObjectWriter;
|
|
|
+
|
|
|
+import org.apache.hadoop.security.UserGroupInformation;
|
|
|
import org.apache.hadoop.thirdparty.com.google.common.base.Preconditions;
|
|
|
import org.apache.hadoop.thirdparty.com.google.common.util.concurrent.AtomicDoubleArray;
|
|
|
import org.apache.commons.lang3.exception.ExceptionUtils;
|
|
@@ -193,6 +195,7 @@ public class DecayRpcScheduler implements RpcScheduler,
|
|
|
private static final double PRECISION = 0.0001;
|
|
|
private MetricsProxy metricsProxy;
|
|
|
private final CostProvider costProvider;
|
|
|
+ private final Map<String, Integer> staticPriorities = new HashMap<>();
|
|
|
private Set<String> serviceUserNames;
|
|
|
|
|
|
/**
|
|
@@ -581,7 +584,10 @@ public class DecayRpcScheduler implements RpcScheduler,
|
|
|
if (isServiceUser((String)identity)) {
|
|
|
return 0;
|
|
|
}
|
|
|
-
|
|
|
+ Integer staticPriority = staticPriorities.get(identity);
|
|
|
+ if (staticPriority != null) {
|
|
|
+ return staticPriority.intValue();
|
|
|
+ }
|
|
|
long totalCallSnapshot = totalDecayedCallCost.get();
|
|
|
|
|
|
double proportion = 0;
|
|
@@ -626,6 +632,15 @@ public class DecayRpcScheduler implements RpcScheduler,
|
|
|
return priority;
|
|
|
}
|
|
|
|
|
|
+ private String getIdentity(Schedulable obj) {
|
|
|
+ String identity = this.identityProvider.makeIdentity(obj);
|
|
|
+ if (identity == null) {
|
|
|
+ // Identity provider did not handle this
|
|
|
+ identity = DECAYSCHEDULER_UNKNOWN_IDENTITY;
|
|
|
+ }
|
|
|
+ return identity;
|
|
|
+ }
|
|
|
+
|
|
|
/**
|
|
|
* Compute the appropriate priority for a schedulable based on past requests.
|
|
|
* @param obj the schedulable obj to query and remember
|
|
@@ -634,15 +649,42 @@ public class DecayRpcScheduler implements RpcScheduler,
|
|
|
@Override
|
|
|
public int getPriorityLevel(Schedulable obj) {
|
|
|
// First get the identity
|
|
|
- String identity = this.identityProvider.makeIdentity(obj);
|
|
|
- if (identity == null) {
|
|
|
- // Identity provider did not handle this
|
|
|
- identity = DECAYSCHEDULER_UNKNOWN_IDENTITY;
|
|
|
- }
|
|
|
+ String identity = getIdentity(obj);
|
|
|
+ // highest priority users may have a negative priority but their
|
|
|
+ // calls will be priority 0.
|
|
|
+ return Math.max(0, cachedOrComputedPriorityLevel(identity));
|
|
|
+ }
|
|
|
|
|
|
+ @VisibleForTesting
|
|
|
+ int getPriorityLevel(UserGroupInformation ugi) {
|
|
|
+ String identity = getIdentity(newSchedulable(ugi));
|
|
|
+ // returns true priority of the user.
|
|
|
return cachedOrComputedPriorityLevel(identity);
|
|
|
}
|
|
|
|
|
|
+ @VisibleForTesting
|
|
|
+ void setPriorityLevel(UserGroupInformation ugi, int priority) {
|
|
|
+ String identity = getIdentity(newSchedulable(ugi));
|
|
|
+ priority = Math.min(numLevels - 1, priority);
|
|
|
+ LOG.info("Setting priority for user:" + identity + "=" + priority);
|
|
|
+ staticPriorities.put(identity, priority);
|
|
|
+ }
|
|
|
+
|
|
|
+ // dummy instance to conform to identity provider api.
|
|
|
+ private static Schedulable newSchedulable(UserGroupInformation ugi) {
|
|
|
+ return new Schedulable() {
|
|
|
+ @Override
|
|
|
+ public UserGroupInformation getUserGroupInformation() {
|
|
|
+ return ugi;
|
|
|
+ }
|
|
|
+
|
|
|
+ @Override
|
|
|
+ public int getPriorityLevel() {
|
|
|
+ return 0;
|
|
|
+ }
|
|
|
+ };
|
|
|
+ }
|
|
|
+
|
|
|
private boolean isServiceUser(String userName) {
|
|
|
return this.serviceUserNames.contains(userName);
|
|
|
}
|