optimize count file bytes and isDone

vesoft-inc · Mar 11, 2022 · 06d7eb2 · 06d7eb2
1 parent ca82699
commit 06d7eb2
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 12 deletions.
diff --git a/pkg/cmd/runner.go b/pkg/cmd/runner.go
@@ -85,9 +85,10 @@ func (r *Runner) Run(yaml *config.YAMLConfig) {
 
 	r.Readers = freaders
 	r.stataMgr = statsMgr
-	r.stataMgr.CountFileBytes(r.Readers)
+
 	<-statsMgr.DoneCh
 
+	r.stataMgr.CountFileBytes(r.Readers)
 	r.Readers = nil
 	r.NumFailed = statsMgr.Stats.NumFailed
 
@@ -99,7 +100,10 @@ func (r *Runner) Run(yaml *config.YAMLConfig) {
 
 func (r *Runner) QueryStats() *stats.Stats {
 	if r.stataMgr != nil {
-		if r.Readers == nil {
+		if r.Readers != nil {
+			r.stataMgr.CountFileBytes(r.Readers)
+		}
+		if r.stataMgr.Done == true {
 			return &r.stataMgr.Stats
 		}
 		r.stataMgr.StatsCh <- base.NewOutputStats()

diff --git a/pkg/csv/reader.go b/pkg/csv/reader.go
@@ -3,12 +3,14 @@ package csv
 import (
 	"bufio"
 	"encoding/csv"
+	"errors"
 	"fmt"
+	"io"
+	"os"
+
 	"github.com/vesoft-inc/nebula-importer/pkg/base"
 	"github.com/vesoft-inc/nebula-importer/pkg/config"
 	"github.com/vesoft-inc/nebula-importer/pkg/logger"
-	"io"
-	"os"
 )
 
 type CSVReader struct {
@@ -19,6 +21,7 @@ type CSVReader struct {
 	br        *bufio.Reader
 	totalBytes int64
 	initComplete bool
+
 }
 
 type recordReader struct {
@@ -50,7 +53,9 @@ func (r *CSVReader) InitReader(file *os.File) {
 		logger.Infof("The stat of %s is wrong, %s", file.Name(), err)
 	}
 	r.totalBytes = stat.Size()
-	r.initComplete = true
+	defer func() {
+		r.initComplete = true
+	}()
 }
 
 func (r *CSVReader) ReadLine() (base.Data, error) {
@@ -88,10 +93,9 @@ func (r *CSVReader) ReadLine() (base.Data, error) {
 	}
 }
 
-func (r *CSVReader) TotalBytes() (int64) {
-	for {
-		if r.initComplete {
-			return r.totalBytes
-		}
+func (r *CSVReader) TotalBytes() (int64, error) {
+	if r.initComplete {
+		return r.totalBytes, nil
 	}
+	return 0, errors.New("init not complete")
 }
diff --git a/pkg/reader/reader.go b/pkg/reader/reader.go
@@ -20,7 +20,7 @@ import (
 type DataFileReader interface {
 	InitReader(*os.File)
 	ReadLine() (base.Data, error)
-	TotalBytes() (int64)
+	TotalBytes() (int64, error)
 }
 
 // FIXME: private fields

diff --git a/pkg/stats/statsmgr.go b/pkg/stats/statsmgr.go
@@ -15,6 +15,8 @@ type StatsMgr struct {
 	StatsCh       chan base.Stats
 	DoneCh        chan bool
 	Stats         Stats
+	Done bool
+	CountFileDone bool
 }
 
 type Stats struct {
@@ -52,6 +54,8 @@ func NewStatsMgr(files []*config.File) *StatsMgr {
 func (s *StatsMgr) Close() {
 	close(s.StatsCh)
 	close(s.DoneCh)
+	close(s.OutputStatsCh)
+	s.Done = true
 }
 
 func (s *StatsMgr) updateStat(stat base.Stats) {
@@ -86,11 +90,25 @@ func (s *StatsMgr) print(prefix string, now time.Time) {
 }
 
 func (s *StatsMgr) CountFileBytes(freaders []*reader.FileReader) {
+	if s.CountFileDone {
+		return
+	}
+	flag := true
+	s.Stats.TotalCount = 0
 	for _, r := range freaders {
 		if r == nil {
 			continue
 		}
-		s.Stats.TotalBytes += r.DataReader.TotalBytes()
+		bytes, err := r.DataReader.TotalBytes()
+		if err != nil {
+			logger.Infof("file %s count file bytes fail, $s", *r.File.Path, err)
+			flag = false
+			continue
+		}
+		s.Stats.TotalBytes += bytes
+	}
+	if flag {
+		s.CountFileDone = true
 	}
 }