162 lines
3.4 KiB
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
|
|
}
|