Examples with RpcWriteAttempt org.apache.beam.sdk.io.gcp.firestore.RpcQos.RpcWriteAttempt used on opensource projects

Example 1 with RpcWriteAttempt

use of org.apache.beam.sdk.io.gcp.firestore.RpcQos.RpcWriteAttempt in project beam by apache.

the class RpcQosTest method doTest_writes_shouldFlush_numBytes.

private void doTest_writes_shouldFlush_numBytes(boolean expectFlush, long numBytes) {
    RpcQosOptions rpcQosOptions = options.toBuilder().withBatchMaxBytes(3000).unsafeBuild();
    RpcQos qos = new RpcQosImpl(rpcQosOptions, random, sleeper, counterFactory, distributionFactory);
    RpcWriteAttempt attempt = qos.newWriteAttempt(RPC_ATTEMPT_CONTEXT);
    FlushBuffer<Element<Write>> accumulator = attempt.newFlushBuffer(monotonicClock.instant());
    assertTrue(accumulator.offer(new FixedSerializationSize<>(newWrite(), numBytes)));
    assertEquals(expectFlush, accumulator.isFull());
    assertEquals(numBytes, accumulator.getBufferedElementsBytes());
    assertEquals(newArrayList(newWrite()), StreamSupport.stream(accumulator.spliterator(), false).map(Element::getValue).collect(Collectors.toList()));
}

Example 2 with RpcWriteAttempt

use of org.apache.beam.sdk.io.gcp.firestore.RpcQos.RpcWriteAttempt in project beam by apache.

the class RpcQosTest method doTest_writes_shouldFlush_numWritesHigherThanBatchCount_existingTimeBucket.

private void doTest_writes_shouldFlush_numWritesHigherThanBatchCount_existingTimeBucket(boolean expectFlush, int batchCount) {
    doTest_shouldFlush_numWritesHigherThanBatchCount(expectFlush, batchCount, (qos) -> {
        RpcWriteAttempt attempt = qos.newWriteAttempt(RPC_ATTEMPT_CONTEXT);
        attempt.recordRequestStart(monotonicClock.instant(), 1);
        attempt.recordWriteCounts(monotonicClock.instant(), 1, 0);
    });
}

Example 3 with RpcWriteAttempt

use of org.apache.beam.sdk.io.gcp.firestore.RpcQos.RpcWriteAttempt in project beam by apache.

the class RpcQosTest method offerOfElementWhichWouldCrossMaxBytesReturnFalse.

@Test
public void offerOfElementWhichWouldCrossMaxBytesReturnFalse() {
    RpcQosOptions rpcQosOptions = options.toBuilder().withBatchMaxBytes(5000).unsafeBuild();
    RpcQos qos = new RpcQosImpl(rpcQosOptions, random, sleeper, counterFactory, distributionFactory);
    RpcWriteAttempt attempt = qos.newWriteAttempt(RPC_ATTEMPT_CONTEXT);
    FlushBuffer<Element<Write>> accumulator = attempt.newFlushBuffer(monotonicClock.instant());
    assertFalse(accumulator.offer(new FixedSerializationSize<>(newWrite(), 5001)));
    assertFalse(accumulator.isFull());
    assertEquals(0, accumulator.getBufferedElementsBytes());
    assertEquals(0, accumulator.getBufferedElementsCount());
}

Example 4 with RpcWriteAttempt

use of org.apache.beam.sdk.io.gcp.firestore.RpcQos.RpcWriteAttempt in project beam by apache.

the class RpcQosTest method writes_blockWhenNotSafeToProceed.

@Test
public void writes_blockWhenNotSafeToProceed() throws InterruptedException {
    RpcQos qos = new RpcQosImpl(options, random, sleeper, counterFactory, distributionFactory);
    // Based on the defined options, 3 failures is the upper bound before the next attempt
    // will have to wait before proceeding
    int numFailures = 3;
    for (int i = 0; i < numFailures; i++) {
        RpcWriteAttempt writeAttempt = qos.newWriteAttempt(RPC_ATTEMPT_CONTEXT);
        Instant start = monotonicClock.instant();
        assertTrue(writeAttempt.awaitSafeToProceed(start));
        writeAttempt.recordRequestStart(start, 1);
        Instant end = monotonicClock.instant();
        writeAttempt.recordWriteCounts(end, 0, 1);
        writeAttempt.recordRequestFailed(end);
    }
    RpcWriteAttempt writeAttempt2 = qos.newWriteAttempt(RPC_ATTEMPT_CONTEXT);
    assertFalse(writeAttempt2.awaitSafeToProceed(monotonicClock.instant()));
    long sleepMillis = options.getInitialBackoff().getMillis();
    verify(sleeper, times(0)).sleep(sleepMillis);
    verify(counterThrottlingMs, times(0)).inc(sleepMillis);
    verify(counterRpcFailures, times(numFailures)).inc();
    verify(counterRpcSuccesses, times(0)).inc();
    verify(counterRpcStreamValueReceived, times(0)).inc();
}

Example 5 with RpcWriteAttempt

use of org.apache.beam.sdk.io.gcp.firestore.RpcQos.RpcWriteAttempt in project beam by apache.

the class RpcQosTest method doTest_shouldFlush_numWritesHigherThanBatchCount.

private void doTest_shouldFlush_numWritesHigherThanBatchCount(boolean expectFlush, int batchCount, Consumer<RpcQos> preAttempt) {
    RpcQosOptions rpcQosOptions = options.toBuilder().withBatchInitialCount(10).withBatchMaxCount(10).unsafeBuild();
    RpcQos qos = new RpcQosImpl(rpcQosOptions, random, sleeper, counterFactory, distributionFactory);
    preAttempt.accept(qos);
    RpcWriteAttempt attempt = qos.newWriteAttempt(RPC_ATTEMPT_CONTEXT);
    FlushBuffer<Element<Write>> accumulator = attempt.newFlushBuffer(monotonicClock.instant());
    for (int i = 0; i < batchCount; i++) {
        accumulator.offer(new WriteElement(i, newWrite(), window));
    }
    if (expectFlush) {
        assertTrue(accumulator.isFull());
        assertEquals(10, accumulator.getBufferedElementsCount());
    } else {
        assertFalse(accumulator.isFull());
        assertEquals(batchCount, accumulator.getBufferedElementsCount());
    }
}