package main import ( "fmt" "net" "bufio" "bitbucket.org/s_l_teichmann/mtredisalize/common" ) type Client struct { conn net.Conn reader *bufio.Reader } func NewClient(network, address string) (client *Client, err error) { var conn net.Conn if conn, err = net.Dial(network, address); err != nil { return } client = &Client{conn: conn, reader: bufio.NewReaderSize(conn, 8*1024)} return } func (client *Client) Close() error { return client.conn.Close() } func (client *Client) writeArray(size int) (err error) { _, err = client.conn.Write([]byte(fmt.Sprintf("*%d\r\n", size))) return } func (client *Client) 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 *Client) 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(common.StringToBytes(p1)); err != nil { return } err = client.writeBulkString(common.StringToBytes(p2)) return } func (client *Client) 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 *Client) 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 *Client) QueryCuboid(cuboid common.Cuboid, fn func(*common.Block)) (err error) { p1 := common.CoordToPlain(cuboid.P1) p2 := common.CoordToPlain(cuboid.P2) if err = client.writeHSpatial(p1, p2); err != nil { return } var ( data = make([]byte, 8*1024) block = common.Block{} size int key int64 ) for { if size, err = client.readBulkString(&data); err != nil { return } if size <= 0 { break } if key, err = common.DecodeStringFromBytes(data[0:size]); err != nil { return } block.Coord = common.PlainToCoord(key) if size, err = client.readBulkString(&data); err != nil { return } block.Data = data[0:size] fn(&block) } _ = size return }