refactor: add internal data structures for transactional connection state

This change adds internal data structures that can be used for transactional
connection state. These data structures also reduces the amount of code that
is needed for each connection property that is added. Connection properties
are currently represented as actual variables in the ConnectionImpl class.
These new data structures removes the need for that.

Only the connection property retryAbortsInternally is refactored to use
the new data structure. All other connection properties will be refactored
in a following change, in order to keep each change as small as possible.

The data structure supports both transactional and non-transactional
connection state. Transactional state is disabled in the current version
in order to be consistent with the current behavior. It will be enabled
in a later change when all connection properties have been refactored
to use the new data structure.

Author: olavloite

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ClientSideStatementValueConverters.java

@@ -72,7 +72,11 @@ private E get(String value) {
 
   /** Converter from string to {@link Boolean} */
   static class BooleanConverter implements ClientSideStatementValueConverter<Boolean> {
+    static final BooleanConverter INSTANCE = new BooleanConverter();
 
+    private BooleanConverter() {}
+
+    /** Constructor that is needed for reflection. */
     public BooleanConverter(String allowedValues) {}
 
     @Override
@@ -141,7 +145,11 @@ public Boolean convert(String value) {
 
   /** Converter from string to a non-negative integer. */
   static class NonNegativeIntegerConverter implements ClientSideStatementValueConverter<Integer> {
+    static final NonNegativeIntegerConverter INSTANCE = new NonNegativeIntegerConverter();
+
+    private NonNegativeIntegerConverter() {}
 
+    /** Constructor needed for reflection. */
     public NonNegativeIntegerConverter(String allowedValues) {}
 
     @Override
@@ -356,9 +364,14 @@ public DirectedReadOptions convert(String value) {
   /** Converter for converting strings to {@link AutocommitDmlMode} values. */
   static class AutocommitDmlModeConverter
       implements ClientSideStatementValueConverter<AutocommitDmlMode> {
+    static final AutocommitDmlModeConverter INSTANCE = new AutocommitDmlModeConverter();
+
     private final CaseInsensitiveEnumMap<AutocommitDmlMode> values =
         new CaseInsensitiveEnumMap<>(AutocommitDmlMode.class);
 
+    private AutocommitDmlModeConverter() {}
+
+    /** Constructor needed for reflection. */
     public AutocommitDmlModeConverter(String allowedValues) {}
 
     @Override
@@ -372,7 +385,35 @@ public AutocommitDmlMode convert(String value) {
     }
   }
 
+  static class ConnectionStateTypeConverter
+      implements ClientSideStatementValueConverter<ConnectionState.Type> {
+    static final ConnectionStateTypeConverter INSTANCE = new ConnectionStateTypeConverter();
+
+    private final CaseInsensitiveEnumMap<ConnectionState.Type> values =
+        new CaseInsensitiveEnumMap<>(ConnectionState.Type.class);
+
+    private ConnectionStateTypeConverter() {}
+
+    /** Constructor that is needed for reflection. */
+    public ConnectionStateTypeConverter(String allowedValues) {}
+
+    @Override
+    public Class<ConnectionState.Type> getParameterClass() {
+      return ConnectionState.Type.class;
+    }
+
+    @Override
+    public ConnectionState.Type convert(String value) {
+      return values.get(value);
+    }
+  }
+
   static class StringValueConverter implements ClientSideStatementValueConverter<String> {
+    static final StringValueConverter INSTANCE = new StringValueConverter();
+
+    private StringValueConverter() {}
+
+    /** Constructor needed for reflection. */
     public StringValueConverter(String allowedValues) {}
 
     @Override

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java

@@ -18,6 +18,7 @@
 
 import static com.google.cloud.spanner.SpannerApiFutures.get;
 import static com.google.cloud.spanner.connection.ConnectionPreconditions.checkValidIdentifier;
+import static com.google.cloud.spanner.connection.CoreConnectionProperties.RETRY_ABORTS_INTERNALLY;
 
 import com.google.api.core.ApiFuture;
 import com.google.api.core.ApiFutures;
@@ -50,6 +51,7 @@
 import com.google.cloud.spanner.TimestampBound.Mode;
 import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
 import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType;
+import com.google.cloud.spanner.connection.ConnectionProperty.Context;
 import com.google.cloud.spanner.connection.StatementExecutor.StatementTimeout;
 import com.google.cloud.spanner.connection.StatementResult.ResultType;
 import com.google.cloud.spanner.connection.UnitOfWork.CallType;
@@ -215,6 +217,7 @@ static UnitOfWorkType of(TransactionMode transactionMode) {
   private final DdlClient ddlClient;
   private final DatabaseClient dbClient;
   private final BatchClient batchClient;
+  private final ConnectionState connectionState;
   private boolean autocommit;
   private boolean readOnly;
   private boolean returnCommitStats;
@@ -236,7 +239,6 @@ static UnitOfWorkType of(TransactionMode transactionMode) {
   private BatchMode batchMode;
   private UnitOfWorkType unitOfWorkType;
   private final Stack<UnitOfWork> transactionStack = new Stack<>();
-  private boolean retryAbortsInternally;
   private final List<TransactionRetryListener> transactionRetryListeners = new ArrayList<>();
   private AutocommitDmlMode autocommitDmlMode = AutocommitDmlMode.TRANSACTIONAL;
   private TimestampBound readOnlyStaleness = TimestampBound.strong();
@@ -307,9 +309,10 @@ static UnitOfWorkType of(TransactionMode transactionMode) {
     this.dbClient = spanner.getDatabaseClient(options.getDatabaseId());
     this.batchClient = spanner.getBatchClient(options.getDatabaseId());
     this.ddlClient = createDdlClient();
+    this.connectionState = new ConnectionState(options.getInitialConnectionPropertyValues());
 
     // (Re)set the state of the connection to the default.
-    reset();
+    reset(Context.STARTUP);
   }
 
   /** Constructor only for test purposes. */
@@ -334,6 +337,7 @@ static UnitOfWorkType of(TransactionMode transactionMode) {
     this.ddlClient = Preconditions.checkNotNull(ddlClient);
     this.dbClient = Preconditions.checkNotNull(dbClient);
     this.batchClient = Preconditions.checkNotNull(batchClient);
+    this.connectionState = new ConnectionState(options.getInitialConnectionPropertyValues());
     setReadOnly(options.isReadOnly());
     setAutocommit(options.isAutocommit());
     setReturnCommitStats(options.isReturnCommitStats());
@@ -423,14 +427,22 @@ public ApiFuture<Void> closeAsync() {
     return ApiFutures.immediateFuture(null);
   }
 
+  private Context getCurrentContext() {
+    return isInTransaction() ? Context.IN_TRANSACTION : Context.OUTSIDE_TRANSACTION;
+  }
+
   /**
    * Resets the state of this connection to the default state in the {@link ConnectionOptions} of
    * this connection.
    */
   public void reset() {
+    reset(getCurrentContext());
+  }
+
+  private void reset(Context context) {
     ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
 
-    this.retryAbortsInternally = options.isRetryAbortsInternally();
+    this.connectionState.resetValue(RETRY_ABORTS_INTERNALLY, context);
     this.readOnly = options.isReadOnly();
     this.autocommit = options.isAutocommit();
     this.queryOptions =
@@ -856,13 +868,14 @@ private void checkSetRetryAbortsInternallyAvailable() {
   @Override
   public boolean isRetryAbortsInternally() {
     ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
-    return retryAbortsInternally;
+    return this.connectionState.getValue(RETRY_ABORTS_INTERNALLY).getValue();
   }
 
   @Override
   public void setRetryAbortsInternally(boolean retryAbortsInternally) {
     checkSetRetryAbortsInternallyAvailable();
-    this.retryAbortsInternally = retryAbortsInternally;
+    this.connectionState.setValue(
+        RETRY_ABORTS_INTERNALLY, retryAbortsInternally, getCurrentContext());
   }
 
   @Override
@@ -1925,7 +1938,8 @@ UnitOfWork createNewUnitOfWork(
               .setDatabaseClient(dbClient)
               .setDelayTransactionStartUntilFirstWrite(delayTransactionStartUntilFirstWrite)
               .setKeepTransactionAlive(keepTransactionAlive)
-              .setRetryAbortsInternally(retryAbortsInternally)
+              .setRetryAbortsInternally(
+                  connectionState.getValue(RETRY_ABORTS_INTERNALLY).getValue())
               .setSavepointSupport(savepointSupport)
               .setReturnCommitStats(returnCommitStats)
               .setMaxCommitDelay(maxCommitDelay)

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionOptions.java

@@ -16,6 +16,8 @@
 
 package com.google.cloud.spanner.connection;
 
+import static com.google.cloud.spanner.connection.CoreConnectionProperties.RETRY_ABORTS_INTERNALLY;
+
 import com.google.api.core.InternalApi;
 import com.google.api.gax.core.CredentialsProvider;
 import com.google.api.gax.rpc.TransportChannelProvider;
@@ -37,6 +39,7 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Strings;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Sets;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
 import io.opentelemetry.api.OpenTelemetry;
@@ -48,6 +51,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -529,6 +533,8 @@ public interface ExternalChannelProvider {
 
   /** Builder for {@link ConnectionOptions} instances. */
   public static class Builder {
+    private final Map<String, ConnectionPropertyValue<?>> connectionPropertyValues =
+        new HashMap<>();
     private String uri;
     private String credentialsUrl;
     private String oauthToken;
@@ -626,6 +632,13 @@ public Builder setUri(String uri) {
       return this;
     }
 
+    <T> Builder setConnectionPropertyValue(
+        com.google.cloud.spanner.connection.ConnectionProperty<T> property, T value) {
+      this.connectionPropertyValues.put(
+          property.getKey(), new ConnectionPropertyValue<>(property, value, value));
+      return this;
+    }
+
     /** Sets the {@link SessionPoolOptions} to use for the connection. */
     public Builder setSessionPoolOptions(SessionPoolOptions sessionPoolOptions) {
       Preconditions.checkNotNull(sessionPoolOptions);
@@ -715,6 +728,7 @@ public static Builder newBuilder() {
     return new Builder();
   }
 
+  private final ConnectionState initialConnectionState;
   private final String uri;
   private final String warnings;
   private final String credentialsUrl;
@@ -756,7 +770,6 @@ public static Builder newBuilder() {
   private final boolean autocommit;
   private final boolean readOnly;
   private final boolean routeToLeader;
-  private final boolean retryAbortsInternally;
   private final boolean useVirtualThreads;
   private final boolean useVirtualGrpcTransportThreads;
   private final OpenTelemetry openTelemetry;
@@ -773,6 +786,11 @@ private ConnectionOptions(Builder builder) {
     this.warnings = checkValidProperties(builder.uri);
 
     this.uri = builder.uri;
+    ImmutableMap<String, ConnectionPropertyValue<?>> connectionPropertyValues =
+        ImmutableMap.<String, ConnectionPropertyValue<?>>builder()
+            .putAll(CoreConnectionProperties.parseValues(builder.uri))
+            .putAll(builder.connectionPropertyValues)
+            .buildKeepingLast();
     this.credentialsUrl =
         builder.credentialsUrl != null ? builder.credentialsUrl : parseCredentials(builder.uri);
     this.encodedCredentials = parseEncodedCredentials(builder.uri);
@@ -866,7 +884,6 @@ private ConnectionOptions(Builder builder) {
     this.autocommit = parseAutocommit(this.uri);
     this.readOnly = parseReadOnly(this.uri);
     this.routeToLeader = parseRouteToLeader(this.uri);
-    this.retryAbortsInternally = parseRetryAbortsInternally(this.uri);
     this.useVirtualThreads = parseUseVirtualThreads(this.uri);
     this.useVirtualGrpcTransportThreads = parseUseVirtualGrpcTransportThreads(this.uri);
     this.openTelemetry = builder.openTelemetry;
@@ -896,6 +913,7 @@ private ConnectionOptions(Builder builder) {
     } else {
       this.sessionPoolOptions = SessionPoolOptions.newBuilder().setAutoDetectDialect(true).build();
     }
+    this.initialConnectionState = new ConnectionState(connectionPropertyValues);
   }
 
   @VisibleForTesting
@@ -993,12 +1011,6 @@ static boolean parseRouteToLeader(String uri) {
     return value != null ? Boolean.parseBoolean(value) : DEFAULT_ROUTE_TO_LEADER;
   }
 
-  @VisibleForTesting
-  static boolean parseRetryAbortsInternally(String uri) {
-    String value = parseUriProperty(uri, RETRY_ABORTS_INTERNALLY_PROPERTY_NAME);
-    return value != null ? Boolean.parseBoolean(value) : DEFAULT_RETRY_ABORTS_INTERNALLY;
-  }
-
   @VisibleForTesting
   static boolean parseUseVirtualThreads(String uri) {
     String value = parseUriProperty(uri, USE_VIRTUAL_THREADS_PROPERTY_NAME);
@@ -1357,6 +1369,11 @@ public String getUri() {
     return uri;
   }
 
+  /** The connection properties that have been pre-set for this {@link ConnectionOptions}. */
+  Map<String, ConnectionPropertyValue<?>> getInitialConnectionPropertyValues() {
+    return this.initialConnectionState.getAllValues();
+  }
+
   /** The credentials URL of this {@link ConnectionOptions} */
   public String getCredentialsUrl() {
     return credentialsUrl;
@@ -1488,7 +1505,7 @@ public boolean isRouteToLeader() {
    * ConnectionOptions}
    */
   public boolean isRetryAbortsInternally() {
-    return retryAbortsInternally;
+    return this.initialConnectionState.getValue(RETRY_ABORTS_INTERNALLY).getValue();
   }
 
   /** Whether connections should use virtual threads for connection executors. */