Examples with RpcWriteAttemptImpl org.apache.beam.sdk.io.gcp.firestore.RpcQosImpl.RpcWriteAttemptImpl used on opensource projects

Example 1 with RpcWriteAttemptImpl

use of org.apache.beam.sdk.io.gcp.firestore.RpcQosImpl.RpcWriteAttemptImpl in project beam by apache.

the class RpcQosTest method doTest_isCodeRetryable.

private void doTest_isCodeRetryable(Code code, boolean shouldBeRetryable) {
    RpcQosOptions options = RpcQosOptions.defaultOptions();
    RpcQosImpl qos = new RpcQosImpl(options, random, sleeper, counterFactory, distributionFactory);
    RpcWriteAttemptImpl attempt = qos.newWriteAttempt(RPC_ATTEMPT_CONTEXT);
    assertEquals(shouldBeRetryable, attempt.isCodeRetryable(code));
}

Example 2 with RpcWriteAttemptImpl

use of org.apache.beam.sdk.io.gcp.firestore.RpcQosImpl.RpcWriteAttemptImpl in project beam by apache.

the class RpcQosSimulationTest method unsafeToProceed.

private void unsafeToProceed(RpcQosImpl qos, Instant t) throws InterruptedException {
    RpcWriteAttemptImpl attempt = qos.newWriteAttempt(rpcAttemptContext);
    assertFalse(msg("verify budget depleted", t, "awaitSafeToProceed was true, expected false"), attempt.awaitSafeToProceed(t));
}

Example 3 with RpcWriteAttemptImpl

use of org.apache.beam.sdk.io.gcp.firestore.RpcQosImpl.RpcWriteAttemptImpl in project beam by apache.

the class RpcQosSimulationTest method safeToProceedAndWithBudgetAndWrite.

private void safeToProceedAndWithBudgetAndWrite(RpcQosImpl qos, Instant t, int expectedBatchMaxCount, int writeCount, String description) throws InterruptedException {
    RpcWriteAttemptImpl attempt = qos.newWriteAttempt(rpcAttemptContext);
    assertTrue(msg(description, t, "awaitSafeToProceed was false, expected true"), attempt.awaitSafeToProceed(t));
    FlushBufferImpl<Object, Element<Object>> buffer = attempt.newFlushBuffer(t);
    assertEquals(msg(description, t, "unexpected batchMaxCount"), expectedBatchMaxCount, buffer.nextBatchMaxCount);
    attempt.recordRequestStart(t, writeCount);
    attempt.recordWriteCounts(t, writeCount, 0);
}

Example 4 with RpcWriteAttemptImpl

use of org.apache.beam.sdk.io.gcp.firestore.RpcQosImpl.RpcWriteAttemptImpl in project beam by apache.

the class RpcQosTest method doTest_initialBatchSizeRelativeToWorkerCount.

private void doTest_initialBatchSizeRelativeToWorkerCount(int hintWorkerCount, int expectedBatchMaxCount) {
    RpcQosOptions options = RpcQosOptions.newBuilder().withHintMaxNumWorkers(hintWorkerCount).withBatchInitialCount(500).build();
    RpcQosImpl qos = new RpcQosImpl(options, random, sleeper, counterFactory, distributionFactory);
    RpcWriteAttemptImpl attempt = qos.newWriteAttempt(RPC_ATTEMPT_CONTEXT);
    FlushBufferImpl<Object, Element<Object>> buffer = attempt.newFlushBuffer(Instant.EPOCH);
    assertEquals(expectedBatchMaxCount, buffer.nextBatchMaxCount);
}