perf: prepare sessions with r/w tx in-process

Preparing sessions with a read/write transaction using a background executor works well as long as that executor is not being overloaded. When the executor has reached its limit, it is more efficient to allow the read/write transaction to be created in-process, as that scales with the number of user threads available, instead of being limited to the fixed thread pool of the background executor. Fixes #151

Author: olavloite

google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java

@@ -1093,6 +1093,7 @@ private static enum Position {
  private final ScheduledExecutorService executor;
  private final ExecutorFactory<ScheduledExecutorService> executorFactory;
  private final ScheduledExecutorService prepareExecutor;
+ private final int prepareThreadPoolSize;
  final PoolMaintainer poolMaintainer;
  private final Clock clock;
  private final Object lock = new Object();
@@ -1137,6 +1138,12 @@ private static enum Position {
  @GuardedBy("lock")
  private long numSessionsReleased = 0;
 
+ @GuardedBy("lock")
+ private long numSessionsInProcessPrepared = 0;
+
+ @GuardedBy("lock")
+ private long numSessionsAsyncPrepared = 0;
+
  private AtomicLong numWaiterTimeouts = new AtomicLong();
 
  @GuardedBy("lock")
@@ -1213,15 +1220,14 @@ private SessionPool(
  this.options = options;
  this.executorFactory = executorFactory;
  this.executor = executor;
- int prepareThreads;
  if (executor instanceof ThreadPoolExecutor) {
- prepareThreads = Math.max(((ThreadPoolExecutor) executor).getCorePoolSize(), 1);
+ prepareThreadPoolSize = Math.max(((ThreadPoolExecutor) executor).getCorePoolSize(), 1);
  } else {
- prepareThreads = 8;
+ prepareThreadPoolSize = 8;
  }
  this.prepareExecutor =
  Executors.newScheduledThreadPool(
- prepareThreads,
+ prepareThreadPoolSize,
  new ThreadFactoryBuilder()
  .setDaemon(true)
  .setNameFormat("session-pool-prepare-%d")
@@ -1232,6 +1238,20 @@ private SessionPool(
  this.initMetricsCollection(metricRegistry, labelValues);
  }
 
+ @VisibleForTesting
+ long getNumberOfSessionsInProcessPrepared() {
+ synchronized (lock) {
+ return numSessionsInProcessPrepared;
+ }
+ }
+
+ @VisibleForTesting
+ long getNumberOfSessionsAsyncPrepared() {
+ synchronized (lock) {
+ return numSessionsAsyncPrepared;
+ }
+ }
+
  @VisibleForTesting
  int getNumberOfAvailableWritePreparedSessions() {
  synchronized (lock) {
@@ -1416,46 +1436,93 @@ PooledSession getReadSession() throws SpannerException {
  PooledSession getReadWriteSession() {
  Span span = Tracing.getTracer().getCurrentSpan();
  span.addAnnotation("Acquiring read write session");
- Waiter waiter = null;
  PooledSession sess = null;
- synchronized (lock) {
- if (closureFuture != null) {
- span.addAnnotation("Pool has been closed");
- throw new IllegalStateException("Pool has been closed");
+ // Loop to retry SessionNotFoundExceptions that might occur during in-process prepare of a
+ // session.
+ while (true) {
+ Waiter waiter = null;
+ boolean inProcessPrepare = false;
+ synchronized (lock) {
+ if (closureFuture != null) {
+ span.addAnnotation("Pool has been closed");
+ throw new IllegalStateException("Pool has been closed");
+ }
+ if (resourceNotFoundException != null) {
+ span.addAnnotation("Database has been deleted");
+ throw SpannerExceptionFactory.newSpannerException(
+ ErrorCode.NOT_FOUND,
+ String.format(
+ "The session pool has been invalidated because a previous RPC returned 'Database not found': %s",
+ resourceNotFoundException.getMessage()),
+ resourceNotFoundException);
+ }
+ sess = writePreparedSessions.poll();
+ if (sess == null) {
+ if (numSessionsBeingPrepared <= prepareThreadPoolSize) {
+ if (numSessionsBeingPrepared <= readWriteWaiters.size()) {
+ PooledSession readSession = readSessions.poll();
+ if (readSession != null) {
+ span.addAnnotation(
+ "Acquired read only session. Preparing for read write transaction");
+ prepareSession(readSession);
+ } else {
+ span.addAnnotation("No session available");
+ maybeCreateSession();
+ }
+ }
+ } else {
+ inProcessPrepare = true;
+ numSessionsInProcessPrepared++;
+ PooledSession readSession = readSessions.poll();
+ if (readSession != null) {
+ // Create a read/write transaction in-process if there is already a queue for prepared
+ // sessions. This is more efficient than doing it asynchronously, as it scales with
+ // the number of user threads. The thread pool for asynchronously preparing sessions
+ // is fixed.
+ span.addAnnotation(
+ "Acquired read only session. Preparing in-process for read write transaction");
+ sess = readSession;
+ } else {
+ span.addAnnotation("No session available");
+ maybeCreateSession();
+ }
+ }
+ if (sess == null) {
+ waiter = new Waiter();
+ if (inProcessPrepare) {
+ readWaiters.add(waiter);
+ } else {
+ readWriteWaiters.add(waiter);
+ }
+ }
+ } else {
+ span.addAnnotation("Acquired read write session");
+ }
  }
- if (resourceNotFoundException != null) {
- span.addAnnotation("Database has been deleted");
- throw SpannerExceptionFactory.newSpannerException(
- ErrorCode.NOT_FOUND,
- String.format(
- "The session pool has been invalidated because a previous RPC returned 'Database not found': %s",
- resourceNotFoundException.getMessage()),
- resourceNotFoundException);
+ if (waiter != null) {
+ logger.log(
+ Level.FINE,
+ "No session available in the pool. Blocking for one to become available/created");
+ span.addAnnotation("Waiting for read write session to be available");
+ sess = waiter.take();
  }
- sess = writePreparedSessions.poll();
- if (sess == null) {
- if (numSessionsBeingPrepared <= readWriteWaiters.size()) {
- PooledSession readSession = readSessions.poll();
- if (readSession != null) {
- span.addAnnotation("Acquired read only session. Preparing for read write transaction");
- prepareSession(readSession);
- } else {
- span.addAnnotation("No session available");
- maybeCreateSession();
+ if (inProcessPrepare) {
+ try {
+ sess.prepareReadWriteTransaction();
+ } catch (Throwable t) {
+ sess = null;
+ SpannerException e = newSpannerException(t);
+ if (!isClosed()) {
+ handlePrepareSessionFailure(e, sess, false);
+ }
+ if (!isSessionNotFound(e)) {
+ throw e;
  }
  }
- waiter = new Waiter();
- readWriteWaiters.add(waiter);
- } else {
- span.addAnnotation("Acquired read write session");
  }
- }
- if (waiter != null) {
- logger.log(
- Level.FINE,
- "No session available in the pool. Blocking for one to become available/created");
- span.addAnnotation("Waiting for read write session to be available");
- sess = waiter.take();
+ if (sess != null) {
+ break;
+ }
  }
  sess.markBusy();
  incrementNumSessionsInUse();
@@ -1583,7 +1650,8 @@ private void handleCreateSessionsFailure(SpannerException e, int count) {
  }
  }
 
- private void handlePrepareSessionFailure(SpannerException e, PooledSession session) {
+ private void handlePrepareSessionFailure(
+ SpannerException e, PooledSession session, boolean informFirstWaiter) {
  synchronized (lock) {
  if (isSessionNotFound(e)) {
  invalidateSession(session);
@@ -1606,7 +1674,7 @@ private void handlePrepareSessionFailure(SpannerException e, PooledSession sessi
  MoreObjects.firstNonNull(
  this.resourceNotFoundException,
  isDatabaseOrInstanceNotFound(e) ? (ResourceNotFoundException) e : null);
- } else if (readWriteWaiters.size() > 0) {
+ } else if (informFirstWaiter && readWriteWaiters.size() > 0) {
  releaseSession(session, Position.FIRST);
  readWriteWaiters.poll().put(e);
  } else {
@@ -1755,6 +1823,7 @@ public void run() {
  sess.prepareReadWriteTransaction();
  logger.log(Level.FINE, "Session prepared");
  synchronized (lock) {
+ numSessionsAsyncPrepared++;
  numSessionsBeingPrepared--;
  if (!isClosed()) {
  if (readWriteWaiters.size() > 0) {
@@ -1770,7 +1839,7 @@ public void run() {
  synchronized (lock) {
  numSessionsBeingPrepared--;
  if (!isClosed()) {
- handlePrepareSessionFailure(newSpannerException(t), sess);
+ handlePrepareSessionFailure(newSpannerException(t), sess, true);
  }
  }
  }