master
go 83 lines 1.47 KB
Raw
1 package framework
2
3 // Batch creates batches from a slice of items
4 func Batch[T any](items []T, size int) <-chan []T {
5 ch := make(chan []T)
6
7 go func() {
8 defer close(ch)
9
10 for i := 0; i < len(items); i += size {
11 end := min(i+size, len(items))
12
13 ch <- items[i:end]
14 }
15 }()
16
17 return ch
18 }
19
20 // BatchWithError creates batches and allows error handling
21 func BatchWithError[T any](items []T, size int, fn func([]T) error) error {
22 for batch := range Batch(items, size) {
23 if err := fn(batch); err != nil {
24 return err
25 }
26 }
27 return nil
28 }
29
30 // ParallelBatch processes batches in parallel with worker pool
31 func ParallelBatch[T any](items []T, batchSize int, workers int, fn func([]T) error) error {
32 type result struct {
33 err error
34 }
35
36 work := make(chan []T, workers)
37 results := make(chan result, workers)
38
39 // Start workers
40 for range workers {
41 go func() {
42 for batch := range work {
43 results <- result{err: fn(batch)}
44 }
45 }()
46 }
47
48 // Send work
49 go func() {
50 for batch := range Batch(items, batchSize) {
51 work <- batch
52 }
53 close(work)
54 }()
55
56 // Collect results
57 var firstErr error
58 batchCount := (len(items) + batchSize - 1) / batchSize
59 for range batchCount {
60 res := <-results
61 if res.err != nil && firstErr == nil {
62 firstErr = res.err
63 }
64 }
65
66 return firstErr
67 }
68
69 // Min returns the minimum of two integers
70 func Min(a, b int) int {
71 if a < b {
72 return a
73 }
74 return b
75 }
76
77 // Max returns the maximum of two integers
78 func Max(a, b int) int {
79 if a > b {
80 return a
81 }
82 return b
83 }