iter_test.go

Documentation: github.com/sourcegraph/conc/iter

     1  package iter
     2  
     3  import (
     4  	"fmt"
     5  	"strconv"
     6  	"sync/atomic"
     7  	"testing"
     8  
     9  	"github.com/stretchr/testify/require"
    10  )
    11  
    12  func ExampleIterator() {
    13  	input := []int{1, 2, 3, 4}
    14  	iterator := Iterator[int]{
    15  		MaxGoroutines: len(input) / 2,
    16  	}
    17  
    18  	iterator.ForEach(input, func(v *int) {
    19  		if *v%2 != 0 {
    20  			*v = -1
    21  		}
    22  	})
    23  
    24  	fmt.Println(input)
    25  	// Output:
    26  	// [-1 2 -1 4]
    27  }
    28  
    29  func TestIterator(t *testing.T) {
    30  	t.Parallel()
    31  
    32  	t.Run("safe for reuse", func(t *testing.T) {
    33  		t.Parallel()
    34  
    35  		iterator := Iterator[int]{MaxGoroutines: 999}
    36  
    37  		// iter.Concurrency > numInput case that updates iter.Concurrency
    38  		iterator.ForEachIdx([]int{1, 2, 3}, func(i int, t *int) {})
    39  
    40  		require.Equal(t, iterator.MaxGoroutines, 999)
    41  	})
    42  
    43  	t.Run("allows more than defaultMaxGoroutines() concurrent tasks", func(t *testing.T) {
    44  		t.Parallel()
    45  
    46  		wantConcurrency := 2 * defaultMaxGoroutines()
    47  
    48  		maxConcurrencyHit := make(chan struct{})
    49  
    50  		tasks := make([]int, wantConcurrency)
    51  		iterator := Iterator[int]{MaxGoroutines: wantConcurrency}
    52  
    53  		var concurrentTasks atomic.Int64
    54  		iterator.ForEach(tasks, func(t *int) {
    55  			n := concurrentTasks.Add(1)
    56  			defer concurrentTasks.Add(-1)
    57  
    58  			if int(n) == wantConcurrency {
    59  				// All our tasks are running concurrently.
    60  				// Signal to the rest of the tasks to stop.
    61  				close(maxConcurrencyHit)
    62  			} else {
    63  				// Wait until we hit max concurrency before exiting.
    64  				// This ensures that all tasks have been started
    65  				// in parallel, despite being a larger input set than
    66  				// defaultMaxGoroutines().
    67  				<-maxConcurrencyHit
    68  			}
    69  		})
    70  	})
    71  }
    72  
    73  func TestForEachIdx(t *testing.T) {
    74  	t.Parallel()
    75  
    76  	t.Run("empty", func(t *testing.T) {
    77  		t.Parallel()
    78  		f := func() {
    79  			ints := []int{}
    80  			ForEachIdx(ints, func(i int, val *int) {
    81  				panic("this should never be called")
    82  			})
    83  		}
    84  		require.NotPanics(t, f)
    85  	})
    86  
    87  	t.Run("panic is propagated", func(t *testing.T) {
    88  		t.Parallel()
    89  		f := func() {
    90  			ints := []int{1}
    91  			ForEachIdx(ints, func(i int, val *int) {
    92  				panic("super bad thing happened")
    93  			})
    94  		}
    95  		require.Panics(t, f)
    96  	})
    97  
    98  	t.Run("mutating inputs is fine", func(t *testing.T) {
    99  		t.Parallel()
   100  		ints := []int{1, 2, 3, 4, 5}
   101  		ForEachIdx(ints, func(i int, val *int) {
   102  			*val += 1
   103  		})
   104  		require.Equal(t, []int{2, 3, 4, 5, 6}, ints)
   105  	})
   106  
   107  	t.Run("huge inputs", func(t *testing.T) {
   108  		t.Parallel()
   109  		ints := make([]int, 10000)
   110  		ForEachIdx(ints, func(i int, val *int) {
   111  			*val = i
   112  		})
   113  		expected := make([]int, 10000)
   114  		for i := 0; i < 10000; i++ {
   115  			expected[i] = i
   116  		}
   117  		require.Equal(t, expected, ints)
   118  	})
   119  }
   120  
   121  func TestForEach(t *testing.T) {
   122  	t.Parallel()
   123  
   124  	t.Run("empty", func(t *testing.T) {
   125  		t.Parallel()
   126  		f := func() {
   127  			ints := []int{}
   128  			ForEach(ints, func(val *int) {
   129  				panic("this should never be called")
   130  			})
   131  		}
   132  		require.NotPanics(t, f)
   133  	})
   134  
   135  	t.Run("panic is propagated", func(t *testing.T) {
   136  		t.Parallel()
   137  		f := func() {
   138  			ints := []int{1}
   139  			ForEach(ints, func(val *int) {
   140  				panic("super bad thing happened")
   141  			})
   142  		}
   143  		require.Panics(t, f)
   144  	})
   145  
   146  	t.Run("mutating inputs is fine", func(t *testing.T) {
   147  		t.Parallel()
   148  		ints := []int{1, 2, 3, 4, 5}
   149  		ForEach(ints, func(val *int) {
   150  			*val += 1
   151  		})
   152  		require.Equal(t, []int{2, 3, 4, 5, 6}, ints)
   153  	})
   154  
   155  	t.Run("huge inputs", func(t *testing.T) {
   156  		t.Parallel()
   157  		ints := make([]int, 10000)
   158  		ForEach(ints, func(val *int) {
   159  			*val = 1
   160  		})
   161  		expected := make([]int, 10000)
   162  		for i := 0; i < 10000; i++ {
   163  			expected[i] = 1
   164  		}
   165  		require.Equal(t, expected, ints)
   166  	})
   167  }
   168  
   169  func BenchmarkForEach(b *testing.B) {
   170  	for _, count := range []int{0, 1, 8, 100, 1000, 10000, 100000} {
   171  		b.Run(strconv.Itoa(count), func(b *testing.B) {
   172  			ints := make([]int, count)
   173  			for i := 0; i < b.N; i++ {
   174  				ForEach(ints, func(i *int) {
   175  					*i = 0
   176  				})
   177  			}
   178  		})
   179  	}
   180  }
   181
View as plain text