mtsatellite/common/redisclient.go

145 lines
3.0 KiB
Go
Raw Normal View History

2015-07-26 16:44:51 +02:00
// 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
2014-09-07 21:46:55 +02:00
import (
"bufio"
2014-09-07 21:46:55 +02:00
"fmt"
"net"
"strconv"
2014-09-07 21:46:55 +02:00
)
type RedisClient struct {
2014-09-07 21:46:55 +02:00
conn net.Conn
reader *bufio.Reader
}
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
}
client = &RedisClient{conn: conn, reader: bufio.NewReaderSize(conn, 8*1024)}
2014-09-07 21:46:55 +02:00
return
}
func (client *RedisClient) Close() error {
2014-09-07 21:46:55 +02:00
return client.conn.Close()
}
func (client *RedisClient) writeArray(size int) (err error) {
2014-09-07 21:46:55 +02:00
_, err = client.conn.Write([]byte(fmt.Sprintf("*%d\r\n", size)))
return
}
func redisLength(prefix byte, s int) []byte {
buf := append(make([]byte, 0, 16), prefix)
return append(strconv.AppendInt(buf, int64(s), 10), '\r', '\n')
}
func (client *RedisClient) writeBulkString(data []byte) (err error) {
if _, err = client.conn.Write(redisLength('$', len(data))); err != nil {
2014-09-07 21:46:55 +02:00
return
}
if _, err = client.conn.Write(data); err != nil {
return
}
_, err = client.conn.Write([]byte("\r\n"))
return
}
func (client *RedisClient) writeHSpatial(p1, p2 int64) (err error) {
2014-09-07 21:46:55 +02:00
if err = client.writeArray(4); err != nil {
return
}
if err = client.writeBulkString([]byte("HSPATIAL")); err != nil {
return
}
if err = client.writeBulkString([]byte("IGNORE")); err != nil {
return
}
if err = client.writeBulkString(StringToBytes(p1)); err != nil {
2014-09-07 21:46:55 +02:00
return
}
err = client.writeBulkString(StringToBytes(p2))
2014-09-07 21:46:55 +02:00
return
}
func (client *RedisClient) readLine() (data []byte, err error) {
2014-09-07 21:46:55 +02:00
return client.reader.ReadBytes('\n')
}
func isError(line []byte) error {
if len(line) > 0 && line[0] == '-' {
return fmt.Errorf("error: %s", line[1:])
}
return nil
}
func (client *RedisClient) readBulkString(data *[]byte) (size int, err error) {
2014-09-07 21:46:55 +02:00
var line []byte
if line, err = client.readLine(); err != nil {
return
}
if err = isError(line); err != nil {
return
}
if _, err = fmt.Sscanf(string(line), "$%d\r\n", &size); err != nil {
return
}
if 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
}
if _, err = client.reader.ReadBytes('\n'); err != nil {
return
}
return
}
func (client *RedisClient) QueryCuboid(cuboid Cuboid, fn func(*Block)) (err error) {
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
}
var (
data = make([]byte, 8*1024)
block = Block{}
2014-09-07 21:46:55 +02:00
size int
key int64
)
for {
if size, err = client.readBulkString(&data); err != nil {
return
}
if size <= 0 {
break
}
if key, err = DecodeStringFromBytes(data[0:size]); err != nil {
2014-09-07 21:46:55 +02:00
return
}
block.Coord = PlainToCoord(key)
2014-09-07 21:46:55 +02:00
if size, err = client.readBulkString(&data); err != nil {
return
}
block.Data = data[0:size]
fn(&block)
}
return
}