Ejemplo n.º 1
0
    public void assignToNonLeafTasks(LinkedList<TaskRequestEvent> taskRequests) {
      Collections.shuffle(taskRequests);
      String queryMasterHostAndPort =
          context
              .getMasterContext()
              .getQueryMasterContext()
              .getWorkerContext()
              .getConnectionInfo()
              .getHostAndQMPort();

      TaskRequestEvent taskRequest;
      while (!taskRequests.isEmpty()) {
        taskRequest = taskRequests.pollFirst();
        LOG.debug("assignToNonLeafTasks: " + taskRequest.getExecutionBlockId());

        TaskAttemptId attemptId;
        // random allocation
        if (nonLeafTasks.size() > 0) {
          synchronized (nonLeafTasks) {
            attemptId = nonLeafTasks.iterator().next();
            nonLeafTasks.remove(attemptId);
          }
          LOG.debug("Assigned based on * match");

          Task task;
          task = stage.getTask(attemptId.getTaskId());

          TaskRequest taskAssign =
              new TaskRequestImpl(
                  attemptId,
                  Lists.newArrayList(task.getAllFragments()),
                  "",
                  false,
                  LogicalNodeSerializer.serialize(task.getLogicalPlan()),
                  context.getMasterContext().getQueryContext(),
                  stage.getDataChannel(),
                  stage.getBlock().getEnforcer(),
                  queryMasterHostAndPort);

          if (checkIfInterQuery(stage.getMasterPlan(), stage.getBlock())) {
            taskAssign.setInterQuery();
          }
          for (Map.Entry<String, Set<FetchImpl>> entry : task.getFetchMap().entrySet()) {
            Collection<FetchImpl> fetches = entry.getValue();
            if (fetches != null) {
              for (FetchImpl fetch : fetches) {
                taskAssign.addFetch(entry.getKey(), fetch);
              }
            }
          }

          WorkerConnectionInfo connectionInfo =
              context.getMasterContext().getWorkerMap().get(taskRequest.getWorkerId());

          // TODO send batch request
          BatchAllocationRequest.Builder requestProto = BatchAllocationRequest.newBuilder();
          requestProto.addTaskRequest(
              TaskAllocationProto.newBuilder()
                  .setResource(taskRequest.getResponseProto().getResource())
                  .setTaskRequest(taskAssign.getProto())
                  .build());

          requestProto.setExecutionBlockId(attemptId.getTaskId().getExecutionBlockId().getProto());
          context
              .getMasterContext()
              .getEventHandler()
              .handle(new TaskAttemptAssignedEvent(attemptId, connectionInfo));

          CallFuture<BatchAllocationResponse> callFuture = new CallFuture<>();

          InetSocketAddress addr = stage.getAssignedWorkerMap().get(connectionInfo.getId());
          if (addr == null)
            addr = new InetSocketAddress(connectionInfo.getHost(), connectionInfo.getPeerRpcPort());

          AsyncRpcClient tajoWorkerRpc;
          try {
            tajoWorkerRpc =
                RpcClientManager.getInstance()
                    .getClient(addr, TajoWorkerProtocol.class, true, rpcParams);
            TajoWorkerProtocol.TajoWorkerProtocolService tajoWorkerRpcClient =
                tajoWorkerRpc.getStub();
            tajoWorkerRpcClient.allocateTasks(
                callFuture.getController(), requestProto.build(), callFuture);

            BatchAllocationResponse responseProto =
                callFuture.get(RpcConstants.FUTURE_TIMEOUT_SECONDS_DEFAULT, TimeUnit.SECONDS);

            if (responseProto.getCancellationTaskCount() > 0) {
              for (TaskAllocationProto proto : responseProto.getCancellationTaskList()) {
                cancel(task.getAttempt(new TaskAttemptId(proto.getTaskRequest().getId())));
                cancellation++;
              }

              if (LOG.isDebugEnabled()) {
                LOG.debug(
                    "Canceled requests: "
                        + responseProto.getCancellationTaskCount()
                        + " from "
                        + addr);
              }
              continue;
            }

            totalAssigned++;
            scheduledObjectNum--;
          } catch (Exception e) {
            LOG.error(e);
          }
        }
      }
    }