func selectDB(redisConn BufioDeadlineReadWriter, dbIndex int, timeout int) error {
redisReader := redisConn.BufioReader()
if err := redisConn.SetWriteDeadline(time.Now().Add(time.Duration(timeout) * time.Second)); err != nil {
return errors.Trace(err)
}
dbStr := strconv.Itoa(dbIndex)
data := []byte(fmt.Sprintf("*2\r\n$6\r\nSELECT\r\n$%d\r\n%s\r\n", len(dbStr), dbStr))
if err := writeBytes2Redis(data, redisConn); err != nil {
return errors.Trace(err)
}
if err := redisConn.SetReadDeadline(time.Now().Add(time.Duration(timeout) * time.Second)); err != nil {
return errors.Trace(err)
}
resp, err := parser.Parse(redisReader)
if err != nil {
return errors.Trace(err)
}
b, err := resp.Bytes()
if err != nil {
return errors.Trace(err)
}
if !bytes.Equal(b, OK_BYTES) {
return errors.Trace(fmt.Errorf("select %d not ok", dbIndex))
}
return nil
}
func write2Client(redisReader *bufio.Reader, clientWriter io.Writer) (redisErr error, clientErr error) {
resp, err := parser.Parse(redisReader)
if err != nil {
return errors.Trace(err), errors.Trace(err)
}
b, err := resp.Bytes()
if err != nil {
return errors.Trace(err), errors.Trace(err)
}
_, err = clientWriter.Write(b)
return nil, errors.Trace(err)
}
func TestForward(t *testing.T) {
client := &fakeDeadlineReadWriter{r: bufio.NewReader(bytes.NewBuffer([]byte(simple_request))),
w: bufio.NewWriter(&bytes.Buffer{})}
redis := &fakeDeadlineReadWriter{r: bufio.NewReader(bytes.NewBuffer([]byte(simple_request))),
w: bufio.NewWriter(&bytes.Buffer{})}
resp, err := parser.Parse(bufio.NewReader(bytes.NewBuffer([]byte(simple_request))))
if err != nil {
t.Error(err)
}
_, clientErr := forward(client, redis, resp, 5)
if clientErr != nil {
t.Error(clientErr)
}
}
func TestWrite2Redis(t *testing.T) {
var result bytes.Buffer
var input bytes.Buffer
input.WriteString(simple_request)
resp, err := parser.Parse(bufio.NewReader(&input))
if err != nil {
t.Error(err)
}
err = write2Redis(resp, &result)
if err != nil {
t.Error(err)
}
if string(result.Bytes()) != simple_request {
t.Error("not match")
}
}
func (s *Server) redisTunnel(c *session) error {
resp, err := parser.Parse(c.r) // read client request
if err != nil {
return errors.Trace(err)
}
op, keys, err := resp.GetOpKeys()
if err != nil {
return errors.Trace(err)
}
opstr := strings.ToUpper(string(op))
var group string
group, keys, err = s.getOpGroupKeys(opstr, keys)
if err != nil {
return errors.Trace(err)
}
if len(keys) == 0 {
keys = [][]byte{[]byte("fakeKey")}
}
start := time.Now()
k := keys[0]
//log.Debugf("op: %s, %s", opstr, keys[0])
next, err := s.filter(opstr, keys, c)
if err != nil {
return errors.Trace(err)
}
s.counter.Add(opstr, 1)
s.counter.Add("ops", 1)
if !next {
return nil
}
//must check multi keys in same slot
i, mkeys, err := checkMigrateKeys(opstr, keys)
if err != nil {
return errors.Trace(err)
}
//i := mapKey2Slot(k)
token := s.concurrentLimiter.Get()
check_state:
s.mu.RLock()
if s.slots[i] == nil {
s.mu.RUnlock()
s.concurrentLimiter.Put(token)
return errors.Errorf("should never happend, slot %d is empty", i)
}
//wait for state change, should be soon
if s.slots[i].slotInfo.State.Status == models.SLOT_STATUS_PRE_MIGRATE {
s.mu.RUnlock()
time.Sleep(10 * time.Millisecond)
goto check_state
}
defer func() {
s.mu.RUnlock()
sec := time.Since(start).Seconds()
if sec > 2 {
log.Warningf("op: %s, key:%s, on: %s, too long %d seconds, client: %s", opstr,
string(k), s.slots[i].dst.Master(), int(sec), c.RemoteAddr().String())
}
recordResponseTime(s.counter, time.Duration(sec)*1000)
s.concurrentLimiter.Put(token)
}()
if err := s.handleMigrateState(i, opstr, group, mkeys); err != nil {
return errors.Trace(err)
}
//get redis connection
redisConn, err := s.pools.GetConn(s.slots[i].dst.Master())
if err != nil {
return errors.Trace(err)
}
if redisConn.(*redispool.PooledConn).DB != i {
if err := selectDB(redisConn.(*redispool.PooledConn), i, s.net_timeout); err != nil {
redisConn.Close()
s.pools.ReleaseConn(redisConn)
return errors.Trace(err)
}
redisConn.(*redispool.PooledConn).DB = i
}
redisErr, clientErr := forward(c, redisConn.(*redispool.PooledConn), resp, s.net_timeout)
if redisErr != nil {
redisConn.Close()
}
s.pools.ReleaseConn(redisConn)
return errors.Trace(clientErr)
}
func (s *Server) handleMigrateState(slotIndex int, op string, group string, keys [][]byte) error {
shd := s.slots[slotIndex]
if shd.slotInfo.State.Status != models.SLOT_STATUS_MIGRATE {
return nil
}
if shd.migrateFrom == nil {
log.Fatalf("migrateFrom not exist %+v", shd)
}
if shd.dst.Master() == shd.migrateFrom.Master() {
log.Fatalf("the same migrate src and dst, %+v", shd)
}
redisConn, err := s.pools.GetConn(shd.migrateFrom.Master())
if err != nil {
return errors.Trace(err)
}
defer s.pools.ReleaseConn(redisConn)
if redisConn.(*redispool.PooledConn).DB != slotIndex {
if err := selectDB(redisConn.(*redispool.PooledConn), slotIndex, s.net_timeout); err != nil {
redisConn.Close()
return errors.Trace(err)
}
redisConn.(*redispool.PooledConn).DB = slotIndex
}
redisReader := redisConn.(*redispool.PooledConn).BufioReader()
//migrate multi keys
for _, key := range keys {
if s.broker == LedisBroker {
err = ledisWriteMigrateKeyCmd(redisConn.(*redispool.PooledConn), shd.dst.Master(), 30*1000, group, key, slotIndex)
} else {
err = writeMigrateKeyCmd(redisConn.(*redispool.PooledConn), shd.dst.Master(), 30*1000, key, slotIndex)
}
if err != nil {
redisConn.Close()
log.Warningf("migrate key %s error", string(key))
return errors.Trace(err)
}
//handle migrate result
resp, err := parser.Parse(redisReader)
if err != nil {
redisConn.Close()
return errors.Trace(err)
}
result, err := resp.Bytes()
log.Debug("migrate", string(key), "from", shd.migrateFrom.Master(), "to", shd.dst.Master(),
string(result))
if resp.Type == parser.ErrorResp {
redisConn.Close()
log.Error(string(key), string(resp.Raw), "migrateFrom", shd.migrateFrom.Master())
return errors.New(string(resp.Raw))
}
s.counter.Add("Migrate", 1)
}
return nil
}
func TestHandleSpecCommand(t *testing.T) {
var tbl = map[string]string{
"PING": "+PONG\r\n",
"QUIT": string(OK_BYTES),
"SELECT": string(OK_BYTES),
"AUTH": string(OK_BYTES),
}
for k, v := range tbl {
resp, err := parser.Parse(bufio.NewReader(bytes.NewBufferString(k + string(parser.NEW_LINE))))
if err != nil {
t.Error(err)
}
_, keys, err := resp.GetOpKeys()
if err != nil {
t.Error(errors.ErrorStack(err))
}
result := &bytes.Buffer{}
w := &fakeDeadlineReadWriter{w: bufio.NewWriter(result)}
_, _, err = handleSpecCommand(k, w, keys, 5)
if err != nil {
t.Error(err)
}
w.w.Flush()
if string(result.Bytes()) != v {
t.Error("result not match", string(result.Bytes()))
}
}
//"ECHO xxxx": "xxxx\r\n",
{
resp, err := parser.Parse(bufio.NewReader(bytes.NewBufferString("ECHO xxxx\r\n")))
if err != nil {
t.Error(errors.ErrorStack(err))
}
result := &bytes.Buffer{}
w := &fakeDeadlineReadWriter{w: bufio.NewWriter(result)}
_, keys, _ := resp.GetOpKeys()
_, _, err = handleSpecCommand("ECHO", w, keys, 5)
if err != nil {
t.Error(errors.ErrorStack(err))
}
w.w.Flush()
if string(result.Bytes()) != "$4\r\nxxxx\r\n" {
t.Error("result not match", string(result.Bytes()))
}
}
//test empty key
{
resp, err := parser.Parse(bufio.NewReader(bytes.NewBufferString("ECHO\r\n")))
if err != nil {
t.Error(errors.ErrorStack(err))
}
result := &bytes.Buffer{}
w := &fakeDeadlineReadWriter{w: bufio.NewWriter(result)}
_, keys, _ := resp.GetOpKeys()
shouldClose, _, err := handleSpecCommand("ECHO", w, keys, 5)
if !shouldClose {
t.Error(errors.ErrorStack(err))
}
}
//test not specific command
{
result := &bytes.Buffer{}
w := &fakeDeadlineReadWriter{w: bufio.NewWriter(result)}
_, handled, err := handleSpecCommand("get", w, nil, 5)
if handled {
t.Error(errors.ErrorStack(err))
}
}
}