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 package org.rocksdb;
6 
7 import static org.assertj.core.api.Assertions.assertThat;
8 
9 import java.nio.ByteBuffer;
10 import org.junit.ClassRule;
11 import org.junit.Rule;
12 import org.junit.Test;
13 import org.junit.rules.TemporaryFolder;
14 
15 public class RocksIteratorTest {
16 
17   @ClassRule
18   public static final RocksNativeLibraryResource ROCKS_NATIVE_LIBRARY_RESOURCE =
19       new RocksNativeLibraryResource();
20 
21   @Rule
22   public TemporaryFolder dbFolder = new TemporaryFolder();
23 
24   @Test
rocksIterator()25   public void rocksIterator() throws RocksDBException {
26     try (final Options options = new Options()
27         .setCreateIfMissing(true)
28         .setCreateMissingColumnFamilies(true);
29          final RocksDB db = RocksDB.open(options,
30              dbFolder.getRoot().getAbsolutePath())) {
31       db.put("key1".getBytes(), "value1".getBytes());
32       db.put("key2".getBytes(), "value2".getBytes());
33 
34       try (final RocksIterator iterator = db.newIterator()) {
35         iterator.seekToFirst();
36         assertThat(iterator.isValid()).isTrue();
37         assertThat(iterator.key()).isEqualTo("key1".getBytes());
38         assertThat(iterator.value()).isEqualTo("value1".getBytes());
39 
40         ByteBuffer key = ByteBuffer.allocateDirect(2);
41         ByteBuffer value = ByteBuffer.allocateDirect(2);
42         assertThat(iterator.key(key)).isEqualTo(4);
43         assertThat(iterator.value(value)).isEqualTo(6);
44 
45         assertThat(key.position()).isEqualTo(0);
46         assertThat(key.limit()).isEqualTo(2);
47         assertThat(value.position()).isEqualTo(0);
48         assertThat(value.limit()).isEqualTo(2);
49 
50         byte[] tmp = new byte[2];
51         key.get(tmp);
52         assertThat(tmp).isEqualTo("ke".getBytes());
53         value.get(tmp);
54         assertThat(tmp).isEqualTo("va".getBytes());
55 
56         key = ByteBuffer.allocateDirect(12);
57         value = ByteBuffer.allocateDirect(12);
58         assertThat(iterator.key(key)).isEqualTo(4);
59         assertThat(iterator.value(value)).isEqualTo(6);
60         assertThat(key.position()).isEqualTo(0);
61         assertThat(key.limit()).isEqualTo(4);
62         assertThat(value.position()).isEqualTo(0);
63         assertThat(value.limit()).isEqualTo(6);
64 
65         tmp = new byte[4];
66         key.get(tmp);
67         assertThat(tmp).isEqualTo("key1".getBytes());
68         tmp = new byte[6];
69         value.get(tmp);
70         assertThat(tmp).isEqualTo("value1".getBytes());
71 
72         iterator.next();
73         assertThat(iterator.isValid()).isTrue();
74         assertThat(iterator.key()).isEqualTo("key2".getBytes());
75         assertThat(iterator.value()).isEqualTo("value2".getBytes());
76         iterator.next();
77         assertThat(iterator.isValid()).isFalse();
78         iterator.seekToLast();
79         iterator.prev();
80         assertThat(iterator.isValid()).isTrue();
81         assertThat(iterator.key()).isEqualTo("key1".getBytes());
82         assertThat(iterator.value()).isEqualTo("value1".getBytes());
83         iterator.seekToFirst();
84         iterator.seekToLast();
85         assertThat(iterator.isValid()).isTrue();
86         assertThat(iterator.key()).isEqualTo("key2".getBytes());
87         assertThat(iterator.value()).isEqualTo("value2".getBytes());
88         iterator.status();
89 
90         key.clear();
91         key.put("key1".getBytes());
92         key.flip();
93         iterator.seek(key);
94         assertThat(iterator.isValid()).isTrue();
95         assertThat(iterator.value()).isEqualTo("value1".getBytes());
96         assertThat(key.position()).isEqualTo(4);
97         assertThat(key.limit()).isEqualTo(4);
98 
99         key.clear();
100         key.put("key2".getBytes());
101         key.flip();
102         iterator.seekForPrev(key);
103         assertThat(iterator.isValid()).isTrue();
104         assertThat(iterator.value()).isEqualTo("value2".getBytes());
105         assertThat(key.position()).isEqualTo(4);
106         assertThat(key.limit()).isEqualTo(4);
107       }
108 
109       try (final RocksIterator iterator = db.newIterator()) {
110         iterator.seek("key0".getBytes());
111         assertThat(iterator.isValid()).isTrue();
112         assertThat(iterator.key()).isEqualTo("key1".getBytes());
113 
114         iterator.seek("key1".getBytes());
115         assertThat(iterator.isValid()).isTrue();
116         assertThat(iterator.key()).isEqualTo("key1".getBytes());
117 
118         iterator.seek("key1.5".getBytes());
119         assertThat(iterator.isValid()).isTrue();
120         assertThat(iterator.key()).isEqualTo("key2".getBytes());
121 
122         iterator.seek("key2".getBytes());
123         assertThat(iterator.isValid()).isTrue();
124         assertThat(iterator.key()).isEqualTo("key2".getBytes());
125 
126         iterator.seek("key3".getBytes());
127         assertThat(iterator.isValid()).isFalse();
128       }
129 
130       try (final RocksIterator iterator = db.newIterator()) {
131         iterator.seekForPrev("key0".getBytes());
132         assertThat(iterator.isValid()).isFalse();
133 
134         iterator.seekForPrev("key1".getBytes());
135         assertThat(iterator.isValid()).isTrue();
136         assertThat(iterator.key()).isEqualTo("key1".getBytes());
137 
138         iterator.seekForPrev("key1.5".getBytes());
139         assertThat(iterator.isValid()).isTrue();
140         assertThat(iterator.key()).isEqualTo("key1".getBytes());
141 
142         iterator.seekForPrev("key2".getBytes());
143         assertThat(iterator.isValid()).isTrue();
144         assertThat(iterator.key()).isEqualTo("key2".getBytes());
145 
146         iterator.seekForPrev("key3".getBytes());
147         assertThat(iterator.isValid()).isTrue();
148         assertThat(iterator.key()).isEqualTo("key2".getBytes());
149       }
150     }
151   }
152 
153   @Test
rocksIteratorReleaseAfterCfClose()154   public void rocksIteratorReleaseAfterCfClose() throws RocksDBException {
155     try (final Options options = new Options()
156         .setCreateIfMissing(true)
157         .setCreateMissingColumnFamilies(true);
158          final RocksDB db = RocksDB.open(options,
159              this.dbFolder.getRoot().getAbsolutePath())) {
160       db.put("key".getBytes(), "value".getBytes());
161 
162       // Test case: release iterator after default CF close
163       try (final RocksIterator iterator = db.newIterator()) {
164         // In fact, calling close() on default CF has no effect
165         db.getDefaultColumnFamily().close();
166 
167         iterator.seekToFirst();
168         assertThat(iterator.isValid()).isTrue();
169         assertThat(iterator.key()).isEqualTo("key".getBytes());
170         assertThat(iterator.value()).isEqualTo("value".getBytes());
171       }
172 
173       // Test case: release iterator after custom CF close
174       ColumnFamilyDescriptor cfd1 = new ColumnFamilyDescriptor("cf1".getBytes());
175       ColumnFamilyHandle cfHandle1 = db.createColumnFamily(cfd1);
176       db.put(cfHandle1, "key1".getBytes(), "value1".getBytes());
177 
178       try (final RocksIterator iterator = db.newIterator(cfHandle1)) {
179         cfHandle1.close();
180 
181         iterator.seekToFirst();
182         assertThat(iterator.isValid()).isTrue();
183         assertThat(iterator.key()).isEqualTo("key1".getBytes());
184         assertThat(iterator.value()).isEqualTo("value1".getBytes());
185       }
186 
187       // Test case: release iterator after custom CF drop & close
188       ColumnFamilyDescriptor cfd2 = new ColumnFamilyDescriptor("cf2".getBytes());
189       ColumnFamilyHandle cfHandle2 = db.createColumnFamily(cfd2);
190       db.put(cfHandle2, "key2".getBytes(), "value2".getBytes());
191 
192       try (final RocksIterator iterator = db.newIterator(cfHandle2)) {
193         db.dropColumnFamily(cfHandle2);
194         cfHandle2.close();
195 
196         iterator.seekToFirst();
197         assertThat(iterator.isValid()).isTrue();
198         assertThat(iterator.key()).isEqualTo("key2".getBytes());
199         assertThat(iterator.value()).isEqualTo("value2".getBytes());
200       }
201     }
202   }
203 }
204