98 lines
2.1 KiB
Go
98 lines
2.1 KiB
Go
|
package main
|
||
|
|
|
||
|
|
import (
|
||
|
|
"bytes"
|
||
|
|
"compress/flate"
|
||
|
|
"encoding/gob"
|
||
|
|
"log"
|
||
|
|
"math/rand"
|
||
|
|
"os"
|
||
|
|
|
||
|
|
"github.com/tecbot/gorocksdb"
|
||
|
|
)
|
||
|
|
|
||
|
|
// CheckErrorPanic 查错误 存在就panic操作
|
||
|
|
func CheckErrorPanic(err error) {
|
||
|
|
if err != nil {
|
||
|
|
log.Panic(err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// OpenDB 打开rocksdb
|
||
|
|
func OpenDB() *gorocksdb.DB {
|
||
|
|
bbto := gorocksdb.NewDefaultBlockBasedTableOptions()
|
||
|
|
|
||
|
|
bbto.SetBlockCache(gorocksdb.NewLRUCache(2 << 30))
|
||
|
|
bbto.SetCacheIndexAndFilterBlocksWithHighPriority(true)
|
||
|
|
|
||
|
|
opts := gorocksdb.NewDefaultOptions()
|
||
|
|
|
||
|
|
bbto.SetFilterPolicy(gorocksdb.NewBloomFilter(16))
|
||
|
|
|
||
|
|
opts.SetBlockBasedTableFactory(bbto)
|
||
|
|
opts.SetCreateIfMissing(true)
|
||
|
|
opts.SetCreateIfMissingColumnFamilies(true)
|
||
|
|
opts.SetCompression(gorocksdb.LZ4Compression)
|
||
|
|
|
||
|
|
db, err := gorocksdb.OpenDb(opts, ".rocksdb")
|
||
|
|
CheckErrorPanic(err)
|
||
|
|
return db
|
||
|
|
}
|
||
|
|
|
||
|
|
var db = OpenDB()
|
||
|
|
var wopts = gorocksdb.NewDefaultWriteOptions()
|
||
|
|
var ropts = gorocksdb.NewDefaultReadOptions()
|
||
|
|
|
||
|
|
// IKeyList key list interface
|
||
|
|
type IKeyList interface {
|
||
|
|
AppendKey(key []byte)
|
||
|
|
GetKeyList() [][]byte
|
||
|
|
GetLength() int
|
||
|
|
}
|
||
|
|
|
||
|
|
// SaveGobFromRocksdb save data from rocksdb keyseek
|
||
|
|
func SaveGobFromRocksdb(fname string, seekKey string, v IKeyList) {
|
||
|
|
i := 0
|
||
|
|
iter := db.NewIterator(ropts)
|
||
|
|
iter.Seek([]byte(seekKey))
|
||
|
|
for ; iter.Valid(); iter.Next() {
|
||
|
|
i++
|
||
|
|
// if i%1000 == 0 {
|
||
|
|
// t.Error(string(iter.Key().Data()))
|
||
|
|
// t.Error(string(iter.Value().Data()))
|
||
|
|
|
||
|
|
// }
|
||
|
|
v.AppendKey(iter.Value().Data())
|
||
|
|
}
|
||
|
|
log.Println(i)
|
||
|
|
|
||
|
|
var keybuf = &bytes.Buffer{}
|
||
|
|
|
||
|
|
enc := gob.NewEncoder(keybuf)
|
||
|
|
enc.Encode(v)
|
||
|
|
log.Println((len(keybuf.Bytes())))
|
||
|
|
|
||
|
|
f, err := os.OpenFile(fname, os.O_CREATE|os.O_CREATE|os.O_RDWR|os.O_TRUNC, os.ModePerm)
|
||
|
|
CheckErrorPanic(err)
|
||
|
|
|
||
|
|
gw, err := flate.NewWriter(f, 5)
|
||
|
|
CheckErrorPanic(err)
|
||
|
|
defer gw.Close()
|
||
|
|
|
||
|
|
gw.Write(keybuf.Bytes())
|
||
|
|
}
|
||
|
|
|
||
|
|
// LoadGob load gob from file
|
||
|
|
func LoadGob(fname string, v IKeyList) {
|
||
|
|
f, err := os.Open(fname)
|
||
|
|
CheckErrorPanic(err)
|
||
|
|
reader := flate.NewReader(f)
|
||
|
|
dec := gob.NewDecoder(reader)
|
||
|
|
dec.Decode(v)
|
||
|
|
}
|
||
|
|
|
||
|
|
// GetRandomKey get ikeylist key by random
|
||
|
|
func GetRandomKey(v IKeyList) []byte {
|
||
|
|
return v.GetKeyList()[rand.Intn(v.GetLength())]
|
||
|
|
}
|