Examples with SessionConsumer com.google.cloud.spanner.SessionClient.SessionConsumer used on opensource projects

Example 1 with SessionConsumer

use of com.google.cloud.spanner.SessionClient.SessionConsumer in project java-spanner by googleapis.

the class SessionClientTest method batchCreateSessionsDistributesMultipleRequestsOverChannels.

/**
 * Tests that multiple consecutive calls to {@link SessionClient#asyncBatchCreateSessions(int,
 * boolean, SessionConsumer)} with distributeOverChannels=false does not distribute one batch over
 * multiple channels, but it does assign each new call to a new channel. This means that multiple
 * calls to this method will still distribute the total set of sessions over all available
 * channels.
 */
@SuppressWarnings("unchecked")
@Test
public void batchCreateSessionsDistributesMultipleRequestsOverChannels() {
    DatabaseId db = DatabaseId.of(dbName);
    final String sessionName = dbName + "/sessions/s%d";
    final Map<String, String> labels = Collections.emptyMap();
    when(spannerOptions.getSessionLabels()).thenReturn(labels);
    final Set<Long> usedChannelHints = Collections.synchronizedSet(new HashSet<>());
    when(rpc.batchCreateSessions(Mockito.eq(dbName), Mockito.anyInt(), Mockito.eq(labels), Mockito.anyMap())).then(invocation -> {
        Map<Option, Object> options = invocation.getArgument(3, Map.class);
        Long channelHint = (Long) options.get(Option.CHANNEL_HINT);
        usedChannelHints.add(channelHint);
        int sessionCount = invocation.getArgument(1, Integer.class);
        List<com.google.spanner.v1.Session> res = new ArrayList<>();
        for (int i = 1; i <= sessionCount; i++) {
            res.add(com.google.spanner.v1.Session.newBuilder().setName(String.format(sessionName, i)).putAllLabels(labels).build());
        }
        return res;
    });
    final AtomicInteger returnedSessionCount = new AtomicInteger();
    SessionConsumer consumer = new SessionConsumer() {

        @Override
        public void onSessionReady(SessionImpl session) {
            assertThat(session.getName()).startsWith(dbName + "/sessions/s");
            returnedSessionCount.incrementAndGet();
            session.close();
        }

        @Override
        public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount) {
        }
    };
    final int numSessions = 10;
    final int numBatches = spannerOptions.getNumChannels() * 2;
    try (SessionClient client = new SessionClient(spanner, db, new TestExecutorFactory())) {
        for (int batch = 0; batch < numBatches; batch++) {
            client.asyncBatchCreateSessions(numSessions, false, consumer);
        }
    }
    assertThat(returnedSessionCount.get()).isEqualTo(numSessions * numBatches);
    assertThat(usedChannelHints.size()).isEqualTo(spannerOptions.getNumChannels() * 2);
    List<Long> expectedChannels = new ArrayList<>();
    for (long l = 0; l < spannerOptions.getNumChannels() * 2; l++) {
        expectedChannels.add(l);
    }
    assertThat(usedChannelHints).containsExactlyElementsIn(expectedChannels);
}

Example 2 with SessionConsumer

use of com.google.cloud.spanner.SessionClient.SessionConsumer in project java-spanner by googleapis.

the class SessionClientTest method batchCreateSessionsServerReturnsLessSessionsPerBatch.

@SuppressWarnings("unchecked")
@Test
public void batchCreateSessionsServerReturnsLessSessionsPerBatch() {
    final int MAX_SESSIONS_PER_BATCH = 5;
    DatabaseId db = DatabaseId.of(dbName);
    final String sessionName = dbName + "/sessions/s%d";
    when(rpc.batchCreateSessions(Mockito.eq(dbName), Mockito.anyInt(), Mockito.anyMap(), Mockito.anyMap())).then(invocation -> {
        int sessionCount = invocation.getArgument(1, Integer.class);
        List<com.google.spanner.v1.Session> res = new ArrayList<>();
        for (int i = 1; i <= Math.min(MAX_SESSIONS_PER_BATCH, sessionCount); i++) {
            res.add(com.google.spanner.v1.Session.newBuilder().setName(String.format(sessionName, i)).build());
        }
        return res;
    });
    final AtomicInteger returnedSessionCount = new AtomicInteger();
    SessionConsumer consumer = new SessionConsumer() {

        @Override
        public void onSessionReady(SessionImpl session) {
            assertThat(session.getName()).startsWith(dbName + "/sessions/s");
            returnedSessionCount.incrementAndGet();
            session.close();
        }

        @Override
        public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount) {
        }
    };
    // We want 100 sessions, but each rpc will only return 5. The consumer should still get 100
    // sessions.
    final int numSessions = 100;
    try (SessionClient client = new SessionClient(spanner, db, new TestExecutorFactory())) {
        client.asyncBatchCreateSessions(numSessions, true, consumer);
    }
    assertThat(returnedSessionCount.get()).isEqualTo(numSessions);
}

Example 3 with SessionConsumer

use of com.google.cloud.spanner.SessionClient.SessionConsumer in project java-spanner by googleapis.

the class SessionClientTest method batchCreateSessionsWithExceptions.

@SuppressWarnings("unchecked")
@Test
public void batchCreateSessionsWithExceptions() {
    for (AddRemoveSetException behavior : AddRemoveSetException.values()) {
        final List<Long> errorOnChannels = new ArrayList<>();
        if (behavior == AddRemoveSetException.REMOVE) {
            for (int c = 0; c < spannerOptions.getNumChannels(); c++) {
                errorOnChannels.add((long) c);
            }
        }
        for (int errorOnChannel = 0; errorOnChannel < spannerOptions.getNumChannels(); errorOnChannel++) {
            switch(behavior) {
                case SET:
                    errorOnChannels.clear();
                case ADD:
                    errorOnChannels.add((long) errorOnChannel);
                    break;
                case REMOVE:
                    errorOnChannels.remove(Long.valueOf(errorOnChannel));
                    break;
                default:
                    throw new IllegalStateException();
            }
            DatabaseId db = DatabaseId.of(dbName);
            final String sessionName = dbName + "/sessions/s%d";
            when(rpc.batchCreateSessions(Mockito.eq(dbName), Mockito.anyInt(), Mockito.anyMap(), Mockito.anyMap())).then(invocation -> {
                Map<Option, Object> options = invocation.getArgument(3, Map.class);
                Long channelHint = (Long) options.get(Option.CHANNEL_HINT);
                if (errorOnChannels.contains(channelHint)) {
                    throw SpannerExceptionFactory.newSpannerException(ErrorCode.RESOURCE_EXHAUSTED, "could not create any more sessions");
                } else {
                    int sessionCount = invocation.getArgument(1, Integer.class);
                    List<com.google.spanner.v1.Session> res = new ArrayList<>();
                    for (int i = 1; i <= sessionCount; i++) {
                        res.add(com.google.spanner.v1.Session.newBuilder().setName(String.format(sessionName, i)).build());
                    }
                    return res;
                }
            });
            final AtomicInteger errorForSessionsCount = new AtomicInteger();
            final AtomicInteger errorCount = new AtomicInteger();
            final AtomicInteger returnedSessionCount = new AtomicInteger();
            SessionConsumer consumer = new SessionConsumer() {

                @Override
                public void onSessionReady(SessionImpl session) {
                    assertThat(session.getName()).startsWith(dbName + "/sessions/s");
                    returnedSessionCount.incrementAndGet();
                    session.close();
                }

                @Override
                public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount) {
                    assertThat(t).isInstanceOf(SpannerException.class);
                    SpannerException e = (SpannerException) t;
                    assertThat(e.getErrorCode()).isEqualTo(ErrorCode.RESOURCE_EXHAUSTED);
                    errorCount.incrementAndGet();
                    errorForSessionsCount.addAndGet(createFailureForSessionCount);
                }
            };
            final int numSessions = 10;
            try (SessionClient client = new SessionClient(spanner, db, new TestExecutorFactory())) {
                client.asyncBatchCreateSessions(numSessions, true, consumer);
            }
            assertThat(errorCount.get()).isEqualTo(errorOnChannels.size());
            assertThat(returnedSessionCount.get()).isAtLeast(numSessions - ((numSessions / spannerOptions.getNumChannels()) * errorOnChannels.size() + numSessions % spannerOptions.getNumChannels()));
            assertThat(returnedSessionCount.get() + errorForSessionsCount.get()).isEqualTo(numSessions);
        }
    }
}

Example 4 with SessionConsumer

use of com.google.cloud.spanner.SessionClient.SessionConsumer in project java-spanner by googleapis.

the class SessionClientTest method batchCreateAndCloseSessions.

@SuppressWarnings("unchecked")
@Test
public void batchCreateAndCloseSessions() {
    DatabaseId db = DatabaseId.of(dbName);
    final String sessionName = dbName + "/sessions/s%d";
    final Map<String, String> labels = new HashMap<>();
    labels.put("env", "dev");
    when(spannerOptions.getSessionLabels()).thenReturn(labels);
    final List<Long> usedChannels = Collections.synchronizedList(new ArrayList<>());
    when(rpc.batchCreateSessions(Mockito.eq(dbName), Mockito.anyInt(), Mockito.eq(labels), Mockito.anyMap())).then(invocation -> {
        Map<Option, Object> options = invocation.getArgument(3, Map.class);
        Long channelHint = (Long) options.get(Option.CHANNEL_HINT);
        usedChannels.add(channelHint);
        int sessionCount = invocation.getArgument(1, Integer.class);
        List<com.google.spanner.v1.Session> res = new ArrayList<>();
        for (int i = 1; i <= sessionCount; i++) {
            res.add(com.google.spanner.v1.Session.newBuilder().setName(String.format(sessionName, i)).putAllLabels(labels).build());
        }
        return res;
    });
    final AtomicInteger returnedSessionCount = new AtomicInteger();
    SessionConsumer consumer = new SessionConsumer() {

        @Override
        public void onSessionReady(SessionImpl session) {
            assertThat(session.getName()).startsWith(dbName + "/sessions/s");
            returnedSessionCount.incrementAndGet();
            session.close();
        }

        @Override
        public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount) {
        }
    };
    final int numSessions = 10;
    try (SessionClient client = new SessionClient(spanner, db, new TestExecutorFactory())) {
        client.asyncBatchCreateSessions(numSessions, true, consumer);
    }
    assertThat(returnedSessionCount.get()).isEqualTo(numSessions);
    assertThat(usedChannels.size()).isEqualTo(spannerOptions.getNumChannels());
    List<Long> expectedChannels = new ArrayList<>();
    for (long l = 0; l < spannerOptions.getNumChannels(); l++) {
        expectedChannels.add(l);
    }
    assertThat(usedChannels).containsExactlyElementsIn(expectedChannels);
}