storm-client/src/jvm/org/apache/storm/grouping/LoadAwareShuffleGrouping.java - storm - Git at Google

 /**
  * Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements.  See the NOTICE file distributed with
  * this work for additional information regarding copyright ownership.  The ASF licenses this file to you under the Apache License, Version
  * 2.0 (the "License"); you may not use this file except in compliance with the License.  You may obtain a copy of the License at
  *
  * http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
  * and limitations under the License.
  */

 package org.apache.storm.grouping;

 import java.io.Serializable;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Random;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import org.apache.storm.Config;
 import org.apache.storm.generated.GlobalStreamId;
 import org.apache.storm.generated.NodeInfo;
 import org.apache.storm.networktopography.DNSToSwitchMapping;
 import org.apache.storm.shade.com.google.common.annotations.VisibleForTesting;
 import org.apache.storm.shade.com.google.common.collect.Sets;
 import org.apache.storm.task.WorkerTopologyContext;
 import org.apache.storm.utils.ObjectReader;
 import org.apache.storm.utils.ReflectionUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 public class LoadAwareShuffleGrouping implements LoadAwareCustomStreamGrouping, Serializable {
     private static final int MAX_WEIGHT = 100;
     private static final Logger LOG = LoggerFactory.getLogger(LoadAwareShuffleGrouping.class);
     private final Map<Integer, IndexAndWeights> orig = new HashMap<>();
     @VisibleForTesting
     List<Integer>[] rets;
     @VisibleForTesting
     volatile int[] choices;
     private int capacity;
     private Random random;
     private volatile int[] prepareChoices;
     private AtomicInteger current;
     private Scope currentScope;
     private NodeInfo sourceNodeInfo;
     private List<Integer> targetTasks;
     private AtomicReference<Map<Integer, NodeInfo>> taskToNodePort;
     private Map<String, Object> conf;
     private DNSToSwitchMapping dnsToSwitchMapping;
     private Map<Scope, List<Integer>> localityGroup;
     private double higherBound;
     private double lowerBound;

     @Override
     public void prepare(WorkerTopologyContext context, GlobalStreamId stream, List<Integer> targetTasks) {
         random = new Random();
         sourceNodeInfo = new NodeInfo(context.getThisWorkerHost(), Sets.newHashSet((long) context.getThisWorkerPort()));
         taskToNodePort = context.getTaskToNodePort();
         this.targetTasks = targetTasks;
         capacity = targetTasks.size() == 1 ? 1 : Math.max(1000, targetTasks.size() * 5);
         conf = context.getConf();
         dnsToSwitchMapping = ReflectionUtils.newInstance((String) conf.get(Config.STORM_NETWORK_TOPOGRAPHY_PLUGIN));
         localityGroup = new HashMap<>();
         currentScope = Scope.WORKER_LOCAL;
         higherBound = ObjectReader.getDouble(conf.get(Config.TOPOLOGY_LOCALITYAWARE_HIGHER_BOUND));
         lowerBound = ObjectReader.getDouble(conf.get(Config.TOPOLOGY_LOCALITYAWARE_LOWER_BOUND));

         rets = (List<Integer>[]) new List<?>[targetTasks.size()];
         int i = 0;
         for (int target : targetTasks) {
             rets[i] = Arrays.asList(target);
             orig.put(target, new IndexAndWeights(i));
             i++;
         }

         // can't leave choices to be empty, so initiate it similar as ShuffleGrouping
         choices = new int[capacity];

         current = new AtomicInteger(0);
         // allocate another array to be switched
         prepareChoices = new int[capacity];
         updateRing(null);
     }

     @Override
     public List<Integer> chooseTasks(int taskId, List<Object> values) {
         int rightNow;
         while (true) {
             rightNow = current.incrementAndGet();
             if (rightNow < capacity) {
                 return rets[choices[rightNow]];
             } else if (rightNow == capacity) {
                 current.set(0);
                 return rets[choices[0]];
             }
             //race condition with another thread, and we lost
             // try again
         }
     }

     @Override
     public void refreshLoad(LoadMapping loadMapping) {
         updateRing(loadMapping);
     }

     private void refreshLocalityGroup() {
         Map<Integer, NodeInfo> cachedTaskToNodePort = taskToNodePort.get();
         Map<String, String> hostToRack = getHostToRackMapping(cachedTaskToNodePort);

         localityGroup.values().stream().forEach(v -> v.clear());

         for (int target : targetTasks) {
             Scope scope = calculateScope(cachedTaskToNodePort, hostToRack, target);
             if (!localityGroup.containsKey(scope)) {
                 localityGroup.put(scope, new ArrayList<>());
             }
             localityGroup.get(scope).add(target);
         }
     }

     private List<Integer> getTargetsInScope(Scope scope) {
         List<Integer> rets = new ArrayList<>();
         List<Integer> targetInScope = localityGroup.get(scope);
         if (null != targetInScope) {
             rets.addAll(targetInScope);
         }
         Scope downgradeScope = Scope.downgrade(scope);
         if (downgradeScope != scope) {
             rets.addAll(getTargetsInScope(downgradeScope));
         }
         return rets;
     }

     private Scope transition(LoadMapping load) {
         List<Integer> targetInScope = getTargetsInScope(currentScope);
         if (targetInScope.isEmpty()) {
             Scope upScope = Scope.upgrade(currentScope);
             if (upScope == currentScope) {
                 throw new RuntimeException("The current scope " + currentScope + " has no target tasks.");
             }
             currentScope = upScope;
             return transition(load);
         }

         if (null == load) {
             return currentScope;
         }

         double avg = targetInScope.stream().mapToDouble((key) -> load.get(key)).average().getAsDouble();
         Scope nextScope;
         if (avg < lowerBound) {
             nextScope = Scope.downgrade(currentScope);
             if (getTargetsInScope(nextScope).isEmpty()) {
                 nextScope = currentScope;
             }
         } else if (avg > higherBound) {
             nextScope = Scope.upgrade(currentScope);
         } else {
             nextScope = currentScope;
         }

         return nextScope;
     }

     private synchronized void updateRing(LoadMapping load) {
         refreshLocalityGroup();
         Scope prevScope = currentScope;
         currentScope = transition(load);
         if (currentScope != prevScope) {
             //reset all the weights
             orig.values().stream().forEach(o -> o.resetWeight());
         }

         List<Integer> targetsInScope = getTargetsInScope(currentScope);

         //We will adjust weights based off of the minimum load
         double min = load == null ? 0 : targetsInScope.stream().mapToDouble((key) -> load.get(key)).min().getAsDouble();
         for (int target : targetsInScope) {
             IndexAndWeights val = orig.get(target);
             double l = load == null ? 0.0 : load.get(target);
             if (l <= min + (0.05)) {
                 //We assume that within 5% of the minimum congestion is still fine.
                 //Not congested we grow (but slowly)
                 val.weight = Math.min(MAX_WEIGHT, val.weight + 1);
             } else {
                 //Congested we contract much more quickly
                 val.weight = Math.max(0, val.weight - 10);
             }
         }
         //Now we need to build the array
         long weightSum = targetsInScope.stream().mapToLong((target) -> orig.get(target).weight).sum();
         //Now we can calculate a percentage

         int currentIdx = 0;
         if (weightSum > 0) {
             for (int target : targetsInScope) {
                 IndexAndWeights indexAndWeights = orig.get(target);
                 int count = (int) ((indexAndWeights.weight / (double) weightSum) * capacity);
                 for (int i = 0; i < count && currentIdx < capacity; i++) {
                     prepareChoices[currentIdx] = indexAndWeights.index;
                     currentIdx++;
                 }
             }

             if (currentIdx > 0) {
                 //in case we didn't fill in enough
                 for (; currentIdx < capacity; currentIdx++) {
                     prepareChoices[currentIdx] = prepareChoices[random.nextInt(currentIdx)];
                 }
             }
         }
         if (currentIdx == 0) {
             //This really should be impossible, because we go off of the min load, and inc anything within 5% of it.
             // But just to be sure it is never an issue, especially with float rounding etc.
             for (; currentIdx < capacity; currentIdx++) {
                 prepareChoices[currentIdx] = currentIdx % rets.length;
             }
         }

         shuffleArray(prepareChoices);

         // swapping two arrays
         int[] tempForSwap = choices;
         choices = prepareChoices;
         prepareChoices = tempForSwap;

         current.set(-1);
     }

     private void shuffleArray(int[] arr) {
         int size = arr.length;
         for (int i = size; i > 1; i--) {
             swap(arr, i - 1, random.nextInt(i));
         }
     }

     private void swap(int[] arr, int i, int j) {
         int tmp = arr[i];
         arr[i] = arr[j];
         arr[j] = tmp;
     }

     private Scope calculateScope(Map<Integer, NodeInfo> taskToNodePort, Map<String, String> hostToRack, int target) {
         NodeInfo targetNodeInfo = taskToNodePort.get(target);

         if (targetNodeInfo == null) {
             return Scope.EVERYTHING;
         }

         String sourceRack = hostToRack.get(sourceNodeInfo.get_node());
         String targetRack = hostToRack.get(targetNodeInfo.get_node());

         if (sourceRack != null && targetRack != null && sourceRack.equals(targetRack)) {
             if (sourceNodeInfo.get_node().equals(targetNodeInfo.get_node())) {
                 if (sourceNodeInfo.get_port().equals(targetNodeInfo.get_port())) {
                     return Scope.WORKER_LOCAL;
                 }
                 return Scope.HOST_LOCAL;
             }
             return Scope.RACK_LOCAL;
         } else {
             return Scope.EVERYTHING;
         }
     }

     private Map<String, String> getHostToRackMapping(Map<Integer, NodeInfo> taskToNodePort) {
         Set<String> hosts = new HashSet();
         for (int task : targetTasks) {
             //if this task containing worker will be killed by a assignments sync,
             //taskToNodePort will be an empty map which is refreshed by WorkerState
             if (taskToNodePort.containsKey(task)) {
                 hosts.add(taskToNodePort.get(task).get_node());
             } else {
                 LOG.error("Could not find task NodeInfo from local cache.");
             }
         }
         hosts.add(sourceNodeInfo.get_node());
         return dnsToSwitchMapping.resolve(new ArrayList<>(hosts));
     }

     //only for test
     public int getCapacity() {
         return capacity;
     }

     enum Scope {
         WORKER_LOCAL, HOST_LOCAL, RACK_LOCAL, EVERYTHING;

         public static Scope downgrade(Scope current) {
             switch (current) {
                 case EVERYTHING:
                     return RACK_LOCAL;
                 case RACK_LOCAL:
                     return HOST_LOCAL;
                 case HOST_LOCAL:
                 case WORKER_LOCAL:
                 default:
                     return WORKER_LOCAL;
             }
         }

         public static Scope upgrade(Scope current) {
             switch (current) {
                 case WORKER_LOCAL:
                     return HOST_LOCAL;
                 case HOST_LOCAL:
                     return RACK_LOCAL;
                 case RACK_LOCAL:
                 case EVERYTHING:
                 default:
                     return EVERYTHING;
             }
         }
     }

     private static class IndexAndWeights {
         final int index;
         int weight;

         IndexAndWeights(int index) {
             this.index = index;
             weight = MAX_WEIGHT;
         }

         void resetWeight() {
             weight = MAX_WEIGHT;
         }
     }
 }
	/**
	* Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with
	* this work for additional information regarding copyright ownership. The ASF licenses this file to you under the Apache License, Version
	* 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS,
	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
	* and limitations under the License.
	*/

	package org.apache.storm.grouping;

	import java.io.Serializable;
	import java.util.ArrayList;
	import java.util.Arrays;
	import java.util.HashMap;
	import java.util.HashSet;
	import java.util.List;
	import java.util.Map;
	import java.util.Random;
	import java.util.Set;
	import java.util.concurrent.atomic.AtomicInteger;
	import java.util.concurrent.atomic.AtomicReference;
	import org.apache.storm.Config;
	import org.apache.storm.generated.GlobalStreamId;
	import org.apache.storm.generated.NodeInfo;
	import org.apache.storm.networktopography.DNSToSwitchMapping;
	import org.apache.storm.shade.com.google.common.annotations.VisibleForTesting;
	import org.apache.storm.shade.com.google.common.collect.Sets;
	import org.apache.storm.task.WorkerTopologyContext;
	import org.apache.storm.utils.ObjectReader;
	import org.apache.storm.utils.ReflectionUtils;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	public class LoadAwareShuffleGrouping implements LoadAwareCustomStreamGrouping, Serializable {
	private static final int MAX_WEIGHT = 100;
	private static final Logger LOG = LoggerFactory.getLogger(LoadAwareShuffleGrouping.class);
	private final Map<Integer, IndexAndWeights> orig = new HashMap<>();
	@VisibleForTesting
	List<Integer>[] rets;
	@VisibleForTesting
	volatile int[] choices;
	private int capacity;
	private Random random;
	private volatile int[] prepareChoices;
	private AtomicInteger current;
	private Scope currentScope;
	private NodeInfo sourceNodeInfo;
	private List<Integer> targetTasks;
	private AtomicReference<Map<Integer, NodeInfo>> taskToNodePort;
	private Map<String, Object> conf;
	private DNSToSwitchMapping dnsToSwitchMapping;
	private Map<Scope, List<Integer>> localityGroup;
	private double higherBound;
	private double lowerBound;

	@Override
	public void prepare(WorkerTopologyContext context, GlobalStreamId stream, List<Integer> targetTasks) {
	random = new Random();
	sourceNodeInfo = new NodeInfo(context.getThisWorkerHost(), Sets.newHashSet((long) context.getThisWorkerPort()));
	taskToNodePort = context.getTaskToNodePort();
	this.targetTasks = targetTasks;
	capacity = targetTasks.size() == 1 ? 1 : Math.max(1000, targetTasks.size() * 5);
	conf = context.getConf();
	dnsToSwitchMapping = ReflectionUtils.newInstance((String) conf.get(Config.STORM_NETWORK_TOPOGRAPHY_PLUGIN));
	localityGroup = new HashMap<>();
	currentScope = Scope.WORKER_LOCAL;
	higherBound = ObjectReader.getDouble(conf.get(Config.TOPOLOGY_LOCALITYAWARE_HIGHER_BOUND));
	lowerBound = ObjectReader.getDouble(conf.get(Config.TOPOLOGY_LOCALITYAWARE_LOWER_BOUND));

	rets = (List<Integer>[]) new List<?>[targetTasks.size()];
	int i = 0;
	for (int target : targetTasks) {
	rets[i] = Arrays.asList(target);
	orig.put(target, new IndexAndWeights(i));
	i++;
	}

	// can't leave choices to be empty, so initiate it similar as ShuffleGrouping
	choices = new int[capacity];

	current = new AtomicInteger(0);
	// allocate another array to be switched
	prepareChoices = new int[capacity];
	updateRing(null);
	}

	@Override
	public List<Integer> chooseTasks(int taskId, List<Object> values) {
	int rightNow;
	while (true) {
	rightNow = current.incrementAndGet();
	if (rightNow < capacity) {
	return rets[choices[rightNow]];
	} else if (rightNow == capacity) {
	current.set(0);
	return rets[choices[0]];
	}
	//race condition with another thread, and we lost
	// try again
	}
	}

	@Override
	public void refreshLoad(LoadMapping loadMapping) {
	updateRing(loadMapping);
	}

	private void refreshLocalityGroup() {
	Map<Integer, NodeInfo> cachedTaskToNodePort = taskToNodePort.get();
	Map<String, String> hostToRack = getHostToRackMapping(cachedTaskToNodePort);

	localityGroup.values().stream().forEach(v -> v.clear());

	for (int target : targetTasks) {
	Scope scope = calculateScope(cachedTaskToNodePort, hostToRack, target);
	if (!localityGroup.containsKey(scope)) {
	localityGroup.put(scope, new ArrayList<>());
	}
	localityGroup.get(scope).add(target);
	}
	}

	private List<Integer> getTargetsInScope(Scope scope) {
	List<Integer> rets = new ArrayList<>();
	List<Integer> targetInScope = localityGroup.get(scope);
	if (null != targetInScope) {
	rets.addAll(targetInScope);
	}
	Scope downgradeScope = Scope.downgrade(scope);
	if (downgradeScope != scope) {
	rets.addAll(getTargetsInScope(downgradeScope));
	}
	return rets;
	}

	private Scope transition(LoadMapping load) {
	List<Integer> targetInScope = getTargetsInScope(currentScope);
	if (targetInScope.isEmpty()) {
	Scope upScope = Scope.upgrade(currentScope);
	if (upScope == currentScope) {
	throw new RuntimeException("The current scope " + currentScope + " has no target tasks.");
	}
	currentScope = upScope;
	return transition(load);
	}

	if (null == load) {
	return currentScope;
	}

	double avg = targetInScope.stream().mapToDouble((key) -> load.get(key)).average().getAsDouble();
	Scope nextScope;
	if (avg < lowerBound) {
	nextScope = Scope.downgrade(currentScope);
	if (getTargetsInScope(nextScope).isEmpty()) {
	nextScope = currentScope;
	}
	} else if (avg > higherBound) {
	nextScope = Scope.upgrade(currentScope);
	} else {
	nextScope = currentScope;
	}

	return nextScope;
	}

	private synchronized void updateRing(LoadMapping load) {
	refreshLocalityGroup();
	Scope prevScope = currentScope;
	currentScope = transition(load);
	if (currentScope != prevScope) {
	//reset all the weights
	orig.values().stream().forEach(o -> o.resetWeight());
	}

	List<Integer> targetsInScope = getTargetsInScope(currentScope);

	//We will adjust weights based off of the minimum load
	double min = load == null ? 0 : targetsInScope.stream().mapToDouble((key) -> load.get(key)).min().getAsDouble();
	for (int target : targetsInScope) {
	IndexAndWeights val = orig.get(target);
	double l = load == null ? 0.0 : load.get(target);
	if (l <= min + (0.05)) {
	//We assume that within 5% of the minimum congestion is still fine.
	//Not congested we grow (but slowly)
	val.weight = Math.min(MAX_WEIGHT, val.weight + 1);
	} else {
	//Congested we contract much more quickly
	val.weight = Math.max(0, val.weight - 10);
	}
	}
	//Now we need to build the array
	long weightSum = targetsInScope.stream().mapToLong((target) -> orig.get(target).weight).sum();
	//Now we can calculate a percentage

	int currentIdx = 0;
	if (weightSum > 0) {
	for (int target : targetsInScope) {
	IndexAndWeights indexAndWeights = orig.get(target);
	int count = (int) ((indexAndWeights.weight / (double) weightSum) * capacity);
	for (int i = 0; i < count && currentIdx < capacity; i++) {
	prepareChoices[currentIdx] = indexAndWeights.index;
	currentIdx++;
	}
	}

	if (currentIdx > 0) {
	//in case we didn't fill in enough
	for (; currentIdx < capacity; currentIdx++) {
	prepareChoices[currentIdx] = prepareChoices[random.nextInt(currentIdx)];
	}
	}
	}
	if (currentIdx == 0) {
	//This really should be impossible, because we go off of the min load, and inc anything within 5% of it.
	// But just to be sure it is never an issue, especially with float rounding etc.
	for (; currentIdx < capacity; currentIdx++) {
	prepareChoices[currentIdx] = currentIdx % rets.length;
	}
	}

	shuffleArray(prepareChoices);

	// swapping two arrays
	int[] tempForSwap = choices;
	choices = prepareChoices;
	prepareChoices = tempForSwap;

	current.set(-1);
	}

	private void shuffleArray(int[] arr) {
	int size = arr.length;
	for (int i = size; i > 1; i--) {
	swap(arr, i - 1, random.nextInt(i));
	}
	}

	private void swap(int[] arr, int i, int j) {
	int tmp = arr[i];
	arr[i] = arr[j];
	arr[j] = tmp;
	}

	private Scope calculateScope(Map<Integer, NodeInfo> taskToNodePort, Map<String, String> hostToRack, int target) {
	NodeInfo targetNodeInfo = taskToNodePort.get(target);

	if (targetNodeInfo == null) {
	return Scope.EVERYTHING;
	}

	String sourceRack = hostToRack.get(sourceNodeInfo.get_node());
	String targetRack = hostToRack.get(targetNodeInfo.get_node());

	if (sourceRack != null && targetRack != null && sourceRack.equals(targetRack)) {
	if (sourceNodeInfo.get_node().equals(targetNodeInfo.get_node())) {
	if (sourceNodeInfo.get_port().equals(targetNodeInfo.get_port())) {
	return Scope.WORKER_LOCAL;
	}
	return Scope.HOST_LOCAL;
	}
	return Scope.RACK_LOCAL;
	} else {
	return Scope.EVERYTHING;
	}
	}

	private Map<String, String> getHostToRackMapping(Map<Integer, NodeInfo> taskToNodePort) {
	Set<String> hosts = new HashSet();
	for (int task : targetTasks) {
	//if this task containing worker will be killed by a assignments sync,
	//taskToNodePort will be an empty map which is refreshed by WorkerState
	if (taskToNodePort.containsKey(task)) {
	hosts.add(taskToNodePort.get(task).get_node());
	} else {
	LOG.error("Could not find task NodeInfo from local cache.");
	}
	}
	hosts.add(sourceNodeInfo.get_node());
	return dnsToSwitchMapping.resolve(new ArrayList<>(hosts));
	}

	//only for test
	public int getCapacity() {
	return capacity;
	}

	enum Scope {
	WORKER_LOCAL, HOST_LOCAL, RACK_LOCAL, EVERYTHING;

	public static Scope downgrade(Scope current) {
	switch (current) {
	case EVERYTHING:
	return RACK_LOCAL;
	case RACK_LOCAL:
	return HOST_LOCAL;
	case HOST_LOCAL:
	case WORKER_LOCAL:
	default:
	return WORKER_LOCAL;
	}
	}

	public static Scope upgrade(Scope current) {
	switch (current) {
	case WORKER_LOCAL:
	return HOST_LOCAL;
	case HOST_LOCAL:
	return RACK_LOCAL;
	case RACK_LOCAL:
	case EVERYTHING:
	default:
	return EVERYTHING;
	}
	}
	}

	private static class IndexAndWeights {
	final int index;
	int weight;

	IndexAndWeights(int index) {
	this.index = index;
	weight = MAX_WEIGHT;
	}

	void resetWeight() {
	weight = MAX_WEIGHT;
	}
	}
	}