Search in sources :

Example 1 with Batcher

use of com.google.api.gax.batching.Batcher 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 Batcher

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

the class MetricsTracerTest method testBatchReadRowsThrottledTime.

@Test
public void testBatchReadRowsThrottledTime() throws Exception {
    doAnswer(new Answer() {

        @Override
        public Object answer(InvocationOnMock invocation) {
            @SuppressWarnings("unchecked") StreamObserver<ReadRowsResponse> observer = (StreamObserver<ReadRowsResponse>) invocation.getArguments()[1];
            observer.onNext(DEFAULT_READ_ROWS_RESPONSES);
            observer.onCompleted();
            return null;
        }
    }).when(mockService).readRows(any(ReadRowsRequest.class), any());
    try (Batcher batcher = stub.newBulkReadRowsBatcher(Query.create(TABLE_ID), GrpcCallContext.createDefault())) {
        batcher.add(ByteString.copyFromUtf8("row1"));
        batcher.sendOutstanding();
        // Give OpenCensus a chance to update the views asynchronously.
        Thread.sleep(100);
        long throttledTimeMetric = StatsTestUtils.getAggregationValueAsLong(localStats, RpcViewConstants.BIGTABLE_BATCH_THROTTLED_TIME_VIEW, ImmutableMap.of(RpcMeasureConstants.BIGTABLE_OP, TagValue.create("Bigtable.ReadRows")), PROJECT_ID, INSTANCE_ID, APP_PROFILE_ID);
        assertThat(throttledTimeMetric).isEqualTo(0);
    }
}
Also used : StreamObserver(io.grpc.stub.StreamObserver) Mockito.doAnswer(org.mockito.Mockito.doAnswer) Answer(org.mockito.stubbing.Answer) ReadRowsResponse(com.google.bigtable.v2.ReadRowsResponse) Batcher(com.google.api.gax.batching.Batcher) InvocationOnMock(org.mockito.invocation.InvocationOnMock) ReadRowsRequest(com.google.bigtable.v2.ReadRowsRequest) Test(org.junit.Test)

Aggregations

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