Skip to content

Commit 18a4509

Browse files
tianlei2tianlei2
committed
Stabilize HBase snapshot import and refactor tests
- Fixed Splittable DoFn contract violation in ReadSnapshotRegion. - Added GCS pagination in SnapshotUtils. - Refactored RegionConfigCoder to use standard Beam coders. - Extracted sharding logic in ReadRegions and added TestPipeline test in ReadRegionsTest. - Fixed potential NPEs and enhanced test assertions.
1 parent b7866da commit 18a4509

11 files changed

Lines changed: 291 additions & 146 deletions

File tree

‎bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/SnapshotUtils.java‎

Lines changed: 18 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -64,8 +64,6 @@ public class SnapshotUtils {
6464
private static final String SNAPSHOT_MANIFEST_DIRECTORY = ".hbase-snapshot";
6565
private static final String GCS_SCHEME = "gs";
6666
private static final Sleeper sleeper = Sleeper.DEFAULT;
67-
private static final Object lock = new Object();
68-
private static volatile Configuration hbaseConfiguration;
6967

7068
private SnapshotUtils() {}
7169

@@ -134,13 +132,7 @@ static Map<String, String> getConfiguration(
134132
}
135133

136134
public static Configuration getHBaseConfiguration(Map<String, String> configurations) {
137-
if (hbaseConfiguration == null) {
138-
synchronized (lock) {
139-
if (hbaseConfiguration == null)
140-
hbaseConfiguration = createHbaseConfiguration(configurations);
141-
}
142-
}
143-
return hbaseConfiguration;
135+
return createHbaseConfiguration(configurations);
144136
}
145137

146138
private static Configuration createHbaseConfiguration(Map<String, String> configurations) {
@@ -229,6 +221,7 @@ public static void setRestorePath(String restorePath, ImportConfig importConfig)
229221
public static Map<String, String> getSnapshotsFromString(String snapshotNames) {
230222
Map<String, String> snapshots = new HashMap<>();
231223
for (String snapshotInfo : snapshotNames.split(",")) {
224+
snapshotInfo = snapshotInfo.trim();
232225
String[] snapshotWithTableName = snapshotInfo.split(":");
233226
if (snapshotWithTableName.length == 2)
234227
snapshots.put(snapshotWithTableName[0], snapshotWithTableName[1]);
@@ -268,20 +261,31 @@ public static Map<String, String> getSnapshotsFromSnapshotPath(
268261
GcsPath gcsPath = GcsPath.fromUri(importSnapshotpath);
269262
Map<String, String> snapshots = new HashMap<>();
270263

271-
List<StorageObject> objects =
272-
gcsUtil.listObjects(gcsPath.getBucket(), gcsPath.getObject(), null).getItems();
273-
if (objects == null)
264+
List<StorageObject> allObjects = new java.util.ArrayList<>();
265+
String pageToken = null;
266+
do {
267+
com.google.api.services.storage.model.Objects objects =
268+
gcsUtil.listObjects(gcsPath.getBucket(), gcsPath.getObject(), pageToken);
269+
if (objects.getItems() == null) {
270+
break;
271+
}
272+
allObjects.addAll(objects.getItems());
273+
pageToken = objects.getNextPageToken();
274+
} while (pageToken != null);
275+
276+
if (allObjects.isEmpty())
274277
throw new IllegalStateException(
275278
String.format("Snapshot path %s does not contain any snapshots", importSnapshotpath));
276279

277280
// Build a pattern for object portion e.g if path is
278281
// gs://sym-bucket/snapshots/20220309230526/.hbase-snapshot
279282
// the object portion would be snapshots/60G/20220309230526/.hbase-snapshot
280-
Pattern pathPattern = Pattern.compile(String.format("%s/(.+?/)", gcsPath.getObject()));
283+
Pattern pathPattern =
284+
Pattern.compile(String.format("%s/(.+?/)", Pattern.quote(gcsPath.getObject())));
281285
Pattern prefixPattern = prefix.equals("*") ? null : Pattern.compile(prefix);
282286
Matcher pathMatcher = null;
283287
String snapshotName = null;
284-
for (StorageObject object : objects) {
288+
for (StorageObject object : allObjects) {
285289
pathMatcher = pathPattern.matcher(object.getId());
286290
if (pathMatcher.find()) {
287291
// Group 1 represents the snapshot directory name along with suffix slash (e.g: snapshot1/)

‎bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/coders/RegionConfigCoder.java‎

Lines changed: 23 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,14 @@
1818
import com.google.api.core.InternalApi;
1919
import com.google.cloud.bigtable.beam.hbasesnapshots.conf.RegionConfig;
2020
import com.google.cloud.bigtable.beam.hbasesnapshots.conf.SnapshotConfig;
21-
import java.io.*;
21+
import java.io.IOException;
22+
import java.io.InputStream;
23+
import java.io.OutputStream;
2224
import java.util.Collections;
2325
import java.util.List;
26+
import org.apache.beam.sdk.coders.ByteArrayCoder;
2427
import org.apache.beam.sdk.coders.Coder;
25-
import org.apache.beam.sdk.coders.CoderException;
28+
import org.apache.beam.sdk.coders.SerializableCoder;
2629
import org.apache.beam.sdk.coders.VarLongCoder;
2730
import org.apache.hadoop.hbase.client.RegionInfo;
2831
import org.apache.hadoop.hbase.client.TableDescriptor;
@@ -35,59 +38,42 @@
3538
public class RegionConfigCoder extends Coder<RegionConfig> {
3639
private static final VarLongCoder longCoder = VarLongCoder.of();
3740

41+
private static final Coder<SnapshotConfig> snapshotConfigCoder =
42+
SerializableCoder.of(SnapshotConfig.class);
43+
private static final Coder<byte[]> byteArrayCoder = ByteArrayCoder.of();
44+
3845
@Override
3946
public void encode(RegionConfig value, OutputStream outStream) throws IOException {
40-
ObjectOutputStream objectOutputStream = new ObjectOutputStream(outStream);
41-
objectOutputStream.writeObject(value.getSnapshotConfig());
47+
snapshotConfigCoder.encode(value.getSnapshotConfig(), outStream);
4248

4349
HBaseProtos.RegionInfo regionInfo = ProtobufUtil.toRegionInfo(value.getRegionInfo());
44-
ByteArrayOutputStream boas1 = new ByteArrayOutputStream();
45-
regionInfo.writeTo(boas1);
46-
objectOutputStream.writeObject(boas1.toByteArray());
50+
byteArrayCoder.encode(regionInfo.toByteArray(), outStream);
4751

4852
HBaseProtos.TableSchema tableSchema = ProtobufUtil.toTableSchema(value.getTableDescriptor());
49-
ByteArrayOutputStream boas2 = new ByteArrayOutputStream();
50-
tableSchema.writeTo(boas2);
51-
objectOutputStream.writeObject(boas2.toByteArray());
53+
byteArrayCoder.encode(tableSchema.toByteArray(), outStream);
5254

5355
longCoder.encode(value.getRegionSize(), outStream);
5456
}
5557

5658
@Override
5759
public RegionConfig decode(InputStream inStream) throws IOException {
58-
ObjectInputStream objectInputStream = new ObjectInputStream(inStream);
59-
SnapshotConfig snapshotConfig;
60-
try {
61-
snapshotConfig = (SnapshotConfig) objectInputStream.readObject();
62-
} catch (ClassNotFoundException e) {
63-
throw new CoderException("Failed to deserialize RestoredSnapshotConfig", e);
64-
}
60+
SnapshotConfig snapshotConfig = snapshotConfigCoder.decode(inStream);
6561

66-
RegionInfo regionInfoProto = null;
67-
try {
68-
regionInfoProto =
69-
ProtobufUtil.toRegionInfo(
70-
HBaseProtos.RegionInfo.parseFrom((byte[]) objectInputStream.readObject()));
71-
} catch (ClassNotFoundException e) {
72-
throw new CoderException("Failed to parse regionInfo", e);
73-
}
62+
byte[] regionInfoBytes = byteArrayCoder.decode(inStream);
63+
RegionInfo regionInfo =
64+
ProtobufUtil.toRegionInfo(HBaseProtos.RegionInfo.parseFrom(regionInfoBytes));
7465

75-
TableDescriptor tableSchema = null;
76-
try {
77-
tableSchema =
78-
ProtobufUtil.toTableDescriptor(
79-
TableSchema.parseFrom((byte[]) objectInputStream.readObject()));
80-
} catch (ClassNotFoundException e) {
81-
throw new CoderException("Failed to parse tableSchema", e);
82-
}
66+
byte[] tableSchemaBytes = byteArrayCoder.decode(inStream);
67+
TableDescriptor tableDescriptor =
68+
ProtobufUtil.toTableDescriptor(TableSchema.parseFrom(tableSchemaBytes));
8369

84-
Long regionsize = longCoder.decode(inStream);
70+
Long regionSize = longCoder.decode(inStream);
8571

8672
return RegionConfig.builder()
8773
.setSnapshotConfig(snapshotConfig)
88-
.setRegionInfo(regionInfoProto)
89-
.setTableDescriptor(tableSchema)
90-
.setRegionSize(regionsize)
74+
.setRegionInfo(regionInfo)
75+
.setTableDescriptor(tableDescriptor)
76+
.setRegionSize(regionSize)
9177
.build();
9278
}
9379

‎bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/conf/SnapshotConfig.java‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616
package com.google.cloud.bigtable.beam.hbasesnapshots.conf;
1717

1818
import com.google.auto.value.AutoValue;
19-
import com.google.auto.value.extension.memoized.Memoized;
2019
import com.google.cloud.bigtable.beam.hbasesnapshots.SnapshotUtils;
2120
import java.io.Serializable;
2221
import java.util.Map;
@@ -38,12 +37,10 @@ public static Builder builder() {
3837

3938
public abstract String getSourceLocation();
4039

41-
@Memoized
4240
public Path getSourcePath() {
4341
return new Path(getSourceLocation());
4442
}
4543

46-
@Memoized
4744
public Path getRestorePath() {
4845
return new Path(getRestoreLocation());
4946
}

‎bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/CreateBigtableMutations.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,7 @@ public void processElement(
7070
@Element KV<SnapshotConfig, Result> element,
7171
OutputReceiver<KV<String, Iterable<Mutation>>> outputReceiver)
7272
throws IOException {
73-
if (element.getValue().listCells().isEmpty()) {
73+
if (element.getValue().listCells() == null || element.getValue().listCells().isEmpty()) {
7474
return;
7575
}
7676
String targetTable = element.getKey().getTableName();
@@ -81,15 +81,15 @@ public void processElement(
8181
List<Mutation> mutations = new ArrayList<>();
8282

8383
boolean logAndSkipIncompatibleRowMutations =
84-
verifyRowMutationThresholds(
84+
convertAndValidateThresholds(
8585
rowKey, element.getValue().listCells(), mutations, snapshotName);
8686

8787
if (!logAndSkipIncompatibleRowMutations && mutations.size() > 0) {
8888
outputReceiver.output(KV.of(targetTable, mutations));
8989
}
9090
}
9191

92-
private boolean verifyRowMutationThresholds(
92+
private boolean convertAndValidateThresholds(
9393
byte[] rowKey, List<Cell> cells, List<Mutation> mutations, String snapshotName)
9494
throws IOException {
9595
boolean logAndSkipIncompatibleRows = false;

‎bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/ReadSnapshotRegion.java‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -54,18 +54,17 @@ public void processElement(
5454
RestrictionTracker<ByteKeyRange, ByteKey> tracker)
5555
throws Exception {
5656

57-
boolean hasSplit = false;
5857
try (HBaseRegionScanner scanner = newScanner(regionConfig, tracker.currentRestriction())) {
5958
for (Result result = scanner.next(); result != null; result = scanner.next()) {
6059
if (tracker.tryClaim(ByteKey.copyFrom(result.getRow()))) {
6160
outputReceiver.output(KV.of(regionConfig.getSnapshotConfig(), result));
6261
} else {
63-
hasSplit = true;
64-
break;
62+
return;
6563
}
6664
}
65+
// Scanner exhausted naturally. Claim the end key to mark done.
66+
tracker.tryClaim(tracker.currentRestriction().getEndKey());
6767
}
68-
tracker.tryClaim(ByteKey.EMPTY);
6968
}
7069

7170
HBaseRegionScanner newScanner(RegionConfig regionConfig, ByteKeyRange byteKeyRange)

‎bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/transforms/ListRegions.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -61,8 +61,8 @@ private Map<Long, Long> computeRegionSize(SnapshotManifest snapshotManifest) {
6161
totalSize = 0;
6262
for (SnapshotProtos.SnapshotRegionManifest.FamilyFiles familyFiles :
6363
regionManifest.getFamilyFilesList()) {
64-
for (SnapshotProtos.SnapshotRegionManifest.StoreFile StoreFile :
65-
familyFiles.getStoreFilesList()) totalSize += StoreFile.getFileSize();
64+
for (SnapshotProtos.SnapshotRegionManifest.StoreFile storeFile :
65+
familyFiles.getStoreFilesList()) totalSize += storeFile.getFileSize();
6666
}
6767
regionsSize.put(regionManifest.getRegionInfo().getRegionId(), totalSize);
6868
}
@@ -104,7 +104,7 @@ public void processElement(
104104
.setSnapshotConfig(snapshotConfig)
105105
.setTableDescriptor(tableDescriptor)
106106
.setRegionInfo(regionInfo)
107-
.setRegionSize(regionsSize.get(regionInfo.getRegionId()))
107+
.setRegionSize(regionsSize.getOrDefault(regionInfo.getRegionId(), 0L))
108108
.build())
109109
.forEach(outputReceiver::output);
110110
}

‎bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/transforms/ReadRegions.java‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,8 @@
3333
import org.apache.beam.sdk.values.KV;
3434
import org.apache.beam.sdk.values.PCollection;
3535
import org.apache.beam.sdk.values.TypeDescriptor;
36-
import org.apache.hadoop.hbase.client.*;
36+
import org.apache.hadoop.hbase.client.Mutation;
37+
import org.apache.hadoop.hbase.client.Result;
3738
import org.slf4j.Logger;
3839
import org.slf4j.LoggerFactory;
3940

@@ -103,6 +104,12 @@ public ReadRegions(
103104
this.shardIndex = shardIndex;
104105
}
105106

107+
static boolean isRegionSelected(RegionConfig rc, int numShards, int shardIndex) {
108+
byte[] regionName = rc.getRegionInfo().getEncodedNameAsBytes();
109+
long remainder = new BigInteger(regionName).mod(BigInteger.valueOf(numShards)).longValue();
110+
return remainder == shardIndex;
111+
}
112+
106113
@Override
107114
public PCollection<KV<String, Iterable<Mutation>>> expand(
108115
PCollection<RegionConfig> regionConfig) {
@@ -123,12 +130,7 @@ public PCollection<KV<String, Iterable<Mutation>>> expand(
123130
"Select regions for shard",
124131
Filter.by(
125132
rc -> {
126-
// encodedName is an MD5 hash of the region info and therefor should be well
127-
// distributed
128-
byte[] regionName = rc.getRegionInfo().getEncodedNameAsBytes();
129-
long remainder =
130-
new BigInteger(regionName).mod(BigInteger.valueOf(numShards)).longValue();
131-
boolean shouldTake = remainder == shardIndex;
133+
boolean shouldTake = isRegionSelected(rc, numShards, shardIndex);
132134
LOG.info(
133135
"Region {} was {} due to sharding",
134136
rc.getRegionInfo().getRegionNameAsString(),

‎bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/EndToEndIT.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -155,7 +155,7 @@ public void setup() throws Exception {
155155

156156
@After
157157
public void teardown() throws IOException {
158-
final List<GcsPath> paths = gcsUtil.expand(GcsPath.fromUri(workDir + "*"));
158+
final List<GcsPath> paths = gcsUtil.expand(GcsPath.fromUri(workDir + "**"));
159159

160160
if (!paths.isEmpty()) {
161161
// TODO: cleanup fails when tests time out. Add a orphan cleaner in the setup()
@@ -520,7 +520,7 @@ public void testHBaseSnapshotImportWithSharding() throws Exception {
520520
Assert.assertEquals(counters.get("ranges_matched"), (Long) 100L);
521521
Assert.assertEquals(counters.get("ranges_not_matched"), (Long) 0L);
522522

523-
final List<GcsPath> paths = gcsUtil.expand(GcsPath.fromUri(sharedRestorePath + "*"));
523+
final List<GcsPath> paths = gcsUtil.expand(GcsPath.fromUri(sharedRestorePath + "**"));
524524
if (!paths.isEmpty()) {
525525
gcsUtil.remove(paths.stream().map(GcsPath::toString).collect(Collectors.toList()));
526526
}

0 commit comments

Comments
 (0)