public synchronized void removeRemoteSource(TaskId sourceTaskId)
    {
        requireNonNull(sourceTaskId, "sourceTaskId is null");

        // Ignore removeRemoteSource call if exchange client is already closed
        if (closed.get()) {
            return;
        }

        removedRemoteSourceTaskIds.add(sourceTaskId);

        URI location = taskIdToLocationMap.get(sourceTaskId);
        if (location == null) {
            return;
        }

        PageBufferClient client = allClients.get(location);
        if (client == null) {
            return;
        }

        closeQuietly(client);
        removedClients.add(client);
        completedClients.add(client);
    }
@Test
    public void testRemoveRemoteSource()
            throws Exception
    {
        DataSize bufferCapacity = new DataSize(1, BYTE);
        DataSize maxResponseSize = new DataSize(1, BYTE);
        MockExchangeRequestProcessor processor = new MockExchangeRequestProcessor(maxResponseSize);

        URI location1 = URI.create("http://localhost:8081/foo.0.0.0.0");
        TaskId taskId1 = TaskId.valueOf("foo.0.0.0.0");
        URI location2 = URI.create("http://localhost:8082/bar.0.0.0.0");
        TaskId taskId2 = TaskId.valueOf("bar.0.0.0.0");

        processor.addPage(location1, createPage(1));
        processor.addPage(location1, createPage(2));
        processor.addPage(location1, createPage(3));

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

        exchangeClient.addLocation(location1, taskId1);
        exchangeClient.addLocation(location2, taskId2);

        assertFalse(exchangeClient.isClosed());

        // Wait until exactly one page is buffered:
        //  * We cannot call ExchangeClient#pollPage() directly, since it will schedule the next request to buffer data,
        //    and this request to buffer data could win the race against the request to remove remote source.
        //  * Buffer capacity is set to 1 byte, so only one page can be buffered.
        waitUntilEquals(() -> exchangeClient.getStatus().getBufferedPages(), 1, new Duration(5, SECONDS));
        assertEquals(exchangeClient.getStatus().getBufferedPages(), 1);

        // remove remote source
        exchangeClient.removeRemoteSource(taskId1);

        // the previously buffered page will still be read out
        assertPageEquals(getNextPage(exchangeClient), createPage(1));

        // client should not receive any further pages from removed remote source
        assertNull(exchangeClient.pollPage());
        assertEquals(exchangeClient.getStatus().getBufferedPages(), 0);

        // add pages to another source
        processor.addPage(location2, createPage(4));
        processor.addPage(location2, createPage(5));
        processor.addPage(location2, createPage(6));
        processor.setComplete(location2);

        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);
        }

        ExchangeClientStatus exchangeClientStatus = exchangeClient.getStatus();
        Optional<PageBufferClientStatus> clientStatusOptional1 = exchangeClientStatus.getPageBufferClientStatuses()
                .stream()
                .filter(pageBufferClientStatus -> pageBufferClientStatus.getUri().equals(location1)).findFirst();
        assertTrue(clientStatusOptional1.isPresent());
        assertStatus(clientStatusOptional1.get(), "closed", "not scheduled");

        Optional<PageBufferClientStatus> clientStatusOptional2 = exchangeClientStatus.getPageBufferClientStatuses()
                .stream()
                .filter(pageBufferClientStatus -> pageBufferClientStatus.getUri().equals(location2)).findFirst();
        assertTrue(clientStatusOptional2.isPresent());
        assertStatus(clientStatusOptional2.get(), "closed", "not scheduled");
    }