1
16
17 package workqueue
18
19 import (
20 "context"
21 "fmt"
22 "sync/atomic"
23 "testing"
24
25 "github.com/google/go-cmp/cmp"
26 )
27
28 type testCase struct {
29 pieces int
30 workers int
31 chunkSize int
32 }
33
34 func (c testCase) String() string {
35 return fmt.Sprintf("pieces:%d,workers:%d,chunkSize:%d", c.pieces, c.workers, c.chunkSize)
36 }
37
38 var cases = []testCase{
39 {
40 pieces: 1000,
41 workers: 10,
42 chunkSize: 1,
43 },
44 {
45 pieces: 1000,
46 workers: 10,
47 chunkSize: 10,
48 },
49 {
50 pieces: 1000,
51 workers: 10,
52 chunkSize: 100,
53 },
54 {
55 pieces: 999,
56 workers: 10,
57 chunkSize: 13,
58 },
59 }
60
61 func TestParallelizeUntil(t *testing.T) {
62 for _, tc := range cases {
63 t.Run(tc.String(), func(t *testing.T) {
64 seen := make([]int32, tc.pieces)
65 ctx := context.Background()
66 ParallelizeUntil(ctx, tc.workers, tc.pieces, func(p int) {
67 atomic.AddInt32(&seen[p], 1)
68 }, WithChunkSize(tc.chunkSize))
69
70 wantSeen := make([]int32, tc.pieces)
71 for i := 0; i < tc.pieces; i++ {
72 wantSeen[i] = 1
73 }
74 if diff := cmp.Diff(wantSeen, seen); diff != "" {
75 t.Errorf("bad number of visits (-want,+got):\n%s", diff)
76 }
77 })
78 }
79 }
80
81 func BenchmarkParallelizeUntil(b *testing.B) {
82 for _, tc := range cases {
83 b.Run(tc.String(), func(b *testing.B) {
84 ctx := context.Background()
85 isPrime := make([]bool, tc.pieces)
86 b.ResetTimer()
87 for c := 0; c < b.N; c++ {
88 ParallelizeUntil(ctx, tc.workers, tc.pieces, func(p int) {
89 isPrime[p] = calPrime(p)
90 }, WithChunkSize(tc.chunkSize))
91 }
92 b.StopTimer()
93 want := []bool{false, false, true, true, false, true, false, true, false, false, false, true}
94 if diff := cmp.Diff(want, isPrime[:len(want)]); diff != "" {
95 b.Errorf("miscalculated isPrime (-want,+got):\n%s", diff)
96 }
97 })
98 }
99 }
100
101 func calPrime(p int) bool {
102 if p <= 1 {
103 return false
104 }
105 for i := 2; i*i <= p; i++ {
106 if p%i == 0 {
107 return false
108 }
109 }
110 return true
111 }
112
View as plain text