qps/rate limit & reduce expire seek range & expire-key hash

bin/titan/main.go

@@ -68,11 +68,19 @@ func main() {
 	}
 
 	svr := metrics.NewServer(&config.Status)
-
+	limitersMgr, err := db.NewLimitersMgr(store, &config.Tikv.RateLimit)
+	if err != nil {
+		zap.L().Fatal("create limitersMgr failed", zap.Error(err))
+		os.Exit(1)
+	}
 	serv := titan.New(&context.ServerContext{
-		RequirePass:      config.Server.Auth,
-		Store:            store,
-		ListZipThreshold: config.Server.ListZipThreshold,
+		RequirePass:       config.Server.Auth,
+		Store:             store,
+		ListZipThreshold:  config.Server.ListZipThreshold,
+		LimitersMgr:       limitersMgr,
+		LimitConnection:   config.Server.LimitConnection,
+		MaxConnection:     config.Server.MaxConnection,
+		MaxConnectionWait: config.Server.MaxConnectionWait,
 	})
 
 	var servOpts, statusOpts []continuous.ServerOption

client.go

@@ -6,7 +6,6 @@ import (
 	"io/ioutil"
 	"net"
 	"strings"
-	"sync"
 	"time"
 
 	"github.com/distributedio/titan/command"
@@ -22,54 +21,43 @@ type client struct {
 	exec   *command.Executor
 	r      *bufio.Reader
 
-	eofLock sync.Mutex //the lock of reading_writing 'eof'
-	eof     bool       //is over when read data from socket
+	remoteClosed bool //is the connection closed by remote peer?
 }
 
 func newClient(cliCtx *context.ClientContext, s *Server, exec *command.Executor) *client {
 	return &client{
-		cliCtx: cliCtx,
-		server: s,
-		exec:   exec,
-		eof:    false,
+		cliCtx:       cliCtx,
+		server:       s,
+		exec:         exec,
+		remoteClosed: false,
 	}
 }
 
-func (c *client) readEof() {
-	c.eofLock.Lock()
-	defer c.eofLock.Unlock()
-
-	c.eof = true
-}
-
-func (c *client) isEof() bool {
-	c.eofLock.Lock()
-	defer c.eofLock.Unlock()
-
-	return c.eof
-}
-
 // Write to conn and log error if needed
 func (c *client) Write(p []byte) (int, error) {
 	zap.L().Debug("write to client", zap.Int64("clientid", c.cliCtx.ID), zap.String("msg", string(p)))
 	n, err := c.conn.Write(p)
 	if err != nil {
-		c.conn.Close()
 		if err == io.EOF {
-			zap.L().Info("close connection", zap.String("addr", c.cliCtx.RemoteAddr),
-				zap.Int64("clientid", c.cliCtx.ID))
+			zap.L().Info("connection was half-closed by remote peer", zap.String("addr", c.cliCtx.RemoteAddr),
+				zap.Int64("clientid", c.cliCtx.ID), zap.String("namespace", c.cliCtx.Namespace))
 		} else {
+			//may be unknown error with message "connection reset by peer"
 			zap.L().Error("write net failed", zap.String("addr", c.cliCtx.RemoteAddr),
 				zap.Int64("clientid", c.cliCtx.ID),
 				zap.String("namespace", c.cliCtx.Namespace),
 				zap.Bool("multi", c.cliCtx.Multi),
 				zap.Bool("watching", c.cliCtx.Txn != nil),
 				zap.String("command", c.cliCtx.LastCmd),
 				zap.String("error", err.Error()))
-			return 0, err
 		}
+		//client.serve() will get the channel close event, close the connection, exit current go routine
+		//if the remote client use pipeline to invoke command, then close the connection(timeout etc), titan still get command from client.bufio.Reader and process
+		//setting client.remoteClosed to true will help client.serve() to interrupt command processing
+		c.remoteClosed = true
 	}
-	return n, nil
+	//return err for above write() error, then replying many times command can break its sending to a half-closed connection, etc BytesArray(lrange invoke it).
+	return n, err
 }
 
 func (c *client) serve(conn net.Conn) error {
@@ -78,21 +66,27 @@ func (c *client) serve(conn net.Conn) error {
 
 	var cmd []string
 	var err error
+	unknownCmdTimes := int(0)
 	for {
 		select {
 		case <-c.cliCtx.Done:
 			return c.conn.Close()
 		default:
+			if c.remoteClosed {
+				zap.L().Info("close connection", zap.String("addr", c.cliCtx.RemoteAddr),
+					zap.Int64("clientid", c.cliCtx.ID), zap.String("namespace", c.cliCtx.Namespace))
+				return c.conn.Close()
+			}
 			cmd, err = c.readCommand()
 			if err != nil {
 				c.conn.Close()
 				if err == io.EOF {
 					zap.L().Info("close connection", zap.String("addr", c.cliCtx.RemoteAddr),
-						zap.Int64("clientid", c.cliCtx.ID))
+						zap.Int64("clientid", c.cliCtx.ID), zap.String("namespace", c.cliCtx.Namespace))
 					return nil
 				}
 				zap.L().Error("read command failed", zap.String("addr", c.cliCtx.RemoteAddr),
-					zap.Int64("clientid", c.cliCtx.ID), zap.Error(err))
+					zap.Int64("clientid", c.cliCtx.ID), zap.String("namespace", c.cliCtx.Namespace), zap.Error(err))
 				return err
 			}
 		}
@@ -102,8 +96,30 @@ func (c *client) serve(conn net.Conn) error {
 			c.server.servCtx.Pause = 0
 		}
 
+		if len(cmd) <= 0 {
+			err := command.ErrEmptyCommand
+			zap.L().Error(err.Error(), zap.String("addr", c.cliCtx.RemoteAddr),
+				zap.Int64("clientid", c.cliCtx.ID))
+			resp.ReplyError(c, err.Error())
+			c.conn.Close()
+			return nil
+		}
+
 		c.cliCtx.Updated = time.Now()
 		c.cliCtx.LastCmd = cmd[0]
+		if !c.exec.CanExecute(c.cliCtx.LastCmd) {
+			err := command.ErrUnKnownCommand(c.cliCtx.LastCmd)
+			zap.L().Error(err.Error(), zap.String("addr", c.cliCtx.RemoteAddr),
+				zap.Int64("clientid", c.cliCtx.ID))
+			resp.ReplyError(c, err.Error())
+			unknownCmdTimes++
+			if unknownCmdTimes >= 3 {
+				c.conn.Close()
+				return nil
+			} else {
+				continue
+			}
+		}
 
 		ctx := &command.Context{
 			Name:    cmd[0],
@@ -114,7 +130,6 @@ func (c *client) serve(conn net.Conn) error {
 		}
 
 		ctx.Context = context.New(c.cliCtx, c.server.servCtx)
-		zap.L().Debug("recv msg", zap.String("command", ctx.Name), zap.Strings("arguments", ctx.Args))
 
 		// Skip reply if necessary
 		if c.cliCtx.SkipN != 0 {
@@ -127,7 +142,8 @@ func (c *client) serve(conn net.Conn) error {
 			env.Write(zap.String("addr", c.cliCtx.RemoteAddr),
 				zap.Int64("clientid", c.cliCtx.ID),
 				zap.String("traceid", ctx.TraceID),
-				zap.String("command", ctx.Name))
+				zap.String("command", ctx.Name),
+				zap.Strings("arguments", ctx.Args))
 		}
 
 		c.exec.Execute(ctx)

command/command.go

@@ -61,28 +61,27 @@ func Integer(w io.Writer, v int64) OnCommit {
 // BytesArray replies a [][]byte when commit
 func BytesArray(w io.Writer, a [][]byte) OnCommit {
 	return func() {
-		start := time.Now()
-		resp.ReplyArray(w, len(a))
-		zap.L().Debug("reply array size", zap.Int64("cost(us)", time.Since(start).Nanoseconds()/1000))
-		start = time.Now()
+		if _, err := resp.ReplyArray(w, len(a)); err != nil {
+			return
+		}
 		for i := range a {
 			if a[i] == nil {
-				resp.ReplyNullBulkString(w)
+				if err := resp.ReplyNullBulkString(w); err != nil {
+					return
+				}
 				continue
 			}
-			resp.ReplyBulkString(w, string(a[i]))
-			if i % 10 == 9 {
-				zap.L().Debug("reply 10 bulk string", zap.Int64("cost(us)", time.Since(start).Nanoseconds()/1000))
-				start = time.Now()
+			if err := resp.ReplyBulkString(w, string(a[i])); err != nil {
+				return
 			}
 		}
 	}
 }
 
 func BytesArrayOnce(w io.Writer, a [][]byte) OnCommit {
-    return func() {
-        resp.ReplyStringArray(w, a)
-    }
+	return func() {
+		resp.ReplyStringArray(w, a)
+	}
 }
 
 // TxnCommand runs a command in transaction
@@ -92,6 +91,10 @@ type TxnCommand func(ctx *Context, txn *db.Transaction) (OnCommit, error)
 func Call(ctx *Context) {
 	ctx.Name = strings.ToLower(ctx.Name)
 
+	if _, ok := txnCommands[ctx.Name]; ok && ctx.Server.LimitersMgr != nil {
+		ctx.Server.LimitersMgr.CheckLimit(ctx.Client.Namespace, ctx.Name, ctx.Args)
+	}
+
 	if ctx.Name != "auth" &&
 		ctx.Server.RequirePass != "" &&
 		ctx.Client.Authenticated == false {
@@ -182,18 +185,16 @@ func AutoCommit(cmd TxnCommand) Command {
 	return func(ctx *Context) {
 		retry.Ensure(ctx, func() error {
 			mt := metrics.GetMetrics()
-            start := time.Now()
+			start := time.Now()
 			txn, err := ctx.Client.DB.Begin()
 			key := ""
 			if len(ctx.Args) > 0 {
 				key = ctx.Args[0]
-				if len(ctx.Args) > 1 {
-					mt.CommandArgsNumHistogramVec.WithLabelValues(ctx.Client.Namespace, ctx.Name).Observe(float64(len(ctx.Args)-1))
-				}
+				mt.CommandArgsNumHistogramVec.WithLabelValues(ctx.Client.Namespace, ctx.Name).Observe(float64(len(ctx.Args)))
 			}
 			cost := time.Since(start).Seconds()
-            zap.L().Debug("transation begin", zap.String("name", ctx.Name), zap.String("key", key), zap.Int64("cost(us)", int64(cost*1000000)))
-            mt.TxnBeginHistogramVec.WithLabelValues(ctx.Client.Namespace, ctx.Name).Observe(cost)
+			mt.TxnBeginHistogramVec.WithLabelValues(ctx.Client.Namespace, ctx.Name).Observe(cost)
+			zap.L().Debug("transation begin", zap.String("name", ctx.Name), zap.String("key", key), zap.Int64("cost(us)", int64(cost*1000000)))
 			if err != nil {
 				mt.TxnFailuresCounterVec.WithLabelValues(ctx.Client.Namespace, ctx.Name).Inc()
 				resp.ReplyError(ctx.Out, "ERR "+err.Error())
@@ -208,8 +209,8 @@ func AutoCommit(cmd TxnCommand) Command {
 			start = time.Now()
 			onCommit, err := cmd(ctx, txn)
 			cost = time.Since(start).Seconds()
-			zap.L().Debug("command done", zap.String("name", ctx.Name), zap.String("key", key), zap.Int64("cost(us)", int64(cost*1000000)))
 			mt.CommandFuncDoneHistogramVec.WithLabelValues(ctx.Client.Namespace, ctx.Name).Observe(cost)
+			zap.L().Debug("command done", zap.String("name", ctx.Name), zap.String("key", key), zap.Int64("cost(us)", int64(cost*1000000)))
 			if err != nil {
 				mt.TxnFailuresCounterVec.WithLabelValues(ctx.Client.Namespace, ctx.Name).Inc()
 				resp.ReplyError(ctx.Out, err.Error())
@@ -257,8 +258,8 @@ func AutoCommit(cmd TxnCommand) Command {
 				onCommit()
 			}
 			cost = time.Since(start).Seconds()
-			zap.L().Debug("onCommit ", zap.String("name", ctx.Name), zap.String("key", key), zap.Int64("cost(us)", int64(cost*1000000)))
 			mt.ReplyFuncDoneHistogramVec.WithLabelValues(ctx.Client.Namespace, ctx.Name).Observe(cost)
+			zap.L().Debug("onCommit ", zap.String("name", ctx.Name), zap.String("key", key), zap.Int64("cost(us)", int64(cost*1000000)))
 			mtFunc()
 			return nil
 		})
@@ -277,9 +278,9 @@ func feedMonitors(ctx *Context) {
 		id := strconv.FormatInt(int64(ctx.Client.DB.ID), 10)
 
 		line := ts + " [" + id + " " + ctx.Client.RemoteAddr + "]" + " " + ctx.Name + " " + strings.Join(ctx.Args, " ")
-        start := time.Now()
+		start := time.Now()
 		err := resp.ReplySimpleString(mCtx.Out, line)
-        zap.L().Debug("feedMonitors reply", zap.String("name", ctx.Name), zap.Int64("cost(us)", time.Since(start).Nanoseconds()/1000))
+		zap.L().Debug("feedMonitors reply", zap.String("name", ctx.Name), zap.Int64("cost(us)", time.Since(start).Nanoseconds()/1000))
 		if err != nil {
 			ctx.Server.Monitors.Delete(k)
 		}
@@ -299,6 +300,12 @@ func NewExecutor() *Executor {
 	return &Executor{txnCommands: txnCommands, commands: commands}
 }
 
+func (e *Executor) CanExecute(cmd string) bool {
+	lowerName := strings.ToLower(cmd)
+	_, ok := commands[lowerName]
+	return ok
+}
+
 // Execute a command
 func (e *Executor) Execute(ctx *Context) {
 	start := time.Now()

command/error.go

@@ -91,6 +91,8 @@ var (
 	// ErrEmptyArray error
 	ErrEmptyArray = errors.New("EmptyArray error")
 
+	ErrEmptyCommand = errors.New("command is empty")
+
 	//ErrExec exec without multi
 	ErrExec = errors.New("ERR EXEC without MULTI")
 

command/init.go

@@ -28,16 +28,16 @@ func init() {
 		"getset":   GetSet,
 		"getrange": GetRange,
 		// "msetnx":   MSetNx,
-		"setnx":    SetNx,
-		"setex":    SetEx,
-		"psetex":   PSetEx,
-		"setrange": SetRange,
-		"setbit":   SetBit,
+		"setnx":  SetNx,
+		"setex":  SetEx,
+		"psetex": PSetEx,
+		//"setrange": SetRange,
+		//"setbit":   SetBit,
 		// "bitop":       BitOp,
 		// "bitfield":    BitField,
-		"getbit":      GetBit,
-		"bitpos":      BitPos,
-		"bitcount":    BitCount,
+		//"getbit":      GetBit,
+		//"bitpos":      BitPos,
+		//"bitcount":    BitCount,
 		"incr":        Incr,
 		"incrby":      IncrBy,
 		"decr":        Decr,
@@ -117,6 +117,8 @@ func init() {
 
 		// transactions, exec and discard should called explicitly, so they are registered here
 		"multi":   Desc{Proc: Multi, Cons: Constraint{1, flags("sF"), 0, 0, 0}},
+		"exec":    Desc{Proc: Exec, Cons: Constraint{1, flags("sF"), 0, 0, 0}},
+		"discard": Desc{Proc: Discard, Cons: Constraint{1, flags("sF"), 0, 0, 0}},
 		"watch":   Desc{Proc: Watch, Cons: Constraint{-2, flags("sF"), 1, -1, 1}},
 		"unwatch": Desc{Proc: Unwatch, Cons: Constraint{1, flags("sF"), 0, 0, 0}},
 
@@ -133,29 +135,29 @@ func init() {
 		"rpushx":  Desc{Proc: AutoCommit(RPushx), Cons: Constraint{-3, flags("wmF"), 1, 1, 1}},
 
 		// strings
-		"get":         Desc{Proc: AutoCommit(Get), Cons: Constraint{2, flags("rF"), 1, 1, 1}},
-		"set":         Desc{Proc: AutoCommit(Set), Cons: Constraint{-3, flags("wm"), 1, 1, 1}},
-		"setnx":       Desc{Proc: AutoCommit(SetNx), Cons: Constraint{3, flags("wmF"), 1, 1, 1}},
-		"setex":       Desc{Proc: AutoCommit(SetEx), Cons: Constraint{4, flags("wm"), 1, 1, 1}},
-		"psetex":      Desc{Proc: AutoCommit(PSetEx), Cons: Constraint{4, flags("wm"), 1, 1, 1}},
-		"mget":        Desc{Proc: AutoCommit(MGet), Cons: Constraint{-2, flags("rF"), 1, -1, 1}},
-		"mset":        Desc{Proc: AutoCommit(MSet), Cons: Constraint{-3, flags("wm"), 1, -1, 2}},
-		"msetnx":      Desc{Proc: AutoCommit(MSetNx), Cons: Constraint{-3, flags("wm"), 1, -1, 2}},
-		"strlen":      Desc{Proc: AutoCommit(Strlen), Cons: Constraint{2, flags("rF"), 1, 1, 1}},
-		"append":      Desc{Proc: AutoCommit(Append), Cons: Constraint{3, flags("wm"), 1, 1, 1}},
-		"setrange":    Desc{Proc: AutoCommit(SetRange), Cons: Constraint{4, flags("wm"), 1, 1, 1}},
+		"get":    Desc{Proc: AutoCommit(Get), Cons: Constraint{2, flags("rF"), 1, 1, 1}},
+		"set":    Desc{Proc: AutoCommit(Set), Cons: Constraint{-3, flags("wm"), 1, 1, 1}},
+		"setnx":  Desc{Proc: AutoCommit(SetNx), Cons: Constraint{3, flags("wmF"), 1, 1, 1}},
+		"setex":  Desc{Proc: AutoCommit(SetEx), Cons: Constraint{4, flags("wm"), 1, 1, 1}},
+		"psetex": Desc{Proc: AutoCommit(PSetEx), Cons: Constraint{4, flags("wm"), 1, 1, 1}},
+		"mget":   Desc{Proc: AutoCommit(MGet), Cons: Constraint{-2, flags("rF"), 1, -1, 1}},
+		"mset":   Desc{Proc: AutoCommit(MSet), Cons: Constraint{-3, flags("wm"), 1, -1, 2}},
+		//"msetnx":      Desc{Proc: AutoCommit(MSetNx), Cons: Constraint{-3, flags("wm"), 1, -1, 2}}, //run test in tests/redis/unit/type/string failed
+		"strlen": Desc{Proc: AutoCommit(Strlen), Cons: Constraint{2, flags("rF"), 1, 1, 1}},
+		"append": Desc{Proc: AutoCommit(Append), Cons: Constraint{3, flags("wm"), 1, 1, 1}},
+		//"setrange":    Desc{Proc: AutoCommit(SetRange), Cons: Constraint{4, flags("wm"), 1, 1, 1}}, //run test in tests/redis/unit/type/string failed
 		"getrange":    Desc{Proc: AutoCommit(GetRange), Cons: Constraint{4, flags("r"), 1, 1, 1}},
 		"incr":        Desc{Proc: AutoCommit(Incr), Cons: Constraint{2, flags("wmF"), 1, 1, 1}},
 		"decr":        Desc{Proc: AutoCommit(Decr), Cons: Constraint{2, flags("wmF"), 1, 1, 1}},
 		"incrby":      Desc{Proc: AutoCommit(IncrBy), Cons: Constraint{3, flags("wmF"), 1, 1, 1}},
 		"decrby":      Desc{Proc: AutoCommit(DecrBy), Cons: Constraint{3, flags("wmF"), 1, 1, 1}},
 		"incrbyfloat": Desc{Proc: AutoCommit(IncrByFloat), Cons: Constraint{3, flags("wmF"), 1, 1, 1}},
-		"setbit":      Desc{Proc: AutoCommit(SetBit), Cons: Constraint{4, flags("wm"), 1, 1, 1}},
+		//"setbit":      Desc{Proc: AutoCommit(SetBit), Cons: Constraint{4, flags("wm"), 1, 1, 1}},
 		// "bitop":       Desc{Proc: AutoCommit(BitOp), Cons: Constraint{-4, flags("wm"), 2, -1, 1}},
 		// "bitfield":    Desc{Proc: AutoCommit(BitField), Cons: Constraint{-2, flags("wm"), 1, 1, 1}},
-		"getbit":   Desc{Proc: AutoCommit(GetBit), Cons: Constraint{3, flags("r"), 1, 1, 1}},
-		"bitcount": Desc{Proc: AutoCommit(BitCount), Cons: Constraint{-2, flags("r"), 1, 1, 1}},
-		"bitpos":   Desc{Proc: AutoCommit(BitPos), Cons: Constraint{-3, flags("r"), 1, 1, 1}},
+		//"getbit":   Desc{Proc: AutoCommit(GetBit), Cons: Constraint{3, flags("r"), 1, 1, 1}},
+		//"bitcount": Desc{Proc: AutoCommit(BitCount), Cons: Constraint{-2, flags("r"), 1, 1, 1}},
+		//"bitpos":   Desc{Proc: AutoCommit(BitPos), Cons: Constraint{-3, flags("r"), 1, 1, 1}},
 
 		// keys
 		"type":      Desc{Proc: AutoCommit(Type), Cons: Constraint{2, flags("rF"), 1, 1, 1}},

command/strings.go

@@ -337,6 +337,9 @@ func PSetEx(ctx *Context, txn *db.Transaction) (OnCommit, error) {
 	if err != nil {
 		return nil, ErrInteger
 	}
+	if ui <= 0 {
+		return nil, ErrExpireSetEx
+	}
 	unit := ui * uint64(time.Millisecond)
 	if err := s.Set([]byte(ctx.Args[2]), int64(unit)); err != nil {
 		return nil, errors.New("ERR " + err.Error())