| // Copyright 2019 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.remote.util; |
| |
| import static com.google.common.truth.Truth.assertThat; |
| import static org.junit.Assert.fail; |
| |
| import com.google.bytestream.ByteStreamGrpc.ByteStreamImplBase; |
| import com.google.bytestream.ByteStreamProto.WriteRequest; |
| import com.google.bytestream.ByteStreamProto.WriteResponse; |
| import com.google.common.util.concurrent.ListenableFuture; |
| import com.google.common.util.concurrent.ListenableScheduledFuture; |
| import com.google.common.util.concurrent.ListeningScheduledExecutorService; |
| import com.google.devtools.build.lib.remote.RemoteRetrier; |
| import com.google.devtools.build.lib.remote.Retrier; |
| import com.google.devtools.build.lib.remote.Retrier.Backoff; |
| import com.google.protobuf.ByteString; |
| import io.grpc.Status; |
| import io.grpc.Status.Code; |
| import io.grpc.stub.StreamObserver; |
| import java.util.Collection; |
| import java.util.List; |
| import java.util.concurrent.Callable; |
| import java.util.concurrent.ExecutionException; |
| import java.util.concurrent.Future; |
| import java.util.concurrent.TimeUnit; |
| import java.util.concurrent.TimeoutException; |
| import java.util.function.Predicate; |
| import java.util.function.Supplier; |
| |
| /** Test utilities */ |
| public class TestUtils { |
| |
| public static RemoteRetrier newRemoteRetrier( |
| Supplier<Backoff> backoff, |
| Predicate<? super Exception> shouldRetry, |
| ListeningScheduledExecutorService retryScheduler) { |
| ZeroDelayListeningScheduledExecutorService zeroDelayRetryScheduler = |
| new ZeroDelayListeningScheduledExecutorService(retryScheduler); |
| return new RemoteRetrier( |
| backoff, |
| (e) -> Status.fromThrowable(e).getCode() != Code.CANCELLED && shouldRetry.test(e), |
| zeroDelayRetryScheduler, |
| Retrier.ALLOW_ALL_CALLS, |
| (millis) -> { |
| /* don't wait in tests */ |
| }); |
| } |
| |
| /** |
| * Wraps around a {@link ListeningScheduledExecutorService} and schedules all tasks with zero |
| * delay. |
| */ |
| private static class ZeroDelayListeningScheduledExecutorService |
| implements ListeningScheduledExecutorService { |
| |
| private final ListeningScheduledExecutorService delegate; |
| |
| ZeroDelayListeningScheduledExecutorService(ListeningScheduledExecutorService delegate) { |
| this.delegate = delegate; |
| } |
| |
| @Override |
| public ListenableScheduledFuture<?> schedule(Runnable runnable, long l, TimeUnit timeUnit) { |
| return delegate.schedule(runnable, 0, timeUnit); |
| } |
| |
| @Override |
| public <V> ListenableScheduledFuture<V> schedule( |
| Callable<V> callable, long l, TimeUnit timeUnit) { |
| return delegate.schedule(callable, 0, timeUnit); |
| } |
| |
| @Override |
| public ListenableScheduledFuture<?> scheduleAtFixedRate( |
| Runnable runnable, long l, long l1, TimeUnit timeUnit) { |
| return delegate.scheduleAtFixedRate(runnable, 0, 0, timeUnit); |
| } |
| |
| @Override |
| public ListenableScheduledFuture<?> scheduleWithFixedDelay( |
| Runnable runnable, long l, long l1, TimeUnit timeUnit) { |
| return delegate.scheduleWithFixedDelay(runnable, 0, 0, timeUnit); |
| } |
| |
| @Override |
| public void shutdown() { |
| delegate.shutdown(); |
| } |
| |
| @Override |
| public List<Runnable> shutdownNow() { |
| return delegate.shutdownNow(); |
| } |
| |
| @Override |
| public boolean isShutdown() { |
| return delegate.isShutdown(); |
| } |
| |
| @Override |
| public boolean isTerminated() { |
| return delegate.isTerminated(); |
| } |
| |
| @Override |
| public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException { |
| return delegate.awaitTermination(timeout, unit); |
| } |
| |
| @Override |
| public <T> ListenableFuture<T> submit(Callable<T> callable) { |
| return delegate.submit(callable); |
| } |
| |
| @Override |
| public ListenableFuture<?> submit(Runnable runnable) { |
| return delegate.submit(runnable); |
| } |
| |
| @Override |
| public <T> ListenableFuture<T> submit(Runnable runnable, T t) { |
| return delegate.submit(runnable, t); |
| } |
| |
| @Override |
| public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks) |
| throws InterruptedException { |
| return delegate.invokeAll(tasks); |
| } |
| |
| @Override |
| public <T> List<Future<T>> invokeAll( |
| Collection<? extends Callable<T>> tasks, long timeout, TimeUnit unit) |
| throws InterruptedException { |
| return delegate.invokeAll(tasks, timeout, unit); |
| } |
| |
| @Override |
| public <T> T invokeAny(Collection<? extends Callable<T>> tasks) |
| throws InterruptedException, ExecutionException { |
| return delegate.invokeAny(tasks); |
| } |
| |
| @Override |
| public <T> T invokeAny(Collection<? extends Callable<T>> tasks, long timeout, TimeUnit unit) |
| throws InterruptedException, ExecutionException, TimeoutException { |
| return delegate.invokeAny(tasks, timeout, unit); |
| } |
| |
| @Override |
| public void execute(Runnable command) { |
| delegate.execute(command); |
| } |
| } |
| |
| public static final ByteStreamImplBase newNoErrorByteStreamService(byte[] blob) { |
| return new ByteStreamImplBase() { |
| @Override |
| public StreamObserver<WriteRequest> write(StreamObserver<WriteResponse> streamObserver) { |
| return new StreamObserver<WriteRequest>() { |
| |
| byte[] receivedData = new byte[blob.length]; |
| long nextOffset = 0; |
| |
| @Override |
| public void onNext(WriteRequest writeRequest) { |
| if (nextOffset == 0) { |
| assertThat(writeRequest.getResourceName()).isNotEmpty(); |
| assertThat(writeRequest.getResourceName()).endsWith(String.valueOf(blob.length)); |
| } else { |
| assertThat(writeRequest.getResourceName()).isEmpty(); |
| } |
| |
| assertThat(writeRequest.getWriteOffset()).isEqualTo(nextOffset); |
| |
| ByteString data = writeRequest.getData(); |
| |
| System.arraycopy(data.toByteArray(), 0, receivedData, (int) nextOffset, data.size()); |
| |
| nextOffset += data.size(); |
| boolean lastWrite = blob.length == nextOffset; |
| assertThat(writeRequest.getFinishWrite()).isEqualTo(lastWrite); |
| } |
| |
| @Override |
| public void onError(Throwable throwable) { |
| fail("onError should never be called."); |
| } |
| |
| @Override |
| public void onCompleted() { |
| assertThat(nextOffset).isEqualTo(blob.length); |
| assertThat(receivedData).isEqualTo(blob); |
| |
| WriteResponse response = |
| WriteResponse.newBuilder().setCommittedSize(nextOffset).build(); |
| streamObserver.onNext(response); |
| streamObserver.onCompleted(); |
| } |
| }; |
| } |
| }; |
| } |
| } |