BrokerStatus { Map<TaskTypeUser, IntCounter> collectMaxAvailableSpacePerUserOnWorker(String workerId, int maxThreadPerUserPerTaskTypePercent, Map<String, Integer> startingAvailableSpace) { Map<TaskTypeUser, IntCounter> res = new HashMap<>(); lock.readLock().lock(); try { this.tasks .values() .stream() .filter(t -> t.getStatus() == Task.STATUS_RUNNING && workerId.equals(t.getWorkerId())) .forEach(t -> { String type = t.getType(); Integer startingMaxAvailableSpacePerUser = startingAvailableSpace.get(type); if (startingMaxAvailableSpacePerUser == null) { startingMaxAvailableSpacePerUser = startingAvailableSpace.get(Task.TASKTYPE_ANY); } if (startingMaxAvailableSpacePerUser != null && startingMaxAvailableSpacePerUser > 0) { TaskTypeUser key = new TaskTypeUser(type, t.getUserId()); IntCounter count = res.get(key); if (count == null) { int effectiveBoundForUser = (startingMaxAvailableSpacePerUser * maxThreadPerUserPerTaskTypePercent) / 100; if (effectiveBoundForUser <= 0) { effectiveBoundForUser = 1; } LOGGER.log(Level.FINEST, "collectMaxAvailableSpacePerUserOnWorker {0} -> for user {1} we are starting from {2} - bound is {3}", new Object[]{workerId, t.getUserId(), startingMaxAvailableSpacePerUser, effectiveBoundForUser}); count = new IntCounter(effectiveBoundForUser); res.put(key, count); } LOGGER.log(Level.FINEST, "collectMaxAvailableSpacePerUserOnWorker {0} -> for user {1} found task {2} - id {3}", new Object[]{workerId, t.getUserId(), t.getType(), t.getTaskId()}); count.count--; } }); } finally { lock.readLock().unlock(); } return res; } BrokerStatus(StatusChangesLog log); boolean isReadonly(); void setReadonly(boolean readonly); WorkerStatus getWorkerStatus(String workerId); Map<String, Long> getActualSlots(); List<TransactionStatus> getAllTransactions(); List<WorkerStatusView> getAllWorkers(); List<TaskStatusView> getAllTasks(); Collection<WorkerStatus> getWorkersAtBoot(); Collection<Task> getTasksAtBoot(); Collection<Transaction> getTransactionsAtBoot(); long nextTaskId(); long nextTransactionId(); int getCheckpointsCount(); void checkpoint(long transactionsTtl); BrokerStatusSnapshot createSnapshot(); void close(); Set<Long> purgeFinishedTasksAndSignalExpiredTasks(int finishedTasksRetention, int maxExpiredPerCycle); void followTheLeader(); TransactionStatus getTransaction(long transactionId); ModificationResult applyModification(StatusEdit edit); BrokerStatusStats getStats(); void recover(); void startWriting(); Task getTask(long taskId); TaskStatusView getTaskStatus(long taskId);  }
@Test public void testCollectMaxAvailableSpacePerUserOnWorker() throws Exception { try (Broker broker = new Broker(new BrokerConfiguration(), new MemoryCommitLog(), new TasksHeap(1000, createTaskPropertiesMapperFunction()));) { broker.startAsWritable(); for (int i = 0; i < 5; i++) { SubmitTaskResult res = broker.getClient().submitTask(new AddTaskRequest(0, TASKTYPE_MYTYPE, userId, "", 1, 0, 0, null, 0, null, null)); assertTrue(res.getTaskId() > 0); } int maxThreadPerTaskType1 = 10; Map<String, Integer> maxThreadsPerTaskType = Collections.singletonMap(TASKTYPE_MYTYPE, maxThreadPerTaskType1); int maxThreadPerUserPerTaskTypePercent = 0; { List<AssignedTask> assignTasksToWorker = broker.assignTasksToWorker(100, maxThreadsPerTaskType, Collections.singletonList(group), Collections.emptySet(), "myworker", new HashMap<>(), new ResourceUsageCounters(), maxThreadPerUserPerTaskTypePercent); assertEquals(5, assignTasksToWorker.size()); } for (int i = 0; i < 10; i++) { SubmitTaskResult res = broker.getClient().submitTask(new AddTaskRequest(0, TASKTYPE_MYTYPE, userId, "", 1, 0, 0, null, 0, null, null)); assertTrue(res.getTaskId() > 0); } maxThreadPerUserPerTaskTypePercent = 80; { List<AssignedTask> assignTasksToWorker = broker.assignTasksToWorker(100, maxThreadsPerTaskType, Collections.singletonList(group), Collections.emptySet(), "myworker", new HashMap<>(), new ResourceUsageCounters(), maxThreadPerUserPerTaskTypePercent); assertEquals(3, assignTasksToWorker.size()); } } }