// 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" "unicode" ) 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 }