src/test/java/com/google/devtools/build/lib/concurrent/RequestBatcherTest.java - bazel - Git at Google


            
            // Copyright 2024 The Bazel Authors. All rights reserved.

            
            //

            
            // 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.google.devtools.build.lib.concurrent;


            import static com.google.common.truth.Truth.assertThat;

            
            import static com.google.common.util.concurrent.Futures.immediateFailedFuture;

            
            import static com.google.common.util.concurrent.Futures.immediateFuture;

            
            import static com.google.common.util.concurrent.Futures.immediateVoidFuture;

            
            import static com.google.common.util.concurrent.MoreExecutors.directExecutor;

            
            import static com.google.devtools.build.lib.concurrent.PaddedAddresses.createPaddedBaseAddress;

            
            import static com.google.devtools.build.lib.concurrent.PaddedAddresses.getAlignedAddress;

            
            import static java.util.concurrent.ForkJoinPool.commonPool;


            import com.google.common.collect.Lists;

            
            import com.google.common.util.concurrent.Futures;

            
            import com.google.common.util.concurrent.ListenableFuture;

            
            import com.google.common.util.concurrent.SettableFuture;

            
            import com.google.devtools.build.lib.concurrent.RequestBatcher.RequestResponse;

            
            import com.google.devtools.build.lib.unsafe.UnsafeProvider;

            
            import java.lang.ref.Cleaner;

            
            import java.util.ArrayList;

            
            import java.util.List;

            
            import java.util.concurrent.ConcurrentLinkedQueue;

            
            import java.util.concurrent.CountDownLatch;

            
            import java.util.concurrent.Executor;

            
            import java.util.concurrent.LinkedBlockingQueue;

            
            import java.util.concurrent.Semaphore;

            
            import org.junit.Test;

            
            import org.junit.runner.RunWith;

            
            import org.junit.runners.JUnit4;

            
            import sun.misc.Unsafe;


            @RunWith(JUnit4.class)

            
            @SuppressWarnings("SunApi") // TODO: b/359688989 - clean this up

            
            public final class RequestBatcherTest {

            
              private static final Cleaner cleaner = Cleaner.create();


              @Test

            
              public void simpleSubmit_executes() throws Exception {

            
                // This covers Step 1A in the documentation.

            
                var batcher =

            
                    RequestBatcher.<Request, Response>create(

            
                        commonPool(),

            
                        requests -> immediateFuture(respondTo(requests)),

            
                        /* targetWorkerCount= */ 1);

            
                ListenableFuture<Response> response = batcher.submit(new Request(1));

            
                assertThat(response.get()).isEqualTo(new Response(1));

            
              }


              @Test

            
              public void queueOverflow_sleeps() throws Exception {

            
                // This covers the overflow case of Step 1B in the documentation.


                var multiplexer = new SettableMultiplexer();

            
                var batcher =

            
                    RequestBatcher.<Request, Response>create(

            
                        commonPool(), multiplexer, /* targetWorkerCount= */ 1);

            
                ListenableFuture<Response> response0 = batcher.submit(new Request(0));

            
                BatchedRequestResponses requestResponses0 = multiplexer.queue.take();

            
                // The first worker is busy until requestResponse0 is populated.


                var responses = new ArrayList<ListenableFuture<Response>>();

            
                // With the single available worker being busy, we can completely fill the queue.

            
                for (int i = 0; i < ConcurrentFifo.MAX_ELEMENTS; i++) {

            
                  responses.add(batcher.submit(new Request(i + 1)));

            
                }


                // The next request triggers a queue overflow. Since this ends up blocking, we do it in another

            
                // thread.

            
                var overflowStarting = new CountDownLatch(1);

            
                var overflowAdded = new CountDownLatch(1);

            
                // A new thread needs must used here instead of commonPool because there are test environments

            
                // where commonPool has only a single thread.

            
                new Thread(

            
                        () -> {

            
                          overflowStarting.countDown();

            
                          responses.add(batcher.submit(new Request(ConcurrentFifo.MAX_ELEMENTS + 1)));

            
                          overflowAdded.countDown();

            
                        })

            
                    .start();


                // The following assertion will occasionally fail if the overflow submit above does not block.

            
                overflowStarting.await();

            
                assertThat(responses).hasSize(ConcurrentFifo.MAX_ELEMENTS);


                // Responding to the first batch enables the overflowing element to enter the queue.

            
                requestResponses0.setSimpleResponses();

            
                assertThat(response0.get()).isEqualTo(new Response(0));

            
                overflowAdded.await();

            
                assertThat(responses).hasSize(ConcurrentFifo.MAX_ELEMENTS + 1);


                // Responds to all remaining batches.

            
                int batchSize = RequestBatcher.BATCH_SIZE + 1;

            
                int batchCount = responses.size() / batchSize;

            
                assertThat(responses).hasSize(batchCount * batchSize);


                for (int i = 0; i < batchCount; i++) {

            
                  multiplexer.queue.take().setSimpleResponses();

            
                }


                for (int i = 0; i < ConcurrentFifo.MAX_ELEMENTS + 1; i++) {

            
                  assertThat(responses.get(i).get()).isEqualTo(new Response(i + 1));

            
                }

            
              }


              @Test

            
              public void submitWithWorkersFull_enqueuesThenExecutes() throws Exception {

            
                // This covers Step 1B, Step 2A and Step 3 of the documentation.


                var multiplexer = new SettableMultiplexer();

            
                var batcher =

            
                    RequestBatcher.<Request, Response>create(

            
                        commonPool(), multiplexer, /* targetWorkerCount= */ 1);

            
                ListenableFuture<Response> response1 = batcher.submit(new Request(1));

            
                BatchedRequestResponses requestResponses1 = multiplexer.queue.take();


                ListenableFuture<Response> response2 = batcher.submit(new Request(2));

            
                // The first batch is not yet complete. The 2nd request waits in an internal queue. Ideally, we

            
                // could make a stronger assertion here, that the 2nd batch executes only after the first one is

            
                // done.

            
                assertThat(multiplexer.queue).isEmpty();


                requestResponses1.setSimpleResponses();


                // With the first batch done, the worker picks up the enqueued 2nd request and executes it.

            
                BatchedRequestResponses requestResponses2 = multiplexer.queue.take();

            
                requestResponses2.setSimpleResponses();


                assertThat(response1.get()).isEqualTo(new Response(1));

            
                assertThat(response2.get()).isEqualTo(new Response(2));

            
              }


              @Test

            
              public void concurrentWorkCompletion_startsNewWorker() throws Exception {

            
                // This covers Step 1B and Step 2B of the documentation.


                // This test uses fakes to achieve the narrow set of conditions needed to reach this code path.

            
                long baseAddress = createPaddedBaseAddress(4);


                var executor = new FakeExecutor();

            
                var fifo =

            
                    new FakeConcurrentFifo(

            
                        getAlignedAddress(baseAddress, /* offset= */ 1),

            
                        getAlignedAddress(baseAddress, /* offset= */ 2),

            
                        getAlignedAddress(baseAddress, /* offset= */ 3));

            
                long countersAddress = getAlignedAddress(baseAddress, /* offset= */ 0);

            
                var batcher =

            
                    new RequestBatcher<Request, Response>(

            
                        executor,

            
                        requests -> immediateFuture(respondTo(requests)),

            
                        /* targetWorkerCount= */ 1,

            
                        countersAddress,

            
                        fifo);

            
                cleaner.register(batcher, new AddressFreer(baseAddress));


                // Submits a request. This starts a worker to run the batch, but it gets blocked on the executor

            
                // and can't continue.

            
                ListenableFuture<Response> response1 = batcher.submit(new Request(1));


                // Submits a 2nd request. This request observes that there are enough active workers so it tries

            
                // to enqueue an element. It gets blocked at the queue.

            
                @SuppressWarnings("FutureReturnValueIgnored")

            
                var response2 = SettableFuture.<ListenableFuture<Response>>create();

            
                commonPool().execute(() -> response2.set(batcher.submit(new Request(2))));

            
                // Waits until the 2nd request starts enqueuing.

            
                fifo.tryAppendTokens.acquireUninterruptibly();


                // Allows the 1st worker to continue. It will go idle. There are two Runnables, one to run the

            
                // continuation logic and one that sets the response callbacks. Runs both.

            
                executor.queue.take().run();

            
                executor.queue.take().run();


                assertThat(response1.get()).isEqualTo(new Response(1));


                // Verifies that there's absolutely nothing inflight in the batcher.

            
                assertThat(UNSAFE.getIntVolatile(null, countersAddress)).isEqualTo(0);


                // Allows the 2nd request to enqueue and complete processing.

            
                fifo.appendPermits.release();

            
                executor.queue.take().run();

            
                executor.queue.take().run();


                assertThat(response2.get().get()).isEqualTo(new Response(2));

            
              }


              @Test

            
              public void randomRaces_executeCorrectly() throws Exception {

            
                var batcher =

            
                    RequestBatcher.<Request, Response>create(

            
                        commonPool(),

            
                        requests -> immediateFuture(respondTo(requests)),

            
                        /* targetWorkerCount= */ 4);


                var results = new ConcurrentLinkedQueue<ListenableFuture<Void>>();

            
                final int requestCount = 4_000_000;

            
                var allStarted = new CountDownLatch(requestCount);

            
                for (int i = 0; i < requestCount; i++) {

            
                  final int iForCapture = i;

            
                  commonPool()

            
                      .execute(

            
                          () -> {

            
                            results.add(

            
                                Futures.transformAsync(

            
                                    batcher.submit(new Request(iForCapture)),

            
                                    response ->

            
                                        response.x() == iForCapture

            
                                            ? immediateVoidFuture()

            
                                            : immediateFailedFuture(

            
                                                new AssertionError(

            
                                                    String.format(

            
                                                        "expected %d got %s", iForCapture, response))),

            
                                    directExecutor()));

            
                            allStarted.countDown();

            
                          });

            
                }

            
                allStarted.await();

            
                // Throws ExecutionException if there are any errors.

            
                var unused = Futures.whenAllSucceed(results).call(() -> null, directExecutor()).get();

            
              }


              private static class FakeExecutor implements Executor {

            
                private final LinkedBlockingQueue<Runnable> queue = new LinkedBlockingQueue<>();


                @Override

            
                public void execute(Runnable runnable) {

            
                  queue.add(runnable);

            
                }

            
              }


              private static class FakeConcurrentFifo

            
                  extends ConcurrentFifo<RequestResponse<Request, Response>> {

            
                private final ConcurrentLinkedQueue<RequestResponse<Request, Response>> queue =

            
                    new ConcurrentLinkedQueue<>();


                private final Semaphore tryAppendTokens = new Semaphore(0);

            
                private final Semaphore appendPermits = new Semaphore(0);


                private FakeConcurrentFifo(long sizeAddress, long appendIndexAddress, long takeIndexAddress) {

            
                  super(RequestResponse.class, sizeAddress, appendIndexAddress, takeIndexAddress);

            
                }


                @Override

            
                boolean tryAppend(RequestResponse<Request, Response> task) {

            
                  tryAppendTokens.release();

            
                  appendPermits.acquireUninterruptibly();

            
                  queue.add(task);

            
                  return true;

            
                }


                @Override

            
                RequestResponse<Request, Response> take() {

            
                  return queue.poll();

            
                }

            
              }


              private static class SettableMultiplexer

            
                  implements RequestBatcher.Multiplexer<Request, Response> {

            
                private final LinkedBlockingQueue<BatchedRequestResponses> queue = new LinkedBlockingQueue<>();


                @Override

            
                public ListenableFuture<List<Response>> execute(List<Request> requests) {

            
                  var responses = SettableFuture.<List<Response>>create();

            
                  queue.add(new BatchedRequestResponses(requests, responses));

            
                  return responses;

            
                }

            
              }


              private record BatchedRequestResponses(

            
                  List<Request> requests, SettableFuture<List<Response>> responses) {

            
                private void setSimpleResponses() {

            
                  responses().set(respondTo(requests()));

            
                }

            
              }


              private record Request(int x) {}


              private record Response(int x) {}


              private static List<Response> respondTo(List<Request> requests) {

            
                return Lists.transform(requests, request -> new Response(request.x()));

            
              }


              private static final Unsafe UNSAFE = UnsafeProvider.unsafe();

            
            }
	// Copyright 2024 The Bazel Authors. All rights reserved.
	//
	// 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.google.devtools.build.lib.concurrent;

	import static com.google.common.truth.Truth.assertThat;
	import static com.google.common.util.concurrent.Futures.immediateFailedFuture;
	import static com.google.common.util.concurrent.Futures.immediateFuture;
	import static com.google.common.util.concurrent.Futures.immediateVoidFuture;
	import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
	import static com.google.devtools.build.lib.concurrent.PaddedAddresses.createPaddedBaseAddress;
	import static com.google.devtools.build.lib.concurrent.PaddedAddresses.getAlignedAddress;
	import static java.util.concurrent.ForkJoinPool.commonPool;

	import com.google.common.collect.Lists;
	import com.google.common.util.concurrent.Futures;
	import com.google.common.util.concurrent.ListenableFuture;
	import com.google.common.util.concurrent.SettableFuture;
	import com.google.devtools.build.lib.concurrent.RequestBatcher.RequestResponse;
	import com.google.devtools.build.lib.unsafe.UnsafeProvider;
	import java.lang.ref.Cleaner;
	import java.util.ArrayList;
	import java.util.List;
	import java.util.concurrent.ConcurrentLinkedQueue;
	import java.util.concurrent.CountDownLatch;
	import java.util.concurrent.Executor;
	import java.util.concurrent.LinkedBlockingQueue;
	import java.util.concurrent.Semaphore;
	import org.junit.Test;
	import org.junit.runner.RunWith;
	import org.junit.runners.JUnit4;
	import sun.misc.Unsafe;

	@RunWith(JUnit4.class)
	@SuppressWarnings("SunApi") // TODO: b/359688989 - clean this up
	public final class RequestBatcherTest {
	private static final Cleaner cleaner = Cleaner.create();

	@Test
	public void simpleSubmit_executes() throws Exception {
	// This covers Step 1A in the documentation.
	var batcher =
	RequestBatcher.<Request, Response>create(
	commonPool(),
	requests -> immediateFuture(respondTo(requests)),
	/* targetWorkerCount= */ 1);
	ListenableFuture<Response> response = batcher.submit(new Request(1));
	assertThat(response.get()).isEqualTo(new Response(1));
	}

	@Test
	public void queueOverflow_sleeps() throws Exception {
	// This covers the overflow case of Step 1B in the documentation.

	var multiplexer = new SettableMultiplexer();
	var batcher =
	RequestBatcher.<Request, Response>create(
	commonPool(), multiplexer, /* targetWorkerCount= */ 1);
	ListenableFuture<Response> response0 = batcher.submit(new Request(0));
	BatchedRequestResponses requestResponses0 = multiplexer.queue.take();
	// The first worker is busy until requestResponse0 is populated.

	var responses = new ArrayList<ListenableFuture<Response>>();
	// With the single available worker being busy, we can completely fill the queue.
	for (int i = 0; i < ConcurrentFifo.MAX_ELEMENTS; i++) {
	responses.add(batcher.submit(new Request(i + 1)));
	}

	// The next request triggers a queue overflow. Since this ends up blocking, we do it in another
	// thread.
	var overflowStarting = new CountDownLatch(1);
	var overflowAdded = new CountDownLatch(1);
	// A new thread needs must used here instead of commonPool because there are test environments
	// where commonPool has only a single thread.
	new Thread(
	() -> {
	overflowStarting.countDown();
	responses.add(batcher.submit(new Request(ConcurrentFifo.MAX_ELEMENTS + 1)));
	overflowAdded.countDown();
	})
	.start();

	// The following assertion will occasionally fail if the overflow submit above does not block.
	overflowStarting.await();
	assertThat(responses).hasSize(ConcurrentFifo.MAX_ELEMENTS);

	// Responding to the first batch enables the overflowing element to enter the queue.
	requestResponses0.setSimpleResponses();
	assertThat(response0.get()).isEqualTo(new Response(0));
	overflowAdded.await();
	assertThat(responses).hasSize(ConcurrentFifo.MAX_ELEMENTS + 1);

	// Responds to all remaining batches.
	int batchSize = RequestBatcher.BATCH_SIZE + 1;
	int batchCount = responses.size() / batchSize;
	assertThat(responses).hasSize(batchCount * batchSize);

	for (int i = 0; i < batchCount; i++) {
	multiplexer.queue.take().setSimpleResponses();
	}

	for (int i = 0; i < ConcurrentFifo.MAX_ELEMENTS + 1; i++) {
	assertThat(responses.get(i).get()).isEqualTo(new Response(i + 1));
	}
	}

	@Test
	public void submitWithWorkersFull_enqueuesThenExecutes() throws Exception {
	// This covers Step 1B, Step 2A and Step 3 of the documentation.

	var multiplexer = new SettableMultiplexer();
	var batcher =
	RequestBatcher.<Request, Response>create(
	commonPool(), multiplexer, /* targetWorkerCount= */ 1);
	ListenableFuture<Response> response1 = batcher.submit(new Request(1));
	BatchedRequestResponses requestResponses1 = multiplexer.queue.take();

	ListenableFuture<Response> response2 = batcher.submit(new Request(2));
	// The first batch is not yet complete. The 2nd request waits in an internal queue. Ideally, we
	// could make a stronger assertion here, that the 2nd batch executes only after the first one is
	// done.
	assertThat(multiplexer.queue).isEmpty();

	requestResponses1.setSimpleResponses();

	// With the first batch done, the worker picks up the enqueued 2nd request and executes it.
	BatchedRequestResponses requestResponses2 = multiplexer.queue.take();
	requestResponses2.setSimpleResponses();

	assertThat(response1.get()).isEqualTo(new Response(1));
	assertThat(response2.get()).isEqualTo(new Response(2));
	}

	@Test
	public void concurrentWorkCompletion_startsNewWorker() throws Exception {
	// This covers Step 1B and Step 2B of the documentation.

	// This test uses fakes to achieve the narrow set of conditions needed to reach this code path.
	long baseAddress = createPaddedBaseAddress(4);

	var executor = new FakeExecutor();
	var fifo =
	new FakeConcurrentFifo(
	getAlignedAddress(baseAddress, /* offset= */ 1),
	getAlignedAddress(baseAddress, /* offset= */ 2),
	getAlignedAddress(baseAddress, /* offset= */ 3));
	long countersAddress = getAlignedAddress(baseAddress, /* offset= */ 0);
	var batcher =
	new RequestBatcher<Request, Response>(
	executor,
	requests -> immediateFuture(respondTo(requests)),
	/* targetWorkerCount= */ 1,
	countersAddress,
	fifo);
	cleaner.register(batcher, new AddressFreer(baseAddress));

	// Submits a request. This starts a worker to run the batch, but it gets blocked on the executor
	// and can't continue.
	ListenableFuture<Response> response1 = batcher.submit(new Request(1));

	// Submits a 2nd request. This request observes that there are enough active workers so it tries
	// to enqueue an element. It gets blocked at the queue.
	@SuppressWarnings("FutureReturnValueIgnored")
	var response2 = SettableFuture.<ListenableFuture<Response>>create();
	commonPool().execute(() -> response2.set(batcher.submit(new Request(2))));
	// Waits until the 2nd request starts enqueuing.
	fifo.tryAppendTokens.acquireUninterruptibly();

	// Allows the 1st worker to continue. It will go idle. There are two Runnables, one to run the
	// continuation logic and one that sets the response callbacks. Runs both.
	executor.queue.take().run();
	executor.queue.take().run();

	assertThat(response1.get()).isEqualTo(new Response(1));

	// Verifies that there's absolutely nothing inflight in the batcher.
	assertThat(UNSAFE.getIntVolatile(null, countersAddress)).isEqualTo(0);

	// Allows the 2nd request to enqueue and complete processing.
	fifo.appendPermits.release();
	executor.queue.take().run();
	executor.queue.take().run();

	assertThat(response2.get().get()).isEqualTo(new Response(2));
	}

	@Test
	public void randomRaces_executeCorrectly() throws Exception {
	var batcher =
	RequestBatcher.<Request, Response>create(
	commonPool(),
	requests -> immediateFuture(respondTo(requests)),
	/* targetWorkerCount= */ 4);

	var results = new ConcurrentLinkedQueue<ListenableFuture<Void>>();
	final int requestCount = 4_000_000;
	var allStarted = new CountDownLatch(requestCount);
	for (int i = 0; i < requestCount; i++) {
	final int iForCapture = i;
	commonPool()
	.execute(
	() -> {
	results.add(
	Futures.transformAsync(
	batcher.submit(new Request(iForCapture)),
	response ->
	response.x() == iForCapture
	? immediateVoidFuture()
	: immediateFailedFuture(
	new AssertionError(
	String.format(
	"expected %d got %s", iForCapture, response))),
	directExecutor()));
	allStarted.countDown();
	});
	}
	allStarted.await();
	// Throws ExecutionException if there are any errors.
	var unused = Futures.whenAllSucceed(results).call(() -> null, directExecutor()).get();
	}

	private static class FakeExecutor implements Executor {
	private final LinkedBlockingQueue<Runnable> queue = new LinkedBlockingQueue<>();

	@Override
	public void execute(Runnable runnable) {
	queue.add(runnable);
	}
	}

	private static class FakeConcurrentFifo
	extends ConcurrentFifo<RequestResponse<Request, Response>> {
	private final ConcurrentLinkedQueue<RequestResponse<Request, Response>> queue =
	new ConcurrentLinkedQueue<>();

	private final Semaphore tryAppendTokens = new Semaphore(0);
	private final Semaphore appendPermits = new Semaphore(0);

	private FakeConcurrentFifo(long sizeAddress, long appendIndexAddress, long takeIndexAddress) {
	super(RequestResponse.class, sizeAddress, appendIndexAddress, takeIndexAddress);
	}

	@Override
	boolean tryAppend(RequestResponse<Request, Response> task) {
	tryAppendTokens.release();
	appendPermits.acquireUninterruptibly();
	queue.add(task);
	return true;
	}

	@Override
	RequestResponse<Request, Response> take() {
	return queue.poll();
	}
	}

	private static class SettableMultiplexer
	implements RequestBatcher.Multiplexer<Request, Response> {
	private final LinkedBlockingQueue<BatchedRequestResponses> queue = new LinkedBlockingQueue<>();

	@Override
	public ListenableFuture<List<Response>> execute(List<Request> requests) {
	var responses = SettableFuture.<List<Response>>create();
	queue.add(new BatchedRequestResponses(requests, responses));
	return responses;
	}
	}

	private record BatchedRequestResponses(
	List<Request> requests, SettableFuture<List<Response>> responses) {
	private void setSimpleResponses() {
	responses().set(respondTo(requests()));
	}
	}

	private record Request(int x) {}

	private record Response(int x) {}

	private static List<Response> respondTo(List<Request> requests) {
	return Lists.transform(requests, request -> new Response(request.x()));
	}

	private static final Unsafe UNSAFE = UnsafeProvider.unsafe();
	}