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"
	"bytes"
	"errors"
	"fmt"
	"net"
	"strconv"
	"strings"
	"unicode"
)

type RedisClientFactory struct {
	proto   string
	address string
}

func NewRedisClientFactory(host string, port int) (*RedisClientFactory, error) {
	var address string
	if strings.ContainsRune(host, '/') {
		address = host
	} else {
		address = fmt.Sprintf("%s:%d", host, port)
	}

	var proto string
	if strings.ContainsRune(address, '/') {
		proto = "unix"
	} else {
		proto = "tcp"
	}
	return &RedisClientFactory{
		proto:   proto,
		address: address,
	}, nil
}

func (rcf *RedisClientFactory) Close() error {
	return nil
}

func (rcf *RedisClientFactory) Create() (DBClient, error) {
	return NewRedisClient(rcf.proto, rcf.address)
}

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
}