BatchCreateSessionsTest.java

/*
 * Copyright 2019 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.google.cloud.spanner;

import static com.google.cloud.spanner.DisableDefaultMtlsProvider.disableDefaultMtlsProvider;
import static org.hamcrest.CoreMatchers.equalTo;
import static org.hamcrest.CoreMatchers.is;
import static org.hamcrest.MatcherAssert.assertThat;

import com.google.api.gax.grpc.testing.LocalChannelProvider;
import com.google.cloud.NoCredentials;
import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime;
import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
import com.google.common.base.Stopwatch;
import com.google.protobuf.ListValue;
import com.google.spanner.v1.ResultSetMetadata;
import com.google.spanner.v1.StructType;
import com.google.spanner.v1.StructType.Field;
import com.google.spanner.v1.TypeCode;
import io.grpc.Server;
import io.grpc.Status;
import io.grpc.inprocess.InProcessServerBuilder;
import java.util.concurrent.TimeUnit;
import org.junit.AfterClass;
import org.junit.Before;
import org.junit.BeforeClass;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;

@RunWith(JUnit4.class)
public class BatchCreateSessionsTest {
  private static final Statement SELECT1AND2 =
      Statement.of("SELECT 1 AS COL1 UNION ALL SELECT 2 AS COL1");
  private static final ResultSetMetadata SELECT1AND2_METADATA =
      ResultSetMetadata.newBuilder()
          .setRowType(
              StructType.newBuilder()
                  .addFields(
                      Field.newBuilder()
                          .setName("COL1")
                          .setType(
                              com.google.spanner.v1.Type.newBuilder()
                                  .setCode(TypeCode.INT64)
                                  .build())
                          .build())
                  .build())
          .build();
  private static final com.google.spanner.v1.ResultSet SELECT1_RESULTSET =
      com.google.spanner.v1.ResultSet.newBuilder()
          .addRows(
              ListValue.newBuilder()
                  .addValues(com.google.protobuf.Value.newBuilder().setStringValue("1").build())
                  .build())
          .addRows(
              ListValue.newBuilder()
                  .addValues(com.google.protobuf.Value.newBuilder().setStringValue("2").build())
                  .build())
          .setMetadata(SELECT1AND2_METADATA)
          .build();

  private static MockSpannerServiceImpl mockSpanner;
  private static Server server;
  private static LocalChannelProvider channelProvider;

  @BeforeClass
  public static void startStaticServer() throws Exception {
    disableDefaultMtlsProvider();
    mockSpanner = new MockSpannerServiceImpl();
    mockSpanner.setAbortProbability(0.0D); // We don't want any unpredictable aborted transactions.
    mockSpanner.putStatementResult(StatementResult.query(SELECT1AND2, SELECT1_RESULTSET));

    String uniqueName = InProcessServerBuilder.generateName();
    server =
        InProcessServerBuilder.forName(uniqueName)
            .directExecutor()
            .addService(mockSpanner)
            .build()
            .start();
    channelProvider = LocalChannelProvider.create(uniqueName);
  }

  @AfterClass
  public static void stopServer() throws InterruptedException {
    server.shutdown();
    server.awaitTermination();
  }

  @Before
  public void setUp() {
    mockSpanner.reset();
    mockSpanner.removeAllExecutionTimes();
  }

  private Spanner createSpanner(int minSessions, int maxSessions) {
    SessionPoolOptions sessionPoolOptions =
        SessionPoolOptions.newBuilder()
            .setMinSessions(minSessions)
            .setMaxSessions(maxSessions)
            .build();
    return SpannerOptions.newBuilder()
        .setProjectId("[PROJECT]")
        .setChannelProvider(channelProvider)
        .setSessionPoolOption(sessionPoolOptions)
        .setCredentials(NoCredentials.getInstance())
        .build()
        .getService();
  }

  @Test
  public void testCreatedMinSessions() throws InterruptedException {
    int minSessions = 1000;
    int maxSessions = 4000;
    try (Spanner spanner = createSpanner(minSessions, maxSessions)) {
      DatabaseClientImpl client =
          (DatabaseClientImpl)
              spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]"));
      Stopwatch watch = Stopwatch.createStarted();
      while (client.pool.totalSessions() < minSessions && watch.elapsed(TimeUnit.SECONDS) < 10) {
        Thread.sleep(10L);
      }
      assertThat(client.pool.totalSessions(), is(equalTo(minSessions)));
    }
  }

  @Test
  public void testClosePoolWhileInitializing() throws InterruptedException {
    int minSessions = 10_000;
    int maxSessions = 10_000;
    DatabaseClientImpl client;
    // Freeze the server to prevent it from creating sessions before we want to.
    mockSpanner.freeze();
    try (Spanner spanner = createSpanner(minSessions, maxSessions)) {
      // Create a database client which will create a session pool.
      // No sessions will be created at the moment as the server is frozen.
      client =
          (DatabaseClientImpl)
              spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]"));
      // Make sure session creation takes a little time to avoid all sessions being created at once.
      mockSpanner.setBatchCreateSessionsExecutionTime(
          SimulatedExecutionTime.ofMinimumAndRandomTime(10, 0));
      // Unfreeze the server to allow session creation to start.
      mockSpanner.unfreeze();
      // Wait until at least one batch of sessions has been created.
      Stopwatch watch = Stopwatch.createStarted();
      while (client.pool.totalSessions() == 0 && watch.elapsed(TimeUnit.SECONDS) < 10) {
        Thread.sleep(1L);
      }
      // Close the Spanner instance which will start to delete sessions while the session pool is
      // still being initialized.
    }
    // Verify that all sessions have been deleted.
    assertThat(client.pool.totalSessions(), is(equalTo(0)));
  }

  @Test
  public void testSpannerReturnsAllAvailableSessionsAndThenNoSessions()
      throws InterruptedException {
    int minSessions = 1000;
    int maxSessions = 1000;
    // Set a maximum number of sessions that will be created by the server.
    // After this the server will return an error when batchCreateSessions is called.
    // This error is not propagated to the client.
    int maxServerSessions = 550;
    DatabaseClientImpl client;
    mockSpanner.setMaxTotalSessions(maxServerSessions);
    try (Spanner spanner = createSpanner(minSessions, maxSessions)) {
      // Create a database client which will create a session pool.
      client =
          (DatabaseClientImpl)
              spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]"));
      Stopwatch watch = Stopwatch.createStarted();
      while (client.pool.totalSessions() < maxServerSessions
          && watch.elapsed(TimeUnit.SECONDS) < 10) {
        Thread.sleep(10L);
      }
      assertThat(client.pool.totalSessions(), is(equalTo(maxServerSessions)));
      // Wait until the pool has given up creating sessions.
      watch = watch.reset();
      watch.start();
      while (client.pool.getNumberOfSessionsBeingCreated() > 0
          && watch.elapsed(TimeUnit.SECONDS) < 10) {
        Thread.sleep(10L);
      }
      // Remove the max server sessions limit.
      mockSpanner.setMaxTotalSessions(Integer.MAX_VALUE);
      // Wait a little. No more sessions should be created, as the previous requests have given up,
      // and no new sessions have been requested from the pool.
      Thread.sleep(20L);
      assertThat(client.pool.totalSessions(), is(equalTo(maxServerSessions)));
    }
    // Verify that all sessions have been deleted.
    assertThat(client.pool.totalSessions(), is(equalTo(0)));
  }

  @Test
  public void testSpannerReturnsFailedPrecondition() throws InterruptedException {
    int minSessions = 100;
    int maxSessions = 1000;
    int expectedSessions;
    DatabaseClientImpl client;
    // Make the first BatchCreateSessions return an error.
    mockSpanner.addException(Status.FAILED_PRECONDITION.asRuntimeException());
    try (Spanner spanner = createSpanner(minSessions, maxSessions)) {
      // Create a database client which will create a session pool.
      client =
          (DatabaseClientImpl)
              spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]"));
      // Wait for the pool to be initialized.
      // The first session creation request will fail.
      expectedSessions = minSessions - minSessions / spanner.getOptions().getNumChannels();
      Stopwatch watch = Stopwatch.createStarted();
      while (client.pool.totalSessions() < expectedSessions
          && watch.elapsed(TimeUnit.SECONDS) < 10) {
        Thread.sleep(10L);
      }
      // Wait a little to allow any additional session creation to finish.
      Thread.sleep(20L);
    }
    // Verify that all sessions have been deleted.
    assertThat(client.pool.totalSessions(), is(equalTo(0)));
  }
}