Пример #1
0
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
}
Пример #2
0
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)
}
Пример #3
0
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)
	}
}
Пример #4
0
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")
	}
}
Пример #5
0
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)
}
Пример #6
0
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
}
Пример #7
0
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))
		}
	}
}