rangefunc_test.go

     1  // Copyright 2024 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  //go:build goexperiment.rangefunc
     6  
     7  package race_test
     8  
     9  import (
    10  	"runtime"
    11  	"sync/atomic"
    12  	"testing"
    13  )
    14  
    15  type Seq2[T1, T2 any] func(yield func(T1, T2) bool)
    16  
    17  // ofSliceIndex returns a Seq over the elements of s. It is equivalent
    18  // to range s, except that it splits s into two halves and iterates
    19  // in two separate goroutines.  This is racy if yield is racy, and yield
    20  // will be racy if it contains an early exit.
    21  func ofSliceIndex[T any, S ~[]T](s S) Seq2[int, T] {
    22  	return func(yield func(int, T) bool) {
    23  		c := make(chan bool, 2)
    24  		var done atomic.Bool
    25  		go func() {
    26  			for i := 0; i < len(s)/2; i++ {
    27  				if !done.Load() && !yield(i, s[i]) {
    28  					done.Store(true)
    29  					c <- false
    30  				}
    31  			}
    32  			c <- true
    33  		}()
    34  		go func() {
    35  			for i := len(s) / 2; i < len(s); i++ {
    36  				if !done.Load() && !yield(i, s[i]) {
    37  					done.Store(true)
    38  					c <- false
    39  				}
    40  			}
    41  			c <- true
    42  			return
    43  		}()
    44  		if !<-c {
    45  			return
    46  		}
    47  		<-c
    48  	}
    49  }
    50  
    51  // foo is racy, or not, depending on the value of v
    52  // (0-4 == racy, otherwise, not racy).
    53  func foo(v int) int64 {
    54  	var asum atomic.Int64
    55  	for i, x := range ofSliceIndex([]int64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}) {
    56  		if i%5 == v {
    57  			break
    58  		}
    59  		asum.Add(x) // don't race on asum
    60  		runtime.Gosched()
    61  	}
    62  	return 100 + asum.Load()
    63  }
    64  
    65  // TestRaceRangeFuncIterator races because x%5 can be equal to 4,
    66  // therefore foo can early exit.
    67  func TestRaceRangeFuncIterator(t *testing.T) {
    68  	x := foo(4)
    69  	t.Logf("foo(4)=%d", x)
    70  }
    71  
    72  // TestNoRaceRangeFuncIterator does not race because x%5 is never 5,
    73  // therefore foo's loop will not exit early, and this it will not race.
    74  func TestNoRaceRangeFuncIterator(t *testing.T) {
    75  	x := foo(5)
    76  	t.Logf("foo(5)=%d", x)
    77  }
    78
View as plain text