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

Example 1 with SessionConsumerImpl

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

the class SessionPoolStressTest method setupSpanner.

private void setupSpanner(DatabaseId db) {
    mockSpanner = mock(SpannerImpl.class);
    spannerOptions = mock(SpannerOptions.class);
    when(spannerOptions.getNumChannels()).thenReturn(4);
    SessionClient sessionClient = mock(SessionClient.class);
    when(mockSpanner.getSessionClient(db)).thenReturn(sessionClient);
    when(mockSpanner.getOptions()).thenReturn(spannerOptions);
    doAnswer(invocation -> {
        createExecutor.submit(() -> {
            int sessionCount = invocation.getArgument(0, Integer.class);
            for (int s = 0; s < sessionCount; s++) {
                SessionImpl session;
                synchronized (lock) {
                    session = mockSession();
                    setupSession(session);
                    sessions.put(session.getName(), false);
                    if (sessions.size() > maxAliveSessions) {
                        maxAliveSessions = sessions.size();
                    }
                }
                SessionConsumerImpl consumer = invocation.getArgument(2, SessionConsumerImpl.class);
                consumer.onSessionReady(session);
            }
        });
        return null;
    }).when(sessionClient).asyncBatchCreateSessions(Mockito.anyInt(), Mockito.anyBoolean(), Mockito.any(SessionConsumer.class));
}

Example 2 with SessionConsumerImpl

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

the class SessionPoolTest method testSessionNotFoundReadOnlyTransaction.

@Test
public void testSessionNotFoundReadOnlyTransaction() {
    Statement statement = Statement.of("SELECT 1");
    final SessionImpl closedSession = mockSession();
    when(closedSession.readOnlyTransaction()).thenThrow(SpannerExceptionFactoryTest.newSessionNotFoundException(sessionName));
    final SessionImpl openSession = mockSession();
    ReadOnlyTransaction openTransaction = mock(ReadOnlyTransaction.class);
    ResultSet openResultSet = mock(ResultSet.class);
    when(openResultSet.next()).thenReturn(true, false);
    when(openTransaction.executeQuery(statement)).thenReturn(openResultSet);
    when(openSession.readOnlyTransaction()).thenReturn(openTransaction);
    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);
    ReadOnlyTransaction transaction = pool.getSession().readOnlyTransaction();
    ResultSet resultSet = transaction.executeQuery(statement);
    assertThat(resultSet.next()).isTrue();
}

Example 3 with SessionConsumerImpl

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

the class SessionPoolTest method poolClosureClosesLeakedSessions.

@Test
public void poolClosureClosesLeakedSessions() throws Exception {
    SessionImpl mockSession1 = mockSession();
    SessionImpl mockSession2 = mockSession();
    final LinkedList<SessionImpl> sessions = new LinkedList<>(Arrays.asList(mockSession1, mockSession2));
    doAnswer(invocation -> {
        executor.submit(() -> {
            SessionConsumerImpl consumer = invocation.getArgument(2, SessionConsumerImpl.class);
            consumer.onSessionReady(sessions.pop());
        });
        return null;
    }).when(sessionClient).asyncBatchCreateSessions(Mockito.eq(1), Mockito.anyBoolean(), any(SessionConsumer.class));
    pool = createPool();
    Session session1 = pool.getSession();
    // Leaked sessions
    PooledSessionFuture leakedSession = pool.getSession();
    // Clear the leaked exception to suppress logging of expected exceptions.
    leakedSession.clearLeakedException();
    session1.close();
    pool.closeAsync(new SpannerImpl.ClosedException()).get(5L, TimeUnit.SECONDS);
    verify(mockSession1).asyncClose();
    verify(mockSession2).asyncClose();
}

Example 4 with SessionConsumerImpl

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

the class SessionPoolTest method poolClosureFailsPendingWriteWaiters.

@Test
public void poolClosureFailsPendingWriteWaiters() throws Exception {
    final CountDownLatch insideCreation = new CountDownLatch(1);
    final CountDownLatch releaseCreation = new CountDownLatch(1);
    final SessionImpl session1 = mockSession();
    final SessionImpl session2 = mockSession();
    doAnswer(invocation -> {
        executor.submit(() -> {
            SessionConsumerImpl consumer = invocation.getArgument(2, SessionConsumerImpl.class);
            consumer.onSessionReady(session1);
        });
        return null;
    }).doAnswer(invocation -> {
        executor.submit(() -> {
            insideCreation.countDown();
            releaseCreation.await();
            SessionConsumerImpl consumer = invocation.getArgument(2, SessionConsumerImpl.class);
            consumer.onSessionReady(session2);
            return null;
        });
        return null;
    }).when(sessionClient).asyncBatchCreateSessions(Mockito.eq(1), Mockito.anyBoolean(), any(SessionConsumer.class));
    pool = createPool();
    PooledSessionFuture leakedSession = pool.getSession();
    // Suppress expected leakedSession warning.
    leakedSession.clearLeakedException();
    AtomicBoolean failed = new AtomicBoolean(false);
    CountDownLatch latch = new CountDownLatch(1);
    getSessionAsync(latch, failed);
    insideCreation.await();
    pool.closeAsync(new SpannerImpl.ClosedException());
    releaseCreation.countDown();
    latch.await();
    assertThat(failed.get()).isTrue();
}

Example 5 with SessionConsumerImpl

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

the class SessionPoolTest method testSessionNotFoundSingleUse.

@Test
public void testSessionNotFoundSingleUse() {
    Statement statement = Statement.of("SELECT 1");
    final SessionImpl closedSession = mockSession();
    ReadContext closedContext = mock(ReadContext.class);
    ResultSet closedResultSet = mock(ResultSet.class);
    when(closedResultSet.next()).thenThrow(SpannerExceptionFactoryTest.newSessionNotFoundException(sessionName));
    when(closedContext.executeQuery(statement)).thenReturn(closedResultSet);
    when(closedSession.singleUse()).thenReturn(closedContext);
    final SessionImpl openSession = mockSession();
    ReadContext openContext = mock(ReadContext.class);
    ResultSet openResultSet = mock(ResultSet.class);
    when(openResultSet.next()).thenReturn(true, false);
    when(openContext.executeQuery(statement)).thenReturn(openResultSet);
    when(openSession.singleUse()).thenReturn(openContext);
    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);
    ReadContext context = pool.getSession().singleUse();
    ResultSet resultSet = context.executeQuery(statement);
    assertThat(resultSet.next()).isTrue();
}