Examples with TrackedFlowController com.google.api.gax.batching.TrackedFlowController used on opensource projects

Example 1 with TrackedFlowController

use of com.google.api.gax.batching.TrackedFlowController in project gax-java by googleapis.

the class BatchingTest method batchingWithFlowControl.

@Test
public void batchingWithFlowControl() throws Exception {
    BatchingSettings batchingSettings = BatchingSettings.newBuilder().setDelayThreshold(Duration.ofSeconds(1)).setElementCountThreshold(4L).setRequestByteThreshold(null).setFlowControlSettings(FlowControlSettings.newBuilder().setLimitExceededBehavior(LimitExceededBehavior.Block).setMaxOutstandingElementCount(10L).setMaxOutstandingRequestBytes(10L).build()).build();
    TrackedFlowController trackedFlowController = new TrackedFlowController(batchingSettings.getFlowControlSettings());
    Truth.assertThat(trackedFlowController.getElementsReserved()).isEqualTo(0);
    Truth.assertThat(trackedFlowController.getElementsReleased()).isEqualTo(0);
    Truth.assertThat(trackedFlowController.getBytesReserved()).isEqualTo(0);
    Truth.assertThat(trackedFlowController.getBytesReleased()).isEqualTo(0);
    Truth.assertThat(trackedFlowController.getCallsToReserve()).isEqualTo(0);
    Truth.assertThat(trackedFlowController.getCallsToRelease()).isEqualTo(0);
    LabeledIntList requestA = new LabeledIntList("one", 1, 2);
    LabeledIntList requestB = new LabeledIntList("one", 3, 4);
    BatchingCallSettings<LabeledIntList, List<Integer>> batchingCallSettings = BatchingCallSettings.newBuilder(SQUARER_BATCHING_DESC).setBatchingSettings(batchingSettings).setFlowController(trackedFlowController).build();
    Callables.BatchingCreateResult<LabeledIntList, List<Integer>> batchingCreateResult = Callables.batchingImpl(callLabeledIntSquarer, batchingCallSettings, clientContext);
    ApiFuture<List<Integer>> f1 = batchingCreateResult.getUnaryCallable().futureCall(requestA, FakeCallContext.createDefault());
    ApiFuture<List<Integer>> f2 = batchingCreateResult.getUnaryCallable().futureCall(requestB, FakeCallContext.createDefault());
    Truth.assertThat(f1.get()).isEqualTo(Arrays.asList(1, 4));
    Truth.assertThat(f2.get()).isEqualTo(Arrays.asList(9, 16));
    batchingCreateResult.getBatcherFactory().getPushingBatcher(SQUARER_BATCHING_DESC.getBatchPartitionKey(requestA)).pushCurrentBatch().get();
    // Check that the number of bytes is correct even when requests are merged, and the merged
    // request consumes fewer bytes.
    Truth.assertThat(trackedFlowController.getElementsReserved()).isEqualTo(4);
    Truth.assertThat(trackedFlowController.getElementsReleased()).isEqualTo(4);
    Truth.assertThat(trackedFlowController.getBytesReserved()).isEqualTo(8);
    Truth.assertThat(trackedFlowController.getBytesReleased()).isEqualTo(8);
    Truth.assertThat(trackedFlowController.getCallsToReserve()).isEqualTo(2);
    Truth.assertThat(trackedFlowController.getCallsToRelease()).isEqualTo(1);
}