-
Notifications
You must be signed in to change notification settings - Fork 22
/
batches.go
104 lines (88 loc) · 2.55 KB
/
batches.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
// Copyright 2021 The Cockroach Authors.
//
// Use of this software is governed by the Business Source License
// included in the file licenses/BSL.txt.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0, included in the file
// licenses/APL.txt.
// Package batches contains support code for working with and testing
// batches of data.
package batches
import (
"flag"
"sync"
"github.com/cockroachdb/cdc-sink/internal/types"
)
const defaultSize = 1000
var batchSize = flag.Int("batchSize", defaultSize, "default size for batched operations")
// Batch is a helper to perform some operation over a large number
// of values in a batch-oriented fashion. The indexes provided to
// the callback function are a half-open range [begin , end).
func Batch(count int, fn func(begin, end int) error) error {
consume := Size()
idx := 0
for {
if consume > count {
consume = count
}
if err := fn(idx, idx+consume); err != nil {
return err
}
if consume == count {
return nil
}
idx += consume
count -= consume
}
}
// Size returns the default size for batch operations. Testing code
// should generally use a multiple of this value to ensure that
// batching has been correctly implemented.
func Size() int {
x := batchSize
if x == nil {
return defaultSize
}
return *x
}
// The Release function must be called to return the underlying array
// back to the pool.
type Release func()
var intPool = &sync.Pool{New: func() interface{} {
x := make([]int, 0, Size())
return &x
}}
// Int returns a slice of Size() capacity.
func Int() ([]int, Release) {
ret := intPool.Get().(*[]int)
return *ret, func() { intPool.Put(ret) }
}
var int64Pool = &sync.Pool{New: func() interface{} {
x := make([]int64, 0, Size())
return &x
}}
// Int64 returns a slice of Size() capacity.
func Int64() ([]int64, Release) {
ret := int64Pool.Get().(*[]int64)
return *ret, func() { int64Pool.Put(ret) }
}
var mutationPool = &sync.Pool{New: func() interface{} {
x := make([]types.Mutation, 0, Size())
return &x
}}
// Mutation returns a slice of Size() capacity.
func Mutation() ([]types.Mutation, Release) {
ret := mutationPool.Get().(*[]types.Mutation)
return *ret, func() { mutationPool.Put(ret) }
}
var stringPool = &sync.Pool{New: func() interface{} {
x := make([]string, 0, Size())
return &x
}}
// String returns a slice of Size() capacity.
func String() ([]string, Release) {
ret := stringPool.Get().(*[]string)
return *ret, func() { stringPool.Put(ret) }
}