Examples with NodeScheduler io.trino.execution.scheduler.NodeScheduler used on opensource projects

Example 1 with NodeScheduler

use of io.trino.execution.scheduler.NodeScheduler in project trino by trinodb.

the class ServerMainModule method setup.

@Override
protected void setup(Binder binder) {
    ServerConfig serverConfig = buildConfigObject(ServerConfig.class);
    if (serverConfig.isCoordinator()) {
        install(new CoordinatorModule());
    } else {
        install(new WorkerModule());
    }
    binder.bind(StartupStatus.class).in(Scopes.SINGLETON);
    configBinder(binder).bindConfigDefaults(HttpServerConfig.class, httpServerConfig -> {
        httpServerConfig.setAdminEnabled(false);
    });
    binder.bind(PreparedStatementEncoder.class).in(Scopes.SINGLETON);
    binder.bind(HttpRequestSessionContextFactory.class).in(Scopes.SINGLETON);
    install(new InternalCommunicationModule());
    configBinder(binder).bindConfig(FeaturesConfig.class);
    configBinder(binder).bindConfig(OptimizerConfig.class);
    configBinder(binder).bindConfig(ProtocolConfig.class);
    binder.bind(SqlParser.class).in(Scopes.SINGLETON);
    jaxrsBinder(binder).bind(ThrowableMapper.class);
    configBinder(binder).bindConfig(QueryManagerConfig.class);
    configBinder(binder).bindConfig(SqlEnvironmentConfig.class);
    newOptionalBinder(binder, ExplainAnalyzeContext.class);
    binder.bind(StatementAnalyzerFactory.class).in(Scopes.SINGLETON);
    // GC Monitor
    binder.bind(GcMonitor.class).to(JmxGcMonitor.class).in(Scopes.SINGLETON);
    newExporter(binder).export(GcMonitor.class).withGeneratedName();
    // session properties
    newSetBinder(binder, SystemSessionPropertiesProvider.class).addBinding().to(SystemSessionProperties.class);
    binder.bind(SessionPropertyManager.class).in(Scopes.SINGLETON);
    binder.bind(SystemSessionProperties.class).in(Scopes.SINGLETON);
    binder.bind(SessionPropertyDefaults.class).in(Scopes.SINGLETON);
    // schema properties
    binder.bind(SchemaPropertyManager.class).in(Scopes.SINGLETON);
    // table properties
    binder.bind(TablePropertyManager.class).in(Scopes.SINGLETON);
    // materialized view properties
    binder.bind(MaterializedViewPropertyManager.class).in(Scopes.SINGLETON);
    // column properties
    binder.bind(ColumnPropertyManager.class).in(Scopes.SINGLETON);
    // analyze properties
    binder.bind(AnalyzePropertyManager.class).in(Scopes.SINGLETON);
    // table procedures properties
    binder.bind(TableProceduresPropertyManager.class).in(Scopes.SINGLETON);
    // node manager
    discoveryBinder(binder).bindSelector("trino");
    binder.bind(DiscoveryNodeManager.class).in(Scopes.SINGLETON);
    binder.bind(InternalNodeManager.class).to(DiscoveryNodeManager.class).in(Scopes.SINGLETON);
    newExporter(binder).export(DiscoveryNodeManager.class).withGeneratedName();
    httpClientBinder(binder).bindHttpClient("node-manager", ForNodeManager.class).withTracing().withConfigDefaults(config -> {
        config.setIdleTimeout(new Duration(30, SECONDS));
        config.setRequestTimeout(new Duration(10, SECONDS));
    });
    // node scheduler
    // TODO: remove from NodePartitioningManager and move to CoordinatorModule
    configBinder(binder).bindConfig(NodeSchedulerConfig.class);
    binder.bind(NodeScheduler.class).in(Scopes.SINGLETON);
    binder.bind(NodeTaskMap.class).in(Scopes.SINGLETON);
    newExporter(binder).export(NodeScheduler.class).withGeneratedName();
    // network topology
    // TODO: move to CoordinatorModule when NodeScheduler is moved
    install(conditionalModule(NodeSchedulerConfig.class, config -> UNIFORM == config.getNodeSchedulerPolicy(), new UniformNodeSelectorModule()));
    install(conditionalModule(NodeSchedulerConfig.class, config -> TOPOLOGY == config.getNodeSchedulerPolicy(), new TopologyAwareNodeSelectorModule()));
    // task execution
    configBinder(binder).bindConfig(FailureInjectionConfig.class);
    binder.bind(FailureInjector.class).in(Scopes.SINGLETON);
    jaxrsBinder(binder).bind(TaskResource.class);
    newExporter(binder).export(TaskResource.class).withGeneratedName();
    jaxrsBinder(binder).bind(TaskExecutorResource.class);
    newExporter(binder).export(TaskExecutorResource.class).withGeneratedName();
    binder.bind(TaskManagementExecutor.class).in(Scopes.SINGLETON);
    binder.bind(SqlTaskManager.class).in(Scopes.SINGLETON);
    binder.bind(TableExecuteContextManager.class).in(Scopes.SINGLETON);
    // memory revoking scheduler
    binder.bind(MemoryRevokingScheduler.class).in(Scopes.SINGLETON);
    // Add monitoring for JVM pauses
    binder.bind(PauseMeter.class).in(Scopes.SINGLETON);
    newExporter(binder).export(PauseMeter.class).withGeneratedName();
    configBinder(binder).bindConfig(MemoryManagerConfig.class);
    configBinder(binder).bindConfig(NodeMemoryConfig.class);
    binder.bind(LocalMemoryManager.class).in(Scopes.SINGLETON);
    binder.bind(LocalMemoryManagerExporter.class).in(Scopes.SINGLETON);
    newOptionalBinder(binder, VersionEmbedder.class).setDefault().to(EmbedVersion.class).in(Scopes.SINGLETON);
    newExporter(binder).export(SqlTaskManager.class).withGeneratedName();
    binder.bind(TaskExecutor.class).in(Scopes.SINGLETON);
    newExporter(binder).export(TaskExecutor.class).withGeneratedName();
    binder.bind(MultilevelSplitQueue.class).in(Scopes.SINGLETON);
    newExporter(binder).export(MultilevelSplitQueue.class).withGeneratedName();
    binder.bind(LocalExecutionPlanner.class).in(Scopes.SINGLETON);
    configBinder(binder).bindConfig(CompilerConfig.class);
    binder.bind(ExpressionCompiler.class).in(Scopes.SINGLETON);
    newExporter(binder).export(ExpressionCompiler.class).withGeneratedName();
    binder.bind(PageFunctionCompiler.class).in(Scopes.SINGLETON);
    newExporter(binder).export(PageFunctionCompiler.class).withGeneratedName();
    configBinder(binder).bindConfig(TaskManagerConfig.class);
    binder.bind(IndexJoinLookupStats.class).in(Scopes.SINGLETON);
    newExporter(binder).export(IndexJoinLookupStats.class).withGeneratedName();
    binder.bind(AsyncHttpExecutionMBean.class).in(Scopes.SINGLETON);
    newExporter(binder).export(AsyncHttpExecutionMBean.class).withGeneratedName();
    binder.bind(JoinFilterFunctionCompiler.class).in(Scopes.SINGLETON);
    newExporter(binder).export(JoinFilterFunctionCompiler.class).withGeneratedName();
    binder.bind(JoinCompiler.class).in(Scopes.SINGLETON);
    newExporter(binder).export(JoinCompiler.class).withGeneratedName();
    binder.bind(OrderingCompiler.class).in(Scopes.SINGLETON);
    newExporter(binder).export(OrderingCompiler.class).withGeneratedName();
    binder.bind(PagesIndex.Factory.class).to(PagesIndex.DefaultFactory.class);
    newOptionalBinder(binder, OperatorFactories.class).setDefault().to(TrinoOperatorFactories.class).in(Scopes.SINGLETON);
    jaxrsBinder(binder).bind(PagesResponseWriter.class);
    // exchange client
    binder.bind(DirectExchangeClientSupplier.class).to(DirectExchangeClientFactory.class).in(Scopes.SINGLETON);
    httpClientBinder(binder).bindHttpClient("exchange", ForExchange.class).withTracing().withFilter(GenerateTraceTokenRequestFilter.class).withConfigDefaults(config -> {
        config.setIdleTimeout(new Duration(30, SECONDS));
        config.setRequestTimeout(new Duration(10, SECONDS));
        config.setMaxConnectionsPerServer(250);
        config.setMaxContentLength(DataSize.of(32, MEGABYTE));
    });
    configBinder(binder).bindConfig(DirectExchangeClientConfig.class);
    binder.bind(ExchangeExecutionMBean.class).in(Scopes.SINGLETON);
    newExporter(binder).export(ExchangeExecutionMBean.class).withGeneratedName();
    // execution
    binder.bind(LocationFactory.class).to(HttpLocationFactory.class).in(Scopes.SINGLETON);
    // memory manager
    jaxrsBinder(binder).bind(MemoryResource.class);
    jsonCodecBinder(binder).bindJsonCodec(MemoryInfo.class);
    // transaction manager
    configBinder(binder).bindConfig(TransactionManagerConfig.class);
    // data stream provider
    binder.bind(PageSourceManager.class).in(Scopes.SINGLETON);
    binder.bind(PageSourceProvider.class).to(PageSourceManager.class).in(Scopes.SINGLETON);
    // page sink provider
    binder.bind(PageSinkManager.class).in(Scopes.SINGLETON);
    binder.bind(PageSinkProvider.class).to(PageSinkManager.class).in(Scopes.SINGLETON);
    // metadata
    binder.bind(StaticCatalogStore.class).in(Scopes.SINGLETON);
    configBinder(binder).bindConfig(StaticCatalogStoreConfig.class);
    binder.bind(MetadataManager.class).in(Scopes.SINGLETON);
    binder.bind(Metadata.class).to(MetadataManager.class).in(Scopes.SINGLETON);
    newOptionalBinder(binder, SystemSecurityMetadata.class).setDefault().to(DisabledSystemSecurityMetadata.class).in(Scopes.SINGLETON);
    binder.bind(GlobalFunctionCatalog.class).in(Scopes.SINGLETON);
    binder.bind(TypeOperatorsCache.class).in(Scopes.SINGLETON);
    newExporter(binder).export(TypeOperatorsCache.class).as(factory -> factory.generatedNameOf(TypeOperators.class));
    binder.bind(BlockTypeOperators.class).in(Scopes.SINGLETON);
    newExporter(binder).export(TypeOperatorsCache.class).withGeneratedName();
    binder.bind(ProcedureRegistry.class).in(Scopes.SINGLETON);
    binder.bind(TableProceduresRegistry.class).in(Scopes.SINGLETON);
    binder.bind(PlannerContext.class).in(Scopes.SINGLETON);
    // function
    binder.bind(FunctionManager.class).in(Scopes.SINGLETON);
    newSetBinder(binder, FunctionBundle.class);
    binder.bind(RegisterFunctionBundles.class).asEagerSingleton();
    // type
    binder.bind(TypeAnalyzer.class).in(Scopes.SINGLETON);
    jsonBinder(binder).addDeserializerBinding(Type.class).to(TypeDeserializer.class);
    jsonBinder(binder).addDeserializerBinding(TypeSignature.class).to(TypeSignatureDeserializer.class);
    jsonBinder(binder).addKeyDeserializerBinding(TypeSignature.class).to(TypeSignatureKeyDeserializer.class);
    binder.bind(TypeRegistry.class).in(Scopes.SINGLETON);
    binder.bind(TypeManager.class).to(InternalTypeManager.class).in(Scopes.SINGLETON);
    newSetBinder(binder, Type.class);
    // split manager
    binder.bind(SplitManager.class).in(Scopes.SINGLETON);
    // node partitioning manager
    binder.bind(NodePartitioningManager.class).in(Scopes.SINGLETON);
    // index manager
    binder.bind(IndexManager.class).in(Scopes.SINGLETON);
    // handle resolver
    binder.install(new HandleJsonModule());
    // connector
    binder.bind(ConnectorManager.class).in(Scopes.SINGLETON);
    // system connector
    binder.install(new SystemConnectorModule());
    // slice
    jsonBinder(binder).addSerializerBinding(Slice.class).to(SliceSerializer.class);
    jsonBinder(binder).addDeserializerBinding(Slice.class).to(SliceDeserializer.class);
    // expression
    jsonBinder(binder).addSerializerBinding(Expression.class).to(ExpressionSerializer.class);
    jsonBinder(binder).addDeserializerBinding(Expression.class).to(ExpressionDeserializer.class);
    // split monitor
    binder.bind(SplitMonitor.class).in(Scopes.SINGLETON);
    // version and announcement
    binder.bind(NodeVersion.class).toInstance(new NodeVersion(nodeVersion));
    discoveryBinder(binder).bindHttpAnnouncement("trino").addProperty("node_version", nodeVersion).addProperty("coordinator", String.valueOf(serverConfig.isCoordinator()));
    // server info resource
    jaxrsBinder(binder).bind(ServerInfoResource.class);
    // node status resource
    jaxrsBinder(binder).bind(StatusResource.class);
    // plugin manager
    binder.bind(PluginManager.class).in(Scopes.SINGLETON);
    newOptionalBinder(binder, PluginsProvider.class).setDefault().to(ServerPluginsProvider.class).in(Scopes.SINGLETON);
    configBinder(binder).bindConfig(ServerPluginsProviderConfig.class);
    binder.bind(CatalogManager.class).in(Scopes.SINGLETON);
    // block encodings
    binder.bind(BlockEncodingManager.class).in(Scopes.SINGLETON);
    jsonBinder(binder).addSerializerBinding(Block.class).to(BlockJsonSerde.Serializer.class);
    jsonBinder(binder).addDeserializerBinding(Block.class).to(BlockJsonSerde.Deserializer.class);
    binder.bind(BlockEncodingSerde.class).to(InternalBlockEncodingSerde.class).in(Scopes.SINGLETON);
    // thread visualizer
    jaxrsBinder(binder).bind(ThreadResource.class);
    // PageSorter
    binder.bind(PageSorter.class).to(PagesIndexPageSorter.class).in(Scopes.SINGLETON);
    // PageIndexer
    binder.bind(PageIndexerFactory.class).to(GroupByHashPageIndexerFactory.class).in(Scopes.SINGLETON);
    // Finalizer
    binder.bind(FinalizerService.class).in(Scopes.SINGLETON);
    // Spiller
    binder.bind(SpillerFactory.class).to(GenericSpillerFactory.class).in(Scopes.SINGLETON);
    binder.bind(SingleStreamSpillerFactory.class).to(FileSingleStreamSpillerFactory.class).in(Scopes.SINGLETON);
    binder.bind(PartitioningSpillerFactory.class).to(GenericPartitioningSpillerFactory.class).in(Scopes.SINGLETON);
    binder.bind(SpillerStats.class).in(Scopes.SINGLETON);
    newExporter(binder).export(SpillerFactory.class).withGeneratedName();
    binder.bind(LocalSpillManager.class).in(Scopes.SINGLETON);
    configBinder(binder).bindConfig(NodeSpillConfig.class);
    // Dynamic Filtering
    configBinder(binder).bindConfig(DynamicFilterConfig.class);
    // dispatcher
    // TODO remove dispatcher fromm ServerMainModule, and bind dependent components only on coordinators
    newOptionalBinder(binder, DispatchManager.class);
    // Added for RuleStatsSystemTable
    // TODO: remove this when system tables are bound separately for coordinator and worker
    newOptionalBinder(binder, RuleStatsRecorder.class);
    // cleanup
    binder.bind(ExecutorCleanup.class).in(Scopes.SINGLETON);
}

Example 2 with NodeScheduler

use of io.trino.execution.scheduler.NodeScheduler in project trino by trinodb.

the class NodePartitioningManager method getNodePartitioningMap.

public NodePartitionMap getNodePartitioningMap(Session session, PartitioningHandle partitioningHandle) {
    requireNonNull(session, "session is null");
    requireNonNull(partitioningHandle, "partitioningHandle is null");
    if (partitioningHandle.getConnectorHandle() instanceof SystemPartitioningHandle) {
        return ((SystemPartitioningHandle) partitioningHandle.getConnectorHandle()).getNodePartitionMap(session, nodeScheduler);
    }
    ConnectorBucketNodeMap connectorBucketNodeMap = getConnectorBucketNodeMap(session, partitioningHandle);
    // safety check for crazy partitioning
    checkArgument(connectorBucketNodeMap.getBucketCount() < 1_000_000, "Too many buckets in partitioning: %s", connectorBucketNodeMap.getBucketCount());
    List<InternalNode> bucketToNode;
    if (connectorBucketNodeMap.hasFixedMapping()) {
        bucketToNode = getFixedMapping(connectorBucketNodeMap);
    } else {
        CatalogName catalogName = partitioningHandle.getConnectorId().orElseThrow(() -> new IllegalArgumentException("No connector ID for partitioning handle: " + partitioningHandle));
        bucketToNode = createArbitraryBucketToNode(nodeScheduler.createNodeSelector(session, Optional.of(catalogName)).allNodes(), connectorBucketNodeMap.getBucketCount());
    }
    int[] bucketToPartition = new int[connectorBucketNodeMap.getBucketCount()];
    BiMap<InternalNode, Integer> nodeToPartition = HashBiMap.create();
    int nextPartitionId = 0;
    for (int bucket = 0; bucket < bucketToNode.size(); bucket++) {
        InternalNode node = bucketToNode.get(bucket);
        Integer partitionId = nodeToPartition.get(node);
        if (partitionId == null) {
            partitionId = nextPartitionId++;
            nodeToPartition.put(node, partitionId);
        }
        bucketToPartition[bucket] = partitionId;
    }
    List<InternalNode> partitionToNode = IntStream.range(0, nodeToPartition.size()).mapToObj(partitionId -> nodeToPartition.inverse().get(partitionId)).collect(toImmutableList());
    return new NodePartitionMap(partitionToNode, bucketToPartition, getSplitToBucket(session, partitioningHandle));
}

Example 3 with NodeScheduler

use of io.trino.execution.scheduler.NodeScheduler in project trino by trinodb.

the class TaskTestUtils method createTestingPlanner.

public static LocalExecutionPlanner createTestingPlanner() {
    PageSourceManager pageSourceManager = new PageSourceManager();
    pageSourceManager.addConnectorPageSourceProvider(CONNECTOR_ID, new TestingPageSourceProvider());
    // we don't start the finalizer so nothing will be collected, which is ok for a test
    FinalizerService finalizerService = new FinalizerService();
    BlockTypeOperators blockTypeOperators = new BlockTypeOperators(PLANNER_CONTEXT.getTypeOperators());
    NodeScheduler nodeScheduler = new NodeScheduler(new UniformNodeSelectorFactory(new InMemoryNodeManager(), new NodeSchedulerConfig().setIncludeCoordinator(true), new NodeTaskMap(finalizerService)));
    NodePartitioningManager nodePartitioningManager = new NodePartitioningManager(nodeScheduler, blockTypeOperators);
    PageFunctionCompiler pageFunctionCompiler = new PageFunctionCompiler(PLANNER_CONTEXT.getFunctionManager(), 0);
    return new LocalExecutionPlanner(PLANNER_CONTEXT, createTestingTypeAnalyzer(PLANNER_CONTEXT), Optional.empty(), pageSourceManager, new IndexManager(), nodePartitioningManager, new PageSinkManager(), new MockDirectExchangeClientSupplier(), new ExpressionCompiler(PLANNER_CONTEXT.getFunctionManager(), pageFunctionCompiler), pageFunctionCompiler, new JoinFilterFunctionCompiler(PLANNER_CONTEXT.getFunctionManager()), new IndexJoinLookupStats(), new TaskManagerConfig(), new GenericSpillerFactory((types, spillContext, memoryContext) -> {
        throw new UnsupportedOperationException();
    }), (types, spillContext, memoryContext) -> {
        throw new UnsupportedOperationException();
    }, (types, partitionFunction, spillContext, memoryContext) -> {
        throw new UnsupportedOperationException();
    }, new PagesIndex.TestingFactory(false), new JoinCompiler(PLANNER_CONTEXT.getTypeOperators()), new TrinoOperatorFactories(), new OrderingCompiler(PLANNER_CONTEXT.getTypeOperators()), new DynamicFilterConfig(), blockTypeOperators, new TableExecuteContextManager(), new ExchangeManagerRegistry(new ExchangeHandleResolver()));
}

Example 4 with NodeScheduler

use of io.trino.execution.scheduler.NodeScheduler in project trino by trinodb.

the class TestLocalExchange method setUp.

@BeforeMethod
public void setUp() {
    NodeScheduler nodeScheduler = new NodeScheduler(new UniformNodeSelectorFactory(new InMemoryNodeManager(), new NodeSchedulerConfig().setIncludeCoordinator(true), new NodeTaskMap(new FinalizerService())));
    nodePartitioningManager = new NodePartitioningManager(nodeScheduler, new BlockTypeOperators(new TypeOperators()));
}

Example 5 with NodeScheduler

use of io.trino.execution.scheduler.NodeScheduler in project trino by trinodb.

the class TestNodeScheduler method testTopologyAwareScheduling.

@Test(timeOut = 60 * 1000)
public void testTopologyAwareScheduling() {
    NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService);
    InMemoryNodeManager nodeManager = new InMemoryNodeManager();
    ImmutableList.Builder<InternalNode> nodeBuilder = ImmutableList.builder();
    nodeBuilder.add(new InternalNode("node1", URI.create("http://host1.rack1:11"), NodeVersion.UNKNOWN, false));
    nodeBuilder.add(new InternalNode("node2", URI.create("http://host2.rack1:12"), NodeVersion.UNKNOWN, false));
    nodeBuilder.add(new InternalNode("node3", URI.create("http://host3.rack2:13"), NodeVersion.UNKNOWN, false));
    ImmutableList<InternalNode> nodes = nodeBuilder.build();
    nodeManager.addNode(CONNECTOR_ID, nodes);
    // contents of taskMap indicate the node-task map for the current stage
    Map<InternalNode, RemoteTask> taskMap = new HashMap<>();
    NodeSchedulerConfig nodeSchedulerConfig = new NodeSchedulerConfig().setMaxSplitsPerNode(25).setIncludeCoordinator(false).setMaxPendingSplitsPerTask(20);
    TestNetworkTopology topology = new TestNetworkTopology();
    NodeSelectorFactory nodeSelectorFactory = new TopologyAwareNodeSelectorFactory(topology, nodeManager, nodeSchedulerConfig, nodeTaskMap, getNetworkTopologyConfig());
    NodeScheduler nodeScheduler = new NodeScheduler(nodeSelectorFactory);
    NodeSelector nodeSelector = nodeScheduler.createNodeSelector(session, Optional.of(CONNECTOR_ID));
    // Fill up the nodes with non-local data
    ImmutableSet.Builder<Split> nonRackLocalBuilder = ImmutableSet.builder();
    for (int i = 0; i < (25 + 11) * 3; i++) {
        nonRackLocalBuilder.add(new Split(CONNECTOR_ID, new TestSplitRemote(HostAddress.fromParts("data.other_rack", 1)), Lifespan.taskWide()));
    }
    Set<Split> nonRackLocalSplits = nonRackLocalBuilder.build();
    Multimap<InternalNode, Split> assignments = nodeSelector.computeAssignments(nonRackLocalSplits, ImmutableList.copyOf(taskMap.values())).getAssignments();
    MockRemoteTaskFactory remoteTaskFactory = new MockRemoteTaskFactory(remoteTaskExecutor, remoteTaskScheduledExecutor);
    int task = 0;
    for (InternalNode node : assignments.keySet()) {
        TaskId taskId = new TaskId(new StageId("test", 1), task, 0);
        task++;
        MockRemoteTaskFactory.MockRemoteTask remoteTask = remoteTaskFactory.createTableScanTask(taskId, node, ImmutableList.copyOf(assignments.get(node)), nodeTaskMap.createPartitionedSplitCountTracker(node, taskId));
        remoteTask.startSplits(25);
        nodeTaskMap.addTask(node, remoteTask);
        taskMap.put(node, remoteTask);
    }
    // Continue assigning to fill up part of the queue
    nonRackLocalSplits = Sets.difference(nonRackLocalSplits, new HashSet<>(assignments.values()));
    assignments = nodeSelector.computeAssignments(nonRackLocalSplits, ImmutableList.copyOf(taskMap.values())).getAssignments();
    for (InternalNode node : assignments.keySet()) {
        RemoteTask remoteTask = taskMap.get(node);
        remoteTask.addSplits(ImmutableMultimap.<PlanNodeId, Split>builder().putAll(new PlanNodeId("sourceId"), assignments.get(node)).build());
    }
    nonRackLocalSplits = Sets.difference(nonRackLocalSplits, new HashSet<>(assignments.values()));
    // Check that 3 of the splits were rejected, since they're non-local
    assertEquals(nonRackLocalSplits.size(), 3);
    // Assign rack-local splits
    ImmutableSet.Builder<Split> rackLocalSplits = ImmutableSet.builder();
    HostAddress dataHost1 = HostAddress.fromParts("data.rack1", 1);
    HostAddress dataHost2 = HostAddress.fromParts("data.rack2", 1);
    for (int i = 0; i < 6 * 2; i++) {
        rackLocalSplits.add(new Split(CONNECTOR_ID, new TestSplitRemote(dataHost1), Lifespan.taskWide()));
    }
    for (int i = 0; i < 6; i++) {
        rackLocalSplits.add(new Split(CONNECTOR_ID, new TestSplitRemote(dataHost2), Lifespan.taskWide()));
    }
    assignments = nodeSelector.computeAssignments(rackLocalSplits.build(), ImmutableList.copyOf(taskMap.values())).getAssignments();
    for (InternalNode node : assignments.keySet()) {
        RemoteTask remoteTask = taskMap.get(node);
        remoteTask.addSplits(ImmutableMultimap.<PlanNodeId, Split>builder().putAll(new PlanNodeId("sourceId"), assignments.get(node)).build());
    }
    Set<Split> unassigned = Sets.difference(rackLocalSplits.build(), new HashSet<>(assignments.values()));
    // Compute the assignments a second time to account for the fact that some splits may not have been assigned due to asynchronous
    // loading of the NetworkLocationCache
    assignments = nodeSelector.computeAssignments(unassigned, ImmutableList.copyOf(taskMap.values())).getAssignments();
    for (InternalNode node : assignments.keySet()) {
        RemoteTask remoteTask = taskMap.get(node);
        remoteTask.addSplits(ImmutableMultimap.<PlanNodeId, Split>builder().putAll(new PlanNodeId("sourceId"), assignments.get(node)).build());
    }
    unassigned = Sets.difference(unassigned, new HashSet<>(assignments.values()));
    assertEquals(unassigned.size(), 3);
    int rack1 = 0;
    int rack2 = 0;
    for (Split split : unassigned) {
        String rack = topology.locate(split.getAddresses().get(0)).getSegments().get(0);
        switch(rack) {
            case "rack1":
                rack1++;
                break;
            case "rack2":
                rack2++;
                break;
            default:
                throw new AssertionError("Unexpected rack: " + rack);
        }
    }
    assertEquals(rack1, 2);
    assertEquals(rack2, 1);
    // Assign local splits
    ImmutableSet.Builder<Split> localSplits = ImmutableSet.builder();
    localSplits.add(new Split(CONNECTOR_ID, new TestSplitRemote(HostAddress.fromParts("host1.rack1", 1)), Lifespan.taskWide()));
    localSplits.add(new Split(CONNECTOR_ID, new TestSplitRemote(HostAddress.fromParts("host2.rack1", 1)), Lifespan.taskWide()));
    localSplits.add(new Split(CONNECTOR_ID, new TestSplitRemote(HostAddress.fromParts("host3.rack2", 1)), Lifespan.taskWide()));
    assignments = nodeSelector.computeAssignments(localSplits.build(), ImmutableList.copyOf(taskMap.values())).getAssignments();
    assertEquals(assignments.size(), 3);
    assertEquals(assignments.keySet().size(), 3);
}