Merge 4b9b91f into 4fa302e

command/init.go

@@ -28,11 +28,11 @@ func init() {
 		"getset":   GetSet,
 		"getrange": GetRange,
 		// "msetnx":   MSetNx,
-		"setnx":    SetNx,
-		"setex":    SetEx,
-		"psetex":   PSetEx,
-		"setrange": SetRange,
-		// "setbit":      SetBit,
+		"setnx":       SetNx,
+		"setex":       SetEx,
+		"psetex":      PSetEx,
+		"setrange":    SetRange,
+		"setbit":      SetBit,
 		"incr":        Incr,
 		"incrby":      IncrBy,
 		"decr":        Decr,
@@ -127,6 +127,7 @@ func init() {
 		"incrby":      Desc{Proc: AutoCommit(IncrBy), Cons: Constraint{3, flags("wmF"), 1, 1, 1}},
 		"decrby":      Desc{Proc: AutoCommit(DecrBy), Cons: Constraint{3, flags("wmF"), 1, 1, 1}},
 		"incrbyfloat": Desc{Proc: AutoCommit(IncrByFloat), Cons: Constraint{3, flags("wmF"), 1, 1, 1}},
+		"setbit":      Desc{Proc: AutoCommit(SetBit), Cons: Constraint{4, flags("wm"), 1, 1, 1}},
 
 		// keys
 		"type":      Desc{Proc: AutoCommit(Type), Cons: Constraint{2, flags("rF"), 1, 1, 1}},

command/strings.go

@@ -469,3 +469,41 @@ func DecrBy(ctx *Context, txn *db.Transaction) (OnCommit, error) {
 	}
 	return Integer(ctx.Out, int64(delta)), nil
 }
+
+// SetBit Sets or clears the bit at offset in the string value stored at key.
+func SetBit(ctx *Context, txn *db.Transaction) (OnCommit, error) {
+	key := []byte(ctx.Args[0])
+	offset, err := strconv.ParseInt(string(ctx.Args[1]), 10, 32)
+	if err != nil {
+		return nil, ErrBitOffset
+	}
+	if offset < 0 {
+		return nil, ErrBitOffset
+	}
+
+	on, err := strconv.ParseInt(string(ctx.Args[2]), 10, 32)
+	if err != nil {
+		return nil, ErrBitInteger
+	}
+
+	// Bits can only be set or cleared...
+	if (on & ^1) != 0 {
+		return nil, ErrBitInteger
+	}
+
+	str, err := txn.String(key)
+	if err != nil {
+		if err == db.ErrTypeMismatch {
+			return nil, ErrTypeMismatch
+		}
+		return nil, errors.New("ERR " + err.Error())
+	}
+	val, err := str.SetBit(uint(on), uint(offset))
+	if err != nil {
+		return nil, errors.New("ERR " + err.Error())
+	}
+	if val != 0 {
+		return Integer(ctx.Out, 1), nil
+	}
+	return Integer(ctx.Out, 0), nil
+}

command/strings_test.go

@@ -497,11 +497,55 @@ func TestStringAppend(t *testing.T) {
 	args[0] = "Append"
 	args[1] = "Ap"
 
-	ctx := ContextTest("append", args...)
-	Call(ctx)
-	assert.Contains(t, ctxString(ctx.Out), strconv.Itoa(len(args[1])))
+	out := CallTest("append", args...)
+	assert.Contains(t, out.String(), strconv.Itoa(len(args[1])))
 
-	ctx = ContextTest("append", args...)
-	Call(ctx)
-	assert.Contains(t, ctxString(ctx.Out), strconv.Itoa(len(args[1])*2))
+	out = CallTest("append", args...)
+	assert.Contains(t, out.String(), strconv.Itoa(len(args[1])*2))
+}
+
+func TestStringSetBit(t *testing.T) {
+	tests := []struct {
+		name string
+		args []string
+		want string
+	}{
+		{
+			name: "1",
+			args: []string{"setbit", "1", "0", "0"},
+			want: "-ERR wrong number of arguments for 'setbit' command",
+		},
+		{
+			name: "2",
+			args: []string{"setbit", "x", "0"},
+			want: ErrBitOffset.Error(),
+		},
+		{
+			name: "3",
+			args: []string{"setbit", "1", "x"},
+			want: ErrBitInteger.Error(),
+		},
+		{
+			name: "4",
+			args: []string{"setbit", "1", "2"},
+			want: ErrBitInteger.Error(),
+		},
+		{
+			name: "5",
+			args: []string{"setbit", "1", "1"},
+			want: ":0",
+		},
+		{
+			name: "6",
+			args: []string{"setbit", "1", "0"},
+			want: ":1",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			out := CallTest("setbit", tt.args...)
+			assert.Contains(t, out.String(), tt.want)
+		})
+	}
 }

db/db_test.go

@@ -1,10 +1,12 @@
 package db
 
 import (
+	"context"
 	"os"
 	"testing"
 
 	"github.com/pingcap/tidb/store/mockstore"
+	"github.com/stretchr/testify/assert"
 )
 
 var mockDB *DB
@@ -22,3 +24,11 @@ func TestMain(m *testing.M) {
 
 	os.Exit(m.Run())
 }
+
+func MockTest(t *testing.T, callFunc func(txn *Transaction)) {
+	txn, err := mockDB.Begin()
+	assert.NoError(t, err)
+	callFunc(txn)
+	err = txn.Commit(context.TODO())
+	assert.NoError(t, err)
+}

db/string.go

@@ -10,14 +10,14 @@ type StringMeta struct {
 	Value []byte
 }
 
-//String object operate tikv
+// String object operate tikv
 type String struct {
 	Meta StringMeta
 	key  []byte
 	txn  *Transaction
 }
 
-//GetString return string object ,
+// GetString return string object ,
 // if key is exist , object load meta
 // otherwise object is null if key is not exist and err is not found
 // otherwise  return err
@@ -40,7 +40,7 @@ func GetString(txn *Transaction, key []byte) (*String, error) {
 	return str, nil
 }
 
-//NewString  create new string object
+// NewString  create new string object
 func NewString(txn *Transaction, key []byte) *String {
 	str := &String{txn: txn, key: key}
 	now := Now()
@@ -53,17 +53,17 @@ func NewString(txn *Transaction, key []byte) *String {
 	return str
 }
 
-//Get the value information for the key from db
+// Get the value information for the key from db
 func (s *String) Get() ([]byte, error) {
 	if !s.Exist() {
 		return nil, ErrKeyNotFound
 	}
 	return s.Meta.Value, nil
 }
 
-//Set set the string value of a key
-//the num of expire slice is not zero and expire[0] is not zero ,the key add exprie queue
-//otherwise the delete expire queue
+// Set set the string value of a key
+// the num of expire slice is not zero and expire[0] is not zero ,the key add exprie queue
+// otherwise the delete expire queue
 func (s *String) Set(val []byte, expire ...int64) error {
 	timestamp := Now()
 	mkey := MetaKey(s.txn.db, s.key)
@@ -82,20 +82,20 @@ func (s *String) Set(val []byte, expire ...int64) error {
 	return s.txn.t.Set(mkey, s.encode())
 }
 
-//Len value len
+// Len value len
 func (s *String) Len() (int, error) {
 	return len(s.Meta.Value), nil
 }
 
-//Exist return ture if key exist
+// Exist return ture if key exist
 func (s *String) Exist() bool {
 	if s.Meta.Value == nil {
 		return false
 	}
 	return true
 }
 
-//Append append a value to key
+// Append append a value to key
 func (s *String) Append(value []byte) (int, error) {
 	s.Meta.Value = append(s.Meta.Value, value...)
 	s.Meta.ExpireAt = 0
@@ -105,7 +105,7 @@ func (s *String) Append(value []byte) (int, error) {
 	return len(s.Meta.Value), nil
 }
 
-//GetSet return old value ,value replace old value
+// GetSet return old value ,value replace old value
 func (s *String) GetSet(value []byte) ([]byte, error) {
 	v := s.Meta.Value
 	if err := s.Set(value); err != nil {
@@ -114,7 +114,7 @@ func (s *String) GetSet(value []byte) ([]byte, error) {
 	return v, nil
 }
 
-//GetRange return string from the absolute of start to the absolute of end
+// GetRange return string from the absolute of start to the absolute of end
 func (s *String) GetRange(start, end int) []byte {
 	vlen := len(s.Meta.Value)
 	if end < 0 {
@@ -135,7 +135,7 @@ func (s *String) GetRange(start, end int) []byte {
 	return s.Meta.Value[start : end+1]
 }
 
-//SetRange Overwrites part of the string stored at key, starting at the specified offset, for the entire length of value.
+// SetRange Overwrites part of the string stored at key, starting at the specified offset, for the entire length of value.
 func (s *String) SetRange(offset int64, value []byte) ([]byte, error) {
 	val := s.Meta.Value
 	if int64(len(val)) < offset+int64(len(value)) {
@@ -149,7 +149,7 @@ func (s *String) SetRange(offset int64, value []byte) ([]byte, error) {
 	return val, nil
 }
 
-//Incr increment the integer value by the given amount
+// Incr increment the integer value by the given amount
 // the old value  must be integer
 func (s *String) Incr(delta int64) (int64, error) {
 	value := s.Meta.Value
@@ -169,7 +169,7 @@ func (s *String) Incr(delta int64) (int64, error) {
 
 }
 
-//Incrf increment the float value by the given amount
+// Incrf increment the float value by the given amount
 // the old value  must be float
 func (s *String) Incrf(delta float64) (float64, error) {
 	value := s.Meta.Value
@@ -188,14 +188,40 @@ func (s *String) Incrf(delta float64) (float64, error) {
 	return delta, nil
 }
 
-//encode because of the value is small size , value and meta decode together
+// SetBit key offset bitvalue
+// return the off postion of value
+func (s *String) SetBit(on, off uint) (uint, error) {
+	val := s.Meta.Value
+	bitoff := off >> 3
+	llen := int(bitoff) - len(val) + 1
+	if llen > 0 {
+		val = append(val, make([]byte, llen)...)
+	}
+
+	/* Get current values */
+	byteval := uint(val[bitoff])
+	bit := uint(7 - (off & 0x7))
+	bitval := byteval & (1 << bit)
+
+	/* Update byte with new bit value and return original value */
+	byteval &= (^(1 << bit))
+	byteval = byteval | ((on & 0x1) << bit)
+	val[bitoff] = byte(byteval)
+
+	if err := s.Set(val); err != nil {
+		return 0, err
+	}
+	return bitval, nil
+}
+
+// encode because of the value is small size , value and meta decode together
 func (s *String) encode() []byte {
 	b := EncodeObject(&s.Meta.Object)
 	b = append(b, s.Meta.Value...)
 	return b
 }
 
-//decode if obj has been existed , stop parse
+// decode if obj has been existed , stop parse
 func (s *String) decode(b []byte) error {
 	obj, err := DecodeObject(b)
 	if err != nil {