github.com/whtcorpsinc/MilevaDB-Prod@v0.0.0-20211104133533-f57f4be3b597/dbs/cmd/benchdb/benchraw/main.go (about)

     1  // Copyright 2020 WHTCORPS INC, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package main
    15  
    16  import (
    17  	"flag"
    18  	"fmt"
    19  	"net/http"
    20  	_ "net/http/pprof"
    21  	"strings"
    22  	"sync"
    23  	"time"
    24  
    25  	"github.com/whtcorpsinc/errors"
    26  	"github.com/whtcorpsinc/log"
    27  	"github.com/whtcorpsinc/BerolinaSQL/terror"
    28  	"github.com/whtcorpsinc/milevadb/config"
    29  	"github.com/whtcorpsinc/milevadb/causetstore/einsteindb"
    30  	"go.uber.org/zap"
    31  )
    32  
    33  var (
    34  	dataCnt   = flag.Int("N", 1000000, "data num")
    35  	workerCnt = flag.Int("C", 100, "concurrent num")
    36  	FIDelAddr    = flag.String("fidel", "localhost:2379", "fidel address:localhost:2379")
    37  	valueSize = flag.Int("V", 5, "value size in byte")
    38  	sslCA     = flag.String("cacert", "", "path of file that contains list of trusted SSL CAs.")
    39  	sslCert   = flag.String("cert", "", "path of file that contains X509 certificate in PEM format.")
    40  	sslKey    = flag.String("key", "", "path of file that contains X509 key in PEM format.")
    41  )
    42  
    43  // batchRawPut blinds put bench.
    44  func batchRawPut(value []byte) {
    45  	cli, err := einsteindb.NewRawKVClient(strings.Split(*FIDelAddr, ","), config.Security{
    46  		ClusterSSLCA:   *sslCA,
    47  		ClusterSSLCert: *sslCert,
    48  		ClusterSSLKey:  *sslKey,
    49  	})
    50  	if err != nil {
    51  		log.Fatal(err.Error())
    52  	}
    53  
    54  	wg := sync.WaitGroup{}
    55  	base := *dataCnt / *workerCnt
    56  	wg.Add(*workerCnt)
    57  	for i := 0; i < *workerCnt; i++ {
    58  		go func(i int) {
    59  			defer wg.Done()
    60  
    61  			for j := 0; j < base; j++ {
    62  				k := base*i + j
    63  				key := fmt.Sprintf("key_%d", k)
    64  				err = cli.Put([]byte(key), value)
    65  				if err != nil {
    66  					log.Fatal("put failed", zap.Error(err))
    67  				}
    68  			}
    69  		}(i)
    70  	}
    71  	wg.Wait()
    72  }
    73  
    74  func main() {
    75  	flag.Parse()
    76  	log.SetLevel(zap.WarnLevel)
    77  	go func() {
    78  		err := http.ListenAndServe(":9191", nil)
    79  		terror.Log(errors.Trace(err))
    80  	}()
    81  
    82  	value := make([]byte, *valueSize)
    83  	t := time.Now()
    84  	batchRawPut(value)
    85  
    86  	fmt.Printf("\nelapse:%v, total %v\n", time.Since(t), *dataCnt)
    87  }