address comment

ddl/ddl_api.go

@@ -1202,14 +1202,26 @@ func (d *ddl) CreateTable(ctx sessionctx.Context, s *ast.CreateTableStmt) (err e
 
 	err = d.doDDLJob(ctx, job)
 	if err == nil {
+		var preSplitAndScatter func()
 		// do pre-split and scatter.
 		if tbInfo.ShardRowIDBits > 0 && tbInfo.PreSplitRegions > 0 {
+			preSplitAndScatter = func() { preSplitTableRegion(d.store, tbInfo, ctx.GetSessionVars().WaitTableSplitFinish) }
+		} else if atomic.LoadUint32(&EnableSplitTableRegion) != 0 {
+			pi := tbInfo.GetPartitionInfo()
+			if pi != nil {
+				preSplitAndScatter = func() { splitPartitionTableRegion(d.store, pi) }
+			} else {
+				preSplitAndScatter = func() { splitTableRegion(d.store, tbInfo.ID) }
+			}
+		}
+		if preSplitAndScatter != nil {
 			if ctx.GetSessionVars().WaitTableSplitFinish {
-				preSplitTableRegion(d.store, tbInfo, true)
+				preSplitAndScatter()
 			} else {
-				go preSplitTableRegion(d.store, tbInfo, false)
+				go preSplitAndScatter()
 			}
 		}
+
 		if tbInfo.AutoIncID > 1 {
 			// Default tableAutoIncID base is 0.
 			// If the first ID is expected to greater than 1, we need to do rebase.

ddl/table.go

@@ -17,7 +17,6 @@ import (
 	"fmt"
 	"strconv"
 	"strings"
-	"sync/atomic"
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/failpoint"
@@ -74,20 +73,6 @@ func onCreateTable(d *ddlCtx, t *meta.Meta, job *model.Job) (ver int64, _ error)
 		if err != nil {
 			return ver, errors.Trace(err)
 		}
-		if atomic.LoadUint32(&EnableSplitTableRegion) != 0 {
-			// TODO: Add restrictions to this operation.
-			pi := tbInfo.GetPartitionInfo()
-			if pi != nil {
-				// Max partition count is 4096, should we sample and just choose some of the partition to split?
-				go func(pi *model.PartitionInfo) {
-					for _, def := range pi.Definitions {
-						splitTableRegion(d.store, def.ID)
-					}
-				}(pi)
-			} else {
-				go splitTableRegion(d.store, tbInfo.ID)
-			}
-		}
 		// Finish this job.
 		job.FinishTableJob(model.JobStateDone, model.StatePublic, ver, tbInfo)
 		asyncNotifyEvent(d, &util.Event{Tp: model.ActionCreateTable, TableInfo: tbInfo})
@@ -350,6 +335,13 @@ type splitableStore interface {
 	WaitScatterRegionFinish(regionID uint64) error
 }
 
+func splitPartitionTableRegion(store kv.Storage, pi *model.PartitionInfo) {
+	// Max partition count is 4096, should we sample and just choose some of the partition to split?
+	for _, def := range pi.Definitions {
+		splitTableRegion(store, def.ID)
+	}
+}
+
 func splitTableRegion(store kv.Storage, tableID int64) {
 	s, ok := store.(splitableStore)
 	if !ok {