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

     1  // Copyright 2020 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  	"testing"
    17  
    18  	"github.com/cockroachdb/cockroach/pkg/col/coldata"
    19  	"github.com/cockroachdb/cockroach/pkg/col/coldatatestutils"
    20  	"github.com/cockroachdb/cockroach/pkg/sql/colcontainer"
    21  	"github.com/cockroachdb/cockroach/pkg/sql/colexecbase"
    22  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    23  	"github.com/cockroachdb/cockroach/pkg/testutils/colcontainerutils"
    24  	"github.com/cockroachdb/cockroach/pkg/util/humanizeutil"
    25  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    26  	"github.com/cockroachdb/cockroach/pkg/util/randutil"
    27  	"github.com/stretchr/testify/require"
    28  )
    29  
    30  func TestSpillingQueue(t *testing.T) {
    31  	defer leaktest.AfterTest(t)()
    32  
    33  	queueCfg, cleanup := colcontainerutils.NewTestingDiskQueueCfg(t, true /* inMem */)
    34  	defer cleanup()
    35  
    36  	rng, _ := randutil.NewPseudoRand()
    37  	for _, rewindable := range []bool{false, true} {
    38  		for _, memoryLimit := range []int64{10 << 10 /* 10 KiB */, 1<<20 + int64(rng.Intn(64<<20)) /* 1 MiB up to 64 MiB */} {
    39  			alwaysCompress := rng.Float64() < 0.5
    40  			diskQueueCacheMode := colcontainer.DiskQueueCacheModeDefault
    41  			// testReuseCache will test the reuse cache modes.
    42  			testReuseCache := rng.Float64() < 0.5
    43  			dequeuedProbabilityBeforeAllEnqueuesAreDone := 0.5
    44  			if testReuseCache {
    45  				dequeuedProbabilityBeforeAllEnqueuesAreDone = 0
    46  				if rng.Float64() < 0.5 {
    47  					diskQueueCacheMode = colcontainer.DiskQueueCacheModeReuseCache
    48  				} else {
    49  					diskQueueCacheMode = colcontainer.DiskQueueCacheModeClearAndReuseCache
    50  				}
    51  			}
    52  			prefix := ""
    53  			if rewindable {
    54  				dequeuedProbabilityBeforeAllEnqueuesAreDone = 0
    55  				prefix = "Rewindable/"
    56  			}
    57  			numBatches := 1 + rng.Intn(1024)
    58  			t.Run(fmt.Sprintf("%sMemoryLimit=%s/DiskQueueCacheMode=%d/AlwaysCompress=%t/NumBatches=%d",
    59  				prefix, humanizeutil.IBytes(memoryLimit), diskQueueCacheMode, alwaysCompress, numBatches), func(t *testing.T) {
    60  				// Create random input.
    61  				batches := make([]coldata.Batch, 0, numBatches)
    62  				op := coldatatestutils.NewRandomDataOp(testAllocator, rng, coldatatestutils.RandomDataOpArgs{
    63  					NumBatches: cap(batches),
    64  					BatchSize:  1 + rng.Intn(coldata.BatchSize()),
    65  					Nulls:      true,
    66  					BatchAccumulator: func(b coldata.Batch, typs []*types.T) {
    67  						batches = append(batches, coldatatestutils.CopyBatch(b, typs, testColumnFactory))
    68  					},
    69  				})
    70  				typs := op.Typs()
    71  
    72  				queueCfg.CacheMode = diskQueueCacheMode
    73  				queueCfg.SetDefaultBufferSizeBytesForCacheMode()
    74  				queueCfg.TestingKnobs.AlwaysCompress = alwaysCompress
    75  
    76  				// Create queue.
    77  				var q *spillingQueue
    78  				if rewindable {
    79  					q = newRewindableSpillingQueue(
    80  						testAllocator, typs, memoryLimit, queueCfg,
    81  						colexecbase.NewTestingSemaphore(2), coldata.BatchSize(),
    82  						testDiskAcc,
    83  					)
    84  				} else {
    85  					q = newSpillingQueue(
    86  						testAllocator, typs, memoryLimit, queueCfg,
    87  						colexecbase.NewTestingSemaphore(2), coldata.BatchSize(),
    88  						testDiskAcc,
    89  					)
    90  				}
    91  
    92  				// Run verification.
    93  				var (
    94  					b   coldata.Batch
    95  					err error
    96  				)
    97  				ctx := context.Background()
    98  				for {
    99  					b = op.Next(ctx)
   100  					require.NoError(t, q.enqueue(ctx, b))
   101  					if b.Length() == 0 {
   102  						break
   103  					}
   104  					if rng.Float64() < dequeuedProbabilityBeforeAllEnqueuesAreDone {
   105  						if b, err = q.dequeue(ctx); err != nil {
   106  							t.Fatal(err)
   107  						} else if b.Length() == 0 {
   108  							t.Fatal("queue incorrectly considered empty")
   109  						}
   110  						coldata.AssertEquivalentBatches(t, batches[0], b)
   111  						batches = batches[1:]
   112  					}
   113  				}
   114  				numReadIterations := 1
   115  				if rewindable {
   116  					numReadIterations = 2
   117  				}
   118  				for i := 0; i < numReadIterations; i++ {
   119  					batchIdx := 0
   120  					for batches[batchIdx].Length() > 0 {
   121  						if b, err = q.dequeue(ctx); err != nil {
   122  							t.Fatal(err)
   123  						} else if b == nil {
   124  							t.Fatal("unexpectedly dequeued nil batch")
   125  						} else if b.Length() == 0 {
   126  							t.Fatal("queue incorrectly considered empty")
   127  						}
   128  						coldata.AssertEquivalentBatches(t, batches[batchIdx], b)
   129  						batchIdx++
   130  					}
   131  
   132  					if b, err := q.dequeue(ctx); err != nil {
   133  						t.Fatal(err)
   134  					} else if b.Length() != 0 {
   135  						t.Fatal("queue should be empty")
   136  					}
   137  
   138  					if rewindable {
   139  						require.NoError(t, q.rewind())
   140  					}
   141  				}
   142  
   143  				// Close queue.
   144  				require.NoError(t, q.close(ctx))
   145  
   146  				// Verify no directories are left over.
   147  				directories, err := queueCfg.FS.List(queueCfg.Path)
   148  				require.NoError(t, err)
   149  				require.Equal(t, 0, len(directories))
   150  			})
   151  		}
   152  	}
   153  }