From 8aadf137fbe721d8d0176fda921ff027c31b0638 Mon Sep 17 00:00:00 2001 From: mmuller Date: Fri, 1 Sep 2017 13:01:25 -0700 Subject: [PATCH] Create an entity record accumulator RecordAccumulator builds a set of datastore Entity records from a set of leveldb logfiles in a directory tree (which is how we receive them for database backup/restore testing). This CL also refactors some of the logfile test code out of LevelDbLogReaderTest so that we can reuse it for building test logs. ------------- Created by MOE: https://github.com/google/moe MOE_MIGRATED_REVID=167313553 --- .../registry/tools/RecordAccumulator.java | 59 ++++++++ javatests/google/registry/tools/BUILD | 1 + .../registry/tools/LevelDbLogReaderTest.java | 47 +----- .../google/registry/tools/LevelDbUtil.java | 87 +++++++++++ .../registry/tools/RecordAccumulatorTest.java | 140 ++++++++++++++++++ 5 files changed, 291 insertions(+), 43 deletions(-) create mode 100644 java/google/registry/tools/RecordAccumulator.java create mode 100644 javatests/google/registry/tools/LevelDbUtil.java create mode 100644 javatests/google/registry/tools/RecordAccumulatorTest.java diff --git a/java/google/registry/tools/RecordAccumulator.java b/java/google/registry/tools/RecordAccumulator.java new file mode 100644 index 00000000000..2a0791b86a4 --- /dev/null +++ b/java/google/registry/tools/RecordAccumulator.java @@ -0,0 +1,59 @@ +// Copyright 2017 The Nomulus Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package google.registry.tools; + +import com.google.appengine.api.datastore.EntityTranslator; +import com.google.common.collect.ImmutableSet; +import com.google.storage.onestore.v3.OnestoreEntity.EntityProto; +import java.io.File; +import java.io.FileInputStream; +import java.io.IOException; + +/** Utility class that accumulates Entity records from level db files. */ +class RecordAccumulator { + private final LevelDbLogReader reader = new LevelDbLogReader(); + + /** Recursively reads all records in the directory. */ + public final RecordAccumulator readDirectory(File dir) { + for (File child : dir.listFiles()) { + if (child.isDirectory()) { + readDirectory(child); + } else if (child.isFile()) { + try { + reader.readFrom(new FileInputStream(child)); + } catch (IOException e) { + throw new RuntimeException("IOException reading from file: " + child, e); + } + } + } + + return this; + } + + /** Creates an entity set from the current set of raw records. */ + ImmutableSet getComparableEntitySet() { + ImmutableSet.Builder builder = new ImmutableSet.Builder<>(); + for (byte[] rawRecord : reader.getRecords()) { + // Parse the entity proto and create an Entity object from it. + EntityProto proto = new EntityProto(); + proto.parseFrom(rawRecord); + ComparableEntity entity = new ComparableEntity(EntityTranslator.createFromPb(proto)); + + builder.add(entity); + } + + return builder.build(); + } +} diff --git a/javatests/google/registry/tools/BUILD b/javatests/google/registry/tools/BUILD index d04f6f8375e..3dd21ff69a6 100644 --- a/javatests/google/registry/tools/BUILD +++ b/javatests/google/registry/tools/BUILD @@ -41,6 +41,7 @@ java_library( "@com_google_api_client", "@com_google_appengine_api_1_0_sdk//:testonly", "@com_google_appengine_remote_api//:link", + "@com_google_auto_value", "@com_google_code_findbugs_jsr305", "@com_google_guava", "@com_google_http_client", diff --git a/javatests/google/registry/tools/LevelDbLogReaderTest.java b/javatests/google/registry/tools/LevelDbLogReaderTest.java index 606c0a76565..ca6daaf7db1 100644 --- a/javatests/google/registry/tools/LevelDbLogReaderTest.java +++ b/javatests/google/registry/tools/LevelDbLogReaderTest.java @@ -15,10 +15,12 @@ package google.registry.tools; import static com.google.common.truth.Truth.assertThat; -import static google.registry.tools.LevelDbLogReader.ChunkType; +import static google.registry.tools.LevelDbUtil.MAX_RECORD; +import static google.registry.tools.LevelDbUtil.addRecord; import com.google.common.collect.ImmutableList; import com.google.common.primitives.Bytes; +import google.registry.tools.LevelDbLogReader.ChunkType; import java.io.ByteArrayInputStream; import java.io.IOException; import java.util.List; @@ -30,8 +32,6 @@ @RunWith(JUnit4.class) public final class LevelDbLogReaderTest { - private static final int MAX_RECORD = LevelDbLogReader.BLOCK_SIZE - LevelDbLogReader.HEADER_SIZE; - // Size of the test record. Any value < 256 will do. private static final int TEST_RECORD_SIZE = 231; @@ -40,51 +40,12 @@ public final class LevelDbLogReaderTest { private static final int MAX_TEST_RECORD_OFFSET = LevelDbLogReader.BLOCK_SIZE - (LevelDbLogReader.HEADER_SIZE + TEST_RECORD_SIZE); - /** - * Adds a record of bytes of 'val' of the given size to bytes. - * - *

This currently doesn't write a real checksum since we're not doing anything with that in the - * leveldb reader. - * - *

Returns the new offset for the next block. - */ - private static int addRecord( - byte[] bytes, int pos, ChunkType type, int size, int val) { - - // Write a bogus checksum. - for (int i = 0; i < 4; ++i) { - bytes[pos++] = -1; - } - - // Write size and type. - bytes[pos++] = (byte) size; - bytes[pos++] = (byte) (size >> 8); - bytes[pos++] = (byte) type.getCode(); - - // Write "size" bytes of data. - for (int i = 0; i < size; ++i) { - bytes[pos + i] = (byte) val; - - // Swap the least significant bytes in val so we can have more than 256 different same-sized - // records. - val = (val >> 8) | ((val & 0xff) << 8); - } - - return pos + size; - } - private TestBlock makeBlockOfRepeatingBytes(int startVal) { byte[] block = new byte[LevelDbLogReader.BLOCK_SIZE]; int pos = 0; int recordCount = 0; while (pos < MAX_TEST_RECORD_OFFSET) { - pos = - addRecord( - block, - pos, - ChunkType.FULL, - TEST_RECORD_SIZE, - 0xffff & (pos + startVal)); + pos = addRecord(block, pos, ChunkType.FULL, TEST_RECORD_SIZE, 0xffff & (pos + startVal)); ++recordCount; } return new TestBlock(block, recordCount); diff --git a/javatests/google/registry/tools/LevelDbUtil.java b/javatests/google/registry/tools/LevelDbUtil.java new file mode 100644 index 00000000000..44086eb1568 --- /dev/null +++ b/javatests/google/registry/tools/LevelDbUtil.java @@ -0,0 +1,87 @@ +// Copyright 2017 The Nomulus Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package google.registry.tools; + +import static com.google.common.base.Preconditions.checkArgument; +import google.registry.tools.LevelDbLogReader.ChunkType; + +class LevelDbUtil { + + public static final int MAX_RECORD = LevelDbLogReader.BLOCK_SIZE - LevelDbLogReader.HEADER_SIZE; + + /** Adds a new record header to "bytes" at "pos", returns the new position. */ + private static int addRecordHeader(byte[] bytes, int pos, ChunkType type, int size) { + // Write a bogus checksum. + for (int i = 0; i < 4; ++i) { + bytes[pos++] = -1; + } + + // Write size and type. + bytes[pos++] = (byte) size; + bytes[pos++] = (byte) (size >> 8); + bytes[pos++] = (byte) type.getCode(); + + return pos; + } + + /** + * Adds a record of repeating bytes of 'val' of the given size to bytes at pos. + * + *

Writes the two least-significant bytes of 'val', alternating the order of them. So if the + * value of 'val' is 0x1234, writes 0x12 0x34 0x34 0x12 0x12 ... If the value is greater than + * 0xffff, it will be truncated to 16 bits. + * + *

This currently doesn't write a real checksum since we're not doing anything with that in the + * leveldb reader. + * + *

Returns the new offset for the next block. + */ + static int addRecord(byte[] bytes, int pos, ChunkType type, int size, int val) { + pos = addRecordHeader(bytes, pos, type, size); + + // Write "size" bytes of data. + for (int i = 0; i < size; ++i) { + bytes[pos + i] = (byte) val; + + // Swap the least significant bytes in val so we can have more than 256 different same-sized + // records. + val = ((val >> 8) & 0xff) | ((val & 0xff) << 8); + } + + return pos + size; + } + + /** + * Adds a record containing "data" to "bytes". + * + *

This currently doesn't write a real checksum since we're not doing anything with that in the + * leveldb reader. + * + *

Returns the new offset for the next block. + */ + static int addRecord(byte[] bytes, int pos, ChunkType type, byte[] data) { + checkArgument( + data.length < MAX_RECORD, "Record length (%d) > max record size (%d)", data.length, + MAX_RECORD); + pos = addRecordHeader(bytes, pos, type, data.length); + + // Write the contents of "data". + for (int i = 0; i < data.length; ++i) { + bytes[pos + i] = data[i]; + } + + return pos + data.length; + } +} diff --git a/javatests/google/registry/tools/RecordAccumulatorTest.java b/javatests/google/registry/tools/RecordAccumulatorTest.java new file mode 100644 index 00000000000..69c2172fef9 --- /dev/null +++ b/javatests/google/registry/tools/RecordAccumulatorTest.java @@ -0,0 +1,140 @@ +// Copyright 2017 The Nomulus Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package google.registry.tools; + +import static com.google.common.truth.Truth.assertThat; +import static google.registry.tools.LevelDbLogReader.BLOCK_SIZE; + +import com.google.appengine.api.datastore.Entity; +import com.google.appengine.api.datastore.EntityTranslator; +import com.google.auto.value.AutoValue; +import com.google.common.collect.ImmutableSet; +import com.google.storage.onestore.v3.OnestoreEntity.EntityProto; +import google.registry.testing.AppEngineRule; +import google.registry.tools.LevelDbLogReader.ChunkType; +import java.io.File; +import java.io.FileNotFoundException; +import java.io.FileOutputStream; +import java.io.IOException; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class RecordAccumulatorTest { + + private static final int BASE_ID = 1001; + private static final String TEST_ENTITY_KIND = "TestEntity"; + + @Rule public final TemporaryFolder tempFs = new TemporaryFolder(); + @Rule public final AppEngineRule appEngine = AppEngineRule.builder().withDatastore().build(); + + @Test + public void testReadDirectory() throws IOException { + File subdir = tempFs.newFolder("folder"); + LevelDbFileBuilder builder = new LevelDbFileBuilder(new File(subdir, "data1")); + + // Note that we need to specify property values as "Long" for property comparisons to work + // correctly because that's how they are deserialized from protos. + ComparableEntity e1 = + builder.addEntityProto( + BASE_ID, + Property.create("eeny", 100L), + Property.create("meeny", 200L), + Property.create("miney", 300L)); + ComparableEntity e2 = + builder.addEntityProto( + BASE_ID + 1, + Property.create("eeny", 100L), + Property.create("meeny", 200L), + Property.create("miney", 300L)); + builder.build(); + + builder = new LevelDbFileBuilder(new File(subdir, "data2")); + + // Duplicate of the record in the other file. + builder.addEntityProto( + BASE_ID, + Property.create("eeny", 100L), + Property.create("meeny", 200L), + Property.create("miney", 300L)); + + ComparableEntity e3 = + builder.addEntityProto( + BASE_ID + 2, + Property.create("moxy", 100L), + Property.create("fruvis", 200L), + Property.create("cortex", 300L)); + builder.build(); + + ImmutableSet entities = + new RecordAccumulator().readDirectory(subdir).getComparableEntitySet(); + assertThat(entities).containsExactly(e1, e2, e3); + } + + /** Utility class for building a leveldb logfile. */ + private static final class LevelDbFileBuilder { + private final FileOutputStream out; + private byte[] currentBlock = new byte[BLOCK_SIZE]; + + // Write position in the current block. + private int currentPos = 0; + + LevelDbFileBuilder(File file) throws FileNotFoundException { + out = new FileOutputStream(file); + } + + /** + * Adds a record containing a new entity protobuf to the file. + * + *

Returns the ComparableEntity object rather than "this" so that we can check for the + * presence of the entity in the result set. + */ + private ComparableEntity addEntityProto(int id, Property... properties) throws IOException { + Entity entity = new Entity(TEST_ENTITY_KIND, id); + for (Property prop : properties) { + entity.setProperty(prop.name(), prop.value()); + } + EntityProto proto = EntityTranslator.convertToPb(entity); + byte[] protoBytes = proto.toByteArray(); + if (protoBytes.length > BLOCK_SIZE - currentPos) { + out.write(currentBlock); + currentBlock = new byte[BLOCK_SIZE]; + } + + currentPos = LevelDbUtil.addRecord(currentBlock, currentPos, ChunkType.FULL, protoBytes); + return new ComparableEntity(entity); + } + + /** Writes all remaining data and closes the block. */ + void build() throws IOException { + out.write(currentBlock); + out.close(); + } + } + + @AutoValue + abstract static class Property { + static Property create(String name, Object value) { + return new AutoValue_RecordAccumulatorTest_Property(name, value); + } + + abstract String name(); + + abstract Object value(); + } +}