Examples with ApiCallContext com.google.api.gax.rpc.ApiCallContext used on opensource projects

Example 1 with ApiCallContext

use of com.google.api.gax.rpc.ApiCallContext in project workbench by all-of-us.

the class ForwardingCloudTasksStub method createTaskCallable.

@Override
public UnaryCallable<CreateTaskRequest, Task> createTaskCallable() {
    return new UnaryCallable<CreateTaskRequest, Task>() {

        @Override
        public ApiFuture<Task> futureCall(CreateTaskRequest request, ApiCallContext context) {
            final QueueName queueName = QueueName.parse(request.getParent());
            final AppEngineHttpRequest gaeReq = request.getTask().getAppEngineHttpRequest();
            final Request apiReq = new Request.Builder().url(baseUrl + gaeReq.getRelativeUri()).headers(Headers.of(gaeReq.getHeadersMap())).addHeader("X-AppEngine-QueueName", queueName.getQueue()).post(RequestBody.create(MediaType.parse("application/json; charset=utf-8"), gaeReq.getBody().toStringUtf8())).build();
            log.info(String.format("asynchronously forwarding task request for queue '%s', to handler '%s'", queueName.getQueue(), apiReq.url()));
            OkHttpClient client = new OkHttpClient();
            client.setReadTimeout(10, TimeUnit.MINUTES);
            client.newCall(apiReq).enqueue(new Callback() {

                @Override
                public void onFailure(Request request, IOException e) {
                    log.log(Level.SEVERE, "task execution failed", e);
                }

                @Override
                public void onResponse(Response response) {
                }
            });
            return ApiFutures.immediateFuture(request.getTask());
        }
    };
}

Example 2 with ApiCallContext

use of com.google.api.gax.rpc.ApiCallContext in project java-docs-samples by GoogleCloudPlatform.

the class VulnerabilityFunctionTest method testAccept.

@Test
public void testAccept() throws IOException {
    AtomicInteger occurrenceCallCount = new AtomicInteger();
    Mockito.when(grafeasStub.getOccurrenceCallable()).thenReturn(new UnaryCallable<>() {

        @Override
        public ApiFuture<Occurrence> futureCall(GetOccurrenceRequest request, ApiCallContext context) {
            occurrenceCallCount.incrementAndGet();
            return ApiFutures.immediateFuture(Occurrence.newBuilder().setKind(NoteKind.VULNERABILITY).setResourceUri("gcr.io/test-project/some-image").setVulnerability(VulnerabilityOccurrence.newBuilder().setSeverity(Severity.CRITICAL).setShortDescription("CVE-CRITICAL").build()).build());
        }
    });
    VulnerabilityFunction function = new VulnerabilityFunction(containerAnalysisClient);
    String notificationPayload = "{\"name\":\"projects/test-project/occurrences/some-uuid\",\"kind\":\"VULNERABILITY\"," + "\"notificationTime\":\"2020-09-04T00:38:25.575543Z\"}";
    PubSubMessage msg = new PubSubMessage();
    msg.setData(Base64.getEncoder().encodeToString(notificationPayload.getBytes(StandardCharsets.UTF_8)));
    function.accept(msg, null);
    Assert.assertEquals(1, occurrenceCallCount.get());
}

Example 3 with ApiCallContext

use of com.google.api.gax.rpc.ApiCallContext 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);
}

Example 4 with ApiCallContext

use of com.google.api.gax.rpc.ApiCallContext in project java-bigtable by googleapis.

the class DynamicFlowControlCallableTest method testDecreasingThresholdsCantGoOverLimit.

@Test
public void testDecreasingThresholdsCantGoOverLimit() throws Exception {
    // set adjusting intervals to 0 so the thresholds can keep getting updated
    callableToTest = new DynamicFlowControlCallable(innerCallable, flowController, stats, TARGET_LATENCY_MS, 0);
    Map<String, List<String>> extraHeaders = new HashMap<>();
    extraHeaders.put(LATENCY_HEADER, Arrays.asList(String.valueOf(TARGET_LATENCY_MS * 4)));
    ApiCallContext newContext = context.withExtraHeaders(extraHeaders);
    List<Future> futures = new ArrayList<>();
    for (int i = 0; i < 3; i++) {
        ApiFuture future = callableToTest.futureCall(request, newContext);
        futures.add(future);
    }
    for (Future f : futures) {
        f.get();
    }
    long expectedStep = Math.round(MAX_ELEMENT * DynamicFlowControlCallable.VERY_HIGH_LATENCY_DECREASE_CONCURRENCY_RATE) * 3;
    assertThat(INITIAL_ELEMENT - expectedStep).isLessThan(MIN_ELEMENT);
    assertThat(flowController.getCurrentElementCountLimit()).isEqualTo(MIN_ELEMENT);
}

Example 5 with ApiCallContext

use of com.google.api.gax.rpc.ApiCallContext in project java-bigtable by googleapis.

the class DynamicFlowControlCallableTest method testTriggeringAdjustingThreshold.

@Test
public void testTriggeringAdjustingThreshold() throws Exception {
    Map<String, List<String>> extraHeaders = new HashMap<>();
    extraHeaders.put(LATENCY_HEADER, Arrays.asList(String.valueOf(TARGET_LATENCY_MS * 4)));
    long currentTimeMs = System.currentTimeMillis();
    ApiCallContext newContext = context.withExtraHeaders(extraHeaders);
    ApiFuture future = callableToTest.futureCall(request, newContext);
    future.get();
    assertThat(stats.getMeanLatency()).isAtLeast(TARGET_LATENCY_MS * DynamicFlowControlCallable.VERY_HIGH_LATENCY_MULTIPLIER);
    assertThat(stats.getLastAdjustedTimestampMs()).isGreaterThan(currentTimeMs);
    long expectedStep = Math.round(MAX_ELEMENT * DynamicFlowControlCallable.VERY_HIGH_LATENCY_DECREASE_CONCURRENCY_RATE);
    assertThat(flowController.getCurrentElementCountLimit()).isEqualTo(INITIAL_ELEMENT - expectedStep);
}