github.com/cockroachdb/pebble@v1.1.1-0.20240513155919-3622ade60459/cmd/pebble/scan.go (about)

     1  // Copyright 2018 The LevelDB-Go and Pebble Authors. All rights reserved. Use
     2  // of this source code is governed by a BSD-style license that can be found in
     3  // the LICENSE file.
     4  
     5  package main
     6  
     7  import (
     8  	"fmt"
     9  	"log"
    10  	"math"
    11  	"sync"
    12  	"sync/atomic"
    13  	"time"
    14  
    15  	"github.com/cockroachdb/pebble"
    16  	"github.com/cockroachdb/pebble/internal/randvar"
    17  	"github.com/spf13/cobra"
    18  	"golang.org/x/exp/rand"
    19  )
    20  
    21  var scanConfig struct {
    22  	reverse bool
    23  	rows    *randvar.Flag
    24  	values  *randvar.BytesFlag
    25  }
    26  
    27  var scanCmd = &cobra.Command{
    28  	Use:   "scan <dir>",
    29  	Short: "run the scan benchmark",
    30  	Long:  ``,
    31  	Args:  cobra.ExactArgs(1),
    32  	Run:   runScan,
    33  }
    34  
    35  func init() {
    36  	scanCmd.Flags().BoolVarP(
    37  		&scanConfig.reverse, "reverse", "r", false, "reverse scan")
    38  	scanConfig.rows = randvar.NewFlag("100")
    39  	scanCmd.Flags().Var(
    40  		scanConfig.rows, "rows", "number of rows to scan in each operation")
    41  	scanConfig.values = randvar.NewBytesFlag("8")
    42  	scanCmd.Flags().Var(
    43  		scanConfig.values, "values",
    44  		"value size distribution [{zipf,uniform}:]min[-max][/<target-compression>]")
    45  }
    46  
    47  func runScan(cmd *cobra.Command, args []string) {
    48  	var (
    49  		bytes       atomic.Int64
    50  		scanned     atomic.Int64
    51  		lastBytes   int64
    52  		lastScanned int64
    53  		lastElapsed time.Duration
    54  	)
    55  
    56  	opts := pebble.Sync
    57  	if disableWAL {
    58  		opts = pebble.NoSync
    59  	}
    60  
    61  	rowDist := scanConfig.rows
    62  
    63  	runTest(args[0], test{
    64  		init: func(d DB, wg *sync.WaitGroup) {
    65  			const count = 100000
    66  			const batch = 1000
    67  
    68  			rng := rand.New(rand.NewSource(1449168817))
    69  			keys := make([][]byte, count)
    70  
    71  			for i := 0; i < count; {
    72  				b := d.NewBatch()
    73  				var value []byte
    74  				for end := i + batch; i < end; i++ {
    75  					keys[i] = mvccEncode(nil, encodeUint32Ascending([]byte("key-"), uint32(i)), uint64(i+1), 0)
    76  					value = scanConfig.values.Bytes(rng, value)
    77  					if err := b.Set(keys[i], value, nil); err != nil {
    78  						log.Fatal(err)
    79  					}
    80  				}
    81  				if err := b.Commit(opts); err != nil {
    82  					log.Fatal(err)
    83  				}
    84  			}
    85  
    86  			if err := d.Flush(); err != nil {
    87  				log.Fatal(err)
    88  			}
    89  
    90  			limiter := maxOpsPerSec.newRateLimiter()
    91  
    92  			wg.Add(concurrency)
    93  			for i := 0; i < concurrency; i++ {
    94  				go func(i int) {
    95  					defer wg.Done()
    96  
    97  					rng := rand.New(rand.NewSource(uint64(i)))
    98  					startKeyBuf := append(make([]byte, 0, 64), []byte("key-")...)
    99  					endKeyBuf := append(make([]byte, 0, 64), []byte("key-")...)
   100  					minTS := encodeUint64Ascending(nil, math.MaxUint64)
   101  
   102  					for {
   103  						wait(limiter)
   104  
   105  						rows := int(rowDist.Uint64(rng))
   106  						startIdx := rng.Int31n(int32(len(keys) - rows))
   107  						startKey := encodeUint32Ascending(startKeyBuf[:4], uint32(startIdx))
   108  						endKey := encodeUint32Ascending(endKeyBuf[:4], uint32(startIdx+int32(rows)))
   109  
   110  						var count int
   111  						var nbytes int64
   112  						if scanConfig.reverse {
   113  							count, nbytes = mvccReverseScan(d, startKey, endKey, minTS)
   114  						} else {
   115  							count, nbytes = mvccForwardScan(d, startKey, endKey, minTS)
   116  						}
   117  
   118  						if count != rows {
   119  							log.Fatalf("scanned %d, expected %d\n", count, rows)
   120  						}
   121  
   122  						bytes.Add(nbytes)
   123  						scanned.Add(int64(count))
   124  					}
   125  				}(i)
   126  			}
   127  		},
   128  
   129  		tick: func(elapsed time.Duration, i int) {
   130  			if i%20 == 0 {
   131  				fmt.Println("_elapsed_______rows/sec_______MB/sec_______ns/row")
   132  			}
   133  
   134  			curBytes := bytes.Load()
   135  			curScanned := scanned.Load()
   136  			dur := elapsed - lastElapsed
   137  			fmt.Printf("%8s %14.1f %12.1f %12.1f\n",
   138  				time.Duration(elapsed.Seconds()+0.5)*time.Second,
   139  				float64(curScanned-lastScanned)/dur.Seconds(),
   140  				float64(curBytes-lastBytes)/(dur.Seconds()*(1<<20)),
   141  				float64(dur)/float64(curScanned-lastScanned),
   142  			)
   143  			lastBytes = curBytes
   144  			lastScanned = curScanned
   145  			lastElapsed = elapsed
   146  		},
   147  
   148  		done: func(elapsed time.Duration) {
   149  			curBytes := bytes.Load()
   150  			curScanned := scanned.Load()
   151  			fmt.Println("\n_elapsed___ops/sec(cum)__MB/sec(cum)__ns/row(avg)")
   152  			fmt.Printf("%7.1fs %14.1f %12.1f %12.1f\n\n",
   153  				elapsed.Seconds(),
   154  				float64(curScanned)/elapsed.Seconds(),
   155  				float64(curBytes)/(elapsed.Seconds()*(1<<20)),
   156  				float64(elapsed)/float64(curScanned),
   157  			)
   158  		},
   159  	})
   160  }