Skip to content

Commit 8556574

Browse files
djyaugcf-owl-bot[bot]
djyau
and
gcf-owl-bot[bot]
authored
feat: add support for awaiting Data Boost (#2329)
* Introduce ConsistencyParams model Change-Id: I7ae07cc4f13e8ffe9ea4a55fb407eae1d64f8547 * Create AwaitConsistencyCallable, a delegate for AwaitReplicationCallable Change-Id: I16d7c69b4e1b9153f93e83f7f846d6c172aae8a6 * Address some PR comments Change-Id: Icdf94f2f6a13f55d2c9204774d4ebbac5aa6a8b3 * Remove unused imports from AwaitReplicationCallable Change-Id: Ia4861f1a5796061ca86844c67c68f54711fdbb94 * Plumb the Consistency callable through to some places, add some tests Change-Id: Ibe60e2a1044933af1008c0cd1b84f757dd6867a8 * Add integration test Change-Id: Ie3b2b2983ca585cb1d6a2cdb8b18b55e81205759 * Rework the ConsistencyRequest model, plumb through RequestContext to BigtableTableAdminClient Change-Id: I840282587d3d6cb4150dfbdd568c347dc32a732d * Fix imports Change-Id: Ic7588b3d04877a56089c23036d6df73a5c9b0cd5 * Fix more imports, fix some tests Change-Id: I2723fd67bd301a4eb3aeae80d91fa663cdd6ab01 * Rename some things Change-Id: Ie1bc8478c418d49b0c2e014edbeb6f56b56b0dd1 * Add tests for ConsistencyRequest model Change-Id: I3548b7aa673be5a92cd4c180e3edb8649657811c * Add newline Change-Id: Icdd22ce2857e5b4316c6fa3f0e139ea9de825178 * Fix broken test Change-Id: Idbd7c0f10ebe575d104ab7ac46a3a1e347e35fe8 * Make request context a final variable in test Change-Id: I81f2a25fe4493021bab150ab0af65d7318ba2399 * Get test working using correct expectations Change-Id: Ie34d5171bd7a472fc695d603849e260054aedfbd * Add a couple of tests for AwaitReplicationCallable Change-Id: I70014db2c0a1d4e74c23b18de7ef591bc70cda2a * Use RequestContextNoAP class Change-Id: I897b343cd1067d43bcc644cac3db44e88bbf1e69 * Make ConsistencyRequest model an AutoValue Change-Id: I9529fb79da69e12a834a2d0fea032d72ae6ea157 * Fix license year, fix some formatting Change-Id: Ibcca1ca9f49988764fdbeeacc59cac5d276ab266 * Run auto formatter Change-Id: I9f5e3f7c7fd79262092c507a523e16a533bc4382 * Rename new RequestContext to TableAdminRequestContext, re run auto formatter Change-Id: Ib3f5918ef0f5b1ac53147baf93dcb72c476d877b * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * Add license header to ConsistencyRequestTest Change-Id: I733d2f8c082647ad32b72b04b218cd5ba79d2377 * Add EnhancedBigtableTableAdminStub to clirr-ignored-differences Change-Id: I7eefeda777305dd3d7c5664097bda87ac63daa72 * Fix IT tests, skip data boost one for now until we run it concurrently Change-Id: I764190b0f91614753080e0a96e7e11e3dfb1fde0 * Run autoformatter Change-Id: Iba4671e4781f1b333279a2410563869f53b284d5 --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
1 parent 3a9b5a6 commit 8556574

File tree

11 files changed

+592
-166
lines changed

11 files changed

+592
-166
lines changed

‎google-cloud-bigtable/clirr-ignored-differences.xml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -259,4 +259,10 @@
259259
<className>com/google/cloud/bigtable/admin/v2/models/Type$Int64$Encoding$BigEndianBytes</className>
260260
<method>*</method>
261261
</difference>
262+
<difference>
263+
<!-- change method args is ok because EnhancedBigtableTableAdminStub is InternalApi -->
264+
<differenceType>7004</differenceType>
265+
<className>com/google/cloud/bigtable/admin/v2/stub/EnhancedBigtableTableAdminStub</className>
266+
<method>*</method>
267+
</difference>
262268
</differences>

‎google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClient.java‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
import com.google.cloud.bigtable.admin.v2.internal.NameUtil;
4747
import com.google.cloud.bigtable.admin.v2.models.AuthorizedView;
4848
import com.google.cloud.bigtable.admin.v2.models.Backup;
49+
import com.google.cloud.bigtable.admin.v2.models.ConsistencyRequest;
4950
import com.google.cloud.bigtable.admin.v2.models.CopyBackupRequest;
5051
import com.google.cloud.bigtable.admin.v2.models.CreateAuthorizedViewRequest;
5152
import com.google.cloud.bigtable.admin.v2.models.CreateBackupRequest;
@@ -61,6 +62,7 @@
6162
import com.google.cloud.bigtable.admin.v2.models.UpdateBackupRequest;
6263
import com.google.cloud.bigtable.admin.v2.models.UpdateTableRequest;
6364
import com.google.cloud.bigtable.admin.v2.stub.EnhancedBigtableTableAdminStub;
65+
import com.google.cloud.bigtable.data.v2.internal.TableAdminRequestContext;
6466
import com.google.common.base.Preconditions;
6567
import com.google.common.collect.ImmutableList;
6668
import com.google.common.collect.ImmutableMap;
@@ -154,8 +156,10 @@ public static BigtableTableAdminClient create(
154156
/** Constructs an instance of BigtableTableAdminClient with the given settings. */
155157
public static BigtableTableAdminClient create(@Nonnull BigtableTableAdminSettings settings)
156158
throws IOException {
159+
TableAdminRequestContext requestContext =
160+
TableAdminRequestContext.create(settings.getProjectId(), settings.getInstanceId());
157161
EnhancedBigtableTableAdminStub stub =
158-
EnhancedBigtableTableAdminStub.createEnhanced(settings.getStubSettings());
162+
EnhancedBigtableTableAdminStub.createEnhanced(settings.getStubSettings(), requestContext);
159163
return create(settings.getProjectId(), settings.getInstanceId(), stub);
160164
}
161165

@@ -917,6 +921,11 @@ public void awaitReplication(String tableId) {
917921
stub.awaitReplicationCallable().futureCall(tableName));
918922
}
919923

924+
public void awaitConsistency(ConsistencyRequest consistencyRequest) {
925+
ApiExceptions.callAndTranslateApiException(
926+
stub.awaitConsistencyCallable().futureCall(consistencyRequest));
927+
}
928+
920929
/**
921930
* Creates a backup with the specified configuration.
922931
*

‎google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/models/ConsistencyRequest.java‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
/*
2+
* Copyright 2024 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.bigtable.admin.v2.models;
17+
18+
import com.google.api.core.InternalApi;
19+
import com.google.auto.value.AutoValue;
20+
import com.google.bigtable.admin.v2.CheckConsistencyRequest;
21+
import com.google.bigtable.admin.v2.DataBoostReadLocalWrites;
22+
import com.google.bigtable.admin.v2.GenerateConsistencyTokenRequest;
23+
import com.google.bigtable.admin.v2.StandardReadRemoteWrites;
24+
import com.google.bigtable.admin.v2.TableName;
25+
import com.google.cloud.bigtable.data.v2.internal.TableAdminRequestContext;
26+
import javax.annotation.Nonnull;
27+
28+
@AutoValue
29+
public abstract class ConsistencyRequest {
30+
@Nonnull
31+
protected abstract String getTableId();
32+
33+
@Nonnull
34+
protected abstract CheckConsistencyRequest.ModeCase getMode();
35+
36+
public static ConsistencyRequest forReplication(String tableId) {
37+
return new AutoValue_ConsistencyRequest(
38+
tableId, CheckConsistencyRequest.ModeCase.STANDARD_READ_REMOTE_WRITES);
39+
}
40+
41+
public static ConsistencyRequest forDataBoost(String tableId) {
42+
return new AutoValue_ConsistencyRequest(
43+
tableId, CheckConsistencyRequest.ModeCase.DATA_BOOST_READ_LOCAL_WRITES);
44+
}
45+
46+
@InternalApi
47+
public CheckConsistencyRequest toCheckConsistencyProto(
48+
TableAdminRequestContext requestContext, String token) {
49+
CheckConsistencyRequest.Builder builder = CheckConsistencyRequest.newBuilder();
50+
TableName tableName =
51+
TableName.of(requestContext.getProjectId(), requestContext.getInstanceId(), getTableId());
52+
53+
if (getMode().equals(CheckConsistencyRequest.ModeCase.STANDARD_READ_REMOTE_WRITES)) {
54+
builder.setStandardReadRemoteWrites(StandardReadRemoteWrites.newBuilder().build());
55+
} else {
56+
builder.setDataBoostReadLocalWrites(DataBoostReadLocalWrites.newBuilder().build());
57+
}
58+
59+
return builder.setName(tableName.toString()).setConsistencyToken(token).build();
60+
}
61+
62+
@InternalApi
63+
public GenerateConsistencyTokenRequest toGenerateTokenProto(
64+
TableAdminRequestContext requestContext) {
65+
GenerateConsistencyTokenRequest.Builder builder = GenerateConsistencyTokenRequest.newBuilder();
66+
TableName tableName =
67+
TableName.of(requestContext.getProjectId(), requestContext.getInstanceId(), getTableId());
68+
69+
return builder.setName(tableName.toString()).build();
70+
}
71+
}

‎google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/stub/AwaitConsistencyCallable.java‎

Lines changed: 195 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,195 @@
1+
/*
2+
* Copyright 2024 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.bigtable.admin.v2.stub;
17+
18+
import com.google.api.core.ApiAsyncFunction;
19+
import com.google.api.core.ApiFunction;
20+
import com.google.api.core.ApiFuture;
21+
import com.google.api.core.ApiFutures;
22+
import com.google.api.gax.retrying.ExponentialPollAlgorithm;
23+
import com.google.api.gax.retrying.NonCancellableFuture;
24+
import com.google.api.gax.retrying.ResultRetryAlgorithm;
25+
import com.google.api.gax.retrying.RetryAlgorithm;
26+
import com.google.api.gax.retrying.RetrySettings;
27+
import com.google.api.gax.retrying.RetryingExecutor;
28+
import com.google.api.gax.retrying.RetryingFuture;
29+
import com.google.api.gax.retrying.ScheduledRetryingExecutor;
30+
import com.google.api.gax.retrying.TimedAttemptSettings;
31+
import com.google.api.gax.rpc.ApiCallContext;
32+
import com.google.api.gax.rpc.ClientContext;
33+
import com.google.api.gax.rpc.UnaryCallable;
34+
import com.google.bigtable.admin.v2.CheckConsistencyRequest;
35+
import com.google.bigtable.admin.v2.CheckConsistencyResponse;
36+
import com.google.bigtable.admin.v2.GenerateConsistencyTokenRequest;
37+
import com.google.bigtable.admin.v2.GenerateConsistencyTokenResponse;
38+
import com.google.cloud.bigtable.admin.v2.models.ConsistencyRequest;
39+
import com.google.cloud.bigtable.data.v2.internal.TableAdminRequestContext;
40+
import com.google.common.annotations.VisibleForTesting;
41+
import com.google.common.util.concurrent.MoreExecutors;
42+
import java.util.concurrent.Callable;
43+
import java.util.concurrent.CancellationException;
44+
45+
/**
46+
* Callable that waits until either replication or Data Boost has caught up to the point it was
47+
* called.
48+
*
49+
* <p>This callable wraps GenerateConsistencyToken and CheckConsistency RPCs. It will generate a
50+
* token then poll until isConsistent is true.
51+
*/
52+
class AwaitConsistencyCallable extends UnaryCallable<ConsistencyRequest, Void> {
53+
private final UnaryCallable<GenerateConsistencyTokenRequest, GenerateConsistencyTokenResponse>
54+
generateCallable;
55+
private final UnaryCallable<CheckConsistencyRequest, CheckConsistencyResponse> checkCallable;
56+
private final RetryingExecutor<CheckConsistencyResponse> executor;
57+
58+
private final TableAdminRequestContext requestContext;
59+
60+
static AwaitConsistencyCallable create(
61+
UnaryCallable<GenerateConsistencyTokenRequest, GenerateConsistencyTokenResponse>
62+
generateCallable,
63+
UnaryCallable<CheckConsistencyRequest, CheckConsistencyResponse> checkCallable,
64+
ClientContext clientContext,
65+
RetrySettings pollingSettings,
66+
TableAdminRequestContext requestContext) {
67+
68+
RetryAlgorithm<CheckConsistencyResponse> retryAlgorithm =
69+
new RetryAlgorithm<>(
70+
new PollResultAlgorithm(),
71+
new ExponentialPollAlgorithm(pollingSettings, clientContext.getClock()));
72+
73+
RetryingExecutor<CheckConsistencyResponse> retryingExecutor =
74+
new ScheduledRetryingExecutor<>(retryAlgorithm, clientContext.getExecutor());
75+
76+
return new AwaitConsistencyCallable(
77+
generateCallable, checkCallable, retryingExecutor, requestContext);
78+
}
79+
80+
@VisibleForTesting
81+
AwaitConsistencyCallable(
82+
UnaryCallable<GenerateConsistencyTokenRequest, GenerateConsistencyTokenResponse>
83+
generateCallable,
84+
UnaryCallable<CheckConsistencyRequest, CheckConsistencyResponse> checkCallable,
85+
RetryingExecutor<CheckConsistencyResponse> executor,
86+
TableAdminRequestContext requestContext) {
87+
this.generateCallable = generateCallable;
88+
this.checkCallable = checkCallable;
89+
this.executor = executor;
90+
this.requestContext = requestContext;
91+
}
92+
93+
@Override
94+
public ApiFuture<Void> futureCall(
95+
final ConsistencyRequest consistencyRequest, final ApiCallContext apiCallContext) {
96+
ApiFuture<GenerateConsistencyTokenResponse> tokenFuture =
97+
generateToken(consistencyRequest.toGenerateTokenProto(requestContext), apiCallContext);
98+
99+
return ApiFutures.transformAsync(
100+
tokenFuture,
101+
new ApiAsyncFunction<GenerateConsistencyTokenResponse, Void>() {
102+
@Override
103+
public ApiFuture<Void> apply(GenerateConsistencyTokenResponse input) {
104+
CheckConsistencyRequest request =
105+
consistencyRequest.toCheckConsistencyProto(
106+
requestContext, input.getConsistencyToken());
107+
return pollToken(request, apiCallContext);
108+
}
109+
},
110+
MoreExecutors.directExecutor());
111+
}
112+
113+
private ApiFuture<GenerateConsistencyTokenResponse> generateToken(
114+
GenerateConsistencyTokenRequest generateRequest, ApiCallContext context) {
115+
return generateCallable.futureCall(generateRequest, context);
116+
}
117+
118+
private ApiFuture<Void> pollToken(CheckConsistencyRequest request, ApiCallContext context) {
119+
AttemptCallable<CheckConsistencyRequest, CheckConsistencyResponse> attemptCallable =
120+
new AttemptCallable<>(checkCallable, request, context);
121+
RetryingFuture<CheckConsistencyResponse> retryingFuture =
122+
executor.createFuture(attemptCallable);
123+
attemptCallable.setExternalFuture(retryingFuture);
124+
attemptCallable.call();
125+
126+
return ApiFutures.transform(
127+
retryingFuture,
128+
new ApiFunction<CheckConsistencyResponse, Void>() {
129+
@Override
130+
public Void apply(CheckConsistencyResponse input) {
131+
return null;
132+
}
133+
},
134+
MoreExecutors.directExecutor());
135+
}
136+
137+
/** A callable representing an attempt to make an RPC call. */
138+
private static class AttemptCallable<RequestT, ResponseT> implements Callable<ResponseT> {
139+
private final UnaryCallable<RequestT, ResponseT> callable;
140+
private final RequestT request;
141+
142+
private volatile RetryingFuture<ResponseT> externalFuture;
143+
private volatile ApiCallContext callContext;
144+
145+
AttemptCallable(
146+
UnaryCallable<RequestT, ResponseT> callable, RequestT request, ApiCallContext callContext) {
147+
this.callable = callable;
148+
this.request = request;
149+
this.callContext = callContext;
150+
}
151+
152+
void setExternalFuture(RetryingFuture<ResponseT> externalFuture) {
153+
this.externalFuture = externalFuture;
154+
}
155+
156+
@Override
157+
public ResponseT call() {
158+
try {
159+
// NOTE: unlike gax's AttemptCallable, this ignores rpc timeouts
160+
externalFuture.setAttemptFuture(new NonCancellableFuture<ResponseT>());
161+
if (externalFuture.isDone()) {
162+
return null;
163+
}
164+
ApiFuture<ResponseT> internalFuture = callable.futureCall(request, callContext);
165+
externalFuture.setAttemptFuture(internalFuture);
166+
} catch (Throwable e) {
167+
externalFuture.setAttemptFuture(ApiFutures.<ResponseT>immediateFailedFuture(e));
168+
}
169+
170+
return null;
171+
}
172+
}
173+
174+
/**
175+
* A polling algorithm for waiting for a consistent {@link CheckConsistencyResponse}. Please note
176+
* that this class doesn't handle retryable errors and expects the underlying callable chain to
177+
* handle this.
178+
*/
179+
private static class PollResultAlgorithm
180+
implements ResultRetryAlgorithm<CheckConsistencyResponse> {
181+
@Override
182+
public TimedAttemptSettings createNextAttempt(
183+
Throwable prevThrowable,
184+
CheckConsistencyResponse prevResponse,
185+
TimedAttemptSettings prevSettings) {
186+
return null;
187+
}
188+
189+
@Override
190+
public boolean shouldRetry(Throwable prevThrowable, CheckConsistencyResponse prevResponse)
191+
throws CancellationException {
192+
return prevResponse != null && !prevResponse.getConsistent();
193+
}
194+
}
195+
}

0 commit comments

Comments
 (0)