github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/workload/ycsb/zipfgenerator_test.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 ycsb
    12  
    13  import (
    14  	"fmt"
    15  	"math"
    16  	"math/rand"
    17  	"sort"
    18  	"testing"
    19  
    20  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    21  	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
    22  )
    23  
    24  type params struct {
    25  	iMin, iMax uint64
    26  	theta      float64
    27  }
    28  
    29  var gens = []params{
    30  	{0, 100, 0.99},
    31  	{0, 100, 1.01},
    32  }
    33  
    34  func TestCreateZipfGenerator(t *testing.T) {
    35  	defer leaktest.AfterTest(t)()
    36  	for _, gen := range gens {
    37  		rng := rand.New(rand.NewSource(timeutil.Now().UnixNano()))
    38  		_, err := NewZipfGenerator(rng, gen.iMin, gen.iMax, gen.theta, false)
    39  		if err != nil {
    40  			t.Fatal(err)
    41  		}
    42  	}
    43  }
    44  
    45  var tests = []struct {
    46  	n        uint64
    47  	theta    float64
    48  	expected float64
    49  }{
    50  	{20, 0.99, 3.64309060779367},
    51  	{200, 0.99, 6.02031118558},
    52  	{1000, 0.99, 7.72895321728},
    53  	{2000, 0.99, 8.47398788329},
    54  	{10000, 0.99, 10.2243614596},
    55  	{100000, 0.99, 12.7783380626},
    56  	{1000000, 0.99, 15.391849746},
    57  	{10000000, 0.99, 18.066242575},
    58  	{100000000, 0.99, 20.80293049},
    59  }
    60  
    61  func TestZetaFromScratch(t *testing.T) {
    62  	defer leaktest.AfterTest(t)()
    63  	if testing.Short() {
    64  		t.Skip("short")
    65  	}
    66  	for _, test := range tests {
    67  		computedZeta, err := computeZetaFromScratch(test.n, test.theta)
    68  		if err != nil {
    69  			t.Fatalf("Failed to compute zeta(%d,%f): %s", test.n, test.theta, err)
    70  		}
    71  		if math.Abs(computedZeta-test.expected) > 0.000000001 {
    72  			t.Fatalf("expected %6.4f, got %6.4f", test.expected, computedZeta)
    73  		}
    74  	}
    75  }
    76  
    77  func TestZetaIncrementally(t *testing.T) {
    78  	defer leaktest.AfterTest(t)()
    79  	if testing.Short() {
    80  		t.Skip("short")
    81  	}
    82  	// Theta cannot be 1 by definition, so this is a safe initial value.
    83  	oldTheta := 1.0
    84  	var oldZetaN float64
    85  	var oldN uint64
    86  	for _, test := range tests {
    87  		// If theta has changed, recompute from scratch
    88  		if test.theta != oldTheta {
    89  			var err error
    90  			oldZetaN, err = computeZetaFromScratch(test.n, test.theta)
    91  			if err != nil {
    92  				t.Fatalf("Failed to compute zeta(%d,%f): %s", test.n, test.theta, err)
    93  			}
    94  			oldN = test.n
    95  			continue
    96  		}
    97  
    98  		computedZeta, err := computeZetaIncrementally(oldN, test.n, test.theta, oldZetaN)
    99  		if err != nil {
   100  			t.Fatalf("Failed to compute zeta(%d,%f) incrementally: %s", test.n, test.theta, err)
   101  		}
   102  		if math.Abs(computedZeta-test.expected) > 0.000000001 {
   103  			t.Fatalf("expected %6.4f, got %6.4f", test.expected, computedZeta)
   104  		}
   105  
   106  		oldZetaN = computedZeta
   107  		oldN = test.n
   108  	}
   109  }
   110  
   111  func runZipfGenerators(t *testing.T, withIncrements bool) {
   112  	gen := gens[0]
   113  	rng := rand.New(rand.NewSource(timeutil.Now().UnixNano()))
   114  	z, err := NewZipfGenerator(rng, gen.iMin, gen.iMax, gen.theta, false)
   115  	if err != nil {
   116  		t.Fatal(err)
   117  	}
   118  
   119  	const ROLLS = 10000
   120  	x := make([]int, ROLLS)
   121  
   122  	for i := 0; i < ROLLS; i++ {
   123  		x[i] = int(z.Uint64())
   124  		z.zipfGenMu.mu.Lock()
   125  		if x[i] < int(z.iMin) || x[i] > int(z.zipfGenMu.iMax) {
   126  			t.Fatalf("zipf(%d,%d,%f) rolled %d at index %d", z.iMin, z.zipfGenMu.iMax, z.theta, x[i], i)
   127  			z.zipfGenMu.mu.Unlock()
   128  			if withIncrements {
   129  				if err := z.IncrementIMax(1); err != nil {
   130  					t.Fatalf("could not increment iMax: %s", err)
   131  				}
   132  			}
   133  		}
   134  		z.zipfGenMu.mu.Unlock()
   135  	}
   136  
   137  	if withIncrements {
   138  		return
   139  	}
   140  
   141  	sort.Ints(x)
   142  
   143  	max := x[ROLLS-1]
   144  	step := max / 20
   145  	index := 0
   146  	count := 0
   147  	for i := 0; i < max; i += step {
   148  		count = 0
   149  		for {
   150  			if x[index] >= i+step {
   151  				break
   152  			}
   153  			index++
   154  			count++
   155  		}
   156  		fmt.Printf("[%10d-%10d)        ", i, i+step)
   157  		for j := 0; j < count; j++ {
   158  			if j%50 == 0 {
   159  				fmt.Printf("%c", '∎')
   160  			}
   161  		}
   162  		fmt.Println()
   163  	}
   164  }
   165  
   166  func TestZipfGenerator(t *testing.T) {
   167  	defer leaktest.AfterTest(t)()
   168  	runZipfGenerators(t, false)
   169  	runZipfGenerators(t, true)
   170  }