├── rockdis.conf ├── README.md ├── utils.go ├── main.go ├── rocks_keys.go ├── rocks_srvs.go ├── rocks_sets.go ├── server.go ├── rocks_lists.go ├── rocks_strings.go ├── proto.go ├── rocks_hashes.go └── rocks.go /rockdis.conf: -------------------------------------------------------------------------------- 1 | ;Rockdis Configuration File Example 2 | 3 | [Server] 4 | Bind = 127.0.0.1 5 | Port = 6379 6 | MonitorLog = false 7 | 8 | [Database] 9 | DbDir = /opt/tmp/rockdis 10 | MaxMemory = 5g 11 | BlockSize = 4k 12 | CreateIfMissing = true 13 | BloomFilter = 10 14 | Compression = snappy ; no/snappy/zlib/bzip2 15 | CompactionStyle = level ;level/universal 16 | MaxOpenFiles = 0 17 | MaxMerge = 5 ; 0 to disable this 18 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | Wrap RocksDB inside a server talks like the REDIS. 2 | 3 | For now, I am using tecbot/gorocksdb Go wrapper for RocksDB which needs his fork of RocksDB. 4 | 5 | Dependencies: 6 | * RocksDB: github.com/tecbot/rocksdb 7 | * RocksDB Go Wrapper: github.com/tecbot/gorocksdb 8 | * go get code.google.com/p/gcfg 9 | 10 | Support commands: 11 | * Keys: del, type, exists, keys 12 | * Strings: getset, get, set, mget, mset, append, incr, incrby, decr, decrby 13 | * Lists: lpush, rpush, lpop, rpop, lrange, lindex, llen, ltrim 14 | * Hashes: hset, hget, hgetall, hexists, hdel, hkeys, hvals, hlen, hmget, hmset 15 | * Sets : sadd, srem, smembers, scard, sismember 16 | 17 | Config: 18 | 19 | Please refer to rockdis.conf for example. 20 | ``` 21 | $ go run *.go -conf=rockdis.conf 22 | ``` 23 | Then, you can just use the redis-cli command line to try the server. 24 | 25 | Some Tests: 26 | * get/set: ops can reach 11000+ (process such redis commands in a second.), set takes average 1.5ms and get takes average 1.8ms. 27 | * lists: ops only can reach 4500+, lpush takes average 4.5ms and get takes average 7ms 28 | 29 | 30 | Thanks for the idea and code from https://github.com/dotcloud/go-redis-server 31 | -------------------------------------------------------------------------------- /utils.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "bytes" 5 | "encoding/gob" 6 | "fmt" 7 | "log" 8 | "reflect" 9 | "strconv" 10 | "strings" 11 | "sync/atomic" 12 | ) 13 | 14 | func encode(value interface{}) ([]byte, error) { 15 | buffer := new(bytes.Buffer) 16 | encoder := gob.NewEncoder(buffer) 17 | if err := encoder.Encode(value); err != nil { 18 | log.Printf("[Encode] Error when encode object, %s", err) 19 | return nil, err 20 | } 21 | return buffer.Bytes(), nil 22 | } 23 | 24 | func decode(data []byte, vType reflect.Type) (interface{}, error) { 25 | v := reflect.New(vType) 26 | buffer := bytes.NewBuffer(data) 27 | decoder := gob.NewDecoder(buffer) 28 | if err := decoder.DecodeValue(v); err != nil { 29 | log.Printf("[Decode] Error when decode object, %s", err) 30 | return nil, err 31 | } 32 | return reflect.Indirect(v).Interface(), nil 33 | } 34 | 35 | func parseComputerSize(size string) (int, error) { 36 | oneKBytes := 1 << 10 37 | oneMBytes := 1 << 20 38 | oneGBytes := 1 << 30 39 | var ( 40 | count int 41 | bits byte 42 | ) 43 | if _, err := fmt.Sscanf(strings.ToLower(size), "%d%c", &count, &bits); err != nil { 44 | return 0, err 45 | } 46 | switch bits { 47 | case 'k': 48 | return count * oneKBytes, nil 49 | case 'm': 50 | return count * oneMBytes, nil 51 | case 'g': 52 | return count * oneGBytes, nil 53 | } 54 | return 0, fmt.Errorf("[Config] Format error") 55 | } 56 | 57 | type AtomicInt int64 58 | 59 | func (i *AtomicInt) Add(n int64) { 60 | atomic.AddInt64((*int64)(i), n) 61 | } 62 | 63 | func (i *AtomicInt) Set(n int64) { 64 | atomic.StoreInt64((*int64)(i), n) 65 | } 66 | 67 | func (i *AtomicInt) Get() int64 { 68 | return atomic.LoadInt64((*int64)(i)) 69 | } 70 | 71 | func (i *AtomicInt) String() string { 72 | return strconv.FormatInt(i.Get(), 10) 73 | } 74 | -------------------------------------------------------------------------------- /main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "code.google.com/p/gcfg" 5 | "flag" 6 | "fmt" 7 | "log" 8 | "math/rand" 9 | "os" 10 | "os/signal" 11 | "runtime" 12 | "syscall" 13 | "time" 14 | ) 15 | 16 | type RockdisConfig struct { 17 | Server struct { 18 | Bind string 19 | Port int 20 | MonitorLog bool 21 | } 22 | Database struct { 23 | DbDir string 24 | MaxMemory string 25 | BlockSize string 26 | CreateIfMissing bool 27 | BloomFilter int 28 | Compression string 29 | CompactionStyle string 30 | MaxOpenFiles int 31 | MaxMerge int 32 | } 33 | } 34 | 35 | func main() { 36 | var confName string 37 | flag.StringVar(&confName, "conf", "rockdis.conf", "Rockdis Configuration file") 38 | flag.Parse() 39 | 40 | var config RockdisConfig 41 | err := gcfg.ReadFileInto(&config, confName) 42 | if err != nil { 43 | log.Fatal(err) 44 | } 45 | globalStat.configFile = confName 46 | globalStat.config = config 47 | globalStat.startTime = time.Now() 48 | globalStat.qpsCommands = 0 49 | globalStat.qpsStart = AtomicInt(time.Now().Unix()) 50 | go func(qpsInterval int) { 51 | tc := time.Tick(time.Minute * time.Duration(qpsInterval)) 52 | for _ = range tc { 53 | globalStat.qpsCommands.Set(0) 54 | globalStat.qpsStart.Set(time.Now().Unix()) 55 | } 56 | }(15) 57 | 58 | rock := NewRocksDBHandler(config) 59 | server := NewServer(config) 60 | defer func() { 61 | rock.Close() 62 | server.Close() 63 | }() 64 | 65 | signalChan := make(chan os.Signal, 1) 66 | signal.Notify(signalChan, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGABRT) 67 | go func() { 68 | s := <-signalChan 69 | log.Printf("[Main] Captured the signal %v", s) 70 | rock.Close() 71 | server.Close() 72 | os.Exit(0) 73 | }() 74 | 75 | if err := server.RegisterHandler(rock); err != nil { 76 | log.Fatalf("[Main] Register Handler error, %s", err) 77 | } 78 | if err := server.ListenAndServe(); err != nil { 79 | log.Fatalf("[Main] ListenAndServe error, %s", err) 80 | } 81 | } 82 | 83 | type Stat struct { 84 | version string 85 | configFile string 86 | config RockdisConfig 87 | startTime time.Time 88 | clients AtomicInt 89 | totalConnections AtomicInt 90 | totalCommands AtomicInt 91 | keyHits AtomicInt 92 | keyMisses AtomicInt 93 | qpsCommands AtomicInt 94 | qpsStart AtomicInt 95 | } 96 | 97 | var globalStat *Stat 98 | 99 | func init() { 100 | globalStat = &Stat{version: "0.0.1"} 101 | log.SetFlags(log.LstdFlags | log.Lshortfile) 102 | runtime.GOMAXPROCS(runtime.NumCPU()) 103 | rand.Seed(time.Now().UnixNano()) 104 | } 105 | 106 | var _ = fmt.Println 107 | -------------------------------------------------------------------------------- /rocks_keys.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "bytes" 5 | rocks "github.com/tecbot/gorocksdb" 6 | ) 7 | 8 | func (rh *RocksDBHandler) RedisDel(key []byte, keys ...[]byte) (int, error) { 9 | if rh.db == nil { 10 | return 0, ErrRocksIsDead 11 | } 12 | if key == nil || len(key) == 0 { 13 | return 0, ErrWrongArgumentsCount 14 | } 15 | 16 | keyData := append([][]byte{key}, keys...) 17 | count := 0 18 | readOptions := rocks.NewDefaultReadOptions() 19 | writeOptions := rocks.NewDefaultWriteOptions() 20 | defer readOptions.Destroy() 21 | defer writeOptions.Destroy() 22 | 23 | for _, dKey := range keyData { 24 | _, err := rh.loadRedisObject(readOptions, dKey) 25 | if err == nil { 26 | batch := rocks.NewWriteBatch() 27 | batch.Delete(rh.getTypeKey(dKey)) 28 | batch.Delete(dKey) 29 | if err := rh.db.Write(writeOptions, batch); err == nil { 30 | count++ 31 | } 32 | batch.Destroy() 33 | } 34 | } 35 | return count, nil 36 | } 37 | 38 | func (rh *RocksDBHandler) RedisType(key []byte) ([]byte, error) { 39 | if rh.db == nil { 40 | return nil, ErrRocksIsDead 41 | } 42 | if key == nil || len(key) == 0 { 43 | return nil, ErrWrongArgumentsCount 44 | } 45 | 46 | options := rocks.NewDefaultReadOptions() 47 | defer options.Destroy() 48 | 49 | obj, err := rh.loadRedisObject(options, key) 50 | if err == nil { 51 | return []byte(obj.Type), nil 52 | } 53 | if err == ErrDoesNotExist { 54 | return []byte("none"), nil 55 | } 56 | return nil, err 57 | } 58 | 59 | func (rh *RocksDBHandler) RedisExists(key []byte) (int, error) { 60 | if rh.db == nil { 61 | return 0, ErrRocksIsDead 62 | } 63 | if key == nil || len(key) == 0 { 64 | return 0, ErrWrongArgumentsCount 65 | } 66 | options := rocks.NewDefaultReadOptions() 67 | defer options.Destroy() 68 | 69 | if _, err := rh.loadRedisObject(options, key); err == nil { 70 | return 1, nil 71 | } else { 72 | if err == ErrDoesNotExist { 73 | return 0, nil 74 | } 75 | return 0, err 76 | } 77 | } 78 | 79 | // Only support key prefix or all keys, e.g. "KEYS *" or "KEYS test*" 80 | func (rh *RocksDBHandler) RedisKeys(pattern []byte) ([][]byte, error) { 81 | if rh.db == nil { 82 | return nil, ErrRocksIsDead 83 | } 84 | if pattern == nil || len(pattern) == 0 { 85 | return nil, ErrWrongArgumentsCount 86 | } 87 | if pattern[len(pattern)-1] == '*' { 88 | pattern = pattern[:len(pattern)-1] 89 | } 90 | 91 | options := rocks.NewDefaultReadOptions() 92 | defer options.Destroy() 93 | options.SetFillCache(false) 94 | 95 | data := make([][]byte, 0) 96 | it := rh.db.NewIterator(options) 97 | defer it.Close() 98 | it.Seek(pattern) 99 | for ; it.Valid(); it.Next() { 100 | key := it.Key() 101 | dKey := rh.copySlice(key, false) 102 | if bytes.HasPrefix(dKey, kTypeKeyPrefix) { 103 | continue 104 | } 105 | if !bytes.HasPrefix(dKey, pattern) { 106 | break 107 | } 108 | data = append(data, dKey) 109 | } 110 | if err := it.Err(); err != nil { 111 | return nil, err 112 | } 113 | return data, nil 114 | } 115 | 116 | // This is a dummy stub, since we are using this on redis. 117 | func (rh *RocksDBHandler) RedisExpire(key []byte, timeout int) (int, error) { 118 | return 1, nil 119 | } 120 | -------------------------------------------------------------------------------- /rocks_srvs.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "fmt" 5 | "os" 6 | "runtime" 7 | "strings" 8 | "time" 9 | ) 10 | 11 | // Server and Connection Command 12 | func (rh *RocksDBHandler) RedisSelect(db int) error { 13 | if rh.db == nil { 14 | return ErrRocksIsDead 15 | } 16 | return nil 17 | } 18 | 19 | func (rh *RocksDBHandler) RedisPing() (*StatusReply, error) { 20 | if rh.db == nil { 21 | return nil, ErrRocksIsDead 22 | } 23 | return &StatusReply{"PONG"}, nil 24 | } 25 | 26 | func (rh *RocksDBHandler) RedisInfo() ([]byte, error) { 27 | if rh.db == nil { 28 | return nil, ErrRocksIsDead 29 | } 30 | 31 | data := make([]string, 0) 32 | 33 | // server section 34 | qpsStart := globalStat.qpsStart.Get() 35 | qps := float64(globalStat.qpsCommands.Get()) / float64(time.Now().Unix()-qpsStart) 36 | serverSection := []string{ 37 | "# Server", 38 | "version: " + globalStat.version, 39 | "os: " + runtime.GOOS, 40 | fmt.Sprintf("process_id: %d", os.Getpid()), 41 | fmt.Sprintf("tcp_port: %d", globalStat.config.Server.Port), 42 | "config_file: " + globalStat.configFile, 43 | fmt.Sprintf("uptime: %s", time.Since(globalStat.startTime)), 44 | fmt.Sprintf("connected_clients: %d", globalStat.clients.Get()), 45 | fmt.Sprintf("total_connections_received: %d", globalStat.totalConnections.Get()), 46 | fmt.Sprintf("total_commands_processed: %d", globalStat.totalCommands.Get()), 47 | fmt.Sprintf("instantaneous_ops_per_sec: %v", qps), 48 | fmt.Sprintf("keyspace_hits: %d", globalStat.keyHits.Get()), 49 | fmt.Sprintf("keyspace_misses: %d", globalStat.keyMisses.Get()), 50 | "", 51 | } 52 | data = append(data, serverSection...) 53 | 54 | // Runtime Section 55 | var memStats runtime.MemStats 56 | runtime.ReadMemStats(&memStats) 57 | runtimeSection := []string{ 58 | "# Runtime", 59 | fmt.Sprintf("cpu_count: %d", runtime.NumCPU()), 60 | fmt.Sprintf("cgo_calls: %d", runtime.NumCgoCall()), 61 | fmt.Sprintf("goroutine_count: %d", runtime.NumGoroutine()), 62 | fmt.Sprintf("mem_alloc: %d", memStats.Alloc), 63 | fmt.Sprintf("mem_sys: %d", memStats.Sys), 64 | fmt.Sprintf("mem_total: %d", memStats.TotalAlloc), 65 | fmt.Sprintf("mem_stack: %d", memStats.StackInuse), 66 | fmt.Sprintf("mem_heap_alloc: %d", memStats.HeapAlloc), 67 | fmt.Sprintf("mem_heap_sys: %d", memStats.HeapSys), 68 | fmt.Sprintf("mem_heap_objects: %d", memStats.HeapObjects), 69 | fmt.Sprintf("mem_gc_count: %d", memStats.NumGC), 70 | fmt.Sprintf("mem_gc_total_pause: %v", float64(memStats.PauseTotalNs)/float64(time.Millisecond)), 71 | "", 72 | } 73 | data = append(data, runtimeSection...) 74 | 75 | // rocksdb section 76 | rocksSection := []string{ 77 | "# Rocksdb Config", 78 | "rocksdb_directory: " + globalStat.config.Database.DbDir, 79 | "max_memory: " + globalStat.config.Database.MaxMemory, 80 | "block_size: " + globalStat.config.Database.BlockSize, 81 | "compression: " + globalStat.config.Database.Compression, 82 | "compaction_style: " + globalStat.config.Database.CompactionStyle, 83 | fmt.Sprintf("max_memtable_merge: %d", globalStat.config.Database.MaxMerge), 84 | "", 85 | "# Rocksdb Stats", 86 | } 87 | rocksStats := strings.Split(rh.db.GetProperty("rocksdb.stats"), "\n") 88 | rocksStats = rocksStats[len(rocksStats)-14:] 89 | for _, rockStat := range rocksStats { 90 | rocksSection = append(rocksSection, rockStat) 91 | } 92 | rocksSection = append(rocksSection, "") 93 | data = append(data, rocksSection...) 94 | 95 | return []byte(strings.Join(data, "\r\n")), nil 96 | } 97 | 98 | var _ = fmt.Println 99 | -------------------------------------------------------------------------------- /rocks_sets.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "encoding/gob" 5 | "fmt" 6 | rocks "github.com/tecbot/gorocksdb" 7 | "reflect" 8 | ) 9 | 10 | func (rh *RocksDBHandler) RedisScard(key []byte) (int, error) { 11 | if err := rh.checkRedisCall(key); err != nil { 12 | return 0, err 13 | } 14 | if err := rh.checkKeyType(key, kRedisSet); err != nil { 15 | return 0, err 16 | } 17 | 18 | setData, err := rh._set_getData(key) 19 | if err != nil { 20 | return 0, err 21 | } 22 | return len(setData), nil 23 | } 24 | 25 | func (rh *RocksDBHandler) RedisSismember(key, member []byte) (int, error) { 26 | if err := rh.checkRedisCall(key); err != nil { 27 | return 0, err 28 | } 29 | if err := rh.checkKeyType(key, kRedisSet); err != nil { 30 | return 0, err 31 | } 32 | 33 | setData, err := rh._set_getData(key) 34 | if err != nil { 35 | return 0, err 36 | } 37 | if _, ok := setData[string(member)]; ok { 38 | return 1, nil 39 | } 40 | return 0, nil 41 | } 42 | 43 | func (rh *RocksDBHandler) RedisSmembers(key []byte) ([][]byte, error) { 44 | if err := rh.checkRedisCall(key); err != nil { 45 | return nil, err 46 | } 47 | if err := rh.checkKeyType(key, kRedisSet); err != nil { 48 | return nil, err 49 | } 50 | 51 | setData, err := rh._set_getData(key) 52 | if err != nil { 53 | return nil, err 54 | } 55 | return __setToBytes(setData), nil 56 | } 57 | 58 | func (rh *RocksDBHandler) RedisSadd(key, value []byte, values ...[]byte) (int, error) { 59 | totalCount, existCount, err := rh._set_doMembership(kSetOpSet, key, value, values...) 60 | if err != nil { 61 | return 0, err 62 | } 63 | return totalCount - existCount, nil 64 | } 65 | 66 | func (rh *RocksDBHandler) RedisSrem(key, value []byte, values ...[]byte) (int, error) { 67 | _, existCount, err := rh._set_doMembership(kSetOpDelete, key, value, values...) 68 | if err != nil { 69 | return 0, err 70 | } 71 | return existCount, nil 72 | } 73 | 74 | func (rh *RocksDBHandler) _set_doMembership(opCode string, key, value []byte, values ...[]byte) (int, int, error) { 75 | if err := rh.checkRedisCall(key, value); err != nil { 76 | return 0, 0, err 77 | } 78 | if err := rh.checkKeyType(key, kRedisSet); err != nil { 79 | return 0, 0, err 80 | } 81 | 82 | setData, err := rh._set_getData(key) 83 | if err != nil { 84 | return 0, 0, err 85 | } 86 | existCount := 0 87 | values = append([][]byte{value}, values...) 88 | for i := range values { 89 | if _, ok := setData[string(values[i])]; ok { 90 | existCount++ 91 | } 92 | } 93 | if err := rh._set_doMerge(key, values, opCode); err != nil { 94 | return 0, 0, err 95 | } 96 | return len(values), existCount, nil 97 | } 98 | 99 | func (rh *RocksDBHandler) _set_doMerge(key []byte, values [][]byte, opCode string) error { 100 | if values == nil || len(values) == 0 { 101 | return ErrWrongArgumentsCount 102 | } 103 | options := rocks.NewDefaultWriteOptions() 104 | defer options.Destroy() 105 | batch := rocks.NewWriteBatch() 106 | defer batch.Destroy() 107 | batch.Put(rh.getTypeKey(key), []byte(kRedisSet)) 108 | for _, value := range values { 109 | operand := SetOperand{opCode, value} 110 | if data, err := encode(operand); err == nil { 111 | batch.Merge(key, data) 112 | } else { 113 | return err 114 | } 115 | } 116 | return rh.db.Write(options, batch) 117 | } 118 | 119 | func (rh *RocksDBHandler) _set_getData(key []byte) (map[string]bool, error) { 120 | setData := make(map[string]bool) 121 | options := rocks.NewDefaultReadOptions() 122 | defer options.Destroy() 123 | if obj, err := rh.loadRedisObject(options, key); err != nil { 124 | if err == ErrDoesNotExist { 125 | return setData, nil 126 | } 127 | return nil, err 128 | } else { 129 | if obj.Type != kRedisSet { 130 | return nil, ErrWrongTypeRedisObject 131 | } 132 | data := obj.Data.([][]byte) 133 | for _, itemData := range data { 134 | setData[string(itemData)] = true 135 | } 136 | return setData, nil 137 | } 138 | } 139 | 140 | const ( 141 | kSetOpSet = "set" 142 | kSetOpDelete = "delete" 143 | ) 144 | 145 | type SetOperand struct { 146 | Command string 147 | Key []byte 148 | } 149 | 150 | func init() { 151 | gob.Register(&SetOperand{}) 152 | } 153 | 154 | type SetMerger struct{} 155 | 156 | func (m *SetMerger) FullMerge(existingObject *RedisObject, operands [][]byte) bool { 157 | rawData, ok := existingObject.Data.([][]byte) 158 | if !ok { 159 | rawData = [][]byte{} 160 | } 161 | setData := make(map[string]bool) 162 | for i := range rawData { 163 | setData[string(rawData[i])] = true 164 | } 165 | for _, operand := range operands { 166 | if obj, err := decode(operand, reflect.TypeOf(SetOperand{})); err == nil { 167 | op := obj.(SetOperand) 168 | switch op.Command { 169 | case kSetOpSet: 170 | setData[string(op.Key)] = true 171 | case kSetOpDelete: 172 | delete(setData, string(op.Key)) 173 | } 174 | } 175 | } 176 | existingObject.Data = __setToBytes(setData) 177 | return true 178 | } 179 | 180 | func (m *SetMerger) PartialMerge(leftOperand, rightOperand []byte) ([]byte, bool) { 181 | obj, err := decode(leftOperand, reflect.TypeOf(SetOperand{})) 182 | if err != nil { 183 | return nil, false 184 | } 185 | leftOp := obj.(SetOperand) 186 | obj, err = decode(rightOperand, reflect.TypeOf(SetOperand{})) 187 | if err != nil { 188 | return nil, false 189 | } 190 | rightOp := obj.(SetOperand) 191 | if string(leftOp.Key) == string(rightOp.Key) { 192 | return rightOperand, true 193 | } 194 | 195 | return nil, false 196 | } 197 | 198 | func __setToBytes(setData map[string]bool) [][]byte { 199 | rawData := make([][]byte, len(setData)) 200 | index := 0 201 | for key := range setData { 202 | rawData[index] = []byte(key) 203 | index++ 204 | } 205 | return rawData 206 | } 207 | 208 | var _ = fmt.Println 209 | -------------------------------------------------------------------------------- /server.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "bytes" 5 | "fmt" 6 | "io" 7 | "log" 8 | "net" 9 | "reflect" 10 | "strings" 11 | "time" 12 | ) 13 | 14 | const ( 15 | kDefaultAddress = ":6379" 16 | ) 17 | 18 | type HandlerFn func(request *Request) (Reply, error) 19 | type GuarderFn func(request *Request) (reflect.Value, *ErrorReply) 20 | 21 | type Server struct { 22 | Address string 23 | Methods map[string]HandlerFn 24 | MonitorLog bool 25 | } 26 | 27 | func (s *Server) RegisterHandler(handler interface{}) error { 28 | hType := reflect.TypeOf(handler) 29 | for i := 0; i < hType.NumMethod(); i++ { 30 | method := hType.Method(i) 31 | if !strings.HasPrefix(method.Name, "Redis") { 32 | continue 33 | } 34 | hFn, err := s.newHandler(handler, &method.Func) 35 | if err != nil { 36 | return err 37 | } 38 | methodName := strings.ToLower(method.Name[5:]) 39 | s.Methods[methodName] = hFn 40 | log.Printf("[RegisterHandler] Registered supported method <%s>", methodName) 41 | } 42 | return nil 43 | } 44 | 45 | func (s *Server) ListenAndServe() error { 46 | addr := s.Address 47 | if addr == "" { 48 | addr = kDefaultAddress 49 | } 50 | l, err := net.Listen("tcp", addr) 51 | if err != nil { 52 | return err 53 | } 54 | defer l.Close() 55 | log.Println("[ListenAndServe] GoRockdis is listening on", addr) 56 | 57 | for { 58 | conn, err := l.Accept() 59 | if err != nil { 60 | return err 61 | } 62 | go s.ServeClient(conn) 63 | } 64 | } 65 | 66 | func (s *Server) ServeClient(conn net.Conn) (err error) { 67 | globalStat.totalConnections.Add(1) 68 | globalStat.clients.Add(1) 69 | clientAddr := conn.RemoteAddr().String() 70 | defer func() { 71 | if err != nil { 72 | log.Printf("[ServeClient] Error in request/reply, will close the connnetion <%s>: %s", clientAddr, err) 73 | fmt.Fprintf(conn, "-ERROR %s\r\n", err) 74 | } 75 | conn.Close() 76 | conn = nil 77 | globalStat.clients.Add(-1) 78 | }() 79 | 80 | for { 81 | conn.SetReadDeadline(time.Now()) 82 | zeroByte := make([]byte, 0) 83 | if _, err := conn.Read(zeroByte); err == io.EOF { 84 | // log.Printf("[ServeClient] Detect a closed connection on %s", clientAddr) 85 | break 86 | } 87 | 88 | conn.SetReadDeadline(time.Now().Add(10 * time.Minute)) 89 | request, err := NewRequest(conn) 90 | if err == io.EOF { 91 | // log.Printf("[ServeClient] Detect a closed connection on %s", clientAddr) 92 | break 93 | } else { 94 | if neterr, ok := err.(net.Error); ok && neterr.Timeout() { 95 | // log.Printf("[ServeClient] Detect a Timeout client on %s, will check if she is closed.", clientAddr) 96 | } else { 97 | if err != nil { 98 | return err 99 | } 100 | globalStat.totalCommands.Add(1) 101 | globalStat.qpsCommands.Add(1) 102 | request.RemoteAddress = clientAddr 103 | if reply, err := s.ServeRequest(request); err != nil { 104 | return err 105 | } else { 106 | if _, err := reply.WriteTo(conn); err != nil { 107 | return err 108 | } 109 | } 110 | } 111 | } 112 | } 113 | return nil 114 | } 115 | 116 | func (s *Server) ServeRequest(request *Request) (Reply, error) { 117 | if fn, ok := s.Methods[strings.ToLower(request.Command)]; ok { 118 | return fn(request) 119 | } else { 120 | return ErrMethodNotSupported, nil 121 | } 122 | } 123 | 124 | func (s *Server) Close() { 125 | log.Printf("[Server] Server Stopped.") 126 | } 127 | 128 | func NewServer(config RockdisConfig) *Server { 129 | s := &Server{} 130 | s.Methods = make(map[string]HandlerFn) 131 | s.Address = fmt.Sprintf("%s:%d", config.Server.Bind, config.Server.Port) 132 | s.MonitorLog = config.Server.MonitorLog 133 | return s 134 | } 135 | 136 | func (s *Server) newHandler(handler interface{}, f *reflect.Value) (HandlerFn, error) { 137 | errType := reflect.TypeOf(s.newHandler).Out(1) // get the error's type 138 | guards, err := s.newHandlerGuards(handler, f) 139 | if err != nil { 140 | return nil, err 141 | } 142 | 143 | fType := f.Type() 144 | if fType.NumOut() == 0 { 145 | return nil, fmt.Errorf("Not enough return values") 146 | } 147 | if fType.NumOut() > 2 { 148 | return nil, fmt.Errorf("Too many return values") 149 | } 150 | if t := fType.Out(fType.NumOut() - 1); t != errType { 151 | return nil, fmt.Errorf("Last return value must be an error type (not %s)", t) 152 | } 153 | 154 | return s.newHandlerFn(handler, f, guards), nil 155 | } 156 | 157 | func (s *Server) newHandlerGuards(handler interface{}, f *reflect.Value) ([]GuarderFn, error) { 158 | guards := []GuarderFn{} 159 | fType := f.Type() 160 | 161 | inputStart := 0 162 | if fType.NumIn() > 0 && fType.In(0).AssignableTo(reflect.TypeOf(handler)) { 163 | // In: [*handlerType, arg0, arg1 ...] 164 | inputStart = 1 165 | } 166 | for i := inputStart; i < fType.NumIn(); i++ { 167 | switch fType.In(i) { 168 | case reflect.TypeOf([]byte{}): 169 | guards = append(guards, guardRequestByteArg(i-inputStart)) 170 | case reflect.TypeOf([][]byte{}): 171 | guards = append(guards, guardRequestByteSliceArg(i-inputStart)) 172 | case reflect.TypeOf(1): 173 | guards = append(guards, guardRequestIntArg(i-inputStart)) 174 | default: 175 | return nil, fmt.Errorf("Argument %d: wrong type %s (%s)", i, fType.In(i), fType.Name()) 176 | } 177 | } 178 | return guards, nil 179 | } 180 | 181 | func (s *Server) newHandlerFn(handler interface{}, f *reflect.Value, guards []GuarderFn) HandlerFn { 182 | return func(request *Request) (Reply, error) { 183 | input := []reflect.Value{reflect.ValueOf(handler)} 184 | for _, guard := range guards { 185 | value, errReply := guard(request) 186 | if errReply != nil { 187 | return errReply, nil 188 | } 189 | input = append(input, value) 190 | } 191 | if f.Type().NumIn() == 0 { 192 | input = []reflect.Value{} 193 | } else if !f.Type().In(0).AssignableTo(reflect.TypeOf(handler)) { 194 | input = input[1:] 195 | } 196 | 197 | var monitorString string 198 | if len(request.Arguments) > 0 { 199 | monitorString = fmt.Sprintf("%.6f [0 %s] \"%s\" \"%s\"", 200 | float64(time.Now().UTC().UnixNano())/1e9, 201 | request.RemoteAddress, 202 | request.Command, 203 | bytes.Join(request.Arguments, []byte{'"', ' ', '"'})) 204 | } else { 205 | monitorString = fmt.Sprintf("%.6f [0 %s] \"%s\"", 206 | float64(time.Now().UTC().UnixNano())/1e9, 207 | request.RemoteAddress, 208 | request.Command) 209 | } 210 | if s.MonitorLog { 211 | log.Printf("[Monitor] %s", monitorString) 212 | } 213 | 214 | var results []reflect.Value 215 | if f.Type().IsVariadic() { 216 | results = f.CallSlice(input) 217 | } else { 218 | results = f.Call(input) 219 | } 220 | if err := results[len(results)-1].Interface(); err != nil { 221 | return &ErrorReply{err.(error).Error()}, nil 222 | } 223 | if len(results) > 1 { 224 | return NewReply(s, request, results[0].Interface()) 225 | } 226 | return &StatusReply{"OK"}, nil 227 | } 228 | } 229 | -------------------------------------------------------------------------------- /rocks_lists.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "encoding/gob" 5 | "fmt" 6 | rocks "github.com/tecbot/gorocksdb" 7 | "reflect" 8 | ) 9 | 10 | func (rh *RocksDBHandler) RedisLlen(key []byte) (int, error) { 11 | if err := rh.checkRedisCall(key); err != nil { 12 | return 0, err 13 | } 14 | if err := rh.checkKeyType(key, kRedisList); err != nil { 15 | return 0, err 16 | } 17 | 18 | data, err := rh._list_getData(key) 19 | if err != nil { 20 | return 0, err 21 | } 22 | return len(data), nil 23 | } 24 | 25 | func (rh *RocksDBHandler) RedisLindex(key []byte, index int) ([]byte, error) { 26 | if err := rh.checkRedisCall(key); err != nil { 27 | return nil, err 28 | } 29 | if err := rh.checkKeyType(key, kRedisList); err != nil { 30 | return nil, err 31 | } 32 | 33 | data, err := rh._list_getData(key) 34 | if err != nil { 35 | return nil, err 36 | } 37 | if index < 0 { 38 | index += len(data) 39 | } 40 | if index < 0 || index >= len(data) { 41 | return []byte{}, nil 42 | } 43 | return data[index], nil 44 | } 45 | 46 | func (rh *RocksDBHandler) RedisLrange(key []byte, start, end int) ([][]byte, error) { 47 | if err := rh.checkRedisCall(key); err != nil { 48 | return nil, err 49 | } 50 | if err := rh.checkKeyType(key, kRedisList); err != nil { 51 | return nil, err 52 | } 53 | 54 | data, err := rh._list_getData(key) 55 | if err != nil { 56 | return nil, err 57 | } 58 | if len(data) == 0 { 59 | return [][]byte{}, nil 60 | } 61 | start = __list_getIndex(start, len(data), false) 62 | end = __list_getIndex(end, len(data), true) 63 | return data[start:end], nil 64 | } 65 | 66 | func (rh *RocksDBHandler) RedisLpop(key []byte) ([]byte, error) { 67 | return rh._list_Pop(0, key) 68 | } 69 | 70 | func (rh *RocksDBHandler) RedisRpop(key []byte) ([]byte, error) { 71 | return rh._list_Pop(-1, key) 72 | } 73 | 74 | func (rh *RocksDBHandler) RedisRpush(key, value []byte, values ...[]byte) (int, error) { 75 | return rh._list_Push(-1, key, value, values...) 76 | } 77 | 78 | func (rh *RocksDBHandler) RedisLpush(key, value []byte, values ...[]byte) (int, error) { 79 | return rh._list_Push(0, key, value, values...) 80 | } 81 | 82 | func (rh *RocksDBHandler) RedisLtrim(key []byte, start, end int) error { 83 | if err := rh.checkRedisCall(key); err != nil { 84 | return err 85 | } 86 | if err := rh._list_doMerge(key, []byte{}, kListOpTrim, start, end); err != nil { 87 | return err 88 | } 89 | return nil 90 | } 91 | 92 | func (rh *RocksDBHandler) _list_Pop(direction int, key []byte) ([]byte, error) { 93 | if err := rh.checkRedisCall(key); err != nil { 94 | return nil, err 95 | } 96 | if err := rh.checkKeyType(key, kRedisList); err != nil { 97 | return nil, err 98 | } 99 | 100 | data, err := rh._list_getData(key) 101 | if err != nil { 102 | return nil, err 103 | } 104 | if len(data) == 0 { 105 | return []byte{}, nil // this is not an error 106 | } 107 | popData := data[0] 108 | if direction == -1 { 109 | popData = data[len(data)-1] 110 | } 111 | if err := rh._list_doMerge(key, popData, kListOpRemove, direction, 0); err != nil { 112 | return nil, err 113 | } 114 | return popData, nil 115 | } 116 | 117 | func (rh *RocksDBHandler) _list_Push(direction int, key, value []byte, values ...[]byte) (int, error) { 118 | if err := rh.checkRedisCall(key, value); err != nil { 119 | return 0, err 120 | } 121 | if err := rh.checkKeyType(key, kRedisList); err != nil { 122 | return 0, err 123 | } 124 | values = append([][]byte{value}, values...) 125 | if err := rh._list_doMerge(key, values, kListOpInsert, direction, 0); err != nil { 126 | return 0, err 127 | } 128 | // This will hurt the performance 129 | // if data, err := rh._list_getData(key); err == nil { 130 | // return len(data), nil 131 | // } else { 132 | // return 0, err 133 | // } 134 | return 1, nil 135 | } 136 | 137 | func (rh *RocksDBHandler) _list_doMerge(key []byte, value interface{}, opCode string, start, end int) error { 138 | var values [][]byte 139 | if d1Slice, ok := value.([]byte); ok { 140 | values = [][]byte{d1Slice} 141 | } 142 | if d2Slice, ok := value.([][]byte); ok { 143 | values = d2Slice 144 | } 145 | if values == nil || len(values) == 0 { 146 | return ErrWrongArgumentsCount 147 | } 148 | 149 | options := rocks.NewDefaultWriteOptions() 150 | defer options.Destroy() 151 | batch := rocks.NewWriteBatch() 152 | defer batch.Destroy() 153 | batch.Put(rh.getTypeKey(key), []byte(kRedisList)) 154 | for _, dValue := range values { 155 | operand := ListOperand{opCode, start, end, dValue} 156 | if data, err := encode(operand); err == nil { 157 | batch.Merge(key, data) 158 | } else { 159 | return err 160 | } 161 | } 162 | return rh.db.Write(options, batch) 163 | } 164 | 165 | func (rh *RocksDBHandler) _list_getData(key []byte) ([][]byte, error) { 166 | options := rocks.NewDefaultReadOptions() 167 | defer options.Destroy() 168 | if obj, err := rh.loadRedisObject(options, key); err != nil { 169 | if err == ErrDoesNotExist { 170 | return [][]byte{}, nil 171 | } 172 | return nil, err 173 | } else { 174 | if obj.Type != kRedisList { 175 | return nil, ErrWrongTypeRedisObject 176 | } 177 | return obj.Data.([][]byte), nil 178 | } 179 | } 180 | 181 | const ( 182 | kListOpInsert = "insert" 183 | kListOpRemove = "remove" 184 | kListOpTrim = "trim" 185 | ) 186 | 187 | type ListOperand struct { 188 | Command string 189 | Start int 190 | End int 191 | Data []byte 192 | } 193 | 194 | func init() { 195 | gob.Register(&ListOperand{}) 196 | } 197 | 198 | type ListMerger struct{} 199 | 200 | func (m *ListMerger) FullMerge(existingObject *RedisObject, operands [][]byte) bool { 201 | listData, ok := existingObject.Data.([][]byte) 202 | if !ok { 203 | listData = [][]byte{} 204 | } 205 | for _, operand := range operands { 206 | if obj, err := decode(operand, reflect.TypeOf(ListOperand{})); err == nil { 207 | op := obj.(ListOperand) 208 | switch op.Command { 209 | case kListOpInsert: 210 | if op.Start == 0 { 211 | listData = append([][]byte{op.Data}, listData...) 212 | } else { 213 | listData = append(listData, op.Data) 214 | } 215 | case kListOpRemove: 216 | if len(listData) > 0 { 217 | if op.Start == 0 { 218 | listData = listData[1:] 219 | } else { 220 | listData = listData[:len(listData)-1] 221 | } 222 | } 223 | case kListOpTrim: 224 | if len(listData) > 0 { 225 | start := __list_getIndex(op.Start, len(listData), false) 226 | end := __list_getIndex(op.End, len(listData), true) 227 | listData = listData[start:end] 228 | } 229 | } 230 | } 231 | } 232 | existingObject.Data = listData 233 | return true 234 | } 235 | 236 | func (m *ListMerger) PartialMerge(leftOperand, rightOperand []byte) ([]byte, bool) { 237 | return nil, false 238 | } 239 | 240 | func __list_getIndex(index, length int, isRightmost bool) int { 241 | if index < 0 { 242 | index += length 243 | if index < 0 { 244 | if isRightmost { 245 | index = -1 246 | } else { 247 | index = 0 248 | } 249 | } 250 | } 251 | if isRightmost { 252 | index++ 253 | } 254 | if index > length { 255 | index = length 256 | } 257 | return index 258 | } 259 | 260 | var _ = fmt.Println 261 | -------------------------------------------------------------------------------- /rocks_strings.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "encoding/gob" 5 | "fmt" 6 | rocks "github.com/tecbot/gorocksdb" 7 | "reflect" 8 | "strconv" 9 | ) 10 | 11 | func (rh *RocksDBHandler) RedisAppend(key, value []byte) (int, error) { 12 | if err := rh.checkRedisCall(key, value); err != nil { 13 | return 0, err 14 | } 15 | if err := rh.checkKeyType(key, kRedisString); err != nil { 16 | return 0, err 17 | } 18 | 19 | if err := rh._string_doMerge(key, value, kStringOpAppend); err != nil { 20 | return 0, err 21 | } 22 | if data, err := rh.RedisGet(key); err == nil { 23 | return len(data), nil 24 | } else { 25 | return 0, err 26 | } 27 | } 28 | 29 | func (rh *RocksDBHandler) RedisDecr(key []byte) ([]byte, error) { 30 | return rh.RedisIncrBy(key, -1) 31 | } 32 | 33 | func (rh *RocksDBHandler) RedisDecrBy(key []byte, value int) ([]byte, error) { 34 | return rh.RedisIncrBy(key, -1*value) 35 | } 36 | 37 | func (rh *RocksDBHandler) RedisIncr(key []byte) ([]byte, error) { 38 | return rh.RedisIncrBy(key, 1) 39 | } 40 | 41 | func (rh *RocksDBHandler) RedisIncrBy(key []byte, value int) ([]byte, error) { 42 | if err := rh.checkRedisCall(key); err != nil { 43 | return nil, err 44 | } 45 | if err := rh.checkKeyType(key, kRedisString); err != nil { 46 | return nil, err 47 | } 48 | 49 | data := []byte(strconv.Itoa(value)) 50 | if err := rh._string_doMerge(key, data, kStringOpIncr); err != nil { 51 | return nil, err 52 | } 53 | return rh.RedisGet(key) 54 | } 55 | 56 | func (rh *RocksDBHandler) RedisGetSet(key, value []byte) ([]byte, error) { 57 | if data, err := rh.RedisGet(key); err != nil { 58 | return nil, err 59 | } else { 60 | if err := rh.RedisSet(key, value); err != nil { 61 | return nil, err 62 | } 63 | return data, nil 64 | } 65 | } 66 | 67 | func (rh *RocksDBHandler) RedisGet(key []byte) ([]byte, error) { 68 | if err := rh.checkRedisCall(key); err != nil { 69 | return nil, err 70 | } 71 | if err := rh.checkKeyType(key, kRedisString); err != nil { 72 | return nil, err 73 | } 74 | 75 | options := rocks.NewDefaultReadOptions() 76 | defer options.Destroy() 77 | if obj, err := rh.loadRedisObject(options, key); err != nil { 78 | if err == ErrDoesNotExist { 79 | return []byte{}, nil 80 | } 81 | return nil, err 82 | } else { 83 | return obj.Data.([]byte), err 84 | } 85 | } 86 | 87 | func (rh *RocksDBHandler) RedisMget(keys [][]byte) ([][]byte, error) { 88 | if rh.db == nil { 89 | return nil, ErrRocksIsDead 90 | } 91 | if keys == nil || len(keys) == 0 { 92 | return nil, ErrWrongArgumentsCount 93 | } 94 | 95 | options := rocks.NewDefaultReadOptions() 96 | defer options.Destroy() 97 | results := make([][]byte, len(keys)) 98 | for i := range results { 99 | results[i] = []byte{} 100 | } 101 | for i := range results { 102 | if obj, err := rh.loadRedisObject(options, keys[i]); err == nil { 103 | if obj.Type == kRedisString { 104 | results[i] = obj.Data.([]byte) 105 | } 106 | } 107 | } 108 | return results, nil 109 | } 110 | 111 | func (rh *RocksDBHandler) RedisSet(key, value []byte) error { 112 | if err := rh.checkRedisCall(key, value); err != nil { 113 | return err 114 | } 115 | 116 | options := rocks.NewDefaultWriteOptions() 117 | defer options.Destroy() 118 | return rh.saveRedisObject(options, key, value, kRedisString) 119 | } 120 | 121 | func (rh *RocksDBHandler) RedisMset(keyValues [][]byte) error { 122 | if rh.db == nil { 123 | return ErrRocksIsDead 124 | } 125 | if keyValues == nil || len(keyValues) == 0 || len(keyValues)%2 != 0 { 126 | return ErrWrongArgumentsCount 127 | } 128 | options := rocks.NewDefaultWriteOptions() 129 | defer options.Destroy() 130 | for i := 0; i < len(keyValues); i += 2 { 131 | err := rh.saveRedisObject(options, keyValues[i], keyValues[i+1], kRedisString) 132 | if err != nil { 133 | return err 134 | } 135 | } 136 | return nil 137 | } 138 | 139 | func (rh *RocksDBHandler) _string_doMerge(key, value []byte, opCode string) error { 140 | options := rocks.NewDefaultWriteOptions() 141 | defer options.Destroy() 142 | batch := rocks.NewWriteBatch() 143 | defer batch.Destroy() 144 | batch.Put(rh.getTypeKey(key), []byte(kRedisString)) 145 | operand := StringOperand{opCode, value} 146 | if data, err := encode(operand); err != nil { 147 | return err 148 | } else { 149 | batch.Merge(key, data) 150 | } 151 | return rh.db.Write(options, batch) 152 | } 153 | 154 | const ( 155 | kStringOpIncr = "incr" 156 | kStringOpAppend = "append" 157 | ) 158 | 159 | type StringOperand struct { 160 | Command string 161 | Data []byte 162 | } 163 | 164 | func init() { 165 | gob.Register(&StringOperand{}) 166 | } 167 | 168 | type StringMerger struct { 169 | } 170 | 171 | func (m *StringMerger) FullMerge(existingObject *RedisObject, operands [][]byte) bool { 172 | appendData, ok := existingObject.Data.([]byte) 173 | if !ok { 174 | appendData = []byte{} 175 | } 176 | incrData, err := strconv.ParseInt(string(appendData), 10, 64) 177 | if err != nil { 178 | incrData = 0 179 | } 180 | 181 | lastOp := "" 182 | for _, operand := range operands { 183 | if op, err := decode(operand, reflect.TypeOf(StringOperand{})); err == nil { 184 | lastOp = op.(StringOperand).Command 185 | switch lastOp { 186 | case kStringOpIncr: 187 | if n, err := strconv.ParseInt(string(op.(StringOperand).Data), 10, 64); err == nil { 188 | incrData += n 189 | appendData = []byte(fmt.Sprintf("%d", incrData)) 190 | } 191 | case kStringOpAppend: 192 | appendData = append(appendData, op.(StringOperand).Data...) 193 | if n, err := strconv.ParseInt(string(appendData), 10, 64); err == nil { 194 | incrData = n 195 | } else { 196 | incrData = 0 197 | } 198 | } 199 | } 200 | } 201 | switch lastOp { 202 | case kStringOpIncr: 203 | existingObject.Data = []byte(fmt.Sprintf("%d", incrData)) 204 | case kStringOpAppend: 205 | existingObject.Data = appendData 206 | } 207 | return true 208 | } 209 | 210 | func (m *StringMerger) PartialMerge(leftOperand, rightOperand []byte) ([]byte, bool) { 211 | obj, err := decode(leftOperand, reflect.TypeOf(StringOperand{})) 212 | if err != nil { 213 | return nil, false 214 | } 215 | leftOp := obj.(StringOperand) 216 | obj, err = decode(rightOperand, reflect.TypeOf(StringOperand{})) 217 | if err != nil { 218 | return nil, false 219 | } 220 | rightOp := obj.(StringOperand) 221 | 222 | if leftOp.Command == rightOp.Command { 223 | mergeOp := StringOperand{Command: leftOp.Command} 224 | merged := false 225 | switch leftOp.Command { 226 | case kStringOpIncr: 227 | left, leftErr := strconv.ParseInt(string(leftOp.Data), 10, 64) 228 | right, rightErr := strconv.ParseInt(string(rightOp.Data), 10, 64) 229 | if leftErr == nil && rightErr == nil { 230 | mergeOp.Data, merged = []byte(fmt.Sprintf("%d", left+right)), true 231 | } else if leftErr == nil { 232 | mergeOp.Data, merged = []byte(fmt.Sprintf("%d", left)), true 233 | } else if rightErr == nil { 234 | mergeOp.Data, merged = []byte(fmt.Sprintf("%d", right)), true 235 | } 236 | case kStringOpAppend: 237 | mergeOp.Data, merged = append(leftOp.Data, rightOp.Data...), true 238 | } 239 | if merged { 240 | if data, err := encode(mergeOp); err == nil { 241 | return data, true 242 | } 243 | } 244 | } 245 | return nil, false 246 | } 247 | -------------------------------------------------------------------------------- /proto.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "bufio" 5 | "fmt" 6 | "io" 7 | "io/ioutil" 8 | "reflect" 9 | "strconv" 10 | "strings" 11 | ) 12 | 13 | type Request struct { 14 | Command string 15 | Arguments [][]byte 16 | RemoteAddress string 17 | Connection io.ReadCloser 18 | } 19 | 20 | func (r *Request) HasArgument(index int) bool { 21 | return index >= 0 && index < len(r.Arguments) 22 | } 23 | 24 | func (r *Request) ExpectArgument(index int) *ErrorReply { 25 | if !r.HasArgument(index) { 26 | return ErrNotEnoughArgs 27 | } 28 | return nil 29 | } 30 | 31 | func (r *Request) GetInt(index int) (int, *ErrorReply) { 32 | if errReply := r.ExpectArgument(index); errReply != nil { 33 | return -1, errReply 34 | } 35 | if n, err := strconv.Atoi(string(r.Arguments[index])); err != nil { 36 | return -1, ErrExpectInteger 37 | } else { 38 | return n, nil 39 | } 40 | } 41 | 42 | func NewRequest(conn io.ReadCloser) (*Request, error) { 43 | reader := bufio.NewReader(conn) 44 | 45 | // *CRLF 46 | line, err := reader.ReadString('\n') 47 | if err != nil { 48 | return nil, err 49 | } 50 | 51 | var argCount int 52 | if line[0] == '*' { 53 | if _, err := fmt.Sscanf(line, "*%d\r", &argCount); err != nil { 54 | return nil, Malformed("*<#Arguments>", line) 55 | } 56 | // $CRLF 57 | // CRLF 58 | command, err := readArgument(reader) 59 | if err != nil { 60 | return nil, err 61 | } 62 | arguments := make([][]byte, argCount-1) 63 | for i := 0; i < argCount-1; i++ { 64 | if arguments[i], err = readArgument(reader); err != nil { 65 | return nil, err 66 | } 67 | } 68 | return &Request{ 69 | Command: strings.ToLower(string(command)), 70 | Arguments: arguments, 71 | Connection: conn, 72 | }, nil 73 | } 74 | 75 | // Inline request: 76 | fields := strings.Split(strings.Trim(line, "\r\n"), " ") 77 | var arguments [][]byte 78 | if len(fields) > 1 { 79 | for _, arg := range fields[1:] { 80 | arguments = append(arguments, []byte(arg)) 81 | } 82 | } 83 | return &Request{ 84 | Command: strings.ToLower(fields[0]), 85 | Arguments: arguments, 86 | Connection: conn, 87 | }, nil 88 | } 89 | 90 | func readArgument(reader *bufio.Reader) ([]byte, error) { 91 | line, err := reader.ReadString('\n') 92 | if err != nil { 93 | return nil, Malformed("$", line) 94 | } 95 | 96 | var argLength int 97 | if _, err := fmt.Sscanf(line, "$%d\r", &argLength); err != nil { 98 | return nil, Malformed("$", line) 99 | } 100 | 101 | data, err := ioutil.ReadAll(io.LimitReader(reader, int64(argLength))) 102 | if err != nil { 103 | return nil, err 104 | } 105 | if len(data) != argLength { 106 | return nil, MalformedLength(argLength, len(data)) 107 | } 108 | if b, err := reader.ReadByte(); err != nil || b != '\r' { 109 | return nil, MalformedMissingCRLF() 110 | } 111 | if b, err := reader.ReadByte(); err != nil || b != '\n' { 112 | return nil, MalformedMissingCRLF() 113 | } 114 | 115 | return data, nil 116 | } 117 | 118 | func Malformed(expected string, got string) error { 119 | return fmt.Errorf("Mailformed request: %s does not match %s", got, expected) 120 | } 121 | 122 | func MalformedLength(expected int, got int) error { 123 | return fmt.Errorf("Mailformed request: argument length %d does not match %d", got, expected) 124 | } 125 | 126 | func MalformedMissingCRLF() error { 127 | return fmt.Errorf("Mailformed request: line should end with CRLF") 128 | } 129 | 130 | type Reply io.WriterTo 131 | 132 | var ( 133 | ErrMethodNotSupported = &ErrorReply{"Method is not supported"} 134 | ErrNotEnoughArgs = &ErrorReply{"Not enough arguments for the command"} 135 | ErrTooMuchArgs = &ErrorReply{"Too many arguments for the command"} 136 | ErrWrongArgsNumber = &ErrorReply{"Wrong number of arguments"} 137 | ErrExpectInteger = &ErrorReply{"Expected integer"} 138 | ErrExpectPositivInteger = &ErrorReply{"Expected positive integer"} 139 | ErrExpectMorePair = &ErrorReply{"Expected at least one key val pair"} 140 | ErrExpectEvenPair = &ErrorReply{"Got uneven number of key val pairs"} 141 | ) 142 | 143 | type ErrorReply struct { 144 | message string 145 | } 146 | 147 | func (er *ErrorReply) WriteTo(w io.Writer) (int64, error) { 148 | n, err := w.Write([]byte("-ERROR " + er.message + "\r\n")) 149 | return int64(n), err 150 | } 151 | 152 | func (er *ErrorReply) Error() string { 153 | return er.message 154 | } 155 | 156 | func guardRequestByteArg(index int) GuarderFn { 157 | return func(request *Request) (reflect.Value, *ErrorReply) { 158 | if err := request.ExpectArgument(index); err != nil { 159 | return reflect.ValueOf([]byte{}), nil 160 | } else { 161 | return reflect.ValueOf(request.Arguments[index]), nil 162 | } 163 | } 164 | } 165 | 166 | func guardRequestByteSliceArg(index int) GuarderFn { 167 | return func(request *Request) (reflect.Value, *ErrorReply) { 168 | if err := request.ExpectArgument(index); err != nil { 169 | return reflect.ValueOf([][]byte{}), nil 170 | } else { 171 | return reflect.ValueOf(request.Arguments[index:]), nil 172 | } 173 | } 174 | } 175 | 176 | func guardRequestIntArg(index int) GuarderFn { 177 | return func(request *Request) (reflect.Value, *ErrorReply) { 178 | n, err := request.GetInt(index) 179 | return reflect.ValueOf(n), err 180 | } 181 | } 182 | 183 | type StatusReply struct { 184 | code string 185 | } 186 | 187 | func (r *StatusReply) WriteTo(w io.Writer) (int64, error) { 188 | n, err := w.Write([]byte("+" + r.code + "\r\n")) 189 | return int64(n), err 190 | } 191 | 192 | type IntReply struct { 193 | number int 194 | } 195 | 196 | func (r *IntReply) WriteTo(w io.Writer) (int64, error) { 197 | n, err := w.Write([]byte(":" + strconv.Itoa(r.number) + "\r\n")) 198 | return int64(n), err 199 | } 200 | 201 | type BulkReply struct { 202 | value []byte 203 | } 204 | 205 | func (r *BulkReply) WriteTo(w io.Writer) (int64, error) { 206 | return writeBytes(r.value, w) 207 | } 208 | 209 | type MultiBulkReply struct { 210 | values [][]byte 211 | } 212 | 213 | func (r *MultiBulkReply) WriteTo(w io.Writer) (int64, error) { 214 | if r.values == nil { 215 | return 0, fmt.Errorf("Multi bulk reply found a nil values") 216 | } 217 | if wrote, err := w.Write([]byte("*" + strconv.Itoa(len(r.values)) + "\r\n")); err != nil { 218 | return int64(wrote), err 219 | } else { 220 | total := int64(wrote) 221 | for _, value := range r.values { 222 | wroteData, err := writeBytes(value, w) 223 | total += wroteData 224 | if err != nil { 225 | return total, err 226 | } 227 | } 228 | return total, nil 229 | } 230 | } 231 | 232 | func NewReply(s *Server, request *Request, value interface{}) (Reply, error) { 233 | switch v := value.(type) { 234 | case []byte: 235 | return &BulkReply{v}, nil 236 | case [][]byte: 237 | return &MultiBulkReply{v}, nil 238 | case int: 239 | return &IntReply{v}, nil 240 | case *StatusReply: 241 | return v, nil 242 | default: 243 | return nil, fmt.Errorf("Unsupported type: %s (%T)", v, v) 244 | } 245 | } 246 | 247 | func writeNullBytes(w io.Writer) (int64, error) { 248 | n, err := w.Write([]byte("$-1\r\n")) 249 | return int64(n), err 250 | } 251 | 252 | func writeBytes(value interface{}, w io.Writer) (int64, error) { 253 | if value == nil { 254 | return writeNullBytes(w) 255 | } 256 | switch v := value.(type) { 257 | case []byte: 258 | if len(v) == 0 { 259 | return writeNullBytes(w) 260 | } 261 | if wrote, err := w.Write([]byte("$" + strconv.Itoa(len(v)) + "\r\n")); err != nil { 262 | return int64(wrote), err 263 | } else if wroteData, err := w.Write([]byte(v)); err != nil { 264 | return int64(wrote + wroteData), err 265 | } else { 266 | wroteCRLF, err := w.Write([]byte("\r\n")) 267 | return int64(wrote + wroteData + wroteCRLF), err 268 | } 269 | case int: 270 | wrote, err := w.Write([]byte(":" + strconv.Itoa(v) + "\r\n")) 271 | return int64(wrote), err 272 | } 273 | return 0, fmt.Errorf("Invalid type sent to WriteBytes") 274 | } 275 | -------------------------------------------------------------------------------- /rocks_hashes.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "encoding/gob" 5 | "fmt" 6 | rocks "github.com/tecbot/gorocksdb" 7 | "reflect" 8 | ) 9 | 10 | func (rh *RocksDBHandler) RedisHkeys(key []byte) ([][]byte, error) { 11 | if err := rh.checkRedisCall(key); err != nil { 12 | return nil, err 13 | } 14 | if err := rh.checkKeyType(key, kRedisHash); err != nil { 15 | return nil, err 16 | } 17 | 18 | hashData, err := rh._hash_getData(key) 19 | if err != nil { 20 | return nil, err 21 | } 22 | data, index := make([][]byte, len(hashData)), 0 23 | for key := range hashData { 24 | data[index] = []byte(key) 25 | index++ 26 | } 27 | return data, nil 28 | } 29 | 30 | func (rh *RocksDBHandler) RedisHvals(key []byte) ([][]byte, error) { 31 | if err := rh.checkRedisCall(key); err != nil { 32 | return nil, err 33 | } 34 | if err := rh.checkKeyType(key, kRedisHash); err != nil { 35 | return nil, err 36 | } 37 | 38 | hashData, err := rh._hash_getData(key) 39 | if err != nil { 40 | return nil, err 41 | } 42 | data, index := make([][]byte, len(hashData)), 0 43 | for _, value := range hashData { 44 | data[index] = value 45 | index++ 46 | } 47 | return data, nil 48 | } 49 | 50 | func (rh *RocksDBHandler) RedisHlen(key []byte) (int, error) { 51 | if err := rh.checkRedisCall(key); err != nil { 52 | return 0, err 53 | } 54 | if err := rh.checkKeyType(key, kRedisHash); err != nil { 55 | return 0, err 56 | } 57 | 58 | hashData, err := rh._hash_getData(key) 59 | if err != nil { 60 | return 0, err 61 | } 62 | return len(hashData), nil 63 | } 64 | 65 | func (rh *RocksDBHandler) RedisHdel(key, field []byte, fields ...[]byte) (int, error) { 66 | if err := rh.checkRedisCall(key, field); err != nil { 67 | return 0, err 68 | } 69 | if err := rh.checkKeyType(key, kRedisHash); err != nil { 70 | return 0, err 71 | } 72 | 73 | hashData, err := rh._hash_getData(key) 74 | if err != nil { 75 | return 0, err 76 | } 77 | 78 | allFields := append([][]byte{field}, fields...) 79 | count := 0 80 | data := make([][]byte, 0) 81 | for _, f := range allFields { 82 | if value, ok := hashData[string(f)]; ok { 83 | count++ 84 | data = append(data, [][]byte{f, value}...) 85 | } 86 | } 87 | if count > 0 { 88 | if err := rh._hash_doMerge(key, data, kHashOpDelete); err != nil { 89 | return 0, err 90 | } 91 | } 92 | return count, nil 93 | } 94 | 95 | func (rh *RocksDBHandler) RedisHexists(key, field []byte) (int, error) { 96 | if err := rh.checkRedisCall(key, field); err != nil { 97 | return 0, err 98 | } 99 | if err := rh.checkKeyType(key, kRedisHash); err != nil { 100 | return 0, err 101 | } 102 | 103 | hashData, err := rh._hash_getData(key) 104 | if err != nil { 105 | return 0, err 106 | } 107 | _, ok := hashData[string(field)] 108 | if ok { 109 | return 1, nil 110 | } 111 | return 0, nil 112 | } 113 | 114 | func (rh *RocksDBHandler) RedisHget(key, field []byte) ([]byte, error) { 115 | if err := rh.checkRedisCall(key, field); err != nil { 116 | return nil, err 117 | } 118 | if err := rh.checkKeyType(key, kRedisHash); err != nil { 119 | return nil, err 120 | } 121 | 122 | hashData, err := rh._hash_getData(key) 123 | if err != nil { 124 | return nil, err 125 | } 126 | if value, ok := hashData[string(field)]; ok { 127 | return value, nil 128 | } else { 129 | return nil, nil 130 | } 131 | } 132 | 133 | func (rh *RocksDBHandler) RedisHmget(key, field []byte, fields ...[]byte) ([][]byte, error) { 134 | if err := rh.checkRedisCall(key, field); err != nil { 135 | return nil, err 136 | } 137 | if err := rh.checkKeyType(key, kRedisHash); err != nil { 138 | return nil, err 139 | } 140 | 141 | hashData, err := rh._hash_getData(key) 142 | if err != nil { 143 | return nil, err 144 | } 145 | allFields := append([][]byte{field}, fields...) 146 | data := make([][]byte, len(allFields)) 147 | for i, f := range allFields { 148 | if value, ok := hashData[string(f)]; ok { 149 | data[i] = value 150 | } else { 151 | data[i] = nil 152 | } 153 | } 154 | return data, nil 155 | } 156 | 157 | func (rh *RocksDBHandler) RedisHset(key, field, value []byte) (int, error) { 158 | if err := rh.checkRedisCall(key, field, value); err != nil { 159 | return 0, err 160 | } 161 | if err := rh.checkKeyType(key, kRedisHash); err != nil { 162 | return 0, err 163 | } 164 | 165 | // This is hurting the performance for 1-5ms 166 | // hashData, err := rh._hash_getData(key) 167 | // if err != nil { 168 | // return 0, err 169 | // } 170 | // _, exists := hashData[string(field)] 171 | if err := rh._hash_doMerge(key, [][]byte{field, value}, kHashOpSet); err != nil { 172 | return 0, err 173 | } 174 | // if exists { 175 | // return 0, nil 176 | // } 177 | return 1, nil 178 | } 179 | 180 | func (rh *RocksDBHandler) RedisHmset(key, field, value []byte, pairs ...[]byte) error { 181 | if err := rh.checkRedisCall(key, field, value); err != nil { 182 | return err 183 | } 184 | if err := rh.checkKeyType(key, kRedisHash); err != nil { 185 | return err 186 | } 187 | 188 | data := append([][]byte{field, value}, pairs...) 189 | if len(data)%2 != 0 { 190 | return ErrWrongArgumentsCount 191 | } 192 | return rh._hash_doMerge(key, data, kHashOpSet) 193 | } 194 | 195 | func (rh *RocksDBHandler) RedisHgetall(key []byte) ([][]byte, error) { 196 | if err := rh.checkRedisCall(key); err != nil { 197 | return nil, err 198 | } 199 | if err := rh.checkKeyType(key, kRedisHash); err != nil { 200 | return nil, err 201 | } 202 | 203 | hashData, err := rh._hash_getData(key) 204 | if err != nil { 205 | return nil, err 206 | } 207 | return __hashToBytes(hashData), nil 208 | } 209 | 210 | func (rh *RocksDBHandler) _hash_doMerge(key []byte, values [][]byte, opCode string) error { 211 | if values == nil || len(values) == 0 || len(values)%2 != 0 { 212 | return ErrWrongArgumentsCount 213 | } 214 | options := rocks.NewDefaultWriteOptions() 215 | defer options.Destroy() 216 | batch := rocks.NewWriteBatch() 217 | defer batch.Destroy() 218 | batch.Put(rh.getTypeKey(key), []byte(kRedisHash)) 219 | for i := 0; i < len(values); i += 2 { 220 | operand := HashOperand{opCode, string(values[i]), values[i+1]} 221 | if data, err := encode(operand); err == nil { 222 | batch.Merge(key, data) 223 | } else { 224 | return err 225 | } 226 | } 227 | return rh.db.Write(options, batch) 228 | } 229 | 230 | func (rh *RocksDBHandler) _hash_getData(key []byte) (map[string][]byte, error) { 231 | hashData := make(map[string][]byte) 232 | options := rocks.NewDefaultReadOptions() 233 | defer options.Destroy() 234 | if obj, err := rh.loadRedisObject(options, key); err != nil { 235 | if err == ErrDoesNotExist { 236 | return hashData, nil 237 | } 238 | return nil, err 239 | } else { 240 | if obj.Type != kRedisHash { 241 | return nil, ErrWrongTypeRedisObject 242 | } 243 | data := obj.Data.([][]byte) 244 | for i := 0; i < len(data); i += 2 { 245 | hashData[string(data[i])] = data[i+1] 246 | } 247 | return hashData, nil 248 | } 249 | } 250 | 251 | const ( 252 | kHashOpNone = "noop" 253 | kHashOpSet = "set" 254 | kHashOpDelete = "delete" 255 | ) 256 | 257 | type HashOperand struct { 258 | Command string 259 | Key string 260 | Value []byte 261 | } 262 | 263 | func init() { 264 | gob.Register(&HashOperand{}) 265 | } 266 | 267 | type HashMerger struct{} 268 | 269 | func (m *HashMerger) FullMerge(existingObject *RedisObject, operands [][]byte) bool { 270 | rawData, ok := existingObject.Data.([][]byte) 271 | if !ok { 272 | rawData = [][]byte{} 273 | } 274 | hashData := make(map[string][]byte) 275 | for i := 0; i < len(rawData); i += 2 { 276 | hashData[string(rawData[i])] = rawData[i+1] 277 | } 278 | for _, operand := range operands { 279 | if obj, err := decode(operand, reflect.TypeOf(HashOperand{})); err == nil { 280 | op := obj.(HashOperand) 281 | switch op.Command { 282 | case kHashOpSet: 283 | hashData[op.Key] = op.Value 284 | case kHashOpDelete: 285 | delete(hashData, op.Key) 286 | } 287 | } 288 | } 289 | 290 | existingObject.Data = __hashToBytes(hashData) 291 | return true 292 | } 293 | 294 | func (m *HashMerger) PartialMerge(leftOperand, rightOperand []byte) ([]byte, bool) { 295 | obj, err := decode(leftOperand, reflect.TypeOf(HashOperand{})) 296 | if err != nil { 297 | return nil, false 298 | } 299 | leftOp := obj.(HashOperand) 300 | obj, err = decode(rightOperand, reflect.TypeOf(HashOperand{})) 301 | if err != nil { 302 | return nil, false 303 | } 304 | rightOp := obj.(HashOperand) 305 | if leftOp.Key == rightOp.Key { 306 | // fmt.Println("PartialMerged", leftOp, rightOp) 307 | return rightOperand, true 308 | } 309 | 310 | return nil, false 311 | } 312 | 313 | func __hashToBytes(hashData map[string][]byte) [][]byte { 314 | rawData := make([][]byte, len(hashData)*2) 315 | index := 0 316 | for key, value := range hashData { 317 | rawData[index] = []byte(key) 318 | rawData[index+1] = value 319 | index += 2 320 | } 321 | return rawData 322 | } 323 | 324 | var _ = fmt.Println 325 | -------------------------------------------------------------------------------- /rocks.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "encoding/gob" 5 | "fmt" 6 | rocks "github.com/tecbot/gorocksdb" 7 | "log" 8 | "reflect" 9 | "strings" 10 | ) 11 | 12 | const ( 13 | kRedisString = "string" 14 | kRedisList = "list" 15 | kRedisHash = "hash" 16 | kRedisSet = "set" 17 | ) 18 | 19 | var ( 20 | kTypeKeyPrefix = []byte("__*type*__") 21 | ) 22 | 23 | type RedisObject struct { 24 | Type string 25 | Data interface{} 26 | } 27 | 28 | func init() { 29 | gob.Register(&RedisObject{}) 30 | gob.Register([][]byte{}) 31 | } 32 | 33 | type DataStructureMerger interface { 34 | FullMerge(existingObject *RedisObject, operands [][]byte) bool 35 | PartialMerge(leftOperand, rightOperand []byte) ([]byte, bool) 36 | } 37 | 38 | func NewRocksDBHandler(config RockdisConfig) *RocksDBHandler { 39 | cacheSize, err := parseComputerSize(config.Database.MaxMemory) 40 | if err != nil { 41 | log.Fatalf("[Config] Format error for [Database] maxmemory=%s", config.Database.MaxMemory) 42 | } 43 | blockSize, err := parseComputerSize(config.Database.BlockSize) 44 | if err != nil { 45 | log.Fatalf("[Config] Format error for [Database] blocksize=%s", config.Database.BlockSize) 46 | } 47 | 48 | handler := &RocksDBHandler{} 49 | handler.dbDir = config.Database.DbDir 50 | handler.cacheSize = cacheSize 51 | handler.blockSize = blockSize 52 | handler.createIfMissing = config.Database.CreateIfMissing 53 | handler.bloomFilter = config.Database.BloomFilter 54 | handler.compression = config.Database.Compression 55 | handler.compactionStyle = config.Database.CompactionStyle 56 | handler.maxOpenFiles = config.Database.MaxOpenFiles 57 | handler.maxMerge = config.Database.MaxMerge 58 | 59 | if err := handler.Init(); err != nil { 60 | log.Fatal(err) 61 | } 62 | return handler 63 | } 64 | 65 | type RocksDBHandler struct { 66 | dbDir string 67 | cacheSize int 68 | blockSize int 69 | createIfMissing bool 70 | bloomFilter int 71 | compression string 72 | compactionStyle string 73 | maxOpenFiles int 74 | maxMerge int 75 | 76 | options *rocks.Options 77 | db *rocks.DB 78 | 79 | dsMergers map[string]DataStructureMerger 80 | } 81 | 82 | func (rh *RocksDBHandler) Init() error { 83 | rh.options = rocks.NewDefaultOptions() 84 | rh.options.SetBlockCache(rocks.NewLRUCache(rh.cacheSize)) 85 | rh.options.SetBlockSize(rh.blockSize) 86 | rh.options.SetCreateIfMissing(rh.createIfMissing) 87 | if rh.bloomFilter > 0 { 88 | rh.options.SetFilterPolicy(rocks.NewBloomFilter(rh.bloomFilter)) 89 | } 90 | if rh.maxOpenFiles > 0 { 91 | rh.options.SetMaxOpenFiles(rh.maxOpenFiles) 92 | } 93 | 94 | switch rh.compression { 95 | case "no": 96 | rh.options.SetCompression(rocks.NoCompression) 97 | case "snappy": 98 | rh.options.SetCompression(rocks.SnappyCompression) 99 | case "zlib": 100 | rh.options.SetCompression(rocks.ZlibCompression) 101 | case "bzip2": 102 | rh.options.SetCompression(rocks.BZip2Compression) 103 | } 104 | 105 | switch rh.compactionStyle { 106 | case "level": 107 | rh.options.SetCompactionStyle(rocks.LevelCompactionStyle) 108 | case "universal": 109 | rh.options.SetCompactionStyle(rocks.UniversalCompactionStyle) 110 | } 111 | 112 | rh.dsMergers = make(map[string]DataStructureMerger) 113 | rh.dsMergers[kRedisString] = &StringMerger{} 114 | rh.dsMergers[kRedisList] = &ListMerger{} 115 | rh.dsMergers[kRedisHash] = &HashMerger{} 116 | rh.dsMergers[kRedisSet] = &SetMerger{} 117 | 118 | if rh.maxMerge > 0 { 119 | rh.options.SetMaxSuccessiveMerges(rh.maxMerge) 120 | } 121 | rh.options.SetMergeOperator(rocks.NewMergeOperator(rh)) 122 | 123 | db, err := rocks.OpenDb(rh.options, rh.dbDir) 124 | if err != nil { 125 | rh.Close() 126 | return err 127 | } 128 | rh.db = db 129 | 130 | infos := []string{ 131 | fmt.Sprintf("dbDir=%s", rh.dbDir), 132 | fmt.Sprintf("cacheSize=%d", rh.cacheSize), 133 | fmt.Sprintf("blockSize=%d", rh.blockSize), 134 | fmt.Sprintf("createIfMissing=%v", rh.createIfMissing), 135 | fmt.Sprintf("bloomFilter=%d", rh.bloomFilter), 136 | fmt.Sprintf("compression=%s", rh.compression), 137 | fmt.Sprintf("compactionStyle=%s", rh.compactionStyle), 138 | fmt.Sprintf("maxOpenFiles=%d", rh.maxOpenFiles), 139 | fmt.Sprintf("maxMerge=%d", rh.maxMerge), 140 | } 141 | log.Printf("[RocksDBHandler] Inited, %s", strings.Join(infos, ", ")) 142 | return nil 143 | } 144 | 145 | func (rh *RocksDBHandler) Close() { 146 | if rh.options != nil { 147 | rh.options.Destroy() 148 | } 149 | if rh.db != nil { 150 | rh.db.Close() 151 | } 152 | log.Printf("[RocksDBHandler] Closed.") 153 | } 154 | 155 | func (rh *RocksDBHandler) getTypeKey(key []byte) []byte { 156 | return append(kTypeKeyPrefix, key...) 157 | } 158 | 159 | func (rh *RocksDBHandler) getKeyType(key []byte) (string, error) { 160 | if rh.db == nil { 161 | return "", ErrRocksIsDead 162 | } 163 | if key == nil || len(key) == 0 { 164 | return "", ErrWrongArgumentsCount 165 | } 166 | 167 | options := rocks.NewDefaultReadOptions() 168 | if slice, err := rh.db.Get(options, rh.getTypeKey(key)); err == nil { 169 | defer slice.Free() 170 | return string(slice.Data()), nil 171 | } else { 172 | return "", err 173 | } 174 | } 175 | 176 | func (rh *RocksDBHandler) FullMerge(key, existingValue []byte, operands [][]byte) ([]byte, bool) { 177 | var redisObj RedisObject 178 | keyType, err := rh.getKeyType(key) 179 | if err != nil || keyType == "" { 180 | return nil, false 181 | } 182 | var emptyData interface{} 183 | switch keyType { 184 | case kRedisString: 185 | emptyData = []byte{} 186 | default: 187 | emptyData = [][]byte{} 188 | } 189 | 190 | if existingValue == nil || len(existingValue) == 0 { 191 | redisObj = RedisObject{keyType, emptyData} 192 | } else { 193 | if obj, err := decode(existingValue, reflect.TypeOf(redisObj)); err != nil { 194 | return nil, false 195 | } else { 196 | redisObj = obj.(RedisObject) 197 | } 198 | if redisObj.Type != keyType { 199 | return nil, false 200 | } 201 | } 202 | 203 | if merger, ok := rh.dsMergers[keyType]; ok { 204 | merged := merger.FullMerge(&redisObj, operands) 205 | if !merged { 206 | return nil, false 207 | } 208 | if data, err := encode(redisObj); err == nil { 209 | return data, true 210 | } 211 | } 212 | 213 | return nil, false 214 | } 215 | 216 | func (rh *RocksDBHandler) PartialMerge(key, leftOperand, rightOperand []byte) ([]byte, bool) { 217 | keyType, err := rh.getKeyType(key) 218 | if err != nil { 219 | return nil, false 220 | } 221 | if merger, ok := rh.dsMergers[keyType]; ok { 222 | return merger.PartialMerge(leftOperand, rightOperand) 223 | } 224 | return nil, false 225 | } 226 | 227 | func (rh *RocksDBHandler) Name() string { 228 | return "GoRockdisMergeOperator" 229 | } 230 | 231 | var ( 232 | ErrRocksIsDead = fmt.Errorf("RocksDB is dead") 233 | ErrDoesNotExist = fmt.Errorf("There is no such object") 234 | ErrWrongArgumentsCount = fmt.Errorf("Wrong number of arguments for command") 235 | ErrWrongTypeRedisObject = fmt.Errorf("Operation against a key holding the wrong kind of value") 236 | ErrNotNumber = fmt.Errorf("value is not an integer or out of range") 237 | ) 238 | 239 | func (rh *RocksDBHandler) copySlice(slice *rocks.Slice, toFree bool) []byte { 240 | data := make([]byte, slice.Size()) 241 | copy(data, slice.Data()) 242 | if toFree { 243 | slice.Free() 244 | } 245 | return data 246 | } 247 | 248 | func (rh *RocksDBHandler) loadRedisObject(options *rocks.ReadOptions, key []byte) (RedisObject, error) { 249 | slice, err := rh.db.Get(options, key) 250 | if err != nil { 251 | log.Printf("[loadRedisObject] Error when GET < RocksDB, %s", err) 252 | return RedisObject{}, err 253 | } 254 | 255 | data := rh.copySlice(slice, true) 256 | if data == nil || len(data) == 0 { 257 | globalStat.keyMisses.Add(1) 258 | return RedisObject{}, ErrDoesNotExist 259 | } 260 | 261 | if obj, err := decode(data, reflect.TypeOf(RedisObject{})); err == nil { 262 | globalStat.keyHits.Add(1) 263 | return obj.(RedisObject), nil 264 | } else { 265 | return RedisObject{}, err 266 | } 267 | } 268 | 269 | func (rh *RocksDBHandler) saveRedisObject(options *rocks.WriteOptions, key []byte, value interface{}, objType string) error { 270 | obj := RedisObject{ 271 | Type: objType, 272 | Data: value, 273 | } 274 | data, err := encode(obj) 275 | if err != nil { 276 | return err 277 | } 278 | 279 | batch := rocks.NewWriteBatch() 280 | defer batch.Destroy() 281 | batch.Put(rh.getTypeKey(key), []byte(objType)) 282 | batch.Put(key, data) 283 | err = rh.db.Write(options, batch) 284 | if err != nil { 285 | log.Printf("[saveRedisObject] Error when PUT > RocksDB, %s", err) 286 | } 287 | return err 288 | } 289 | 290 | func (rh *RocksDBHandler) checkRedisCall(args ...[]byte) error { 291 | if rh.db == nil { 292 | return ErrRocksIsDead 293 | } 294 | if len(args) > 0 { 295 | for _, arg := range args { 296 | if arg == nil || len(arg) == 0 { 297 | return ErrWrongArgumentsCount 298 | } 299 | } 300 | } 301 | return nil 302 | } 303 | 304 | func (rh *RocksDBHandler) checkKeyType(key []byte, assertType string) error { 305 | if keyType, err := rh.getKeyType(key); err != nil { 306 | return err 307 | } else { 308 | if keyType != "" && keyType != assertType { 309 | return ErrWrongTypeRedisObject 310 | } 311 | } 312 | return nil 313 | } 314 | --------------------------------------------------------------------------------