Start line:  
End line:  

Snippet Preview

Snippet HTML Code

Stack Overflow Questions
  /*
   * Licensed 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 com.facebook.presto.execution;
 
 
 import java.net.URI;
 
 import static com.facebook.presto.OutputBuffers.INITIAL_EMPTY_OUTPUT_BUFFERS;
 import static com.facebook.presto.SessionTestUtils.TEST_SESSION;
 import static com.facebook.presto.execution.TaskTestUtils.PLAN_FRAGMENT;
 import static com.facebook.presto.execution.TaskTestUtils.SPLIT;
 import static com.facebook.presto.execution.TaskTestUtils.TABLE_SCAN_NODE_ID;
 import static com.facebook.presto.execution.TaskTestUtils.createTestingPlanner;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertNotEquals;
 import static org.testng.Assert.assertNotNull;
 import static org.testng.Assert.assertNull;
 
 public class TestSqlTaskManager
 {
     private static final TaskId TASK_ID = new TaskId("query""stage""task");
     public static final TaskId OUT = new TaskId("query""stage""out");
 
     private final TaskExecutor taskExecutor;
     private final LocalMemoryManager localMemoryManager;
 
     public TestSqlTaskManager()
     {
          = new TaskExecutor(8, 16);
         .start();
     }
 
     @AfterClass
     public void tearDown()
             throws Exception
     {
         .stop();
         .destroy();
     }
 
     @Test
     public void testEmptyQuery()
             throws Exception
     {
         try (SqlTaskManager sqlTaskManager = createSqlTaskManager(new TaskManagerConfig())) {
             TaskId taskId = ;
             TaskInfo taskInfo = sqlTaskManager.updateTask(,
                     taskId,
                     ,
                     ImmutableList.<TaskSource>of(),
                     );
             assertEquals(taskInfo.getState(), .);
 
             taskInfo = sqlTaskManager.getTaskInfo(taskId);
             assertEquals(taskInfo.getState(), .);
 
             taskInfo = sqlTaskManager.updateTask(,
                     taskId,
                     ,
                     ImmutableList.of(new TaskSource(, ImmutableSet.<ScheduledSplit>of(), true)),
                     .withNoMoreBufferIds());
             assertEquals(taskInfo.getState(), .);
            taskInfo = sqlTaskManager.getTaskInfo(taskId);
            assertEquals(taskInfo.getState(), .);
        }
    }
    @Test
    public void testSimpleQuery()
            throws Exception
    {
        try (SqlTaskManager sqlTaskManager = createSqlTaskManager(new TaskManagerConfig())) {
            TaskId taskId = ;
            TaskInfo taskInfo = sqlTaskManager.updateTask(,
                    taskId,
                    ,
                    ImmutableList.of(new TaskSource(, ImmutableSet.of(), true)),
            assertEquals(taskInfo.getState(), .);
            taskInfo = sqlTaskManager.getTaskInfo(taskId);
            assertEquals(taskInfo.getState(), .);
            BufferResult results = sqlTaskManager.getTaskResults(taskId, 0, new DataSize(1, .)).get();
            assertEquals(results.isBufferClosed(), false);
            assertEquals(results.getPages().size(), 1);
            assertEquals(results.getPages().get(0).getPositionCount(), 1);
            results = sqlTaskManager.getTaskResults(taskIdresults.getToken() + results.getPages().size(), new DataSize(1, .)).get();
            assertEquals(results.isBufferClosed(), true);
            assertEquals(results.getPages().size(), 0);
            taskInfo = sqlTaskManager.getTaskInfo(taskIdtaskInfo.getState()).get(1, .);
            assertEquals(taskInfo.getState(), .);
            taskInfo = sqlTaskManager.getTaskInfo(taskId);
            assertEquals(taskInfo.getState(), .);
        }
    }
    @Test
    public void testCancel()
            throws Exception
    {
        try (SqlTaskManager sqlTaskManager = createSqlTaskManager(new TaskManagerConfig())) {
            TaskId taskId = ;
            TaskInfo taskInfo = sqlTaskManager.updateTask(,
                    taskId,
                    ,
                    ImmutableList.<TaskSource>of(),
                    );
            assertEquals(taskInfo.getState(), .);
            assertNull(taskInfo.getStats().getEndTime());
            taskInfo = sqlTaskManager.getTaskInfo(taskId);
            assertEquals(taskInfo.getState(), .);
            assertNull(taskInfo.getStats().getEndTime());
            taskInfo = sqlTaskManager.cancelTask(taskId);
            assertEquals(taskInfo.getState(), .);
            assertNotNull(taskInfo.getStats().getEndTime());
            taskInfo = sqlTaskManager.getTaskInfo(taskId);
            assertEquals(taskInfo.getState(), .);
            assertNotNull(taskInfo.getStats().getEndTime());
        }
    }
    @Test
    public void testAbort()
            throws Exception
    {
        try (SqlTaskManager sqlTaskManager = createSqlTaskManager(new TaskManagerConfig())) {
            TaskId taskId = ;
            TaskInfo taskInfo = sqlTaskManager.updateTask(,
                    taskId,
                    ,
                    ImmutableList.<TaskSource>of(),
                    );
            assertEquals(taskInfo.getState(), .);
            assertNull(taskInfo.getStats().getEndTime());
            taskInfo = sqlTaskManager.getTaskInfo(taskId);
            assertEquals(taskInfo.getState(), .);
            assertNull(taskInfo.getStats().getEndTime());
            taskInfo = sqlTaskManager.abortTask(taskId);
            assertEquals(taskInfo.getState(), .);
            assertNotNull(taskInfo.getStats().getEndTime());
            taskInfo = sqlTaskManager.getTaskInfo(taskId);
            assertEquals(taskInfo.getState(), .);
            assertNotNull(taskInfo.getStats().getEndTime());
        }
    }
    @Test
    public void testAbortResults()
            throws Exception
    {
        try (SqlTaskManager sqlTaskManager = createSqlTaskManager(new TaskManagerConfig())) {
            TaskId taskId = ;
            TaskInfo taskInfo = sqlTaskManager.updateTask(,
                    taskId,
                    ,
                    ImmutableList.of(new TaskSource(, ImmutableSet.of(), true)),
            assertEquals(taskInfo.getState(), .);
            taskInfo = sqlTaskManager.getTaskInfo(taskId);
            assertEquals(taskInfo.getState(), .);
            sqlTaskManager.abortTaskResults(taskId);
            taskInfo = sqlTaskManager.getTaskInfo(taskIdtaskInfo.getState()).get(1, .);
            assertEquals(taskInfo.getState(), .);
            taskInfo = sqlTaskManager.getTaskInfo(taskId);
            assertEquals(taskInfo.getState(), .);
        }
    }
    @Test
    public void testRemoveOldTasks()
            throws Exception
    {
        try (SqlTaskManager sqlTaskManager = createSqlTaskManager(new TaskManagerConfig().setInfoMaxAge(new Duration(5, .)))) {
            TaskId taskId = ;
            TaskInfo taskInfo = sqlTaskManager.updateTask(,
                    taskId,
                    ,
                    ImmutableList.<TaskSource>of(),
                    );
            assertEquals(taskInfo.getState(), .);
            taskInfo = sqlTaskManager.cancelTask(taskId);
            assertEquals(taskInfo.getState(), .);
            taskInfo = sqlTaskManager.getTaskInfo(taskId);
            assertEquals(taskInfo.getState(), .);
            Thread.sleep(100);
            sqlTaskManager.removeOldTasks();
            for (TaskInfo info : sqlTaskManager.getAllTaskInfo()) {
                assertNotEquals(info.getTaskId(), taskId);
            }
        }
    }
    {
        return new SqlTaskManager(
                createTestingPlanner(),
                new MockLocationFactory(),
                ,
                new QueryMonitor(new ObjectMapperProvider().get(), new NullEventClient(), new NodeInfo("test"), new NodeVersion("testVersion")),
                new NodeInfo("test"),
                ,
                config,
                new MemoryManagerConfig());
    }
    public static class MockExchangeClientSupplier
            implements Supplier<ExchangeClient>
    {
        @Override
        public ExchangeClient get()
        {
            throw new UnsupportedOperationException();
        }
    }
    public static class MockLocationFactory
            implements LocationFactory
    {
        @Override
        public URI createQueryLocation(QueryId queryId)
        {
            return URI.create("fake://query/" + queryId);
        }
        @Override
        public URI createStageLocation(StageId stageId)
        {
            return URI.create("fake://stage/" + stageId);
        }
        @Override
        public URI createLocalTaskLocation(TaskId taskId)
        {
            return URI.create("fake://task/" + taskId);
        }
        @Override
        public URI createTaskLocation(Node nodeTaskId taskId)
        {
            return URI.create("fake://task/" + node.getNodeIdentifier() + "/" + taskId);
        }
        @Override
        public URI createMemoryInfoLocation(Node node)
        {
            return URI.create("fake://" + node.getNodeIdentifier() + "/memory");
        }
    }
New to GrepCode? Check out our FAQ X