github.com/whtcorpsinc/milevadb-prod@v0.0.0-20211104133533-f57f4be3b597/dbs/cmd/benchdb/benchkv/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  	"context"
    18  	"flag"
    19  	"fmt"
    20  	"io/ioutil"
    21  	"net/http"
    22  	_ "net/http/pprof"
    23  	"sync"
    24  	"time"
    25  
    26  	_ "github.com/go-allegrosql-driver/allegrosql"
    27  	"github.com/prometheus/client_golang/prometheus"
    28  	"github.com/prometheus/client_golang/prometheus/promhttp"
    29  	"github.com/whtcorpsinc/BerolinaSQL/terror"
    30  	"github.com/whtcorpsinc/errors"
    31  	"github.com/whtcorpsinc/log"
    32  	"github.com/whtcorpsinc/milevadb/causetstore/einsteindb"
    33  	"github.com/whtcorpsinc/milevadb/ekv"
    34  	"go.uber.org/zap"
    35  )
    36  
    37  var (
    38  	causetstore ekv.CausetStorage
    39  	dataCnt     = flag.Int("N", 1000000, "data num")
    40  	workerCnt   = flag.Int("C", 400, "concurrent num")
    41  	FIDelAddr   = flag.String("fidel", "localhost:2379", "fidel address:localhost:2379")
    42  	valueSize   = flag.Int("V", 5, "value size in byte")
    43  
    44  	txnCounter = prometheus.NewCounterVec(
    45  		prometheus.CounterOpts{
    46  			Namespace: "einsteindb",
    47  			Subsystem: "txn",
    48  			Name:      "total",
    49  			Help:      "Counter of txns.",
    50  		}, []string{"type"})
    51  
    52  	txnRolledbackCounter = prometheus.NewCounterVec(
    53  		prometheus.CounterOpts{
    54  			Namespace: "einsteindb",
    55  			Subsystem: "txn",
    56  			Name:      "failed_total",
    57  			Help:      "Counter of rolled back txns.",
    58  		}, []string{"type"})
    59  
    60  	txnDurations = prometheus.NewHistogramVec(
    61  		prometheus.HistogramOpts{
    62  			Namespace: "einsteindb",
    63  			Subsystem: "txn",
    64  			Name:      "durations_histogram_seconds",
    65  			Help:      "Txn latency distributions.",
    66  			Buckets:   prometheus.ExponentialBuckets(0.0005, 2, 13),
    67  		}, []string{"type"})
    68  )
    69  
    70  // Init initializes information.
    71  func Init() {
    72  	driver := einsteindb.Driver{}
    73  	var err error
    74  	causetstore, err = driver.Open(fmt.Sprintf("einsteindb://%s?cluster=1", *FIDelAddr))
    75  	terror.MustNil(err)
    76  
    77  	prometheus.MustRegister(txnCounter)
    78  	prometheus.MustRegister(txnRolledbackCounter)
    79  	prometheus.MustRegister(txnDurations)
    80  	http.Handle("/metrics", promhttp.Handler())
    81  
    82  	go func() {
    83  		err1 := http.ListenAndServe(":9191", nil)
    84  		terror.Log(errors.Trace(err1))
    85  	}()
    86  }
    87  
    88  // batchRW makes sure conflict free.
    89  func batchRW(value []byte) {
    90  	wg := sync.WaitGroup{}
    91  	base := *dataCnt / *workerCnt
    92  	wg.Add(*workerCnt)
    93  	for i := 0; i < *workerCnt; i++ {
    94  		go func(i int) {
    95  			defer wg.Done()
    96  			for j := 0; j < base; j++ {
    97  				txnCounter.WithLabelValues("txn").Inc()
    98  				start := time.Now()
    99  				k := base*i + j
   100  				txn, err := causetstore.Begin()
   101  				if err != nil {
   102  					log.Fatal(err.Error())
   103  				}
   104  				key := fmt.Sprintf("key_%d", k)
   105  				err = txn.Set([]byte(key), value)
   106  				terror.Log(errors.Trace(err))
   107  				err = txn.Commit(context.Background())
   108  				if err != nil {
   109  					txnRolledbackCounter.WithLabelValues("txn").Inc()
   110  					terror.Call(txn.Rollback)
   111  				}
   112  
   113  				txnDurations.WithLabelValues("txn").Observe(time.Since(start).Seconds())
   114  			}
   115  		}(i)
   116  	}
   117  	wg.Wait()
   118  }
   119  
   120  func main() {
   121  	flag.Parse()
   122  	log.SetLevel(zap.ErrorLevel)
   123  	Init()
   124  
   125  	value := make([]byte, *valueSize)
   126  	t := time.Now()
   127  	batchRW(value)
   128  	resp, err := http.Get("http://localhost:9191/metrics")
   129  	terror.MustNil(err)
   130  
   131  	defer terror.Call(resp.Body.Close)
   132  	text, err1 := ioutil.ReadAll(resp.Body)
   133  	terror.Log(errors.Trace(err1))
   134  
   135  	fmt.Println(string(text))
   136  
   137  	fmt.Printf("\nelapse:%v, total %v\n", time.Since(t), *dataCnt)
   138  }