Github user arina-ielchiieva commented on a diff in the pull request: https://github.com/apache/drill/pull/1051#discussion_r154117721 --- Diff: exec/java-exec/src/main/java/org/apache/drill/exec/work/foreman/FragmentsRunner.java --- @@ -0,0 +1,439 @@ +/* + * 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.drill.exec.work.foreman; + +import com.google.common.base.Preconditions; +import com.google.common.collect.ArrayListMultimap; +import com.google.common.collect.Multimap; +import com.google.common.collect.Sets; +import io.netty.buffer.ByteBuf; +import org.apache.drill.common.concurrent.ExtendedLatch; +import org.apache.drill.common.exceptions.ExecutionSetupException; +import org.apache.drill.common.exceptions.UserException; +import org.apache.drill.exec.ops.FragmentContext; +import org.apache.drill.exec.physical.base.FragmentRoot; +import org.apache.drill.exec.proto.BitControl; +import org.apache.drill.exec.proto.BitControl.PlanFragment; +import org.apache.drill.exec.proto.CoordinationProtos; +import org.apache.drill.exec.proto.GeneralRPCProtos; +import org.apache.drill.exec.proto.UserBitShared.QueryId; +import org.apache.drill.exec.proto.UserBitShared.QueryResult.QueryState; +import org.apache.drill.exec.rpc.RpcException; +import org.apache.drill.exec.rpc.UserClientConnection; +import org.apache.drill.exec.rpc.control.Controller; +import org.apache.drill.exec.server.DrillbitContext; +import org.apache.drill.exec.testing.ControlsInjector; +import org.apache.drill.exec.testing.ControlsInjectorFactory; +import org.apache.drill.exec.work.EndpointListener; +import org.apache.drill.exec.work.WorkManager.WorkerBee; +import org.apache.drill.exec.work.fragment.FragmentExecutor; +import org.apache.drill.exec.work.fragment.FragmentStatusReporter; +import org.apache.drill.exec.work.fragment.NonRootFragmentManager; +import org.apache.drill.exec.work.fragment.RootFragmentManager; + + +import java.util.ArrayList; +import java.util.Collection; +import java.util.LinkedList; +import java.util.List; +import java.util.Set; +import java.util.concurrent.CountDownLatch; + +/** + * Is responsible for submitting query fragments for running (locally and remotely). + */ +public class FragmentsRunner { + + private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(FragmentsRunner.class); + private static final ControlsInjector injector = ControlsInjectorFactory.getInjector(FragmentsRunner.class); + + private static final long RPC_WAIT_IN_MSECS_PER_FRAGMENT = 5000; + + private final WorkerBee bee; + private final UserClientConnection initiatingClient; + private final DrillbitContext drillbitContext; + private final Foreman foreman; + + private List<PlanFragment> planFragments; + private PlanFragment rootPlanFragment; + private FragmentRoot rootOperator; + + public FragmentsRunner(WorkerBee bee, UserClientConnection initiatingClient, DrillbitContext drillbitContext, Foreman foreman) { + this.bee = bee; + this.initiatingClient = initiatingClient; + this.drillbitContext = drillbitContext; + this.foreman = foreman; + } + + public WorkerBee getBee() { + return bee; + } + + public void setPlanFragments(List<PlanFragment> planFragments) { + this.planFragments = planFragments; + } + + public void setRootPlanFragment(PlanFragment rootPlanFragment) { + this.rootPlanFragment = rootPlanFragment; + } + + public void setRootOperator(FragmentRoot rootOperator) { + this.rootOperator = rootOperator; + } + + /** + * Submits root and non-root fragments fragments for running. + * In case of success move query to the running state. + */ + public void submit() { + try { + assert planFragments != null; + assert rootPlanFragment != null; + assert rootOperator != null; + + QueryId queryId = foreman.getQueryId(); + assert queryId == rootPlanFragment.getHandle().getQueryId(); + + QueryManager queryManager = foreman.getQueryManager(); + + try { + drillbitContext.getWorkBus().addFragmentStatusListener(queryId, queryManager.getFragmentStatusListener()); + drillbitContext.getClusterCoordinator().addDrillbitStatusListener(queryManager.getDrillbitStatusListener()); + + logger.debug("Submitting fragments to run."); + // set up the root fragment first so we'll have incoming buffers available. + setupRootFragment(rootPlanFragment, rootOperator); + setupNonRootFragments(planFragments); + + } catch (ExecutionSetupException e) { --- End diff -- Removed try-catch block. Will catch exception in Foreman.
---