mtsatellite/common/redisclient.go
2017-03-03 16:03:08 +01:00

160 lines
3.4 KiB
Go

// Copyright 2014, 2015 by Sascha L. Teichmann
// Use of this source code is governed by the MIT license
// that can be found in the LICENSE file.
package common
import (
"bufio"
"bytes"
"errors"
"fmt"
"net"
"strconv"
"sync"
"unicode"
)
type RedisClient struct {
conn net.Conn
reader *bufio.Reader
scratch [80]byte
}
func NewRedisClient(network, address string) (client *RedisClient, err error) {
var conn net.Conn
if conn, err = net.Dial(network, address); err != nil {
return
}
client = &RedisClient{conn: conn, reader: bufio.NewReaderSize(conn, 8*1024)}
return
}
func (client *RedisClient) Close() error {
return client.conn.Close()
}
var (
writeArray4 = []byte("*4\r\n")
hspatial = []byte("HSPATIAL")
nl = []byte("\r\n")
ignore = []byte("IGNORE")
)
func (client *RedisClient) writeBulkString(buf []byte, data []byte) (err error) {
buf = append(buf, '$')
buf = strconv.AppendInt(buf, int64(len(data)), 10)
buf = append(buf, nl...)
buf = append(buf, data...)
buf = append(buf, nl...)
_, err = client.conn.Write(buf)
return
}
func (client *RedisClient) writeHSpatial(p1, p2 int64) (err error) {
if _, err = client.conn.Write(writeArray4); err != nil {
return
}
b1 := client.scratch[:0:40]
b2 := client.scratch[40:40:80]
if err = client.writeBulkString(b1, hspatial); err != nil {
return
}
if err = client.writeBulkString(b1, ignore); err != nil {
return
}
if err = client.writeBulkString(b1, keyToBytes(p1, b2)); err != nil {
return
}
err = client.writeBulkString(b1, keyToBytes(p2, b2))
return
}
func isError(line []byte) error {
if len(line) > 0 && line[0] == '-' {
return fmt.Errorf("error: %s", line[1:])
}
return nil
}
// parseSize is a cheaper replacement for fmt.Sscanf(string(line), "$%d\r\n", &size).
func parseSize(line []byte) (int, error) {
if len(line) < 1 || line[0] != '$' {
return 0, errors.New("Missing '$' at begin of line")
}
line = bytes.TrimFunc(line[1:], unicode.IsSpace)
v, err := strconv.ParseInt(string(line), 10, 0)
return int(v), err
}
func (client *RedisClient) readBulkString(data *[]byte) (size int, err error) {
var line []byte
if line, err = client.reader.ReadBytes('\n'); err != nil {
return
}
if err = isError(line); err != nil {
return
}
if size, err = parseSize(line); err != nil || size <= 0 {
return
}
if cap(*data) < size {
*data = make([]byte, size)
}
for rest := size; rest > 0; {
var n int
if n, err = client.reader.Read((*data)[size-rest : size]); err != nil {
return
}
rest -= n
}
_, err = client.reader.ReadBytes('\n')
return
}
var dataPool = sync.Pool{
New: func() interface{} {
return make([]byte, 8*1024)
},
}
func (client *RedisClient) QueryCuboid(cuboid Cuboid, fn func(*Block)) (count int, err error) {
p1 := CoordToPlain(cuboid.P1)
p2 := CoordToPlain(cuboid.P2)
if err = client.writeHSpatial(p1, p2); err != nil {
return
}
data := dataPool.Get().([]byte)
defer dataPool.Put(data[:0])
var (
block = Block{}
size int
key int64
)
for s := client.scratch[:]; ; count++ {
p := &s
if size, err = client.readBulkString(p); err != nil {
return
}
if size <= 0 {
break
}
if key, err = DecodeStringFromBytes((*p)[:size]); err != nil {
return
}
block.Coord = PlainToCoord(key)
if size, err = client.readBulkString(&data); err != nil || size < 0 {
return
}
block.Data = data[:size]
fn(&block)
}
return
}