github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/util/quotapool/notify_queue_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 quotapool
    12  
    13  import (
    14  	"math/rand"
    15  	"testing"
    16  
    17  	"github.com/stretchr/testify/assert"
    18  )
    19  
    20  func BenchmarkNotifyQueue(b *testing.B) {
    21  	testNotifyQueue(b, b.N)
    22  }
    23  
    24  func TestNotifyQueue(t *testing.T) {
    25  	testNotifyQueue(t, 10000)
    26  }
    27  
    28  type op bool
    29  
    30  const (
    31  	enqueue op = true
    32  	dequeue op = false
    33  )
    34  
    35  func testNotifyQueue(t testing.TB, N int) {
    36  	b, _ := t.(*testing.B)
    37  	var q notifyQueue
    38  	initializeNotifyQueue(&q)
    39  	n := q.peek()
    40  	assert.Nil(t, n)
    41  	q.dequeue()
    42  	assert.Equal(t, 0, int(q.len))
    43  	chans := make([]chan struct{}, N)
    44  	for i := 0; i < N; i++ {
    45  		chans[i] = make(chan struct{})
    46  	}
    47  	in := chans
    48  	out := make([]chan struct{}, 0, N)
    49  	ops := make([]op, (4*N)/3)
    50  	for i := 0; i < N; i++ {
    51  		ops[i] = enqueue
    52  	}
    53  	rand.Shuffle(len(ops), func(i, j int) {
    54  		ops[i], ops[j] = ops[j], ops[i]
    55  	})
    56  	if b != nil {
    57  		b.ResetTimer()
    58  	}
    59  	l := 0 // only used if b == nil
    60  	for _, op := range ops {
    61  		switch op {
    62  		case enqueue:
    63  			q.enqueue(in[0])
    64  			in = in[1:]
    65  			if b == nil {
    66  				l++
    67  			}
    68  		case dequeue:
    69  			// Only test Peek if we're not benchmarking.
    70  			if b == nil {
    71  				if n := q.peek(); n != nil {
    72  					out = append(out, n.c)
    73  					q.dequeue()
    74  					l--
    75  					assert.Equal(t, l, int(q.len))
    76  				}
    77  			} else {
    78  				if n := q.peek(); n != nil {
    79  					out = append(out, n.c)
    80  					q.dequeue()
    81  				}
    82  			}
    83  		}
    84  	}
    85  	for n := q.peek(); n != nil; n = q.peek() {
    86  		out = append(out, n.c)
    87  		q.dequeue()
    88  	}
    89  	if b != nil {
    90  		b.StopTimer()
    91  	}
    92  	assert.EqualValues(t, chans, out)
    93  }