Examples with SessionConsumerImpl com.google.cloud.spanner.SessionPool.SessionConsumerImpl used on opensource projects

Example 11 with SessionConsumerImpl

use of com.google.cloud.spanner.SessionPool.SessionConsumerImpl in project java-spanner by googleapis.

the class SessionPoolTest method testSessionNotFoundWrite.

@Test
public void testSessionNotFoundWrite() {
    SpannerException sessionNotFound = SpannerExceptionFactoryTest.newSessionNotFoundException(sessionName);
    List<Mutation> mutations = Collections.singletonList(Mutation.newInsertBuilder("FOO").build());
    final SessionImpl closedSession = mockSession();
    when(closedSession.writeWithOptions(mutations)).thenThrow(sessionNotFound);
    final SessionImpl openSession = mockSession();
    com.google.cloud.spanner.CommitResponse response = mock(com.google.cloud.spanner.CommitResponse.class);
    when(response.getCommitTimestamp()).thenReturn(Timestamp.now());
    when(openSession.writeWithOptions(mutations)).thenReturn(response);
    doAnswer(invocation -> {
        executor.submit(() -> {
            SessionConsumerImpl consumer = invocation.getArgument(2, SessionConsumerImpl.class);
            consumer.onSessionReady(closedSession);
        });
        return null;
    }).doAnswer(invocation -> {
        executor.submit(() -> {
            SessionConsumerImpl consumer = invocation.getArgument(2, SessionConsumerImpl.class);
            consumer.onSessionReady(openSession);
        });
        return null;
    }).when(sessionClient).asyncBatchCreateSessions(Mockito.eq(1), Mockito.anyBoolean(), any(SessionConsumer.class));
    FakeClock clock = new FakeClock();
    clock.currentTimeMillis = System.currentTimeMillis();
    pool = createPool(clock);
    DatabaseClientImpl impl = new DatabaseClientImpl(pool);
    assertThat(impl.write(mutations)).isNotNull();
}

Example 12 with SessionConsumerImpl

use of com.google.cloud.spanner.SessionPool.SessionConsumerImpl in project java-spanner by googleapis.

the class SessionPoolTest method poolClosureFailsNewRequests.

@Test
public void poolClosureFailsNewRequests() {
    final SessionImpl session = mockSession();
    doAnswer(invocation -> {
        executor.submit(() -> {
            SessionConsumerImpl consumer = invocation.getArgument(2, SessionConsumerImpl.class);
            consumer.onSessionReady(session);
        });
        return null;
    }).when(sessionClient).asyncBatchCreateSessions(Mockito.eq(1), Mockito.anyBoolean(), any(SessionConsumer.class));
    pool = createPool();
    PooledSessionFuture leakedSession = pool.getSession();
    leakedSession.get();
    // Suppress expected leakedSession warning.
    leakedSession.clearLeakedException();
    pool.closeAsync(new SpannerImpl.ClosedException());
    IllegalStateException e = assertThrows(IllegalStateException.class, () -> pool.getSession());
    assertNotNull(e.getMessage());
}

Example 13 with SessionConsumerImpl

use of com.google.cloud.spanner.SessionPool.SessionConsumerImpl in project java-spanner by googleapis.

the class SessionPoolTest method keepAlive.

@Test
public void keepAlive() throws Exception {
    options = SessionPoolOptions.newBuilder().setMinSessions(2).setMaxSessions(3).build();
    final SessionImpl session = mockSession();
    mockKeepAlive(session);
    // This is cheating as we are returning the same session each but it makes the verification
    // easier.
    doAnswer(invocation -> {
        executor.submit(() -> {
            int sessionCount = invocation.getArgument(0, Integer.class);
            SessionConsumerImpl consumer = invocation.getArgument(2, SessionConsumerImpl.class);
            for (int i = 0; i < sessionCount; i++) {
                consumer.onSessionReady(session);
            }
        });
        return null;
    }).when(sessionClient).asyncBatchCreateSessions(anyInt(), Mockito.anyBoolean(), any(SessionConsumer.class));
    FakeClock clock = new FakeClock();
    clock.currentTimeMillis = System.currentTimeMillis();
    pool = createPool(clock);
    PooledSessionFuture session1 = pool.getSession();
    PooledSessionFuture session2 = pool.getSession();
    session1.get();
    session2.get();
    session1.close();
    session2.close();
    runMaintenanceLoop(clock, pool, pool.poolMaintainer.numKeepAliveCycles);
    verify(session, never()).singleUse(any(TimestampBound.class));
    runMaintenanceLoop(clock, pool, pool.poolMaintainer.numKeepAliveCycles);
    verify(session, times(2)).singleUse(any(TimestampBound.class));
    clock.currentTimeMillis += clock.currentTimeMillis + (options.getKeepAliveIntervalMinutes() + 5) * 60 * 1000;
    session1 = pool.getSession();
    session1.writeAtLeastOnceWithOptions(new ArrayList<>());
    session1.close();
    runMaintenanceLoop(clock, pool, pool.poolMaintainer.numKeepAliveCycles);
    // The session pool only keeps MinSessions + MaxIdleSessions alive.
    verify(session, times(options.getMinSessions() + options.getMaxIdleSessions())).singleUse(any(TimestampBound.class));
    pool.closeAsync(new SpannerImpl.ClosedException()).get(5L, TimeUnit.SECONDS);
}

Example 14 with SessionConsumerImpl

use of com.google.cloud.spanner.SessionPool.SessionConsumerImpl in project java-spanner by googleapis.

the class SessionPoolTest method poolClosesEvenIfCreationFails.

@Test
public void poolClosesEvenIfCreationFails() throws Exception {
    final CountDownLatch insideCreation = new CountDownLatch(1);
    final CountDownLatch releaseCreation = new CountDownLatch(1);
    doAnswer(invocation -> {
        executor.submit(() -> {
            insideCreation.countDown();
            releaseCreation.await();
            SessionConsumerImpl consumer = invocation.getArgument(2, SessionConsumerImpl.class);
            consumer.onSessionCreateFailure(SpannerExceptionFactory.newSpannerException(new RuntimeException()), 1);
            return null;
        });
        return null;
    }).when(sessionClient).asyncBatchCreateSessions(Mockito.eq(1), Mockito.anyBoolean(), any(SessionConsumer.class));
    pool = createPool();
    AtomicBoolean failed = new AtomicBoolean(false);
    CountDownLatch latch = new CountDownLatch(1);
    getSessionAsync(latch, failed);
    insideCreation.await();
    ListenableFuture<Void> f = pool.closeAsync(new SpannerImpl.ClosedException());
    releaseCreation.countDown();
    f.get();
    assertThat(f.isDone()).isTrue();
}

Example 15 with SessionConsumerImpl

use of com.google.cloud.spanner.SessionPool.SessionConsumerImpl in project java-spanner by googleapis.

the class SessionPoolMaintainerTest method setupMockSessionCreation.

private void setupMockSessionCreation() {
    doAnswer(invocation -> {
        executor.submit(() -> {
            int sessionCount = invocation.getArgument(0, Integer.class);
            SessionConsumerImpl consumer = invocation.getArgument(2, SessionConsumerImpl.class);
            for (int i = 0; i < sessionCount; i++) {
                consumer.onSessionReady(setupMockSession(mockSession()));
            }
        });
        return null;
    }).when(sessionClient).asyncBatchCreateSessions(Mockito.anyInt(), Mockito.anyBoolean(), any(SessionConsumer.class));
}