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 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407
|
// Copyright 2016 Google LLC.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Package bundler supports bundling (batching) of items. Bundling amortizes an
// action with fixed costs over multiple items. For example, if an API provides
// an RPC that accepts a list of items as input, but clients would prefer
// adding items one at a time, then a Bundler can accept individual items from
// the client and bundle many of them into a single RPC.
//
// This package is experimental and subject to change without notice.
package bundler
import (
"context"
"errors"
"reflect"
"sync"
"time"
"golang.org/x/sync/semaphore"
)
type mode int
const (
DefaultDelayThreshold = time.Second
DefaultBundleCountThreshold = 10
DefaultBundleByteThreshold = 1e6 // 1M
DefaultBufferedByteLimit = 1e9 // 1G
)
const (
none mode = iota
add
addWait
)
var (
// ErrOverflow indicates that Bundler's stored bytes exceeds its BufferedByteLimit.
ErrOverflow = errors.New("bundler reached buffered byte limit")
// ErrOversizedItem indicates that an item's size exceeds the maximum bundle size.
ErrOversizedItem = errors.New("item size exceeds bundle byte limit")
// errMixedMethods indicates that mutually exclusive methods has been
// called subsequently.
errMixedMethods = errors.New("calls to Add and AddWait cannot be mixed")
)
// A Bundler collects items added to it into a bundle until the bundle
// exceeds a given size, then calls a user-provided function to handle the
// bundle.
//
// The exported fields are only safe to modify prior to the first call to Add
// or AddWait.
type Bundler struct {
// Starting from the time that the first message is added to a bundle, once
// this delay has passed, handle the bundle. The default is DefaultDelayThreshold.
DelayThreshold time.Duration
// Once a bundle has this many items, handle the bundle. Since only one
// item at a time is added to a bundle, no bundle will exceed this
// threshold, so it also serves as a limit. The default is
// DefaultBundleCountThreshold.
BundleCountThreshold int
// Once the number of bytes in current bundle reaches this threshold, handle
// the bundle. The default is DefaultBundleByteThreshold. This triggers handling,
// but does not cap the total size of a bundle.
BundleByteThreshold int
// The maximum size of a bundle, in bytes. Zero means unlimited.
BundleByteLimit int
// The maximum number of bytes that the Bundler will keep in memory before
// returning ErrOverflow. The default is DefaultBufferedByteLimit.
BufferedByteLimit int
// The maximum number of handler invocations that can be running at once.
// The default is 1.
HandlerLimit int
handler func(interface{}) // called to handle a bundle
itemSliceZero reflect.Value // nil (zero value) for slice of items
mu sync.Mutex // guards access to fields below
flushTimer *time.Timer // implements DelayThreshold
handlerCount int // # of bundles currently being handled (i.e. handler is invoked on them)
sem *semaphore.Weighted // enforces BufferedByteLimit
semOnce sync.Once // guards semaphore initialization
// The current bundle we're adding items to. Not yet in the queue.
// Appended to the queue once the flushTimer fires or the bundle
// thresholds/limits are reached. If curBundle is nil and tail is
// not, we first try to add items to tail. Once tail is full or handled,
// we create a new curBundle for the incoming item.
curBundle *bundle
// The next bundle in the queue to be handled. Nil if the queue is
// empty.
head *bundle
// The last bundle in the queue to be handled. Nil if the queue is
// empty. If curBundle is nil and tail isn't, we attempt to add new
// items to the tail until if becomes full or has been passed to the
// handler.
tail *bundle
curFlush *sync.WaitGroup // counts outstanding bundles since last flush
prevFlush chan bool // signal used to wait for prior flush
// The first call to Add or AddWait, mode will be add or addWait respectively.
// If there wasn't call yet then mode is none.
mode mode
// TODO: consider alternative queue implementation for head/tail bundle. see:
// https://code-review.googlesource.com/c/google-api-go-client/+/47991/4/support/bundler/bundler.go#74
}
// A bundle is a group of items that were added individually and will be passed
// to a handler as a slice.
type bundle struct {
items reflect.Value // slice of T
size int // size in bytes of all items
next *bundle // bundles are handled in order as a linked list queue
flush *sync.WaitGroup // the counter that tracks flush completion
}
// add appends item to this bundle and increments the total size. It requires
// that b.mu is locked.
func (bu *bundle) add(item interface{}, size int) {
bu.items = reflect.Append(bu.items, reflect.ValueOf(item))
bu.size += size
}
// NewBundler creates a new Bundler.
//
// itemExample is a value of the type that will be bundled. For example, if you
// want to create bundles of *Entry, you could pass &Entry{} for itemExample.
//
// handler is a function that will be called on each bundle. If itemExample is
// of type T, the argument to handler is of type []T. handler is always called
// sequentially for each bundle, and never in parallel.
//
// Configure the Bundler by setting its thresholds and limits before calling
// any of its methods.
func NewBundler(itemExample interface{}, handler func(interface{})) *Bundler {
b := &Bundler{
DelayThreshold: DefaultDelayThreshold,
BundleCountThreshold: DefaultBundleCountThreshold,
BundleByteThreshold: DefaultBundleByteThreshold,
BufferedByteLimit: DefaultBufferedByteLimit,
HandlerLimit: 1,
handler: handler,
itemSliceZero: reflect.Zero(reflect.SliceOf(reflect.TypeOf(itemExample))),
curFlush: &sync.WaitGroup{},
}
return b
}
func (b *Bundler) initSemaphores() {
// Create the semaphores lazily, because the user may set limits
// after NewBundler.
b.semOnce.Do(func() {
b.sem = semaphore.NewWeighted(int64(b.BufferedByteLimit))
})
}
// enqueueCurBundle moves curBundle to the end of the queue. The bundle may be
// handled immediately if we are below HandlerLimit. It requires that b.mu is
// locked.
func (b *Bundler) enqueueCurBundle() {
// We don't require callers to check if there is a pending bundle. It
// may have already been appended to the queue. If so, return early.
if b.curBundle == nil {
return
}
// If we are below the HandlerLimit, the queue must be empty. Handle
// immediately with a new goroutine.
if b.handlerCount < b.HandlerLimit {
b.handlerCount++
go b.handle(b.curBundle)
} else if b.tail != nil {
// There are bundles on the queue, so append to the end
b.tail.next = b.curBundle
b.tail = b.curBundle
} else {
// The queue is empty, so initialize the queue
b.head = b.curBundle
b.tail = b.curBundle
}
b.curBundle = nil
if b.flushTimer != nil {
b.flushTimer.Stop()
b.flushTimer = nil
}
}
// setMode sets the state of Bundler's mode. If mode was defined before
// and passed state is different from it then return an error.
func (b *Bundler) setMode(m mode) error {
b.mu.Lock()
defer b.mu.Unlock()
if b.mode == m || b.mode == none {
b.mode = m
return nil
}
return errMixedMethods
}
// canFit returns true if bu can fit an additional item of size bytes based
// on the limits of Bundler b.
func (b *Bundler) canFit(bu *bundle, size int) bool {
return (b.BundleByteLimit <= 0 || bu.size+size <= b.BundleByteLimit) &&
(b.BundleCountThreshold <= 0 || bu.items.Len() < b.BundleCountThreshold)
}
// Add adds item to the current bundle. It marks the bundle for handling and
// starts a new one if any of the thresholds or limits are exceeded.
// The type of item must be assignable to the itemExample parameter of the NewBundler
// method, otherwise there will be a panic.
//
// If the item's size exceeds the maximum bundle size (Bundler.BundleByteLimit), then
// the item can never be handled. Add returns ErrOversizedItem in this case.
//
// If adding the item would exceed the maximum memory allowed
// (Bundler.BufferedByteLimit) or an AddWait call is blocked waiting for
// memory, Add returns ErrOverflow.
//
// Add never blocks.
func (b *Bundler) Add(item interface{}, size int) error {
if err := b.setMode(add); err != nil {
return err
}
// If this item exceeds the maximum size of a bundle,
// we can never send it.
if b.BundleByteLimit > 0 && size > b.BundleByteLimit {
return ErrOversizedItem
}
// If adding this item would exceed our allotted memory
// footprint, we can't accept it.
// (TryAcquire also returns false if anything is waiting on the semaphore,
// so calls to Add and AddWait shouldn't be mixed.)
b.initSemaphores()
if !b.sem.TryAcquire(int64(size)) {
return ErrOverflow
}
b.mu.Lock()
defer b.mu.Unlock()
return b.add(item, size)
}
// add adds item to the tail of the bundle queue or curBundle depending on space
// and nil-ness (see inline comments). It marks curBundle for handling (by
// appending it to the queue) if any of the thresholds or limits are exceeded.
// curBundle is lazily initialized. It requires that b.mu is locked.
func (b *Bundler) add(item interface{}, size int) error {
// If we don't have a curBundle, see if we can add to the queue tail.
if b.tail != nil && b.curBundle == nil && b.canFit(b.tail, size) {
b.tail.add(item, size)
return nil
}
// If we can't fit in the existing curBundle, move it onto the queue.
if b.curBundle != nil && !b.canFit(b.curBundle, size) {
b.enqueueCurBundle()
}
// Create a curBundle if we don't have one.
if b.curBundle == nil {
b.curFlush.Add(1)
b.curBundle = &bundle{
items: b.itemSliceZero,
flush: b.curFlush,
}
}
// Add the item.
b.curBundle.add(item, size)
// If curBundle is ready for handling, move it to the queue.
if b.curBundle.size >= b.BundleByteThreshold ||
b.curBundle.items.Len() == b.BundleCountThreshold {
b.enqueueCurBundle()
}
// If we created a new bundle and it wasn't immediately handled, set a timer
if b.curBundle != nil && b.flushTimer == nil {
b.flushTimer = time.AfterFunc(b.DelayThreshold, b.tryHandleBundles)
}
return nil
}
// tryHandleBundles is the timer callback that handles or queues any current
// bundle after DelayThreshold time, even if the bundle isn't completely full.
func (b *Bundler) tryHandleBundles() {
b.mu.Lock()
b.enqueueCurBundle()
b.mu.Unlock()
}
// next returns the next bundle that is ready for handling and removes it from
// the internal queue. It requires that b.mu is locked.
func (b *Bundler) next() *bundle {
if b.head == nil {
return nil
}
out := b.head
b.head = b.head.next
if b.head == nil {
b.tail = nil
}
out.next = nil
return out
}
// handle calls the user-specified handler on the given bundle. handle is
// intended to be run as a goroutine. After the handler returns, we update the
// byte total. handle continues processing additional bundles that are ready.
// If no more bundles are ready, the handler count is decremented and the
// goroutine ends.
func (b *Bundler) handle(bu *bundle) {
for bu != nil {
b.handler(bu.items.Interface())
bu = b.postHandle(bu)
}
}
func (b *Bundler) postHandle(bu *bundle) *bundle {
b.mu.Lock()
defer b.mu.Unlock()
b.sem.Release(int64(bu.size))
bu.flush.Done()
bu = b.next()
if bu == nil {
b.handlerCount--
}
return bu
}
// AddWait adds item to the current bundle. It marks the bundle for handling and
// starts a new one if any of the thresholds or limits are exceeded.
//
// If the item's size exceeds the maximum bundle size (Bundler.BundleByteLimit), then
// the item can never be handled. AddWait returns ErrOversizedItem in this case.
//
// If adding the item would exceed the maximum memory allowed (Bundler.BufferedByteLimit),
// AddWait blocks until space is available or ctx is done.
//
// Calls to Add and AddWait should not be mixed on the same Bundler.
func (b *Bundler) AddWait(ctx context.Context, item interface{}, size int) error {
if err := b.setMode(addWait); err != nil {
return err
}
// If this item exceeds the maximum size of a bundle,
// we can never send it.
if b.BundleByteLimit > 0 && size > b.BundleByteLimit {
return ErrOversizedItem
}
// If adding this item would exceed our allotted memory footprint, block
// until space is available. The semaphore is FIFO, so there will be no
// starvation.
b.initSemaphores()
if err := b.sem.Acquire(ctx, int64(size)); err != nil {
return err
}
b.mu.Lock()
defer b.mu.Unlock()
return b.add(item, size)
}
// Flush invokes the handler for all remaining items in the Bundler and waits
// for it to return.
func (b *Bundler) Flush() {
b.mu.Lock()
// If a curBundle is pending, move it to the queue.
b.enqueueCurBundle()
// Store a pointer to the WaitGroup that counts outstanding bundles
// in the current flush and create a new one to track the next flush.
wg := b.curFlush
b.curFlush = &sync.WaitGroup{}
// Flush must wait for all prior, outstanding flushes to complete.
// We use a channel to communicate completion between each flush in
// the sequence.
prev := b.prevFlush
next := make(chan bool)
b.prevFlush = next
b.mu.Unlock()
// Wait until the previous flush is finished.
if prev != nil {
<-prev
}
// Wait until this flush is finished.
wg.Wait()
// Allow the next flush to finish.
close(next)
}
|