/**
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you 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 com.alibaba.jstorm.cache;
import static org.assertj.core.api.Assertions.assertThat;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.concurrent.TimeUnit;
import junit.framework.Assert;
import org.rocksdb.ColumnFamilyDescriptor;
import org.rocksdb.ColumnFamilyHandle;
import org.rocksdb.DBOptions;
import org.rocksdb.RocksDB;
import org.rocksdb.RocksDBException;
import org.rocksdb.TtlDB;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import com.alibaba.jstorm.utils.JStormUtils;
import com.alibaba.jstorm.utils.OSInfo;
import com.alibaba.jstorm.utils.RunCounter;
public class RocksDBTest {
private static Logger LOG = LoggerFactory.getLogger(RocksDBTest.class);
public static String rootDir = "rocksdb_test";
// ATTENTIME
// DUE to windows can't load RocksDB jni libray,
// So these junit tests can't run under window
//
RocksTTLDBCache cache = new RocksTTLDBCache();
{
Map conf = new HashMap();
// TemporaryFolder cacheRoot = new TemporaryFolder()
conf.put(RocksDBCache.ROCKSDB_ROOT_DIR, rootDir);
conf.put(RocksDBCache.ROCKSDB_RESET, true);
List<Integer> list = new ArrayList<Integer>() {
{
add(60);
add(120);
}
};
conf.put(RocksDBCache.TAG_TIMEOUT_LIST, list);
try {
cache.init(conf);
} catch (Exception e) {
// TODO Auto-generated catch block
e.printStackTrace();
}
}
public void testSimple() {
List<Map<String, Object>> dataList = new ArrayList<Map<String, Object>>();
for (int i = 0; i < 4; i++) {
dataList.add(new HashMap<String, Object>());
}
for (int i = 0; i < 10; i++) {
String key = "key" + i;
for (int j = 0; j < 4; j++) {
Map<String, Object> data = dataList.get(j);
String value = "value--" + i + "--" + j;
data.put(key, value);
}
}
Map<String, Object> data = dataList.get(0);
for (Entry<String, Object> entry : data.entrySet()) {
cache.put(entry.getKey(), entry.getValue());
}
for (Entry<String, Object> entry : data.entrySet()) {
String rawValue = (String) entry.getValue();
String fetchValue = (String) cache.get(entry.getKey());
Assert.assertEquals(rawValue, fetchValue);
}
for (Entry<String, Object> entry : data.entrySet()) {
cache.remove(entry.getKey());
}
for (Entry<String, Object> entry : data.entrySet()) {
String fetchValue = (String) cache.get(entry.getKey());
Assert.assertEquals(fetchValue, null);
}
Map<String, Object> data1 = dataList.get(1);
cache.putBatch(data1);
Map<String, Object> getBatch = new HashMap<String, Object>();
for (Entry<String, Object> entry : data1.entrySet()) {
getBatch.put(entry.getKey(), null);
}
Map<String, Object> emptyBatch = new HashMap<String, Object>();
emptyBatch.putAll(getBatch);
cache.getBatch(getBatch);
for (String key : getBatch.keySet()) {
String rawValue = (String) data1.get(key);
String fetchValue = (String) getBatch.get(key);
Assert.assertEquals(rawValue, fetchValue);
}
cache.removeBatch(data1.keySet());
getBatch.clear();
getBatch.putAll(emptyBatch);
cache.getBatch(getBatch);
for (String key : getBatch.keySet()) {
String rawValue = null;
String fetchValue = (String) getBatch.get(key);
Assert.assertEquals(rawValue, fetchValue);
}
Map<String, Object> data2 = dataList.get(2);
for (Entry<String, Object> entry : data2.entrySet()) {
cache.put(entry.getKey(), entry.getValue(), 60);
}
for (Entry<String, Object> entry : data2.entrySet()) {
String rawValue = (String) entry.getValue();
String fetchValue = (String) cache.get(entry.getKey());
Assert.assertEquals(rawValue, fetchValue);
}
for (Entry<String, Object> entry : data2.entrySet()) {
cache.remove(entry.getKey());
}
for (Entry<String, Object> entry : data2.entrySet()) {
String fetchValue = (String) cache.get(entry.getKey());
Assert.assertEquals(fetchValue, null);
}
Map<String, Object> data3 = dataList.get(3);
cache.putBatch(data3, 60);
for (Entry<String, Object> entry : data3.entrySet()) {
getBatch.put(entry.getKey(), null);
}
cache.getBatch(getBatch);
for (String key : getBatch.keySet()) {
String rawValue = (String) data3.get(key);
String fetchValue = (String) getBatch.get(key);
Assert.assertEquals(rawValue, fetchValue);
}
cache.removeBatch(data3.keySet());
getBatch.clear();
getBatch.putAll(emptyBatch);
cache.getBatch(getBatch);
for (String key : getBatch.keySet()) {
String rawValue = null;
String fetchValue = (String) getBatch.get(key);
Assert.assertEquals(rawValue, fetchValue);
}
for (Entry<String, Object> entry : data.entrySet()) {
cache.put(entry.getKey(), entry.getValue());
}
cache.putBatch(data3, 60);
for (Entry<String, Object> entry : data3.entrySet()) {
String rawValue = (String) entry.getValue();
String fetchValue = (String) cache.get(entry.getKey());
Assert.assertEquals(rawValue, fetchValue);
}
cache.removeBatch(data3.keySet());
for (Entry<String, Object> entry : data3.entrySet()) {
String rawValue = null;
String fetchValue = (String) cache.get(entry.getKey());
Assert.assertEquals(rawValue, fetchValue);
}
// double delete
cache.removeBatch(data3.keySet());
}
public void visitorAccross() throws RocksDBException, InterruptedException {
DBOptions dbOptions = null;
TtlDB ttlDB = null;
List<ColumnFamilyDescriptor> cfNames = new ArrayList<ColumnFamilyDescriptor>();
List<ColumnFamilyHandle> columnFamilyHandleList = new ArrayList<ColumnFamilyHandle>();
cfNames.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY));
cfNames.add(new ColumnFamilyDescriptor("new_cf".getBytes()));
List<Integer> ttlValues = new ArrayList<Integer>();
// new column family with 1 second ttl
ttlValues.add(1);
// Default column family with infinite lifetime
ttlValues.add(0);
try {
System.out.println("Begin to open db");
dbOptions = new DBOptions().setCreateMissingColumnFamilies(true).setCreateIfMissing(true);
ttlDB = TtlDB.open(dbOptions, rootDir, cfNames, columnFamilyHandleList, ttlValues, false);
System.out.println("Successfully open db " + rootDir);
List<String> keys = new ArrayList<String>();
keys.add("key");
ttlDB.put("key".getBytes(), "key".getBytes());
for (int i = 0; i < 2; i++) {
String key = "key" + i;
keys.add(key);
ttlDB.put(columnFamilyHandleList.get(i), key.getBytes(), key.getBytes());
}
try {
byte[] value = ttlDB.get("others".getBytes());
if (value != null) {
System.out.println("Raw get :" + new String(value));
} else {
System.out.println("No value of other");
}
} catch (Exception e) {
System.out.println("Occur exception other");
}
for (String key : keys) {
try {
byte[] value = ttlDB.get(key.getBytes());
if (value != null) {
System.out.println("Raw get :" + new String(value));
} else {
System.out.println("No value of " + key);
}
} catch (Exception e) {
System.out.println("Occur exception " + key + ", Raw");
}
for (int i = 0; i < 2; i++) {
try {
byte[] value = ttlDB.get(columnFamilyHandleList.get(i), key.getBytes());
if (value != null) {
System.out.println("handler index" + i + " get :" + new String(value));
} else {
System.out.println("No value of index" + i + " get :" + key);
}
} catch (Exception e) {
System.out.println("Occur exception " + key + ", handler index:" + i);
}
}
}
} finally {
for (ColumnFamilyHandle columnFamilyHandle : columnFamilyHandleList) {
columnFamilyHandle.dispose();
}
if (ttlDB != null) {
ttlDB.close();
}
if (dbOptions != null) {
dbOptions.dispose();
}
}
}
public void ttlDbOpenWithColumnFamilies() throws RocksDBException, InterruptedException {
DBOptions dbOptions = null;
TtlDB ttlDB = null;
List<ColumnFamilyDescriptor> cfNames = new ArrayList<ColumnFamilyDescriptor>();
List<ColumnFamilyHandle> columnFamilyHandleList = new ArrayList<ColumnFamilyHandle>();
cfNames.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY));
cfNames.add(new ColumnFamilyDescriptor("new_cf".getBytes()));
List<Integer> ttlValues = new ArrayList<Integer>();
// Default column family with infinite lifetime
ttlValues.add(0);
// new column family with 1 second ttl
ttlValues.add(1);
try {
System.out.println("Begin to open db");
dbOptions = new DBOptions().setCreateMissingColumnFamilies(true).setCreateIfMissing(true);
ttlDB = TtlDB.open(dbOptions, rootDir, cfNames, columnFamilyHandleList, ttlValues, false);
System.out.println("Successfully open db " + rootDir);
ttlDB.put("key".getBytes(), "value".getBytes());
assertThat(ttlDB.get("key".getBytes())).isEqualTo("value".getBytes());
ttlDB.put(columnFamilyHandleList.get(1), "key".getBytes(), "value".getBytes());
assertThat(ttlDB.get(columnFamilyHandleList.get(1), "key".getBytes())).isEqualTo("value".getBytes());
TimeUnit.SECONDS.sleep(2);
ttlDB.compactRange();
ttlDB.compactRange(columnFamilyHandleList.get(1));
assertThat(ttlDB.get("key".getBytes())).isNotNull();
assertThat(ttlDB.get(columnFamilyHandleList.get(1), "key".getBytes())).isNull();
} finally {
for (ColumnFamilyHandle columnFamilyHandle : columnFamilyHandleList) {
columnFamilyHandle.dispose();
}
if (ttlDB != null) {
ttlDB.close();
}
if (dbOptions != null) {
dbOptions.dispose();
}
}
}
public void performanceTest() {
final LinkedBlockingDeque<String> getQueue = new LinkedBlockingDeque<String>();
final LinkedBlockingDeque<String> rmQueue = new LinkedBlockingDeque<String>();
final int TIMES = 10000000;
final String END_TAG = "end";
Thread putThread = new Thread(new Runnable() {
RunCounter runCounter = new RunCounter("Put");
@Override
public void run() {
// TODO Auto-generated method stub
for (int i = 0; i < TIMES; i++) {
String key = "Key" + i;
cache.put(key, key);
runCounter.count(1l);
getQueue.offer(key);
}
getQueue.offer(END_TAG);
System.out.println("Shutdown ");
}
});
Thread getThread = new Thread(new Runnable() {
RunCounter runCounter = new RunCounter("Get");
@Override
public void run() {
// TODO Auto-generated method stub
while (true) {
String key = null;
try {
key = getQueue.take();
} catch (InterruptedException e) {
// TODO Auto-generated catch block
e.printStackTrace();
}
if (key == null || key.equals(END_TAG)) {
break;
}
cache.get(key);
runCounter.count(1l);
rmQueue.offer(key);
}
System.out.println("Shutdown get");
}
});
Thread rmThread = new Thread(new Runnable() {
RunCounter runCounter = new RunCounter("Rm");
@Override
public void run() {
// TODO Auto-generated method stub
while (true) {
String key = null;
try {
key = rmQueue.take();
} catch (InterruptedException e) {
// TODO Auto-generated catch block
e.printStackTrace();
}
if (key == null) {
break;
}
cache.remove(key);
runCounter.count(1l);
}
System.out.println("Shutdown rm");
}
});
System.out.println("Start performance test");
putThread.start();
getThread.start();
rmThread.start();
JStormUtils.sleepMs(1000000);
}
public static void main(String[] args) throws Exception {
if (OSInfo.isLinux() || OSInfo.isMac()) {
RocksDB.loadLibrary();
RocksDBTest instance = new RocksDBTest();
// instance.testSimple();
// instance.ttlDbOpenWithColumnFamilies();
// instance.visitorAccross();
instance.performanceTest();
}
}
}