tidb,config: limit statement count in a transaction (#5754)

* tidb,config: limit statement count in a transaction
* ast: do statement is read only (#5752)

ast/read_only_checker.go

@@ -27,7 +27,7 @@ func IsReadOnly(node Node) bool {
 
 		node.Accept(&checker)
 		return checker.readOnly
-	case *ExplainStmt:
+	case *ExplainStmt, *DoStmt:
 		return true
 	default:
 		return false

ast/read_only_checker_test.go

@@ -38,4 +38,7 @@ func (s *testCacheableSuite) TestCacheable(c *C) {
 
 	stmt = &ExplainStmt{}
 	c.Assert(IsReadOnly(stmt), IsTrue)
+
+	stmt = &DoStmt{}
+	c.Assert(IsReadOnly(stmt), IsTrue)
 }

config/config.go

@@ -80,6 +80,7 @@ type Performance struct {
 	CrossJoin       bool   `toml:"cross-join" json:"cross-join"`
 	StatsLease      string `toml:"stats-lease" json:"stats-lease"`
 	RunAutoAnalyze  bool   `toml:"run-auto-analyze" json:"run-auto-analyze"`
+	StmtCountLimit  int    `toml:"stmt-count-limit" json:"stmt-count-limit"`
 }
 
 // XProtocol is the XProtocol section of the config.
@@ -126,6 +127,7 @@ var defaultConf = Config{
 		CrossJoin:       true,
 		StatsLease:      "3s",
 		RunAutoAnalyze:  true,
+		StmtCountLimit:  5000,
 	},
 	XProtocol: XProtocol{
 		XHost: "0.0.0.0",

config/config.toml.example

@@ -92,6 +92,8 @@ metrics-interval = 15
 [performance]
 # Set keep alive option for tcp connection.
 tcp-keep-alive = true
+# StmtCountLimit limits the max count of statement inside a transaction.
+stmt-count-limit = 5000
 
 # The maximum number of retries when commit a transaction.
 retry-limit = 10

new_session_test.go

@@ -22,6 +22,7 @@ import (
 
 	. "github.com/pingcap/check"
 	"github.com/pingcap/tidb"
+	"github.com/pingcap/tidb/config"
 	"github.com/pingcap/tidb/context"
 	"github.com/pingcap/tidb/domain"
 	"github.com/pingcap/tidb/executor"
@@ -1351,6 +1352,29 @@ type testSchemaSuite struct {
 	checkLeak func()
 }
 
+func (s *testSessionSuite) TestStatementCountLimit(c *C) {
+	tk := testkit.NewTestKitWithInit(c, s.store)
+	tk.MustExec("create table stmt_count_limit (id int)")
+	saved := config.GetGlobalConfig().Performance.StmtCountLimit
+	config.GetGlobalConfig().Performance.StmtCountLimit = 3
+	defer func() {
+		config.GetGlobalConfig().Performance.StmtCountLimit = saved
+	}()
+	tk.MustExec("begin")
+	tk.MustExec("insert into stmt_count_limit values (1)")
+	tk.MustExec("insert into stmt_count_limit values (2)")
+	_, err := tk.Exec("insert into stmt_count_limit values (3)")
+	c.Assert(err, NotNil)
+
+	// begin is counted into history but this one is not.
+	tk.MustExec("SET SESSION autocommit = false")
+	tk.MustExec("insert into stmt_count_limit values (1)")
+	tk.MustExec("insert into stmt_count_limit values (2)")
+	tk.MustExec("insert into stmt_count_limit values (3)")
+	_, err = tk.Exec("insert into stmt_count_limit values (4)")
+	c.Assert(err, NotNil)
+}
+
 func (s *testSchemaSuite) TearDownTest(c *C) {
 	tk := testkit.NewTestKitWithInit(c, s.store)
 	r := tk.MustQuery("show tables")

tidb.go

@@ -26,6 +26,7 @@ import (
 	log "github.com/Sirupsen/logrus"
 	"github.com/juju/errors"
 	"github.com/pingcap/tidb/ast"
+	"github.com/pingcap/tidb/config"
 	"github.com/pingcap/tidb/context"
 	"github.com/pingcap/tidb/domain"
 	"github.com/pingcap/tidb/executor"
@@ -170,6 +171,16 @@ func runStmt(ctx context.Context, s ast.Statement) (ast.RecordSet, error) {
 		} else {
 			err = se.CommitTxn()
 		}
+	} else {
+		// If the user insert, insert, insert ... but never commit, TiDB would OOM.
+		// So we limit the statement count in a transaction here.
+		history := GetHistory(ctx)
+		if history.Count() > config.GetGlobalConfig().Performance.StmtCountLimit {
+			err1 := se.RollbackTxn()
+			terror.Log(errors.Trace(err1))
+			return rs, errors.Errorf("statement count %d exceeds the transaction limitation, autocommit = %t",
+				history.Count(), ctx.GetSessionVars().IsAutocommit())
+		}
 	}
 	return rs, errors.Trace(err)
 }