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

     1  // Copyright 2019 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  	"sync"
    16  	"sync/atomic"
    17  	"testing"
    18  	"time"
    19  
    20  	"github.com/cockroachdb/cockroach/pkg/col/coldata"
    21  	"github.com/cockroachdb/cockroach/pkg/col/coldatatestutils"
    22  	"github.com/cockroachdb/cockroach/pkg/sql/colexecbase"
    23  	"github.com/cockroachdb/cockroach/pkg/sql/colexecbase/colexecerror"
    24  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    25  	"github.com/cockroachdb/cockroach/pkg/testutils"
    26  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    27  	"github.com/cockroachdb/cockroach/pkg/util/randutil"
    28  	"github.com/stretchr/testify/require"
    29  )
    30  
    31  func TestParallelUnorderedSynchronizer(t *testing.T) {
    32  	defer leaktest.AfterTest(t)()
    33  
    34  	const (
    35  		maxInputs  = 16
    36  		maxBatches = 16
    37  	)
    38  
    39  	var (
    40  		rng, _     = randutil.NewPseudoRand()
    41  		typs       = []*types.T{types.Int}
    42  		numInputs  = rng.Intn(maxInputs) + 1
    43  		numBatches = rng.Intn(maxBatches) + 1
    44  	)
    45  
    46  	inputs := make([]colexecbase.Operator, numInputs)
    47  	for i := range inputs {
    48  		source := colexecbase.NewRepeatableBatchSource(
    49  			testAllocator,
    50  			coldatatestutils.RandomBatch(testAllocator, rng, typs, coldata.BatchSize(), 0 /* length */, rng.Float64()),
    51  			typs,
    52  		)
    53  		source.ResetBatchesToReturn(numBatches)
    54  		inputs[i] = source
    55  	}
    56  
    57  	var wg sync.WaitGroup
    58  	s := NewParallelUnorderedSynchronizer(inputs, typs, &wg)
    59  
    60  	ctx, cancelFn := context.WithCancel(context.Background())
    61  	var cancel bool
    62  	if rng.Float64() < 0.5 {
    63  		cancel = true
    64  	}
    65  	if cancel {
    66  		wg.Add(1)
    67  		sleepTime := time.Duration(rng.Intn(500)) * time.Microsecond
    68  		go func() {
    69  			time.Sleep(sleepTime)
    70  			cancelFn()
    71  			wg.Done()
    72  		}()
    73  	} else {
    74  		// Appease the linter complaining about context leaks.
    75  		defer cancelFn()
    76  	}
    77  
    78  	batchesReturned := 0
    79  	for {
    80  		var b coldata.Batch
    81  		if err := colexecerror.CatchVectorizedRuntimeError(func() { b = s.Next(ctx) }); err != nil {
    82  			if cancel {
    83  				require.True(t, testutils.IsError(err, "context canceled"), err)
    84  				break
    85  			} else {
    86  				t.Fatal(err)
    87  			}
    88  		}
    89  		if b.Length() == 0 {
    90  			break
    91  		}
    92  		batchesReturned++
    93  	}
    94  	if !cancel {
    95  		require.Equal(t, numInputs*numBatches, batchesReturned)
    96  	}
    97  	wg.Wait()
    98  }
    99  
   100  func TestUnorderedSynchronizerNoLeaksOnError(t *testing.T) {
   101  	defer leaktest.AfterTest(t)()
   102  
   103  	const expectedErr = "first input error"
   104  
   105  	inputs := make([]colexecbase.Operator, 6)
   106  	inputs[0] = &colexecbase.CallbackOperator{NextCb: func(context.Context) coldata.Batch {
   107  		colexecerror.InternalError(expectedErr)
   108  		// This code is unreachable, but the compiler cannot infer that.
   109  		return nil
   110  	}}
   111  	for i := 1; i < len(inputs); i++ {
   112  		inputs[i] = &colexecbase.CallbackOperator{
   113  			NextCb: func(ctx context.Context) coldata.Batch {
   114  				<-ctx.Done()
   115  				colexecerror.InternalError(ctx.Err())
   116  				// This code is unreachable, but the compiler cannot infer that.
   117  				return nil
   118  			},
   119  		}
   120  	}
   121  
   122  	var (
   123  		ctx = context.Background()
   124  		wg  sync.WaitGroup
   125  	)
   126  	s := NewParallelUnorderedSynchronizer(inputs, []*types.T{types.Int}, &wg)
   127  	err := colexecerror.CatchVectorizedRuntimeError(func() { _ = s.Next(ctx) })
   128  	// This is the crux of the test: assert that all inputs have finished.
   129  	require.Equal(t, len(inputs), int(atomic.LoadUint32(&s.numFinishedInputs)))
   130  	require.True(t, testutils.IsError(err, expectedErr), err)
   131  }
   132  
   133  func BenchmarkParallelUnorderedSynchronizer(b *testing.B) {
   134  	const numInputs = 6
   135  
   136  	typs := []*types.T{types.Int}
   137  	inputs := make([]colexecbase.Operator, numInputs)
   138  	for i := range inputs {
   139  		batch := testAllocator.NewMemBatchWithSize(typs, coldata.BatchSize())
   140  		batch.SetLength(coldata.BatchSize())
   141  		inputs[i] = colexecbase.NewRepeatableBatchSource(testAllocator, batch, typs)
   142  	}
   143  	var wg sync.WaitGroup
   144  	ctx, cancelFn := context.WithCancel(context.Background())
   145  	s := NewParallelUnorderedSynchronizer(inputs, typs, &wg)
   146  	b.SetBytes(8 * int64(coldata.BatchSize()))
   147  	b.ResetTimer()
   148  	for i := 0; i < b.N; i++ {
   149  		s.Next(ctx)
   150  	}
   151  	b.StopTimer()
   152  	cancelFn()
   153  	wg.Wait()
   154  }