mirror of
https://bitbucket.org/s_l_teichmann/mtsatellite
synced 2024-09-06 04:40:18 +02:00
139 lines
2.9 KiB
Go
139 lines
2.9 KiB
Go
// Copyright 2014 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"
|
|
"fmt"
|
|
"net"
|
|
)
|
|
|
|
type RedisClient struct {
|
|
conn net.Conn
|
|
reader *bufio.Reader
|
|
}
|
|
|
|
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()
|
|
}
|
|
|
|
func (client *RedisClient) writeArray(size int) (err error) {
|
|
_, err = client.conn.Write([]byte(fmt.Sprintf("*%d\r\n", size)))
|
|
return
|
|
}
|
|
|
|
func (client *RedisClient) writeBulkString(data []byte) (err error) {
|
|
if _, err = client.conn.Write([]byte(fmt.Sprintf("$%d\r\n", len(data)))); err != nil {
|
|
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) {
|
|
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 {
|
|
return
|
|
}
|
|
err = client.writeBulkString(StringToBytes(p2))
|
|
return
|
|
}
|
|
|
|
func (client *RedisClient) readLine() (data []byte, err error) {
|
|
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) {
|
|
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)
|
|
if err = client.writeHSpatial(p1, p2); err != nil {
|
|
return
|
|
}
|
|
var (
|
|
data = make([]byte, 8*1024)
|
|
block = Block{}
|
|
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 {
|
|
return
|
|
}
|
|
block.Coord = PlainToCoord(key)
|
|
if size, err = client.readBulkString(&data); err != nil {
|
|
return
|
|
}
|
|
block.Data = data[0:size]
|
|
fn(&block)
|
|
}
|
|
|
|
return
|
|
}
|