2015-07-26 16:44:51 +02:00
|
|
|
// Copyright 2014, 2015 by Sascha L. Teichmann
|
2014-09-10 17:33:13 +02:00
|
|
|
// Use of this source code is governed by the MIT license
|
|
|
|
// that can be found in the LICENSE file.
|
|
|
|
|
2014-09-13 19:18:12 +02:00
|
|
|
package common
|
2014-09-07 21:46:55 +02:00
|
|
|
|
|
|
|
import (
|
2014-09-13 19:18:12 +02:00
|
|
|
"bufio"
|
2017-02-25 19:29:27 +01:00
|
|
|
"bytes"
|
|
|
|
"errors"
|
2014-09-07 21:46:55 +02:00
|
|
|
"fmt"
|
|
|
|
"net"
|
2015-06-28 14:41:07 +02:00
|
|
|
"strconv"
|
2017-02-25 17:00:17 +01:00
|
|
|
"sync"
|
2017-02-25 19:29:27 +01:00
|
|
|
"unicode"
|
2014-09-07 21:46:55 +02:00
|
|
|
)
|
|
|
|
|
2014-09-13 19:18:12 +02:00
|
|
|
type RedisClient struct {
|
2017-03-02 10:57:03 +01:00
|
|
|
conn net.Conn
|
|
|
|
reader *bufio.Reader
|
|
|
|
scratch [80]byte
|
2014-09-07 21:46:55 +02:00
|
|
|
}
|
|
|
|
|
2014-09-13 19:18:12 +02:00
|
|
|
func NewRedisClient(network, address string) (client *RedisClient, err error) {
|
2014-09-07 21:46:55 +02:00
|
|
|
var conn net.Conn
|
|
|
|
if conn, err = net.Dial(network, address); err != nil {
|
|
|
|
return
|
|
|
|
}
|
|
|
|
|
2014-09-13 19:18:12 +02:00
|
|
|
client = &RedisClient{conn: conn, reader: bufio.NewReaderSize(conn, 8*1024)}
|
2014-09-07 21:46:55 +02:00
|
|
|
return
|
|
|
|
}
|
|
|
|
|
2014-09-13 19:18:12 +02:00
|
|
|
func (client *RedisClient) Close() error {
|
2014-09-07 21:46:55 +02:00
|
|
|
return client.conn.Close()
|
|
|
|
}
|
|
|
|
|
2017-02-24 19:49:08 +01:00
|
|
|
var (
|
|
|
|
writeArray4 = []byte("*4\r\n")
|
|
|
|
hspatial = []byte("HSPATIAL")
|
|
|
|
nl = []byte("\r\n")
|
|
|
|
ignore = []byte("IGNORE")
|
|
|
|
)
|
|
|
|
|
2017-02-28 18:50:45 +01:00
|
|
|
func (client *RedisClient) writeBulkString(buf []byte, data []byte) (err error) {
|
2017-02-25 17:00:17 +01:00
|
|
|
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)
|
2014-09-07 21:46:55 +02:00
|
|
|
return
|
|
|
|
}
|
|
|
|
|
2014-09-13 19:18:12 +02:00
|
|
|
func (client *RedisClient) writeHSpatial(p1, p2 int64) (err error) {
|
2017-02-24 19:49:08 +01:00
|
|
|
if _, err = client.conn.Write(writeArray4); err != nil {
|
2014-09-07 21:46:55 +02:00
|
|
|
return
|
|
|
|
}
|
2017-03-02 10:57:03 +01:00
|
|
|
b1 := client.scratch[:0:40]
|
|
|
|
b2 := client.scratch[40:40:80]
|
|
|
|
|
|
|
|
if err = client.writeBulkString(b1, hspatial); err != nil {
|
|
|
|
return
|
2014-09-07 21:46:55 +02:00
|
|
|
}
|
2017-03-02 10:57:03 +01:00
|
|
|
if err = client.writeBulkString(b1, ignore); err != nil {
|
|
|
|
return
|
2014-09-07 21:46:55 +02:00
|
|
|
}
|
2017-03-02 10:57:03 +01:00
|
|
|
if err = client.writeBulkString(b1, keyToBytes(p1, b2)); err != nil {
|
|
|
|
return
|
2014-09-07 21:46:55 +02:00
|
|
|
}
|
2017-03-02 10:57:03 +01:00
|
|
|
err = client.writeBulkString(b1, keyToBytes(p2, b2))
|
2014-09-07 21:46:55 +02:00
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
func isError(line []byte) error {
|
|
|
|
if len(line) > 0 && line[0] == '-' {
|
|
|
|
return fmt.Errorf("error: %s", line[1:])
|
|
|
|
}
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
2017-02-25 19:29:27 +01:00
|
|
|
// 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
|
|
|
|
}
|
|
|
|
|
2014-09-13 19:18:12 +02:00
|
|
|
func (client *RedisClient) readBulkString(data *[]byte) (size int, err error) {
|
2014-09-07 21:46:55 +02:00
|
|
|
var line []byte
|
2017-02-24 21:23:04 +01:00
|
|
|
if line, err = client.reader.ReadBytes('\n'); err != nil {
|
2014-09-07 21:46:55 +02:00
|
|
|
return
|
|
|
|
}
|
|
|
|
if err = isError(line); err != nil {
|
|
|
|
return
|
|
|
|
}
|
2017-02-25 19:29:27 +01:00
|
|
|
if size, err = parseSize(line); err != nil || size <= 0 {
|
2014-09-07 21:46:55 +02:00
|
|
|
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
|
|
|
|
}
|
2017-02-24 19:54:19 +01:00
|
|
|
_, err = client.reader.ReadBytes('\n')
|
2014-09-07 21:46:55 +02:00
|
|
|
return
|
|
|
|
}
|
|
|
|
|
2017-02-28 00:33:59 +01:00
|
|
|
var dataPool = sync.Pool{
|
|
|
|
New: func() interface{} {
|
|
|
|
return make([]byte, 8*1024)
|
|
|
|
},
|
|
|
|
}
|
|
|
|
|
2017-02-26 12:57:38 +01:00
|
|
|
func (client *RedisClient) QueryCuboid(cuboid Cuboid, fn func(*Block)) (count int, err error) {
|
2014-09-13 19:18:12 +02:00
|
|
|
p1 := CoordToPlain(cuboid.P1)
|
|
|
|
p2 := CoordToPlain(cuboid.P2)
|
2014-09-07 21:46:55 +02:00
|
|
|
if err = client.writeHSpatial(p1, p2); err != nil {
|
|
|
|
return
|
|
|
|
}
|
2017-02-28 00:33:59 +01:00
|
|
|
|
|
|
|
data := dataPool.Get().([]byte)
|
|
|
|
defer dataPool.Put(data[:0])
|
|
|
|
|
2014-09-07 21:46:55 +02:00
|
|
|
var (
|
2014-09-13 19:18:12 +02:00
|
|
|
block = Block{}
|
2014-09-07 21:46:55 +02:00
|
|
|
size int
|
|
|
|
key int64
|
|
|
|
)
|
|
|
|
|
2017-02-26 12:57:38 +01:00
|
|
|
for ; ; count++ {
|
2014-09-07 21:46:55 +02:00
|
|
|
if size, err = client.readBulkString(&data); err != nil {
|
|
|
|
return
|
|
|
|
}
|
|
|
|
if size <= 0 {
|
|
|
|
break
|
|
|
|
}
|
2017-02-28 00:33:59 +01:00
|
|
|
if key, err = DecodeStringFromBytes(data[:size]); err != nil {
|
2014-09-07 21:46:55 +02:00
|
|
|
return
|
|
|
|
}
|
2014-09-13 19:18:12 +02:00
|
|
|
block.Coord = PlainToCoord(key)
|
2017-02-25 19:29:27 +01:00
|
|
|
if size, err = client.readBulkString(&data); err != nil || size < 0 {
|
2014-09-07 21:46:55 +02:00
|
|
|
return
|
|
|
|
}
|
2017-02-28 00:33:59 +01:00
|
|
|
block.Data = data[:size]
|
2014-09-07 21:46:55 +02:00
|
|
|
fn(&block)
|
|
|
|
}
|
|
|
|
|
|
|
|
return
|
|
|
|
}
|