Skip to content

Commit f4c6768

Browse files
odeke-emodeke-em
committed
chore(x-goog-spanner-request-id): plumb for BatchCreateSessions
This change plumbs x-goog-spanner-request-id into BatchCreateSessions and asserts that the header is present for that method. Updates #3537
1 parent e97b92e commit f4c6768

12 files changed

+560
-33
lines changed

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

+74-7
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,14 @@
2323
import com.google.cloud.spanner.SessionPool.PooledSessionFuture;
2424
import com.google.cloud.spanner.SpannerImpl.ClosedException;
2525
import com.google.common.annotations.VisibleForTesting;
26-
import com.google.common.base.Function;
2726
import com.google.common.util.concurrent.ListenableFuture;
2827
import com.google.spanner.v1.BatchWriteResponse;
2928
import io.opentelemetry.api.common.Attributes;
29+
import java.util.ArrayList;
30+
import java.util.Arrays;
31+
import java.util.Objects;
32+
import java.util.concurrent.atomic.AtomicInteger;
33+
import java.util.function.BiFunction;
3034
import javax.annotation.Nullable;
3135

3236
class DatabaseClientImpl implements DatabaseClient {
@@ -40,6 +44,8 @@ class DatabaseClientImpl implements DatabaseClient {
4044
@VisibleForTesting final MultiplexedSessionDatabaseClient multiplexedSessionDatabaseClient;
4145
@VisibleForTesting final boolean useMultiplexedSessionPartitionedOps;
4246
@VisibleForTesting final boolean useMultiplexedSessionForRW;
47+
private final int dbId;
48+
private final AtomicInteger nthRequest;
4349

4450
final boolean useMultiplexedSessionBlindWrite;
4551

@@ -86,6 +92,18 @@ class DatabaseClientImpl implements DatabaseClient {
8692
this.tracer = tracer;
8793
this.useMultiplexedSessionForRW = useMultiplexedSessionForRW;
8894
this.commonAttributes = commonAttributes;
95+
96+
this.dbId = this.dbIdFromClientId(this.clientId);
97+
this.nthRequest = new AtomicInteger(0);
98+
}
99+
100+
private int dbIdFromClientId(String clientId) {
101+
int i = clientId.indexOf("-");
102+
String strWithValue = clientId.substring(i + 1);
103+
if (Objects.equals(strWithValue, "")) {
104+
strWithValue = "0";
105+
}
106+
return Integer.parseInt(strWithValue);
89107
}
90108

91109
@VisibleForTesting
@@ -159,7 +177,11 @@ public CommitResponse writeWithOptions(
159177
if (canUseMultiplexedSessionsForRW() && getMultiplexedSessionDatabaseClient() != null) {
160178
return getMultiplexedSessionDatabaseClient().writeWithOptions(mutations, options);
161179
}
162-
return runWithSessionRetry(session -> session.writeWithOptions(mutations, options));
180+
181+
return runWithSessionRetry(
182+
(session, reqId) -> {
183+
return session.writeWithOptions(mutations, withReqId(reqId, options));
184+
});
163185
} catch (RuntimeException e) {
164186
span.setStatus(e);
165187
throw e;
@@ -177,14 +199,23 @@ public Timestamp writeAtLeastOnce(final Iterable<Mutation> mutations) throws Spa
177199
public CommitResponse writeAtLeastOnceWithOptions(
178200
final Iterable<Mutation> mutations, final TransactionOption... options)
179201
throws SpannerException {
202+
return doWriteAtLeastOnceWithOptions(mutations, options);
203+
}
204+
205+
private CommitResponse doWriteAtLeastOnceWithOptions(
206+
final Iterable<Mutation> mutations, final TransactionOption... options)
207+
throws SpannerException {
180208
ISpan span = tracer.spanBuilder(READ_WRITE_TRANSACTION, commonAttributes, options);
181209
try (IScope s = tracer.withSpan(span)) {
182210
if (useMultiplexedSessionBlindWrite && getMultiplexedSessionDatabaseClient() != null) {
183211
return getMultiplexedSessionDatabaseClient()
184212
.writeAtLeastOnceWithOptions(mutations, options);
185213
}
214+
186215
return runWithSessionRetry(
187-
session -> session.writeAtLeastOnceWithOptions(mutations, options));
216+
(session, reqId) -> {
217+
return session.writeAtLeastOnceWithOptions(mutations, withReqId(reqId, options));
218+
});
188219
} catch (RuntimeException e) {
189220
span.setStatus(e);
190221
throw e;
@@ -193,6 +224,10 @@ public CommitResponse writeAtLeastOnceWithOptions(
193224
}
194225
}
195226

227+
private int nextNthRequest() {
228+
return this.nthRequest.incrementAndGet();
229+
}
230+
196231
@Override
197232
public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
198233
final Iterable<MutationGroup> mutationGroups, final TransactionOption... options)
@@ -202,7 +237,9 @@ public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
202237
if (canUseMultiplexedSessionsForRW() && getMultiplexedSessionDatabaseClient() != null) {
203238
return getMultiplexedSessionDatabaseClient().batchWriteAtLeastOnce(mutationGroups, options);
204239
}
205-
return runWithSessionRetry(session -> session.batchWriteAtLeastOnce(mutationGroups, options));
240+
return runWithSessionRetry(
241+
(session, reqId) ->
242+
session.batchWriteAtLeastOnce(mutationGroups, withReqId(reqId, options)));
206243
} catch (RuntimeException e) {
207244
span.setStatus(e);
208245
throw e;
@@ -346,27 +383,57 @@ public long executePartitionedUpdate(final Statement stmt, final UpdateOption...
346383
return executePartitionedUpdateWithPooledSession(stmt, options);
347384
}
348385

386+
private UpdateOption[] withReqId(
387+
final XGoogSpannerRequestId reqId, final UpdateOption... options) {
388+
if (reqId == null) {
389+
return options;
390+
}
391+
ArrayList<UpdateOption> allOptions = new ArrayList(Arrays.asList(options));
392+
allOptions.add(new Options.RequestIdOption(reqId));
393+
return allOptions.toArray(new UpdateOption[0]);
394+
}
395+
396+
private TransactionOption[] withReqId(
397+
final XGoogSpannerRequestId reqId, final TransactionOption... options) {
398+
if (reqId == null) {
399+
return options;
400+
}
401+
ArrayList<TransactionOption> allOptions = new ArrayList(Arrays.asList(options));
402+
allOptions.add(new Options.RequestIdOption(reqId));
403+
return allOptions.toArray(new TransactionOption[0]);
404+
}
405+
349406
private long executePartitionedUpdateWithPooledSession(
350407
final Statement stmt, final UpdateOption... options) {
351408
ISpan span = tracer.spanBuilder(PARTITION_DML_TRANSACTION, commonAttributes);
352409
try (IScope s = tracer.withSpan(span)) {
353-
return runWithSessionRetry(session -> session.executePartitionedUpdate(stmt, options));
410+
return runWithSessionRetry(
411+
(session, reqId) -> {
412+
return session.executePartitionedUpdate(stmt, withReqId(reqId, options));
413+
});
354414
} catch (RuntimeException e) {
355415
span.setStatus(e);
356416
span.end();
357417
throw e;
358418
}
359419
}
360420

361-
private <T> T runWithSessionRetry(Function<Session, T> callable) {
421+
private <T> T runWithSessionRetry(BiFunction<Session, XGoogSpannerRequestId, T> callable) {
362422
PooledSessionFuture session = getSession();
423+
XGoogSpannerRequestId reqId =
424+
XGoogSpannerRequestId.of(
425+
this.dbId, Long.valueOf(session.getChannel()), this.nextNthRequest(), 0);
363426
while (true) {
364427
try {
365-
return callable.apply(session);
428+
reqId.incrementAttempt();
429+
return callable.apply(session, reqId);
366430
} catch (SessionNotFoundException e) {
367431
session =
368432
(PooledSessionFuture)
369433
pool.getPooledSessionReplacementHandler().replaceSession(e, session);
434+
reqId =
435+
XGoogSpannerRequestId.of(
436+
this.dbId, Long.valueOf(session.getChannel()), this.nextNthRequest(), 0);
370437
}
371438
}
372439
}

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

+47-1
Original file line numberDiff line numberDiff line change
@@ -177,6 +177,10 @@ public static UpdateTransactionOption excludeTxnFromChangeStreams() {
177177
return EXCLUDE_TXN_FROM_CHANGE_STREAMS_OPTION;
178178
}
179179

180+
public static RequestIdOption requestId(XGoogSpannerRequestId reqId) {
181+
return new RequestIdOption(reqId);
182+
}
183+
180184
/**
181185
* Specifying this will cause the read to yield at most this many rows. This should be greater
182186
* than 0.
@@ -535,6 +539,7 @@ void appendToOptions(Options options) {
535539
private RpcLockHint lockHint;
536540
private Boolean lastStatement;
537541
private IsolationLevel isolationLevel;
542+
private XGoogSpannerRequestId reqId;
538543

539544
// Construction is via factory methods below.
540545
private Options() {}
@@ -599,6 +604,14 @@ String filter() {
599604
return filter;
600605
}
601606

607+
boolean hasReqId() {
608+
return reqId != null;
609+
}
610+
611+
XGoogSpannerRequestId reqId() {
612+
return reqId;
613+
}
614+
602615
boolean hasPriority() {
603616
return priority != null;
604617
}
@@ -756,6 +769,9 @@ public String toString() {
756769
if (isolationLevel != null) {
757770
b.append("isolationLevel: ").append(isolationLevel).append(' ');
758771
}
772+
if (reqId != null) {
773+
b.append("requestId: ").append(reqId.toString());
774+
}
759775
return b.toString();
760776
}
761777

@@ -798,7 +814,8 @@ public boolean equals(Object o) {
798814
&& Objects.equals(orderBy(), that.orderBy())
799815
&& Objects.equals(isLastStatement(), that.isLastStatement())
800816
&& Objects.equals(lockHint(), that.lockHint())
801-
&& Objects.equals(isolationLevel(), that.isolationLevel());
817+
&& Objects.equals(isolationLevel(), that.isolationLevel())
818+
&& Objects.equals(reqId(), that.reqId());
802819
}
803820

804821
@Override
@@ -867,6 +884,9 @@ public int hashCode() {
867884
if (isolationLevel != null) {
868885
result = 31 * result + isolationLevel.hashCode();
869886
}
887+
if (reqId != null) {
888+
result = 31 * result + reqId.hashCode();
889+
}
870890
return result;
871891
}
872892

@@ -1052,4 +1072,30 @@ public boolean equals(Object o) {
10521072
return o instanceof LastStatementUpdateOption;
10531073
}
10541074
}
1075+
1076+
static final class RequestIdOption extends InternalOption
1077+
implements TransactionOption, UpdateOption {
1078+
private final XGoogSpannerRequestId reqId;
1079+
1080+
RequestIdOption(XGoogSpannerRequestId reqId) {
1081+
this.reqId = reqId;
1082+
}
1083+
1084+
@Override
1085+
void appendToOptions(Options options) {
1086+
options.reqId = this.reqId;
1087+
}
1088+
1089+
@Override
1090+
public int hashCode() {
1091+
return RequestIdOption.class.hashCode();
1092+
}
1093+
1094+
@Override
1095+
public boolean equals(Object o) {
1096+
// TODO: Examine why the precedent for LastStatementUpdateOption
1097+
// does not check against the actual value.
1098+
return o instanceof RequestIdOption;
1099+
}
1100+
}
10551101
}

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

+33-10
Original file line numberDiff line numberDiff line change
@@ -31,10 +31,11 @@
3131
import java.util.List;
3232
import java.util.Map;
3333
import java.util.concurrent.ScheduledExecutorService;
34+
import java.util.concurrent.atomic.AtomicInteger;
3435
import javax.annotation.concurrent.GuardedBy;
3536

3637
/** Client for creating single sessions and batches of sessions. */
37-
class SessionClient implements AutoCloseable {
38+
class SessionClient implements AutoCloseable, XGoogSpannerRequestId.RequestIdCreator {
3839
static class SessionId {
3940
private static final PathTemplate NAME_TEMPLATE =
4041
PathTemplate.create(
@@ -174,6 +175,12 @@ interface SessionConsumer {
174175
private final DatabaseId db;
175176
private final Attributes commonAttributes;
176177

178+
// SessionClient is created long before a DatabaseClientImpl is created,
179+
// as batch sessions are firstly created then later attached to each Client.
180+
private static AtomicInteger NTH_ID = new AtomicInteger(0);
181+
private final int nthId;
182+
private final AtomicInteger nthRequest;
183+
177184
@GuardedBy("this")
178185
private volatile long sessionChannelCounter;
179186

@@ -186,6 +193,8 @@ interface SessionConsumer {
186193
this.executorFactory = executorFactory;
187194
this.executor = executorFactory.get();
188195
this.commonAttributes = spanner.getTracer().createCommonAttributes(db);
196+
this.nthId = SessionClient.NTH_ID.incrementAndGet();
197+
this.nthRequest = new AtomicInteger(0);
189198
}
190199

191200
@Override
@@ -201,28 +210,38 @@ DatabaseId getDatabaseId() {
201210
return db;
202211
}
203212

213+
@Override
214+
public XGoogSpannerRequestId nextRequestId(long channelId, int attempt) {
215+
return XGoogSpannerRequestId.of(this.nthId, this.nthRequest.incrementAndGet(), channelId, 1);
216+
}
217+
204218
/** Create a single session. */
205219
SessionImpl createSession() {
206220
// The sessionChannelCounter could overflow, but that will just flip it to Integer.MIN_VALUE,
207221
// which is also a valid channel hint.
208222
final Map<SpannerRpc.Option, ?> options;
223+
final long channelId;
209224
synchronized (this) {
210225
options = optionMap(SessionOption.channelHint(sessionChannelCounter++));
226+
channelId = sessionChannelCounter;
211227
}
212228
ISpan span = spanner.getTracer().spanBuilder(SpannerImpl.CREATE_SESSION, this.commonAttributes);
213229
try (IScope s = spanner.getTracer().withSpan(span)) {
230+
XGoogSpannerRequestId reqId = this.nextRequestId(channelId, 1);
214231
com.google.spanner.v1.Session session =
215232
spanner
216233
.getRpc()
217234
.createSession(
218235
db.getName(),
219236
spanner.getOptions().getDatabaseRole(),
220237
spanner.getOptions().getSessionLabels(),
221-
options);
238+
reqId.withOptions(options));
222239
SessionReference sessionReference =
223240
new SessionReference(
224241
session.getName(), session.getCreateTime(), session.getMultiplexed(), options);
225-
return new SessionImpl(spanner, sessionReference);
242+
SessionImpl sessionImpl = new SessionImpl(spanner, sessionReference);
243+
sessionImpl.setRequestIdCreator(this);
244+
return sessionImpl;
226245
} catch (RuntimeException e) {
227246
span.setStatus(e);
228247
throw e;
@@ -273,6 +292,7 @@ SessionImpl createMultiplexedSession() {
273292
spanner,
274293
new SessionReference(
275294
session.getName(), session.getCreateTime(), session.getMultiplexed(), null));
295+
sessionImpl.setRequestIdCreator(this);
276296
span.addAnnotation(
277297
String.format("Request for %d multiplexed session returned %d session", 1, 1));
278298
return sessionImpl;
@@ -387,6 +407,8 @@ private List<SessionImpl> internalBatchCreateSessions(
387407
.spanBuilderWithExplicitParent(SpannerImpl.BATCH_CREATE_SESSIONS_REQUEST, parent);
388408
span.addAnnotation(String.format("Requesting %d sessions", sessionCount));
389409
try (IScope s = spanner.getTracer().withSpan(span)) {
410+
XGoogSpannerRequestId reqId =
411+
XGoogSpannerRequestId.of(this.nthId, this.nthRequest.incrementAndGet(), channelHint, 1);
390412
List<com.google.spanner.v1.Session> sessions =
391413
spanner
392414
.getRpc()
@@ -395,21 +417,20 @@ private List<SessionImpl> internalBatchCreateSessions(
395417
sessionCount,
396418
spanner.getOptions().getDatabaseRole(),
397419
spanner.getOptions().getSessionLabels(),
398-
options);
420+
reqId.withOptions(options));
399421
span.addAnnotation(
400422
String.format(
401423
"Request for %d sessions returned %d sessions", sessionCount, sessions.size()));
402424
span.end();
403425
List<SessionImpl> res = new ArrayList<>(sessionCount);
404426
for (com.google.spanner.v1.Session session : sessions) {
405-
res.add(
427+
SessionImpl sessionImpl =
406428
new SessionImpl(
407429
spanner,
408430
new SessionReference(
409-
session.getName(),
410-
session.getCreateTime(),
411-
session.getMultiplexed(),
412-
options)));
431+
session.getName(), session.getCreateTime(), session.getMultiplexed(), options));
432+
sessionImpl.setRequestIdCreator(this);
433+
res.add(sessionImpl);
413434
}
414435
return res;
415436
} catch (RuntimeException e) {
@@ -425,6 +446,8 @@ SessionImpl sessionWithId(String name) {
425446
synchronized (this) {
426447
options = optionMap(SessionOption.channelHint(sessionChannelCounter++));
427448
}
428-
return new SessionImpl(spanner, new SessionReference(name, options));
449+
SessionImpl sessionImpl = new SessionImpl(spanner, new SessionReference(name, options));
450+
sessionImpl.setRequestIdCreator(this);
451+
return sessionImpl;
429452
}
430453
}

0 commit comments

Comments
 (0)