// 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" "sync" "unicode" ) 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() } var ( writeArray4 = []byte("*4\r\n") hspatial = []byte("HSPATIAL") nl = []byte("\r\n") ignore = []byte("IGNORE") ) // stringPool is a pool to help recycle bulk strings // for writing. Experimented with sync.Pool and // channel based leaky buffers but the mutex based // version performs best in this case. type stringPool struct { list [][]byte sync.Mutex } func (sp *stringPool) alloc() (l []byte) { sp.Lock() if n := len(sp.list); n > 0 { l = sp.list[n-1] sp.list[n-1] = nil sp.list = sp.list[:n-1] sp.Unlock() } else { sp.Unlock() l = make([]byte, 0, 32) } return } func (sp *stringPool) free(b []byte) { b = b[:0] sp.Lock() sp.list = append(sp.list, b) sp.Unlock() } var spool stringPool func (client *RedisClient) writeBulkString(data []byte) (err error) { buf := spool.alloc() buf = append(buf, '$') buf = strconv.AppendInt(buf, int64(len(data)), 10) buf = append(buf, nl...) buf = append(buf, data...) buf = append(buf, nl...) _, err = client.conn.Write(buf) spool.free(buf) 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 } // 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) 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 = 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)) (count int, 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 ; ; count++ { 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 || size < 0 { return } block.Data = data[0:size] fn(&block) } return }