github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/colexec/main_test.go (about)

     1  // Copyright 2018 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 colexec
    12  
    13  import (
    14  	"context"
    15  	"fmt"
    16  	"os"
    17  	"testing"
    18  
    19  	"github.com/cockroachdb/cockroach/pkg/col/coldata"
    20  	"github.com/cockroachdb/cockroach/pkg/col/coldataext"
    21  	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
    22  	"github.com/cockroachdb/cockroach/pkg/sql/colexecbase/colexecerror"
    23  	"github.com/cockroachdb/cockroach/pkg/sql/colmem"
    24  	"github.com/cockroachdb/cockroach/pkg/sql/execinfra"
    25  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    26  	"github.com/cockroachdb/cockroach/pkg/util/envutil"
    27  	"github.com/cockroachdb/cockroach/pkg/util/mon"
    28  	"github.com/cockroachdb/cockroach/pkg/util/randutil"
    29  )
    30  
    31  //go:generate ../../util/leaktest/add-leaktest.sh *_test.go
    32  
    33  var (
    34  	// testAllocator is an Allocator with an unlimited budget for use in tests.
    35  	testAllocator     *colmem.Allocator
    36  	testColumnFactory coldata.ColumnFactory
    37  
    38  	// testMemMonitor and testMemAcc are a test monitor with an unlimited budget
    39  	// and a memory account bound to it for use in tests.
    40  	testMemMonitor *mon.BytesMonitor
    41  	testMemAcc     *mon.BoundAccount
    42  
    43  	// testDiskMonitor and testDiskAcc are a test monitor with an unlimited budget
    44  	// and a disk account bound to it for use in tests.
    45  	testDiskMonitor *mon.BytesMonitor
    46  	testDiskAcc     *mon.BoundAccount
    47  )
    48  
    49  func TestMain(m *testing.M) {
    50  	randutil.SeedForTests()
    51  	os.Exit(func() int {
    52  		ctx := context.Background()
    53  		st := cluster.MakeTestingClusterSettings()
    54  		testMemMonitor = execinfra.NewTestMemMonitor(ctx, st)
    55  		defer testMemMonitor.Stop(ctx)
    56  		memAcc := testMemMonitor.MakeBoundAccount()
    57  		testMemAcc = &memAcc
    58  		evalCtx := tree.MakeTestingEvalContext(st)
    59  		testColumnFactory = coldataext.NewExtendedColumnFactory(&evalCtx)
    60  		testAllocator = colmem.NewAllocator(ctx, testMemAcc, testColumnFactory)
    61  		defer testMemAcc.Close(ctx)
    62  
    63  		testDiskMonitor = execinfra.NewTestDiskMonitor(ctx, st)
    64  		defer testDiskMonitor.Stop(ctx)
    65  		diskAcc := testDiskMonitor.MakeBoundAccount()
    66  		testDiskAcc = &diskAcc
    67  		defer testDiskAcc.Close(ctx)
    68  
    69  		// Pick a random batch size in [minBatchSize, coldata.MaxBatchSize]
    70  		// range. The randomization can be disabled using COCKROACH_RANDOMIZE_BATCH_SIZE=false.
    71  		randomBatchSize := generateBatchSize()
    72  		fmt.Printf("coldata.BatchSize() is set to %d\n", randomBatchSize)
    73  		if err := coldata.SetBatchSizeForTests(randomBatchSize); err != nil {
    74  			colexecerror.InternalError(err)
    75  		}
    76  		return m.Run()
    77  	}())
    78  }
    79  
    80  // minBatchSize is the minimum acceptable size of batches for tests in this
    81  // package.
    82  const minBatchSize = 3
    83  
    84  func generateBatchSize() int {
    85  	randomizeBatchSize := envutil.EnvOrDefaultBool("COCKROACH_RANDOMIZE_BATCH_SIZE", true)
    86  	if randomizeBatchSize {
    87  		rng, _ := randutil.NewPseudoRand()
    88  		return minBatchSize + rng.Intn(coldata.MaxBatchSize-minBatchSize)
    89  	}
    90  	return coldata.BatchSize()
    91  }