Skip to content

Commit e1476e0

Browse files
pratickchokhaniolavloite
authored and
surbhigarg92
committed
feat(spanner): support multiplexed session for Partitioned operations (googleapis#3231)
* feat(spanner): support multiplexed session for Partitioned read or query. * chore(spanner): lint fixes * feat(spanner): support multiplexed session for Partitioned DML operations. * lint(spanner): javadoc fixes. * feat(spanner): Updated unit tests of Partitioned operations for Multiplexed Session. * feat(spanner): Updated unit tests of Partitioned operations for Multiplexed Session. * lint(spanner): Apply suggestions from code review Co-authored-by: Knut Olav Løite <[email protected]> * lint(spanner): Apply suggestions from code review Co-authored-by: Knut Olav Løite <[email protected]> * feat(spanner): Modified BatchClientImpl to store multiplexed session and create fresh session after expiration date. * feat(spanner): Removed env variable for Partitioned Ops ensuring that Multiplexed Session for Partitioned Ops is not available to customers. * lint(spanner): Removed unused variables. --------- Co-authored-by: Knut Olav Løite <[email protected]>
1 parent 02e8b2f commit e1476e0

16 files changed

+230
-49
lines changed

‎.github/workflows/ci.yaml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ jobs:
5353
env:
5454
JOB_TYPE: test
5555
GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS: true
56+
GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS: true
5657
units-java8:
5758
# Building using Java 17 and run the tests with Java 8 runtime
5859
name: "units (8)"
@@ -92,6 +93,7 @@ jobs:
9293
env:
9394
JOB_TYPE: test
9495
GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS: true
96+
GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS: true
9597
windows:
9698
runs-on: windows-latest
9799
steps:

‎.kokoro/presubmit/integration-multiplexed-sessions-enabled.cfg‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,3 +36,8 @@ env_vars: {
3636
key: "GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS"
3737
value: "true"
3838
}
39+
40+
env_vars: {
41+
key: "GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS"
42+
value: "true"
43+
}

‎google-cloud-spanner-executor/src/main/java/com/google/cloud/spanner/SessionPoolOptionsHelper.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,4 +44,12 @@ public static SessionPoolOptions.Builder setUseMultiplexedSessionForRW(
4444
SessionPoolOptions.Builder sessionPoolOptionsBuilder, boolean useMultiplexedSessionForRW) {
4545
return sessionPoolOptionsBuilder.setUseMultiplexedSessionForRW(useMultiplexedSessionForRW);
4646
}
47+
48+
// TODO: Remove when multiplexed session for partitioned operations are released.
49+
public static SessionPoolOptions.Builder setUseMultiplexedSessionForPartitionedOperations(
50+
SessionPoolOptions.Builder sessionPoolOptionsBuilder,
51+
boolean useMultiplexedSessionForPartitionedOps) {
52+
return sessionPoolOptionsBuilder.setUseMultiplexedSessionPartitionedOps(
53+
useMultiplexedSessionForPartitionedOps);
54+
}
4755
}

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractMultiplexedSessionDatabaseClient.java‎

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@
1919
import com.google.api.gax.rpc.ServerStream;
2020
import com.google.cloud.Timestamp;
2121
import com.google.cloud.spanner.Options.TransactionOption;
22-
import com.google.cloud.spanner.Options.UpdateOption;
2322
import com.google.spanner.v1.BatchWriteResponse;
2423

2524
/**
@@ -51,9 +50,4 @@ public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
5150
throws SpannerException {
5251
throw new UnsupportedOperationException();
5352
}
54-
55-
@Override
56-
public long executePartitionedUpdate(Statement stmt, UpdateOption... options) {
57-
throw new UnsupportedOperationException();
58-
}
5953
}

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java‎

Lines changed: 55 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,16 +30,50 @@
3030
import com.google.spanner.v1.PartitionReadRequest;
3131
import com.google.spanner.v1.PartitionResponse;
3232
import com.google.spanner.v1.TransactionSelector;
33+
import java.time.Clock;
34+
import java.time.Duration;
35+
import java.time.Instant;
3336
import java.util.List;
3437
import java.util.Map;
38+
import java.util.concurrent.atomic.AtomicReference;
39+
import java.util.concurrent.locks.ReentrantLock;
3540
import javax.annotation.Nullable;
41+
import javax.annotation.concurrent.GuardedBy;
3642

3743
/** Default implementation for Batch Client interface. */
3844
public class BatchClientImpl implements BatchClient {
3945
private final SessionClient sessionClient;
4046

41-
BatchClientImpl(SessionClient sessionClient) {
47+
private final boolean isMultiplexedSessionEnabled;
48+
49+
/** Lock to protect the multiplexed session. */
50+
private final ReentrantLock multiplexedSessionLock = new ReentrantLock();
51+
52+
/** The duration before we try to replace the multiplexed session. The default is 7 days. */
53+
private final Duration sessionExpirationDuration;
54+
55+
/** The expiration date/time of the current multiplexed session. */
56+
@GuardedBy("multiplexedSessionLock")
57+
private final AtomicReference<Instant> expirationDate;
58+
59+
@GuardedBy("multiplexedSessionLock")
60+
private final AtomicReference<SessionImpl> multiplexedSessionReference;
61+
62+
BatchClientImpl(SessionClient sessionClient, boolean isMultiplexedSessionEnabled) {
4263
this.sessionClient = checkNotNull(sessionClient);
64+
this.isMultiplexedSessionEnabled = isMultiplexedSessionEnabled;
65+
this.sessionExpirationDuration =
66+
Duration.ofMillis(
67+
sessionClient
68+
.getSpanner()
69+
.getOptions()
70+
.getSessionPoolOptions()
71+
.getMultiplexedSessionMaintenanceDuration()
72+
.toMillis());
73+
// Initialize the expiration date to the start of time to avoid unnecessary null checks.
74+
// This also ensured that a new session is created on first request.
75+
this.expirationDate = new AtomicReference<>(Instant.MIN);
76+
this.multiplexedSessionReference = new AtomicReference<>();
4377
}
4478

4579
@Override
@@ -50,7 +84,12 @@ public String getDatabaseRole() {
5084

5185
@Override
5286
public BatchReadOnlyTransaction batchReadOnlyTransaction(TimestampBound bound) {
53-
SessionImpl session = sessionClient.createSession();
87+
SessionImpl session;
88+
if (isMultiplexedSessionEnabled) {
89+
session = getMultiplexedSession();
90+
} else {
91+
session = sessionClient.createSession();
92+
}
5493
return new BatchReadOnlyTransactionImpl(
5594
MultiUseReadOnlyTransaction.newBuilder()
5695
.setSession(session)
@@ -92,6 +131,20 @@ public BatchReadOnlyTransaction batchReadOnlyTransaction(BatchTransactionId batc
92131
batchTransactionId);
93132
}
94133

134+
private SessionImpl getMultiplexedSession() {
135+
this.multiplexedSessionLock.lock();
136+
try {
137+
if (Clock.systemUTC().instant().isAfter(this.expirationDate.get())
138+
|| this.multiplexedSessionReference.get() == null) {
139+
this.multiplexedSessionReference.set(this.sessionClient.createMultiplexedSession());
140+
this.expirationDate.set(Clock.systemUTC().instant().plus(this.sessionExpirationDuration));
141+
}
142+
return this.multiplexedSessionReference.get();
143+
} finally {
144+
this.multiplexedSessionLock.unlock();
145+
}
146+
}
147+
95148
private static class BatchReadOnlyTransactionImpl extends MultiUseReadOnlyTransaction
96149
implements BatchReadOnlyTransaction {
97150
private final String sessionName;

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ class DatabaseClientImpl implements DatabaseClient {
3636
@VisibleForTesting final String clientId;
3737
@VisibleForTesting final SessionPool pool;
3838
@VisibleForTesting final MultiplexedSessionDatabaseClient multiplexedSessionDatabaseClient;
39+
@VisibleForTesting final boolean useMultiplexedSessionPartitionedOps;
3940
@VisibleForTesting final boolean useMultiplexedSessionForRW;
4041

4142
final boolean useMultiplexedSessionBlindWrite;
@@ -47,6 +48,7 @@ class DatabaseClientImpl implements DatabaseClient {
4748
pool,
4849
/* useMultiplexedSessionBlindWrite = */ false,
4950
/* multiplexedSessionDatabaseClient = */ null,
51+
/* useMultiplexedSessionPartitionedOps= */ false,
5052
tracer,
5153
/* useMultiplexedSessionForRW = */ false);
5254
}
@@ -58,6 +60,7 @@ class DatabaseClientImpl implements DatabaseClient {
5860
pool,
5961
/* useMultiplexedSessionBlindWrite = */ false,
6062
/* multiplexedSessionDatabaseClient = */ null,
63+
/* useMultiplexedSessionPartitionedOps= */ false,
6164
tracer,
6265
/* useMultiplexedSessionForRW = */ false);
6366
}
@@ -67,12 +70,14 @@ class DatabaseClientImpl implements DatabaseClient {
6770
SessionPool pool,
6871
boolean useMultiplexedSessionBlindWrite,
6972
@Nullable MultiplexedSessionDatabaseClient multiplexedSessionDatabaseClient,
73+
boolean useMultiplexedSessionPartitionedOps,
7074
TraceWrapper tracer,
7175
boolean useMultiplexedSessionForRW) {
7276
this.clientId = clientId;
7377
this.pool = pool;
7478
this.useMultiplexedSessionBlindWrite = useMultiplexedSessionBlindWrite;
7579
this.multiplexedSessionDatabaseClient = multiplexedSessionDatabaseClient;
80+
this.useMultiplexedSessionPartitionedOps = useMultiplexedSessionPartitionedOps;
7681
this.tracer = tracer;
7782
this.useMultiplexedSessionForRW = useMultiplexedSessionForRW;
7883
}
@@ -309,6 +314,14 @@ public AsyncTransactionManager transactionManagerAsync(TransactionOption... opti
309314

310315
@Override
311316
public long executePartitionedUpdate(final Statement stmt, final UpdateOption... options) {
317+
if (useMultiplexedSessionPartitionedOps) {
318+
return getMultiplexedSession().executePartitionedUpdate(stmt, options);
319+
}
320+
return executePartitionedUpdateWithPooledSession(stmt, options);
321+
}
322+
323+
private long executePartitionedUpdateWithPooledSession(
324+
final Statement stmt, final UpdateOption... options) {
312325
ISpan span = tracer.spanBuilder(PARTITION_DML_TRANSACTION);
313326
try (IScope s = tracer.withSpan(span)) {
314327
return runWithSessionRetry(session -> session.executePartitionedUpdate(stmt, options));

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedMultiplexedSessionTransaction.java‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import com.google.cloud.spanner.DelayedReadContext.DelayedReadOnlyTransaction;
2525
import com.google.cloud.spanner.MultiplexedSessionDatabaseClient.MultiplexedSessionTransaction;
2626
import com.google.cloud.spanner.Options.TransactionOption;
27+
import com.google.cloud.spanner.Options.UpdateOption;
2728
import com.google.common.util.concurrent.MoreExecutors;
2829
import java.util.concurrent.ExecutionException;
2930

@@ -224,4 +225,16 @@ private SessionReference getSessionReference() {
224225
throw SpannerExceptionFactory.propagateInterrupt(interruptedException);
225226
}
226227
}
228+
229+
/**
230+
* Execute `stmt` within PARTITIONED_DML transaction using multiplexed session. This method is a
231+
* blocking call as the interface expects to return the output of the `stmt`.
232+
*/
233+
@Override
234+
public long executePartitionedUpdate(Statement stmt, UpdateOption... options) {
235+
SessionReference sessionReference = getSessionReference();
236+
return new MultiplexedSessionTransaction(
237+
client, span, sessionReference, NO_CHANNEL_HINT, /* singleUse = */ true)
238+
.executePartitionedUpdate(stmt, options);
239+
}
227240
}

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import com.google.api.core.SettableApiFuture;
2525
import com.google.cloud.Timestamp;
2626
import com.google.cloud.spanner.Options.TransactionOption;
27+
import com.google.cloud.spanner.Options.UpdateOption;
2728
import com.google.cloud.spanner.SessionClient.SessionConsumer;
2829
import com.google.cloud.spanner.SpannerException.ResourceNotFoundException;
2930
import com.google.common.annotations.VisibleForTesting;
@@ -553,6 +554,12 @@ public AsyncTransactionManager transactionManagerAsync(TransactionOption... opti
553554
.transactionManagerAsync(options);
554555
}
555556

557+
@Override
558+
public long executePartitionedUpdate(Statement stmt, UpdateOption... options) {
559+
return createMultiplexedSessionTransaction(/* singleUse = */ true)
560+
.executePartitionedUpdate(stmt, options);
561+
}
562+
556563
/**
557564
* It is enough with one executor to maintain the multiplexed sessions in all the clients, as they
558565
* do not need to be updated often, and the maintenance task is light. The core pool size is set

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java‎

Lines changed: 18 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -237,6 +237,19 @@ SessionImpl createSession() {
237237
* @param consumer The {@link SessionConsumer} to use for callbacks when sessions are available.
238238
*/
239239
void createMultiplexedSession(SessionConsumer consumer) {
240+
try {
241+
SessionImpl sessionImpl = createMultiplexedSession();
242+
consumer.onSessionReady(sessionImpl);
243+
} catch (Throwable t) {
244+
consumer.onSessionCreateFailure(t, 1);
245+
}
246+
}
247+
248+
/**
249+
* Creates a multiplexed session and returns it. A multiplexed session is not affiliated with any
250+
* GRPC channel. In case of an error during the gRPC calls, an exception will be thrown.
251+
*/
252+
SessionImpl createMultiplexedSession() {
240253
ISpan span = spanner.getTracer().spanBuilder(SpannerImpl.CREATE_MULTIPLEXED_SESSION);
241254
try (IScope s = spanner.getTracer().withSpan(span)) {
242255
com.google.spanner.v1.Session session =
@@ -253,10 +266,12 @@ void createMultiplexedSession(SessionConsumer consumer) {
253266
spanner,
254267
new SessionReference(
255268
session.getName(), session.getCreateTime(), session.getMultiplexed(), null));
256-
consumer.onSessionReady(sessionImpl);
269+
span.addAnnotation(
270+
String.format("Request for %d multiplexed session returned %d session", 1, 1));
271+
return sessionImpl;
257272
} catch (Throwable t) {
258273
span.setStatus(t);
259-
consumer.onSessionCreateFailure(t, 1);
274+
throw t;
260275
} finally {
261276
span.end();
262277
}
@@ -289,31 +304,7 @@ private CreateMultiplexedSessionsRunnable(SessionConsumer consumer) {
289304

290305
@Override
291306
public void run() {
292-
ISpan span = spanner.getTracer().spanBuilder(SpannerImpl.CREATE_MULTIPLEXED_SESSION);
293-
try (IScope s = spanner.getTracer().withSpan(span)) {
294-
com.google.spanner.v1.Session session =
295-
spanner
296-
.getRpc()
297-
.createSession(
298-
db.getName(),
299-
spanner.getOptions().getDatabaseRole(),
300-
spanner.getOptions().getSessionLabels(),
301-
null,
302-
true);
303-
SessionImpl sessionImpl =
304-
new SessionImpl(
305-
spanner,
306-
new SessionReference(
307-
session.getName(), session.getCreateTime(), session.getMultiplexed(), null));
308-
span.addAnnotation(
309-
String.format("Request for %d multiplexed session returned %d session", 1, 1));
310-
consumer.onSessionReady(sessionImpl);
311-
} catch (Throwable t) {
312-
span.setStatus(t);
313-
consumer.onSessionCreateFailure(t, 1);
314-
} finally {
315-
span.end();
316-
}
307+
createMultiplexedSession(consumer);
317308
}
318309
}
319310

0 commit comments

Comments
 (0)