| /** | |
| * 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 java.io.File; | |
| import java.io.IOException; | |
| import java.util.ArrayList; | |
| import java.util.Collection; | |
| import java.util.HashSet; | |
| import java.util.List; | |
| import java.util.Map; | |
| import java.util.Map.Entry; | |
| import java.util.Set; | |
| import java.util.TreeMap; | |
| import org.apache.commons.lang.StringUtils; | |
| import org.rocksdb.ColumnFamilyDescriptor; | |
| import org.rocksdb.ColumnFamilyHandle; | |
| import org.rocksdb.DBOptions; | |
| import org.rocksdb.RocksDB; | |
| import org.rocksdb.TtlDB; | |
| import org.rocksdb.WriteBatch; | |
| import org.rocksdb.WriteOptions; | |
| import org.slf4j.Logger; | |
| import org.slf4j.LoggerFactory; | |
| import backtype.storm.utils.Utils; | |
| import com.alibaba.jstorm.client.ConfigExtension; | |
| import com.alibaba.jstorm.utils.JStormUtils; | |
| import com.alibaba.jstorm.utils.PathUtils; | |
| public class RocksTTLDBCache implements JStormCache { | |
| private static final long serialVersionUID = 705938812240167583L; | |
| private static Logger LOG = LoggerFactory.getLogger(RocksTTLDBCache.class); | |
| static { | |
| RocksDB.loadLibrary(); | |
| } | |
| public static final String ROCKSDB_ROOT_DIR = "rocksdb.root.dir"; | |
| public static final String ROCKSDB_RESET = "rocksdb.reset"; | |
| protected TtlDB ttlDB; | |
| protected String rootDir; | |
| protected TreeMap<Integer, ColumnFamilyHandle> windowHandlers = new TreeMap<Integer, ColumnFamilyHandle>(); | |
| public void initDir(Map<Object, Object> conf) { | |
| String confDir = (String) conf.get(ROCKSDB_ROOT_DIR); | |
| if (StringUtils.isBlank(confDir) == true) { | |
| throw new RuntimeException("Doesn't set rootDir of rocksDB"); | |
| } | |
| boolean clean = ConfigExtension.getNimbusCacheReset(conf); | |
| LOG.info("RocksDB reset is " + clean); | |
| if (clean == true) { | |
| try { | |
| PathUtils.rmr(confDir); | |
| } catch (IOException e) { | |
| // TODO Auto-generated catch block | |
| throw new RuntimeException("Failed to cleanup rooDir of rocksDB " + confDir); | |
| } | |
| } | |
| File file = new File(confDir); | |
| if (file.exists() == false) { | |
| try { | |
| PathUtils.local_mkdirs(confDir); | |
| file = new File(confDir); | |
| } catch (IOException e) { | |
| // TODO Auto-generated catch block | |
| throw new RuntimeException("Failed to mkdir rooDir of rocksDB " + confDir); | |
| } | |
| } | |
| rootDir = file.getAbsolutePath(); | |
| } | |
| public void initDb(List<Integer> list) throws Exception { | |
| LOG.info("Begin to init rocksDB of {}", rootDir); | |
| DBOptions dbOptions = null; | |
| List<ColumnFamilyDescriptor> columnFamilyNames = new ArrayList<ColumnFamilyDescriptor>(); | |
| columnFamilyNames.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY)); | |
| for (Integer timeout : list) { | |
| columnFamilyNames.add(new ColumnFamilyDescriptor(String.valueOf(timeout).getBytes())); | |
| } | |
| List<Integer> ttlValues = new ArrayList<Integer>(); | |
| // Default column family with infinite lifetime | |
| // ATTENSION, the first must be 0, RocksDB.java API has this limitation | |
| ttlValues.add(0); | |
| // new column family with list second ttl | |
| ttlValues.addAll(list); | |
| try { | |
| dbOptions = new DBOptions().setCreateMissingColumnFamilies(true).setCreateIfMissing(true); | |
| List<ColumnFamilyHandle> columnFamilyHandleList = new ArrayList<ColumnFamilyHandle>(); | |
| ttlDB = TtlDB.open(dbOptions, rootDir, columnFamilyNames, columnFamilyHandleList, ttlValues, false); | |
| for (int i = 0; i < ttlValues.size(); i++) { | |
| windowHandlers.put(ttlValues.get(i), columnFamilyHandleList.get(i)); | |
| } | |
| LOG.info("Successfully init rocksDB of {}", rootDir); | |
| } finally { | |
| if (dbOptions != null) { | |
| dbOptions.dispose(); | |
| } | |
| } | |
| } | |
| @Override | |
| public void init(Map<Object, Object> conf) throws Exception { | |
| // TODO Auto-generated method stub | |
| initDir(conf); | |
| List<Integer> list = new ArrayList<Integer>(); | |
| if (conf.get(TAG_TIMEOUT_LIST) != null) { | |
| for (Object obj : (List) ConfigExtension.getCacheTimeoutList(conf)) { | |
| Integer timeoutSecond = JStormUtils.parseInt(obj); | |
| if (timeoutSecond == null || timeoutSecond <= 0) { | |
| continue; | |
| } | |
| list.add(timeoutSecond); | |
| } | |
| } | |
| // Add retry logic | |
| boolean isSuccess = false; | |
| for (int i = 0; i < 3; i++) { | |
| try { | |
| initDb(list); | |
| isSuccess = true; | |
| break; | |
| } catch (Exception e) { | |
| LOG.warn("Failed to init rocksDB " + rootDir, e); | |
| try { | |
| PathUtils.rmr(rootDir); | |
| } catch (IOException e1) { | |
| // TODO Auto-generated catch block | |
| } | |
| } | |
| } | |
| if (isSuccess == false) { | |
| throw new RuntimeException("Failed to init rocksDB " + rootDir); | |
| } | |
| } | |
| @Override | |
| public void cleanup() { | |
| LOG.info("Begin to close rocketDb of {}", rootDir); | |
| for (ColumnFamilyHandle columnFamilyHandle : windowHandlers.values()) { | |
| columnFamilyHandle.dispose(); | |
| } | |
| if (ttlDB != null) { | |
| ttlDB.close(); | |
| } | |
| LOG.info("Successfully closed rocketDb of {}", rootDir); | |
| } | |
| @Override | |
| public Object get(String key) { | |
| // TODO Auto-generated method stub | |
| for (Entry<Integer, ColumnFamilyHandle> entry : windowHandlers.entrySet()) { | |
| try { | |
| byte[] data = ttlDB.get(entry.getValue(), key.getBytes()); | |
| if (data != null) { | |
| try { | |
| return Utils.javaDeserialize(data); | |
| } catch (Exception e) { | |
| LOG.error("Failed to deserialize obj of " + key); | |
| ttlDB.remove(entry.getValue(), key.getBytes()); | |
| return null; | |
| } | |
| } | |
| } catch (Exception e) { | |
| } | |
| } | |
| return null; | |
| } | |
| @Override | |
| public void getBatch(Map<String, Object> map) { | |
| List<byte[]> lookupKeys = new ArrayList<byte[]>(); | |
| for (String key : map.keySet()) { | |
| lookupKeys.add(key.getBytes()); | |
| } | |
| for (Entry<Integer, ColumnFamilyHandle> entry : windowHandlers.entrySet()) { | |
| List<ColumnFamilyHandle> cfHandlers = new ArrayList<ColumnFamilyHandle>(); | |
| for (String key : map.keySet()) { | |
| cfHandlers.add(entry.getValue()); | |
| } | |
| try { | |
| Map<byte[], byte[]> results = ttlDB.multiGet(cfHandlers, lookupKeys); | |
| if (results == null || results.size() == 0) { | |
| continue; | |
| } | |
| for (Entry<byte[], byte[]> resultEntry : results.entrySet()) { | |
| byte[] keyByte = resultEntry.getKey(); | |
| byte[] valueByte = resultEntry.getValue(); | |
| if (keyByte == null || valueByte == null) { | |
| continue; | |
| } | |
| Object value = null; | |
| try { | |
| value = Utils.javaDeserialize(valueByte); | |
| } catch (Exception e) { | |
| LOG.error("Failed to deserialize obj of " + new String(keyByte)); | |
| ttlDB.remove(entry.getValue(), keyByte); | |
| continue; | |
| } | |
| map.put(new String(keyByte), value); | |
| } | |
| return; | |
| } catch (Exception e) { | |
| LOG.error("Failed to query " + map.keySet() + ", in window: " + entry.getKey()); | |
| } | |
| } | |
| return; | |
| } | |
| @Override | |
| public void remove(String key) { | |
| for (Entry<Integer, ColumnFamilyHandle> entry : windowHandlers.entrySet()) { | |
| try { | |
| ttlDB.remove(entry.getValue(), key.getBytes()); | |
| } catch (Exception e) { | |
| LOG.error("Failed to remove " + key); | |
| } | |
| } | |
| } | |
| @Override | |
| public void removeBatch(Collection<String> keys) { | |
| // TODO Auto-generated method stub | |
| for (String key : keys) { | |
| remove(key); | |
| } | |
| } | |
| protected void put(String key, Object value, Entry<Integer, ColumnFamilyHandle> entry) { | |
| byte[] data = Utils.javaSerialize(value); | |
| try { | |
| ttlDB.put(entry.getValue(), key.getBytes(), data); | |
| } catch (Exception e) { | |
| LOG.error("Failed put into cache, " + key, e); | |
| return; | |
| } | |
| for (Entry<Integer, ColumnFamilyHandle> removeEntry : windowHandlers.entrySet()) { | |
| if (removeEntry.getKey().equals(entry.getKey())) { | |
| continue; | |
| } | |
| try { | |
| ttlDB.remove(removeEntry.getValue(), key.getBytes()); | |
| } catch (Exception e) { | |
| // TODO Auto-generated catch block | |
| LOG.warn("Failed to remove other " + key); | |
| } | |
| } | |
| } | |
| protected Entry<Integer, ColumnFamilyHandle> getHandler(int timeoutSecond) { | |
| ColumnFamilyHandle cfHandler = null; | |
| Entry<Integer, ColumnFamilyHandle> ceilingEntry = windowHandlers.ceilingEntry(timeoutSecond); | |
| if (ceilingEntry != null) { | |
| return ceilingEntry; | |
| } else { | |
| return windowHandlers.firstEntry(); | |
| } | |
| } | |
| @Override | |
| public void put(String key, Object value, int timeoutSecond) { | |
| // TODO Auto-generated method stub | |
| put(key, value, getHandler(timeoutSecond)); | |
| } | |
| @Override | |
| public void put(String key, Object value) { | |
| put(key, value, windowHandlers.firstEntry()); | |
| } | |
| protected void putBatch(Map<String, Object> map, Entry<Integer, ColumnFamilyHandle> putEntry) { | |
| // TODO Auto-generated method stub | |
| WriteOptions writeOpts = null; | |
| WriteBatch writeBatch = null; | |
| Set<byte[]> putKeys = new HashSet<byte[]>(); | |
| try { | |
| writeOpts = new WriteOptions(); | |
| writeBatch = new WriteBatch(); | |
| for (Entry<String, Object> entry : map.entrySet()) { | |
| String key = entry.getKey(); | |
| Object value = entry.getValue(); | |
| byte[] data = Utils.javaSerialize(value); | |
| if (StringUtils.isBlank(key) || data == null || data.length == 0) { | |
| continue; | |
| } | |
| byte[] keyByte = key.getBytes(); | |
| writeBatch.put(putEntry.getValue(), keyByte, data); | |
| putKeys.add(keyByte); | |
| } | |
| ttlDB.write(writeOpts, writeBatch); | |
| } catch (Exception e) { | |
| LOG.error("Failed to putBatch into DB, " + map.keySet(), e); | |
| } finally { | |
| if (writeOpts != null) { | |
| writeOpts.dispose(); | |
| } | |
| if (writeBatch != null) { | |
| writeBatch.dispose(); | |
| } | |
| } | |
| for (Entry<Integer, ColumnFamilyHandle> entry : windowHandlers.entrySet()) { | |
| if (entry.getKey().equals(putEntry.getKey())) { | |
| continue; | |
| } | |
| for (byte[] keyByte : putKeys) { | |
| try { | |
| ttlDB.remove(entry.getValue(), keyByte); | |
| } catch (Exception e) { | |
| LOG.error("Failed to remove other's " + new String(keyByte)); | |
| } | |
| } | |
| } | |
| } | |
| @Override | |
| public void putBatch(Map<String, Object> map) { | |
| // TODO Auto-generated method stub | |
| putBatch(map, windowHandlers.firstEntry()); | |
| } | |
| @Override | |
| public void putBatch(Map<String, Object> map, int timeoutSeconds) { | |
| // TODO Auto-generated method stub | |
| putBatch(map, getHandler(timeoutSeconds)); | |
| } | |
| // public void put() throws Exception { | |
| // } | |
| // | |
| // public void write() throws Exception { | |
| // Options options = null; | |
| // WriteBatch wb1 = null; | |
| // WriteBatch wb2 = null; | |
| // WriteOptions opts = null; | |
| // try { | |
| // options = new Options().setMergeOperator(new StringAppendOperator()).setCreateIfMissing(true); | |
| // db = RocksDB.open(options, dbFolder.getRoot().getAbsolutePath()); | |
| // opts = new WriteOptions(); | |
| // wb1 = new WriteBatch(); | |
| // wb1.put("key1".getBytes(), "aa".getBytes()); | |
| // wb1.merge("key1".getBytes(), "bb".getBytes()); | |
| // wb2 = new WriteBatch(); | |
| // wb2.put("key2".getBytes(), "xx".getBytes()); | |
| // wb2.merge("key2".getBytes(), "yy".getBytes()); | |
| // db.write(opts, wb1); | |
| // db.write(opts, wb2); | |
| // assertThat(db.get("key1".getBytes())).isEqualTo("aa,bb".getBytes()); | |
| // assertThat(db.get("key2".getBytes())).isEqualTo("xx,yy".getBytes()); | |
| // } finally { | |
| // if (db != null) { | |
| // db.close(); | |
| // } | |
| // if (wb1 != null) { | |
| // wb1.dispose(); | |
| // } | |
| // if (wb2 != null) { | |
| // wb2.dispose(); | |
| // } | |
| // if (options != null) { | |
| // options.dispose(); | |
| // } | |
| // if (opts != null) { | |
| // opts.dispose(); | |
| // } | |
| // } | |
| // } | |
| // | |
| // | |
| // public void remove() throws Exception { | |
| // RocksDB db = null; | |
| // WriteOptions wOpt; | |
| // try { | |
| // wOpt = new WriteOptions(); | |
| // db = RocksDB.open(dbFolder.getRoot().getAbsolutePath()); | |
| // db.put("key1".getBytes(), "value".getBytes()); | |
| // db.put("key2".getBytes(), "12345678".getBytes()); | |
| // assertThat(db.get("key1".getBytes())).isEqualTo("value".getBytes()); | |
| // assertThat(db.get("key2".getBytes())).isEqualTo("12345678".getBytes()); | |
| // db.remove("key1".getBytes()); | |
| // db.remove(wOpt, "key2".getBytes()); | |
| // assertThat(db.get("key1".getBytes())).isNull(); | |
| // assertThat(db.get("key2".getBytes())).isNull(); | |
| // } finally { | |
| // if (db != null) { | |
| // db.close(); | |
| // } | |
| // } | |
| // } | |
| // | |
| // public void ttlDbOpenWithColumnFamilies() throws Exception, InterruptedException { | |
| // | |
| // } | |
| } |