add szet

Author: latermoon

redis/command.go

@@ -23,14 +23,6 @@ func (c Command) Bytes() []byte {
 	return buf.Bytes()
 }
 
-func (c Command) Size() int {
-	sum := 0
-	for i := 0; i < len(c); i++ {
-		sum += len(c[i])
-	}
-	return sum
-}
-
 func (c Command) String() string {
 	arr := make([]string, len(c))
 	for i := range c {

redis/session.go

@@ -7,16 +7,14 @@ import (
 	"io"
 	"net"
 	"strconv"
-	"sync"
 )
 
 // cmd, err := session.ReadCommand()
 // session.WriteReply(reply)
 // session.Write(reply.Bytes())
 type Session struct {
 	net.Conn
-	rd    *bufio.Reader
-	rlock sync.Mutex
+	rd *bufio.Reader
 }
 
 func NewSession(conn net.Conn) *Session {
@@ -27,8 +25,6 @@ func NewSession(conn net.Conn) *Session {
 }
 
 func (s *Session) ReadCommand() (Command, error) {
-	s.rlock.Lock()
-	defer s.rlock.Unlock()
 	// Read ( *<number of arguments> CR LF )
 	if err := s.skipByte('*'); err != nil { // io.EOF
 		return nil, err

rocks/db.go

@@ -37,6 +37,8 @@ func (d *DB) objFromCache(key []byte, e ElementType) interface{} {
 			obj = NewHashElement(d, key)
 		case LIST:
 			obj = NewListElement(d, key)
+		case SORTEDSET:
+			obj = NewSortedSetElement(d, key)
 		}
 		d.caches.Add(skey, obj)
 	}
@@ -51,6 +53,10 @@ func (d *DB) List(key []byte) *ListElement {
 	return d.objFromCache(key, LIST).(*ListElement)
 }
 
+func (d *DB) SortedSet(key []byte) *SortedSetElement {
+	return d.objFromCache(key, SORTEDSET).(*SortedSetElement)
+}
+
 func (d *DB) Delete(key []byte) error {
 	return nil
 }

rocks/list.go

@@ -208,6 +208,6 @@ func (l *ListElement) indexKey(i int64) []byte {
 
 // split l[key]index into index
 func (l *ListElement) indexInKey(key []byte) int64 {
-	idxbuf := bytes.TrimLeft(key, string(l.keyPrefix()))
+	idxbuf := bytes.TrimPrefix(key, l.keyPrefix())
 	return BytesToInt64(idxbuf[1:]) // skip sign "0/1"
 }

rocks/util.go

@@ -20,11 +20,11 @@ var (
 type ElementType byte
 
 const (
-	STRING ElementType = 's'
-	HASH               = 'h'
-	LIST               = 'l'
-	ZSET               = 'z'
-	NONE               = '0'
+	STRING    ElementType = 's'
+	HASH                  = 'h'
+	LIST                  = 'l'
+	SORTEDSET             = 'z'
+	NONE                  = '0'
 )
 
 func (e ElementType) String() string {
@@ -49,8 +49,11 @@ const (
 	IterBackward
 )
 
-// 字节最大范围
-const MAXBYTE byte = math.MaxUint8
+// 字节范围
+const (
+	MINBYTE byte = 0
+	MAXBYTE byte = math.MaxUint8
+)
 
 func rawKey(key []byte, t ElementType) []byte {
 	return bytes.Join([][]byte{KEY, key, SEP, []byte{byte(t)}}, nil)

rocks/zset.go

@@ -0,0 +1,121 @@
+package rocks
+
+import (
+	"bytes"
+	"errors"
+	"github.com/tecbot/gorocksdb"
+	"sync"
+)
+
+// SortedSet
+// +key,z = ""
+// z[key]m member = score
+// z[key]s score member = ""
+type SortedSetElement struct {
+	db  *DB
+	key []byte
+	mu  sync.RWMutex
+}
+
+type SortedSetEnumerateFunc func(i int, score, member []byte, quit *bool)
+
+func NewSortedSetElement(db *DB, key []byte) *SortedSetElement {
+	s := &SortedSetElement{db: db, key: key}
+	return s
+}
+
+// http://redis.io/commands/zadd#return-value
+func (s *SortedSetElement) Add(scoreMembers ...[]byte) (int, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	count := len(scoreMembers)
+	if count < 2 || count%2 != 0 {
+		return 0, errors.New("invalid score/member pairs")
+	}
+
+	batch := gorocksdb.NewWriteBatch()
+	defer batch.Destroy()
+
+	added := 0
+	for i := 0; i < count; i += 2 {
+		score, member := scoreMembers[i], scoreMembers[i+1]
+		skey, mkey := s.scoreKey(score, member), s.memberKey(member)
+
+		// remove old score key
+		oldscore, err := s.db.RawGet(mkey)
+		if err != nil {
+			return 0, err
+		} else if oldscore != nil {
+			batch.Delete(skey)
+		} else {
+			added++
+		}
+
+		// put new value
+		batch.Put(mkey, score)
+		batch.Put(skey, nil)
+	}
+
+	return added, s.db.WriteBatch(batch)
+}
+
+func (s *SortedSetElement) Score(member []byte) ([]byte, error) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	return s.score(member)
+}
+
+func (s *SortedSetElement) score(member []byte) ([]byte, error) {
+	return s.db.RawGet(s.memberKey(member))
+}
+
+func (s *SortedSetElement) Remove(members ...[]byte) (int, error) {
+	return 0, nil
+}
+
+// +key,z = ""
+func (s *SortedSetElement) rawKey() []byte {
+	return rawKey(s.key, SORTEDSET)
+}
+
+// z[key]
+func (s *SortedSetElement) keyPrefix() []byte {
+	return bytes.Join([][]byte{[]byte{SORTEDSET}, SOK, s.key, EOK}, nil)
+}
+
+func (s *SortedSetElement) memberKey(member []byte) []byte {
+	return bytes.Join([][]byte{s.keyPrefix(), []byte{'m'}, member}, nil)
+}
+
+func (s *SortedSetElement) scoreKey(score, member []byte) []byte {
+	return bytes.Join([][]byte{s.keyPrefix(), []byte{'s'}, score, []byte{' '}, member}, nil)
+}
+
+// split (z[key]s score member) into (score, member)
+func (s *SortedSetElement) splitScoreKey(skey []byte) ([]byte, []byte, error) {
+	buf := bytes.TrimPrefix(skey, s.keyPrefix())
+	pairs := bytes.Split(buf[1:], []byte{' '}) // skip score mark 's'
+	if len(pairs) != 2 {
+		return nil, nil, errors.New("invalid score/member key: " + string(skey))
+	}
+	return pairs[0], pairs[1], nil
+}
+
+// split (z[key]m member) into (member)
+func (s *SortedSetElement) splitMemberKey(mkey []byte) ([]byte, error) {
+	buf := bytes.TrimPrefix(mkey, s.keyPrefix())
+	return buf[1:], nil // skip member mark 'm'
+}
+
+func (s *SortedSetElement) RemoveByScore(min, max []byte) (int, error) {
+	return 0, nil
+}
+
+func (s *SortedSetElement) RangeByScore(min, max []byte, fn SortedSetEnumerateFunc) error {
+	return nil
+}
+
+func (s *SortedSetElement) RangeByMember(min, max []byte, fu SortedSetEnumerateFunc) error {
+	return nil
+}