feat: enhance the data split logic

cmd/main.go

@@ -7,6 +7,7 @@ import (
 	"os/signal"
 	"sync"
 	"syscall"
+	"time"
 
 	"github.com/databendcloud/db-archiver/config"
 	"github.com/databendcloud/db-archiver/ingester"
@@ -15,6 +16,9 @@ import (
 )
 
 func main() {
+	start := fmt.Sprintf("start time: %s", time.Now().Format("2006-01-02 15:04:05"))
+	fmt.Println(start)
+	startTime := time.Now()
 	ctx, cancel := context.WithCancel(context.Background())
 	go func() {
 		sigch := make(chan os.Signal, 1)
@@ -37,6 +41,9 @@ func main() {
 		wg.Done()
 	}()
 	wg.Wait()
+	endTime := fmt.Sprintf("end time: %s", time.Now().Format("2006-01-02 15:04:05"))
+	fmt.Println(endTime)
+	fmt.Println(fmt.Sprintf("total time: %s", time.Since(startTime)))
 }
 
 func parseConfigWithFile() *config.Config {

config/conf.json

@@ -16,5 +16,6 @@
   "copyForce": false,
   "disableVariantCheck": false,
   "userStage": "~",
-  "deleteAfterSync": false
+  "deleteAfterSync": false,
+  "maxThread": 2
 }

config/config.go

@@ -30,6 +30,7 @@ type Config struct {
 	DisableVariantCheck bool   `json:"disableVariantCheck" default:"false"`
 	UserStage           string `json:"userStage" default:"~"`
 	DeleteAfterSync     bool   `json:"deleteAfterSync" default:"false"`
+	MaxThread           int    `json:"maxThread" default:"2"`
 }
 
 func LoadConfig() (*Config, error) {

source/source.go

@@ -96,6 +96,21 @@ func (s *Source) GerMinMaxSplitKey() (int, int, error) {
 	return minSplitKey, maxSplitKey, nil
 }
 
+func (s *Source) SlimCondition(minSplitKey, maxSplitKey int) [][]int {
+	var conditions [][]int
+	rangeSize := (maxSplitKey - minSplitKey) / s.cfg.MaxThread
+	for i := 0; i < s.cfg.MaxThread; i++ {
+		lowerBound := minSplitKey + rangeSize*i
+		upperBound := lowerBound + rangeSize
+		if i == s.cfg.MaxThread-1 {
+			// Ensure the last condition includes maxSplitKey
+			upperBound = maxSplitKey
+		}
+		conditions = append(conditions, []int{lowerBound, upperBound})
+	}
+	return conditions
+}
+
 func (s *Source) DeleteAfterSync() error {
 	if s.cfg.DeleteAfterSync {
 		_, err := s.db.Exec(fmt.Sprintf("delete from %s.%s where %s", s.cfg.SourceDB,
@@ -182,6 +197,26 @@ func (s *Source) SplitCondition(minSplitKey, maxSplitKey int) []string {
 	return conditions
 }
 
+func (s *Source) SplitConditionAccordingMaxGoRoutine(minSplitKey, maxSplitKey, allMax int) []string {
+	var conditions []string
+	for {
+		if minSplitKey >= maxSplitKey {
+			if minSplitKey > allMax {
+				return conditions
+			}
+			if maxSplitKey == allMax {
+				conditions = append(conditions, fmt.Sprintf("(%s >= %d and %s <= %d)", s.cfg.SourceSplitKey, minSplitKey, s.cfg.SourceSplitKey, maxSplitKey))
+			} else {
+				conditions = append(conditions, fmt.Sprintf("(%s >= %d and %s < %d)", s.cfg.SourceSplitKey, minSplitKey, s.cfg.SourceSplitKey, maxSplitKey))
+			}
+			break
+		}
+		conditions = append(conditions, fmt.Sprintf("(%s >= %d and %s < %d)", s.cfg.SourceSplitKey, minSplitKey, s.cfg.SourceSplitKey, minSplitKey+s.cfg.BatchSize-1))
+		minSplitKey += s.cfg.BatchSize - 1
+	}
+	return conditions
+}
+
 func GenerateJSONFile(columns []string, data [][]interface{}) (string, int, error) {
 	l := logrus.WithFields(logrus.Fields{"tardatabend": "IngestData"})
 	var batchJsonData []string

worker/worker.go

@@ -51,12 +51,43 @@ func (w *Worker) stepBatchWithCondition(conditionSql string) error {
 	return nil
 }
 
+func (w *Worker) IsSplitAccordingMaxGoRoutine(minSplitKey, maxSplitKey, batchSize int) bool {
+	return (maxSplitKey-minSplitKey)/batchSize > w.cfg.MaxThread
+}
+
 func (w *Worker) stepBatch() error {
 	wg := &sync.WaitGroup{}
 	minSplitKey, maxSplitKey, err := w.src.GerMinMaxSplitKey()
 	if err != nil {
 		return err
 	}
+	fmt.Println("minSplitKey", minSplitKey, "maxSplitKey", maxSplitKey)
+
+	if w.IsSplitAccordingMaxGoRoutine(minSplitKey, maxSplitKey, w.cfg.BatchSize) {
+		fmt.Println("split according maxGoRoutine", w.cfg.MaxThread)
+		slimedRange := w.src.SlimCondition(minSplitKey, maxSplitKey)
+		fmt.Println("slimedRange", slimedRange)
+		wg.Add(w.cfg.MaxThread)
+		for i := 0; i < w.cfg.MaxThread; i++ {
+			go func(idx int) {
+				defer wg.Done()
+				conditions := w.src.SplitConditionAccordingMaxGoRoutine(slimedRange[idx][0], slimedRange[idx][1], maxSplitKey)
+				logrus.Infof("conditions in one routine: %v", conditions)
+				if err != nil {
+					logrus.Errorf("stepBatchWithCondition failed: %v", err)
+				}
+				for _, condition := range conditions {
+					logrus.Infof("condition: %s", condition)
+					err := w.stepBatchWithCondition(condition)
+					if err != nil {
+						logrus.Errorf("stepBatchWithCondition failed: %v", err)
+					}
+				}
+			}(i)
+		}
+		wg.Wait()
+		return nil
+	}
 	conditions := w.src.SplitCondition(minSplitKey, maxSplitKey)
 	for _, condition := range conditions {
 		wg.Add(1)