@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); }
private WorkerTopologyContext mockContext(List<Integer> availableTaskIds) { Map<String, Object> conf = new HashMap<>(); conf.put(Config.STORM_NETWORK_TOPOGRAPHY_PLUGIN, "org.apache.storm.networktopography.DefaultRackDNSToSwitchMapping"); conf.put(Config.TOPOLOGY_LOCALITYAWARE_HIGHER_BOUND, 0.8); conf.put(Config.TOPOLOGY_LOCALITYAWARE_LOWER_BOUND, 0.2); WorkerTopologyContext context = mock(WorkerTopologyContext.class); when(context.getConf()).thenReturn(conf); Map<Integer, NodeInfo> taskNodeToPort = new HashMap<>(); NodeInfo nodeInfo = new NodeInfo("node-id", Sets.newHashSet(6700L)); availableTaskIds.forEach(e -> taskNodeToPort.put(e, nodeInfo)); when(context.getTaskToNodePort()).thenReturn(new AtomicReference<>(taskNodeToPort)); when(context.getThisWorkerHost()).thenReturn("node-id"); when(context.getThisWorkerPort()).thenReturn(6700); return context; }