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)) } } }