diff --git a/benchmark/build.gradle b/benchmark/build.gradle index 5b1bfac0f..8417aa01f 100644 --- a/benchmark/build.gradle +++ b/benchmark/build.gradle @@ -21,6 +21,7 @@ dependencies { compile project(':gax') compile project(':gax-grpc') compile "io.grpc:grpc-netty:${grpcVersion}" + compile 'org.openjdk.jmh:jmh-core:1.17.4' compile 'com.google.api.grpc:grpc-google-cloud-bigtable-v2:0.1.28' compile 'com.google.api.grpc:grpc-google-cloud-pubsub-v1:0.1.28' @@ -30,3 +31,4 @@ dependencies { if (project.properties.containsKey('include')) { jmh.include = [project.properties.get('include')] } +jmh.forceGC = true \ No newline at end of file diff --git a/benchmark/src/jmh/java/com/google/api/gax/grpc/BatchingBenchmark.java b/benchmark/src/jmh/java/com/google/api/gax/grpc/BatchingBenchmark.java new file mode 100644 index 000000000..453aab659 --- /dev/null +++ b/benchmark/src/jmh/java/com/google/api/gax/grpc/BatchingBenchmark.java @@ -0,0 +1,541 @@ +/* + * Copyright 2018, Google LLC All rights reserved. + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are + * met: + * + * * Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above + * copyright notice, this list of conditions and the following disclaimer + * in the documentation and/or other materials provided with the + * distribution. + * * Neither the name of Google LLC nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS + * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT + * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR + * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT + * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, + * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT + * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, + * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY + * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + */ +package com.google.api.gax.grpc; + +import com.google.api.core.ApiFuture; +import com.google.api.core.CurrentMillisClock; +import com.google.api.gax.batching.BatchingFlowController; +import com.google.api.gax.batching.BatchingSettings; +import com.google.api.gax.batching.BatchingThreshold; +import com.google.api.gax.batching.ElementCounter; +import com.google.api.gax.batching.FlowControlSettings; +import com.google.api.gax.batching.FlowController; +import com.google.api.gax.batching.FlowController.LimitExceededBehavior; +import com.google.api.gax.batching.NumericThreshold; +import com.google.api.gax.batching.PartitionKey; +import com.google.api.gax.batching.RequestBuilder; +import com.google.api.gax.batching.ThresholdBatcher; +import com.google.api.gax.grpc.batching.Batcher; +import com.google.api.gax.rpc.Batch; +import com.google.api.gax.rpc.BatchExecutor; +import com.google.api.gax.rpc.BatchedFuture; +import com.google.api.gax.rpc.BatchedRequestIssuer; +import com.google.api.gax.rpc.BatcherFactory; +import com.google.api.gax.rpc.BatchingCallSettings; +import com.google.api.gax.rpc.BatchingDescriptor; +import com.google.api.gax.rpc.ClientContext; +import com.google.api.gax.rpc.UnaryCallSettings; +import com.google.api.gax.rpc.UnaryCallable; +import com.google.common.base.Preconditions; +import com.google.common.collect.Lists; +import com.google.common.util.concurrent.MoreExecutors; +import com.google.protobuf.ByteString; +import com.google.pubsub.v1.PublishRequest; +import com.google.pubsub.v1.PublishResponse; +import com.google.pubsub.v1.PublishResponse.Builder; +import com.google.pubsub.v1.PublisherGrpc; +import com.google.pubsub.v1.PubsubMessage; +import io.grpc.CallOptions; +import io.grpc.ManagedChannel; +import io.grpc.ManagedChannelBuilder; +import io.grpc.Server; +import io.grpc.ServerBuilder; +import io.grpc.stub.StreamObserver; +import java.io.IOException; +import java.net.ServerSocket; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.Random; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.Semaphore; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OperationsPerInvocation; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.TearDown; +import org.openjdk.jmh.annotations.Warmup; +import org.openjdk.jmh.infra.BenchmarkParams; +import org.openjdk.jmh.infra.Blackhole; +import org.threeten.bp.Duration; + +/** + * Exploratory benchmarks for batching. + * + *

This benchmark tries to measure the overhead of gax's batching infrastructure. It uses a fake + * pubsub server to sink the requests. This is done to take netty context switchs into account. It + * tries to pin done performance sensitive portions. + * + *

Each iteration will send {@code OperationsPerInvocation} entries, split by {@code + * elementsPerBatch}, with at most {@code maxOutstandingElements} elements unsent. When the {@code + * maxOutstandingElements} is reached all benchmarks will block. + */ +@Fork(value = 1) +@BenchmarkMode(Mode.Throughput) +@Warmup(iterations = 10 ) +@Measurement(iterations = 10) +@State(Scope.Benchmark) +@OperationsPerInvocation(10_000) +@OutputTimeUnit(TimeUnit.SECONDS) +public class BatchingBenchmark { + private static final Logger LOG = Logger.getLogger(BatchingBenchmark.class.getName()); + private static final String TOPIC = "projects/fake-project/topics/fake-topic"; + private static final Random RANDOM = new Random(123); + + @SuppressWarnings("WeakerAccess") + @Param("1000") + int maxOutstandingElements; + + @SuppressWarnings("WeakerAccess") + @Param("100") + int elementsPerBatch; + + @SuppressWarnings("WeakerAccess") + @Param("1024") + int messageSize; + + private ScheduledExecutorService executor; + private Server fakeServer; + private ManagedChannel grpcChannel; + private ClientContext clientContext; + private ByteString[] payloads; + + private UnaryCallable baseCallable; + private UnaryCallable batchingCallable; + private ThresholdBatcher> pushingBatcher; + private Batcher impl1; + + @Setup + public void setup(BenchmarkParams benchmarkParams, Blackhole blackhole) throws IOException { + Preconditions.checkState(elementsPerBatch <= maxOutstandingElements); + Preconditions.checkState(maxOutstandingElements % elementsPerBatch == 0); + Preconditions.checkState(benchmarkParams.getOpsPerInvocation() % elementsPerBatch == 0); + + byte[] buffer = new byte[messageSize]; + payloads = new ByteString[100]; + for (int i = 0; i < payloads.length; i++) { + RANDOM.nextBytes(buffer); + payloads[i] = ByteString.copyFrom(buffer); + } + + executor = Executors.newScheduledThreadPool(4); + + final int availablePort; + try (ServerSocket ss = new ServerSocket(0)) { + availablePort = ss.getLocalPort(); + } + fakeServer = + ServerBuilder.forPort(availablePort).addService(new FakePubSub(blackhole)).build().start(); + + grpcChannel = + ManagedChannelBuilder.forAddress("localhost", availablePort).usePlaintext(true).build(); + + clientContext = + ClientContext.newBuilder() + .setExecutor(executor) + .setClock(CurrentMillisClock.getDefaultClock()) + .setDefaultCallContext( + GrpcCallContext.of( + grpcChannel, CallOptions.DEFAULT.withDeadlineAfter(1, TimeUnit.HOURS))) + .setTransportChannel(GrpcTransportChannel.create(grpcChannel)) + .build(); + + GrpcCallSettings grpcSettings = + GrpcCallSettings.newBuilder() + .setMethodDescriptor(PublisherGrpc.METHOD_PUBLISH) + .build(); + + UnaryCallSettings callSettings = + UnaryCallSettings.newUnaryCallSettingsBuilder() + .setSimpleTimeoutNoRetries(Duration.ofMinutes(1)) + .build(); + + baseCallable = + GrpcCallableFactory.createUnaryCallable(grpcSettings, callSettings, clientContext); + + BatchingCallSettings.Builder batchingCallSettings = + BatchingCallSettings.newBuilder(new FakeBatchingDescriptor()) + .setBatchingSettings( + BatchingSettings.newBuilder() + .setIsEnabled(true) + .setFlowControlSettings( + FlowControlSettings.newBuilder() + .setLimitExceededBehavior(LimitExceededBehavior.Block) + .setMaxOutstandingElementCount((long) maxOutstandingElements) + .build()) + // Not actually used because we generate messages faster + .setDelayThreshold(Duration.ofSeconds(5)) + .setElementCountThreshold((long) elementsPerBatch) + .build()); + batchingCallSettings.setSimpleTimeoutNoRetries(Duration.ofSeconds(10)); + + batchingCallable = + GrpcCallableFactory.createBatchingCallable( + grpcSettings, batchingCallSettings.build(), clientContext); + + BatcherFactory batcherFactory = + new BatcherFactory<>( + new FakeBatchingDescriptor(), + BatchingSettings.newBuilder() + .setIsEnabled(true) + .setFlowControlSettings( + FlowControlSettings.newBuilder() + .setLimitExceededBehavior(LimitExceededBehavior.Block) + .setMaxOutstandingElementCount((long) maxOutstandingElements) + .build()) + // Not actually used because we generate messages faster + .setDelayThreshold(Duration.ofSeconds(5)) + .setElementCountThreshold((long) elementsPerBatch) + .build(), + executor, + new FlowController( + FlowControlSettings.newBuilder() + .setLimitExceededBehavior(LimitExceededBehavior.Block) + .setMaxOutstandingElementCount((long) maxOutstandingElements) + .build())); + pushingBatcher = batcherFactory.getPushingBatcher(new PartitionKey(TOPIC)); + + ArrayList> batchingThresholds = Lists.newArrayList(); + batchingThresholds.add( + new NumericThreshold<>(elementsPerBatch, + new ElementCounter() { + @Override + public long count(PublishRequest element) { + return element.getMessagesCount(); + } + }) + ); + + FakeBatchingDescriptor batchingDescriptor = new FakeBatchingDescriptor(); + PartitionKey partitionKey = new PartitionKey(TOPIC); + FlowController flowController = new FlowController( + FlowControlSettings.newBuilder() + .setLimitExceededBehavior(LimitExceededBehavior.Block) + .setMaxOutstandingElementCount((long) maxOutstandingElements) + .build() + ); + BatchingFlowController batchingFlowController = new BatchingFlowController<>( + flowController, + new ElementCounter() { + @Override + public long count(PublishRequest element) { + return element.getMessagesCount(); + } + }, + new ElementCounter() { + @Override + public long count(PublishRequest element) { + return element.getSerializedSize(); + } + } + ); + impl1 = Batcher.newBuilder() + .setDescriptor(new FakeBatchingDescriptor()) + .setInnerCallable(baseCallable) + .setThresholds(batchingThresholds) + .setExecutor(executor) + .setMaxDelay(Duration.ofSeconds(5)) + .setReceiver(new BatchExecutor<>(batchingDescriptor,partitionKey)) + .setFlowController(batchingFlowController) + .build(); + } + + @TearDown + public void teardown() { + try { + grpcChannel.shutdown(); + if (!grpcChannel.awaitTermination(10, TimeUnit.SECONDS)) { + throw new TimeoutException(); + } + } catch (Exception e) { + LOG.log(Level.SEVERE, "Failed to close the grpc channel", e); + } + + try { + fakeServer.shutdown(); + if (!fakeServer.awaitTermination(10, TimeUnit.SECONDS)) { + throw new TimeoutException(); + } + } catch (Exception e) { + LOG.log(Level.SEVERE, "Failed to shutdown the fake server", e); + } + + try { + executor.shutdown(); + if (!executor.awaitTermination(10, TimeUnit.SECONDS)) { + throw new TimeoutException(); + } + } catch (Exception e) { + LOG.log(Level.SEVERE, "Failed to shutdown the the executor", e); + } + } + + /** Simple, direct implementation that should represent minimal overhead */ + @Benchmark + public void manualBaseline(BenchmarkParams benchmarkParams) throws Exception { + int messageCount = benchmarkParams.getOpsPerInvocation(); + int batchCount = messageCount / elementsPerBatch; + + final Semaphore outstandingElements = new Semaphore(maxOutstandingElements); + FinishLine finishLine = new FinishLine(batchCount) { + @Override + public void run() { + outstandingElements.release(elementsPerBatch); + super.run(); + } + }; + + for (int i = 0; i < batchCount; i++) { + outstandingElements.acquire(elementsPerBatch); + PublishRequest request = buildRequest(elementsPerBatch); + + final ApiFuture batchFuture = baseCallable.futureCall(request); + batchFuture.addListener(finishLine, MoreExecutors.directExecutor()); + } + + finishLine.waitForArrival(); + } + + /** + * The full gax batching stack. + * + *

Note: since gax doesn't support flushing, {@link OperationsPerInvocation} must be a multiple + * of {@link #elementsPerBatch}. + */ + @Benchmark + public void batchingCallableBenchmark(BenchmarkParams benchmarkParams) throws Exception { + int messageCount = benchmarkParams.getOpsPerInvocation(); + FinishLine finishLine = new FinishLine(messageCount); + + for (int i = 0; i < messageCount; i++) { + PublishRequest request = buildRequest(1); + + ApiFuture future = batchingCallable.futureCall(request); + future.addListener(finishLine, MoreExecutors.directExecutor()); + } + + finishLine.waitForArrival(); + } + + /** + * One layer down into the gax batching infrastructure: skips the {@link + * com.google.api.gax.rpc.BatchingCallable} and the {@link PartitionKey} routing. + */ + @Benchmark + public void pushingBatcherBenchmark(BenchmarkParams benchmarkParams) throws Exception { + int messageCount = benchmarkParams.getOpsPerInvocation(); + FinishLine finishLine = new FinishLine(messageCount); + + FakeBatchingDescriptor descriptor = new FakeBatchingDescriptor(); + + for (int i = 0; i < messageCount; i++) { + PublishRequest request = buildRequest(1); + + BatchedFuture future = new BatchedFuture<>(); + pushingBatcher.add(new Batch<>(descriptor, request, baseCallable, future)); + future.addListener(finishLine, MoreExecutors.directExecutor()); + } + + pushingBatcher.pushCurrentBatch(); + finishLine.waitForArrival(); + } + + /** + * 2 layers deep into the gax batching infrastructure that tries to avoid re-wrapping the request + * protos by pre-sizing the batches. + */ + @Benchmark + public void pushingBatcher2Benchmark(BenchmarkParams benchmarkParams) throws Exception { + int messageCount = benchmarkParams.getOpsPerInvocation(); + int batchCount = messageCount / elementsPerBatch; + FinishLine finishLine = new FinishLine(batchCount); + + FakeBatchingDescriptor descriptor = new FakeBatchingDescriptor(); + + for (int i = 0; i < batchCount; i++) { + PublishRequest request = buildRequest(elementsPerBatch); + BatchedFuture future = new BatchedFuture<>(); + pushingBatcher.add(new Batch<>(descriptor, request, baseCallable, future)); + future.addListener(finishLine, MoreExecutors.directExecutor()); + } + + pushingBatcher.pushCurrentBatch(); + finishLine.waitForArrival(); + } + + // Skip batch merging + @Benchmark + public void noBatchMerging(BenchmarkParams benchmarkParams) throws Exception { + int messageCount = benchmarkParams.getOpsPerInvocation(); + FinishLine finishLine = new FinishLine(messageCount); + + for (int i = 0; i < messageCount; i++) { + PublishRequest request = buildRequest(1); + ApiFuture future = impl1.add(request); + future.addListener(finishLine, MoreExecutors.directExecutor()); + } + + pushingBatcher.pushCurrentBatch(); + finishLine.waitForArrival(); + } + + // Helpers ------- + private PublishRequest buildRequest(int entryCount) { + PublishRequest.Builder builder = PublishRequest.newBuilder() + .setTopic(TOPIC); + + for (int i = 0; i < entryCount; i++) { + builder.addMessages( + PubsubMessage.newBuilder() + .setData(payloads[RANDOM.nextInt(payloads.length)]) + .setMessageId("message-" + i) + .build()); + } + + return builder.build(); + } + + static class FinishLine implements Runnable { + private final CountDownLatch countDownLatch; + + FinishLine(int eventCount) { + countDownLatch = new CountDownLatch(eventCount); + } + + @Override + public void run() { + countDownLatch.countDown(); + } + + void waitForArrival() throws TimeoutException, InterruptedException { + if (!countDownLatch.await(10, TimeUnit.MINUTES)) { + throw new TimeoutException("Timed out waiting for all elements to finish"); + } + } + } + + static class FakePubSub extends com.google.pubsub.v1.PublisherGrpc.PublisherImplBase { + private final Blackhole blackhole; + + FakePubSub(Blackhole blackhole) { + this.blackhole = blackhole; + } + + @Override + public void publish(PublishRequest request, StreamObserver responseObserver) { + blackhole.consume(request); + Builder responseBuilder = PublishResponse.newBuilder(); + for (PubsubMessage msg : request.getMessagesList()) { + responseBuilder.addMessageIds(msg.getMessageId()); + } + + responseObserver.onNext(responseBuilder.build()); + responseObserver.onCompleted(); + } + } + + static class FakeBatchingDescriptor + implements BatchingDescriptor { + @Override + public PartitionKey getBatchPartitionKey(PublishRequest request) { + return new PartitionKey(request.getTopic()); + } + + @Override + public RequestBuilder getRequestBuilder() { + return new RequestBuilder() { + private PublishRequest.Builder builder; + + @Override + public void appendRequest(PublishRequest request) { + if (builder == null) { + builder = request.toBuilder(); + } else { + builder.addAllMessages(request.getMessagesList()); + } + } + + @Override + public PublishRequest build() { + return builder.build(); + } + }; + } + + @Override + public void splitResponse( + PublishResponse batchResponse, + Collection> batch) { + int batchMessageIndex = 0; + for (BatchedRequestIssuer responder : batch) { + List subresponseElements = new ArrayList<>(); + long subresponseCount = responder.getMessageCount(); + for (int i = 0; i < subresponseCount; i++) { + subresponseElements.add(batchResponse.getMessageIds(batchMessageIndex)); + batchMessageIndex += 1; + } + PublishResponse response = + PublishResponse.newBuilder().addAllMessageIds(subresponseElements).build(); + responder.setResponse(response); + } + } + + @Override + public void splitException( + Throwable throwable, Collection> batch) { + for (BatchedRequestIssuer responder : batch) { + responder.setException(throwable); + } + } + + @Override + public long countElements(PublishRequest request) { + return request.getMessagesCount(); + } + + @Override + public long countBytes(PublishRequest request) { + return request.getSerializedSize(); + } + } +} diff --git a/benchmark/src/jmh/java/com/google/api/gax/grpc/batching/Batcher.java b/benchmark/src/jmh/java/com/google/api/gax/grpc/batching/Batcher.java new file mode 100644 index 000000000..b32b7901a --- /dev/null +++ b/benchmark/src/jmh/java/com/google/api/gax/grpc/batching/Batcher.java @@ -0,0 +1,283 @@ +/* + * Copyright 2018, Google LLC All rights reserved. + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are + * met: + * + * * Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above + * copyright notice, this list of conditions and the following disclaimer + * in the documentation and/or other materials provided with the + * distribution. + * * Neither the name of Google LLC nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS + * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT + * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR + * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT + * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, + * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT + * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, + * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY + * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + */ +package com.google.api.gax.grpc.batching; + +import com.google.api.core.ApiFunction; +import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutures; +import com.google.api.gax.batching.BatchingFlowController; +import com.google.api.gax.batching.BatchingThreshold; +import com.google.api.gax.batching.FlowController.FlowControlException; +import com.google.api.gax.batching.RequestBuilder; +import com.google.api.gax.batching.ThresholdBatchReceiver; +import com.google.api.gax.rpc.Batch; +import com.google.api.gax.rpc.BatchedFuture; +import com.google.api.gax.rpc.BatchedRequestIssuer; +import com.google.api.gax.rpc.BatchingDescriptor; +import com.google.api.gax.rpc.UnaryCallable; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.collect.Lists; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.concurrent.Future; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.locks.ReentrantLock; +import org.threeten.bp.Duration; + +/** + * Queues up elements until either a duration of time has passed or any threshold in a given set of + * thresholds is breached, and then delivers the elements in a batch to the consumer. + */ +public class Batcher { + + private class ReleaseResourcesFunction implements ApiFunction { + private final RequestT request; + + private ReleaseResourcesFunction(RequestT request) { + this.request = request; + } + + @Override + public Void apply(T input) { + flowController.release(request); + return null; + } + } + + private final Runnable pushCurrentBatchRunnable = + new Runnable() { + @Override + public void run() { + pushCurrentBatch(); + } + }; + + private final ArrayList> thresholds; + private final ScheduledExecutorService executor; + private final Duration maxDelay; + private final ThresholdBatchReceiver> receiver; + private final BatchingFlowController flowController; + + private final ReentrantLock lock = new ReentrantLock(); + + private RequestBuilder requestBuilder; + final BatchingDescriptor descriptor; + final UnaryCallable innerCallable; + List> requestIssuers = Lists.newArrayList(); + long requestSize = 0; + + private Future currentAlarmFuture; + + + private Batcher(Builder builder) { + this.thresholds = Lists.newArrayList(builder.thresholds); + this.executor = builder.executor; + this.maxDelay = builder.maxDelay; + this.receiver = builder.receiver; + this.flowController = builder.flowController; + this.requestBuilder = builder.descriptor.getRequestBuilder(); + this.descriptor = builder.descriptor; + this.innerCallable = builder.innerCallable; + + resetThresholds(); + } + + /** Builder for a Batcher. */ + public static class Builder { + private Collection> thresholds; + private ScheduledExecutorService executor; + private Duration maxDelay; + private ThresholdBatchReceiver> receiver; + private BatchingFlowController flowController; + + private BatchingDescriptor descriptor; + private UnaryCallable innerCallable; + + private Builder() {} + + /** Set the executor for the ThresholdBatcher. */ + public Builder setExecutor(ScheduledExecutorService executor) { + this.executor = executor; + return this; + } + + /** Set the max delay for a batch. This is counted from the first item added to a batch. */ + public Builder setMaxDelay(Duration maxDelay) { + this.maxDelay = maxDelay; + return this; + } + + /** Set the thresholds for the ThresholdBatcher. */ + public Builder setThresholds(Collection> thresholds) { + this.thresholds = Lists.newArrayList(thresholds); + return this; + } + + /** Set the threshold batch receiver for the ThresholdBatcher. */ + public Builder setReceiver(ThresholdBatchReceiver> receiver) { + this.receiver = receiver; + return this; + } + + /** Set the flow controller for the ThresholdBatcher. */ + public Builder setFlowController(BatchingFlowController flowController) { + this.flowController = flowController; + return this; + } + + public Builder setDescriptor(BatchingDescriptor descriptor) { + this.descriptor = descriptor; + return this; + } + + public Builder setInnerCallable(UnaryCallable innerCallable) { + this.innerCallable = innerCallable; + return this; + } + + /** Build the ThresholdBatcher. */ + public Batcher build() { + return new Batcher<>(this); + } + } + + /** Get a new builder for a Batcher. */ + public static Builder newBuilder() { + return new Builder<>(); + } + + /** + * Adds an element to the batcher. If the element causes the collection to go past any of the + * thresholds, the batch will be sent to the {@code ThresholdBatchReceiver}. + */ + public ApiFuture add(RequestT r) throws FlowControlException { + // We need to reserve resources from flowController outside the lock, so that they can be + // released by pushCurrentBatch(). + flowController.reserve(r); + lock.lock(); + try { + boolean anyThresholdReached = isAnyThresholdReached(r); + + requestBuilder.appendRequest(r); + BatchedFuture future = new BatchedFuture<>(); + requestIssuers.add(new BatchedRequestIssuer<>(future, descriptor.countElements(r))); + requestSize += descriptor.countBytes(r); + + if (currentAlarmFuture == null) { + // Schedule a job only when no thresholds have been exceeded, otherwise it will be + // immediately cancelled + if (!anyThresholdReached) { + currentAlarmFuture = + executor.schedule( + pushCurrentBatchRunnable, maxDelay.toMillis(), TimeUnit.MILLISECONDS); + } + } + + if (anyThresholdReached) { + pushCurrentBatch(); + } + + return future; + } finally { + lock.unlock(); + } + } + + /** * Package-private for use in testing. */ + @VisibleForTesting + boolean isEmpty() { + lock.lock(); + try { + return requestSize == 0; + } finally { + lock.unlock(); + } + } + + /** + * Push the current batch to the batch receiver. Returns an ApiFuture that completes once the + * batch has been processed by the batch receiver and the flow controller resources have been + * released. + * + *

Note that this future can complete for the current batch before previous batches have + * completed, so it cannot be depended upon for flushing. + */ + @VisibleForTesting + public ApiFuture pushCurrentBatch() { + final Batch batch = removeBatch(); + if (batch == null) { + return ApiFutures.immediateFuture(null); + } else { + return ApiFutures.transform( + receiver.processBatch(batch), new ReleaseResourcesFunction<>(batch.getRequest())); + } + } + + private Batch removeBatch() { + lock.lock(); + try { + RequestT request = requestBuilder.build(); + if (currentAlarmFuture != null) { + currentAlarmFuture.cancel(false); + currentAlarmFuture = null; + } + + requestBuilder = descriptor.getRequestBuilder(); + resetThresholds(); + + Batch batch = + new Batch<>(request, requestIssuers, innerCallable, requestSize); + requestSize = 0; + requestIssuers = Lists.newArrayList(); + resetThresholds(); + return batch; + } finally { + lock.unlock(); + } + } + + private boolean isAnyThresholdReached(RequestT e) { + for (BatchingThreshold threshold : thresholds) { + threshold.accumulate(e); + if (threshold.isThresholdReached()) { + return true; + } + } + return false; + } + + private void resetThresholds() { + for (int i = 0; i < thresholds.size(); i++) { + thresholds.set(i, thresholds.get(i).copyWithZeroedValue()); + } + } +} diff --git a/gax/src/main/java/com/google/api/gax/rpc/Batch.java b/gax/src/main/java/com/google/api/gax/rpc/Batch.java index 2a6188814..6a06f8a93 100644 --- a/gax/src/main/java/com/google/api/gax/rpc/Batch.java +++ b/gax/src/main/java/com/google/api/gax/rpc/Batch.java @@ -56,6 +56,24 @@ public class Batch { private UnaryCallable callable; private long byteCount; + public Batch(final RequestT request, List> requestIssuerList, UnaryCallable callable, long byteCount) { + requestBuilder = new RequestBuilder() { + @Override + public void appendRequest(RequestT request) { + throw new UnsupportedOperationException(); + } + + @Override + public RequestT build() { + return request; + } + }; + + this.requestIssuerList =requestIssuerList; + this.callable = callable; + this.byteCount = byteCount; + } + public Batch( BatchingDescriptor descriptor, RequestT request, diff --git a/gax/src/main/java/com/google/api/gax/rpc/BatchExecutor.java b/gax/src/main/java/com/google/api/gax/rpc/BatchExecutor.java index 0398b2e1e..2e6c84fc6 100644 --- a/gax/src/main/java/com/google/api/gax/rpc/BatchExecutor.java +++ b/gax/src/main/java/com/google/api/gax/rpc/BatchExecutor.java @@ -48,7 +48,7 @@ * *

Package-private for internal use. */ -class BatchExecutor +public class BatchExecutor implements ThresholdBatchReceiver> { private final BatchingDescriptor batchingDescriptor;