public TaskContext build() { MemoryPool memoryPool = new MemoryPool(new MemoryPoolId("test"), memoryPoolSize); SpillSpaceTracker spillSpaceTracker = new SpillSpaceTracker(maxSpillSize); DefaultQueryContext queryContext = new DefaultQueryContext( queryId, queryMaxMemory, queryMaxTotalMemory, memoryPool, GC_MONITOR, notificationExecutor, yieldExecutor, queryMaxSpillSize, spillSpaceTracker); return createTaskContext(queryContext, session, taskStateMachine); } }
private static TaskContext createTaskContext(DefaultQueryContext queryContext, Session session, TaskStateMachine taskStateMachine) { return queryContext.addTaskContext( taskStateMachine, session, true, true, OptionalInt.empty()); }
private synchronized ListenableFuture<?> updateSystemMemory(String allocationTag, long delta) { // We call memoryPool.getQueryMemoryReservation(queryId) instead of calling queryMemoryContext.getUserMemory() to // calculate the total memory size. // // Calling the latter can result in a deadlock: // * A thread doing a user allocation will acquire locks in this order: // 1. monitor of queryMemoryContext.userAggregateMemoryContext // 2. monitor of this (QueryContext) // * The current thread doing a system allocation will acquire locks in this order: // 1. monitor of this (QueryContext) // 2. monitor of queryMemoryContext.userAggregateMemoryContext // Deadlock is possible for concurrent user and system allocations when updateSystemMemory()/updateUserMemory // calls queryMemoryContext.getUserMemory()/queryMemoryContext.getSystemMemory(), respectively. For concurrent // allocations of the same type (e.g., tryUpdateUserMemory/updateUserMemory) it is not possible as they share // the same RootAggregatedMemoryContext instance, and one of the threads will be blocked on the monitor of that // RootAggregatedMemoryContext instance even before calling the QueryContext methods (the monitors of // RootAggregatedMemoryContext instance and this will be acquired in the same order). long totalMemory = memoryPool.getQueryMemoryReservation(queryId); if (delta >= 0) { enforceTotalMemoryLimit(totalMemory, delta, maxTotalMemory); return memoryPool.reserve(queryId, allocationTag, delta); } memoryPool.free(queryId, allocationTag, -delta); return NOT_BLOCKED; }
SpillSpaceTracker spillSpaceTracker = new SpillSpaceTracker(new DataSize(1, GIGABYTE)); TaskContext taskContext = new DefaultQueryContext( new QueryId("test"), new DataSize(256, MEGABYTE), new DataSize(256, MEGABYTE), spillSpaceTracker) .addTaskContext(new TaskStateMachine(new TaskId("query", 0, 0), localQueryRunner.getExecutor()), session, false,
DefaultQueryContext queryContext = new DefaultQueryContext( new QueryId("query"), new DataSize(10, BYTE), queryContext.getQueryMemoryContext().initializeLocalMemoryContexts("test"); LocalMemoryContext userMemoryContext = queryContext.getQueryMemoryContext().localUserMemoryContext(); LocalMemoryContext revocableMemoryContext = queryContext.getQueryMemoryContext().localRevocableMemoryContext(); assertTrue(userMemoryContext.setBytes(3).isDone()); assertTrue(revocableMemoryContext.setBytes(5).isDone()); queryContext.setMemoryPool(reservedPool);
@GuardedBy("this") private void enforceUserMemoryLimit(long allocated, long delta, long maxMemory) { if (allocated + delta > maxMemory) { throw exceededLocalUserMemoryLimit(succinctBytes(maxMemory), getAdditionalFailureInfo(allocated, delta)); } }
/** * Deadlock is possible for concurrent user and system allocations when updateSystemMemory()/updateUserMemory * calls queryMemoryContext.getUserMemory()/queryMemoryContext.getSystemMemory(), respectively. * * @see this##updateSystemMemory(long) for details. */ private synchronized ListenableFuture<?> updateUserMemory(String allocationTag, long delta) { if (delta >= 0) { enforceUserMemoryLimit(queryMemoryContext.getUserMemory(), delta, maxUserMemory); return memoryPool.reserve(queryId, allocationTag, delta); } memoryPool.free(queryId, allocationTag, -delta); return NOT_BLOCKED; }
@BeforeMethod public void setUpTest() { memoryPool = new MemoryPool(new MemoryPoolId("test"), memoryPoolSize); queryContext = new DefaultQueryContext( new QueryId("test_query"), queryMaxMemory, queryMaxTotalMemory, memoryPool, new TestingGcMonitor(), notificationExecutor, yieldExecutor, queryMaxSpillSize, spillSpaceTracker); taskContext = queryContext.addTaskContext( new TaskStateMachine(new TaskId("query", 0, 0), notificationExecutor), testSessionBuilder().build(), true, true, OptionalInt.empty()); pipelineContext = taskContext.addPipelineContext(0, true, true, false); driverContext = pipelineContext.addDriverContext(); operatorContext = driverContext.addOperatorContext(1, new PlanNodeId("a"), "test-operator"); }
@GuardedBy("this") private void enforceTotalMemoryLimit(long allocated, long delta, long maxMemory) { if (allocated + delta > maxMemory) { throw exceededLocalTotalMemoryLimit(succinctBytes(maxMemory), getAdditionalFailureInfo(allocated, delta)); } }
private TaskContext newTestingTaskContext(ScheduledExecutorService taskNotificationExecutor, ScheduledExecutorService driverYieldExecutor, TaskStateMachine taskStateMachine) { DefaultQueryContext queryContext = new DefaultQueryContext( new QueryId("queryid"), new DataSize(1, MEGABYTE), new DataSize(2, MEGABYTE), new MemoryPool(new MemoryPoolId("test"), new DataSize(1, GIGABYTE)), new TestingGcMonitor(), taskNotificationExecutor, driverYieldExecutor, new DataSize(1, MEGABYTE), new SpillSpaceTracker(new DataSize(1, GIGABYTE))); return queryContext.addTaskContext(taskStateMachine, TEST_SESSION, false, false, OptionalInt.empty()); }
return new DefaultQueryContext( queryId, maxQueryUserMemoryPerNode,
DefaultQueryContext queryContext = new DefaultQueryContext(taskId.getQueryId(), new DataSize(1, MEGABYTE), new DataSize(2, MEGABYTE), new DataSize(1, MEGABYTE), spillSpaceTracker); this.taskContext = queryContext.addTaskContext(taskStateMachine, TEST_SESSION, true, true, totalPartitions);
private SqlTask newSqlTask() { TaskId taskId = new TaskId("query", 0, idGeneator.incrementAndGet()); URI location = URI.create("fake://task/" + taskId); return createSqlTask( taskId, location, "fake", new DefaultQueryContext(new QueryId("query"), new DataSize(1, MEGABYTE), new DataSize(2, MEGABYTE), memoryPool, new TestingGcMonitor(), executor, scheduledExecutor, new DataSize(1, GIGABYTE), spillSpaceTracker), sqlTaskExecutionFactory, executor, Functions.identity(), new DataSize(32, MEGABYTE), new CounterStat()); } }
public SqlTask createInitialTask() { TaskId taskId = new TaskId("query", 0, nextTaskId.incrementAndGet()); URI location = URI.create("fake://task/" + taskId); DefaultQueryContext queryContext = new DefaultQueryContext(new QueryId("query"), new DataSize(1, MEGABYTE), new DataSize(2, MEGABYTE), new MemoryPool(new MemoryPoolId("test"), new DataSize(1, GIGABYTE)), new TestingGcMonitor(), taskNotificationExecutor, driverYieldExecutor, new DataSize(1, MEGABYTE), new SpillSpaceTracker(new DataSize(1, GIGABYTE))); queryContext.addTaskContext(new TaskStateMachine(taskId, taskNotificationExecutor), testSessionBuilder().build(), false, false, OptionalInt.empty()); return createSqlTask( taskId, location, "fake", queryContext, sqlTaskExecutionFactory, taskNotificationExecutor, Functions.identity(), new DataSize(32, MEGABYTE), new CounterStat()); } }
private void setUp(Supplier<List<Driver>> driversSupplier) { checkState(localQueryRunner == null, "Already set up"); Session session = testSessionBuilder() .setCatalog("tpch") .setSchema("tiny") .setSystemProperty("task_default_concurrency", "1") .build(); localQueryRunner = queryRunnerWithInitialTransaction(session); // add tpch localQueryRunner.createCatalog("tpch", new TpchConnectorFactory(1), ImmutableMap.of()); userPool = new MemoryPool(new MemoryPoolId("test"), TEN_MEGABYTES); fakeQueryId = new QueryId("fake"); SpillSpaceTracker spillSpaceTracker = new SpillSpaceTracker(new DataSize(1, GIGABYTE)); DefaultQueryContext queryContext = new DefaultQueryContext(new QueryId("query"), TEN_MEGABYTES, new DataSize(20, MEGABYTE), userPool, new TestingGcMonitor(), localQueryRunner.getExecutor(), localQueryRunner.getScheduler(), TEN_MEGABYTES, spillSpaceTracker); taskContext = createTaskContext(queryContext, localQueryRunner.getExecutor(), session); drivers = driversSupplier.get(); }
DefaultQueryContext queryContext = new DefaultQueryContext( new QueryId("test"), new DataSize(1, GIGABYTE), .addTaskContext(new TaskStateMachine(new TaskId("query", 0, 0), localQueryRunner.getExecutor()), localQueryRunner.getDefaultSession(), false,
DefaultQueryContext queryContext = new DefaultQueryContext( queryId, new DataSize(512, MEGABYTE),