/** {@inheritDoc} */
@Override public ClusterNode getBalancedNode(ComputeTaskSession ses, List<ClusterNode> top, ComputeJob job) {
    A.notNull(ses, "ses");
    A.notNull(top, "top");
    A.notNull(job, "job");

    IgniteBiTuple<Boolean, WeightedTopology> weightedTop = taskTops.get(ses.getId());

    // Create new cached topology if there is no one. Do not
    // use cached topology after task has been mapped.
    if (weightedTop == null)
        // Called from ComputeTask#map(). Put new topology and false as not mapped yet.
        taskTops.put(ses.getId(), weightedTop = F.t(false, new WeightedTopology(top)));
    // We have topology - check if task has been mapped.
    else if (weightedTop.get1())
        // Do not use cache after ComputeTask#map().
        return new WeightedTopology(top).pickWeightedNode();

    return weightedTop.get2().pickWeightedNode();
}
 

/** {@inheritDoc} */
@Override public ClusterNode getBalancedNode(ComputeTaskSession ses, List<ClusterNode> top, ComputeJob job) {
    A.notNull(ses, "ses", top, "top");

    if (isPerTask) {
        // Note that every session operates from single thread which
        // allows us to use concurrent map and avoid synchronization.
        RoundRobinPerTaskLoadBalancer taskBalancer = perTaskBalancers.get(ses.getId());

        if (taskBalancer == null)
            perTaskBalancers.put(ses.getId(), taskBalancer = new RoundRobinPerTaskLoadBalancer());

        return taskBalancer.getBalancedNode(top);
    }

    return balancer.getBalancedNode(top);
}
 

/**
 * THIS METHOD IS USED ONLY FOR TESTING.
 *
 * @param ses Task session.
 * @return Internal list of nodes.
 */
List<UUID> getNodeIds(ComputeTaskSession ses) {
    if (isPerTask) {
        RoundRobinPerTaskLoadBalancer balancer = perTaskBalancers.get(ses.getId());

        if (balancer == null)
            return Collections.emptyList();

        List<UUID> ids = new ArrayList<>();

        for (ClusterNode node : balancer.getNodes()) {
            ids.add(node.id());
        }

        return ids;
    }

    return balancer.getNodeIds();
}
 

/** {@inheritDoc} */
@Override public ClusterNode getBalancedNode(ComputeTaskSession ses, List<ClusterNode> top, ComputeJob job) {
    A.notNull(ses, "ses");
    A.notNull(top, "top");
    A.notNull(job, "job");

    // Optimization for non-weighted randomization.
    if (!isUseWeights)
        return top.get(RAND.nextInt(top.size()));

    IgniteBiTuple<Boolean, WeightedTopology> weightedTop = taskTops.get(ses.getId());

    // Create new cached topology if there is no one. Do not
    // use cached topology after task has been mapped.
    if (weightedTop == null) {
        // Called from ComputeTask#map(). Put new topology and false as not mapped yet.
        taskTops.put(ses.getId(), weightedTop = F.t(false, new WeightedTopology(top)));
    }
    // We have topology - check if task has been mapped.
    else if (weightedTop.get1()) {
        // Do not use cache after ComputeTask#map().
        return new WeightedTopology(top).pickWeightedNode();
    }

    return weightedTop.get2().pickWeightedNode();
}
 

/**
 * Injects held resources into given {@code job}.
 *
 * @param dep Deployment.
 * @param taskCls Task class.
 * @param job Grid job to inject resources to.
 * @param ses Current task session.
 * @param jobCtx Job context.
 * @throws IgniteCheckedException Thrown in case of any errors.
 */
public void inject(GridDeployment dep, Class<?> taskCls, ComputeJob job, ComputeTaskSession ses,
    GridJobContextImpl jobCtx) throws IgniteCheckedException {
    if (log.isDebugEnabled())
        log.debug(S.toString("Injecting resources", "job", job, true));

    // Unwrap Proxy object.
    Object obj = unwrapTarget(job);

    injectToJob(dep, taskCls, obj, ses, jobCtx);

    if (obj instanceof GridInternalWrapper) {
        Object usrObj = ((GridInternalWrapper<?>)obj).userObject();

        if (usrObj != null)
            injectToJob(dep, taskCls, usrObj, ses, jobCtx);
    }
}
 

/**
 * @param level Level.
 * @param jobSes Job session.
 * @return Always returns {@code 1}.
 */
@SuppressWarnings("unused")
public int executeLoadTestJob(int level, ComputeTaskSession jobSes) {
    assert level > 0;
    assert jobSes != null;

    try {
        assert "1".equals(jobSes.waitForAttribute("1st", 10000));

        assert "2".equals(jobSes.waitForAttribute("2nd", 10000));
    }
    catch (InterruptedException e) {
        // Fail.
        throw new IgniteException("Failed to wait for attribute.", e);
    }

    return 1;
}
 

/**
 * @throws Exception If test failed.
 */
@Test
public void testMultipleTaskSessions() throws Exception {
    ComputeTaskSession ses1 = new GridTestTaskSession(IgniteUuid.randomUuid());
    ComputeTaskSession ses2 = new GridTestTaskSession(IgniteUuid.randomUuid());

    List<ClusterNode> allNodes = (List<ClusterNode>)getSpiContext().nodes();

    List<UUID> orderedNodes = getSpi().getNodeIds(ses1);

    assertEquals("Balancer doesn't use all available nodes", orderedNodes.size(), allNodes.size());

    checkCyclicBalancing(getSpi(), allNodes, orderedNodes, ses1, ses2);

    getSpiContext().triggerEvent(new TaskEvent(
        null, null, EVT_TASK_FINISHED, ses1.getId(), null, null, false, null));
    getSpiContext().triggerEvent(new TaskEvent(
        null, null, EVT_TASK_FAILED, ses2.getId(), null, null, false, null));
}
 

/** */
@Test
public void testNodeNotInTopology() throws Exception {
    ComputeTaskSession ses = new GridTestTaskSession();

    ClusterNode node = new GridTestNode(UUID.randomUUID());

    List<ClusterNode> notInTop = Arrays.asList(node);

    try {
        getSpi().getBalancedNode(ses, notInTop, new GridTestJob());
    }
    catch (IgniteException e) {
        assertTrue(e.getMessage().contains("Task topology does not have alive nodes"));
    }
}
 

/**
 * Performs two full cycles by round robin routine for check correct order.
 *
 * @param spi Load balancing SPI.
 * @param allNodes Topology nodes.
 * @param orderedNodes Balancing nodes.
 * @param ses Task session.
 */
static void checkCyclicBalancing(RoundRobinLoadBalancingSpi spi, List<ClusterNode> allNodes,
    List<UUID> orderedNodes, ComputeTaskSession ses) {

    ClusterNode firstNode = spi.getBalancedNode(ses, allNodes, new GridTestJob());

    int startIdx = firstBalancedNodeIndex(firstNode, orderedNodes);

    // Two full cycles by round robin routine.
    for (int i = 0; i < allNodes.size() * 2; i++) {
        int actualIdx = (startIdx + i + 1) % allNodes.size();

        ClusterNode nextNode = spi.getBalancedNode(ses, allNodes, new GridTestJob());

        assertEquals("Balancer returns node out of order", nextNode.id(), orderedNodes.get(actualIdx));
    }
}
 

/**
 * Performs two full cycles by round robin routine for check correct order.
 * Switches between two task sessions by turns.
 *
 * @param spi Load balancing SPI.
 * @param allNodes Topology nodes.
 * @param orderedNodes Balancing nodes.
 * @param ses1 First task session.
 * @param ses2 Second task session.
 */
static void checkCyclicBalancing(RoundRobinLoadBalancingSpi spi, List<ClusterNode> allNodes,
    List<UUID> orderedNodes, ComputeTaskSession ses1, ComputeTaskSession ses2) {

    ClusterNode firstNode = spi.getBalancedNode(ses1, allNodes, new GridTestJob());

    int startIdx = firstBalancedNodeIndex(firstNode, orderedNodes);

    // Two full cycles by round robin routine.
    for (int i = 0; i < allNodes.size() * 2; i++) {
        int actualIdx = (startIdx + i + 1) % allNodes.size();

        ClusterNode nextNode = spi.getBalancedNode(i % 2 == 0 ? ses1 : ses2, allNodes, new GridTestJob());

        assertEquals("Balancer returns node out of order", nextNode.id(), orderedNodes.get(actualIdx));
    }
}
 

/** {@inheritDoc} */
@Override protected Collection<? extends ComputeJob> split(int gridSize, Void arg) {
    return F.asSet(new ComputeJobAdapter() {
        @TaskSessionResource
        private ComputeTaskSession ses;

        @Override public Object execute() {
            CNT.incrementAndGet();

            if (fail)
                throw new ComputeExecutionRejectedException("Expected error.");

            return ses.getTaskName();
        }
    });
}
 

/**
 * @param level Level.
 * @param jobSes Job session.
 * @return ALways returns {@code 1}.
 */
@SuppressWarnings("unused")
@Gridify(taskClass = GridifyLoadTestTask.class, timeout = 10000)
public int executeLoadTestJob(int level, ComputeTaskSession jobSes) {
    assert level > 0;
    assert jobSes != null;

    jobSes.setAttribute("1st", 10000);
    jobSes.setAttribute("2nd", 10000);

    return 1;
}
 

/**
 * Gets name task executed by current thread.
 *
 * @return Task name or {@code null} if security is disabled.
 */
public String currentTaskName() {
    if (!ctx.security().enabled())
        return null;

    ComputeTaskSession ses = currSess.get();

    if (ses == null)
        return null;

    return ses.getTaskName();
}
 

/**
 * @param ses Task session instance.
 * @param ctx Kernal context.
 */
public ComputeTaskInternalFuture(ComputeTaskSession ses, GridKernalContext ctx) {
    assert ses != null;
    assert ctx != null;

    this.ses = ses;
    this.ctx = ctx;

    userFut = new ComputeFuture<>(this);

    log = ctx.log(ComputeTaskInternalFuture.class);
}
 

/**
 * Gets task timeout.
 *
 * @return Task timeout.
 */
public ComputeTaskSession getTaskSession() {
    if (ses == null)
        throw new IllegalStateException("Cannot access task session after future has been deserialized.");

    return ses;
}
 

/**
 * @param ses Session.
 */
public GridTestCollisionJobContext(ComputeTaskSession ses) {
    this.ses = ses;

    ctx = new GridTestJobContext();

    job = new GridTestJob();
}
 

/**
 * @param ses Session.
 * @param idx Index.
 */
public GridTestCollisionJobContext(ComputeTaskSession ses, int idx) {
    this.idx = idx;
    this.ses = ses;

    ctx = new GridTestJobContext();

    job = new GridTestJob();
}
 

/**
 * @param ses Session.
 * @param idx Index in wrapping collection.
 * @param onActivate Closure to be called when current task get activated.
 */
public GridTestCollisionJobContext(ComputeTaskSession ses, int idx,
    IgniteInClosure<GridTestCollisionJobContext> onActivate) {
    this(ses, idx);

    this.onActivate = onActivate;
}
 

/**
 * @param ses Task session.
 * @param jobId Job ID.
 */
public GridTestCollisionJobContext(ComputeTaskSession ses, IgniteUuid jobId) {
    this.ses = ses;

    ctx = new GridTestJobContext(jobId);

    job = new GridTestJob();
}
 

/**
 * @throws Exception If test failed.
 */
@Test
public void testBalancingOneNode() throws Exception {
    ComputeTaskSession ses = new GridTestTaskSession();

    List<ClusterNode> allNodes = (List<ClusterNode>)getSpiContext().nodes();

    List<ClusterNode> balancedNode = Arrays.asList(allNodes.get(0));

    ClusterNode firstNode = getSpi().getBalancedNode(ses, balancedNode, new GridTestJob());
    ClusterNode secondNode = getSpi().getBalancedNode(ses, balancedNode, new GridTestJob());

    assertEquals(firstNode, secondNode);
}
 

/** {@inheritDoc} */
@NotNull @Override public Map<? extends ComputeJob, ClusterNode> map(List<ClusterNode> subgrid,
    @Nullable Void arg) {
    return F.asMap(new ComputeJobAdapter() {
        @TaskSessionResource
        private ComputeTaskSession ses;

        @Override public Object execute() {
            return ses.getTaskName();
        }
    }, F.rand(subgrid));
}
 

/** {@inheritDoc} */
@Override public ClusterNode getBalancedNode(ComputeTaskSession ses, List<ClusterNode> top,
    ComputeJob job) {
    if (getName().equals(expName))
        isTaskLoadBalancingCalled = true;
    else
        isWrongTaskLoadBalancingCalled = true;

    return super.getBalancedNode(ses, top, job);
}
 

/**
 * @throws Exception If failed.
 */
@Test
public void testReleasePartitionJobImplementMasterLeave() throws Exception {
    final int orgId = primaryKey(grid(0).cache(Organization.class.getSimpleName()));

    try {
        grid(1).compute().affinityRunAsync(
            Arrays.asList(Organization.class.getSimpleName(), Person.class.getSimpleName()),
            new Integer(orgId),
            new RunnableWithMasterLeave() {
                @IgniteInstanceResource
                private Ignite ignite;

                @Override public void onMasterNodeLeft(ComputeTaskSession ses) throws IgniteException {
                    // No-op.
                }

                @Override public void run() {
                    try {
                        checkPartitionsReservations((IgniteEx)ignite, orgId, 1);
                    }
                    catch (Exception e) {
                        e.printStackTrace();
                        fail("Unexpected exception");
                    }

                    try {
                        Thread.sleep(1000);
                    }
                    catch (InterruptedException ignored) {
                        // No-op.
                    }
                }
            });

        stopGrid(1, true);

        Thread.sleep(3000);

        awaitPartitionMapExchange();

        checkPartitionsReservations(grid(0), orgId, 0);
    }
    finally {
        startGrid(1);
        awaitPartitionMapExchange();
    }
}
 

/** {@inheritDoc} */
@Override public ComputeTaskSession getTaskSession() {
    return taskSes;
}
 

/** {@inheritDoc} */
@Override public void onMasterNodeLeft(ComputeTaskSession ses) {
    ((ComputeJobMasterLeaveAware)job).onMasterNodeLeft(ses);
}
 

/** {@inheritDoc} */
@Override public void onMasterNodeLeft(ComputeTaskSession ses) {
    ((ComputeJobMasterLeaveAware)c).onMasterNodeLeft(ses);
}
 

/** {@inheritDoc} */
@Override public void onMasterNodeLeft(ComputeTaskSession ses) {
    ((ComputeJobMasterLeaveAware)r).onMasterNodeLeft(ses);
}
 

/** {@inheritDoc} */
@Override public ComputeTaskSession getTaskSession() {
    return ((ComputeTaskInternalFuture<R>)fut).getTaskSession();
}
 

/** {@inheritDoc} */
@Override public ComputeTaskSession getTaskSession() {
    return ses;
}
 

/**
 * @param ses Session.
 */
public void setTaskSession(ComputeTaskSession ses) {
    this.ses = ses;
}