mtsatellite/common/redisclient.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"
	"fmt"
	"net"
	"strconv"
)

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 redisLength(prefix byte, s int) []byte {
	buf := append(make([]byte, 0, 16), prefix)
	return append(strconv.AppendInt(buf, int64(s), 10), '\r', '\n')
}

var (
	writeArray4 = []byte("*4\r\n")
	hspatial    = []byte("HSPATIAL")
	nl          = []byte("\r\n")
	ignore      = []byte("IGNORE")
)

func (client *RedisClient) writeBulkString(data []byte) (err error) {
	if _, err = client.conn.Write(redisLength('$', len(data))); err != nil {
		return
	}
	if _, err = client.conn.Write(data); err != nil {
		return
	}
	_, err = client.conn.Write(nl)
	return
}

func (client *RedisClient) writeHSpatial(p1, p2 int64) (err error) {
	if _, err = client.conn.Write(writeArray4); err != nil {
		return
	}
	if err = client.writeBulkString(hspatial); err != nil {
		return
	}
	if err = client.writeBulkString(ignore); err != nil {
		return
	}
	if err = client.writeBulkString(StringToBytes(p1)); err != nil {
		return
	}
	err = client.writeBulkString(StringToBytes(p2))
	return
}

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.reader.ReadBytes('\n'); 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
	}
	_, err = client.reader.ReadBytes('\n')
	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
}