├── .gitignore ├── .travis.yml ├── LICENSE ├── README.md ├── cluster.go ├── command_test.go ├── commands.go ├── db.go ├── hash.go ├── key_buffer.go ├── list.go ├── lockring └── lockring.go ├── protocol.go ├── protocol_test.go ├── rdb.go ├── set.go ├── string.go ├── util.go └── zset.go /.gitignore: -------------------------------------------------------------------------------- 1 | # Compiled Object files, Static and Dynamic libs (Shared Objects) 2 | *.o 3 | *.a 4 | *.so 5 | 6 | # Folders 7 | _obj 8 | _test 9 | 10 | # Architecture specific extensions/prefixes 11 | *.[568vq] 12 | [568vq].out 13 | 14 | *.cgo1.go 15 | *.cgo2.c 16 | _cgo_defun.c 17 | _cgo_gotypes.go 18 | _cgo_export.* 19 | 20 | _testmain.go 21 | 22 | *.exe 23 | 24 | # SetDB-specific files 25 | db/ 26 | setdb 27 | .vagrant 28 | Vagrantfile 29 | tmp/ 30 | *.test 31 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: go 2 | go: 1.1 3 | before_install: 4 | - cd 5 | - wget https://leveldb.googlecode.com/files/leveldb-1.10.0.tar.gz 6 | - tar xzf leveldb-1.10.0.tar.gz 7 | - cd leveldb-1.10.0 8 | - make 9 | - export CGO_CFLAGS="-I`pwd`/include" 10 | - export CGO_LDFLAGS="-L`pwd`" 11 | - export LD_LIBRARY_PATH="`pwd`" 12 | - cd $TRAVIS_BUILD_DIR 13 | - go get launchpad.net/gocheck 14 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2013 Apollic Software, LLC. All rights reserved. 2 | 3 | Redistribution and use in source and binary forms, with or without 4 | modification, are permitted provided that the following conditions are 5 | met: 6 | 7 | * Redistributions of source code must retain the above copyright 8 | notice, this list of conditions and the following disclaimer. 9 | * Redistributions in binary form must reproduce the above 10 | copyright notice, this list of conditions and the following disclaimer 11 | in the documentation and/or other materials provided with the 12 | distribution. 13 | * Neither the name of Apollic Software, LLC nor the names of its 14 | contributors may be used to endorse or promote products derived from 15 | this software without specific prior written permission. 16 | 17 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 18 | "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 19 | LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR 20 | A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT 21 | OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, 22 | SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT 23 | LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, 24 | DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY 25 | THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 26 | (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 27 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 28 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # SetDB 2 | 3 | SetDB is an implementation of the Redis protocol that persists to LevelDB on 4 | disk instead of memory. This allows the dataset size to grow beyond the bounds 5 | of memory. 6 | 7 | If you have leveldb in a non-standard path (git clone), you need to explicitly set CGO_{C,LD}FLAGS. 8 | `CGO_CFLAGS="-I/path/to/leveldb/include" CGO_LDFLAGS="-L/path/to/leveldb" go get github.com/cupcake/setdb` 9 | 10 | ## Running 11 | 12 | Get the package 13 | 14 | go get github.com/cupcake/setdb 15 | 16 | Run it 17 | 18 | $GOPATH/bin/setdb 19 | 20 | Connect with redis client 21 | 22 | redis-cli -p 12345 23 | -------------------------------------------------------------------------------- /cluster.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "bytes" 5 | "fmt" 6 | "time" 7 | 8 | "github.com/cupcake/rdb" 9 | "github.com/garyburd/redigo/redis" 10 | "github.com/jmhodges/levigo" 11 | "github.com/titanous/bconv" 12 | ) 13 | 14 | func Restore(args [][]byte, wb *levigo.WriteBatch) interface{} { 15 | ttl, err := bconv.ParseInt(args[1], 10, 64) 16 | if err != nil { 17 | return InvalidIntError 18 | } 19 | err = rdb.DecodeDump(args[2], 0, args[0], ttl, &rdbDecoder{wb: wb}) 20 | if err != nil { 21 | return err 22 | } 23 | return ReplyOK 24 | } 25 | 26 | func Dump(args [][]byte, wb *levigo.WriteBatch) interface{} { 27 | res, err := dumpKey(args[0]) 28 | if err != nil { 29 | return err 30 | } 31 | return res 32 | } 33 | 34 | func dumpKey(key []byte) ([]byte, error) { 35 | buf := &bytes.Buffer{} 36 | e := &rdbEncoder{rdb.NewEncoder(buf)} 37 | err := e.encodeKey(key, true) 38 | if err != nil { 39 | return nil, err 40 | } 41 | if buf.Len() == 0 { 42 | return nil, nil 43 | } 44 | return buf.Bytes(), nil 45 | } 46 | 47 | func Migrate(args [][]byte, wb *levigo.WriteBatch) interface{} { 48 | timeout, err := bconv.ParseInt(args[4], 10, 64) 49 | if err != nil { 50 | return InvalidIntError 51 | } 52 | 53 | data, err := dumpKey(args[2]) 54 | if err != nil { 55 | return err 56 | } 57 | if data == nil { 58 | return ReplyNOKEY 59 | } 60 | 61 | t := time.Duration(timeout) * time.Millisecond 62 | r, err := redis.DialTimeout("tcp", string(args[0])+":"+string(args[1]), t, t, t) 63 | defer r.Close() 64 | if err != nil { 65 | return IOError{fmt.Errorf("error or timeout connecting to target instance")} 66 | } 67 | 68 | res, err := redis.String(r.Do("SELECT", args[3])) 69 | if _, ok := err.(redis.Error); ok { 70 | return fmt.Errorf("Target instance replied with error: %s", err) 71 | } 72 | if err != nil || res != "OK" { 73 | return IOError{fmt.Errorf("error or timeout performing SELECT of database %s on target instance", args[3])} 74 | } 75 | 76 | res, err = redis.String(r.Do("RESTORE", args[2], "0", data)) 77 | if _, ok := err.(redis.Error); ok { 78 | return fmt.Errorf("Target instance replied with error: %s", err) 79 | } 80 | if err != nil || res != "OK" { 81 | return IOError{fmt.Errorf("error or timeout performing RESTORE of key on target instance")} 82 | } 83 | 84 | _, err = delKey(metaKey(args[2]), wb) 85 | if err != nil { 86 | return IOError{fmt.Errorf("error deleting key from local instance: %s", err)} 87 | } 88 | 89 | return ReplyOK 90 | } 91 | -------------------------------------------------------------------------------- /command_test.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "bytes" 5 | "encoding/base64" 6 | "os" 7 | "testing" 8 | 9 | "github.com/jmhodges/levigo" 10 | . "launchpad.net/gocheck" 11 | ) 12 | 13 | // Hook gocheck into the gotest runner. 14 | func Test(t *testing.T) { TestingT(t) } 15 | 16 | type CommandSuite struct{} 17 | 18 | var _ = Suite(&CommandSuite{}) 19 | 20 | func (s CommandSuite) SetUpSuite(c *C) { 21 | os.RemoveAll("db") 22 | openDB() 23 | } 24 | 25 | var stringDump, _ = base64.StdEncoding.DecodeString("AAVIZWxsbwYAKiSfaMhQwxA=") 26 | var hashDump, _ = base64.StdEncoding.DecodeString("BAIGZmllbGQxBUhlbGxvBmZpZWxkMgVXb3JsZAYAgLPX8O51AWE=") 27 | var setDump, _ = base64.StdEncoding.DecodeString("AgIFSGVsbG8FV29ybGQGADOH1ks5wqwF") 28 | var zsetDump, _ = base64.StdEncoding.DecodeString("AwMDdHdvATMDdW5vATEDb25lATEGAGHZ/WW1zUvC") 29 | var listDump, _ = base64.StdEncoding.DecodeString("AQIFSGVsbG8FV29ybGQGACzxp+PtJo1E") 30 | 31 | var tests = []struct { 32 | command string 33 | args string 34 | response interface{} 35 | }{ 36 | {"ping", "", "PONG"}, 37 | {"echo", "foo", []byte("foo")}, 38 | {"zadd", "foo 1 bar", uint32(1)}, 39 | {"zadd", "foo 1 bar", uint32(0)}, 40 | {"zadd", "foo 2 bar", uint32(0)}, 41 | {"zadd", "foo 1 baz", uint32(1)}, 42 | {"zadd", "foo 1 baz 2 bar", uint32(0)}, 43 | {"zadd", "foo 5.1 asdf 2 buzz 1 baz 2 bar", uint32(2)}, 44 | {"zadd", "asdf 0.1 bar", uint32(1)}, 45 | {"zadd", "fooz 4e29 bar 0.2 baz", uint32(2)}, 46 | {"zscore", "foo bar", []byte("2")}, 47 | {"zscore", "foo baz", []byte("1")}, 48 | {"zscore", "asdf bar", []byte("0.1")}, 49 | {"zscore", "fooz bar", []byte("4e+29")}, 50 | {"zscore", "fooz bag", nil}, 51 | {"zincrby", "foo 0.1 bar", []byte("2.1")}, 52 | {"zincrby", "foo 1.1 bazz", []byte("1.1")}, 53 | {"zcard", "foo", uint32(5)}, 54 | {"zcard", "fooz", uint32(2)}, 55 | {"zcard", "asdf", uint32(1)}, 56 | {"zcard", "asdfa", uint32(0)}, 57 | {"zunionstore", "dz 2 foo fooz WEIGHTS 2 4 aggregate sum", uint32(5)}, 58 | {"zrange", "dz 0 -1 withscores", []interface{}{[]byte("bazz"), []byte("2.2"), []byte("baz"), []byte("2.8"), []byte("buzz"), []byte("4"), []byte("asdf"), []byte("10.2"), []byte("bar"), []byte("1.6e+30")}}, 59 | {"zinterstore", "dz 2 foo fooz WEIGHTS 2 4 aggregate min", uint32(2)}, 60 | {"zrange", "dz 0 -1 withscores", []interface{}{[]byte("baz"), []byte("0.8"), []byte("bar"), []byte("4.2")}}, 61 | {"sadd", "zs bar", uint32(1)}, 62 | {"zinterstore", "dz 2 foo zs aggregate max", uint32(1)}, 63 | {"zrange", "dz 0 -1 withscores", []interface{}{[]byte("bar"), []byte("2.1")}}, 64 | {"zrem", "foo bar baz", uint32(2)}, 65 | {"zrem", "foo bar", uint32(0)}, 66 | {"zrem", "asdfa bar", 0}, 67 | {"zcard", "foo", uint32(3)}, 68 | {"zrem", "asdf bar", uint32(1)}, 69 | {"zcard", "asdf", uint32(0)}, 70 | {"exists", "asdf", 0}, 71 | {"zrange", "foozzzz 0 1", []interface{}{}}, 72 | {"zrange", "foo 0 1", []interface{}{[]byte("bazz"), []byte("buzz")}}, 73 | {"zrange", "foo 0 -1 withscores", []interface{}{[]byte("bazz"), []byte("1.1"), []byte("buzz"), []byte("2"), []byte("asdf"), []byte("5.1")}}, 74 | {"zrange", "foo 0 -2 WITHSCORES", []interface{}{[]byte("bazz"), []byte("1.1"), []byte("buzz"), []byte("2")}}, 75 | {"zrange", "foo -1 -1", []interface{}{[]byte("asdf")}}, 76 | {"zrange", "foo 10 12", []interface{}{}}, 77 | {"zrange", "foo 2 1", []interface{}{}}, 78 | {"zrange", "foo -10 -1 withscores", []interface{}{[]byte("bazz"), []byte("1.1"), []byte("buzz"), []byte("2"), []byte("asdf"), []byte("5.1")}}, 79 | {"zrevrange", "foozzzz 0 1", []interface{}{}}, 80 | {"zrevrange", "foo 0 1", []interface{}{[]byte("asdf"), []byte("buzz")}}, 81 | {"zrevrange", "foo 0 -1 withscores", []interface{}{[]byte("asdf"), []byte("5.1"), []byte("buzz"), []byte("2"), []byte("bazz"), []byte("1.1")}}, 82 | {"zrevrange", "foo 0 -2 WITHSCORES", []interface{}{[]byte("asdf"), []byte("5.1"), []byte("buzz"), []byte("2")}}, 83 | {"zrevrange", "foo -1 -1", []interface{}{[]byte("bazz")}}, 84 | {"zrevrange", "foo 10 12", []interface{}{}}, 85 | {"zrevrange", "foo 2 1", []interface{}{}}, 86 | {"zrevrange", "foo -10 -1 withscores", []interface{}{[]byte("asdf"), []byte("5.1"), []byte("buzz"), []byte("2"), []byte("bazz"), []byte("1.1")}}, 87 | {"zrangebyscore", "foozzzz +inf -inf withscores", []interface{}{}}, 88 | {"zrangebyscore", "foo -inf +inf withscores", []interface{}{[]byte("bazz"), []byte("1.1"), []byte("buzz"), []byte("2"), []byte("asdf"), []byte("5.1")}}, 89 | {"zrangebyscore", "foo -inf +inf withscores limit 0 2", []interface{}{[]byte("bazz"), []byte("1.1"), []byte("buzz"), []byte("2")}}, 90 | {"zrangebyscore", "foo -inf +inf limit 1 1", []interface{}{[]byte("buzz")}}, 91 | {"zrangebyscore", "foo 1.1 1.2", []interface{}{[]byte("bazz")}}, 92 | {"zrangebyscore", "foo (2 +inf", []interface{}{[]byte("asdf")}}, 93 | {"zrangebyscore", "foo 2 (3", []interface{}{[]byte("buzz")}}, 94 | {"zrevrangebyscore", "foozzzz +inf -inf withscores", []interface{}{}}, 95 | {"zrevrangebyscore", "foo +inf -inf withscores", []interface{}{[]byte("asdf"), []byte("5.1"), []byte("buzz"), []byte("2"), []byte("bazz"), []byte("1.1")}}, 96 | {"zrevrangebyscore", "foo +inf -inf withscores limit 0 2", []interface{}{[]byte("asdf"), []byte("5.1"), []byte("buzz"), []byte("2")}}, 97 | {"zrevrangebyscore", "foo +inf -inf limit 1 1", []interface{}{[]byte("buzz")}}, 98 | {"zrevrangebyscore", "foo 1.2 1.1", []interface{}{[]byte("bazz")}}, 99 | {"zrevrangebyscore", "foo +inf (2", []interface{}{[]byte("asdf")}}, 100 | {"zrevrangebyscore", "foo (3 2", []interface{}{[]byte("buzz")}}, 101 | {"zcount", "foozzz -inf +inf", 0}, 102 | {"zcount", "foo -inf +inf", uint32(3)}, 103 | {"zcount", "foo 1.1 (1.3", uint32(1)}, 104 | {"zrank", "foo asdf", 2}, 105 | {"zrevrank", "foo asdf", 0}, 106 | {"zrank", "foo foobar", nil}, 107 | {"zrevrank", "foo foobar", nil}, 108 | {"zrank", "foozzz foobar", nil}, 109 | {"zrevrank", "foozzz foobar", nil}, 110 | {"zadd", "deletetest 1 one 2 two 3 three", uint32(3)}, 111 | {"zremrangebyscore", "deletetesting 1 2", 0}, 112 | {"zremrangebyscore", "deletetest 1 2", uint32(2)}, 113 | {"zcard", "deletetest", uint32(1)}, 114 | {"zrange", "deletetest 0 -1", []interface{}{[]byte("three")}}, 115 | {"zremrangebyscore", "deletetest -inf +inf", uint32(1)}, 116 | {"exists", "deletetest", 0}, 117 | {"zadd", "asdf 1 bar", uint32(1)}, 118 | {"del", "foo asdf", 2}, 119 | {"del", "foo asdf", 0}, 120 | {"zcard", "foo", uint32(0)}, 121 | {"exists", "foo", 0}, 122 | {"zcard", "asdf", uint32(0)}, 123 | {"exists", "asdf", 0}, 124 | {"set", "foo bar", "OK"}, 125 | {"get", "foo", []byte("bar")}, 126 | {"set", "foo baz", "OK"}, 127 | {"get", "foo", []byte("baz")}, 128 | {"append", "foo qux", 6}, 129 | {"get", "foo", []byte("bazqux")}, 130 | {"del", "foo", 1}, 131 | {"exists", "foo", 0}, 132 | {"zadd", "asdf 1 bar", uint32(1)}, 133 | {"set", "asdf foo", "OK"}, 134 | {"get", "asdf", []byte("foo")}, 135 | {"sadd", "aset 1 2 3 4 5", uint32(5)}, 136 | {"sadd", "set2 1 a 3", uint32(3)}, 137 | {"sadd", "set3 1 b 4", uint32(3)}, 138 | {"sunion", "aset set2 set3", []interface{}{[]byte("1"), []byte("2"), []byte("3"), []byte("4"), []byte("5"), []byte("a"), []byte("b")}}, 139 | {"sinter", "aset set2 set3", []interface{}{[]byte("1")}}, 140 | {"sinter", "aset bset set2 set3", []interface{}{}}, 141 | {"sdiff", "aset set2 set3", []interface{}{[]byte("2"), []byte("5")}}, 142 | {"sunionstore", "destset aset set2 set3", uint32(7)}, 143 | {"smembers", "destset", []interface{}{[]byte("1"), []byte("2"), []byte("3"), []byte("4"), []byte("5"), []byte("a"), []byte("b")}}, 144 | {"sinterstore", "destset aset set2 set3", uint32(1)}, 145 | {"smembers", "destset", []interface{}{[]byte("1")}}, 146 | {"sdiffstore", "destset aset set2 set3", uint32(2)}, 147 | {"smembers", "destset", []interface{}{[]byte("2"), []byte("5")}}, 148 | {"sadd", "aset 1", uint32(0)}, 149 | {"scard", "aset", uint32(5)}, 150 | {"sadd", "aset 6", uint32(1)}, 151 | {"scard", "aset", uint32(6)}, 152 | {"srem", "aset 4 5", uint32(2)}, 153 | {"srem", "aset 4 5", uint32(0)}, 154 | {"scard", "aset", uint32(4)}, 155 | {"sismember", "aset 6", 1}, 156 | {"sismember", "aset 7", 0}, 157 | {"smembers", "aset", []interface{}{[]byte("1"), []byte("2"), []byte("3"), []byte("6")}}, 158 | {"smembers", "bset", []interface{}{}}, 159 | {"smove", "aset newset 1", 1}, 160 | {"sismember", "aset 1", 0}, 161 | {"sismember", "newset 1", 1}, 162 | {"smove", "aset newset 1", 0}, 163 | {"del", "aset", 1}, 164 | {"exists", "aset", 0}, 165 | {"scard", "aset", uint32(0)}, 166 | {"sadd", "bset a", uint32(1)}, 167 | {"spop", "bset", []byte("a")}, 168 | {"scard", "bset", uint32(0)}, 169 | {"hset", "hash foo bar", 1}, 170 | {"hget", "hash foo", []byte("bar")}, 171 | {"hget", "hash0 baz", []byte(nil)}, 172 | {"hset", "hash foo baz", 0}, 173 | {"hget", "hash foo", []byte("baz")}, 174 | {"hlen", "hash", uint32(1)}, 175 | {"hlen", "haz", uint32(0)}, 176 | {"hset", "hash bar baz", 1}, 177 | {"hlen", "hash", uint32(2)}, 178 | {"hmset", "hash2 foo 1 bar 2 baz 3", "OK"}, 179 | {"hmget", "hash2 foo bar test", []interface{}{[]byte("1"), []byte("2"), []byte(nil)}}, 180 | {"hgetall", "hash2", []interface{}{[]byte("bar"), []byte("2"), []byte("baz"), []byte("3"), []byte("foo"), []byte("1")}}, 181 | {"hgetall", "hash3", []interface{}{}}, 182 | {"hkeys", "hash2", []interface{}{[]byte("bar"), []byte("baz"), []byte("foo")}}, 183 | {"hkeys", "hash3", []interface{}{}}, 184 | {"hvals", "hash2", []interface{}{[]byte("2"), []byte("3"), []byte("1")}}, 185 | {"hvals", "hash3", []interface{}{}}, 186 | {"hexists", "hash2 bar", 1}, 187 | {"hexists", "hash2 bax", 0}, 188 | {"hsetnx", "hash2 foox 3", 1}, 189 | {"hget", "hash2 foox", []byte("3")}, 190 | {"hsetnx", "hash2 foox 4", 0}, 191 | {"hget", "hash2 foox", []byte("3")}, 192 | {"hincrby", "hash2 fooz 3", []byte("3")}, 193 | {"hget", "hash2 fooz", []byte("3")}, 194 | {"hincrby", "hash2 fooz -20", []byte("-17")}, 195 | {"hget", "hash2 fooz", []byte("-17")}, 196 | {"hincrbyfloat", "hash2 fooa 10.50", []byte("10.5")}, 197 | {"hget", "hash2 fooa", []byte("10.5")}, 198 | {"hset", "hash2 fooa 5.0e3", 0}, 199 | {"hincrbyfloat", "hash2 fooa 2.0e2", []byte("5200")}, 200 | {"hget", "hash2 fooa", []byte("5200")}, 201 | {"keys", "hash*", []interface{}{[]byte("hash"), []byte("hash2")}}, 202 | {"del", "hash2", 1}, 203 | {"hlen", "hash2", uint32(0)}, 204 | {"exists", "hash", 1}, 205 | {"exists", "hash2", 0}, 206 | {"type", "hash", "hash"}, 207 | {"type", "asdf", "string"}, 208 | {"type", "newset", "set"}, 209 | {"type", "fooz", "zset"}, 210 | {"type", "aaaaa", "none"}, 211 | {"lpush", "mylist world hello", uint32(2)}, 212 | {"llen", "mylist", uint32(2)}, 213 | {"lpop", "mylist", []byte("hello")}, 214 | {"llen", "mylist", uint32(1)}, 215 | {"lpop", "mylist", []byte("world")}, 216 | {"llen", "mylist", uint32(0)}, 217 | {"rpush", "mylist hello world test", uint32(3)}, 218 | {"rpushx", "mylist test2", uint32(4)}, 219 | {"rpushx", "myotherlist test", uint32(0)}, 220 | {"lpushx", "mylist test", uint32(5)}, 221 | {"lpushx", "myotherlist test", uint32(0)}, 222 | {"rpop", "mylist", []byte("test2")}, 223 | {"rpoplpush", "mylist myotherlist", []byte("test")}, 224 | {"rpop", "myotherlist", []byte("test")}, 225 | {"llen", "myotherlist", uint32(0)}, 226 | {"exists", "myotherlist", 0}, 227 | {"lrange", "mylist 0 -1", []interface{}{[]byte("test"), []byte("hello"), []byte("world")}}, 228 | {"lrange", "mylist -1 -1", []interface{}{[]byte("world")}}, 229 | {"lrange", "mylist -2 -5", []interface{}{}}, 230 | {"lrange", "mylist 1 2", []interface{}{[]byte("hello"), []byte("world")}}, 231 | {"restore", "r 0 " + string(stringDump), "OK"}, 232 | {"dump", "r", stringDump}, 233 | {"get", "r", []byte("Hello")}, 234 | {"restore", "r 0 " + string(hashDump), "OK"}, 235 | {"dump", "r", hashDump}, 236 | {"hlen", "r", uint32(2)}, 237 | {"hgetall", "r", []interface{}{[]byte("field1"), []byte("Hello"), []byte("field2"), []byte("World")}}, 238 | {"restore", "r 0 " + string(setDump), "OK"}, 239 | {"dump", "r", setDump}, 240 | {"scard", "r", uint32(2)}, 241 | {"smembers", "r", []interface{}{[]byte("Hello"), []byte("World")}}, 242 | {"restore", "r 0 " + string(zsetDump), "OK"}, 243 | {"dump", "r", zsetDump}, 244 | {"zcard", "r", uint32(3)}, 245 | {"zrange", "r 0 -1 withscores", []interface{}{[]byte("one"), []byte("1"), []byte("uno"), []byte("1"), []byte("two"), []byte("3")}}, 246 | {"restore", "r 0 " + string(listDump), "OK"}, 247 | {"dump", "r", listDump}, 248 | {"llen", "r", uint32(2)}, 249 | {"lrange", "r 0 -1", []interface{}{[]byte("Hello"), []byte("World")}}, 250 | } 251 | 252 | func (s CommandSuite) TestCommands(c *C) { 253 | for _, t := range tests { 254 | cmd := commands[t.command] 255 | var wb *levigo.WriteBatch 256 | if cmd.writes { 257 | wb = levigo.NewWriteBatch() 258 | } 259 | var args [][]byte 260 | if t.args != "" { 261 | if cmd.arity > 0 { 262 | args = bytes.SplitN([]byte(t.args), []byte(" "), cmd.arity) 263 | } else { 264 | args = bytes.Split([]byte(t.args), []byte(" ")) 265 | } 266 | } 267 | cmd.lockKeys(args) 268 | res := cmd.function(args, wb) 269 | if cmd.writes { 270 | err := DB.Write(DefaultWriteOptions, wb) 271 | c.Assert(err, IsNil) 272 | wb.Close() 273 | } 274 | cmd.unlockKeys(args) 275 | if stream, ok := res.(*cmdReplyStream); ok { 276 | items := make([]interface{}, 0, int(stream.size)) 277 | for item := range stream.items { 278 | items = append(items, item) 279 | } 280 | res = items 281 | } 282 | if reply, ok := res.(rawReply); ok { 283 | if reply[0] == '+' { 284 | res = string(reply[1 : len(reply)-2]) 285 | } 286 | } 287 | c.Assert(res, DeepEquals, t.response, Commentf("%s %s, obtained=%s expected=%s", t.command, t.args, res, t.response)) 288 | } 289 | } 290 | -------------------------------------------------------------------------------- /commands.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "bytes" 5 | "fmt" 6 | "path/filepath" 7 | "strconv" 8 | "time" 9 | 10 | "github.com/cupcake/setdb/lockring" 11 | "github.com/jmhodges/levigo" 12 | "github.com/titanous/bconv" 13 | ) 14 | 15 | var KeyMutex = lockring.New(1024) 16 | 17 | // Key/Value type identifiers, only append to this list 18 | const ( 19 | MetaKey byte = iota 20 | StringKey 21 | HashKey 22 | ListKey 23 | SetKey 24 | ZSetKey 25 | ZScoreKey 26 | StringLengthValue 27 | HashLengthValue 28 | ListLengthValue 29 | SetCardValue 30 | ZCardValue 31 | ) 32 | 33 | var ( 34 | InvalidKeyTypeError = fmt.Errorf("Operation against a key holding the wrong kind of value") 35 | InvalidDataError = fmt.Errorf("Invalid data") 36 | InvalidIntError = fmt.Errorf("value is not an integer or out of range") 37 | SyntaxError = fmt.Errorf("syntax error") 38 | ) 39 | 40 | var ( 41 | ReplyOK = rawReply("+OK\r\n") 42 | ReplyPONG = rawReply("+PONG\r\n") 43 | ReplyNOKEY = rawReply("+NOKEY\r\n") 44 | ) 45 | 46 | type IOError struct{ error } 47 | 48 | // if the number of items is known before the items, 49 | // they do not need to be buffered into memory, and can be streamed over a channel 50 | type cmdReplyStream struct { 51 | size int64 // the number of items that will be sent 52 | items chan interface{} // a multi-bulk reply item, one of nil, []byte, or int 53 | } 54 | 55 | // A raw reply will be returned verbatim to the client 56 | type rawReply []byte 57 | 58 | // cmdFunc response to Redis protocol conversion: 59 | // 60 | // string - single line reply, automatically prefixed with "+" 61 | // error - error message, automatically prefixed with "-" 62 | // int - integer number, automatically encoded and prefixed with ":" 63 | // []byte - bulk reply, automatically prefixed with the length like "$3\r\n" 64 | // nil, nil []byte - nil response, encoded as "$-1\r\n" 65 | // rawReply - no serialization, returned verbatim 66 | // []interface{} - multi-bulk reply, automatically serialized, members can be nil, []byte, or int 67 | // nil []interface{} - nil multi-bulk reply, serialized as "*-1\r\n" 68 | // map[string]bool - multi-bulk reply (used by SUNION) 69 | // *cmdReplyStream - multi-bulk reply sent over a channel 70 | type cmdFunc func(args [][]byte, wb *levigo.WriteBatch) interface{} 71 | 72 | type cmdDesc struct { 73 | name string 74 | function cmdFunc 75 | arity int // the number of required arguments, -n means >= n 76 | writes bool // false if the command doesn't write data (the WriteBatch will not be passed in) 77 | firstKey int // first argument that is a key (-1 for none) 78 | lastKey int // last argument that is a key (-1 for unbounded) 79 | keyStep int // step to get all the keys from first to last. For instance MSET is 2 since the arguments are KEY VAL KEY VAL... 80 | keyLookup func([][]byte) [][]byte // function that extracts the keys from the args 81 | } 82 | 83 | var commandList = []cmdDesc{ 84 | {"del", Del, -1, true, 0, -1, 1, nil}, 85 | {"echo", Echo, 1, false, -1, 0, 0, nil}, 86 | {"exists", Exists, 1, false, 0, 0, 0, nil}, 87 | {"get", Get, 1, false, 0, 0, 0, nil}, 88 | {"hdel", Hdel, -2, true, 0, 0, 0, nil}, 89 | {"hexists", Hexists, 2, false, 0, 0, 0, nil}, 90 | {"hget", Hget, 2, false, 0, 0, 0, nil}, 91 | {"hgetall", Hgetall, 1, false, 0, 0, 0, nil}, 92 | {"hincrby", Hincrby, 3, true, 0, 0, 0, nil}, 93 | {"hincrbyfloat", Hincrbyfloat, 3, true, 0, 0, 0, nil}, 94 | {"hkeys", Hkeys, 1, false, 0, 0, 0, nil}, 95 | {"hlen", Hlen, 1, false, 0, 0, 0, nil}, 96 | {"hmget", Hmget, -2, false, 0, 0, 0, nil}, 97 | {"hmset", Hmset, -3, true, 0, 0, 0, nil}, 98 | {"hset", Hset, 3, true, 0, 0, 0, nil}, 99 | {"hsetnx", Hsetnx, 3, true, 0, 0, 0, nil}, 100 | {"hvals", Hvals, 1, false, 0, 0, 0, nil}, 101 | {"keys", Keys, 1, false, -1, 0, 0, nil}, 102 | {"llen", Llen, 1, false, 0, 0, 0, nil}, 103 | {"lpush", Lpush, -2, true, 0, 0, 0, nil}, 104 | {"lpushx", Lpushx, 2, true, 0, 0, 0, nil}, 105 | {"rpush", Rpush, -2, true, 0, 0, 0, nil}, 106 | {"rpushx", Rpushx, 2, true, 0, 0, 0, nil}, 107 | {"lpop", Lpop, 1, true, 0, 0, 0, nil}, 108 | {"rpop", Rpop, 1, true, 0, 0, 0, nil}, 109 | {"rpoplpush", Rpoplpush, 2, true, 0, 1, 0, nil}, 110 | {"lrange", Lrange, 3, false, 0, 0, 0, nil}, 111 | {"ping", Ping, 0, false, -1, 0, 0, nil}, 112 | {"append", Append, 2, true, 0, 0, 0, nil}, 113 | {"set", Set, 2, true, 0, 0, 0, nil}, 114 | {"sadd", Sadd, -2, true, 0, 0, 0, nil}, 115 | {"scard", Scard, 1, false, 0, 0, 0, nil}, 116 | {"sismember", Sismember, 2, false, 0, 0, 0, nil}, 117 | {"smembers", Smembers, 1, false, 0, 0, 0, nil}, 118 | {"smove", Smove, 3, true, 0, 1, 0, nil}, 119 | {"spop", Spop, 1, true, 0, 0, 0, nil}, 120 | {"srem", Srem, -2, true, 0, 0, 0, nil}, 121 | {"sunion", Sunion, -1, false, 0, -1, 1, nil}, 122 | {"sunionstore", Sunionstore, -2, true, 0, -1, 1, nil}, 123 | {"sinter", Sinter, -1, false, 0, -1, 1, nil}, 124 | {"sinterstore", Sinterstore, -2, true, 0, -1, 1, nil}, 125 | {"sdiff", Sdiff, -1, false, 0, -1, 1, nil}, 126 | {"sdiffstore", Sdiffstore, -2, true, 0, -1, 1, nil}, 127 | {"time", Time, 0, false, -1, 0, 0, nil}, 128 | {"type", Type, 1, false, 0, 0, 0, nil}, 129 | {"zadd", Zadd, -3, true, 0, 0, 0, nil}, 130 | {"zcard", Zcard, 1, false, 0, 0, 0, nil}, 131 | {"zincrby", Zincrby, 3, true, 0, 0, 0, nil}, 132 | {"zrange", Zrange, -3, false, 0, 0, 0, nil}, 133 | {"zrem", Zrem, -2, true, 0, 0, 0, nil}, 134 | {"zrevrange", Zrevrange, -3, false, 0, 0, 0, nil}, 135 | {"zrangebyscore", Zrangebyscore, -3, false, 0, 0, 0, nil}, 136 | {"zrevrangebyscore", Zrevrangebyscore, -3, false, 0, 0, 0, nil}, 137 | {"zremrangebyscore", Zremrangebyscore, 3, true, 0, 0, 0, nil}, 138 | {"zcount", Zcount, 3, false, 0, 0, 0, nil}, 139 | {"zscore", Zscore, 2, false, 0, 0, 0, nil}, 140 | {"zrank", Zrank, 2, false, 0, 0, 0, nil}, 141 | {"zrevrank", Zrevrank, 2, false, 0, 0, 0, nil}, 142 | {"zunionstore", Zunionstore, -3, true, 0, 0, 0, ZunionInterKeys}, 143 | {"zinterstore", Zinterstore, -3, true, 0, 0, 0, ZunionInterKeys}, 144 | {"restore", Restore, 3, true, 0, 0, 0, nil}, 145 | {"dump", Dump, 1, false, 0, 0, 0, nil}, 146 | {"migrate", Migrate, 5, true, 2, 2, 0, nil}, 147 | {"select", Select, 1, false, 0, 0, 0, nil}, 148 | } 149 | 150 | // extract the keys from the command args 151 | func (c *cmdDesc) getKeys(args [][]byte) [][]byte { 152 | // if a key lookup function was specified, use it 153 | if c.keyLookup != nil { 154 | return c.keyLookup(args) 155 | } 156 | // if no keys are expected, or the argument with the first key doesn't exist 157 | if c.firstKey < 0 || len(args) <= c.firstKey { 158 | return nil 159 | } 160 | // shortcut: if the keystep is 0 or 1, we can slice the array 161 | if c.keyStep <= 1 { 162 | return args[c.firstKey : c.lastKey+1] 163 | } 164 | keys := make([][]byte, 0, 1) 165 | keyloop: 166 | for i := c.firstKey; i < len(args) && (c.lastKey == -1 || i <= c.lastKey); i += c.keyStep { 167 | for _, k := range keys { 168 | // skip keys that are already in the array 169 | if bytes.Equal(k, args[i]) { 170 | continue keyloop 171 | } 172 | } 173 | keys = append(keys, args[i]) 174 | } 175 | return keys 176 | } 177 | 178 | // acquires a read or write lock for the keys in arguments using the cmdDesc 179 | func (c *cmdDesc) lockKeys(args [][]byte) { 180 | if !c.writes { 181 | return 182 | } 183 | for _, k := range c.getKeys(args) { 184 | KeyMutex.Lock(k) 185 | } 186 | } 187 | 188 | func (c *cmdDesc) unlockKeys(args [][]byte) { 189 | if !c.writes { 190 | return 191 | } 192 | for _, k := range c.getKeys(args) { 193 | KeyMutex.Unlock(k) 194 | } 195 | } 196 | 197 | var commands = make(map[string]cmdDesc, len(commandList)) 198 | 199 | func Ping(args [][]byte, wb *levigo.WriteBatch) interface{} { 200 | return ReplyPONG 201 | } 202 | 203 | func Echo(args [][]byte, wb *levigo.WriteBatch) interface{} { 204 | return args[0] 205 | } 206 | 207 | func Time(args [][]byte, wb *levigo.WriteBatch) interface{} { 208 | now := time.Now() 209 | secs := strconv.AppendInt(nil, now.Unix(), 10) 210 | micros := strconv.AppendInt(nil, int64(now.Nanosecond()/1000), 10) 211 | return []interface{}{secs, micros} 212 | } 213 | 214 | func Exists(args [][]byte, wb *levigo.WriteBatch) interface{} { 215 | res, err := DB.Get(DefaultReadOptions, metaKey(args[0])) 216 | if err != nil { 217 | return err 218 | } 219 | if res == nil { 220 | return 0 221 | } 222 | return 1 223 | } 224 | 225 | func Type(args [][]byte, wb *levigo.WriteBatch) interface{} { 226 | res, err := DB.Get(DefaultReadOptions, metaKey(args[0])) 227 | if err != nil { 228 | return err 229 | } 230 | if res == nil { 231 | return "none" 232 | } 233 | if len(res) == 0 { 234 | return InvalidDataError 235 | } 236 | switch res[0] { 237 | case StringLengthValue: 238 | return "string" 239 | case HashLengthValue: 240 | return "hash" 241 | case ListLengthValue: 242 | return "list" 243 | case SetCardValue: 244 | return "set" 245 | case ZCardValue: 246 | return "zset" 247 | } 248 | panic("unknown type") 249 | } 250 | 251 | func Keys(args [][]byte, wb *levigo.WriteBatch) interface{} { 252 | it := DB.NewIterator(ReadWithoutCacheFill) 253 | defer it.Close() 254 | keys := []interface{}{} 255 | pattern := string(args[0]) 256 | 257 | for it.Seek([]byte{MetaKey}); it.Valid(); it.Next() { 258 | k := it.Key() 259 | // if the first byte isn't MetaKey, we've reached the end 260 | if len(k) < 2 || k[0] != MetaKey { 261 | break 262 | } 263 | // filepatch.Match() implements the same pattern syntax as we want 264 | matched, err := filepath.Match(pattern, string(k[1:])) 265 | if err != nil { 266 | return fmt.Errorf("invalid pattern for 'keys' command") 267 | } 268 | if matched { 269 | keys = append(keys, k[1:]) 270 | } 271 | } 272 | return keys 273 | } 274 | 275 | // No-op for now 276 | func Select(args [][]byte, wb *levigo.WriteBatch) interface{} { 277 | return ReplyOK 278 | } 279 | 280 | func Del(args [][]byte, wb *levigo.WriteBatch) interface{} { 281 | deleted := 0 282 | k := make([]byte, 1, len(args[0])) // make a reusable slice with room for the first metakey 283 | 284 | for _, key := range args { 285 | k = bufMetaKey(k, key) 286 | d, err := delKey(k, wb) 287 | if err != nil { 288 | return err 289 | } 290 | if d { 291 | deleted++ 292 | } 293 | } 294 | return deleted 295 | } 296 | 297 | func delKey(key []byte, wb *levigo.WriteBatch) (deleted bool, err error) { 298 | res, err := DB.Get(ReadWithoutCacheFill, key) 299 | if err != nil { 300 | return 301 | } 302 | if res == nil { 303 | return 304 | } 305 | if len(res) < 1 { 306 | return false, InvalidDataError 307 | } 308 | del(key[1:], res[0], wb) 309 | wb.Delete(key) 310 | return true, nil 311 | } 312 | 313 | func del(key []byte, t byte, wb *levigo.WriteBatch) { 314 | switch t { 315 | case StringLengthValue: 316 | DelString(key, wb) 317 | case HashLengthValue: 318 | DelHash(key, wb) 319 | case SetCardValue: 320 | DelSet(key, wb) 321 | case ZCardValue: 322 | DelZset(key, wb) 323 | default: 324 | panic("unknown key type") 325 | } 326 | } 327 | 328 | // set buf to the metaKey for key 329 | func bufMetaKey(buf []byte, key []byte) []byte { 330 | buf[0] = MetaKey 331 | return append(buf[:1], key...) 332 | } 333 | 334 | func metaKey(k []byte) []byte { 335 | key := make([]byte, 1+len(k)) 336 | key[0] = MetaKey 337 | copy(key[1:], k) 338 | return key 339 | } 340 | 341 | func parseRange(args [][]byte, length int64) (int64, int64, error) { 342 | start, err := bconv.ParseInt(args[0], 10, 64) 343 | end, err2 := bconv.ParseInt(args[1], 10, 64) 344 | if err != nil || err2 != nil { 345 | return 0, 0, InvalidIntError 346 | } 347 | 348 | // if the index is negative, it is counting from the end, 349 | // so add it to the length to get the absolute index 350 | if start < 0 { 351 | start += length 352 | } 353 | if end < 0 { 354 | end += length 355 | } 356 | 357 | if end > length { // limit the end to the last member 358 | end = length - 1 359 | } 360 | 361 | return start, end, nil 362 | } 363 | 364 | func init() { 365 | for _, c := range commandList { 366 | commands[c.name] = c 367 | } 368 | } 369 | 370 | // Keys 371 | // EXPIRE 372 | // EXPIREAT 373 | // MOVE 374 | // OBJECT? 375 | // PERSIST 376 | // PEXPIRE 377 | // PEXPIREAT 378 | // PTTL 379 | // RANDOMKEY 380 | // RENAME 381 | // RENAMENX 382 | // SORT 383 | // TTL 384 | // TYPE 385 | // 386 | // Pub/Sub 387 | // PSUBSCRIBE 388 | // PUNSUBSCRIBE 389 | // UNSUBSCRIBE 390 | // PUBLISH 391 | // SUBSCRIBE 392 | // 393 | // Transactions 394 | // DISCARD 395 | // MULTI 396 | // EXEC 397 | // WATCH 398 | // UNWATCH 399 | // 400 | // Scripting 401 | // EVAL 402 | // EVALSHA 403 | // SCRIPT EXISTS 404 | // SCRIPT KILL 405 | // SCRIPT FLUSH 406 | // SCRIPT LOAD 407 | // 408 | // Connection 409 | // AUTH 410 | // SELECT 411 | // QUIT 412 | // 413 | // Server 414 | // FLUSHALL 415 | // FLUSHDB 416 | // SYNC 417 | // CONFIG RESETSTAT 418 | // INFO 419 | // DBSIZE 420 | // CLIENT LIST 421 | // CLIENT KILL 422 | // MONITOR 423 | // CONFIG GET 424 | // CONFIG SET 425 | // SLAVEOF 426 | // SHUTDOWN 427 | // SAVE 428 | // SLOWLOG 429 | -------------------------------------------------------------------------------- /db.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "fmt" 5 | "log" 6 | "net/http" 7 | _ "net/http/pprof" 8 | "os" 9 | "runtime" 10 | 11 | "github.com/jmhodges/levigo" 12 | ) 13 | 14 | var DB *levigo.DB 15 | var DefaultReadOptions = levigo.NewReadOptions() 16 | var DefaultWriteOptions = levigo.NewWriteOptions() 17 | var ReadWithoutCacheFill = levigo.NewReadOptions() 18 | 19 | func openDB() { 20 | opts := levigo.NewOptions() 21 | cache := levigo.NewLRUCache(128 * 1024 * 1024) // 128MB cache 22 | opts.SetCache(cache) 23 | filter := levigo.NewBloomFilter(10) 24 | opts.SetFilterPolicy(filter) 25 | opts.SetCreateIfMissing(true) 26 | 27 | var err error 28 | DB, err = levigo.Open("db", opts) 29 | maybeFatal(err) 30 | } 31 | 32 | func maybeFatal(err error) { 33 | if err != nil { 34 | fmt.Printf("Fatal error: %s\n", err) 35 | os.Exit(1) 36 | } 37 | } 38 | 39 | func main() { 40 | runtime.GOMAXPROCS(runtime.NumCPU()) 41 | openDB() 42 | go func() { 43 | log.Println(http.ListenAndServe("localhost:6060", nil)) 44 | }() 45 | listen() 46 | } 47 | 48 | func init() { 49 | ReadWithoutCacheFill.SetFillCache(false) 50 | } 51 | -------------------------------------------------------------------------------- /hash.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "encoding/binary" 5 | "fmt" 6 | "strconv" 7 | 8 | "github.com/jmhodges/levigo" 9 | "github.com/titanous/bconv" 10 | ) 11 | 12 | // Keys stored in LevelDB for hashes 13 | // 14 | // MetadataKey | key = HashLengthValue | count of fields uint32 15 | // 16 | // For each field: 17 | // HashKey | key length uint32 | key | field = value 18 | 19 | func Hset(args [][]byte, wb *levigo.WriteBatch) interface{} { 20 | return hset(args, true, wb) 21 | } 22 | 23 | func Hsetnx(args [][]byte, wb *levigo.WriteBatch) interface{} { 24 | return hset(args, false, wb) 25 | } 26 | 27 | func hset(args [][]byte, overwrite bool, wb *levigo.WriteBatch) interface{} { 28 | mk := metaKey(args[0]) 29 | length, err := hlen(mk, nil) 30 | if err != nil { 31 | return err 32 | } 33 | var res []byte 34 | key := NewKeyBufferWithSuffix(HashKey, args[0], args[1]).Key() 35 | if length > 0 { 36 | res, err = DB.Get(DefaultReadOptions, key) 37 | if err != nil { 38 | return err 39 | } 40 | } 41 | if overwrite || res == nil { 42 | wb.Put(key, args[2]) 43 | } 44 | if res == nil { 45 | setHlen(mk, length+1, wb) 46 | return 1 47 | } 48 | return 0 49 | } 50 | 51 | func Hget(args [][]byte, wb *levigo.WriteBatch) interface{} { 52 | res, err := DB.Get(DefaultReadOptions, NewKeyBufferWithSuffix(HashKey, args[0], args[1]).Key()) 53 | if err != nil { 54 | return err 55 | } 56 | return res 57 | } 58 | 59 | func Hexists(args [][]byte, wb *levigo.WriteBatch) interface{} { 60 | res, err := DB.Get(DefaultReadOptions, NewKeyBufferWithSuffix(HashKey, args[0], args[1]).Key()) 61 | if err != nil { 62 | return err 63 | } 64 | if res == nil { 65 | return 0 66 | } 67 | return 1 68 | } 69 | 70 | func Hlen(args [][]byte, wb *levigo.WriteBatch) interface{} { 71 | length, err := hlen(metaKey(args[0]), nil) 72 | if err != nil { 73 | return err 74 | } 75 | return length 76 | } 77 | 78 | func Hdel(args [][]byte, wb *levigo.WriteBatch) interface{} { 79 | mk := metaKey(args[0]) 80 | length, err := hlen(mk, nil) 81 | if err != nil { 82 | return err 83 | } 84 | if length == 0 { 85 | return 0 86 | } 87 | 88 | var deleted uint32 89 | key := NewKeyBuffer(HashKey, args[0], len(args[1])) 90 | for _, field := range args[1:] { 91 | key.SetSuffix(field) 92 | res, err := DB.Get(ReadWithoutCacheFill, key.Key()) 93 | if err != nil { 94 | return err 95 | } 96 | if res == nil { 97 | continue 98 | } 99 | wb.Delete(key.Key()) 100 | deleted++ 101 | } 102 | if deleted == length { 103 | wb.Delete(mk) 104 | } else if deleted > 0 { 105 | setHlen(mk, length-deleted, wb) 106 | } 107 | return deleted 108 | } 109 | 110 | func Hmset(args [][]byte, wb *levigo.WriteBatch) interface{} { 111 | if (len(args)-1)%2 != 0 { 112 | return fmt.Errorf("wrong number of arguments for 'hmset' command") 113 | } 114 | 115 | mk := metaKey(args[0]) 116 | length, err := hlen(mk, nil) 117 | if err != nil { 118 | return err 119 | } 120 | 121 | var added uint32 122 | key := NewKeyBuffer(HashKey, args[0], len(args[1])) 123 | for i := 1; i < len(args); i += 2 { 124 | key.SetSuffix(args[i]) 125 | var res []byte 126 | if length > 0 { 127 | res, err = DB.Get(DefaultReadOptions, key.Key()) 128 | if err != nil { 129 | return err 130 | } 131 | } 132 | if res == nil { 133 | added++ 134 | } 135 | wb.Put(key.Key(), args[i+1]) 136 | } 137 | if added > 0 { 138 | setHlen(mk, length+added, wb) 139 | } 140 | return ReplyOK 141 | } 142 | 143 | func Hmget(args [][]byte, wb *levigo.WriteBatch) interface{} { 144 | stream := &cmdReplyStream{int64(len(args) - 1), make(chan interface{})} 145 | go func() { 146 | defer close(stream.items) 147 | key := NewKeyBuffer(HashKey, args[0], len(args[1])) 148 | for _, field := range args[1:] { 149 | key.SetSuffix(field) 150 | res, err := DB.Get(DefaultReadOptions, key.Key()) 151 | if err != nil { 152 | stream.items <- err 153 | continue 154 | } 155 | stream.items <- res 156 | } 157 | }() 158 | return stream 159 | } 160 | 161 | func Hincrby(args [][]byte, wb *levigo.WriteBatch) interface{} { 162 | mk := metaKey(args[0]) 163 | length, err := hlen(mk, nil) 164 | if err != nil { 165 | return err 166 | } 167 | key := NewKeyBufferWithSuffix(HashKey, args[0], args[1]).Key() 168 | res, err := DB.Get(DefaultReadOptions, key) 169 | if err != nil { 170 | return err 171 | } 172 | 173 | var current int64 174 | if res != nil { 175 | current, err = bconv.ParseInt(res, 10, 64) 176 | if err != nil { 177 | return InvalidIntError 178 | } 179 | } 180 | increment, err := bconv.ParseInt(args[2], 10, 64) 181 | if err != nil { 182 | return InvalidIntError 183 | } 184 | result := strconv.AppendInt(nil, current+increment, 10) 185 | wb.Put(key, result) 186 | 187 | // if is a new key, increment the hash length 188 | if res == nil { 189 | setHlen(mk, length+1, wb) 190 | } 191 | return result 192 | } 193 | 194 | func Hincrbyfloat(args [][]byte, wb *levigo.WriteBatch) interface{} { 195 | mk := metaKey(args[0]) 196 | length, err := hlen(mk, nil) 197 | if err != nil { 198 | return err 199 | } 200 | key := NewKeyBufferWithSuffix(HashKey, args[0], args[1]).Key() 201 | res, err := DB.Get(DefaultReadOptions, key) 202 | if err != nil { 203 | return err 204 | } 205 | 206 | var current float64 207 | if res != nil { 208 | current, err = bconv.ParseFloat(res, 64) 209 | if err != nil { 210 | return fmt.Errorf("hash value is not a valid float") 211 | } 212 | } 213 | increment, err := bconv.ParseFloat(args[2], 64) 214 | if err != nil { 215 | return fmt.Errorf("value is not a valid float") 216 | } 217 | result := strconv.AppendFloat(nil, current+increment, 'f', -1, 64) 218 | wb.Put(key, result) 219 | 220 | // if is a new key, increment the hash length 221 | if res == nil { 222 | setHlen(mk, length+1, wb) 223 | } 224 | return result 225 | } 226 | 227 | func Hgetall(args [][]byte, wb *levigo.WriteBatch) interface{} { 228 | return hgetall(args[0], true, true) 229 | } 230 | 231 | func Hkeys(args [][]byte, wb *levigo.WriteBatch) interface{} { 232 | return hgetall(args[0], true, false) 233 | } 234 | 235 | func Hvals(args [][]byte, wb *levigo.WriteBatch) interface{} { 236 | return hgetall(args[0], false, true) 237 | } 238 | 239 | func hgetall(key []byte, fields bool, values bool) interface{} { 240 | // use a snapshot so that the length is consistent with the iterator 241 | snapshot := DB.NewSnapshot() 242 | opts := levigo.NewReadOptions() 243 | opts.SetSnapshot(snapshot) 244 | 245 | length, err := hlen(metaKey(key), opts) 246 | if err != nil { 247 | return err 248 | DB.ReleaseSnapshot(snapshot) 249 | opts.Close() 250 | } 251 | if length == 0 { 252 | return []interface{}{} 253 | DB.ReleaseSnapshot(snapshot) 254 | opts.Close() 255 | } 256 | 257 | if fields && values { 258 | length *= 2 259 | } 260 | 261 | stream := &cmdReplyStream{int64(length), make(chan interface{})} 262 | go func() { 263 | defer close(stream.items) 264 | iterKey := NewKeyBuffer(HashKey, key, 0) 265 | it := DB.NewIterator(opts) 266 | for it.Seek(iterKey.Key()); it.Valid(); it.Next() { 267 | k := it.Key() 268 | if !iterKey.IsPrefixOf(k) { 269 | break 270 | } 271 | if fields { 272 | stream.items <- k[len(iterKey.Key()):] 273 | } 274 | if values { 275 | stream.items <- it.Value() 276 | } 277 | } 278 | DB.ReleaseSnapshot(snapshot) 279 | opts.Close() 280 | }() 281 | return stream 282 | } 283 | 284 | func DelHash(key []byte, wb *levigo.WriteBatch) { 285 | it := DB.NewIterator(ReadWithoutCacheFill) 286 | defer it.Close() 287 | iterKey := NewKeyBuffer(HashKey, key, 0) 288 | for it.Seek(iterKey.Key()); it.Valid(); it.Next() { 289 | k := it.Key() 290 | if !iterKey.IsPrefixOf(k) { 291 | break 292 | } 293 | wb.Delete(k) 294 | } 295 | } 296 | 297 | func hlen(key []byte, opts *levigo.ReadOptions) (uint32, error) { 298 | if opts == nil { 299 | opts = DefaultReadOptions 300 | } 301 | res, err := DB.Get(opts, key) 302 | if err != nil { 303 | return 0, err 304 | } 305 | if res == nil { 306 | return 0, nil 307 | } 308 | if len(res) > 0 && res[0] != HashLengthValue { 309 | return 0, InvalidKeyTypeError 310 | } 311 | if len(res) < 5 { 312 | return 0, InvalidDataError 313 | } 314 | return binary.BigEndian.Uint32(res[1:]), nil 315 | } 316 | 317 | func setHlen(key []byte, length uint32, wb *levigo.WriteBatch) { 318 | data := make([]byte, 5) 319 | data[0] = HashLengthValue 320 | binary.BigEndian.PutUint32(data[1:], length) 321 | wb.Put(key, data) 322 | } 323 | -------------------------------------------------------------------------------- /key_buffer.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "bytes" 5 | "encoding/binary" 6 | ) 7 | 8 | const keyPrefixSize = 5 9 | 10 | // KeyBuffer is a reusable key for LevelDB 11 | type KeyBuffer struct { 12 | buf []byte 13 | keyLen int 14 | reverseKey bool 15 | } 16 | 17 | // Create a new key of type t with key and room for extra bytes. 18 | func NewKeyBuffer(t byte, key []byte, extra int) *KeyBuffer { 19 | k := &KeyBuffer{make([]byte, keyPrefixSize, keyPrefixSize+len(key)+extra), len(key), false} 20 | k.buf[0] = t 21 | k.SetKey(key) 22 | return k 23 | } 24 | 25 | func NewKeyBufferWithSuffix(t byte, key []byte, suffix []byte) *KeyBuffer { 26 | k := &KeyBuffer{make([]byte, keyPrefixSize, keyPrefixSize+len(key)+len(suffix)), len(key), false} 27 | k.buf[0] = t 28 | k.SetKey(key) 29 | k.SetSuffix(suffix) 30 | return k 31 | } 32 | 33 | func (b *KeyBuffer) SetKey(key []byte) { 34 | if len(key) == 0 { 35 | return 36 | } 37 | b.keyLen = len(key) 38 | binary.BigEndian.PutUint32(b.buf[1:], uint32(len(key))) 39 | b.buf = append(b.buf[:keyPrefixSize], key...) 40 | } 41 | 42 | // Add extra after the key, will overwrite any existing extra 43 | func (b *KeyBuffer) SetSuffix(s []byte) { 44 | b.buf = append(b.buf[:keyPrefixSize+b.keyLen], s...) 45 | } 46 | 47 | // Return a slice of size n suitable for using with an io.Reader 48 | // The read bytes will overwrite the first n bytes of suffix (if they exist) 49 | func (b *KeyBuffer) SuffixForRead(n int) []byte { 50 | start := keyPrefixSize + b.keyLen 51 | if len(b.buf) < start+n { 52 | b.buf = append(b.buf[:start], make([]byte, n)...) // resize the slice to be large enough 53 | } 54 | return b.buf[start : start+n] 55 | } 56 | 57 | // Check if k starts with the key (without the suffix) 58 | func (b *KeyBuffer) IsPrefixOf(k []byte) bool { 59 | keyLen := keyPrefixSize + b.keyLen 60 | // the last byte is 0xff, so truncate a byte early 61 | if b.reverseKey { 62 | keyLen-- 63 | } 64 | if len(k) > keyLen && bytes.Equal(b.buf[:keyLen], k[:keyLen]) { 65 | return true 66 | } 67 | return false 68 | } 69 | 70 | // Change the key so that it sorts to come after the last prefix key 71 | // 72 | // To get a key that will sort *after* given prefix, we increment the last 73 | // byte that is not 0xff and truncated after the byte that was incremented 74 | func (b *KeyBuffer) ReverseIterKey() { 75 | b.reverseKey = true 76 | for i := len(b.buf) - 1; i >= 0; i-- { 77 | if b.buf[i] == 0xff { 78 | continue 79 | } 80 | b.buf[i] += 1 81 | b.buf = b.buf[:i+1] 82 | break 83 | } 84 | } 85 | 86 | func (b *KeyBuffer) Type() byte { 87 | return b.buf[0] 88 | } 89 | 90 | func (b *KeyBuffer) Key() []byte { 91 | return b.buf 92 | } 93 | -------------------------------------------------------------------------------- /list.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "encoding/binary" 5 | "math" 6 | 7 | "github.com/jmhodges/levigo" 8 | ) 9 | 10 | // Keys stored in LevelDB for lists 11 | // 12 | // MetaKey | key = ListLengthValue | uint32 list length | 1 byte flags | int64 sequence number of the leftmost element | int64 sequence number of the rightmost element 13 | // 14 | // For each list item: 15 | // ListKey | key length uint32 | key | int64 sequence number = value 16 | 17 | type listDetails struct { 18 | flags byte 19 | length uint32 20 | left int64 21 | right int64 22 | } 23 | 24 | const ( 25 | listLooseSeq byte = 1 << iota 26 | ) 27 | 28 | func Lrange(args [][]byte, wb *levigo.WriteBatch) interface{} { 29 | snapshot := DB.NewSnapshot() 30 | opts := levigo.NewReadOptions() 31 | opts.SetSnapshot(snapshot) 32 | 33 | l, err := llen(metaKey(args[0]), opts) 34 | if err != nil { 35 | DB.ReleaseSnapshot(snapshot) 36 | opts.Close() 37 | return err 38 | } 39 | if l.length == 0 { 40 | DB.ReleaseSnapshot(snapshot) 41 | opts.Close() 42 | return []interface{}{} 43 | } 44 | 45 | start, end, err := parseRange(args[1:], int64(l.length)) 46 | if err != nil { 47 | return err 48 | } 49 | // the start comes after the end, so we're not going to find anything 50 | if start > end { 51 | DB.ReleaseSnapshot(snapshot) 52 | opts.Close() 53 | return []interface{}{} 54 | } 55 | 56 | count := end + 1 - start 57 | stream := &cmdReplyStream{count, make(chan interface{})} 58 | 59 | go func() { 60 | defer close(stream.items) 61 | it := DB.NewIterator(opts) 62 | defer it.Close() 63 | 64 | iterKey := NewKeyBuffer(ListKey, args[0], 8) 65 | binary.BigEndian.PutUint64(iterKey.SuffixForRead(8), uint64(l.left+1+start-math.MinInt64)) 66 | it.Seek(iterKey.Key()) 67 | for i := int64(0); it.Valid() && i < count; i++ { 68 | stream.items <- it.Value() 69 | it.Next() 70 | } 71 | DB.ReleaseSnapshot(snapshot) 72 | opts.Close() 73 | }() 74 | return stream 75 | } 76 | 77 | func Llen(args [][]byte, wb *levigo.WriteBatch) interface{} { 78 | l, err := llen(metaKey(args[0]), nil) 79 | if err != nil { 80 | return err 81 | } 82 | return l.length 83 | } 84 | 85 | // A LPUSH onto a list takes the seq number of the leftmost element, 86 | // decrements it and inserts the item. 87 | func Lpush(args [][]byte, wb *levigo.WriteBatch) interface{} { 88 | res, err := lpush(args, true, true, wb) 89 | if err != nil { 90 | return err 91 | } 92 | return res 93 | } 94 | 95 | func Lpushx(args [][]byte, wb *levigo.WriteBatch) interface{} { 96 | res, err := lpush(args, true, false, wb) 97 | if err != nil { 98 | return err 99 | } 100 | return res 101 | } 102 | 103 | func Rpush(args [][]byte, wb *levigo.WriteBatch) interface{} { 104 | res, err := lpush(args, false, true, wb) 105 | if err != nil { 106 | return err 107 | } 108 | return res 109 | } 110 | 111 | func Rpushx(args [][]byte, wb *levigo.WriteBatch) interface{} { 112 | res, err := lpush(args, false, false, wb) 113 | if err != nil { 114 | return err 115 | } 116 | return res 117 | } 118 | 119 | func lpush(args [][]byte, left bool, create bool, wb *levigo.WriteBatch) (interface{}, error) { 120 | mk := metaKey(args[0]) 121 | l, err := llen(mk, nil) 122 | if err != nil { 123 | return nil, err 124 | } 125 | if create || l.length > 0 { 126 | key := NewKeyBuffer(ListKey, args[0], 8) 127 | for _, value := range args[1:] { 128 | l.length++ 129 | var seq int64 130 | if left { 131 | seq = l.left 132 | l.left-- 133 | } else { 134 | seq = l.right 135 | l.right++ 136 | } 137 | // To sort negative ints in order before positive, we subtract math.MinInt64 138 | // which wraps the numbers around and sorts correctly 139 | binary.BigEndian.PutUint64(key.SuffixForRead(8), uint64(seq-math.MinInt64)) 140 | wb.Put(key.Key(), value) 141 | } 142 | setLlen(mk, l, wb) 143 | } 144 | return l.length, nil 145 | } 146 | 147 | func Lpop(args [][]byte, wb *levigo.WriteBatch) interface{} { 148 | res, err := lpop(args[0], true, wb) 149 | if err != nil { 150 | return err 151 | } 152 | return res 153 | } 154 | 155 | func Rpop(args [][]byte, wb *levigo.WriteBatch) interface{} { 156 | res, err := lpop(args[0], false, wb) 157 | if err != nil { 158 | return err 159 | } 160 | return res 161 | } 162 | 163 | func Rpoplpush(args [][]byte, wb *levigo.WriteBatch) interface{} { 164 | res, err := lpop(args[0], false, wb) 165 | if err != nil { 166 | return err 167 | } 168 | if res == nil { 169 | return nil 170 | } 171 | _, err = lpush([][]byte{args[1], res.([]byte)}, true, true, wb) 172 | if err != nil { 173 | return err 174 | } 175 | return res 176 | } 177 | 178 | func lpop(key []byte, left bool, wb *levigo.WriteBatch) (interface{}, error) { 179 | mk := metaKey(key) 180 | l, err := llen(mk, nil) 181 | if err != nil { 182 | return nil, err 183 | } 184 | if l.length == 0 { 185 | return nil, nil 186 | } 187 | 188 | iterKey := NewKeyBuffer(ListKey, key, 0) 189 | it := DB.NewIterator(ReadWithoutCacheFill) 190 | defer it.Close() 191 | if !left { 192 | iterKey.ReverseIterKey() 193 | } 194 | it.Seek(iterKey.Key()) 195 | if !left { 196 | it.Prev() 197 | } 198 | if !it.Valid() { 199 | return nil, nil 200 | } 201 | k := it.Key() 202 | if !iterKey.IsPrefixOf(k) { 203 | return nil, nil 204 | } 205 | res := it.Value() 206 | 207 | wb.Delete(k) 208 | l.length-- 209 | if l.length == 0 { 210 | wb.Delete(mk) 211 | } else { 212 | // decode the sequence number from the list item key 213 | seq := int64(binary.BigEndian.Uint64(k[len(key)+5:])) + math.MinInt64 214 | if left { 215 | l.left = seq 216 | } else { 217 | l.right = seq 218 | } 219 | setLlen(mk, l, wb) 220 | } 221 | 222 | return res, nil 223 | } 224 | 225 | func llen(key []byte, opts *levigo.ReadOptions) (*listDetails, error) { 226 | if opts == nil { 227 | opts = DefaultReadOptions 228 | } 229 | res, err := DB.Get(opts, key) 230 | if err != nil { 231 | return nil, err 232 | } 233 | l := &listDetails{right: 1} 234 | if res == nil { 235 | return l, nil 236 | } 237 | if len(res) < 22 || res[0] != ListLengthValue { 238 | return nil, InvalidDataError 239 | } 240 | l.length = binary.BigEndian.Uint32(res[1:]) 241 | l.flags = res[5] 242 | l.left = int64(binary.BigEndian.Uint64(res[6:])) 243 | l.right = int64(binary.BigEndian.Uint64(res[14:])) 244 | return l, nil 245 | } 246 | 247 | func setLlen(key []byte, l *listDetails, wb *levigo.WriteBatch) { 248 | data := make([]byte, 22) 249 | data[0] = ListLengthValue 250 | binary.BigEndian.PutUint32(data[1:], l.length) 251 | data[5] = l.flags 252 | binary.BigEndian.PutUint64(data[6:], uint64(l.left)) 253 | binary.BigEndian.PutUint64(data[14:], uint64(l.right)) 254 | wb.Put(key, data) 255 | } 256 | 257 | // BLPOP 258 | // BRPOP 259 | // BRPOPLPUSH 260 | // LINDEX 261 | // LINSERT 262 | // LRANGE 263 | // LREM 264 | // LSET 265 | // LTRIM 266 | -------------------------------------------------------------------------------- /lockring/lockring.go: -------------------------------------------------------------------------------- 1 | package lockring 2 | 3 | import ( 4 | "hash/crc32" 5 | "sync" 6 | ) 7 | 8 | type LockRing struct { 9 | size uint32 10 | locks []sync.Mutex 11 | } 12 | 13 | func New(n uint32) *LockRing { 14 | return &LockRing{n, make([]sync.Mutex, n)} 15 | } 16 | 17 | func (l *LockRing) Lock(k []byte) { 18 | l.lockForKey(k).Lock() 19 | } 20 | 21 | func (l *LockRing) Unlock(k []byte) { 22 | l.lockForKey(k).Unlock() 23 | } 24 | 25 | func (l *LockRing) lockForKey(k []byte) *sync.Mutex { 26 | return &l.locks[int(crc32.ChecksumIEEE(k)%l.size)] 27 | } 28 | -------------------------------------------------------------------------------- /protocol.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "bufio" 5 | "bytes" 6 | "io" 7 | "net" 8 | "strconv" 9 | "sync" 10 | 11 | "github.com/jmhodges/levigo" 12 | "github.com/titanous/bconv" 13 | ) 14 | 15 | type client struct { 16 | cn net.Conn 17 | r *bufio.Reader 18 | w chan []byte 19 | 20 | writeQueueSize int // current queue size in bytes 21 | } 22 | 23 | var ( 24 | clients = make(map[string]*client) // ip:port -> client mapping 25 | clientsMtx = &sync.RWMutex{} 26 | ) 27 | 28 | func listen() { 29 | l, err := net.Listen("tcp", ":12345") 30 | maybeFatal(err) 31 | for { 32 | conn, err := l.Accept() 33 | if err != nil { 34 | // TODO: log error 35 | continue 36 | } 37 | go handleClient(conn) 38 | } 39 | } 40 | 41 | func handleClient(cn net.Conn) { 42 | c := &client{cn: cn, r: bufio.NewReader(cn), w: make(chan []byte)} 43 | defer close(c.w) 44 | 45 | addr := cn.RemoteAddr().String() 46 | clientsMtx.Lock() 47 | clients[addr] = c 48 | clientsMtx.Unlock() 49 | defer func() { 50 | clientsMtx.Lock() 51 | delete(clients, addr) 52 | clientsMtx.Unlock() 53 | }() 54 | 55 | o := make(chan []byte) 56 | go responseQueue(c, o) 57 | go responseWriter(c, o) 58 | 59 | protocolHandler(c) 60 | } 61 | 62 | func protocolHandler(c *client) { 63 | // Read a length (looks like "$3\r\n") 64 | readLength := func(prefix byte) (length int, err error) { 65 | b, err := c.r.ReadByte() 66 | if err != nil { 67 | return 68 | } 69 | if b != prefix { 70 | writeProtocolError(c.w, "invalid length") 71 | return 72 | } 73 | l, overflowed, err := c.r.ReadLine() // Read bytes will look like "123" 74 | if err != nil { 75 | return 76 | } 77 | if overflowed { 78 | writeProtocolError(c.w, "length line too long") 79 | return 80 | } 81 | if len(l) == 0 { 82 | writeProtocolError(c.w, "missing length") 83 | return 84 | } 85 | length, err = bconv.Atoi(l) 86 | if err != nil { 87 | writeProtocolError(c.w, "length is not a valid integer") 88 | return 89 | } 90 | return 91 | } 92 | 93 | runCommand := func(args [][]byte) (err error) { 94 | if len(args) == 0 { 95 | writeProtocolError(c.w, "missing command") 96 | return 97 | } 98 | 99 | // lookup the command 100 | command, ok := commands[UnsafeBytesToString(bytes.ToLower(args[0]))] 101 | if !ok { 102 | writeError(c.w, "unknown command '"+string(args[0])+"'") 103 | return 104 | } 105 | 106 | // check command arity, negative arity means >= n 107 | if (command.arity < 0 && len(args)-1 < -command.arity) || (command.arity >= 0 && len(args)-1 < command.arity) { 108 | writeError(c.w, "wrong number of arguments for '"+string(args[0])+"' command") 109 | return 110 | } 111 | 112 | // call the command and respond 113 | var wb *levigo.WriteBatch 114 | if command.writes { 115 | wb = levigo.NewWriteBatch() 116 | defer wb.Close() 117 | } 118 | command.lockKeys(args[1:]) 119 | res := command.function(args[1:], wb) 120 | if command.writes { 121 | if _, ok := res.(error); !ok { // only write the batch if the return value is not an error 122 | err = DB.Write(DefaultWriteOptions, wb) 123 | } 124 | if err != nil { 125 | writeError(c.w, "data write error: "+err.Error()) 126 | return 127 | } 128 | } 129 | command.unlockKeys(args[1:]) 130 | writeReply(c.w, res) 131 | 132 | return 133 | } 134 | 135 | processInline := func() error { 136 | line, err := c.r.ReadBytes('\n') 137 | if err != nil { 138 | return err 139 | } 140 | return runCommand(bytes.Split(line[:len(line)-2], []byte(" "))) 141 | } 142 | 143 | scratch := make([]byte, 2) 144 | args := [][]byte{} 145 | // Client event loop, each iteration handles a command 146 | for { 147 | // check if we're using the old inline protocol 148 | b, err := c.r.Peek(1) 149 | if err != nil { 150 | return 151 | } 152 | if b[0] != '*' { 153 | err = processInline() 154 | if err != nil { 155 | return 156 | } 157 | continue 158 | } 159 | 160 | // Step 1: get the number of arguments 161 | argCount, err := readLength('*') 162 | if err != nil { 163 | return 164 | } 165 | 166 | // read the arguments 167 | for i := 0; i < argCount; i++ { 168 | length, err := readLength('$') 169 | if err != nil { 170 | return 171 | } 172 | 173 | // Read the argument bytes 174 | args = append(args, make([]byte, length)) 175 | _, err = io.ReadFull(c.r, args[i]) 176 | if err != nil { 177 | return 178 | } 179 | 180 | // The argument has a trailing \r\n that we need to discard 181 | c.r.Read(scratch) // TODO: make sure these bytes are read 182 | } 183 | 184 | err = runCommand(args) 185 | if err != nil { 186 | return 187 | } 188 | 189 | // Truncate arguments for the next run 190 | args = args[:0] 191 | } 192 | } 193 | 194 | func responseQueue(c *client, out chan<- []byte) { 195 | defer close(out) 196 | 197 | queue := [][]byte{} 198 | 199 | receive: 200 | for { 201 | // ensure that the queue always has an item in it 202 | if len(queue) == 0 { 203 | v, ok := <-c.w 204 | if !ok { 205 | break // in is closed, we're done 206 | } 207 | queue = append(queue, v) 208 | c.writeQueueSize += len(v) 209 | } 210 | 211 | select { 212 | case out <- queue[0]: 213 | c.writeQueueSize -= len(queue[0]) 214 | queue = queue[1:] 215 | case v, ok := <-c.w: 216 | if !ok { 217 | break receive // in is closed, we're done 218 | } 219 | queue = append(queue, v) 220 | c.writeQueueSize += len(v) 221 | } 222 | } 223 | 224 | for _, v := range queue { 225 | out <- v 226 | } 227 | } 228 | 229 | func responseWriter(c *client, out <-chan []byte) { 230 | for v := range out { 231 | c.cn.Write(v) 232 | } 233 | } 234 | 235 | func writeReply(w chan<- []byte, reply interface{}) { 236 | if _, ok := reply.([]interface{}); !ok && reply == nil { 237 | w <- []byte("$-1\r\n") 238 | return 239 | } 240 | switch r := reply.(type) { 241 | case rawReply: 242 | w <- r 243 | case string: 244 | w <- []byte("+" + r + "\r\n") 245 | case []byte: 246 | writeBulk(w, r) 247 | case int: 248 | writeInt(w, int64(r)) 249 | case int64: 250 | writeInt(w, r) 251 | case uint32: 252 | writeInt(w, int64(r)) 253 | case IOError: 254 | w <- []byte("-IOERR " + reply.(IOError).Error() + "\r\n") 255 | case error: 256 | writeError(w, r.Error()) 257 | case []interface{}: 258 | writeMultibulk(w, r) 259 | case *cmdReplyStream: 260 | writeMultibulkStream(w, r) 261 | case map[string]bool: 262 | writeMultibulkStringMap(w, r) 263 | default: 264 | panic("Invalid reply type") 265 | } 266 | } 267 | 268 | func writeProtocolError(w chan<- []byte, msg string) { 269 | writeError(w, "Protocol error: "+msg) 270 | } 271 | 272 | func writeInt(w chan<- []byte, n int64) { 273 | w <- append(strconv.AppendInt([]byte{':'}, n, 10), "\r\n"...) 274 | } 275 | 276 | func writeBulk(w chan<- []byte, b []byte) { 277 | if b == nil { 278 | w <- []byte("$-1\r\n") 279 | } 280 | // TODO: find a more efficient way of doing this 281 | w <- append(strconv.AppendInt([]byte{'$'}, int64(len(b)), 10), "\r\n"...) 282 | w <- b 283 | w <- []byte("\r\n") 284 | } 285 | 286 | func writeMultibulkStream(w chan<- []byte, reply *cmdReplyStream) { 287 | writeMultibulkLength(w, reply.size) 288 | for r := range reply.items { 289 | writeReply(w, r) 290 | } 291 | } 292 | 293 | func writeMultibulk(w chan<- []byte, reply []interface{}) { 294 | if reply == nil { 295 | writeMultibulkLength(w, -1) 296 | } 297 | writeMultibulkLength(w, int64(len(reply))) 298 | for _, r := range reply { 299 | writeReply(w, r) 300 | } 301 | } 302 | 303 | func writeMultibulkStringMap(w chan<- []byte, reply map[string]bool) { 304 | writeMultibulkLength(w, int64(len(reply))) 305 | for r, _ := range reply { 306 | writeBulk(w, []byte(r)) 307 | } 308 | } 309 | 310 | func writeMultibulkLength(w chan<- []byte, n int64) { 311 | w <- append(strconv.AppendInt([]byte{'*'}, n, 10), "\r\n"...) 312 | } 313 | 314 | func writeError(w chan<- []byte, msg string) { 315 | w <- []byte("-ERR " + msg + "\r\n") 316 | } 317 | -------------------------------------------------------------------------------- /protocol_test.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "net" 5 | "testing" 6 | 7 | . "launchpad.net/gocheck" 8 | ) 9 | 10 | type ProtocolSuite struct{} 11 | 12 | var _ = Suite(&ProtocolSuite{}) 13 | 14 | func (s ProtocolSuite) TestPing(c *C) { 15 | a, b := net.Pipe() 16 | defer a.Close() 17 | go handleClient(b) 18 | 19 | a.Write([]byte("*1\r\n$4\r\nPING\r\n")) 20 | res := make([]byte, 7) 21 | a.Read(res) 22 | c.Assert(res, DeepEquals, []byte("+PONG\r\n")) 23 | } 24 | 25 | func (s ProtocolSuite) TestArity(c *C) { 26 | a, b := net.Pipe() 27 | defer a.Close() 28 | go handleClient(b) 29 | 30 | tests := []struct { 31 | cmd string 32 | expected string 33 | }{ 34 | // invalid positive arity 35 | { 36 | "*1\r\n$6\r\nLRANGE\r\n", 37 | "-ERR wrong number of arguments for 'LRANGE' command\r\n", 38 | }, 39 | // valid positive arity 40 | { 41 | "*4\r\n$6\r\nLRANGE\r\n$3\r\nfoo\r\n$1\r\n0\r\n$2\r\n-1\r\n", 42 | "*0\r\n", 43 | }, 44 | // invalid negative arity 45 | { 46 | "*1\r\n$5\r\nLPUSH\r\n", 47 | "-ERR wrong number of arguments for 'LPUSH' command\r\n", 48 | }, 49 | // valid negative arity 50 | { 51 | "*3\r\n$5\r\nLPUSH\r\n$3\r\nfoo\r\n$1\r\nA\r\n", 52 | ":1\r\n", 53 | }, 54 | } 55 | 56 | var res []byte 57 | for _, t := range tests { 58 | a.Write([]byte(t.cmd)) 59 | res = make([]byte, len(t.expected)) 60 | a.Read(res) 61 | c.Assert(res, DeepEquals, []byte(t.expected)) 62 | } 63 | } 64 | 65 | func BenchmarkProtocolParserSimple(b *testing.B) { 66 | b.StopTimer() 67 | client, server := net.Pipe() 68 | defer client.Close() 69 | go handleClient(server) 70 | 71 | // eat anything that gets written to the pipe 72 | res := make([]byte, 1024) 73 | go func() { 74 | for { 75 | _, err := client.Read(res) 76 | if err != nil { 77 | break 78 | } 79 | } 80 | }() 81 | 82 | b.StartTimer() 83 | 84 | for i := 0; i < b.N; i++ { 85 | client.Write([]byte("*1\r\n$4\r\nPING\r\n")) 86 | } 87 | } 88 | 89 | func BenchmarkProtocolParserInline(b *testing.B) { 90 | b.StopTimer() 91 | client, server := net.Pipe() 92 | defer client.Close() 93 | go handleClient(server) 94 | 95 | // eat anything that gets written to the pipe 96 | res := make([]byte, 1024) 97 | go func() { 98 | for { 99 | _, err := client.Read(res) 100 | if err != nil { 101 | break 102 | } 103 | } 104 | }() 105 | 106 | b.StartTimer() 107 | 108 | for i := 0; i < b.N; i++ { 109 | client.Write([]byte("PING\r\n")) 110 | } 111 | } 112 | -------------------------------------------------------------------------------- /rdb.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "encoding/binary" 5 | "math" 6 | 7 | "github.com/cupcake/rdb" 8 | "github.com/cupcake/rdb/nopdecoder" 9 | "github.com/jmhodges/levigo" 10 | ) 11 | 12 | type rdbDecoder struct { 13 | wb *levigo.WriteBatch 14 | i int64 15 | nopdecoder.NopDecoder 16 | } 17 | 18 | func (p *rdbDecoder) Set(key, value []byte, expiry int64) { 19 | Del([][]byte{key}, p.wb) 20 | setStringLen(metaKey(key), len(value), p.wb) 21 | p.wb.Put(stringKey(key), value) 22 | } 23 | 24 | func (p *rdbDecoder) StartHash(key []byte, length, expiry int64) { 25 | Del([][]byte{key}, p.wb) 26 | setHlen(metaKey(key), uint32(length), p.wb) 27 | } 28 | 29 | func (p *rdbDecoder) Hset(key, field, value []byte) { 30 | p.wb.Put(NewKeyBufferWithSuffix(HashKey, key, field).Key(), value) 31 | } 32 | 33 | func (p *rdbDecoder) StartSet(key []byte, cardinality, expiry int64) { 34 | Del([][]byte{key}, p.wb) 35 | setCard(metaKey(key), uint32(cardinality), p.wb) 36 | } 37 | 38 | func (p *rdbDecoder) Sadd(key, member []byte) { 39 | p.wb.Put(NewKeyBufferWithSuffix(SetKey, key, member).Key(), []byte{}) 40 | } 41 | 42 | func (p *rdbDecoder) StartList(key []byte, length, expiry int64) { 43 | p.i = 0 44 | Del([][]byte{key}, p.wb) 45 | setLlen(metaKey(key), &listDetails{length: uint32(length), right: length + 2}, p.wb) 46 | } 47 | 48 | func (p *rdbDecoder) Rpush(key, value []byte) { 49 | p.i++ 50 | k := NewKeyBuffer(ListKey, key, 8) 51 | binary.BigEndian.PutUint64(k.SuffixForRead(8), uint64(p.i-math.MinInt64)) 52 | p.wb.Put(k.Key(), value) 53 | } 54 | 55 | func (p *rdbDecoder) StartZSet(key []byte, cardinality, expiry int64) { 56 | Del([][]byte{key}, p.wb) 57 | setZcard(metaKey(key), uint32(cardinality), p.wb) 58 | } 59 | 60 | func (p *rdbDecoder) Zadd(key []byte, score float64, member []byte) { 61 | scoreBytes := make([]byte, 8) 62 | binary.BigEndian.PutUint64(scoreBytes, math.Float64bits(score)) 63 | scoreKey := NewKeyBuffer(ZScoreKey, key, len(member)+8) 64 | setZScoreKeyMember(scoreKey, member) 65 | setZScoreKeyScore(scoreKey, score) 66 | p.wb.Put(NewKeyBufferWithSuffix(ZSetKey, key, member).Key(), scoreBytes) 67 | p.wb.Put(scoreKey.Key(), []byte{}) 68 | } 69 | 70 | type rdbEncoder struct { 71 | r *rdb.Encoder 72 | } 73 | 74 | func (e *rdbEncoder) encodeKey(key []byte, dump bool) error { 75 | snapshot := DB.NewSnapshot() 76 | opts := levigo.NewReadOptions() 77 | opts.SetSnapshot(snapshot) 78 | defer DB.ReleaseSnapshot(snapshot) 79 | defer opts.Close() 80 | 81 | res, err := DB.Get(opts, metaKey(key)) 82 | if err != nil { 83 | return err 84 | } 85 | if res == nil { 86 | return nil 87 | } 88 | if len(res) < 5 { 89 | return InvalidDataError 90 | } 91 | 92 | length := binary.BigEndian.Uint32(res[1:]) 93 | switch res[0] { 94 | case StringLengthValue: 95 | e.r.EncodeType(rdb.TypeString) 96 | case HashLengthValue: 97 | e.r.EncodeType(rdb.TypeHash) 98 | case SetCardValue: 99 | e.r.EncodeType(rdb.TypeSet) 100 | case ZCardValue: 101 | e.r.EncodeType(rdb.TypeZSet) 102 | case ListLengthValue: 103 | e.r.EncodeType(rdb.TypeList) 104 | default: 105 | panic("unknown key type") 106 | } 107 | 108 | if !dump { 109 | e.r.EncodeString(key) 110 | } 111 | 112 | switch res[0] { 113 | case StringLengthValue: 114 | e.encodeString(key, opts) 115 | case HashLengthValue: 116 | e.encodeHash(key, length, opts) 117 | case SetCardValue: 118 | e.encodeSet(key, length, opts) 119 | case ZCardValue: 120 | e.encodeZSet(key, length, opts) 121 | case ListLengthValue: 122 | e.encodeList(key, length, opts) 123 | } 124 | 125 | if dump { 126 | e.r.EncodeDumpFooter() 127 | } 128 | 129 | return nil 130 | } 131 | 132 | func (e *rdbEncoder) encodeString(key []byte, opts *levigo.ReadOptions) error { 133 | res, err := DB.Get(opts, stringKey(key)) 134 | if err != nil { 135 | return err 136 | } 137 | return e.r.EncodeString(res) 138 | } 139 | 140 | func (e *rdbEncoder) encodeHash(key []byte, length uint32, opts *levigo.ReadOptions) error { 141 | err := e.r.EncodeLength(length) 142 | if err != nil { 143 | return err 144 | } 145 | 146 | iterKey := NewKeyBuffer(HashKey, key, 0) 147 | it := DB.NewIterator(opts) 148 | defer it.Close() 149 | 150 | for it.Seek(iterKey.Key()); it.Valid(); it.Next() { 151 | k := it.Key() 152 | if !iterKey.IsPrefixOf(k) { 153 | break 154 | } 155 | err = e.r.EncodeString(k[len(iterKey.Key()):]) 156 | if err != nil { 157 | return err 158 | } 159 | err = e.r.EncodeString(it.Value()) 160 | if err != nil { 161 | return err 162 | } 163 | } 164 | 165 | return nil 166 | } 167 | 168 | func (e *rdbEncoder) encodeList(key []byte, length uint32, opts *levigo.ReadOptions) error { 169 | err := e.r.EncodeLength(length) 170 | if err != nil { 171 | return err 172 | } 173 | 174 | iterKey := NewKeyBuffer(ListKey, key, 0) 175 | it := DB.NewIterator(opts) 176 | defer it.Close() 177 | 178 | for it.Seek(iterKey.Key()); it.Valid(); it.Next() { 179 | if !iterKey.IsPrefixOf(it.Key()) { 180 | break 181 | } 182 | err = e.r.EncodeString(it.Value()) 183 | if err != nil { 184 | return err 185 | } 186 | } 187 | 188 | return nil 189 | } 190 | 191 | func (e *rdbEncoder) encodeSet(key []byte, cardinality uint32, opts *levigo.ReadOptions) error { 192 | err := e.r.EncodeLength(cardinality) 193 | if err != nil { 194 | return err 195 | } 196 | 197 | iterKey := NewKeyBuffer(SetKey, key, 0) 198 | it := DB.NewIterator(opts) 199 | defer it.Close() 200 | 201 | for it.Seek(iterKey.Key()); it.Valid(); it.Next() { 202 | k := it.Key() 203 | if !iterKey.IsPrefixOf(k) { 204 | break 205 | } 206 | err = e.r.EncodeString(k[len(iterKey.Key()):]) 207 | if err != nil { 208 | return err 209 | } 210 | } 211 | 212 | return nil 213 | } 214 | 215 | func (e *rdbEncoder) encodeZSet(key []byte, cardinality uint32, opts *levigo.ReadOptions) error { 216 | err := e.r.EncodeLength(cardinality) 217 | if err != nil { 218 | return err 219 | } 220 | 221 | iterKey := NewKeyBuffer(ZScoreKey, key, 0) 222 | iterKey.ReverseIterKey() 223 | it := DB.NewIterator(opts) 224 | defer it.Close() 225 | it.Seek(iterKey.Key()) 226 | 227 | for it.Prev(); it.Valid(); it.Prev() { 228 | k := it.Key() 229 | if !iterKey.IsPrefixOf(k) { 230 | break 231 | } 232 | score, member := parseZScoreKey(k, len(iterKey.Key())-5) 233 | err = e.r.EncodeString(member) 234 | if err != nil { 235 | return err 236 | } 237 | err = e.r.EncodeFloat(score) 238 | if err != nil { 239 | return err 240 | } 241 | } 242 | 243 | return nil 244 | } 245 | -------------------------------------------------------------------------------- /set.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "bytes" 5 | "crypto/rand" 6 | "encoding/binary" 7 | "sort" 8 | 9 | "github.com/jmhodges/levigo" 10 | ) 11 | 12 | // Keys stored in LevelDB for sets 13 | // 14 | // MetadataKey | key = SetCardValue | count of members uint32 15 | // 16 | // For each member: 17 | // SetKey | key length uint32 | key | member = empty 18 | 19 | func Sadd(args [][]byte, wb *levigo.WriteBatch) interface{} { 20 | var newMembers uint32 21 | key := NewKeyBuffer(SetKey, args[0], len(args[1])) 22 | mk := metaKey(args[0]) 23 | card, err := scard(mk, nil) 24 | if err != nil { 25 | return err 26 | } 27 | 28 | for _, member := range args[1:] { 29 | key.SetSuffix(member) 30 | if card > 0 { 31 | res, err := DB.Get(DefaultReadOptions, key.Key()) 32 | if err != nil { 33 | return err 34 | } 35 | if res != nil { 36 | continue 37 | } 38 | } 39 | wb.Put(key.Key(), []byte{}) 40 | newMembers++ 41 | } 42 | if newMembers > 0 { 43 | setCard(mk, card+newMembers, wb) 44 | } 45 | return newMembers 46 | } 47 | 48 | func Scard(args [][]byte, wb *levigo.WriteBatch) interface{} { 49 | card, err := scard(metaKey(args[0]), nil) 50 | if err != nil { 51 | return err 52 | } 53 | return card 54 | } 55 | 56 | func Srem(args [][]byte, wb *levigo.WriteBatch) interface{} { 57 | mk := metaKey(args[0]) 58 | card, err := scard(mk, nil) 59 | if err != nil { 60 | return err 61 | } 62 | if card == 0 { 63 | return card 64 | } 65 | var deleted uint32 66 | key := NewKeyBuffer(SetKey, args[0], len(args[1])) 67 | for _, member := range args[1:] { 68 | key.SetSuffix(member) 69 | res, err := DB.Get(ReadWithoutCacheFill, key.Key()) 70 | if err != nil { 71 | return err 72 | } 73 | if res == nil { 74 | continue 75 | } 76 | wb.Delete(key.Key()) 77 | deleted++ 78 | } 79 | if deleted == card { 80 | wb.Delete(mk) 81 | } else if deleted > 0 { // decrement the cardinality 82 | setCard(mk, card-deleted, wb) 83 | } 84 | return deleted 85 | } 86 | 87 | func Sismember(args [][]byte, wb *levigo.WriteBatch) interface{} { 88 | res, err := DB.Get(DefaultReadOptions, NewKeyBufferWithSuffix(SetKey, args[0], args[1]).Key()) 89 | if err != nil { 90 | return err 91 | } 92 | if res == nil { 93 | return 0 94 | } 95 | return 1 96 | } 97 | 98 | func Smembers(args [][]byte, wb *levigo.WriteBatch) interface{} { 99 | // use a snapshot so that the cardinality is consistent with the iterator 100 | snapshot := DB.NewSnapshot() 101 | opts := levigo.NewReadOptions() 102 | opts.SetSnapshot(snapshot) 103 | 104 | card, err := scard(metaKey(args[0]), opts) 105 | if err != nil { 106 | return err 107 | DB.ReleaseSnapshot(snapshot) 108 | opts.Close() 109 | } 110 | if card == 0 { 111 | return []interface{}{} 112 | DB.ReleaseSnapshot(snapshot) 113 | opts.Close() 114 | } 115 | 116 | // send the reply back over a channel since there could be a lot of items 117 | stream := &cmdReplyStream{int64(card), make(chan interface{})} 118 | go func() { 119 | defer close(stream.items) 120 | it := DB.NewIterator(opts) 121 | defer it.Close() 122 | iterKey := NewKeyBuffer(SetKey, args[0], 0) 123 | 124 | for it.Seek(iterKey.Key()); it.Valid(); it.Next() { 125 | // If the prefix of the current key doesn't match the iteration key, 126 | // we have reached the end of the set 127 | key := it.Key() 128 | if !iterKey.IsPrefixOf(key) { 129 | break 130 | } 131 | stream.items <- parseMemberFromSetKey(key) 132 | } 133 | DB.ReleaseSnapshot(snapshot) 134 | opts.Close() 135 | }() 136 | return stream 137 | } 138 | 139 | func Spop(args [][]byte, wb *levigo.WriteBatch) interface{} { 140 | mk := metaKey(args[0]) 141 | card, err := scard(mk, nil) 142 | if err != nil { 143 | return err 144 | } 145 | if card == 0 { 146 | return nil 147 | } 148 | key := NewKeyBuffer(SetKey, args[0], 1) 149 | member := srand(key) 150 | if member == nil { 151 | return nil 152 | } 153 | key.SetSuffix(member) 154 | wb.Delete(key.Key()) 155 | if card == 1 { // we're removing the last remaining member 156 | wb.Delete(mk) 157 | } else { 158 | setCard(mk, card-1, wb) 159 | } 160 | return member 161 | } 162 | 163 | func Smove(args [][]byte, wb *levigo.WriteBatch) interface{} { 164 | resp, err := DB.Get(DefaultReadOptions, NewKeyBufferWithSuffix(SetKey, args[0], args[2]).Key()) 165 | if err != nil { 166 | return err 167 | } 168 | if resp == nil { 169 | return 0 170 | } 171 | 172 | res := Srem([][]byte{args[0], args[2]}, wb) 173 | if err, ok := res.(error); ok { 174 | return err 175 | } 176 | res = Sadd([][]byte{args[1], args[2]}, wb) 177 | if err, ok := res.(error); ok { 178 | return err 179 | } 180 | return 1 181 | } 182 | 183 | const ( 184 | setUnion int = iota 185 | setInter 186 | setDiff 187 | ) 188 | 189 | func Sunion(args [][]byte, wb *levigo.WriteBatch) interface{} { 190 | return combineSet(args, setUnion, nil) 191 | } 192 | 193 | func Sinter(args [][]byte, wb *levigo.WriteBatch) interface{} { 194 | return combineSet(args, setInter, nil) 195 | } 196 | 197 | func Sdiff(args [][]byte, wb *levigo.WriteBatch) interface{} { 198 | return combineSet(args, setDiff, nil) 199 | } 200 | 201 | func Sunionstore(args [][]byte, wb *levigo.WriteBatch) interface{} { 202 | return combineSet(args, setUnion, wb) 203 | } 204 | 205 | func Sinterstore(args [][]byte, wb *levigo.WriteBatch) interface{} { 206 | return combineSet(args, setInter, wb) 207 | } 208 | 209 | func Sdiffstore(args [][]byte, wb *levigo.WriteBatch) interface{} { 210 | return combineSet(args, setDiff, wb) 211 | } 212 | 213 | func combineSet(keys [][]byte, op int, wb *levigo.WriteBatch) interface{} { 214 | var count uint32 215 | res := []interface{}{} 216 | members := make(chan *iterSetMember) 217 | var storeKey *KeyBuffer 218 | var mk []byte 219 | 220 | if wb != nil { 221 | mk = metaKey(keys[0]) 222 | _, err := delKey(mk, wb) 223 | if err != nil { 224 | return err 225 | } 226 | storeKey = NewKeyBuffer(SetKey, keys[0], 0) 227 | keys = keys[1:] 228 | } 229 | 230 | go multiSetIter(keys, members, op != setUnion) 231 | 232 | combine: 233 | for m := range members { 234 | switch op { 235 | case setInter: 236 | for _, k := range m.exists { 237 | if !k { 238 | continue combine 239 | } 240 | } 241 | case setDiff: 242 | for i, k := range m.exists { 243 | if i == 0 && !k || i > 0 && k { 244 | continue combine 245 | } 246 | } 247 | } 248 | if wb != nil { 249 | storeKey.SetSuffix(m.member) 250 | wb.Put(storeKey.Key(), []byte{}) 251 | count++ 252 | } else { 253 | res = append(res, m.member) 254 | } 255 | } 256 | 257 | if wb != nil { 258 | if count > 0 { 259 | setCard(mk, count, wb) 260 | } 261 | return count 262 | } 263 | 264 | return res 265 | } 266 | 267 | type iterSetMember struct { 268 | member []byte 269 | exists []bool 270 | } 271 | 272 | type setMember struct { 273 | member []byte 274 | key int 275 | } 276 | 277 | type setMembers []*setMember 278 | 279 | func (m setMembers) Len() int { return len(m) } 280 | func (m setMembers) Less(i, j int) bool { return bytes.Compare(m[i].member, m[j].member) == -1 } 281 | func (m setMembers) Swap(i, j int) { m[i], m[j] = m[j], m[i] } 282 | 283 | // Take a list of keys and send a member to out for each unique member with a list 284 | // of the keys that have that member. 285 | // An iterator for each key is created (members are lexographically sorted), and 286 | // the first member from each key is added to a list. This list is sorted, so 287 | // that the member that would iterate first is at the beginnning of the list. 288 | // The member list is then checked for any other keys that have the same member. 289 | // The first member is sent to out, and all keys that had that member are iterated 290 | // forward. This is repeated until all keys have run out of members. 291 | func multiSetIter(keys [][]byte, out chan<- *iterSetMember, stopEarly bool) { 292 | defer close(out) 293 | // Set up a snapshot so that we have a consistent view of the data 294 | snapshot := DB.NewSnapshot() 295 | opts := levigo.NewReadOptions() 296 | defer opts.Close() 297 | opts.SetSnapshot(snapshot) 298 | defer DB.ReleaseSnapshot(snapshot) 299 | 300 | members := make(setMembers, len(keys)) // a list of the current member for each key iterator 301 | iterKeys := make([]*KeyBuffer, len(keys)) 302 | iterators := make([]*levigo.Iterator, len(keys)) 303 | for i, k := range keys { 304 | iterKeys[i] = NewKeyBuffer(SetKey, k, 0) 305 | it := DB.NewIterator(opts) 306 | defer it.Close() 307 | it.Seek(iterKeys[i].Key()) 308 | iterators[i] = it 309 | } 310 | 311 | getMember := func(i int) []byte { 312 | // If the iterator is done, we remove the iterator and ignore it in future runs 313 | if iterators[i] == nil || !iterators[i].Valid() { 314 | iterators[i] = nil 315 | return nil 316 | } 317 | k := iterators[i].Key() 318 | if !iterKeys[i].IsPrefixOf(k) { 319 | iterators[i] = nil 320 | return nil 321 | } 322 | // Strip the key prefix from the key and return the member 323 | return k[len(iterKeys[i].Key()):] 324 | } 325 | 326 | // Initialize the members list 327 | for i := 0; i < len(members); i++ { 328 | members[i] = &setMember{getMember(i), i} 329 | } 330 | 331 | // This loop runs until we run out of keys 332 | iter: 333 | for { 334 | im := &iterSetMember{exists: make([]bool, len(members))} 335 | first := true 336 | sort.Sort(members) 337 | 338 | for _, m := range members { 339 | // The member will be nil if the key it is from has no more members 340 | if m.member == nil { 341 | if m.key == 0 && stopEarly { 342 | break iter 343 | } 344 | continue 345 | } 346 | // The first member is the one that we will send out on this iteration 347 | if first { 348 | im.member = m.member 349 | first = false 350 | } 351 | if first || bytes.Compare(im.member, m.member) == 0 { 352 | im.exists[m.key] = true 353 | iterators[m.key].Next() 354 | m.member = getMember(m.key) 355 | } else { 356 | // If the member isn't first or the same as the one we are 357 | // looking for, it's not in the list 358 | break 359 | } 360 | } 361 | // When the result member is nil, there are no members left in any of the sets 362 | if im.member == nil { 363 | break 364 | } 365 | out <- im 366 | } 367 | } 368 | 369 | func DelSet(key []byte, wb *levigo.WriteBatch) { 370 | it := DB.NewIterator(ReadWithoutCacheFill) 371 | defer it.Close() 372 | iterKey := NewKeyBuffer(SetKey, key, 0) 373 | for it.Seek(iterKey.Key()); it.Valid(); it.Next() { 374 | k := it.Key() 375 | // If the prefix of the current key doesn't match the iteration key, 376 | // we have reached the end of the set 377 | if !iterKey.IsPrefixOf(k) { 378 | break 379 | } 380 | wb.Delete(k) 381 | } 382 | } 383 | 384 | func scard(key []byte, opts *levigo.ReadOptions) (uint32, error) { 385 | if opts == nil { 386 | opts = DefaultReadOptions 387 | } 388 | res, err := DB.Get(opts, key) 389 | if err != nil { 390 | return 0, err 391 | } 392 | if res == nil { 393 | return 0, nil 394 | } 395 | if len(res) > 0 && res[0] != SetCardValue { 396 | return 0, InvalidKeyTypeError 397 | } 398 | if len(res) < 5 { 399 | return 0, InvalidDataError 400 | } 401 | return binary.BigEndian.Uint32(res[1:]), nil 402 | } 403 | 404 | func srand(key *KeyBuffer) []byte { 405 | it := DB.NewIterator(DefaultReadOptions) 406 | defer it.Close() 407 | rand.Read(key.SuffixForRead(1)) 408 | it.Seek(key.Key()) 409 | if !it.Valid() { 410 | return nil 411 | } 412 | k := it.Key() 413 | // check if we are in the set 414 | // if we aren't it's possible that we ended up at the end, so go back a key 415 | if !key.IsPrefixOf(k) { 416 | it.Prev() 417 | k = it.Key() 418 | if !key.IsPrefixOf(k) { 419 | return nil 420 | } 421 | } 422 | return parseMemberFromSetKey(k) 423 | } 424 | 425 | func setCard(key []byte, card uint32, wb *levigo.WriteBatch) { 426 | data := make([]byte, 5) 427 | data[0] = SetCardValue 428 | binary.BigEndian.PutUint32(data[1:], card) 429 | wb.Put(key, data) 430 | } 431 | 432 | func parseMemberFromSetKey(key []byte) []byte { 433 | keyLen := binary.BigEndian.Uint32(key[1:]) 434 | return key[5+int(keyLen):] 435 | } 436 | 437 | // SRANDMEMBER 438 | -------------------------------------------------------------------------------- /string.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "encoding/binary" 5 | 6 | "github.com/jmhodges/levigo" 7 | ) 8 | 9 | // Keys stored in LevelDB for strings 10 | // 11 | // MetadataKey | key = StringLengthValue | string length uint32 12 | // 13 | // For each key: 14 | // StringKey | key = value 15 | 16 | func Set(args [][]byte, wb *levigo.WriteBatch) interface{} { 17 | err := set(args[0], args[1], wb) 18 | if err != nil { 19 | return err 20 | } 21 | return ReplyOK 22 | } 23 | 24 | func Get(args [][]byte, wb *levigo.WriteBatch) interface{} { 25 | res, err := DB.Get(DefaultReadOptions, stringKey(args[0])) 26 | if err != nil { 27 | return err 28 | } 29 | return res 30 | } 31 | 32 | func DelString(key []byte, wb *levigo.WriteBatch) { 33 | wb.Delete(stringKey(key)) 34 | } 35 | 36 | func setStringLen(key []byte, length int, wb *levigo.WriteBatch) { 37 | meta := make([]byte, 5) 38 | meta[0] = StringLengthValue 39 | binary.BigEndian.PutUint32(meta[1:], uint32(length)) 40 | wb.Put(key, meta) 41 | } 42 | 43 | func stringKey(k []byte) []byte { 44 | key := make([]byte, 5+len(k)) 45 | key[0] = StringKey 46 | binary.BigEndian.PutUint32(key[1:], uint32(len(k))) 47 | copy(key[5:], k) 48 | return key 49 | } 50 | 51 | func set(k []byte, v []byte, wb *levigo.WriteBatch) error { 52 | mk := metaKey(k) 53 | res, err := DB.Get(DefaultReadOptions, mk) 54 | if err != nil { 55 | return err 56 | } 57 | 58 | // If there is a non-string key here, let's delete it first 59 | if len(res) > 0 && res[0] != StringLengthValue { 60 | del(k, res[0], wb) 61 | } 62 | 63 | setStringLen(mk, len(v), wb) 64 | wb.Put(stringKey(k), v) 65 | 66 | return nil 67 | } 68 | 69 | // APPEND 70 | func Append(args [][]byte, wb *levigo.WriteBatch) interface{} { 71 | k := args[0] 72 | appendVal := args[1] 73 | 74 | res, err := DB.Get(DefaultReadOptions, stringKey(k)) 75 | if err != nil { 76 | return err 77 | } 78 | concat := append(res, appendVal...) 79 | err = set(k, concat, wb) 80 | if err != nil { 81 | return err 82 | } 83 | 84 | return len(concat) 85 | } 86 | 87 | // BITCOUNT 88 | // BITOP 89 | // DECR 90 | // DECRBY 91 | // GETBIT 92 | // GETRANGE 93 | // GETSET 94 | // INCR 95 | // INCRBY 96 | // INCRBYFLOAT 97 | // MGET 98 | // MSET 99 | // MSETNX 100 | // PSETEX 101 | // SETBIT 102 | // SETEX 103 | // SETNX 104 | // SETRANGE 105 | // STRLEN 106 | -------------------------------------------------------------------------------- /util.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "reflect" 5 | "unsafe" 6 | ) 7 | 8 | // DANGEROUS! Only use when you know that b will never be modified. 9 | func UnsafeBytesToString(b []byte) string { 10 | bytesHeader := (*reflect.SliceHeader)(unsafe.Pointer(&b)) 11 | strHeader := &reflect.StringHeader{bytesHeader.Data, bytesHeader.Len} 12 | return *(*string)(unsafe.Pointer(strHeader)) 13 | } 14 | 15 | func EqualIgnoreCase(b1, b2 []byte) bool { 16 | if len(b1) != len(b2) { 17 | return false 18 | } 19 | for i := 0; i < len(b1); i++ { 20 | c1 := b1[i] 21 | if 'A' <= c1 && c1 <= 'Z' { 22 | c1 += 'a' - 'A' 23 | } 24 | c2 := b2[i] 25 | if 'A' <= c2 && c2 <= 'Z' { 26 | c2 += 'a' - 'A' 27 | } 28 | if c1 != c2 { 29 | return false 30 | } 31 | } 32 | return true 33 | } 34 | -------------------------------------------------------------------------------- /zset.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "bytes" 5 | "encoding/binary" 6 | "fmt" 7 | "math" 8 | "sort" 9 | "strconv" 10 | 11 | "github.com/jmhodges/levigo" 12 | "github.com/titanous/bconv" 13 | ) 14 | 15 | // Keys stored in LevelDB for zsets 16 | // 17 | // MetadataKey | key = ZCardValue | count of members uint32 18 | // 19 | // For each member: 20 | // ZSetKey | key length uint32 | key | member = score float64 21 | // ZScoreKey | key length uint32 | key | score float64 | member = empty 22 | 23 | func Zadd(args [][]byte, wb *levigo.WriteBatch) interface{} { 24 | if (len(args)-1)%2 != 0 { 25 | return fmt.Errorf("wrong number of arguments for 'zadd' command") 26 | } 27 | return zadd(args, wb, false) 28 | } 29 | 30 | func Zincrby(args [][]byte, wb *levigo.WriteBatch) interface{} { 31 | return zadd(args, wb, true) 32 | } 33 | 34 | func zadd(args [][]byte, wb *levigo.WriteBatch, incr bool) interface{} { 35 | var newMembers uint32 36 | var score float64 37 | scoreBytes := make([]byte, 8) 38 | setKey := NewKeyBuffer(ZSetKey, args[0], len(args[2])) 39 | scoreKey := NewKeyBuffer(ZScoreKey, args[0], 8+len(args[2])) 40 | 41 | mk := metaKey(args[0]) 42 | card, err := zcard(mk, nil) 43 | if err != nil { 44 | return err 45 | } 46 | 47 | // Iterate through each of the score/member pairs 48 | for i := 1; i < len(args); i += 2 { 49 | var err error 50 | score, err = bconv.ParseFloat(args[i], 64) 51 | if err != nil { 52 | return fmt.Errorf("'%s' is not a valid float", string(args[1])) 53 | } 54 | 55 | // Check if the member exists 56 | setKey.SetSuffix(args[i+1]) 57 | var res []byte 58 | if card > 0 { 59 | res, err = DB.Get(DefaultReadOptions, setKey.Key()) 60 | if err != nil { 61 | return err 62 | } 63 | } 64 | 65 | // set the score key with 8 empty bytes before the member for the score 66 | setZScoreKeyMember(scoreKey, args[i+1]) 67 | if res != nil { // We got a score from the db, so the member already exists 68 | if len(res) != 8 { 69 | return InvalidDataError 70 | } 71 | actualScore := math.Float64frombits(binary.BigEndian.Uint64(res)) 72 | if incr { // this is a ZINCRBY, so increment the score 73 | score += actualScore 74 | } 75 | if score == actualScore { // Member already exists with the same score, do nothing 76 | continue 77 | } 78 | 79 | // Delete score key for member 80 | setZScoreKeyScore(scoreKey, actualScore) 81 | wb.Delete(scoreKey.Key()) 82 | } else { // No score found, we're adding a new member 83 | newMembers++ 84 | } 85 | 86 | // Store the set and score keys 87 | binary.BigEndian.PutUint64(scoreBytes, math.Float64bits(score)) 88 | setZScoreKeyScore(scoreKey, score) 89 | wb.Put(setKey.Key(), scoreBytes) 90 | wb.Put(scoreKey.Key(), []byte{}) // The score key is only used for sorting, the value is empty 91 | } 92 | 93 | // Update the set metadata with the new cardinality 94 | if newMembers > 0 { 95 | setZcard(mk, card+newMembers, wb) 96 | } 97 | 98 | if incr { // This is a ZINCRBY, return the new score 99 | return ftoa(score) 100 | } 101 | return newMembers 102 | } 103 | 104 | func Zscore(args [][]byte, wb *levigo.WriteBatch) interface{} { 105 | res, err := DB.Get(DefaultReadOptions, NewKeyBufferWithSuffix(ZSetKey, args[0], args[1]).Key()) 106 | if err != nil { 107 | return err 108 | } 109 | if res == nil { 110 | return nil 111 | } 112 | if len(res) != 8 { 113 | return InvalidDataError 114 | } 115 | 116 | return ftoa(btof(res)) 117 | } 118 | 119 | func Zcard(args [][]byte, wb *levigo.WriteBatch) interface{} { 120 | c, err := zcard(metaKey(args[0]), nil) 121 | if err != nil { 122 | return err 123 | } 124 | return c 125 | } 126 | 127 | func zcard(key []byte, opts *levigo.ReadOptions) (uint32, error) { 128 | if opts == nil { 129 | opts = DefaultReadOptions 130 | } 131 | res, err := DB.Get(opts, key) 132 | if err != nil { 133 | return 0, err 134 | } 135 | if res == nil { 136 | return 0, nil 137 | } 138 | if len(res) > 0 && res[0] != ZCardValue { 139 | return 0, InvalidKeyTypeError 140 | } 141 | if len(res) < 5 { 142 | return 0, InvalidDataError 143 | } 144 | return binary.BigEndian.Uint32(res[1:]), nil 145 | } 146 | 147 | func Zrem(args [][]byte, wb *levigo.WriteBatch) interface{} { 148 | mk := metaKey(args[0]) 149 | card, err := zcard(mk, nil) 150 | if err != nil { 151 | return err 152 | } 153 | if card == 0 { 154 | return 0 155 | } 156 | 157 | var deleted uint32 158 | setKey := NewKeyBuffer(ZSetKey, args[0], len(args[1])) 159 | scoreKey := NewKeyBuffer(ZScoreKey, args[0], 8+len(args[1])) 160 | // Delete each of the members 161 | for _, member := range args[1:] { 162 | setKey.SetSuffix(member) 163 | res, err := DB.Get(ReadWithoutCacheFill, setKey.Key()) 164 | if err != nil { 165 | return nil 166 | } 167 | if res == nil { 168 | continue 169 | } 170 | if len(res) != 8 { 171 | return InvalidDataError 172 | } 173 | 174 | score := btof(res) 175 | setZScoreKeyMember(scoreKey, member) 176 | setZScoreKeyScore(scoreKey, score) 177 | wb.Delete(setKey.Key()) 178 | wb.Delete(scoreKey.Key()) 179 | deleted++ 180 | } 181 | if deleted == card { // We deleted all of the members, so delete the meta key 182 | wb.Delete(mk) 183 | } else if deleted > 0 { // Decrement the cardinality 184 | setZcard(mk, card-deleted, wb) 185 | } 186 | 187 | return deleted 188 | } 189 | 190 | func Zunionstore(args [][]byte, wb *levigo.WriteBatch) interface{} { 191 | return combineZset(args, zsetUnion, wb) 192 | } 193 | 194 | func Zinterstore(args [][]byte, wb *levigo.WriteBatch) interface{} { 195 | return combineZset(args, zsetInter, wb) 196 | } 197 | 198 | const ( 199 | zsetUnion int = iota 200 | zsetInter 201 | ) 202 | 203 | const ( 204 | zsetAggSum int = iota 205 | zsetAggMin 206 | zsetAggMax 207 | ) 208 | 209 | func combineZset(args [][]byte, op int, wb *levigo.WriteBatch) interface{} { 210 | var count uint32 211 | res := []interface{}{} 212 | members := make(chan *iterZsetMember) 213 | var setKey, scoreKey *KeyBuffer 214 | scoreBytes := make([]byte, 8) 215 | mk := metaKey(args[0]) 216 | 217 | if wb != nil { 218 | _, err := delKey(mk, wb) 219 | if err != nil { 220 | return err 221 | } 222 | setKey = NewKeyBuffer(ZSetKey, args[0], 0) 223 | scoreKey = NewKeyBuffer(ZScoreKey, args[0], 0) 224 | } 225 | 226 | numKeys, err := bconv.Atoi(args[1]) 227 | if err != nil { 228 | return InvalidIntError 229 | } 230 | 231 | aggregate := zsetAggSum 232 | weights := make([]float64, numKeys) 233 | scores := make([]float64, 0, numKeys) 234 | for i := 0; i < numKeys; i++ { 235 | weights[i] = 1 236 | } 237 | 238 | argOffset := 2 + numKeys 239 | if len(args) > argOffset { 240 | if len(args) > argOffset+numKeys { 241 | if EqualIgnoreCase(args[argOffset], []byte("weights")) { 242 | argOffset += numKeys + 1 243 | for i, w := range args[numKeys+3 : argOffset] { 244 | weights[i], err = bconv.ParseFloat(w, 64) 245 | if err != nil { 246 | return fmt.Errorf("weight value is not a float") 247 | } 248 | } 249 | } else { 250 | return SyntaxError 251 | } 252 | } 253 | if len(args) > argOffset { 254 | if len(args) == argOffset+2 && EqualIgnoreCase(args[argOffset], []byte("aggregate")) { 255 | agg := bytes.ToLower(args[argOffset+1]) 256 | switch { 257 | case bytes.Equal(agg, []byte("sum")): 258 | aggregate = zsetAggSum 259 | case bytes.Equal(agg, []byte("min")): 260 | aggregate = zsetAggMin 261 | case bytes.Equal(agg, []byte("max")): 262 | aggregate = zsetAggMax 263 | default: 264 | return SyntaxError 265 | } 266 | } else { 267 | return SyntaxError 268 | } 269 | } 270 | } 271 | 272 | go multiZsetIter(args[2:numKeys+2], members, op != zsetUnion) 273 | 274 | combine: 275 | for m := range members { 276 | if op == zsetInter { 277 | for _, k := range m.exists { 278 | if !k { 279 | continue combine 280 | } 281 | } 282 | } 283 | 284 | scores = scores[:0] 285 | for i, k := range m.exists { 286 | if k { 287 | scores = append(scores, m.scores[i]) 288 | } 289 | } 290 | var score float64 291 | for i, s := range scores { 292 | scores[i] = s * weights[i] 293 | } 294 | switch aggregate { 295 | case zsetAggSum: 296 | for _, s := range scores { 297 | score += s 298 | } 299 | case zsetAggMin: 300 | sort.Float64s(scores) 301 | score = scores[0] 302 | case zsetAggMax: 303 | sort.Float64s(scores) 304 | score = scores[len(scores)-1] 305 | } 306 | 307 | if wb != nil { 308 | setKey.SetSuffix(m.member) 309 | binary.BigEndian.PutUint64(scoreBytes, math.Float64bits(score)) 310 | setZScoreKeyMember(scoreKey, m.member) 311 | setZScoreKeyScore(scoreKey, score) 312 | wb.Put(setKey.Key(), scoreBytes) 313 | wb.Put(scoreKey.Key(), []byte{}) 314 | count++ 315 | } else { 316 | res = append(res, m.member, ftoa(score)) 317 | } 318 | } 319 | 320 | if wb != nil { 321 | if count > 0 { 322 | setZcard(mk, count, wb) 323 | } 324 | return count 325 | } 326 | 327 | return res 328 | } 329 | 330 | type iterZsetMember struct { 331 | member []byte 332 | exists []bool 333 | scores []float64 334 | } 335 | 336 | type zsetMember struct { 337 | member []byte 338 | score float64 339 | key int 340 | } 341 | 342 | type zsetMembers []*zsetMember 343 | 344 | func (m zsetMembers) Len() int { return len(m) } 345 | func (m zsetMembers) Less(i, j int) bool { return bytes.Compare(m[i].member, m[j].member) == -1 } 346 | func (m zsetMembers) Swap(i, j int) { m[i], m[j] = m[j], m[i] } 347 | 348 | // See set.go's multiSetIter() for details on how this works 349 | func multiZsetIter(keys [][]byte, out chan<- *iterZsetMember, stopEarly bool) { 350 | defer close(out) 351 | snapshot := DB.NewSnapshot() 352 | opts := levigo.NewReadOptions() 353 | defer opts.Close() 354 | opts.SetSnapshot(snapshot) 355 | defer DB.ReleaseSnapshot(snapshot) 356 | 357 | members := make(zsetMembers, len(keys)) 358 | iterKeys := make([]*KeyBuffer, len(keys)) 359 | iterators := make([]*levigo.Iterator, len(keys)) 360 | for i, k := range keys { 361 | if card, _ := zcard(metaKey(k), opts); card > 0 { 362 | iterKeys[i] = NewKeyBuffer(ZSetKey, k, 0) 363 | } else { 364 | // If the zset is not found, we'll assume that it is actually a set. 365 | // There is a slight edge case that an error could be raised by 366 | // zcard(), but since this function is run by a goroutine, there 367 | // isn't a clean way of handling it. 368 | iterKeys[i] = NewKeyBuffer(SetKey, k, 0) 369 | } 370 | it := DB.NewIterator(opts) 371 | defer it.Close() 372 | it.Seek(iterKeys[i].Key()) 373 | iterators[i] = it 374 | } 375 | 376 | getMember := func(i int) ([]byte, float64) { 377 | // If the iterator is done, we remove the iterator and ignore it in future runs 378 | if iterators[i] == nil || !iterators[i].Valid() { 379 | iterators[i] = nil 380 | return nil, 0 381 | } 382 | k := iterators[i].Key() 383 | if !iterKeys[i].IsPrefixOf(k) { 384 | iterators[i] = nil 385 | return nil, 0 386 | } 387 | 388 | // Default non-sorted set members to score 1.0 389 | var score float64 = 1 390 | if iterKeys[i].Type() == ZSetKey { 391 | score = btof(iterators[i].Value()) 392 | } 393 | // Strip the key prefix from the key and return the member and score 394 | return k[len(iterKeys[i].Key()):], score 395 | } 396 | 397 | // Initialize the members list 398 | for i := 0; i < len(members); i++ { 399 | m := &zsetMember{key: i} 400 | m.member, m.score = getMember(i) 401 | members[i] = m 402 | } 403 | 404 | iter: 405 | for { 406 | im := &iterZsetMember{exists: make([]bool, len(members)), scores: make([]float64, len(members))} 407 | first := true 408 | sort.Sort(members) 409 | 410 | for _, m := range members { 411 | // The member will be nil if the key it is from has no more members 412 | if m.member == nil { 413 | if m.key == 0 && stopEarly { 414 | break iter 415 | } 416 | continue 417 | } 418 | // The first member is the one that we will send out on this iteration 419 | if first { 420 | im.member = m.member 421 | first = false 422 | } 423 | if first || bytes.Compare(im.member, m.member) == 0 { 424 | im.exists[m.key] = true 425 | im.scores[m.key] = m.score 426 | iterators[m.key].Next() 427 | m.member, m.score = getMember(m.key) 428 | } else { 429 | // If the member isn't first or the same as the one we are 430 | // looking for, it's not in the list 431 | break 432 | } 433 | } 434 | // When the result member is nil, there are no members left in any of the sets 435 | if im.member == nil { 436 | break 437 | } 438 | out <- im 439 | } 440 | } 441 | 442 | func Zrange(args [][]byte, wb *levigo.WriteBatch) interface{} { 443 | return zrange(args, false) 444 | } 445 | 446 | func Zrevrange(args [][]byte, wb *levigo.WriteBatch) interface{} { 447 | return zrange(args, true) 448 | } 449 | 450 | func zrange(args [][]byte, reverse bool) interface{} { 451 | // use a snapshot for this read so that the zcard is consistent 452 | snapshot := DB.NewSnapshot() 453 | opts := levigo.NewReadOptions() 454 | opts.SetSnapshot(snapshot) 455 | 456 | count, err := zcard(metaKey(args[0]), opts) 457 | if err != nil { 458 | DB.ReleaseSnapshot(snapshot) 459 | opts.Close() 460 | return err 461 | } 462 | if count == 0 { 463 | DB.ReleaseSnapshot(snapshot) 464 | opts.Close() 465 | return []interface{}{} 466 | } 467 | 468 | start, end, err := parseRange(args[1:], int64(count)) 469 | if err != nil { 470 | return err 471 | } 472 | // the start comes after the end, so we're not going to find anything 473 | if start > end { 474 | DB.ReleaseSnapshot(snapshot) 475 | opts.Close() 476 | return []interface{}{} 477 | } 478 | 479 | var withscores bool 480 | items := end + 1 - start 481 | if len(args) >= 4 { 482 | if !EqualIgnoreCase(args[3], []byte("withscores")) || len(args) > 4 { 483 | DB.ReleaseSnapshot(snapshot) 484 | opts.Close() 485 | return SyntaxError 486 | } 487 | withscores = true 488 | items *= 2 489 | } 490 | stream := &cmdReplyStream{items, make(chan interface{})} 491 | 492 | go func() { 493 | defer close(stream.items) 494 | it := DB.NewIterator(opts) 495 | defer it.Close() 496 | 497 | var i int64 498 | iterKey := NewKeyBuffer(ZScoreKey, args[0], 0) 499 | if reverse { 500 | iterKey.ReverseIterKey() 501 | } 502 | for it.Seek(iterKey.Key()); it.Valid() && i <= end; i++ { 503 | if reverse { 504 | it.Prev() 505 | } 506 | if i >= start { 507 | score, member := parseZScoreKey(it.Key(), len(args[0])) 508 | stream.items <- member 509 | if withscores { 510 | stream.items <- ftoa(score) 511 | } 512 | } 513 | if !reverse { 514 | it.Next() 515 | } 516 | } 517 | DB.ReleaseSnapshot(snapshot) 518 | opts.Close() 519 | }() 520 | 521 | return stream 522 | } 523 | 524 | type zrangeFlag byte 525 | 526 | const ( 527 | zrangeForward zrangeFlag = iota 528 | zrangeReverse 529 | zrangeDelete 530 | zrangeCount 531 | ) 532 | 533 | func Zrangebyscore(args [][]byte, wb *levigo.WriteBatch) interface{} { 534 | return zrangebyscore(args, zrangeForward, wb) 535 | } 536 | 537 | func Zrevrangebyscore(args [][]byte, wb *levigo.WriteBatch) interface{} { 538 | return zrangebyscore(args, zrangeReverse, wb) 539 | } 540 | 541 | func Zremrangebyscore(args [][]byte, wb *levigo.WriteBatch) interface{} { 542 | return zrangebyscore(args, zrangeDelete, wb) 543 | } 544 | 545 | func Zcount(args [][]byte, wb *levigo.WriteBatch) interface{} { 546 | return zrangebyscore(args, zrangeCount, wb) 547 | } 548 | 549 | func zrangebyscore(args [][]byte, flag zrangeFlag, wb *levigo.WriteBatch) interface{} { 550 | // use a snapshot for this read so that the zcard is consistent 551 | snapshot := DB.NewSnapshot() 552 | opts := levigo.NewReadOptions() 553 | opts.SetSnapshot(snapshot) 554 | defer opts.Close() 555 | defer DB.ReleaseSnapshot(snapshot) 556 | 557 | mk := metaKey(args[0]) 558 | card, err := zcard(mk, opts) 559 | if err != nil { 560 | return err 561 | } 562 | 563 | var minExclusive, maxExclusive bool 564 | if args[1][0] == '(' { 565 | minExclusive = true 566 | args[1] = args[1][1:] 567 | } 568 | if args[2][0] == '(' { 569 | maxExclusive = true 570 | args[2] = args[2][1:] 571 | } 572 | min, err := bconv.ParseFloat(args[1], 64) 573 | max, err2 := bconv.ParseFloat(args[2], 64) 574 | if err != nil || err2 != nil { 575 | return fmt.Errorf("min or max is not a float") 576 | } 577 | if flag == zrangeReverse { 578 | min, max = max, min 579 | minExclusive, maxExclusive = maxExclusive, minExclusive 580 | } 581 | 582 | if card == 0 || min > max { 583 | if flag <= zrangeReverse { 584 | return []interface{}{} 585 | } else { 586 | return 0 587 | } 588 | } 589 | 590 | // shortcut for zcount of -Inf to +Inf 591 | if flag == zrangeCount && math.IsInf(min, -1) && math.IsInf(max, 1) { 592 | return card 593 | } 594 | 595 | var withscores bool 596 | var offset, total int64 = 0, -1 597 | if flag <= zrangeReverse && len(args) > 3 { 598 | if len(args) == 4 || len(args) == 7 { 599 | if EqualIgnoreCase(args[3], []byte("withscores")) { 600 | withscores = true 601 | } 602 | } else if len(args) != 6 { 603 | return SyntaxError 604 | } 605 | if len(args) >= 6 { 606 | if EqualIgnoreCase(args[len(args)-3], []byte("limit")) { 607 | offset, err = bconv.ParseInt(args[len(args)-2], 10, 64) 608 | total, err2 = bconv.ParseInt(args[len(args)-1], 10, 64) 609 | if err != nil || err2 != nil { 610 | return InvalidIntError 611 | } 612 | if offset < 0 || total < 1 { 613 | return []interface{}{} 614 | } 615 | } else { 616 | return SyntaxError 617 | } 618 | } 619 | } 620 | 621 | it := DB.NewIterator(opts) 622 | defer it.Close() 623 | 624 | var deleted, count uint32 625 | var deleteKey *KeyBuffer 626 | if flag == zrangeDelete { 627 | deleteKey = NewKeyBuffer(ZSetKey, args[0], 0) 628 | } 629 | 630 | res := []interface{}{} 631 | iterKey := NewKeyBuffer(ZScoreKey, args[0], 8) 632 | if flag != zrangeReverse { 633 | setZScoreKeyScore(iterKey, min) 634 | } else { 635 | setZScoreKeyScore(iterKey, max) 636 | iterKey.ReverseIterKey() 637 | } 638 | it.Seek(iterKey.Key()) 639 | for i := int64(0); it.Valid(); i++ { 640 | if flag == zrangeReverse { 641 | it.Prev() 642 | } 643 | if i >= offset { 644 | if total > -1 && i-offset >= total { 645 | break 646 | } 647 | k := it.Key() 648 | if !iterKey.IsPrefixOf(k) { 649 | break 650 | } 651 | score, member := parseZScoreKey(it.Key(), len(args[0])) 652 | if (!minExclusive && score < min) || (minExclusive && score <= min) { 653 | if flag != zrangeReverse { 654 | it.Next() 655 | } 656 | continue 657 | } 658 | if (!maxExclusive && score > max) || (maxExclusive && score >= max) { 659 | break 660 | } 661 | if flag <= zrangeReverse { 662 | res = append(res, member) 663 | if withscores { 664 | res = append(res, ftoa(score)) 665 | } 666 | } 667 | if flag == zrangeDelete { 668 | deleteKey.SetSuffix(member) 669 | wb.Delete(k) 670 | wb.Delete(deleteKey.Key()) 671 | deleted++ 672 | } 673 | if flag == zrangeCount { 674 | count++ 675 | } 676 | } 677 | if flag != zrangeReverse { 678 | it.Next() 679 | } 680 | } 681 | 682 | if flag == zrangeDelete && deleted == card { 683 | wb.Delete(mk) 684 | } else if deleted > 0 { 685 | setZcard(mk, card-deleted, wb) 686 | } 687 | if flag == zrangeDelete { 688 | return deleted 689 | } 690 | if flag == zrangeCount { 691 | return count 692 | } 693 | 694 | return res 695 | } 696 | 697 | func Zrank(args [][]byte, wb *levigo.WriteBatch) interface{} { 698 | return zrank(args, false) 699 | } 700 | 701 | func Zrevrank(args [][]byte, wb *levigo.WriteBatch) interface{} { 702 | return zrank(args, true) 703 | } 704 | 705 | func zrank(args [][]byte, reverse bool) interface{} { 706 | // use a snapshot for this read so that the zcard is consistent 707 | snapshot := DB.NewSnapshot() 708 | opts := levigo.NewReadOptions() 709 | opts.SetSnapshot(snapshot) 710 | defer opts.Close() 711 | defer DB.ReleaseSnapshot(snapshot) 712 | 713 | card, err := zcard(metaKey(args[0]), opts) 714 | if err != nil { 715 | return err 716 | } 717 | if card == 0 { 718 | return nil 719 | } 720 | 721 | iterKey := NewKeyBuffer(ZScoreKey, args[0], 0) 722 | it := DB.NewIterator(opts) 723 | defer it.Close() 724 | 725 | if reverse { 726 | iterKey.ReverseIterKey() 727 | } 728 | i := 0 729 | for it.Seek(iterKey.Key()); it.Valid(); i++ { 730 | if reverse { 731 | it.Prev() 732 | } 733 | k := it.Key() 734 | if !iterKey.IsPrefixOf(k) { 735 | break 736 | } 737 | if bytes.Equal(args[1], k[len(args[0])+13:]) { 738 | return i 739 | } 740 | if !reverse { 741 | it.Next() 742 | } 743 | } 744 | 745 | return nil 746 | } 747 | 748 | func DelZset(key []byte, wb *levigo.WriteBatch) { 749 | // TODO: count keys to verify everything works as expected? 750 | it := DB.NewIterator(ReadWithoutCacheFill) 751 | defer it.Close() 752 | iterKey := NewKeyBuffer(ZSetKey, key, 0) 753 | scoreKey := NewKeyBuffer(ZScoreKey, key, 0) 754 | 755 | for it.Seek(iterKey.Key()); it.Valid(); it.Next() { 756 | k := it.Key() 757 | // If the prefix of the current key doesn't match the iteration key, 758 | // we have reached the end of the zset 759 | if !iterKey.IsPrefixOf(k) { 760 | break 761 | } 762 | wb.Delete(k) 763 | setZScoreKeyMember(scoreKey, k[len(iterKey.Key()):]) 764 | setZScoreKeyScore(scoreKey, btof(it.Value())) 765 | wb.Delete(scoreKey.Key()) 766 | } 767 | } 768 | 769 | func setZcard(key []byte, card uint32, wb *levigo.WriteBatch) { 770 | data := make([]byte, 5) 771 | data[0] = ZCardValue 772 | binary.BigEndian.PutUint32(data[1:], card) 773 | wb.Put(key, data) 774 | } 775 | 776 | func setZScoreKeyScore(key *KeyBuffer, score float64) { 777 | writeByteSortableFloat(key.SuffixForRead(8), score) 778 | } 779 | 780 | func setZScoreKeyMember(key *KeyBuffer, member []byte) { 781 | key.SetSuffix(append(make([]byte, 8, 8+len(member)), member...)) 782 | } 783 | 784 | func parseZScoreKey(b []byte, keyLen int) (float64, []byte) { 785 | return readByteSortableFloat(b[keyLen+5:]), b[keyLen+13:] 786 | } 787 | 788 | func btof(b []byte) float64 { 789 | return math.Float64frombits(binary.BigEndian.Uint64(b)) 790 | } 791 | 792 | func ftoa(f float64) []byte { 793 | b := strconv.AppendFloat(nil, f, 'g', -1, 64) 794 | if len(b) > 1 && b[1] == 'I' { // -Inf/+Inf to lowercase 795 | b[1] = 'i' 796 | } 797 | if b[0] == 'N' { // NaN to lowercase 798 | b[0], b[2] = 'n', 'n' 799 | } 800 | if b[0] == '+' { // +inf to inf 801 | b = b[1:] 802 | } 803 | 804 | return b 805 | } 806 | 807 | /* Natural sorting of floating point numbers 808 | * 809 | * ENCODING 810 | * If the number is positive, flip the sign (first bit to 1) 811 | * Else the number is negative, flip all the bits 812 | * 813 | * DECODING 814 | * If the first byte is >= 0x80 (0b10000000), decode float, flip the sign 815 | * Else flip all the bits, decode float 816 | */ 817 | 818 | func writeByteSortableFloat(b []byte, f float64) { 819 | if math.Signbit(f) { 820 | binary.BigEndian.PutUint64(b, math.Float64bits(f)) 821 | for i, v := range b[:8] { 822 | b[i] = v ^ 255 823 | } 824 | } else { 825 | binary.BigEndian.PutUint64(b, math.Float64bits(-f)) 826 | } 827 | } 828 | 829 | func readByteSortableFloat(b []byte) float64 { 830 | if b[0] < 0x80 { 831 | for i, v := range b[:8] { 832 | b[i] = v ^ 255 833 | } 834 | return math.Float64frombits(binary.BigEndian.Uint64(b)) 835 | } 836 | return -math.Float64frombits(binary.BigEndian.Uint64(b)) 837 | } 838 | 839 | func ZunionInterKeys(args [][]byte) [][]byte { 840 | numKeys, err := bconv.Atoi(args[1]) 841 | // don't return any keys if the response will be a syntax error 842 | if err != nil || len(args) < 2+numKeys { 843 | return nil 844 | } 845 | keys := make([][]byte, 1, 1+numKeys) 846 | keys[0] = args[0] 847 | keyloop: 848 | for _, k := range args[2 : 2+numKeys] { 849 | for _, key := range keys { 850 | // skip keys that are already in the array 851 | if bytes.Equal(k, key) { 852 | continue keyloop 853 | } 854 | } 855 | keys = append(keys, k) 856 | } 857 | return keys 858 | } 859 | 860 | // ZREMRANGEBYRANK 861 | --------------------------------------------------------------------------------