package main import ( "bytes" "compress/flate" "encoding/gob" "log" "math/rand" "os" "reflect" ) 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 gob序列化存数据 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()) } // SaveGob 用于解析rocksdb源数据, 爬取数据存储在rocksdb后取出 序列化存储 // 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) // } // fv := data.FieldByName(kfs[n].FieldName) // fv.Set(reflect.ValueOf(v.Data())) // } // kl.AppendKey(data.Interface()) // } // SaveData(fname, kl) // } // LoadGob load gob from file func LoadGob(fname string, v IKeyList) { f, err := os.Open(fname) log.Println("load data: ", fname) CheckErrorPanic(err) reader := flate.NewReader(f) dec := gob.NewDecoder(reader) CheckErrorPanic(dec.Decode(v)) } // GetRandomKey get ikeylist key by random func GetRandomKey(v IKeyList) interface{} { ret := v.GetKeys()[rand.Intn(v.GetLength())] return ret } // GetRandomKeyBySlice get ikeylist key by random func GetRandomKeyBySlice(v interface{}) interface{} { sli := reflect.ValueOf(v) return sli.Index(rand.Intn(sli.Len())).Interface() }