Search in sources :

Example 1 with FlowController

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

the class MetricsTracerTest method testBatchMutateRowsThrottledTime.

@Test
public void testBatchMutateRowsThrottledTime() throws Exception {
    FlowController flowController = Mockito.mock(FlowController.class);
    BatchingDescriptor batchingDescriptor = Mockito.mock(MutateRowsBatchingDescriptor.class);
    // Mock throttling
    final long throttled = 50;
    doAnswer(new Answer() {

        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            Thread.sleep(throttled);
            return null;
        }
    }).when(flowController).reserve(any(Long.class), any(Long.class));
    when(flowController.getMaxElementCountLimit()).thenReturn(null);
    when(flowController.getMaxRequestBytesLimit()).thenReturn(null);
    when(batchingDescriptor.countBytes(any())).thenReturn(1l);
    when(batchingDescriptor.newRequestBuilder(any())).thenCallRealMethod();
    doAnswer(new Answer() {

        @Override
        public Object answer(InvocationOnMock invocation) {
            @SuppressWarnings("unchecked") StreamObserver<MutateRowsResponse> observer = (StreamObserver<MutateRowsResponse>) invocation.getArguments()[1];
            observer.onNext(MutateRowsResponse.getDefaultInstance());
            observer.onCompleted();
            return null;
        }
    }).when(mockService).mutateRows(any(MutateRowsRequest.class), any());
    ApiCallContext defaultContext = GrpcCallContext.createDefault();
    Batcher batcher = new BatcherImpl(batchingDescriptor, stub.bulkMutateRowsCallable().withDefaultCallContext(defaultContext), BulkMutation.create(TABLE_ID), settings.getStubSettings().bulkMutateRowsSettings().getBatchingSettings(), Executors.newSingleThreadScheduledExecutor(), flowController, defaultContext);
    batcher.add(RowMutationEntry.create("key"));
    batcher.sendOutstanding();
    Thread.sleep(100);
    long throttledTimeMetric = StatsTestUtils.getAggregationValueAsLong(localStats, RpcViewConstants.BIGTABLE_BATCH_THROTTLED_TIME_VIEW, ImmutableMap.of(RpcMeasureConstants.BIGTABLE_OP, TagValue.create("Bigtable.MutateRows")), PROJECT_ID, INSTANCE_ID, APP_PROFILE_ID);
    assertThat(throttledTimeMetric).isAtLeast(throttled);
}
Also used : StreamObserver(io.grpc.stub.StreamObserver) MutateRowsBatchingDescriptor(com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsBatchingDescriptor) BatchingDescriptor(com.google.api.gax.batching.BatchingDescriptor) ApiCallContext(com.google.api.gax.rpc.ApiCallContext) BatcherImpl(com.google.api.gax.batching.BatcherImpl) Mockito.doAnswer(org.mockito.Mockito.doAnswer) Answer(org.mockito.stubbing.Answer) MutateRowsResponse(com.google.bigtable.v2.MutateRowsResponse) MutateRowsRequest(com.google.bigtable.v2.MutateRowsRequest) Batcher(com.google.api.gax.batching.Batcher) InvocationOnMock(org.mockito.invocation.InvocationOnMock) FlowController(com.google.api.gax.batching.FlowController) Test(org.junit.Test)

Example 2 with FlowController

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

the class DynamicFlowControlCallableTest method setup.

@Before
public void setup() {
    flowController = new FlowController(DynamicFlowControlSettings.newBuilder().setInitialOutstandingElementCount(INITIAL_ELEMENT).setMaxOutstandingElementCount(MAX_ELEMENT).setMinOutstandingElementCount(MIN_ELEMENT).setInitialOutstandingRequestBytes(15L).setMaxOutstandingRequestBytes(15L).setMinOutstandingRequestBytes(15L).setLimitExceededBehavior(LimitExceededBehavior.Block).build());
    flowControlEvents = flowController.getFlowControlEventStats();
    stats = new DynamicFlowControlStats();
    context = GrpcCallContext.createDefault();
    innerCallable = new MockInnerCallable();
    request = MutateRowsRequest.newBuilder().addEntries(MutateRowsRequest.Entry.getDefaultInstance()).build();
    callableToTest = new DynamicFlowControlCallable(innerCallable, flowController, stats, TARGET_LATENCY_MS, ADJUSTING_INTERVAL_MS);
}
Also used : FlowController(com.google.api.gax.batching.FlowController) Before(org.junit.Before)

Aggregations

FlowController (com.google.api.gax.batching.FlowController)2 Batcher (com.google.api.gax.batching.Batcher)1 BatcherImpl (com.google.api.gax.batching.BatcherImpl)1 BatchingDescriptor (com.google.api.gax.batching.BatchingDescriptor)1 ApiCallContext (com.google.api.gax.rpc.ApiCallContext)1 MutateRowsRequest (com.google.bigtable.v2.MutateRowsRequest)1 MutateRowsResponse (com.google.bigtable.v2.MutateRowsResponse)1 MutateRowsBatchingDescriptor (com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsBatchingDescriptor)1 StreamObserver (io.grpc.stub.StreamObserver)1 Before (org.junit.Before)1 Test (org.junit.Test)1 Mockito.doAnswer (org.mockito.Mockito.doAnswer)1 InvocationOnMock (org.mockito.invocation.InvocationOnMock)1 Answer (org.mockito.stubbing.Answer)1