github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/workload/ledger/worker.go (about)

     1  // Copyright 2017 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package ledger
    12  
    13  import (
    14  	"context"
    15  	gosql "database/sql"
    16  	"math/rand"
    17  	"strconv"
    18  	"strings"
    19  
    20  	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
    21  	"github.com/cockroachdb/cockroach/pkg/workload/histogram"
    22  	"github.com/cockroachdb/errors"
    23  )
    24  
    25  type worker struct {
    26  	config *ledger
    27  	hists  *histogram.Histograms
    28  	db     *gosql.DB
    29  
    30  	rng      *rand.Rand
    31  	deckPerm []int
    32  	permIdx  int
    33  }
    34  
    35  type ledgerTx interface {
    36  	run(config *ledger, db *gosql.DB, rng *rand.Rand) (interface{}, error)
    37  }
    38  
    39  type tx struct {
    40  	ledgerTx
    41  	weight int    // percent likelihood that each transaction type is run
    42  	name   string // display name
    43  }
    44  
    45  var allTxs = [...]tx{
    46  	{
    47  		ledgerTx: balance{}, name: "balance",
    48  	},
    49  	{
    50  		ledgerTx: withdrawal{}, name: "withdrawal",
    51  	},
    52  	{
    53  		ledgerTx: deposit{}, name: "deposit",
    54  	},
    55  	{
    56  		ledgerTx: reversal{}, name: "reversal",
    57  	},
    58  }
    59  
    60  func initializeMix(config *ledger) error {
    61  	config.txs = append([]tx(nil), allTxs[0:]...)
    62  	nameToTx := make(map[string]int, len(allTxs))
    63  	for i, tx := range config.txs {
    64  		nameToTx[tx.name] = i
    65  	}
    66  
    67  	items := strings.Split(config.mix, `,`)
    68  	totalWeight := 0
    69  	for _, item := range items {
    70  		kv := strings.Split(item, `=`)
    71  		if len(kv) != 2 {
    72  			return errors.Errorf(`Invalid mix %s: %s is not a k=v pair`, config.mix, item)
    73  		}
    74  		txName, weightStr := kv[0], kv[1]
    75  
    76  		weight, err := strconv.Atoi(weightStr)
    77  		if err != nil {
    78  			return errors.Errorf(
    79  				`Invalid percentage mix %s: %s is not an integer`, config.mix, weightStr)
    80  		}
    81  
    82  		i, ok := nameToTx[txName]
    83  		if !ok {
    84  			return errors.Errorf(
    85  				`Invalid percentage mix %s: no such transaction %s`, config.mix, txName)
    86  		}
    87  
    88  		config.txs[i].weight = weight
    89  		totalWeight += weight
    90  	}
    91  
    92  	config.deck = make([]int, 0, totalWeight)
    93  	for i, t := range config.txs {
    94  		for j := 0; j < t.weight; j++ {
    95  			config.deck = append(config.deck, i)
    96  		}
    97  	}
    98  
    99  	return nil
   100  }
   101  
   102  func (w *worker) run(ctx context.Context) error {
   103  	if w.permIdx == len(w.deckPerm) {
   104  		rand.Shuffle(len(w.deckPerm), func(i, j int) {
   105  			w.deckPerm[i], w.deckPerm[j] = w.deckPerm[j], w.deckPerm[i]
   106  		})
   107  		w.permIdx = 0
   108  	}
   109  	// Move through our permutation slice until its exhausted, using each value to
   110  	// to index into our deck of transactions, which contains indexes into the
   111  	// txs slice.
   112  	opIdx := w.deckPerm[w.permIdx]
   113  	t := w.config.txs[opIdx]
   114  	w.permIdx++
   115  
   116  	start := timeutil.Now()
   117  	if _, err := t.run(w.config, w.db, w.rng); err != nil {
   118  		return errors.Wrapf(err, "error in %s", t.name)
   119  	}
   120  	elapsed := timeutil.Since(start)
   121  	w.hists.Get(t.name).Record(elapsed)
   122  	return nil
   123  }