├── .gitignore ├── LICENSE ├── README.md ├── error.go ├── go.mod ├── go.sum ├── protocol ├── reader.go ├── reader_test.go ├── writer.go └── writer_test.go ├── redis.go ├── redis_test.go └── util ├── strconv.go └── unsafe.go /.gitignore: -------------------------------------------------------------------------------- 1 | dump.rdb 2 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 bilibili 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # redis 2 | 3 | 最新的[go-redis](https://github.com/go-redis/redis)已经支持调用方法传入ctx。 4 | 5 | 新项目不要使用这个库了。 6 | 7 | 轻量 redis 客户端 8 | 9 | 支持传入 ctx 对象。纯单机版 SDK，集群相关内容请使用 envoy 等中间件。 10 | -------------------------------------------------------------------------------- /error.go: -------------------------------------------------------------------------------- 1 | package redis 2 | 3 | import ( 4 | "io" 5 | "net" 6 | "strings" 7 | 8 | "github.com/go-kiss/redis/protocol" 9 | ) 10 | 11 | const Nil = protocol.Nil 12 | 13 | func IsRetryableError(err error, retryTimeout bool) bool { 14 | if err == nil { 15 | return false 16 | } 17 | if err == io.EOF { 18 | return true 19 | } 20 | if netErr, ok := err.(net.Error); ok { 21 | if netErr.Timeout() { 22 | return retryTimeout 23 | } 24 | return true 25 | } 26 | s := err.Error() 27 | if s == "ERR max number of clients reached" { 28 | return true 29 | } 30 | if strings.HasPrefix(s, "LOADING ") { 31 | return true 32 | } 33 | if strings.HasPrefix(s, "READONLY ") { 34 | return true 35 | } 36 | if strings.HasPrefix(s, "CLUSTERDOWN ") { 37 | return true 38 | } 39 | return false 40 | } 41 | 42 | func IsRedisError(err error) bool { 43 | _, ok := err.(protocol.RedisError) 44 | return ok 45 | } 46 | 47 | func IsBadConn(err error, allowTimeout bool) bool { 48 | if err == nil { 49 | return false 50 | } 51 | 52 | if IsRedisError(err) { 53 | // #790 54 | return IsReadOnlyError(err) 55 | } 56 | 57 | if allowTimeout { 58 | if netErr, ok := err.(net.Error); ok && netErr.Timeout() { 59 | return false 60 | } 61 | } 62 | 63 | return true 64 | } 65 | 66 | func IsMovedError(err error) (moved bool, ask bool, addr string) { 67 | if !IsRedisError(err) { 68 | return 69 | } 70 | 71 | s := err.Error() 72 | if strings.HasPrefix(s, "MOVED ") { 73 | moved = true 74 | } else if strings.HasPrefix(s, "ASK ") { 75 | ask = true 76 | } else { 77 | return 78 | } 79 | 80 | ind := strings.LastIndex(s, " ") 81 | if ind == -1 { 82 | return false, false, "" 83 | } 84 | addr = s[ind+1:] 85 | return 86 | } 87 | 88 | func IsLoadingError(err error) bool { 89 | return strings.HasPrefix(err.Error(), "LOADING ") 90 | } 91 | 92 | func IsReadOnlyError(err error) bool { 93 | return strings.HasPrefix(err.Error(), "READONLY ") 94 | } 95 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/go-kiss/redis 2 | 3 | go 1.12 4 | 5 | require github.com/go-kiss/net/pool v0.0.0-20210719091328-f4192f07e5b8 6 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/fsnotify/fsnotify v1.4.7 h1:IXs+QLmnXW2CcXuY+8Mzv/fWEsPGWxqefPtCP5CnV9I= 2 | github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo= 3 | github.com/go-kiss/net/pool v0.0.0-20210719091328-f4192f07e5b8 h1:nK1Z1n6dgScorbt/cHcSxlZcx7jHcxaXpukc/BQC37c= 4 | github.com/go-kiss/net/pool v0.0.0-20210719091328-f4192f07e5b8/go.mod h1:PiyXlXMtzPZ+JimCzlIHhSQJ+CvC2eveYBjrh3Uupuw= 5 | github.com/golang/protobuf v1.2.0 h1:P3YflyNX/ehuJFLhxviNdFxQPkGK5cDcApsge1SqnvM= 6 | github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= 7 | github.com/hpcloud/tail v1.0.0 h1:nfCOvKYfkgYP8hkirhJocXT2+zOD8yUNjXaWfTlyFKI= 8 | github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU= 9 | github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= 10 | github.com/onsi/ginkgo v1.8.0 h1:VkHVNpR4iVnU8XQR6DBm8BqYjN7CRzw+xKUbVVbbW9w= 11 | github.com/onsi/ginkgo v1.8.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= 12 | github.com/onsi/gomega v1.5.0 h1:izbySO9zDPmjJ8rDjLvkA2zJHIo+HkYXHnf7eN7SSyo= 13 | github.com/onsi/gomega v1.5.0/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY= 14 | golang.org/x/net v0.0.0-20180906233101-161cd47e91fd h1:nTDtHvHSdCn1m6ITfMRqtOd/9+7a3s8RBNOZ3eYZzJA= 15 | golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= 16 | golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f h1:wMNYb4v58l5UBM7MYRLPG6ZhfOqbKu7X5eyFl8ZhKvA= 17 | golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= 18 | golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e h1:o3PsSEY8E4eXWkXrIP9YJALUkVZqzHJT5DOasTyn8Vs= 19 | golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= 20 | golang.org/x/text v0.3.0 h1:g61tztE5qeGQ89tm6NTjjM9VPIm088od1l6aSorWRWg= 21 | golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= 22 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= 23 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= 24 | gopkg.in/fsnotify.v1 v1.4.7 h1:xOHLXZwVvI9hhs+cLKq5+I5onOuwQLhQwiu63xxlHs4= 25 | gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys= 26 | gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 h1:uRGJdciOHaEIrze2W8Q3AKkepLTh2hOroT7a+7czfdQ= 27 | gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw= 28 | gopkg.in/yaml.v2 v2.2.1 h1:mUhvW9EsL+naU5Q3cakzfE91YhliOondGd6ZrsDBHQE= 29 | gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= 30 | -------------------------------------------------------------------------------- /protocol/reader.go: -------------------------------------------------------------------------------- 1 | package protocol 2 | 3 | import ( 4 | "bufio" 5 | "fmt" 6 | "io" 7 | "strconv" 8 | 9 | "github.com/go-kiss/redis/util" 10 | ) 11 | 12 | const ( 13 | ErrorReply = '-' 14 | StatusReply = '+' 15 | IntReply = ':' 16 | StringReply = '$' 17 | ArrayReply = '*' 18 | ) 19 | 20 | //------------------------------------------------------------------------------ 21 | 22 | const Nil = RedisError("redis: nil") 23 | 24 | type RedisError string 25 | 26 | func (e RedisError) Error() string { return string(e) } 27 | 28 | //------------------------------------------------------------------------------ 29 | 30 | type MultiBulkParse func(*Reader, int64) (interface{}, error) 31 | 32 | type Reader struct { 33 | rd *bufio.Reader 34 | _buf []byte 35 | } 36 | 37 | func NewReader(rd io.Reader) *Reader { 38 | return &Reader{ 39 | rd: bufio.NewReader(rd), 40 | _buf: make([]byte, 64), 41 | } 42 | } 43 | 44 | func (r *Reader) Reset(rd io.Reader) { 45 | r.rd.Reset(rd) 46 | } 47 | 48 | func (r *Reader) readLine() ([]byte, error) { 49 | line, isPrefix, err := r.rd.ReadLine() 50 | if err != nil { 51 | return nil, err 52 | } 53 | if isPrefix { 54 | return nil, bufio.ErrBufferFull 55 | } 56 | if len(line) == 0 { 57 | return nil, fmt.Errorf("redis: reply is empty") 58 | } 59 | if isNilReply(line) { 60 | return nil, Nil 61 | } 62 | return line, nil 63 | } 64 | 65 | func (r *Reader) ReadInterfaceReply() (interface{}, error) { 66 | line, err := r.readLine() 67 | if err != nil { 68 | return nil, err 69 | } 70 | return r.parseInterface(line) 71 | } 72 | 73 | func (r *Reader) parseInterface(line []byte) (interface{}, error) { 74 | switch line[0] { 75 | case ErrorReply: 76 | return nil, ParseErrorReply(line) 77 | case IntReply: 78 | return util.ParseInt(line[1:], 10, 64) 79 | case StatusReply: 80 | return string(line[1:]), nil 81 | case StringReply: 82 | b, err := r.readBytes(line) 83 | if err != nil { 84 | return nil, err 85 | } 86 | return string(b), nil 87 | case ArrayReply: 88 | vs, err := r.readArray(line) 89 | return vs, err 90 | default: 91 | return nil, fmt.Errorf("redis: can't parse int reply: %.100q", line) 92 | } 93 | } 94 | 95 | func (r *Reader) readArray(line []byte) ([]interface{}, error) { 96 | if isNilReply(line) { 97 | return nil, Nil 98 | } 99 | 100 | arrayLen, err := parseArrayLen(line) 101 | if err != nil { 102 | return nil, err 103 | } 104 | 105 | vs := make([]interface{}, 0, arrayLen) 106 | 107 | for i := 0; i < arrayLen; i++ { 108 | line, err := r.readLine() 109 | if err != nil { 110 | return nil, err 111 | } 112 | 113 | v, err := r.parseInterface(line) 114 | if err != nil { 115 | return nil, err 116 | } 117 | 118 | vs = append(vs, v) 119 | } 120 | 121 | return vs, nil 122 | } 123 | 124 | func (r *Reader) ReadIntReply() (int64, error) { 125 | line, err := r.readLine() 126 | if err != nil { 127 | return 0, err 128 | } 129 | switch line[0] { 130 | case ErrorReply: 131 | return 0, ParseErrorReply(line) 132 | case IntReply: 133 | return util.ParseInt(line[1:], 10, 64) 134 | default: 135 | return 0, fmt.Errorf("redis: can't parse int reply: %.100q", line) 136 | } 137 | } 138 | 139 | func (r *Reader) ReadStatusReply() (string, error) { 140 | line, err := r.readLine() 141 | if err != nil { 142 | return "", err 143 | } 144 | switch line[0] { 145 | case ErrorReply: 146 | return "", ParseErrorReply(line) 147 | case StatusReply: 148 | return string(line[1:]), nil 149 | case StringReply: 150 | buf, err := r.ReadBytesReply() 151 | if err != nil { 152 | return "", err 153 | } 154 | 155 | return string(buf), nil 156 | default: 157 | return "", fmt.Errorf("redis: can't parse reply=%.100q reading string", line) 158 | } 159 | } 160 | 161 | func (r *Reader) ReadBytesReply() ([]byte, error) { 162 | line, err := r.readLine() 163 | if err != nil { 164 | return nil, err 165 | } 166 | switch line[0] { 167 | case ErrorReply: 168 | return nil, ParseErrorReply(line) 169 | case StringReply: 170 | return r.readBytes(line) 171 | default: 172 | return nil, fmt.Errorf("redis: can't parse string reply: %.100q", line) 173 | } 174 | } 175 | 176 | func (r *Reader) readBytes(line []byte) ([]byte, error) { 177 | if isNilReply(line) { 178 | return nil, Nil 179 | } 180 | 181 | replyLen, err := strconv.Atoi(string(line[1:])) 182 | if err != nil { 183 | return nil, err 184 | } 185 | 186 | b := make([]byte, replyLen+2) 187 | _, err = io.ReadFull(r.rd, b) 188 | if err != nil { 189 | return nil, err 190 | } 191 | 192 | return b[:replyLen], nil 193 | } 194 | 195 | func (r *Reader) ReadArrayLenReply() (int, error) { 196 | line, err := r.readLine() 197 | if err != nil { 198 | return 0, err 199 | } 200 | switch line[0] { 201 | case ErrorReply: 202 | return 0, ParseErrorReply(line) 203 | case ArrayReply: 204 | return parseArrayLen(line) 205 | default: 206 | return 0, fmt.Errorf("redis: can't parse array reply: %.100q", line) 207 | } 208 | } 209 | 210 | func (r *Reader) ReadInt() (int64, error) { 211 | b, err := r.readTmpBytesReply() 212 | if err != nil { 213 | return 0, err 214 | } 215 | return util.ParseInt(b, 10, 64) 216 | } 217 | 218 | func (r *Reader) ReadUint() (uint64, error) { 219 | b, err := r.readTmpBytesReply() 220 | if err != nil { 221 | return 0, err 222 | } 223 | return util.ParseUint(b, 10, 64) 224 | } 225 | 226 | func (r *Reader) ReadFloat() (float64, error) { 227 | b, err := r.readTmpBytesReply() 228 | if err != nil { 229 | return 0, err 230 | } 231 | return util.ParseFloat(b, 64) 232 | } 233 | 234 | func (r *Reader) readTmpBytesReply() ([]byte, error) { 235 | line, err := r.readLine() 236 | if err != nil { 237 | return nil, err 238 | } 239 | switch line[0] { 240 | case ErrorReply: 241 | return nil, ParseErrorReply(line) 242 | case StringReply: 243 | return r._readTmpBytesReply(line) 244 | case StatusReply: 245 | return line[1:], nil 246 | default: 247 | return nil, fmt.Errorf("redis: can't parse string reply: %.100q", line) 248 | } 249 | } 250 | 251 | func (r *Reader) _readTmpBytesReply(line []byte) ([]byte, error) { 252 | if isNilReply(line) { 253 | return nil, Nil 254 | } 255 | 256 | replyLen, err := strconv.Atoi(string(line[1:])) 257 | if err != nil { 258 | return nil, err 259 | } 260 | 261 | buf := r.buf(replyLen + 2) 262 | _, err = io.ReadFull(r.rd, buf) 263 | if err != nil { 264 | return nil, err 265 | } 266 | 267 | return buf[:replyLen], nil 268 | } 269 | 270 | func (r *Reader) buf(n int) []byte { 271 | if d := n - cap(r._buf); d > 0 { 272 | r._buf = append(r._buf, make([]byte, d)...) 273 | } 274 | return r._buf[:n] 275 | } 276 | 277 | func isNilReply(b []byte) bool { 278 | return len(b) == 3 && 279 | (b[0] == StringReply || b[0] == ArrayReply) && 280 | b[1] == '-' && b[2] == '1' 281 | } 282 | 283 | func ParseErrorReply(line []byte) error { 284 | return RedisError(string(line[1:])) 285 | } 286 | 287 | func parseArrayLen(line []byte) (int, error) { 288 | if isNilReply(line) { 289 | return 0, Nil 290 | } 291 | return util.Atoi(line[1:]) 292 | } 293 | -------------------------------------------------------------------------------- /protocol/reader_test.go: -------------------------------------------------------------------------------- 1 | package protocol 2 | 3 | import ( 4 | "bytes" 5 | "testing" 6 | ) 7 | 8 | func BenchmarkReaderParseReplyStatus(b *testing.B) { 9 | buf := new(bytes.Buffer) 10 | for i := 0; i < b.N; i++ { 11 | buf.WriteString("+OK\r\n") 12 | } 13 | p := NewReader(buf) 14 | b.ResetTimer() 15 | 16 | for i := 0; i < b.N; i++ { 17 | if m, err := p.ReadStatusReply(); err != nil || m != "OK" { 18 | b.Fatal(err) 19 | } 20 | } 21 | } 22 | 23 | func BenchmarkReaderParseReplyInt(b *testing.B) { 24 | buf := new(bytes.Buffer) 25 | for i := 0; i < b.N; i++ { 26 | buf.WriteString(":1\r\n") 27 | } 28 | p := NewReader(buf) 29 | b.ResetTimer() 30 | 31 | for i := 0; i < b.N; i++ { 32 | if i, err := p.ReadIntReply(); err != nil || i != 1 { 33 | b.Fatal(err) 34 | } 35 | } 36 | } 37 | 38 | func BenchmarkReaderParseReplyError(b *testing.B) { 39 | buf := new(bytes.Buffer) 40 | for i := 0; i < b.N; i++ { 41 | buf.WriteString("-Error message\r\n") 42 | } 43 | p := NewReader(buf) 44 | b.ResetTimer() 45 | 46 | for i := 0; i < b.N; i++ { 47 | if _, err := p.ReadInt(); err.Error() != "Error message" { 48 | b.Fatal(err) 49 | } 50 | } 51 | } 52 | 53 | func BenchmarkReaderParseReplyString(b *testing.B) { 54 | buf := new(bytes.Buffer) 55 | for i := 0; i < b.N; i++ { 56 | buf.WriteString("$5\r\nhello\r\n") 57 | } 58 | p := NewReader(buf) 59 | b.ResetTimer() 60 | 61 | for i := 0; i < b.N; i++ { 62 | if buf, err := p.ReadBytesReply(); err != nil || !bytes.Equal(buf, []byte("hello")) { 63 | b.Fatal(err) 64 | } 65 | } 66 | } 67 | 68 | func BenchmarkReaderParseReplyArray(b *testing.B) { 69 | buf := new(bytes.Buffer) 70 | for i := 0; i < b.N; i++ { 71 | buf.WriteString("*3\r\n$5\r\nhello\r\n$3\r\nfoo\r\n$3\r\nbar\r\n") 72 | } 73 | p := NewReader(buf) 74 | b.ResetTimer() 75 | 76 | bufs := [][]byte{[]byte("hello"), []byte("foo"), []byte("bar")} 77 | 78 | for i := 0; i < b.N; i++ { 79 | l, err := p.ReadArrayLenReply() 80 | if err != nil || l != 3 { 81 | b.Fatal(err) 82 | } 83 | 84 | for j := 0; j < l; j++ { 85 | if buf, err := p.ReadBytesReply(); err != nil || !bytes.Equal(buf, bufs[j]) { 86 | b.Fatal(err) 87 | } 88 | } 89 | 90 | } 91 | } 92 | -------------------------------------------------------------------------------- /protocol/writer.go: -------------------------------------------------------------------------------- 1 | package protocol 2 | 3 | import ( 4 | "bufio" 5 | "fmt" 6 | "io" 7 | "strconv" 8 | 9 | "github.com/go-kiss/redis/util" 10 | ) 11 | 12 | type Writer struct { 13 | wr *bufio.Writer 14 | 15 | lenBuf []byte 16 | numBuf []byte 17 | } 18 | 19 | func NewWriter(wr io.Writer) *Writer { 20 | return &Writer{ 21 | wr: bufio.NewWriter(wr), 22 | 23 | lenBuf: make([]byte, 64), 24 | numBuf: make([]byte, 64), 25 | } 26 | } 27 | 28 | func (w *Writer) WriteArgs(args []interface{}) error { 29 | err := w.wr.WriteByte(ArrayReply) 30 | if err != nil { 31 | return err 32 | } 33 | 34 | err = w.writeLen(len(args)) 35 | if err != nil { 36 | return err 37 | } 38 | 39 | for _, arg := range args { 40 | err := w.writeArg(arg) 41 | if err != nil { 42 | return err 43 | } 44 | } 45 | 46 | return nil 47 | } 48 | 49 | func (w *Writer) writeLen(n int) error { 50 | w.lenBuf = strconv.AppendUint(w.lenBuf[:0], uint64(n), 10) 51 | w.lenBuf = append(w.lenBuf, '\r', '\n') 52 | _, err := w.wr.Write(w.lenBuf) 53 | return err 54 | } 55 | 56 | func (w *Writer) writeArg(v interface{}) error { 57 | switch v := v.(type) { 58 | case nil: 59 | return w.string("") 60 | case string: 61 | return w.string(v) 62 | case []byte: 63 | return w.bytes(v) 64 | case int: 65 | return w.int(int64(v)) 66 | case int8: 67 | return w.int(int64(v)) 68 | case int16: 69 | return w.int(int64(v)) 70 | case int32: 71 | return w.int(int64(v)) 72 | case int64: 73 | return w.int(v) 74 | case uint: 75 | return w.uint(uint64(v)) 76 | case uint8: 77 | return w.uint(uint64(v)) 78 | case uint16: 79 | return w.uint(uint64(v)) 80 | case uint32: 81 | return w.uint(uint64(v)) 82 | case uint64: 83 | return w.uint(v) 84 | case float32: 85 | return w.float(float64(v)) 86 | case float64: 87 | return w.float(v) 88 | case bool: 89 | vv := int64(0) 90 | if v { 91 | vv = 1 92 | } 93 | return w.int(vv) 94 | default: 95 | return fmt.Errorf( 96 | "redis: can't marshal %T (implement encoding.BinaryMarshaler)", v) 97 | } 98 | } 99 | 100 | func (w *Writer) bytes(b []byte) error { 101 | err := w.wr.WriteByte(StringReply) 102 | if err != nil { 103 | return err 104 | } 105 | 106 | err = w.writeLen(len(b)) 107 | if err != nil { 108 | return err 109 | } 110 | 111 | _, err = w.wr.Write(b) 112 | if err != nil { 113 | return err 114 | } 115 | 116 | return w.crlf() 117 | } 118 | 119 | func (w *Writer) string(s string) error { 120 | return w.bytes(util.StringToBytes(s)) 121 | } 122 | 123 | func (w *Writer) uint(n uint64) error { 124 | w.numBuf = strconv.AppendUint(w.numBuf[:0], n, 10) 125 | return w.bytes(w.numBuf) 126 | } 127 | 128 | func (w *Writer) int(n int64) error { 129 | w.numBuf = strconv.AppendInt(w.numBuf[:0], n, 10) 130 | return w.bytes(w.numBuf) 131 | } 132 | 133 | func (w *Writer) float(f float64) error { 134 | w.numBuf = strconv.AppendFloat(w.numBuf[:0], f, 'f', -1, 64) 135 | return w.bytes(w.numBuf) 136 | } 137 | 138 | func (w *Writer) crlf() error { 139 | err := w.wr.WriteByte('\r') 140 | if err != nil { 141 | return err 142 | } 143 | return w.wr.WriteByte('\n') 144 | } 145 | 146 | func (w *Writer) Reset(wr io.Writer) { 147 | w.wr.Reset(wr) 148 | } 149 | 150 | func (w *Writer) Flush() error { 151 | return w.wr.Flush() 152 | } 153 | -------------------------------------------------------------------------------- /protocol/writer_test.go: -------------------------------------------------------------------------------- 1 | package protocol 2 | 3 | import ( 4 | "bytes" 5 | "io/ioutil" 6 | "testing" 7 | ) 8 | 9 | func TestWriteArgs(t *testing.T) { 10 | buf := new(bytes.Buffer) 11 | wr := NewWriter(buf) 12 | 13 | err := wr.WriteArgs([]interface{}{ 14 | "string", 15 | 12, 16 | 34.56, 17 | []byte{'b', 'y', 't', 'e', 's'}, 18 | true, 19 | nil, 20 | }) 21 | 22 | expected := []byte("*6\r\n" + 23 | "$6\r\nstring\r\n" + 24 | "$2\r\n12\r\n" + 25 | "$5\r\n34.56\r\n" + 26 | "$5\r\nbytes\r\n" + 27 | "$1\r\n1\r\n" + 28 | "$0\r\n" + 29 | "\r\n") 30 | 31 | if err != nil || bytes.Equal(buf.Bytes(), expected) { 32 | t.Fatal("WriteArgs faild") 33 | } 34 | } 35 | 36 | func BenchmarkWriteBuffer_Append(b *testing.B) { 37 | buf := NewWriter(ioutil.Discard) 38 | 39 | for i := 0; i < b.N; i++ { 40 | err := buf.WriteArgs("hello", "world", "foo", "bar") 41 | if err != nil { 42 | panic(err) 43 | } 44 | 45 | err = buf.Flush() 46 | if err != nil { 47 | panic(err) 48 | } 49 | } 50 | } 51 | -------------------------------------------------------------------------------- /redis.go: -------------------------------------------------------------------------------- 1 | package redis 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "net" 7 | "time" 8 | 9 | "github.com/go-kiss/net/pool" 10 | "github.com/go-kiss/redis/protocol" 11 | ) 12 | 13 | const ( 14 | FlagNX = 1 << 0 15 | FlagXX = 1 << 1 16 | FlagCH = 1 << 2 17 | ) 18 | 19 | type EvalReturn struct { 20 | val interface{} 21 | } 22 | 23 | func (e EvalReturn) Int64() (int64, error) { 24 | v, ok := e.val.(int64) 25 | if !ok { 26 | return 0, fmt.Errorf("redis: unexpected type=%T for int64", e.val) 27 | } 28 | return v, nil 29 | } 30 | 31 | func (e EvalReturn) String() (string, error) { 32 | v, ok := e.val.(string) 33 | if !ok { 34 | return "", fmt.Errorf("redis: unexpected type=%T for string", e.val) 35 | } 36 | return v, nil 37 | } 38 | 39 | func (e EvalReturn) Array() ([]interface{}, error) { 40 | v, ok := e.val.([]interface{}) 41 | if !ok { 42 | return nil, fmt.Errorf("redis: unexpected type=%T for Array", e.val) 43 | } 44 | return v, nil 45 | } 46 | 47 | func (e EvalReturn) Interface() interface{} { 48 | return e.val 49 | } 50 | 51 | type Options struct { 52 | Address string 53 | PoolSize int 54 | MinIdleConns int 55 | 56 | MaxConnAge time.Duration 57 | PoolTimeout time.Duration 58 | IdleTimeout time.Duration 59 | 60 | IdleCheckFrequency time.Duration 61 | 62 | OnPreCmd func(context.Context, []interface{}) context.Context 63 | OnPostCmd func(context.Context, error) 64 | } 65 | 66 | type Client struct { 67 | opts Options 68 | pool pool.Pooler 69 | } 70 | 71 | func New(opts Options) Client { 72 | poolOpts := pool.Options{ 73 | PoolSize: opts.PoolSize, 74 | MinIdleConns: opts.MinIdleConns, 75 | MaxConnAge: opts.MaxConnAge, 76 | PoolTimeout: opts.PoolTimeout, 77 | IdleTimeout: opts.IdleTimeout, 78 | IdleCheckFrequency: opts.IdleCheckFrequency, 79 | } 80 | 81 | poolOpts.Dialer = func(ctx context.Context) (pool.Closer, error) { 82 | d := net.Dialer{} 83 | conn, err := d.DialContext(ctx, "tcp", opts.Address) 84 | if err != nil { 85 | return nil, err 86 | } 87 | 88 | rw := redisConn{ 89 | c: conn, 90 | r: protocol.NewReader(conn), 91 | w: protocol.NewWriter(conn), 92 | } 93 | 94 | return &rw, nil 95 | } 96 | 97 | return Client{pool: pool.New(poolOpts), opts: opts} 98 | } 99 | 100 | type redisConn struct { 101 | c net.Conn 102 | r *protocol.Reader 103 | w *protocol.Writer 104 | } 105 | 106 | func (rc *redisConn) Close() error { 107 | return rc.c.Close() 108 | } 109 | 110 | type ZSetValue struct { 111 | Member string 112 | Score float64 113 | } 114 | 115 | type Item struct { 116 | // Key is the Item's key (250 bytes maximum). 117 | Key string 118 | 119 | // Value is the Item's value. 120 | Value []byte 121 | 122 | ZSetValues map[string]float64 123 | 124 | // Flags 一些 redis 标记位，请参考 Flag 开头的常量定义 125 | Flags uint32 126 | 127 | // TTL 缓存时间，秒，0 表示不过期 128 | TTL int32 129 | } 130 | 131 | var noDeadline = time.Time{} 132 | 133 | // PoolStats 返回连接池状态 134 | func (c *Client) PoolStats() *pool.Stats { 135 | return c.pool.Stats() 136 | } 137 | 138 | func (c *Client) do(ctx context.Context, args []interface{}, fn func(conn *redisConn) error) error { 139 | if c.opts.OnPreCmd != nil { 140 | ctx = c.opts.OnPreCmd(ctx, args) 141 | } 142 | 143 | conn, err := c.pool.Get(ctx) 144 | if err != nil { 145 | return err 146 | } 147 | 148 | defer func() { 149 | if IsBadConn(err, false) { 150 | c.pool.Remove(conn) 151 | } else { 152 | c.pool.Put(conn) 153 | } 154 | 155 | }() 156 | 157 | rc := conn.C.(*redisConn) 158 | 159 | if t, ok := ctx.Deadline(); ok { 160 | err = rc.c.SetDeadline(t) 161 | } else { 162 | err = rc.c.SetDeadline(noDeadline) 163 | } 164 | 165 | if err != nil { 166 | return err 167 | } 168 | 169 | // 此处赋值给 defer 函数用的，不要去掉 170 | err = fn(rc) 171 | 172 | if c.opts.OnPostCmd != nil { 173 | c.opts.OnPostCmd(ctx, err) 174 | } 175 | 176 | return err 177 | } 178 | 179 | func (c *Client) Get(ctx context.Context, key string) (item *Item, err error) { 180 | args := []interface{}{"get", key} 181 | err = c.do(ctx, args, func(conn *redisConn) error { 182 | if err := conn.w.WriteArgs(args); err != nil { 183 | return err 184 | } 185 | 186 | if err := conn.w.Flush(); err != nil { 187 | return err 188 | } 189 | 190 | var b []byte 191 | if b, err = conn.r.ReadBytesReply(); err != nil { 192 | item = nil 193 | return err 194 | } 195 | 196 | item = &Item{Value: b} 197 | 198 | return nil 199 | }) 200 | return 201 | } 202 | 203 | func (c *Client) MGet(ctx context.Context, keys []string) (items map[string]*Item, err error) { 204 | args := make([]interface{}, 0, len(keys)+1) 205 | 206 | args = append(args, "mget") 207 | for _, key := range keys { 208 | args = append(args, key) 209 | } 210 | 211 | err = c.do(ctx, args, func(conn *redisConn) error { 212 | if err := conn.w.WriteArgs(args); err != nil { 213 | return err 214 | } 215 | 216 | if err := conn.w.Flush(); err != nil { 217 | return err 218 | } 219 | 220 | l, err := conn.r.ReadArrayLenReply() 221 | if err != nil { 222 | return err 223 | } 224 | 225 | items = make(map[string]*Item, l) 226 | 227 | for i := 0; i < l; i++ { 228 | b, err := conn.r.ReadBytesReply() 229 | if err == Nil { 230 | continue 231 | } 232 | if err != nil { 233 | return err 234 | } 235 | 236 | key := keys[i] 237 | 238 | items[key] = &Item{Value: b} 239 | } 240 | 241 | return nil 242 | }) 243 | return 244 | } 245 | 246 | func (c *Client) Eval(ctx context.Context, script string, keys []string, argvs ...interface{}) (result *EvalReturn, err error) { 247 | args := make([]interface{}, 0, len(keys)+len(argvs)+2) 248 | args = append(args, "eval", script, len(keys)) 249 | for _, v := range keys { 250 | args = append(args, v) 251 | } 252 | args = append(args, argvs...) 253 | 254 | err = c.do(ctx, args, func(conn *redisConn) error { 255 | if err := conn.w.WriteArgs(args); err != nil { 256 | return err 257 | } 258 | 259 | if err := conn.w.Flush(); err != nil { 260 | return err 261 | } 262 | 263 | b, err := conn.r.ReadInterfaceReply() 264 | if err != nil { 265 | return err 266 | } 267 | 268 | result = &EvalReturn{ 269 | val: b, 270 | } 271 | 272 | return nil 273 | }) 274 | 275 | return 276 | } 277 | 278 | func (c *Client) Set(ctx context.Context, item *Item) error { 279 | args := make([]interface{}, 0, 6) 280 | args = append(args, "set", item.Key, item.Value) 281 | 282 | if item.TTL > 0 { 283 | args = append(args, "EX", item.TTL) 284 | } 285 | 286 | if item.Flags&FlagNX > 0 { 287 | args = append(args, "NX") 288 | } else if item.Flags&FlagXX > 0 { 289 | args = append(args, "XX") 290 | } 291 | 292 | return c.do(ctx, args, func(conn *redisConn) error { 293 | if err := conn.w.WriteArgs(args); err != nil { 294 | return err 295 | } 296 | 297 | if err := conn.w.Flush(); err != nil { 298 | return err 299 | } 300 | 301 | _, err := conn.r.ReadStatusReply() 302 | if err != nil { 303 | return err 304 | } 305 | 306 | return nil 307 | }) 308 | } 309 | 310 | func (c *Client) Del(ctx context.Context, keys ...string) error { 311 | args := make([]interface{}, 0, 1+len(keys)) 312 | 313 | args = append(args, "del") 314 | for _, key := range keys { 315 | args = append(args, key) 316 | } 317 | 318 | return c.do(ctx, args, func(conn *redisConn) error { 319 | if err := conn.w.WriteArgs(args); err != nil { 320 | return err 321 | } 322 | 323 | if err := conn.w.Flush(); err != nil { 324 | return err 325 | } 326 | 327 | _, err := conn.r.ReadIntReply() 328 | 329 | return err 330 | }) 331 | } 332 | 333 | func (c *Client) IncrBy(ctx context.Context, key string, by int64) (i int64, err error) { 334 | args := []interface{}{"incrby", key, by} 335 | 336 | err = c.do(ctx, args, func(conn *redisConn) error { 337 | 338 | if err = conn.w.WriteArgs(args); err != nil { 339 | return err 340 | } 341 | 342 | if err := conn.w.Flush(); err != nil { 343 | return err 344 | } 345 | 346 | i, err = conn.r.ReadIntReply() 347 | 348 | return err 349 | }) 350 | 351 | return 352 | } 353 | 354 | func (c *Client) DecrBy(ctx context.Context, key string, by int64) (int64, error) { 355 | return c.IncrBy(ctx, key, -by) 356 | } 357 | 358 | func (c *Client) Expire(ctx context.Context, key string, ttl int32) error { 359 | args := []interface{}{"expire", key, ttl} 360 | 361 | return c.do(ctx, args, func(conn *redisConn) error { 362 | if err := conn.w.WriteArgs(args); err != nil { 363 | return err 364 | } 365 | 366 | if err := conn.w.Flush(); err != nil { 367 | return err 368 | } 369 | 370 | _, err := conn.r.ReadIntReply() 371 | 372 | return err 373 | }) 374 | } 375 | 376 | func (c *Client) TTL(ctx context.Context, key string) (ttl int32, err error) { 377 | args := []interface{}{"ttl", key} 378 | 379 | err = c.do(ctx, args, func(conn *redisConn) error { 380 | if err := conn.w.WriteArgs(args); err != nil { 381 | return err 382 | } 383 | 384 | if err := conn.w.Flush(); err != nil { 385 | return err 386 | } 387 | 388 | var i int64 389 | i, err = conn.r.ReadIntReply() 390 | if err != nil { 391 | return err 392 | } 393 | 394 | if i == -2 { 395 | err = Nil 396 | return err 397 | } 398 | 399 | ttl = int32(i) 400 | 401 | return err 402 | }) 403 | 404 | return 405 | } 406 | 407 | func (c *Client) ZAdd(ctx context.Context, item *Item) (added int64, err error) { 408 | args := make([]interface{}, 0, 4+len(item.ZSetValues)) 409 | args = append(args, "zadd", item.Key) 410 | 411 | if item.Flags&FlagNX > 0 { 412 | args = append(args, "NX") 413 | } else if item.Flags&FlagXX > 0 { 414 | args = append(args, "XX") 415 | } 416 | 417 | if item.Flags&FlagCH > 0 { 418 | args = append(args, "CH") 419 | } 420 | 421 | for member, score := range item.ZSetValues { 422 | args = append(args, score, member) 423 | } 424 | 425 | err = c.do(ctx, args, func(conn *redisConn) error { 426 | if err := conn.w.WriteArgs(args); err != nil { 427 | return err 428 | } 429 | 430 | if err := conn.w.Flush(); err != nil { 431 | return err 432 | } 433 | 434 | added, err = conn.r.ReadIntReply() 435 | if err != nil { 436 | return err 437 | } 438 | 439 | return nil 440 | }) 441 | 442 | return 443 | } 444 | 445 | func (c *Client) ZIncrBy(ctx context.Context, key, member string, by float64) error { 446 | args := []interface{}{"zincrby", key, by, member} 447 | 448 | return c.do(ctx, args, func(conn *redisConn) error { 449 | if err := conn.w.WriteArgs(args); err != nil { 450 | return err 451 | } 452 | 453 | if err := conn.w.Flush(); err != nil { 454 | return err 455 | } 456 | 457 | _, err := conn.r.ReadFloat() 458 | return err 459 | }) 460 | } 461 | 462 | func (c *Client) ZRange(ctx context.Context, key string, start, stop int64) (values []*ZSetValue, err error) { 463 | return c.zrange(ctx, "zrange", key, start, stop, 0, 0) 464 | } 465 | 466 | func (c *Client) ZRevRange(ctx context.Context, key string, start, stop int64) (values []*ZSetValue, err error) { 467 | return c.zrange(ctx, "zrevrange", key, start, stop, 0, 0) 468 | } 469 | 470 | func (c *Client) ZRangeByScore(ctx context.Context, key string, min, max float64, offset, count int64) (values []*ZSetValue, err error) { 471 | return c.zrange(ctx, "zrangebyscore", key, min, max, offset, count) 472 | } 473 | 474 | func (c *Client) ZRevRangeByScore(ctx context.Context, key string, max, min float64, offset, count int64) (values []*ZSetValue, err error) { 475 | return c.zrange(ctx, "zrevrangebyscore", key, max, min, offset, count) 476 | } 477 | 478 | func (c *Client) zrange(ctx context.Context, cmd, key string, start, stop interface{}, offset, count int64) (values []*ZSetValue, err error) { 479 | args := []interface{}{cmd, key, start, stop, "WITHSCORES"} 480 | if count > 0 { 481 | args = append(args, "LIMIT", offset, count) 482 | } 483 | 484 | err = c.do(ctx, args, func(conn *redisConn) error { 485 | if err := conn.w.WriteArgs(args); err != nil { 486 | return err 487 | } 488 | 489 | if err := conn.w.Flush(); err != nil { 490 | return err 491 | } 492 | 493 | l, err := conn.r.ReadArrayLenReply() 494 | if err != nil { 495 | return err 496 | } 497 | 498 | values = make([]*ZSetValue, 0, l) 499 | for i := 0; i < l/2; i++ { 500 | b, err := conn.r.ReadBytesReply() 501 | if err != nil { 502 | return err 503 | } 504 | f, err := conn.r.ReadFloat() 505 | if err != nil { 506 | return err 507 | } 508 | 509 | values = append(values, &ZSetValue{Member: string(b), Score: f}) 510 | } 511 | 512 | return nil 513 | }) 514 | 515 | return 516 | } 517 | 518 | func (c *Client) ZRank(ctx context.Context, key, member string) (rank int64, err error) { 519 | return c.zrank(ctx, "zrank", key, member) 520 | } 521 | func (c *Client) ZRevRank(ctx context.Context, key, member string) (rank int64, err error) { 522 | return c.zrank(ctx, "zrevrank", key, member) 523 | } 524 | 525 | func (c *Client) zrank(ctx context.Context, cmd, key, member string) (rank int64, err error) { 526 | args := []interface{}{cmd, key, member} 527 | 528 | err = c.do(ctx, args, func(conn *redisConn) error { 529 | if err := conn.w.WriteArgs(args); err != nil { 530 | return err 531 | } 532 | 533 | if err := conn.w.Flush(); err != nil { 534 | return err 535 | } 536 | 537 | rank, err = conn.r.ReadIntReply() 538 | return err 539 | }) 540 | 541 | return 542 | } 543 | 544 | func (c *Client) ZScore(ctx context.Context, key, member string) (score float64, err error) { 545 | args := []interface{}{"zscore", key, member} 546 | 547 | err = c.do(ctx, args, func(conn *redisConn) error { 548 | if err := conn.w.WriteArgs(args); err != nil { 549 | return err 550 | } 551 | 552 | if err := conn.w.Flush(); err != nil { 553 | return err 554 | } 555 | 556 | score, err = conn.r.ReadFloat() 557 | return err 558 | }) 559 | 560 | return 561 | } 562 | 563 | func (c *Client) ZCard(ctx context.Context, key string) (card int64, err error) { 564 | args := []interface{}{"zcard", key} 565 | 566 | err = c.do(ctx, args, func(conn *redisConn) error { 567 | if err := conn.w.WriteArgs(args); err != nil { 568 | return err 569 | } 570 | 571 | if err := conn.w.Flush(); err != nil { 572 | return err 573 | } 574 | 575 | card, err = conn.r.ReadIntReply() 576 | 577 | return err 578 | }) 579 | return 580 | } 581 | 582 | func (c *Client) ZCount(ctx context.Context, key, min, max string) (i int64, err error) { 583 | args := []interface{}{"zcount", key, min, max} 584 | 585 | err = c.do(ctx, args, func(conn *redisConn) error { 586 | if err := conn.w.WriteArgs(args); err != nil { 587 | return err 588 | } 589 | 590 | if err := conn.w.Flush(); err != nil { 591 | return err 592 | } 593 | 594 | i, err = conn.r.ReadIntReply() 595 | 596 | return err 597 | }) 598 | return 599 | } 600 | 601 | func (c *Client) ZRem(ctx context.Context, keys ...string) error { 602 | args := make([]interface{}, 0, 1+len(keys)) 603 | 604 | args = append(args, "zrem") 605 | for _, key := range keys { 606 | args = append(args, key) 607 | } 608 | 609 | return c.do(ctx, args, func(conn *redisConn) error { 610 | if err := conn.w.WriteArgs(args); err != nil { 611 | return err 612 | } 613 | 614 | if err := conn.w.Flush(); err != nil { 615 | return err 616 | } 617 | 618 | _, err := conn.r.ReadIntReply() 619 | 620 | return err 621 | }) 622 | } 623 | 624 | func (c *Client) ZRemRangeByRank(ctx context.Context, key string, start, stop int64) (i int64, err error) { 625 | args := []interface{}{"zremrangebyrank", key, start, stop} 626 | 627 | err = c.do(ctx, args, func(conn *redisConn) error { 628 | if err := conn.w.WriteArgs(args); err != nil { 629 | return err 630 | } 631 | 632 | if err := conn.w.Flush(); err != nil { 633 | return err 634 | } 635 | 636 | i, err = conn.r.ReadIntReply() 637 | 638 | return err 639 | }) 640 | return 641 | } 642 | 643 | func (c *Client) ZRemRangeByScore(ctx context.Context, key, min, max string) (i int64, err error) { 644 | args := []interface{}{"zremrangebyscore", key, min, max} 645 | 646 | err = c.do(ctx, args, func(conn *redisConn) error { 647 | if err := conn.w.WriteArgs(args); err != nil { 648 | return err 649 | } 650 | 651 | if err := conn.w.Flush(); err != nil { 652 | return err 653 | } 654 | 655 | i, err = conn.r.ReadIntReply() 656 | 657 | return err 658 | }) 659 | return 660 | } 661 | 662 | func (c *Client) Stats() *pool.Stats { 663 | return c.pool.Stats() 664 | } 665 | 666 | func (c *Client) SAdd(ctx context.Context, key string, data ...[]byte) (err error) { 667 | args := []interface{}{"sadd", key} 668 | for _, d := range data { 669 | args = append(args, d) 670 | } 671 | 672 | err = c.do(ctx, args, func(conn *redisConn) error { 673 | if err := conn.w.WriteArgs(args); err != nil { 674 | return err 675 | } 676 | 677 | if err := conn.w.Flush(); err != nil { 678 | return err 679 | } 680 | 681 | _, err = conn.r.ReadIntReply() 682 | 683 | return err 684 | }) 685 | return 686 | } 687 | 688 | func (c *Client) SPop(ctx context.Context, key string, cnt int32) (data [][]byte, err error) { 689 | args := []interface{}{"spop", key, cnt} 690 | 691 | err = c.do(ctx, args, func(conn *redisConn) error { 692 | if err := conn.w.WriteArgs(args); err != nil { 693 | return err 694 | } 695 | 696 | if err := conn.w.Flush(); err != nil { 697 | return err 698 | } 699 | 700 | l, err := conn.r.ReadArrayLenReply() 701 | if err != nil { 702 | return err 703 | } 704 | 705 | data = make([][]byte, l) 706 | for i := 0; i < l; i++ { 707 | b, err := conn.r.ReadBytesReply() 708 | if err != nil { 709 | return err 710 | } 711 | data[i] = b 712 | } 713 | 714 | return nil 715 | }) 716 | return 717 | } 718 | 719 | func (c *Client) SCard(ctx context.Context, key string) (card int64, err error) { 720 | args := []interface{}{"scard", key} 721 | 722 | err = c.do(ctx, args, func(conn *redisConn) error { 723 | if err := conn.w.WriteArgs(args); err != nil { 724 | return err 725 | } 726 | 727 | if err := conn.w.Flush(); err != nil { 728 | return err 729 | } 730 | 731 | card, err = conn.r.ReadIntReply() 732 | 733 | return err 734 | }) 735 | return 736 | } 737 | 738 | func (c *Client) SIsMember(ctx context.Context, key string, data []byte) (result bool, err error) { 739 | args := []interface{}{"sismember", key, data} 740 | 741 | var resultInt int64 742 | err = c.do(ctx, args, func(conn *redisConn) error { 743 | if err := conn.w.WriteArgs(args); err != nil { 744 | return err 745 | } 746 | 747 | if err := conn.w.Flush(); err != nil { 748 | return err 749 | } 750 | 751 | resultInt, err = conn.r.ReadIntReply() 752 | 753 | return err 754 | }) 755 | if resultInt == 1 { 756 | result = true 757 | } 758 | return 759 | } 760 | 761 | func (c *Client) SRem(ctx context.Context, key string, data ...[]byte) (result int64, err error) { 762 | args := []interface{}{"srem", key} 763 | for _, d := range data { 764 | args = append(args, d) 765 | } 766 | 767 | err = c.do(ctx, args, func(conn *redisConn) error { 768 | if err := conn.w.WriteArgs(args); err != nil { 769 | return err 770 | } 771 | 772 | if err := conn.w.Flush(); err != nil { 773 | return err 774 | } 775 | 776 | result, err = conn.r.ReadIntReply() 777 | 778 | return err 779 | }) 780 | return 781 | } 782 | 783 | func (c *Client) SMembers(ctx context.Context, key string) (items [][]byte, err error) { 784 | args := []interface{}{"smembers", key} 785 | 786 | err = c.do(ctx, args, func(conn *redisConn) error { 787 | if err := conn.w.WriteArgs(args); err != nil { 788 | return err 789 | } 790 | 791 | if err := conn.w.Flush(); err != nil { 792 | return err 793 | } 794 | 795 | l, err := conn.r.ReadArrayLenReply() 796 | if err != nil { 797 | return err 798 | } 799 | 800 | items = make([][]byte, l) 801 | for i := 0; i < l; i++ { 802 | b, err := conn.r.ReadBytesReply() 803 | if err != nil { 804 | return err 805 | } 806 | items[i] = b 807 | } 808 | 809 | return nil 810 | }) 811 | return 812 | } 813 | -------------------------------------------------------------------------------- /redis_test.go: -------------------------------------------------------------------------------- 1 | package redis 2 | 3 | import ( 4 | "bytes" 5 | "context" 6 | "os" 7 | "testing" 8 | "time" 9 | ) 10 | 11 | var ctx = context.Background() 12 | 13 | func TestBasic(t *testing.T) { 14 | c := New(Options{ 15 | Address: os.Getenv("REDIS_HOST"), 16 | PoolSize: 1, 17 | }) 18 | 19 | if err := c.Del(ctx, "foo"); err != nil { 20 | t.Fatal("start faild") 21 | } 22 | 23 | c.Set(ctx, &Item{Key: "foo", Value: []byte("hello")}) 24 | 25 | item, _ := c.Get(ctx, "foo") 26 | if !bytes.Equal(item.Value, []byte("hello")) { 27 | t.Fatal("get foo failed") 28 | } 29 | 30 | // set with expire 31 | c.Set(ctx, &Item{Key: "set_with_expire", Value: []byte("test"), TTL: 86400}) 32 | item, _ = c.Get(ctx, "set_with_expire") 33 | if !bytes.Equal(item.Value, []byte("test")) { 34 | t.Fatal("get foo failed") 35 | } 36 | // 获取 TTL 37 | if ttl, err := c.TTL(ctx, "set_with_expire"); err != nil { 38 | t.Fatal("ttl failed") 39 | } else if ttl <= 0 { 40 | t.Fatal("set with expire failed") 41 | } 42 | 43 | // add 44 | if err := c.Set(ctx, &Item{Key: "foo", Value: []byte("123"), Flags: FlagNX}); err != Nil { 45 | t.Fatal("add foo failed") 46 | } 47 | 48 | // replace 49 | if err := c.Set(ctx, &Item{Key: "foo", Value: []byte("123"), Flags: FlagXX}); err == Nil { 50 | t.Fatal("replace foo failed") 51 | } 52 | 53 | item, _ = c.Get(ctx, "foo") 54 | if !bytes.Equal(item.Value, []byte("123")) { 55 | t.Fatal("replace foo failed") 56 | } 57 | 58 | c.Del(ctx, "foo") 59 | if _, err := c.Get(ctx, "foo"); err != Nil { 60 | t.Fatal("del foo failed") 61 | } 62 | 63 | if _, err := c.TTL(ctx, "foo"); err != Nil { 64 | t.Fatal("ttl foo failed") 65 | } 66 | 67 | c.IncrBy(ctx, "foo", 3) 68 | item, _ = c.Get(ctx, "foo") 69 | if !bytes.Equal(item.Value, []byte("3")) { 70 | t.Fatal("increment foo failed") 71 | } 72 | 73 | c.DecrBy(ctx, "foo", 4) 74 | item, _ = c.Get(ctx, "foo") 75 | if !bytes.Equal(item.Value, []byte("-1")) { 76 | t.Fatal("decrement foo failed") 77 | } 78 | 79 | if ttl, _ := c.TTL(ctx, "foo"); ttl != -1 { 80 | t.Fatal("ttl foo failed") 81 | } 82 | 83 | c.Expire(ctx, "foo", 10) 84 | if ttl, _ := c.TTL(ctx, "foo"); ttl != 10 { 85 | t.Fatal("ttl foo failed") 86 | } 87 | time.Sleep(1 * time.Second) 88 | if ttl, _ := c.TTL(ctx, "foo"); ttl != 9 { 89 | t.Fatal("ttl foo failed") 90 | } 91 | } 92 | 93 | func TestZSet(t *testing.T) { 94 | c := New(Options{ 95 | Address: os.Getenv("REDIS_HOST"), 96 | PoolSize: 1, 97 | }) 98 | 99 | if err := c.Del(ctx, "foo"); err != nil { 100 | t.Fatal("start faild") 101 | } 102 | 103 | c.ZAdd(ctx, &Item{Key: "foo", ZSetValues: map[string]float64{"a": 1, "b": 2}}) 104 | 105 | values, _ := c.ZRange(ctx, "foo", 0, -1) 106 | if values[0].Member != "a" || 107 | values[0].Score != 1 || 108 | values[1].Member != "b" || 109 | values[1].Score != 2 { 110 | 111 | t.Fatal("zrange faild") 112 | } 113 | 114 | values, _ = c.ZRevRange(ctx, "foo", 0, -1) 115 | if values[1].Member != "a" || 116 | values[1].Score != 1 || 117 | values[0].Member != "b" || 118 | values[0].Score != 2 { 119 | 120 | t.Fatal("zrange faild") 121 | } 122 | 123 | values, _ = c.ZRangeByScore(ctx, "foo", 0, 1, 0, 0) 124 | if values[0].Member != "a" || 125 | values[0].Score != 1 { 126 | 127 | t.Fatal("zrangebyscore faild") 128 | } 129 | 130 | values, _ = c.ZRevRangeByScore(ctx, "foo", 2, 1, 0, 0) 131 | if values[1].Member != "a" || 132 | values[1].Score != 1 || 133 | values[0].Member != "b" || 134 | values[0].Score != 2 { 135 | 136 | t.Fatal("zrevrangebyscore faild") 137 | } 138 | 139 | values, _ = c.ZRevRangeByScore(ctx, "foo", 2, 1, 1, 1) 140 | if values[0].Member != "a" || 141 | values[0].Score != 1 { 142 | 143 | t.Fatal("zrevrangebyscore faild") 144 | } 145 | 146 | if c, _ := c.ZCard(ctx, "foo"); c != 2 { 147 | t.Fatal("zcard faild") 148 | } 149 | 150 | c.ZAdd(ctx, &Item{Key: "foo", ZSetValues: map[string]float64{"c": 2}}) 151 | 152 | if c, _ := c.ZCount(ctx, "foo", "(1", "+inf"); c != 2 { 153 | t.Fatal("zcount faild") 154 | } 155 | 156 | c.ZIncrBy(ctx, "foo", "a", 4.05) 157 | if s, _ := c.ZScore(ctx, "foo", "a"); s-5.05 >= 0.000001 { 158 | t.Fatal("zincrby faild") 159 | } 160 | 161 | if r, _ := c.ZRank(ctx, "foo", "a"); r != 2 { 162 | t.Fatal("zrank faild") 163 | } 164 | 165 | if r, err := c.ZRevRank(ctx, "foo", "a"); err != nil || r != 0 { 166 | t.Fatal("zrevrank faild") 167 | } 168 | 169 | c.ZRem(ctx, "foo", "a", "b", "c") 170 | if c, err := c.ZCard(ctx, "foo"); err != nil || c != 0 { 171 | t.Fatal("zrem faild") 172 | } 173 | 174 | c.ZAdd(ctx, &Item{Key: "foo", ZSetValues: map[string]float64{"a": 1, "b": 2, "c": 3, "d": 4}}) 175 | c.ZRemRangeByRank(ctx, "foo", 2, -1) 176 | values, _ = c.ZRange(ctx, "foo", 0, -1) 177 | if values[0].Member != "a" || values[1].Member != "b" { 178 | t.Fatal("zremrangebyrank faild") 179 | } 180 | 181 | c.ZAdd(ctx, &Item{Key: "foo", ZSetValues: map[string]float64{"a": 1, "b": 2, "c": 3, "d": 4}}) 182 | c.ZRemRangeByScore(ctx, "foo", "0", "2") 183 | values, _ = c.ZRange(ctx, "foo", 0, -1) 184 | if values[0].Member != "c" || values[1].Member != "d" { 185 | t.Fatal("zremrangebyscore faild") 186 | } 187 | } 188 | 189 | func TestOnCmd(t *testing.T) { 190 | c := New(Options{ 191 | Address: os.Getenv("REDIS_HOST"), 192 | PoolSize: 1, 193 | OnPreCmd: func(ctx context.Context, args []interface{}) context.Context { 194 | if len(args) == 0 { 195 | t.Fatal("OnPre failed") 196 | } 197 | 198 | return context.WithValue(ctx, "foo", "bar") 199 | }, 200 | OnPostCmd: func(ctx context.Context, err error) { 201 | if ctx.Value("foo").(string) != "bar" { 202 | t.Fatal("OnPostCmd failed") 203 | } 204 | }, 205 | }) 206 | 207 | if err := c.Del(ctx, "foo"); err != nil { 208 | t.Fatal("start faild") 209 | } 210 | 211 | c.Set(ctx, &Item{Key: "foo", Value: []byte("123"), Flags: FlagXX}) 212 | } 213 | 214 | func TestMget(t *testing.T) { 215 | c := New(Options{ 216 | Address: os.Getenv("REDIS_HOST"), 217 | PoolSize: 1, 218 | }) 219 | 220 | err := c.Set(ctx, &Item{Key: "key_m_1", Value: []byte("value_m_1")}) 221 | if err != nil { 222 | t.Fatal("Set Failed") 223 | } 224 | err = c.Set(ctx, &Item{Key: "key_m_2", Value: []byte("value_m_2")}) 225 | if err != nil { 226 | t.Fatal("Set Failed") 227 | } 228 | // 删除数据 229 | defer c.Del(ctx, "key_m_1", "key_m_2") 230 | 231 | // 批量获取 232 | items, err := c.MGet(ctx, []string{"key_m_1", "key_m_2", "key_m_3"}) 233 | if err != nil { 234 | t.Fatal("MGet Failed") 235 | } 236 | 237 | // 校验获取的值与插入的一致 238 | if string(items["key_m_1"].Value) != "value_m_1" { 239 | t.Fatal("MGet Failed") 240 | } 241 | 242 | if string(items["key_m_2"].Value) != "value_m_2" { 243 | t.Fatal("MGet Failed") 244 | } 245 | 246 | if _, ok := items["key_m_3"]; ok { 247 | t.Fatal("MGet Failed") 248 | } 249 | } 250 | 251 | func TestEval(t *testing.T) { 252 | c := New(Options{ 253 | Address: os.Getenv("REDIS_HOST"), 254 | PoolSize: 1, 255 | }) 256 | 257 | if err := c.Del(ctx, "foo"); err != nil { 258 | t.Fatal("start faild") 259 | } 260 | 261 | // array 262 | val, _ := c.Eval(ctx, "return {\"abcc\",1, {\"b\"} }", []string{}) 263 | if v, err := val.Array(); v[0].(string) != "abcc" || 264 | v[1].(int64) != 1 || 265 | v[2].([]interface{})[0].(string) != "b" { 266 | t.Fatal("eval faild", err) 267 | } 268 | 269 | // int64 270 | val, _ = c.Eval(ctx, "return 64", []string{}) 271 | if v, err := val.Int64(); v != 64 { 272 | t.Fatal("eval faild", err) 273 | } 274 | // string 275 | val, _ = c.Eval(ctx, "return ARGV[1]", []string{}, "a\nb\nc") 276 | if v, err := val.String(); v != "a\nb\nc" { 277 | t.Fatal("eval faild", err) 278 | } 279 | 280 | // status 281 | val, _ = c.Eval(ctx, "return redis.call('set',KEYS[1],ARGV[1])", []string{"foo"}, "hello") 282 | c.Del(ctx, "foo") 283 | if v, err := val.String(); v != "OK" { 284 | t.Fatal("eval faild", err) 285 | } 286 | 287 | // with no params 288 | c.Eval(ctx, "return redis.call('set',KEYS[1],ARGV[1])", []string{"foo"}, "hello") 289 | defer c.Del(ctx, "foo") 290 | 291 | item, _ := c.Get(ctx, "foo") 292 | if !bytes.Equal(item.Value, []byte("hello")) { 293 | t.Fatal("eval failed") 294 | } 295 | } 296 | 297 | func TestSet(t *testing.T) { 298 | c := New(Options{ 299 | Address: os.Getenv("REDIS_HOST"), 300 | PoolSize: 1, 301 | }) 302 | var err error 303 | 304 | if err := c.Del(ctx, "foo"); err != nil { 305 | t.Fatal("start faild") 306 | } 307 | 308 | // sadd 309 | if err := c.SAdd(ctx, "foo", []byte("aaa")); err != nil { 310 | t.Fatalf("add foo aaa failed, err: %v", err) 311 | } 312 | if err := c.SAdd(ctx, "foo", []byte("bbb"), []byte("ccc")); err != nil { 313 | t.Fatalf("add foo ddd,eee,fff failed, err: %v", err) 314 | } 315 | 316 | // scard 317 | if card, err := c.SCard(ctx, "foo"); err != nil || card != 3 { 318 | t.Fatalf("Key: foo, Scard: %d; Failed, err: %v", card, err) 319 | } 320 | 321 | // smembers 322 | items, err := c.SMembers(ctx, "foo") 323 | if err != nil { 324 | t.Fatalf("get foo's members failed, err: %v", err) 325 | } 326 | t.Logf("items: %#v", items) 327 | 328 | // sismember 329 | if result, err := c.SIsMember(ctx, "foo", []byte("bbb")); err != nil || result != true { 330 | t.Fatalf("Key: foo, SIsMember: %t; Failed, err: %v", result, err) 331 | } 332 | 333 | // spop 334 | item, _ := c.SPop(ctx, "foo", 1) 335 | if len(item) < 1 { 336 | t.Fatal("spop foo empty") 337 | } 338 | set := map[string]bool{"aaa": true, "bbb": true, "ccc": true} 339 | exists := set[string(item[0])] 340 | if !exists { 341 | t.Fatal("spop foo failed") 342 | } 343 | 344 | // srem 345 | if err := c.SAdd(ctx, "foo", []byte("ddd"), []byte("eee"), []byte("fff")); err != nil { 346 | t.Fatalf("add foo ddd failed, err: %v", err) 347 | } 348 | if result, err := c.SRem(ctx, "foo", []byte("ddd"), []byte("fff")); err != nil || result == 0 { 349 | t.Fatalf("Key: foo, SRem: %d; Failed, err: %v", result, err) 350 | } 351 | 352 | // srem 删除不存在的值 353 | if result, err := c.SRem(ctx, "foo", []byte("zzz"), []byte("xxx")); err != nil || result == 0 { 354 | t.Logf("Key: foo, SRem: %d; Failed, err: %v", result, err) 355 | } 356 | } 357 | -------------------------------------------------------------------------------- /util/strconv.go: -------------------------------------------------------------------------------- 1 | package util 2 | 3 | import "strconv" 4 | 5 | func Atoi(b []byte) (int, error) { 6 | return strconv.Atoi(BytesToString(b)) 7 | } 8 | 9 | func ParseInt(b []byte, base int, bitSize int) (int64, error) { 10 | return strconv.ParseInt(BytesToString(b), base, bitSize) 11 | } 12 | 13 | func ParseUint(b []byte, base int, bitSize int) (uint64, error) { 14 | return strconv.ParseUint(BytesToString(b), base, bitSize) 15 | } 16 | 17 | func ParseFloat(b []byte, bitSize int) (float64, error) { 18 | return strconv.ParseFloat(BytesToString(b), bitSize) 19 | } 20 | -------------------------------------------------------------------------------- /util/unsafe.go: -------------------------------------------------------------------------------- 1 | package util 2 | 3 | import ( 4 | "unsafe" 5 | ) 6 | 7 | // BytesToString converts byte slice to string. 8 | func BytesToString(b []byte) string { 9 | return *(*string)(unsafe.Pointer(&b)) 10 | } 11 | 12 | // StringToBytes converts string to byte slice. 13 | func StringToBytes(s string) []byte { 14 | return *(*[]byte)(unsafe.Pointer( 15 | &struct { 16 | string 17 | Cap int 18 | }{s, len(s)}, 19 | )) 20 | } 21 | --------------------------------------------------------------------------------