Merge pull request #8 from youzan/more-cmd-support

More bit cmd support

Author: absolute8511

.travis.yml

@@ -1,7 +1,7 @@
 language: go
 go:
-  - 1.8.x
   - 1.9.x
+  - 1.10.x
 env:
   - GOARCH=amd64 TEST_RACE=false
   - GOARCH=amd64 TEST_RACE=true

README.md

@@ -38,7 +38,7 @@ CGO_CFLAGS="-I/path/to/rocksdb/include" CGO_LDFLAGS="-L/path/to/rocksdb -lrocksd
 
 use the `dep ensure` to install other dependencies
 
-Build zankv and placedriver from the source (only support go version 1.8+, gcc 4.9+ or xcode-command-line-tools on Mac):
+Build zankv and placedriver from the source (only support go version 1.9+, gcc 4.9+ or xcode-command-line-tools on Mac):
 <pre>
 make
 </pre>

node/keys.go

@@ -1,11 +1,13 @@
 package node
 
 import (
+	"fmt"
 	"strconv"
 	"time"
 
-	"github.com/youzan/ZanRedisDB/common"
 	"github.com/absolute8511/redcon"
+	"github.com/youzan/ZanRedisDB/common"
+	"github.com/youzan/ZanRedisDB/rockredis"
 )
 
 func (nd *KVNode) Lookup(key []byte) ([]byte, error) {
@@ -32,6 +34,51 @@ func (nd *KVNode) existsCommand(cmd redcon.Command) (interface{}, error) {
 	return val, err
 }
 
+func (nd *KVNode) getbitCommand(conn redcon.Conn, cmd redcon.Command) {
+	if len(cmd.Args) < 3 {
+		conn.WriteError(errWrongNumberArgs.Error())
+		return
+	}
+	offset, err := strconv.ParseInt(string(cmd.Args[2]), 10, 64)
+	if err != nil {
+		conn.WriteError(err.Error())
+		return
+	}
+	val, err := nd.store.BitGet(cmd.Args[1], offset)
+	if err != nil {
+		conn.WriteError(err.Error())
+		return
+	}
+	conn.WriteInt64(val)
+}
+
+func (nd *KVNode) bitcountCommand(conn redcon.Conn, cmd redcon.Command) {
+	if len(cmd.Args) != 2 && len(cmd.Args) != 4 {
+		conn.WriteError(errWrongNumberArgs.Error())
+		return
+	}
+	start, end := int64(0), int64(-1)
+	var err error
+	if len(cmd.Args) == 4 {
+		start, err = strconv.ParseInt(string(cmd.Args[2]), 10, 64)
+		if err != nil {
+			conn.WriteError(err.Error())
+			return
+		}
+		end, err = strconv.ParseInt(string(cmd.Args[3]), 10, 64)
+		if err != nil {
+			conn.WriteError(err.Error())
+			return
+		}
+	}
+	val, err := nd.store.BitCount(cmd.Args[1], int(start), int(end))
+	if err != nil {
+		conn.WriteError(err.Error())
+		return
+	}
+	conn.WriteInt64(val)
+}
+
 func (nd *KVNode) mgetCommand(conn redcon.Conn, cmd redcon.Command) {
 	vals, _ := nd.store.MGet(cmd.Args[1:]...)
 	conn.WriteArray(len(vals))
@@ -67,6 +114,42 @@ func (nd *KVNode) setnxCommand(conn redcon.Conn, cmd redcon.Command, v interface
 	}
 }
 
+func (nd *KVNode) setbitCommand(conn redcon.Conn, cmd redcon.Command) {
+	if len(cmd.Args) != 4 {
+		conn.WriteError(errWrongNumberArgs.Error())
+		return
+	}
+
+	offset, err := strconv.ParseInt(string(cmd.Args[2]), 10, 64)
+	if err != nil {
+		conn.WriteError(err.Error())
+		return
+	}
+	on, err := strconv.ParseInt(string(cmd.Args[3]), 10, 64)
+	if err != nil {
+		conn.WriteError(err.Error())
+		return
+	}
+	if offset > rockredis.MaxBitOffset {
+		conn.WriteError(rockredis.ErrBitOverflow.Error())
+		return
+	}
+	if (on & ^1) != 0 {
+		conn.WriteError(fmt.Errorf("bit should be 0 or 1, got %d", on).Error())
+		return
+	}
+
+	_, v, ok := rebuildFirstKeyAndPropose(nd, conn, cmd)
+	if !ok {
+		return
+	}
+	if rsp, ok := v.(int64); ok {
+		conn.WriteInt64(rsp)
+	} else {
+		conn.WriteError(errInvalidResponse.Error())
+	}
+}
+
 func (nd *KVNode) msetCommand(cmd redcon.Command, v interface{}) (interface{}, error) {
 	return nil, nil
 }
@@ -158,3 +241,15 @@ func (kvsm *kvStoreSM) localPFAddCommand(cmd redcon.Command, ts int64) (interfac
 	v, err := kvsm.store.PFAdd(ts, cmd.Args[1], cmd.Args[2:]...)
 	return v, err
 }
+
+func (kvsm *kvStoreSM) localBitSetCommand(cmd redcon.Command, ts int64) (interface{}, error) {
+	offset, err := strconv.ParseInt(string(cmd.Args[2]), 10, 64)
+	if err != nil {
+		return 0, err
+	}
+	on, err := strconv.ParseInt(string(cmd.Args[3]), 10, 64)
+	if err != nil {
+		return 0, err
+	}
+	return kvsm.store.BitSet(ts, cmd.Args[1], offset, int(on))
+}

node/keys_test.go

@@ -11,12 +11,12 @@ import (
 	"testing"
 	"time"
 
+	"github.com/absolute8511/redcon"
+	"github.com/stretchr/testify/assert"
 	"github.com/youzan/ZanRedisDB/common"
 	"github.com/youzan/ZanRedisDB/rockredis"
 	"github.com/youzan/ZanRedisDB/stats"
 	"github.com/youzan/ZanRedisDB/transport/rafthttp"
-	"github.com/absolute8511/redcon"
-	"github.com/stretchr/testify/assert"
 )
 
 func getTestKVNode(t *testing.T) (*KVNode, string, chan struct{}) {
@@ -150,6 +150,7 @@ func TestKVNode_kvCommand(t *testing.T) {
 	testKey2 := []byte("default:test:2")
 	testKey2Value := []byte("2")
 	testPFKey := []byte("default:test:pf1")
+	testBitKey := []byte("default:test:bit1")
 	tests := []struct {
 		name string
 		args redcon.Command
@@ -169,6 +170,9 @@ func TestKVNode_kvCommand(t *testing.T) {
 		{"exists", buildCommand([][]byte{[]byte("exists"), testKey})},
 		{"pfadd", buildCommand([][]byte{[]byte("pfadd"), testPFKey, testKeyValue})},
 		{"pfcount", buildCommand([][]byte{[]byte("pfcount"), testPFKey})},
+		{"setbit", buildCommand([][]byte{[]byte("setbit"), testBitKey, []byte("1"), []byte("1")})},
+		{"getbit", buildCommand([][]byte{[]byte("getbit"), testBitKey, []byte("1")})},
+		{"bitcount", buildCommand([][]byte{[]byte("bitcount"), testBitKey, []byte("1"), []byte("2")})},
 	}
 	defer os.RemoveAll(dataDir)
 	defer nd.Stop()

node/namespace.go

@@ -337,10 +337,20 @@ func (nsm *NamespaceMgr) InitNamespaceNode(conf *NamespaceConfig, raftID uint64,
 	if err != nil {
 		return nil, err
 	}
-	if _, ok := nsm.nsMetas[conf.BaseName]; !ok {
+	if oldMeta, ok := nsm.nsMetas[conf.BaseName]; !ok {
 		nsm.nsMetas[conf.BaseName] = NamespaceMeta{
 			PartitionNum: conf.PartitionNum,
 		}
+		nodeLog.Infof("namespace meta init: %v", conf)
+	} else {
+		if oldMeta.PartitionNum != conf.PartitionNum {
+			nodeLog.Errorf("namespace meta mismatch: %v, old: %v", conf, oldMeta)
+			// update the meta if mismatch, it may happen if create the same namespace with different
+			// config for old deleted namespace
+			nsm.nsMetas[conf.BaseName] = NamespaceMeta{
+				PartitionNum: conf.PartitionNum,
+			}
+		}
 	}
 
 	n := &NamespaceNode{

node/node.go

@@ -31,6 +31,7 @@ var (
 	errUnknownData          = errors.New("unknown request data type")
 	errTooMuchBatchSize     = errors.New("the batch size exceed the limit")
 	errRaftNotReadyForWrite = errors.New("ERR_CLUSTER_CHANGED: the raft is not ready for write")
+	errWrongNumberArgs      = errors.New("ERR wrong number of arguments for redis command")
 )
 
 const (

node/node_cmd_reg.go

@@ -5,6 +5,7 @@ func (kvsm *kvStoreSM) registerHandlers() {
 	// kv
 	kvsm.router.RegisterInternal("del", kvsm.localDelCommand)
 	kvsm.router.RegisterInternal("set", kvsm.localSetCommand)
+	kvsm.router.RegisterInternal("setbit", kvsm.localBitSetCommand)
 	kvsm.router.RegisterInternal("setnx", kvsm.localSetnxCommand)
 	kvsm.router.RegisterInternal("mset", kvsm.localMSetCommand)
 	kvsm.router.RegisterInternal("incr", kvsm.localIncrCommand)
@@ -72,8 +73,11 @@ func (nd *KVNode) registerHandler() {
 	}
 	// for kv
 	nd.router.Register(false, "get", wrapReadCommandK(nd.getCommand))
+	nd.router.Register(false, "getbit", wrapReadCommandKAnySubkeyN(nd.getbitCommand, 1))
+	nd.router.Register(false, "bitcount", wrapReadCommandKAnySubkey(nd.bitcountCommand))
 	nd.router.Register(false, "mget", wrapReadCommandKK(nd.mgetCommand))
 	nd.router.Register(true, "set", wrapWriteCommandKV(nd, nd.setCommand))
+	nd.router.Register(true, "setbit", nd.setbitCommand)
 	nd.router.Register(true, "setnx", wrapWriteCommandKV(nd, nd.setnxCommand))
 	nd.router.Register(true, "incr", wrapWriteCommandK(nd, nd.incrCommand))
 	nd.router.Register(true, "incrby", wrapWriteCommandKV(nd, nd.incrbyCommand))

rockredis/t_kv.go

@@ -19,7 +19,7 @@ const (
 
 var errKVKey = errors.New("invalid encode kv key")
 var errInvalidDBValue = errors.New("invalide db value")
-var errBitOverflow = errors.New("bit offset overflowed")
+var ErrBitOverflow = errors.New("bit offset overflowed")
 
 func convertRedisKeyToDBKVKey(key []byte) ([]byte, []byte, error) {
 	table, _, _ := extractTableFromRedisKey(key)
@@ -530,7 +530,7 @@ func (db *RockDB) BitSet(ts int64, key []byte, offset int64, on int) (int64, err
 		return 0, err
 	}
 	if offset > MaxBitOffset {
-		return 0, errBitOverflow
+		return 0, ErrBitOverflow
 	}
 
 	if (on & ^1) != 0 {
@@ -625,6 +625,9 @@ func (db *RockDB) BitCount(key []byte, start, end int) (int64, error) {
 		return 0, err
 	}
 	start, end = getRange(start, end, len(v))
+	if start > end {
+		return 0, nil
+	}
 	v = v[start : end+1]
 	return popcountBytes(v), nil
 }

server/redis_api_test.go

@@ -332,6 +332,81 @@ func TestKVIncrDecr(t *testing.T) {
 	}
 }
 
+func TestKVBitOp(t *testing.T) {
+	c := getTestConn(t)
+	defer c.Close()
+
+	key := "default:test:kv_bitop"
+	if n, err := goredis.Int64(c.Do("bitcount", key)); err != nil {
+		t.Fatal(err)
+	} else if n != 0 {
+		t.Fatal(n)
+	}
+
+	if n, err := goredis.Int64(c.Do("getbit", key, 1)); err != nil {
+		t.Fatal(err)
+	} else if n != 0 {
+		t.Fatal(n)
+	}
+
+	if n, err := goredis.Int64(c.Do("setbit", key, 100, 1)); err != nil {
+		t.Fatal(err)
+	} else if n != 0 {
+		t.Fatal(n)
+	}
+
+	if n, err := goredis.Int64(c.Do("getbit", key, 100)); err != nil {
+		t.Fatal(err)
+	} else if n != 1 {
+		t.Fatal(n)
+	}
+	if n, err := goredis.Int64(c.Do("bitcount", key)); err != nil {
+		t.Fatal(err)
+	} else if n != 1 {
+		t.Fatal(n)
+	}
+	if n, err := goredis.Int64(c.Do("setbit", key, 1, 1)); err != nil {
+		t.Fatal(err)
+	} else if n != 0 {
+		t.Fatal(n)
+	}
+	if n, err := goredis.Int64(c.Do("bitcount", key)); err != nil {
+		t.Fatal(err)
+	} else if n != 2 {
+		t.Fatal(n)
+	}
+	if n, err := goredis.Int64(c.Do("bitcount", key, 0, 0)); err != nil {
+		t.Fatal(err)
+	} else if n != 1 {
+		t.Fatal(n)
+	}
+	if n, err := goredis.Int64(c.Do("setbit", key, 8, 1)); err != nil {
+		t.Fatal(err)
+	} else if n != 0 {
+		t.Fatal(n)
+	}
+	if n, err := goredis.Int64(c.Do("bitcount", key, 0, 0)); err != nil {
+		t.Fatal(err)
+	} else if n != 1 {
+		t.Fatal(n)
+	}
+	if n, err := goredis.Int64(c.Do("setbit", key, 7, 1)); err != nil {
+		t.Fatal(err)
+	} else if n != 0 {
+		t.Fatal(n)
+	}
+	if n, err := goredis.Int64(c.Do("bitcount", key, 0, 0)); err != nil {
+		t.Fatal(err)
+	} else if n != 2 {
+		t.Fatal(n)
+	}
+	if n, err := goredis.Int64(c.Do("bitcount", key, 0, 1)); err != nil {
+		t.Fatal(err)
+	} else if n != 3 {
+		t.Fatal(n)
+	}
+}
+
 func TestKVBatch(t *testing.T) {
 
 	var wg sync.WaitGroup
@@ -546,6 +621,24 @@ func TestKVErrorParams(t *testing.T) {
 
 	_, err = c.Do("mget")
 	assert.NotNil(t, err)
+
+	_, err = c.Do("getbit")
+	assert.NotNil(t, err)
+
+	_, err = c.Do("getbit", key1)
+	assert.NotNil(t, err)
+
+	_, err = c.Do("setbit", key1)
+	assert.NotNil(t, err)
+
+	_, err = c.Do("setbit")
+	assert.NotNil(t, err)
+
+	_, err = c.Do("bitcount")
+	assert.NotNil(t, err)
+
+	_, err = c.Do("bitcount", key1, "0")
+	assert.NotNil(t, err)
 }
 
 func TestPFOp(t *testing.T) {