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

Example 11 with PooledSessionFuture

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

the class DatabaseClientImplTest method testTransactionManagerAsync_usesOptions.

@Test
public void testTransactionManagerAsync_usesOptions() {
    SessionPool pool = mock(SessionPool.class);
    PooledSessionFuture session = mock(PooledSessionFuture.class);
    when(pool.getSession()).thenReturn(session);
    TransactionOption option = mock(TransactionOption.class);
    DatabaseClientImpl client = new DatabaseClientImpl(pool);
    client.transactionManagerAsync(option);
    verify(session).transactionManagerAsync(option);
}

Example 12 with PooledSessionFuture

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

the class SessionPoolTest method testSessionNotFoundReadWriteTransaction.

@SuppressWarnings("unchecked")
@Test
public void testSessionNotFoundReadWriteTransaction() {
    final Statement queryStatement = Statement.of("SELECT 1");
    final Statement updateStatement = Statement.of("UPDATE FOO SET BAR=1 WHERE ID=2");
    final SpannerException sessionNotFound = SpannerExceptionFactoryTest.newSessionNotFoundException(sessionName);
    for (ReadWriteTransactionTestStatementType statementType : ReadWriteTransactionTestStatementType.values()) {
        final ReadWriteTransactionTestStatementType executeStatementType = statementType;
        SpannerRpc.StreamingCall closedStreamingCall = mock(SpannerRpc.StreamingCall.class);
        doThrow(sessionNotFound).when(closedStreamingCall).request(Mockito.anyInt());
        SpannerRpc rpc = mock(SpannerRpc.class);
        when(rpc.asyncDeleteSession(Mockito.anyString(), Mockito.anyMap())).thenReturn(ApiFutures.immediateFuture(Empty.getDefaultInstance()));
        when(rpc.executeQuery(any(ExecuteSqlRequest.class), any(ResultStreamConsumer.class), any(Map.class))).thenReturn(closedStreamingCall);
        when(rpc.executeQuery(any(ExecuteSqlRequest.class), any(Map.class))).thenThrow(sessionNotFound);
        when(rpc.executeBatchDml(any(ExecuteBatchDmlRequest.class), any(Map.class))).thenThrow(sessionNotFound);
        when(rpc.commitAsync(any(CommitRequest.class), any(Map.class))).thenReturn(ApiFutures.<CommitResponse>immediateFailedFuture(sessionNotFound));
        when(rpc.rollbackAsync(any(RollbackRequest.class), any(Map.class))).thenReturn(ApiFutures.<Empty>immediateFailedFuture(sessionNotFound));
        final SessionImpl closedSession = mock(SessionImpl.class);
        when(closedSession.getName()).thenReturn("projects/dummy/instances/dummy/database/dummy/sessions/session-closed");
        final TransactionContextImpl closedTransactionContext = TransactionContextImpl.newBuilder().setSession(closedSession).setOptions(Options.fromTransactionOptions()).setRpc(rpc).build();
        when(closedSession.asyncClose()).thenReturn(ApiFutures.immediateFuture(Empty.getDefaultInstance()));
        when(closedSession.newTransaction(Options.fromTransactionOptions())).thenReturn(closedTransactionContext);
        when(closedSession.beginTransactionAsync()).thenThrow(sessionNotFound);
        TransactionRunnerImpl closedTransactionRunner = new TransactionRunnerImpl(closedSession);
        closedTransactionRunner.setSpan(mock(Span.class));
        when(closedSession.readWriteTransaction()).thenReturn(closedTransactionRunner);
        final SessionImpl openSession = mock(SessionImpl.class);
        when(openSession.asyncClose()).thenReturn(ApiFutures.immediateFuture(Empty.getDefaultInstance()));
        when(openSession.getName()).thenReturn("projects/dummy/instances/dummy/database/dummy/sessions/session-open");
        final TransactionContextImpl openTransactionContext = mock(TransactionContextImpl.class);
        when(openSession.newTransaction(Options.fromTransactionOptions())).thenReturn(openTransactionContext);
        when(openSession.beginTransactionAsync()).thenReturn(ApiFutures.immediateFuture(ByteString.copyFromUtf8("open-txn")));
        TransactionRunnerImpl openTransactionRunner = new TransactionRunnerImpl(openSession);
        openTransactionRunner.setSpan(mock(Span.class));
        when(openSession.readWriteTransaction()).thenReturn(openTransactionRunner);
        ResultSet openResultSet = mock(ResultSet.class);
        when(openResultSet.next()).thenReturn(true, false);
        ResultSet planResultSet = mock(ResultSet.class);
        when(planResultSet.getStats()).thenReturn(ResultSetStats.getDefaultInstance());
        when(openTransactionContext.executeQuery(queryStatement)).thenReturn(openResultSet);
        when(openTransactionContext.analyzeQuery(queryStatement, QueryAnalyzeMode.PLAN)).thenReturn(planResultSet);
        when(openTransactionContext.executeUpdate(updateStatement)).thenReturn(1L);
        when(openTransactionContext.batchUpdate(Arrays.asList(updateStatement, updateStatement))).thenReturn(new long[] { 1L, 1L });
        SpannerImpl spanner = mock(SpannerImpl.class);
        SessionClient sessionClient = mock(SessionClient.class);
        when(spanner.getSessionClient(db)).thenReturn(sessionClient);
        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));
        SessionPoolOptions options = SessionPoolOptions.newBuilder().setMinSessions(// The pool should not auto-create any sessions
        0).setMaxSessions(2).setIncStep(1).setBlockIfPoolExhausted().build();
        SpannerOptions spannerOptions = mock(SpannerOptions.class);
        when(spannerOptions.getSessionPoolOptions()).thenReturn(options);
        when(spannerOptions.getNumChannels()).thenReturn(4);
        when(spanner.getOptions()).thenReturn(spannerOptions);
        SessionPool pool = SessionPool.createPool(options, new TestExecutorFactory(), spanner.getSessionClient(db));
        try (PooledSessionFuture readWriteSession = pool.getSession()) {
            TransactionRunner runner = readWriteSession.readWriteTransaction();
            try {
                runner.run(new TransactionCallable<Integer>() {

                    private int callNumber = 0;

                    @Override
                    public Integer run(TransactionContext transaction) {
                        callNumber++;
                        if (callNumber == 1) {
                            assertThat(transaction).isEqualTo(closedTransactionContext);
                        } else {
                            assertThat(transaction).isEqualTo(openTransactionContext);
                        }
                        switch(executeStatementType) {
                            case QUERY:
                                ResultSet resultSet = transaction.executeQuery(queryStatement);
                                assertThat(resultSet.next()).isTrue();
                                break;
                            case ANALYZE:
                                ResultSet planResultSet = transaction.analyzeQuery(queryStatement, QueryAnalyzeMode.PLAN);
                                assertThat(planResultSet.next()).isFalse();
                                assertThat(planResultSet.getStats()).isNotNull();
                                break;
                            case UPDATE:
                                long updateCount = transaction.executeUpdate(updateStatement);
                                assertThat(updateCount).isEqualTo(1L);
                                break;
                            case BATCH_UPDATE:
                                long[] updateCounts = transaction.batchUpdate(Arrays.asList(updateStatement, updateStatement));
                                assertThat(updateCounts).isEqualTo(new long[] { 1L, 1L });
                                break;
                            case WRITE:
                                transaction.buffer(Mutation.delete("FOO", Key.of(1L)));
                                break;
                            case EXCEPTION:
                                throw new RuntimeException("rollback at call " + callNumber);
                            default:
                                fail("Unknown statement type: " + executeStatementType);
                        }
                        return callNumber;
                    }
                });
            } catch (Exception e) {
                // The rollback will also cause a SessionNotFoundException, but this is caught, logged
                // and further ignored by the library, meaning that the session will not be re-created
                // for retry. Hence rollback at call 1.
                assertThat(executeStatementType).isEqualTo(ReadWriteTransactionTestStatementType.EXCEPTION);
                assertThat(e.getMessage()).contains("rollback at call 1");
            }
        }
        pool.closeAsync(new SpannerImpl.ClosedException());
    }
}

Example 13 with PooledSessionFuture

use of com.google.cloud.spanner.SessionPool.PooledSessionFuture 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 14 with PooledSessionFuture

use of com.google.cloud.spanner.SessionPool.PooledSessionFuture 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 15 with PooledSessionFuture

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

the class SessionPoolTest method testSessionMetrics.

@Test
public void testSessionMetrics() throws Exception {
    // Create a session pool with max 2 session and a low timeout for waiting for a session.
    options = SessionPoolOptions.newBuilder().setMinSessions(1).setMaxSessions(2).setMaxIdleSessions(0).setInitialWaitForSessionTimeoutMillis(50L).build();
    FakeClock clock = new FakeClock();
    clock.currentTimeMillis = System.currentTimeMillis();
    FakeMetricRegistry metricRegistry = new FakeMetricRegistry();
    List<LabelValue> labelValues = Arrays.asList(LabelValue.create("client1"), LabelValue.create("database1"), LabelValue.create("instance1"), LabelValue.create("1.0.0"));
    setupMockSessionCreation();
    pool = createPool(clock, metricRegistry, labelValues);
    PooledSessionFuture session1 = pool.getSession();
    PooledSessionFuture session2 = pool.getSession();
    session1.get();
    session2.get();
    MetricsRecord record = metricRegistry.pollRecord();
    assertThat(record.getMetrics().size()).isEqualTo(6);
    List<PointWithFunction> maxInUseSessions = record.getMetrics().get(MetricRegistryConstants.MAX_IN_USE_SESSIONS);
    assertThat(maxInUseSessions.size()).isEqualTo(1);
    assertThat(maxInUseSessions.get(0).value()).isEqualTo(2L);
    assertThat(maxInUseSessions.get(0).keys()).isEqualTo(SPANNER_LABEL_KEYS);
    assertThat(maxInUseSessions.get(0).values()).isEqualTo(labelValues);
    List<PointWithFunction> getSessionsTimeouts = record.getMetrics().get(MetricRegistryConstants.GET_SESSION_TIMEOUTS);
    assertThat(getSessionsTimeouts.size()).isEqualTo(1);
    assertThat(getSessionsTimeouts.get(0).value()).isAtMost(1L);
    assertThat(getSessionsTimeouts.get(0).keys()).isEqualTo(SPANNER_LABEL_KEYS);
    assertThat(getSessionsTimeouts.get(0).values()).isEqualTo(labelValues);
    List<PointWithFunction> numAcquiredSessions = record.getMetrics().get(MetricRegistryConstants.NUM_ACQUIRED_SESSIONS);
    assertThat(numAcquiredSessions.size()).isEqualTo(1);
    assertThat(numAcquiredSessions.get(0).value()).isEqualTo(2L);
    assertThat(numAcquiredSessions.get(0).keys()).isEqualTo(SPANNER_LABEL_KEYS);
    assertThat(numAcquiredSessions.get(0).values()).isEqualTo(labelValues);
    List<PointWithFunction> numReleasedSessions = record.getMetrics().get(MetricRegistryConstants.NUM_RELEASED_SESSIONS);
    assertThat(numReleasedSessions.size()).isEqualTo(1);
    assertThat(numReleasedSessions.get(0).value()).isEqualTo(0);
    assertThat(numReleasedSessions.get(0).keys()).isEqualTo(SPANNER_LABEL_KEYS);
    assertThat(numReleasedSessions.get(0).values()).isEqualTo(labelValues);
    List<PointWithFunction> maxAllowedSessions = record.getMetrics().get(MetricRegistryConstants.MAX_ALLOWED_SESSIONS);
    assertThat(maxAllowedSessions.size()).isEqualTo(1);
    assertThat(maxAllowedSessions.get(0).value()).isEqualTo(options.getMaxSessions());
    assertThat(maxAllowedSessions.get(0).keys()).isEqualTo(SPANNER_LABEL_KEYS);
    assertThat(maxAllowedSessions.get(0).values()).isEqualTo(labelValues);
    List<PointWithFunction> numSessionsInPool = record.getMetrics().get(MetricRegistryConstants.NUM_SESSIONS_IN_POOL);
    assertThat(numSessionsInPool.size()).isEqualTo(4);
    PointWithFunction beingPrepared = numSessionsInPool.get(0);
    List<LabelValue> labelValuesWithBeingPreparedType = new ArrayList<>(labelValues);
    labelValuesWithBeingPreparedType.add(NUM_SESSIONS_BEING_PREPARED);
    assertThat(beingPrepared.value()).isEqualTo(0L);
    assertThat(beingPrepared.keys()).isEqualTo(SPANNER_LABEL_KEYS_WITH_TYPE);
    assertThat(beingPrepared.values()).isEqualTo(labelValuesWithBeingPreparedType);
    PointWithFunction numSessionsInUse = numSessionsInPool.get(1);
    List<LabelValue> labelValuesWithInUseType = new ArrayList<>(labelValues);
    labelValuesWithInUseType.add(NUM_IN_USE_SESSIONS);
    assertThat(numSessionsInUse.value()).isEqualTo(2L);
    assertThat(numSessionsInUse.keys()).isEqualTo(SPANNER_LABEL_KEYS_WITH_TYPE);
    assertThat(numSessionsInUse.values()).isEqualTo(labelValuesWithInUseType);
    PointWithFunction readSessions = numSessionsInPool.get(2);
    List<LabelValue> labelValuesWithReadType = new ArrayList<>(labelValues);
    labelValuesWithReadType.add(NUM_READ_SESSIONS);
    assertThat(readSessions.value()).isEqualTo(0L);
    assertThat(readSessions.keys()).isEqualTo(SPANNER_LABEL_KEYS_WITH_TYPE);
    assertThat(readSessions.values()).isEqualTo(labelValuesWithReadType);
    PointWithFunction writePreparedSessions = numSessionsInPool.get(3);
    List<LabelValue> labelValuesWithWriteType = new ArrayList<>(labelValues);
    labelValuesWithWriteType.add(NUM_WRITE_SESSIONS);
    assertThat(writePreparedSessions.value()).isEqualTo(0L);
    assertThat(writePreparedSessions.keys()).isEqualTo(SPANNER_LABEL_KEYS_WITH_TYPE);
    assertThat(writePreparedSessions.values()).isEqualTo(labelValuesWithWriteType);
    final CountDownLatch latch = new CountDownLatch(1);
    // Try asynchronously to take another session. This attempt should time out.
    Future<Void> fut = executor.submit(() -> {
        latch.countDown();
        Session session = pool.getSession();
        session.close();
        return null;
    });
    // Wait until the background thread is actually waiting for a session.
    latch.await();
    // Wait until the request has timed out.
    int waitCount = 0;
    while (pool.getNumWaiterTimeouts() == 0L && waitCount < 1000) {
        Thread.sleep(5L);
        waitCount++;
    }
    // Return the checked out session to the pool so the async request will get a session and
    // finish.
    session2.close();
    // Verify that the async request also succeeds.
    fut.get(10L, TimeUnit.SECONDS);
    executor.shutdown();
    session1.close();
    numAcquiredSessions = record.getMetrics().get(MetricRegistryConstants.NUM_ACQUIRED_SESSIONS);
    assertThat(numAcquiredSessions.size()).isEqualTo(1);
    assertThat(numAcquiredSessions.get(0).value()).isEqualTo(3L);
    numReleasedSessions = record.getMetrics().get(MetricRegistryConstants.NUM_RELEASED_SESSIONS);
    assertThat(numReleasedSessions.size()).isEqualTo(1);
    assertThat(numReleasedSessions.get(0).value()).isEqualTo(3L);
    maxInUseSessions = record.getMetrics().get(MetricRegistryConstants.MAX_IN_USE_SESSIONS);
    assertThat(maxInUseSessions.size()).isEqualTo(1);
    assertThat(maxInUseSessions.get(0).value()).isEqualTo(2L);
    numSessionsInPool = record.getMetrics().get(MetricRegistryConstants.NUM_SESSIONS_IN_POOL);
    assertThat(numSessionsInPool.size()).isEqualTo(4);
    beingPrepared = numSessionsInPool.get(0);
    assertThat(beingPrepared.value()).isEqualTo(0L);
    numSessionsInUse = numSessionsInPool.get(1);
    assertThat(numSessionsInUse.value()).isEqualTo(0L);
    readSessions = numSessionsInPool.get(2);
    assertThat(readSessions.value()).isEqualTo(2L);
    writePreparedSessions = numSessionsInPool.get(3);
    assertThat(writePreparedSessions.value()).isEqualTo(0L);
}