public synchronized void addLocation(URI location, TaskId remoteSourceTaskId)
    {
        requireNonNull(location, "location is null");

        // Ignore new locations after close
        // NOTE: this MUST happen before checking no more locations is checked
        if (closed.get()) {
            return;
        }

        // ignore duplicate locations
        if (allClients.containsKey(location)) {
            return;
        }

        // already removed
        if (removedRemoteSourceTaskIds.contains(remoteSourceTaskId)) {
            return;
        }

        checkState(!noMoreLocations, "No more locations already set");

        RpcShuffleClient resultClient;
        switch (location.getScheme().toLowerCase(Locale.ENGLISH)) {
            case "http":
            case "https":
                resultClient = new HttpRpcShuffleClient(httpClient, location);
                break;
            case "thrift":
                resultClient = new ThriftRpcShuffleClient(driftClient, location);
                break;
            default:
                throw new PrestoException(GENERIC_INTERNAL_ERROR, "unsupported task result client scheme " + location.getScheme());
        }

        PageBufferClient client = new PageBufferClient(
                resultClient,
                maxErrorDuration,
                acknowledgePages,
                location,
                new ExchangeClientCallback(),
                scheduler,
                pageBufferClientCallbackExecutor);
        allClients.put(location, client);
        checkState(taskIdToLocationMap.put(remoteSourceTaskId, location) == null, "Duplicate remoteSourceTaskId: " + remoteSourceTaskId);
        queuedClients.add(client);

        scheduleRequestIfNecessary();
    }
@Test(timeOut = 10000)
    public void testAddLocation()
            throws Exception
    {
        DataSize bufferCapacity = new DataSize(32, MEGABYTE);
        DataSize maxResponseSize = new DataSize(10, MEGABYTE);
        MockExchangeRequestProcessor processor = new MockExchangeRequestProcessor(maxResponseSize);

        ExchangeClient exchangeClient = createExchangeClient(processor, bufferCapacity, maxResponseSize);

        URI location1 = URI.create("http://localhost:8081/foo");
        processor.addPage(location1, createPage(1));
        processor.addPage(location1, createPage(2));
        processor.addPage(location1, createPage(3));
        processor.setComplete(location1);
        exchangeClient.addLocation(location1, TaskId.valueOf("foo.0.0.0.0"));

        assertFalse(exchangeClient.isClosed());
        assertPageEquals(getNextPage(exchangeClient), createPage(1));
        assertFalse(exchangeClient.isClosed());
        assertPageEquals(getNextPage(exchangeClient), createPage(2));
        assertFalse(exchangeClient.isClosed());
        assertPageEquals(getNextPage(exchangeClient), createPage(3));

        assertFalse(tryGetFutureValue(exchangeClient.isBlocked(), 10, MILLISECONDS).isPresent());
        assertFalse(exchangeClient.isClosed());

        URI location2 = URI.create("http://localhost:8082/bar");
        processor.addPage(location2, createPage(4));
        processor.addPage(location2, createPage(5));
        processor.addPage(location2, createPage(6));
        processor.setComplete(location2);
        exchangeClient.addLocation(location2, TaskId.valueOf("bar.0.0.0.0"));

        assertFalse(exchangeClient.isClosed());
        assertPageEquals(getNextPage(exchangeClient), createPage(4));
        assertFalse(exchangeClient.isClosed());
        assertPageEquals(getNextPage(exchangeClient), createPage(5));
        assertFalse(exchangeClient.isClosed());
        assertPageEquals(getNextPage(exchangeClient), createPage(6));

        assertFalse(tryGetFutureValue(exchangeClient.isBlocked(), 10, MILLISECONDS).isPresent());
        assertFalse(exchangeClient.isClosed());

        exchangeClient.noMoreLocations();
        // The transition to closed may happen asynchronously, since it requires that all the HTTP clients
        // receive a final GONE response, so just spin until it's closed or the test times out.
        while (!exchangeClient.isClosed()) {
            Thread.sleep(1);
        }

        ImmutableMap<URI, PageBufferClientStatus> statuses = uniqueIndex(exchangeClient.getStatus().getPageBufferClientStatuses(), PageBufferClientStatus::getUri);
        assertStatus(statuses.get(location1), location1, "closed", 3, 3, 3, "not scheduled");
        assertStatus(statuses.get(location2), location2, "closed", 3, 3, 3, "not scheduled");
    }