teleport/lib/utils/concurrentqueue/queue.go

/*
Copyright 2021 Gravitational, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package concurrentqueue

import (
	"sync"
)

type config struct {
	workers   int
	capacity  int
	inputBuf  int
	outputBuf int
}

// Option is a queue configuration option.
type Option func(*config)

// Workers is the number of concurrent workers to be used for the
// queue.  Defaults to 4.
func Workers(w int) Option {
	return func(cfg *config) {
		cfg.workers = w
	}
}

// Capacity is the amount of "in flight" items the queue can hold, not including
// any extra capacity in the input/output buffers. Assuming unbuffered input/output,
// this is the number of items that can be pushed to a queue before it begins to exhibit
// backpressure.  A value of 8-16x the number of workers is probably a reasonable choice
// for most applications.  Note that a queue always has a capacity at least equal to the
// number of workers, so `New(fn,Workers(7),Capacity(3))` results in a queue with capacity
// `7`. Defaults to 64.
func Capacity(c int) Option {
	return func(cfg *config) {
		cfg.capacity = c
	}
}

// InputBuf is the amount of buffer to be used in the input/push channel. Defaults to 0 (unbuffered).
func InputBuf(b int) Option {
	return func(cfg *config) {
		cfg.inputBuf = b
	}
}

// OutputBuf is the amount of buffer to be used in the output/pop channel.  Allocating output
// buffer space may improve performance when items are able to be popped in quick succession.
// Defaults to 0 (unbuffered).
func OutputBuf(b int) Option {
	return func(cfg *config) {
		cfg.outputBuf = b
	}
}

// item is the internal "work item" used by the queue.  it holds a value, and a nonce indicating the
// order in which the value was received.
type item struct {
	value interface{}
	nonce uint64
}

// Queue is a data processing helper which uses a worker pool to apply a closure to a series of
// values concurrently, preserving the correct ordering of results.  It is essentially the concurrent
// equivalent of this:
//    for msg := range inputChannel {
//        outputChannel <- workFunction(msg)
//    }
// In order to prevent indefinite memory growth within the queue due to slow consumption and/or
// workers, the queue will exert backpressure over its input channel once a configurable capacity
// is reached.
type Queue struct {
	input     chan interface{}
	output    chan interface{}
	closeOnce sync.Once
	done      chan struct{}
}

// Push accesses the queue's input channel.  The type of sent values must match
// that expected by the queue's work function.  If the queue was configured with
// a buffered input/push channel, non-blocking sends can be used as a heuristic for
// detecting backpressure due to queue capacity.  This is not a perfect test, but
// the rate of false positives will be extremely low for a queue with a decent
// capacity and non-trivial work function.
func (q *Queue) Push() chan<- interface{} {
	return q.input
}

// Pop accesses the queue's output channel.  The type of the received value
// will match the output of the work function.
func (q *Queue) Pop() <-chan interface{} {
	return q.output
}

// Done signals closure of the queue.
func (q *Queue) Done() <-chan struct{} {
	return q.done
}

// Close permanently terminates all background operations.  If the queue is not drained before
// closure, items may be lost.
func (q *Queue) Close() error {
	q.closeOnce.Do(func() {
		close(q.done)
	})
	return nil
}

// New builds a new queue instance around the supplied work function.
func New(workfn func(interface{}) interface{}, opts ...Option) *Queue {
	const defaultWorkers = 4
	const defaultCapacity = 64

	var cfg config
	for _, opt := range opts {
		opt(&cfg)
	}

	if cfg.workers < 1 {
		cfg.workers = defaultWorkers
	}
	if cfg.capacity < 1 {
		cfg.capacity = defaultCapacity
	}

	// capacity must be at least equal to the number of workers or else workers
	// will always be idle.
	if cfg.capacity < cfg.workers {
		cfg.capacity = cfg.workers
	}

	q := &Queue{
		input:  make(chan interface{}, cfg.inputBuf),
		output: make(chan interface{}, cfg.outputBuf),
		done:   make(chan struct{}),
	}

	go q.run(workfn, cfg)

	return q
}

// run spawns background tasks and then blocks on collection/reordering routine.
func (q *Queue) run(workfn func(interface{}) interface{}, cfg config) {
	// internal worker input/output channels. due to the semaphore channel below,
	// sends on these channels never block, as they are each allocated with sufficient
	// capacity to hold all in-flight items.
	workerIn, workerOut := make(chan item, cfg.capacity), make(chan item, cfg.capacity)

	// semaphore channel used to limit the number of "in flight" items.  a message is added prior to accepting
	// every input and removed upon emission of every output.  this allows us to exert backpressure and prevent
	// uncapped memory growth due to a slow worker.  this also keeps the queue's "capacity" consistent, regardless
	// of whether we are experiencing general slowness or a "head of line blocking" type scenario.
	sem := make(chan struct{}, cfg.capacity)

	// spawn workers
	for i := 0; i < cfg.workers; i++ {
		go func() {
			for {
				var itm item
				select {
				case itm = <-workerIn:
				case <-q.done:
					return
				}

				itm.value = workfn(itm.value)

				select {
				case workerOut <- itm:
				default:
					panic("cq worker output channel already full (semaphore violation)")
				}
			}
		}()
	}

	go q.distribute(workerIn, sem)

	q.collect(workerOut, sem)
}

// distribute takes inbound work items, applies a nonce, and then distributes
// them to the workers.
func (q *Queue) distribute(workerIn chan<- item, sem chan struct{}) {
	var nonce uint64
	for {
		// we are about to accept an input, add an item to the in-flight semaphore channel
		select {
		case sem <- struct{}{}:
		case <-q.done:
			return
		}

		var value interface{}
		select {
		case value = <-q.input:
		case <-q.done:
			return
		}

		select {
		case workerIn <- item{value: value, nonce: nonce}:
		default:
			panic("cq worker input channel already full (semaphore violation)")
		}

		nonce++
	}
}

// collect takes the potentially disordered worker output and unifies it into
// an ordered output.
func (q *Queue) collect(workerOut <-chan item, sem chan struct{}) {
	// items that cannot be emitted yet (due to arriving out of order),
	// stored in mapping of nonce => value.
	queue := make(map[uint64]interface{})

	// the nonce of the item we need to emit next.  incremented upon
	// successful emission.
	var nonce uint64

	// output value to be emitted (if any).  note that nil is a valid
	// output value, so we cannot inspect this value directly to
	// determine our state.
	var out interface{}

	// emit indicates whether or not we should be attempting to emit
	// the output value.
	var emit bool

	for {
		outc := q.output
		if !emit {
			// we do not have the next output item yet, do not attempt to send
			outc = nil
		}

		select {
		case itm := <-workerOut:
			if itm.nonce == nonce {
				// item matches current nonce, proceed directly to emitting state
				out, emit = itm.value, true
			} else {
				// item does not match current nonce, store it in queue
				queue[itm.nonce] = itm.value
			}
		case outc <- out:
			// successfully sent current item, increment nonce and setup next
			// output if it is present.
			nonce++
			out, emit = queue[nonce]
			delete(queue, nonce)

			// event has been emitted, remove an item from in-flight semaphore channel
			select {
			case <-sem:
			default:
				panic("cq sem channel already empty (semaphore violation)")
			}
		case <-q.done:
			return
		}

	}
}