Skip to content

Commit e7463fb

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 1afd815 commit e7463fb

13 files changed

+555
-33
lines changed

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

+66-7
Original file line numberDiff line numberDiff line change
@@ -24,14 +24,18 @@
2424
import com.google.cloud.spanner.SpannerImpl.ClosedException;
2525
import com.google.cloud.spanner.Statement.StatementFactory;
2626
import com.google.common.annotations.VisibleForTesting;
27-
import com.google.common.base.Function;
2827
import com.google.common.util.concurrent.ListenableFuture;
2928
import com.google.spanner.v1.BatchWriteResponse;
3029
import io.opentelemetry.api.common.Attributes;
30+
import java.util.ArrayList;
31+
import java.util.Arrays;
32+
import java.util.Objects;
3133
import java.util.concurrent.ExecutionException;
3234
import java.util.concurrent.Future;
3335
import java.util.concurrent.TimeUnit;
3436
import java.util.concurrent.TimeoutException;
37+
import java.util.concurrent.atomic.AtomicInteger;
38+
import java.util.function.BiFunction;
3539
import javax.annotation.Nullable;
3640

3741
class DatabaseClientImpl implements DatabaseClient {
@@ -45,6 +49,8 @@ class DatabaseClientImpl implements DatabaseClient {
4549
@VisibleForTesting final MultiplexedSessionDatabaseClient multiplexedSessionDatabaseClient;
4650
@VisibleForTesting final boolean useMultiplexedSessionPartitionedOps;
4751
@VisibleForTesting final boolean useMultiplexedSessionForRW;
52+
private final int dbId;
53+
private final AtomicInteger nthRequest;
4854

4955
final boolean useMultiplexedSessionBlindWrite;
5056

@@ -91,6 +97,18 @@ class DatabaseClientImpl implements DatabaseClient {
9197
this.tracer = tracer;
9298
this.useMultiplexedSessionForRW = useMultiplexedSessionForRW;
9399
this.commonAttributes = commonAttributes;
100+
101+
this.dbId = this.dbIdFromClientId(this.clientId);
102+
this.nthRequest = new AtomicInteger(0);
103+
}
104+
105+
private int dbIdFromClientId(String clientId) {
106+
int i = clientId.indexOf("-");
107+
String strWithValue = clientId.substring(i + 1);
108+
if (Objects.equals(strWithValue, "")) {
109+
strWithValue = "0";
110+
}
111+
return Integer.parseInt(strWithValue);
94112
}
95113

96114
@VisibleForTesting
@@ -188,7 +206,11 @@ public CommitResponse writeWithOptions(
188206
if (canUseMultiplexedSessionsForRW() && getMultiplexedSessionDatabaseClient() != null) {
189207
return getMultiplexedSessionDatabaseClient().writeWithOptions(mutations, options);
190208
}
191-
return runWithSessionRetry(session -> session.writeWithOptions(mutations, options));
209+
210+
return runWithSessionRetry(
211+
(session, reqId) -> {
212+
return session.writeWithOptions(mutations, withReqId(reqId, options));
213+
});
192214
} catch (RuntimeException e) {
193215
span.setStatus(e);
194216
throw e;
@@ -213,7 +235,8 @@ public CommitResponse writeAtLeastOnceWithOptions(
213235
.writeAtLeastOnceWithOptions(mutations, options);
214236
}
215237
return runWithSessionRetry(
216-
session -> session.writeAtLeastOnceWithOptions(mutations, options));
238+
(session, reqId) ->
239+
session.writeAtLeastOnceWithOptions(mutations, withReqId(reqId, options)));
217240
} catch (RuntimeException e) {
218241
span.setStatus(e);
219242
throw e;
@@ -222,6 +245,10 @@ public CommitResponse writeAtLeastOnceWithOptions(
222245
}
223246
}
224247

248+
private int nextNthRequest() {
249+
return this.nthRequest.incrementAndGet();
250+
}
251+
225252
@Override
226253
public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
227254
final Iterable<MutationGroup> mutationGroups, final TransactionOption... options)
@@ -231,7 +258,9 @@ public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
231258
if (canUseMultiplexedSessionsForRW() && getMultiplexedSessionDatabaseClient() != null) {
232259
return getMultiplexedSessionDatabaseClient().batchWriteAtLeastOnce(mutationGroups, options);
233260
}
234-
return runWithSessionRetry(session -> session.batchWriteAtLeastOnce(mutationGroups, options));
261+
return runWithSessionRetry(
262+
(session, reqId) ->
263+
session.batchWriteAtLeastOnce(mutationGroups, withReqId(reqId, options)));
235264
} catch (RuntimeException e) {
236265
span.setStatus(e);
237266
throw e;
@@ -383,27 +412,57 @@ private Future<Dialect> getDialectAsync() {
383412
return pool.getDialectAsync();
384413
}
385414

415+
private UpdateOption[] withReqId(
416+
final XGoogSpannerRequestId reqId, final UpdateOption... options) {
417+
if (reqId == null) {
418+
return options;
419+
}
420+
ArrayList<UpdateOption> allOptions = new ArrayList(Arrays.asList(options));
421+
allOptions.add(new Options.RequestIdOption(reqId));
422+
return allOptions.toArray(new UpdateOption[0]);
423+
}
424+
425+
private TransactionOption[] withReqId(
426+
final XGoogSpannerRequestId reqId, final TransactionOption... options) {
427+
if (reqId == null) {
428+
return options;
429+
}
430+
ArrayList<TransactionOption> allOptions = new ArrayList(Arrays.asList(options));
431+
allOptions.add(new Options.RequestIdOption(reqId));
432+
return allOptions.toArray(new TransactionOption[0]);
433+
}
434+
386435
private long executePartitionedUpdateWithPooledSession(
387436
final Statement stmt, final UpdateOption... options) {
388437
ISpan span = tracer.spanBuilder(PARTITION_DML_TRANSACTION, commonAttributes);
389438
try (IScope s = tracer.withSpan(span)) {
390-
return runWithSessionRetry(session -> session.executePartitionedUpdate(stmt, options));
439+
return runWithSessionRetry(
440+
(session, reqId) -> {
441+
return session.executePartitionedUpdate(stmt, withReqId(reqId, options));
442+
});
391443
} catch (RuntimeException e) {
392444
span.setStatus(e);
393445
span.end();
394446
throw e;
395447
}
396448
}
397449

398-
private <T> T runWithSessionRetry(Function<Session, T> callable) {
450+
private <T> T runWithSessionRetry(BiFunction<Session, XGoogSpannerRequestId, T> callable) {
399451
PooledSessionFuture session = getSession();
452+
XGoogSpannerRequestId reqId =
453+
XGoogSpannerRequestId.of(
454+
this.dbId, Long.valueOf(session.getChannel()), this.nextNthRequest(), 0);
400455
while (true) {
401456
try {
402-
return callable.apply(session);
457+
reqId.incrementAttempt();
458+
return callable.apply(session, reqId);
403459
} catch (SessionNotFoundException e) {
404460
session =
405461
(PooledSessionFuture)
406462
pool.getPooledSessionReplacementHandler().replaceSession(e, session);
463+
reqId =
464+
XGoogSpannerRequestId.of(
465+
this.dbId, Long.valueOf(session.getChannel()), this.nextNthRequest(), 0);
407466
}
408467
}
409468
}

‎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 ReadOption, 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)