149 lines
3.0 KiB
Go
149 lines
3.0 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"compress/flate"
|
|
"encoding/gob"
|
|
fmt "fmt"
|
|
"log"
|
|
"math/rand"
|
|
"os"
|
|
"reflect"
|
|
|
|
"github.com/tecbot/gorocksdb"
|
|
)
|
|
|
|
type KeyKind struct {
|
|
Key string
|
|
FieldName string
|
|
}
|
|
|
|
func NewKeyKind(key string, fieldname string) KeyKind {
|
|
return KeyKind{Key: key, FieldName: fieldname}
|
|
}
|
|
|
|
type KeyList struct {
|
|
Keys []interface{}
|
|
}
|
|
|
|
func (kl *KeyList) GetKeys() []interface{} {
|
|
return kl.Keys
|
|
}
|
|
|
|
func (kl *KeyList) AppendKey(key interface{}) {
|
|
kl.Keys = append(kl.Keys, key)
|
|
}
|
|
|
|
func (kl *KeyList) GetLength() int {
|
|
return len(kl.Keys)
|
|
}
|
|
|
|
// 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 interface{})
|
|
GetKeys() []interface{}
|
|
GetLength() int
|
|
}
|
|
|
|
// SaveData 存数据
|
|
func SaveData(fname string, v interface{}) {
|
|
var keybuf = &bytes.Buffer{}
|
|
|
|
enc := gob.NewEncoder(keybuf)
|
|
enc.Encode(v)
|
|
|
|
f, err := os.OpenFile(fname, os.O_CREATE|os.O_RDWR|os.O_TRUNC, os.ModePerm)
|
|
CheckErrorPanic(err)
|
|
|
|
gw, err := flate.NewWriter(f, 7)
|
|
CheckErrorPanic(err)
|
|
defer gw.Close()
|
|
|
|
gw.Write(keybuf.Bytes())
|
|
}
|
|
|
|
func SaveGob(fname string, datatype reflect.Type, kfs ...KeyKind) {
|
|
|
|
// cl := &CountryList{}
|
|
kl := &KeyList{}
|
|
|
|
i := 0
|
|
iter := db.NewIterator(ropts)
|
|
bkey := []byte(kfs[0].Key)
|
|
iter.Seek(bkey)
|
|
|
|
for ; iter.ValidForPrefix(bkey); iter.Next() {
|
|
i++
|
|
first := iter.Value().Data()
|
|
data := reflect.New(datatype).Elem()
|
|
|
|
data.FieldByName(kfs[0].FieldName).Set(reflect.ValueOf(first))
|
|
|
|
// country := &Country{}
|
|
// country.Name = name
|
|
for n := 1; n < len(kfs); n++ {
|
|
v, err := db.Get(ropts, []byte(fmt.Sprintf(kfs[n].Key, first)))
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
// t.Error(string(iter.Key().Data()), len(pic.Data()))
|
|
// country.Pic = pic.Data()
|
|
fv := data.FieldByName(kfs[n].FieldName)
|
|
fv.Set(reflect.ValueOf(v.Data()))
|
|
// cl.KeyList = append(cl.KeyList, country)
|
|
|
|
}
|
|
kl.AppendKey(data.Interface())
|
|
}
|
|
// t.Error(i)
|
|
// // SaveData("./data/country.gob", cl)
|
|
SaveData(fname, kl)
|
|
// t.Error()
|
|
}
|
|
|
|
// 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) interface{} {
|
|
return v.GetKeys()[rand.Intn(v.GetLength())]
|
|
}
|