rocksdb-api/main.go

package main

import (
	"fmt"
	"log"
	"os"
	"strconv"
	"sync"
	"time"

	"github.com/Pallinder/go-randomdata"
	"github.com/tecbot/gorocksdb"
)

func putData(wg *sync.WaitGroup, db *gorocksdb.DB, cfs []*gorocksdb.ColumnFamilyHandle) {
	defer wg.Done()

	now := time.Now().Unix()

	wo := gorocksdb.NewDefaultWriteOptions()
	COUNT := 1000000

	for i := 0; i < COUNT; i++ {
		name := strconv.Itoa(i) + "-" + randomdata.FirstName(-1) + "-" + randomdata.FullName(-1) + strconv.Itoa(randomdata.Number(0, 10000000))
		age := randomdata.Number(0, 105)
		level := randomdata.Number(0, 10)
		task := randomdata.Email() + randomdata.Adjective() + randomdata.City()

		// tx.Exec("insert into user(name, age,level, task, attr, taskid) values(?, ?, ?, ?, ?, ?)", name, age, level, task, attr, taskid)
		putbuf := []byte(name)
		keybuf := append(putbuf, '\x01', byte(age), '\x01', byte(level), '\x01')

		putbuf = append(keybuf, []byte(task)...)
		db.PutCF(wo, cfs[1], keybuf, putbuf)
	}

	log.Println(int64(COUNT) / (time.Now().Unix() - now))
}

func OpenDataBase() (*gorocksdb.DB, []*gorocksdb.ColumnFamilyHandle) {
	bbto := gorocksdb.NewDefaultBlockBasedTableOptions()

	bbto.SetBlockCache(gorocksdb.NewLRUCache(3 << 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)

	year, month, day := time.Now().Date()

	timeCFStr := fmt.Sprintf("%d-%d-%d", year, int64(month), day)

	f, err := os.OpenFile("./log", os.O_CREATE|os.O_RDWR|os.O_TRUNC, 0660)
	if err != nil {
		panic(err)
	}
	log.SetOutput(f)

	log.Println(timeCFStr)
	names, err := gorocksdb.ListColumnFamilies(opts, ".rocksdb")
	if err != nil {
		log.Println(".rocksdb 文件不存在")
		names = append(names, "default")
	}

	isadd := true

	var opslist []*gorocksdb.Options
	for _, name := range names {
		opslist = append(opslist, opts)
		if timeCFStr == name {
			isadd = false
		}
	}

	if isadd {
		names = append(names, timeCFStr)
		opslist = append(opslist, opts)
	}

	log.Println(names)

	db, cfs, err := gorocksdb.OpenDbColumnFamilies(opts, ".rocksdb", names, opslist)
	if err != nil {
		panic(err)
	}
	return db, cfs
}

func CreateData() {

	db, cfs := OpenDataBase()

	GCOUNT := 10

	wg := &sync.WaitGroup{}
	wg.Add(GCOUNT)

	for i := 0; i < GCOUNT; i++ {
		go putData(wg, db, cfs)
	}

	wg.Wait()
}

func queryData(wg *sync.WaitGroup, db *gorocksdb.DB, cfs []*gorocksdb.ColumnFamilyHandle) {
	defer wg.Done()

	// now := time.Now().Unix()

	ro := gorocksdb.NewDefaultReadOptions()

	COUNT := 10

	iter := db.NewIteratorCF(ro, cfs[2])
	iter.SeekForPrev([]byte("100"))

	for i := 0; i < COUNT; i++ {
		key := string(iter.Key().Data())
		value := string(iter.Value().Data())
		log.Println("\nkey:", key, "\n", "value:", value)
		iter.Next()
	}

	// log.Println(int64(COUNT) / (time.Now().Unix() - now))
}

func QueryData() {
	db, cfs := OpenDataBase()

	GCOUNT := 1

	wg := &sync.WaitGroup{}
	wg.Add(GCOUNT)

	for i := 0; i < GCOUNT; i++ {
		go queryData(wg, db, cfs)
	}

	wg.Wait()
}

func main() {

	// QueryData()
	CreateData()

	//  db.Put(wo, []byte("foo"), []byte("bar"))

	// ro := gorocksdb.NewDefaultReadOptions()
	// db, err := gorocksdb.OpenDb(opts, "./.rocksdb")
	// if err != nil {
	// 	panic(err)
	// }
	// s, err := db.Get(ro, []byte("foo"))
	// if err != nil {
	// 	panic(err)
	// }
	// log.Println(string(s.Data()))

}