Support SST Decoder

pom.xml

@@ -61,6 +61,7 @@
         <maven.compiler.target>1.8</maven.compiler.target>
         <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
         <project.reporting.outputEncoding>UTF-8</project.reporting.outputEncoding>
+        <rocksdb.version>6.22.1.1</rocksdb.version>
         <protobuf.version>3.5.1</protobuf.version>
         <log4j.version>1.2.17</log4j.version>
         <slf4j.version>1.7.16</slf4j.version>
@@ -79,6 +80,11 @@
     </properties>
 
     <dependencies>
+        <dependency>
+            <groupId>org.rocksdb</groupId>
+            <artifactId>rocksdbjni</artifactId>
+            <version>${rocksdb.version}</version>
+        </dependency>
         <dependency>
             <groupId>org.antlr</groupId>
             <artifactId>antlr4-runtime</artifactId>

scripts/proto.sh

@@ -14,7 +14,7 @@
 #   limitations under the License.
 #
 
-kvproto_hash=2ac2a7984b2d01b96ed56fd8474f4bf80fa33c51
+kvproto_hash=465fa4c7b42e644d5aacaf79d06c75733dc12eb3
 raft_rs_hash=b9891b673573fad77ebcf9bbe0969cf945841926
 tipb_hash=c4d518eb1d60c21f05b028b36729e64610346dac
 

src/main/java/org/tikv/br/BackupDecoder.java

@@ -0,0 +1,64 @@
+/*
+ *
+ * Copyright 2021 PingCAP, Inc.
+ *
+ * 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,
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package org.tikv.br;
+
+import org.rocksdb.Options;
+import org.rocksdb.ReadOptions;
+import org.tikv.common.exception.SSTDecodeException;
+import org.tikv.kvproto.Brpb;
+
+public class BackupDecoder {
+  private final Brpb.BackupMeta backupMeta;
+  private final boolean ttlEnabled;
+  private final KVDecoder kvDecoder;
+
+  public BackupDecoder(Brpb.BackupMeta backupMeta) throws SSTDecodeException {
+    this.backupMeta = backupMeta;
+    this.ttlEnabled = false;
+    this.kvDecoder = initKVDecoder();
+  }
+
+  public BackupDecoder(Brpb.BackupMeta backupMeta, boolean ttlEnabled) throws SSTDecodeException {
+    this.backupMeta = backupMeta;
+    this.ttlEnabled = ttlEnabled;
+    this.kvDecoder = initKVDecoder();
+  }
+
+  private KVDecoder initKVDecoder() throws SSTDecodeException {
+    // Currently only v1 is supported.
+    // V2 will be added after https://github.com/tikv/tikv/issues/10938.
+    if (backupMeta.getIsRawKv()) {
+      // TODO: ttl_enable should be witten to BackupMeta
+      return new RawKVDecoderV1(ttlEnabled);
+    } else {
+      throw new SSTDecodeException("TxnKV is not supported yet!");
+    }
+  }
+
+  public SSTDecoder decodeSST(String sstFilePath) {
+    return decodeSST(sstFilePath, new Options(), new ReadOptions());
+  }
+
+  public SSTDecoder decodeSST(String sstFilePath, Options options, ReadOptions readOptions) {
+    return new SSTDecoder(sstFilePath, kvDecoder, options, readOptions);
+  }
+
+  public Brpb.BackupMeta getBackupMeta() {
+    return backupMeta;
+  }
+}

src/main/java/org/tikv/br/BackupMetaDecoder.java

@@ -0,0 +1,41 @@
+/*
+ *
+ * Copyright 2021 PingCAP, Inc.
+ *
+ * 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,
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package org.tikv.br;
+
+import com.google.protobuf.InvalidProtocolBufferException;
+import java.io.File;
+import java.io.IOException;
+import java.nio.file.Files;
+import org.tikv.kvproto.Brpb;
+
+public class BackupMetaDecoder {
+  private final Brpb.BackupMeta backupMeta;
+
+  public BackupMetaDecoder(byte[] data) throws InvalidProtocolBufferException {
+    this.backupMeta = Brpb.BackupMeta.parseFrom(data);
+  }
+
+  public Brpb.BackupMeta getBackupMeta() {
+    return backupMeta;
+  }
+
+  public static BackupMetaDecoder parse(String backupMetaFilePath) throws IOException {
+    byte[] data = Files.readAllBytes(new File(backupMetaFilePath).toPath());
+    return new BackupMetaDecoder(data);
+  }
+}

src/main/java/org/tikv/br/KVDecoder.java

@@ -0,0 +1,26 @@
+/*
+ *
+ * Copyright 2021 PingCAP, Inc.
+ *
+ * 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,
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package org.tikv.br;
+
+import com.google.protobuf.ByteString;
+
+public interface KVDecoder {
+  ByteString decodeKey(byte[] key);
+
+  ByteString decodeValue(byte[] value);
+}

src/main/java/org/tikv/br/RawKVDecoderV1.java

@@ -0,0 +1,56 @@
+/*
+ *
+ * Copyright 2021 PingCAP, Inc.
+ *
+ * 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,
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package org.tikv.br;
+
+import com.google.protobuf.ByteString;
+import java.util.Arrays;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class RawKVDecoderV1 implements KVDecoder {
+  private static final Logger logger = LoggerFactory.getLogger(SSTIterator.class);
+
+  private final boolean ttlEnabled;
+
+  public RawKVDecoderV1(boolean ttlEnabled) {
+    this.ttlEnabled = ttlEnabled;
+  }
+
+  @Override
+  public ByteString decodeKey(byte[] key) {
+    if (key == null || key.length == 0) {
+      logger.warn(
+          "skip Key-Value pair because key == null || key.length == 0, key = "
+              + Arrays.toString(key));
+      return null;
+    } else if (key[0] != 'z') {
+      logger.warn("skip Key-Value pair because key[0] != 'z', key = " + Arrays.toString(key));
+      return null;
+    }
+    return ByteString.copyFrom(key, 1, key.length - 1);
+  }
+
+  @Override
+  public ByteString decodeValue(byte[] value) {
+    if (!ttlEnabled) {
+      return ByteString.copyFrom(value);
+    } else {
+      return ByteString.copyFrom(value).substring(0, value.length - 8);
+    }
+  }
+}

src/main/java/org/tikv/br/SSTDecoder.java

@@ -0,0 +1,93 @@
+/*
+ *
+ * Copyright 2021 PingCAP, Inc.
+ *
+ * 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,
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package org.tikv.br;
+
+import com.google.protobuf.ByteString;
+import java.util.Iterator;
+import org.rocksdb.Options;
+import org.rocksdb.ReadOptions;
+import org.rocksdb.RocksDBException;
+import org.rocksdb.SstFileReader;
+import org.rocksdb.SstFileReaderIterator;
+import org.tikv.common.util.Pair;
+
+public class SSTDecoder {
+  private final String filePath;
+  private final KVDecoder kvDecoder;
+  private final Options options;
+  private final ReadOptions readOptions;
+
+  private SstFileReader sstFileReader;
+  private SstFileReaderIterator iterator;
+
+  public SSTDecoder(String sstFilePath, KVDecoder kvDecoder) {
+    this.filePath = sstFilePath;
+    this.kvDecoder = kvDecoder;
+    this.options = new Options();
+    this.readOptions = new ReadOptions();
+  }
+
+  public SSTDecoder(
+      String filePath, KVDecoder kvDecoder, Options options, ReadOptions readOptions) {
+    this.filePath = filePath;
+    this.kvDecoder = kvDecoder;
+    this.options = options;
+    this.readOptions = readOptions;
+  }
+
+  public synchronized Iterator<Pair<ByteString, ByteString>> getIterator() throws RocksDBException {
+    if (sstFileReader != null || iterator != null) {
+      throw new RocksDBException("File already opened!");
+    }
+
+    sstFileReader = new SstFileReader(new Options());
+    sstFileReader.open(filePath);
+    iterator = sstFileReader.newIterator(new ReadOptions());
+    return new SSTIterator(iterator, kvDecoder);
+  }
+
+  public synchronized void close() {
+    try {
+      if (iterator != null) {
+        iterator.close();
+      }
+    } finally {
+      iterator = null;
+    }
+
+    try {
+      if (sstFileReader != null) {
+        sstFileReader.close();
+      }
+    } finally {
+      sstFileReader = null;
+    }
+  }
+
+  public String getFilePath() {
+    return filePath;
+  }
+
+  public Options getOptions() {
+    return options;
+  }
+
+  public ReadOptions getReadOptions() {
+    return readOptions;
+  }
+}

src/main/java/org/tikv/br/SSTIterator.java

@@ -0,0 +1,64 @@
+/*
+ *
+ * Copyright 2021 PingCAP, Inc.
+ *
+ * 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,
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package org.tikv.br;
+
+import com.google.protobuf.ByteString;
+import java.util.Iterator;
+import org.rocksdb.SstFileReaderIterator;
+import org.tikv.common.util.Pair;
+
+public class SSTIterator implements Iterator<Pair<ByteString, ByteString>> {
+  private final SstFileReaderIterator iterator;
+  private final KVDecoder kvDecoder;
+
+  private Pair<ByteString, ByteString> nextPair;
+
+  public SSTIterator(SstFileReaderIterator iterator, KVDecoder kvDecoder) {
+    this.iterator = iterator;
+    this.kvDecoder = kvDecoder;
+    this.iterator.seekToFirst();
+    this.nextPair = processNext();
+  }
+
+  @Override
+  public boolean hasNext() {
+    return nextPair != null;
+  }
+
+  @Override
+  public Pair<ByteString, ByteString> next() {
+    Pair<ByteString, ByteString> result = nextPair;
+    nextPair = processNext();
+    return result;
+  }
+
+  private Pair<ByteString, ByteString> processNext() {
+    if (iterator.isValid()) {
+      ByteString key = kvDecoder.decodeKey(iterator.key());
+      ByteString value = kvDecoder.decodeValue(iterator.value());
+      iterator.next();
+      if (key != null) {
+        return Pair.create(key, value);
+      } else {
+        return processNext();
+      }
+    } else {
+      return null;
+    }
+  }
+}