mtsatellite/common/redisclient.go

162 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"
"unicode"
)
type RedisClient struct {
conn net.Conn
reader *bufio.Reader
arena []byte
scratch [130]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 writeBulkString(buf []byte, data []byte) []byte {
buf = append(buf, '$')
buf = strconv.AppendInt(buf, int64(len(data)), 10)
buf = append(buf, nl...)
buf = append(buf, data...)
buf = append(buf, nl...)
return buf
}
func (client *RedisClient) writeHSpatial(p1, p2 int64) error {
tmp := client.scratch[:0:40]
buf := client.scratch[40:40]
buf = append(buf, writeArray4...)
buf = writeBulkString(buf, hspatial)
buf = writeBulkString(buf, ignore)
buf = writeBulkString(buf, keyToBytes(p1, tmp))
buf = writeBulkString(buf, keyToBytes(p2, tmp))
_, err := client.conn.Write(buf)
return err
}
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) alloc(size int) []byte {
a := client.arena
if len(a) < size {
a = make([]byte, Max(size, 16*1024))
}
x := a[:size:size]
client.arena = a[size:]
return x
}
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 = client.alloc(size)
} else {
*data = (*data)[: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
}
func (client *RedisClient) QueryCuboid(cuboid Cuboid, fn func(*Block) *Block) (count int, err error) {
p1 := CoordToPlain(cuboid.P1)
p2 := CoordToPlain(cuboid.P2)
if err = client.writeHSpatial(p1, p2); err != nil {
return
}
var (
block *Block
size int
key int64
data []byte
)
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); err != nil {
return
}
if size, err = client.readBulkString(&data); err != nil || size < 0 {
return
}
if block == nil {
block = &Block{Coord: PlainToCoord(key), Data: data}
} else {
*block = Block{Coord: PlainToCoord(key), Data: data}
}
if block = fn(block); block != nil {
data = block.Data[:0]
} else {
data = nil
}
}
return
}