Add arg scanner-buffer-size, by defautl is 1024*1024, for reading lon…

…g lines

cmd/load_influx/main.go

@@ -12,7 +12,7 @@ import (
 	"strings"
 	"sync"
 	"time"
-	
+
 	"github.com/blagojts/viper"
 	"github.com/cnosdb/tsdb-comparisons/internal/utils"
 	"github.com/cnosdb/tsdb-comparisons/load"
@@ -30,6 +30,7 @@ var (
 	useGzip           bool
 	doAbortOnExist    bool
 	consistency       string
+	scannerBufferSize int
 )
 
 // Global vars
@@ -57,41 +58,49 @@ func init() {
 	config.AddToFlagSet(pflag.CommandLine)
 	target.TargetSpecificFlags("", pflag.CommandLine)
 	var csvDaemonURLs string
-	
+
 	pflag.Parse()
-	
+
 	err := utils.SetupConfigFile()
-	
+
 	if err != nil {
 		panic(fmt.Errorf("fatal error config file: %s", err))
 	}
-	
+
 	if err := viper.Unmarshal(&config); err != nil {
 		panic(fmt.Errorf("unable to decode config: %s", err))
 	}
-	
+
 	csvDaemonURLs = viper.GetString("urls")
 	replicationFactor = viper.GetInt("replication-factor")
 	consistency = viper.GetString("consistency")
 	backoff = viper.GetDuration("backoff")
 	useGzip = viper.GetBool("gzip")
-	
+
 	if _, ok := consistencyChoices[consistency]; !ok {
 		log.Fatalf("invalid consistency settings")
 	}
-	
+
 	daemonURLs = strings.Split(csvDaemonURLs, ",")
 	if len(daemonURLs) == 0 {
 		log.Fatal("missing 'urls' flag")
 	}
 	config.HashWorkers = false
 	loader = load.GetBenchmarkRunner(config)
+
+	scannerBufferSize = viper.GetInt("scanner-buffer-size")
+	if scannerBufferSize == 0 {
+		scannerBufferSize = 1024 * 1024
+	}
 }
 
 type benchmark struct{}
 
 func (b *benchmark) GetDataSource() targets.DataSource {
-	return &fileDataSource{scanner: bufio.NewScanner(load.GetBufferedReader(config.FileName))}
+	scanner := bufio.NewScanner(load.GetBufferedReader(config.FileName))
+	buf := make([]byte, 0, scannerBufferSize)
+	scanner.Buffer(buf, scannerBufferSize*4)
+	return &fileDataSource{scanner: scanner}
 }
 
 func (b *benchmark) GetBatchFactory() targets.BatchFactory {
@@ -116,6 +125,6 @@ func main() {
 			return bytes.NewBuffer(make([]byte, 0, 4*1024*1024))
 		},
 	}
-	
+
 	loader.RunBenchmark(&benchmark{})
 }