1 // Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
2 //  This source code is licensed under both the GPLv2 (found in the
3 //  COPYING file in the root directory) and Apache 2.0 License
4 //  (found in the LICENSE.Apache file in the root directory).
5 
6 package org.rocksdb;
7 
8 import static org.assertj.core.api.Assertions.assertThat;
9 import static org.junit.Assert.fail;
10 
11 import java.io.File;
12 import java.io.IOException;
13 import java.nio.ByteBuffer;
14 import java.util.ArrayList;
15 import java.util.Arrays;
16 import java.util.List;
17 import org.junit.ClassRule;
18 import org.junit.Rule;
19 import org.junit.Test;
20 import org.junit.rules.TemporaryFolder;
21 import org.rocksdb.util.BytewiseComparator;
22 
23 public class SstFileWriterTest {
24   private static final String SST_FILE_NAME = "test.sst";
25   private static final String DB_DIRECTORY_NAME = "test_db";
26 
27   @ClassRule
28   public static final RocksNativeLibraryResource ROCKS_NATIVE_LIBRARY_RESOURCE
29       = new RocksNativeLibraryResource();
30 
31   @Rule public TemporaryFolder parentFolder = new TemporaryFolder();
32 
33   enum OpType { PUT, PUT_BYTES, PUT_DIRECT, MERGE, MERGE_BYTES, DELETE, DELETE_BYTES }
34 
35   class KeyValueWithOp {
KeyValueWithOp(String key, String value, OpType opType)36     KeyValueWithOp(String key, String value, OpType opType) {
37       this.key = key;
38       this.value = value;
39       this.opType = opType;
40     }
41 
getKey()42     String getKey() {
43       return key;
44     }
45 
getValue()46     String getValue() {
47       return value;
48     }
49 
getOpType()50     OpType getOpType() {
51       return opType;
52     }
53 
54     private String key;
55     private String value;
56     private OpType opType;
57   };
58 
newSstFile(final List<KeyValueWithOp> keyValues, boolean useJavaBytewiseComparator)59   private File newSstFile(final List<KeyValueWithOp> keyValues,
60       boolean useJavaBytewiseComparator) throws IOException, RocksDBException {
61     final EnvOptions envOptions = new EnvOptions();
62     final StringAppendOperator stringAppendOperator = new StringAppendOperator();
63     final Options options = new Options().setMergeOperator(stringAppendOperator);
64     SstFileWriter sstFileWriter = null;
65     ComparatorOptions comparatorOptions = null;
66     BytewiseComparator comparator = null;
67     if (useJavaBytewiseComparator) {
68       comparatorOptions = new ComparatorOptions().setUseDirectBuffer(false);
69       comparator = new BytewiseComparator(comparatorOptions);
70       options.setComparator(comparator);
71       sstFileWriter = new SstFileWriter(envOptions, options);
72     } else {
73       sstFileWriter = new SstFileWriter(envOptions, options);
74     }
75 
76     final File sstFile = parentFolder.newFile(SST_FILE_NAME);
77     try {
78       sstFileWriter.open(sstFile.getAbsolutePath());
79       assertThat(sstFileWriter.fileSize()).isEqualTo(0);
80       for (KeyValueWithOp keyValue : keyValues) {
81         Slice keySlice = new Slice(keyValue.getKey());
82         Slice valueSlice = new Slice(keyValue.getValue());
83         byte[] keyBytes = keyValue.getKey().getBytes();
84         byte[] valueBytes = keyValue.getValue().getBytes();
85         ByteBuffer keyDirect = ByteBuffer.allocateDirect(keyBytes.length);
86         keyDirect.put(keyBytes);
87         keyDirect.flip();
88         ByteBuffer valueDirect = ByteBuffer.allocateDirect(valueBytes.length);
89         valueDirect.put(valueBytes);
90         valueDirect.flip();
91         switch (keyValue.getOpType()) {
92           case PUT:
93             sstFileWriter.put(keySlice, valueSlice);
94             break;
95           case PUT_BYTES:
96             sstFileWriter.put(keyBytes, valueBytes);
97             break;
98           case PUT_DIRECT:
99             sstFileWriter.put(keyDirect, valueDirect);
100             assertThat(keyDirect.position()).isEqualTo(keyBytes.length);
101             assertThat(keyDirect.limit()).isEqualTo(keyBytes.length);
102             assertThat(valueDirect.position()).isEqualTo(valueBytes.length);
103             assertThat(valueDirect.limit()).isEqualTo(valueBytes.length);
104             break;
105           case MERGE:
106             sstFileWriter.merge(keySlice, valueSlice);
107             break;
108           case MERGE_BYTES:
109             sstFileWriter.merge(keyBytes, valueBytes);
110             break;
111           case DELETE:
112             sstFileWriter.delete(keySlice);
113             break;
114           case DELETE_BYTES:
115             sstFileWriter.delete(keyBytes);
116             break;
117           default:
118             fail("Unsupported op type");
119         }
120         keySlice.close();
121         valueSlice.close();
122       }
123       sstFileWriter.finish();
124       assertThat(sstFileWriter.fileSize()).isGreaterThan(100);
125     } finally {
126       assertThat(sstFileWriter).isNotNull();
127       sstFileWriter.close();
128       options.close();
129       envOptions.close();
130       if (comparatorOptions != null) {
131         comparatorOptions.close();
132       }
133       if (comparator != null) {
134         comparator.close();
135       }
136     }
137     return sstFile;
138   }
139 
140   @Test
generateSstFileWithJavaComparator()141   public void generateSstFileWithJavaComparator()
142       throws RocksDBException, IOException {
143     final List<KeyValueWithOp> keyValues = new ArrayList<>();
144     keyValues.add(new KeyValueWithOp("key1", "value1", OpType.PUT));
145     keyValues.add(new KeyValueWithOp("key2", "value2", OpType.PUT));
146     keyValues.add(new KeyValueWithOp("key3", "value3", OpType.MERGE));
147     keyValues.add(new KeyValueWithOp("key4", "value4", OpType.MERGE));
148     keyValues.add(new KeyValueWithOp("key5", "", OpType.DELETE));
149 
150     newSstFile(keyValues, true);
151   }
152 
153   @Test
generateSstFileWithNativeComparator()154   public void generateSstFileWithNativeComparator()
155       throws RocksDBException, IOException {
156     final List<KeyValueWithOp> keyValues = new ArrayList<>();
157     keyValues.add(new KeyValueWithOp("key1", "value1", OpType.PUT));
158     keyValues.add(new KeyValueWithOp("key2", "value2", OpType.PUT));
159     keyValues.add(new KeyValueWithOp("key3", "value3", OpType.MERGE));
160     keyValues.add(new KeyValueWithOp("key4", "value4", OpType.MERGE));
161     keyValues.add(new KeyValueWithOp("key5", "", OpType.DELETE));
162 
163     newSstFile(keyValues, false);
164   }
165 
166   @Test
ingestSstFile()167   public void ingestSstFile() throws RocksDBException, IOException {
168     final List<KeyValueWithOp> keyValues = new ArrayList<>();
169     keyValues.add(new KeyValueWithOp("key1", "value1", OpType.PUT));
170     keyValues.add(new KeyValueWithOp("key2", "value2", OpType.PUT_DIRECT));
171     keyValues.add(new KeyValueWithOp("key3", "value3", OpType.PUT_BYTES));
172     keyValues.add(new KeyValueWithOp("key4", "value4", OpType.MERGE));
173     keyValues.add(new KeyValueWithOp("key5", "value5", OpType.MERGE_BYTES));
174     keyValues.add(new KeyValueWithOp("key6", "", OpType.DELETE));
175     keyValues.add(new KeyValueWithOp("key7", "", OpType.DELETE));
176 
177 
178     final File sstFile = newSstFile(keyValues, false);
179     final File dbFolder = parentFolder.newFolder(DB_DIRECTORY_NAME);
180     try(final StringAppendOperator stringAppendOperator =
181             new StringAppendOperator();
182         final Options options = new Options()
183             .setCreateIfMissing(true)
184             .setMergeOperator(stringAppendOperator);
185         final RocksDB db = RocksDB.open(options, dbFolder.getAbsolutePath());
186         final IngestExternalFileOptions ingestExternalFileOptions =
187             new IngestExternalFileOptions()) {
188       db.ingestExternalFile(Arrays.asList(sstFile.getAbsolutePath()),
189           ingestExternalFileOptions);
190 
191       assertThat(db.get("key1".getBytes())).isEqualTo("value1".getBytes());
192       assertThat(db.get("key2".getBytes())).isEqualTo("value2".getBytes());
193       assertThat(db.get("key3".getBytes())).isEqualTo("value3".getBytes());
194       assertThat(db.get("key4".getBytes())).isEqualTo("value4".getBytes());
195       assertThat(db.get("key5".getBytes())).isEqualTo("value5".getBytes());
196       assertThat(db.get("key6".getBytes())).isEqualTo(null);
197       assertThat(db.get("key7".getBytes())).isEqualTo(null);
198     }
199   }
200 
201   @Test
ingestSstFile_cf()202   public void ingestSstFile_cf() throws RocksDBException, IOException {
203     final List<KeyValueWithOp> keyValues = new ArrayList<>();
204     keyValues.add(new KeyValueWithOp("key1", "value1", OpType.PUT));
205     keyValues.add(new KeyValueWithOp("key2", "value2", OpType.PUT));
206     keyValues.add(new KeyValueWithOp("key3", "value3", OpType.MERGE));
207     keyValues.add(new KeyValueWithOp("key4", "", OpType.DELETE));
208 
209     final File sstFile = newSstFile(keyValues, false);
210     final File dbFolder = parentFolder.newFolder(DB_DIRECTORY_NAME);
211     try(final StringAppendOperator stringAppendOperator =
212             new StringAppendOperator();
213         final Options options = new Options()
214             .setCreateIfMissing(true)
215             .setCreateMissingColumnFamilies(true)
216             .setMergeOperator(stringAppendOperator);
217         final RocksDB db = RocksDB.open(options, dbFolder.getAbsolutePath());
218         final IngestExternalFileOptions ingestExternalFileOptions =
219             new IngestExternalFileOptions()) {
220 
221       try(final ColumnFamilyOptions cf_opts = new ColumnFamilyOptions()
222               .setMergeOperator(stringAppendOperator);
223           final ColumnFamilyHandle cf_handle = db.createColumnFamily(
224               new ColumnFamilyDescriptor("new_cf".getBytes(), cf_opts))) {
225 
226         db.ingestExternalFile(cf_handle,
227             Arrays.asList(sstFile.getAbsolutePath()),
228             ingestExternalFileOptions);
229 
230         assertThat(db.get(cf_handle,
231             "key1".getBytes())).isEqualTo("value1".getBytes());
232         assertThat(db.get(cf_handle,
233             "key2".getBytes())).isEqualTo("value2".getBytes());
234         assertThat(db.get(cf_handle,
235             "key3".getBytes())).isEqualTo("value3".getBytes());
236         assertThat(db.get(cf_handle,
237             "key4".getBytes())).isEqualTo(null);
238       }
239     }
240   }
241 }
242