<commit_before>// Copyright 2011 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 oglemock

import (
	"errors"
	"fmt"
	"math"
	"reflect"
	"sync"
)

// PartialExpecation is a function that should be called exactly once with
// expected arguments or matchers in order to set up an expected method call.
// See Controller.ExpectMethodCall below. It returns an expectation that can be
// further modified (e.g. by calling WillOnce).
//
// If the arguments are of the wrong type, the function reports a fatal error
// and returns nil.
type PartialExpecation func(...interface{}) Expectation

// Controller represents an object that implements the central logic of
// oglemock: recording and verifying expectations, responding to mock method
// calls, and so on.
type Controller interface {
	// ExpectCall expresses an expectation that the method of the given name
	// should be called on the supplied mock object. It returns a function that
	// should be called with the expected arguments, matchers for the arguments,
	// or a mix of both.
	//
	// fileName and lineNumber should indicate the line on which the expectation
	// was made, if known.
	//
	// For example:
	//
	//     mockWriter := [...]
	//     controller.ExpectCall(mockWriter, "Write", "foo.go", 17)(ElementsAre(0x1))
	//         .WillOnce(Return(1, nil))
	//
	// If the mock object doesn't have a method of the supplied name, the
	// function reports a fatal error and returns nil.
	ExpectCall(
		o MockObject,
		methodName string,
		fileName string,
		lineNumber int) PartialExpecation

	// Finish causes the controller to check for any unsatisfied expectations,
	// and report them as errors if they exist.
	//
	// The controller may panic if any of its methods (including this one) are
	// called after Finish is called.
	Finish()

	// HandleMethodCall looks for a registered expectation matching the call of
	// the given method on mock object o, invokes the appropriate action (if
	// any), and returns the values returned by that action (if any).
	//
	// If the action returns nothing, the controller returns zero values. If
	// there is no matching expectation, the controller reports an error and
	// returns zero values.
	//
	// If the mock object doesn't have a method of the supplied name, the
	// arguments are of the wrong type, or the action returns the wrong types,
	// the function reports a fatal error.
	//
	// HandleMethodCall is exported for the sake of mock implementations, and
	// should not be used directly.
	HandleMethodCall(
		o MockObject,
		methodName string,
		fileName string,
		lineNumber int,
		args []interface{}) []interface{}
}

// methodMap represents a map from method name to set of expectations for that
// method.
type methodMap map[string][]*InternalExpectation

// objectMap represents a map from mock object ID to a methodMap for that object.
type objectMap map[uintptr]methodMap

// NewController sets up a fresh controller, without any expectations set, and
// configures the controller to use the supplied error reporter.
func NewController(reporter ErrorReporter) Controller {
	return &controllerImpl{reporter, sync.RWMutex{}, objectMap{}}
}

type controllerImpl struct {
	reporter ErrorReporter

	mutex sync.RWMutex
	expectationsByObject objectMap  // Protected by mutex
}

// Return the list of registered expectations for the named method of the
// supplied object, or an empty slice if none have been registered. When this
// method returns, it is guaranteed that c.expectationsByObject has an entry
// for the object.
//
// c.mutex must be held for reading.
func (c *controllerImpl) getExpectationsLocked(
	o MockObject,
	methodName string) []*InternalExpectation {
	id := o.Oglemock_Id()

	// Look up the mock object.
	expectationsByMethod, ok := c.expectationsByObject[id]
	if !ok {
		expectationsByMethod = methodMap{}
		c.expectationsByObject[id] = expectationsByMethod
	}

	result, ok := expectationsByMethod[methodName]
	if !ok {
		return []*InternalExpectation{}
	}

	return result
}

// Add an expectation to the list registered for the named method of the
// supplied mock object.
//
// c.mutex must be held for writing.
func (c *controllerImpl) addExpectationLocked(
	o MockObject,
	methodName string,
	exp *InternalExpectation) {
	// Get the existing list.
	existing := c.getExpectationsLocked(o, methodName)

	// Store a modified list.
	id := o.Oglemock_Id()
	c.expectationsByObject[id][methodName] = append(existing, exp)
}

func (c *controllerImpl) ExpectCall(
	o MockObject,
	methodName string,
	fileName string,
	lineNumber int) PartialExpecation {
	// Find the signature for the requested method.
	ov := reflect.ValueOf(o)
	method := ov.MethodByName(methodName)
	if method.Kind() == reflect.Invalid {
		c.reporter.ReportFatalError(
			fileName,
			lineNumber,
			errors.New("Unknown method: " + methodName))
		return nil
	}

	partialAlreadyCalled := false  // Protected by c.mutex
	return func(args ...interface{}) Expectation {
		c.mutex.Lock()
		defer c.mutex.Unlock()

		// This function should only be called once.
		if partialAlreadyCalled {
			c.reporter.ReportFatalError(
				fileName,
				lineNumber,
				errors.New("Partial expectation called more than once."))
			return nil
		}

		partialAlreadyCalled = true

		// Make sure that the number of args is legal. Keep in mind that the
		// method's type has an extra receiver arg.
		if len(args) != method.Type().NumIn() {
			c.reporter.ReportFatalError(
				fileName,
				lineNumber,
				errors.New(
					fmt.Sprintf(
						"Expectation for %s given wrong number of arguments: " +
						"expected %d, got %d.",
						methodName,
						method.Type().NumIn(),
						len(args))))
			return nil
		}

		// Create an expectation and insert it into the controller's map.
		exp := InternalNewExpectation(
			c.reporter,
			method.Type(),
			args,
			fileName,
			lineNumber)

		c.addExpectationLocked(o, methodName, exp)

		// Return the expectation to the user.
		return exp
	}
}

func (c *controllerImpl) Finish() {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	// Check whether the minimum cardinality for each registered expectation has
	// been satisfied.
	for _, expectationsByMethod := range c.expectationsByObject {
		for methodName, expectations := range expectationsByMethod {
			for _, exp := range expectations {
				exp.mutex.Lock()
				defer exp.mutex.Unlock()

				minCardinality, _ := computeCardinalityLocked(exp)
				if exp.NumMatches < minCardinality {
					c.reporter.ReportError(
						exp.FileName,
						exp.LineNumber,
						errors.New(
							fmt.Sprintf(
								"Unsatisfied expectation; expected %s to be called " +
								"at least %d times; called %d times.",
								methodName,
								minCardinality,
								exp.NumMatches)))
				}
			}
		}
	}
}

// expectationMatches checks the matchers for the expectation against the
// supplied arguments.
func expectationMatches(exp *InternalExpectation, args []interface{}) bool {
	matchers := exp.ArgMatchers
	if len(args) != len(matchers) {
		panic("expectationMatches: len(args)")
	}

	// Check each matcher.
	for i, matcher := range matchers {
		if err := matcher.Matches(args[i]); err != nil {
			return false
		}
	}

	return true
}

// Return the expectation that matches the supplied arguments. If there is more
// than one such expectation, the one furthest along in the list for the method
// is returned. If there is no such expectation, nil is returned.
//
// c.mutex must be held for reading.
func (c *controllerImpl) chooseExpectationLocked(
	o MockObject,
	methodName string,
	args []interface{}) *InternalExpectation {
	// Do we have any expectations for this method?
	expectations := c.getExpectationsLocked(o, methodName)
	if len(expectations) == 0 {
		return nil
	}

	for i := len(expectations) - 1; i >= 0; i-- {
		if (expectationMatches(expectations[i], args)) {
			return expectations[i]
		}
	}

	return nil
}

// makeZeroReturnValues creates a []interface{} containing appropriate zero
// values for returning from the supplied method type.
func makeZeroReturnValues(signature reflect.Type) []interface{} {
	result := make([]interface{}, signature.NumOut())

	for i, _ := range result {
		outType := signature.Out(i)
		zeroVal := reflect.Zero(outType)
		result[i] = zeroVal.Interface()
	}

	return result
}

// computeCardinality decides on the [min, max] range of the number of expected
// matches for the supplied expectations, according to the rules documented in
// expectation.go.
//
// exp.mutex must be held for reading.
func computeCardinalityLocked(exp *InternalExpectation) (min, max uint) {
	// Explicit cardinality.
	if exp.ExpectedNumMatches >= 0 {
		min = uint(exp.ExpectedNumMatches)
		max = min
		return
	}

	// Implicit count based on one-time actions.
	if len(exp.OneTimeActions) != 0 {
		min = uint(len(exp.OneTimeActions))
		max = min

		// If there is a fallback action, this is only a lower bound.
		if exp.FallbackAction != nil {
			max = math.MaxUint32
		}

		return
	}

	// Implicit lack of restriction based on a fallback action being configured.
	if exp.FallbackAction != nil {
		min = 0
		max = math.MaxUint32
		return
	}

	// Implicit cardinality of one.
	min = 1
	max = 1
	return
}

// chooseAction returns the action that should be invoked for the i'th match to
// the supplied expectation (counting from zero). If the implicit "return zero
// values" action should be used, it returns nil.
//
// exp.mutex must be held for reading.
func chooseActionLocked(i uint, exp *InternalExpectation) Action {
	// Exhaust one-time actions first.
	if i < uint(len(exp.OneTimeActions)) {
		return exp.OneTimeActions[i]
	}

	// Fallback action (or nil if none is configured).
	return exp.FallbackAction
}

func (c *controllerImpl) HandleMethodCall(
	o MockObject,
	methodName string,
	fileName string,
	lineNumber int,
	args []interface{}) []interface{} {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	// Find the signature for the requested method.
	ov := reflect.ValueOf(o)
	method := ov.MethodByName(methodName)
	if method.Kind() == reflect.Invalid {
		c.reporter.ReportFatalError(
			fileName,
			lineNumber,
			errors.New("Unknown method: " + methodName))
		return nil
	}

	// HACK(jacobsa): Make sure we got the correct number of arguments. This will
	// need to be refined when issue #5 (variadic methods) is handled.
	if len(args) != method.Type().NumIn() {
		c.reporter.ReportFatalError(
			fileName,
			lineNumber,
			errors.New(
				fmt.Sprintf(
					"Wrong number of arguments: expected %d; got %d",
					method.Type().NumIn(),
					len(args))))
		return nil
	}

	// Find an expectation matching this call.
	expectation := c.chooseExpectationLocked(o, methodName, args)
	if expectation == nil {
		c.reporter.ReportError(
			fileName,
			lineNumber,
			errors.New(
				fmt.Sprintf("Unexpected call to %s with args: %v", methodName, args)))

		return makeZeroReturnValues(method.Type())
	}

	expectation.mutex.Lock()
	defer expectation.mutex.Unlock()

	// Increase the number of matches recorded, and check whether we're over the
	// number expected.
	expectation.NumMatches++
	_, maxCardinality := computeCardinalityLocked(expectation)
	if expectation.NumMatches > maxCardinality {
		c.reporter.ReportError(
			expectation.FileName,
			expectation.LineNumber,
			errors.New(
				fmt.Sprintf(
					"Unexpected call to %s: " +
						"expected to be called at most %d times; called %d times.",
					methodName,
					maxCardinality,
					expectation.NumMatches)))

		return makeZeroReturnValues(method.Type())
	}

	// Choose an action to invoke. If there is none, just return zero values.
	action := chooseActionLocked(expectation.NumMatches - 1, expectation)
	if action == nil {
		return makeZeroReturnValues(method.Type())
	}

	// Let the action take over.
	return action.Invoke(args)
}
<commit_msg>Updated implementation.<commit_after>// Copyright 2011 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 oglemock

import (
	"errors"
	"fmt"
	"math"
	"reflect"
	"sync"
)

// PartialExpecation is a function that should be called exactly once with
// expected arguments or matchers in order to set up an expected method call.
// See Controller.ExpectMethodCall below. It returns an expectation that can be
// further modified (e.g. by calling WillOnce).
//
// If the arguments are of the wrong type, the function reports a fatal error
// and returns nil.
type PartialExpecation func(...interface{}) Expectation

// Controller represents an object that implements the central logic of
// oglemock: recording and verifying expectations, responding to mock method
// calls, and so on.
type Controller interface {
	// ExpectCall expresses an expectation that the method of the given name
	// should be called on the supplied mock object. It returns a function that
	// should be called with the expected arguments, matchers for the arguments,
	// or a mix of both.
	//
	// fileName and lineNumber should indicate the line on which the expectation
	// was made, if known.
	//
	// For example:
	//
	//     mockWriter := [...]
	//     controller.ExpectCall(mockWriter, "Write", "foo.go", 17)(ElementsAre(0x1))
	//         .WillOnce(Return(1, nil))
	//
	// If the mock object doesn't have a method of the supplied name, the
	// function reports a fatal error and returns nil.
	ExpectCall(
		o MockObject,
		methodName string,
		fileName string,
		lineNumber int) PartialExpecation

	// Finish causes the controller to check for any unsatisfied expectations,
	// and report them as errors if they exist.
	//
	// The controller may panic if any of its methods (including this one) are
	// called after Finish is called.
	Finish()

	// HandleMethodCall looks for a registered expectation matching the call of
	// the given method on mock object o, invokes the appropriate action (if
	// any), and returns the values returned by that action (if any).
	//
	// If the action returns nothing, the controller returns zero values. If
	// there is no matching expectation, the controller reports an error and
	// returns zero values.
	//
	// If the mock object doesn't have a method of the supplied name, the
	// arguments are of the wrong type, or the action returns the wrong types,
	// the function reports a fatal error.
	//
	// HandleMethodCall is exported for the sake of mock implementations, and
	// should not be used directly.
	HandleMethodCall(
		o MockObject,
		methodName string,
		fileName string,
		lineNumber int,
		args []interface{}) []interface{}
}

// methodMap represents a map from method name to set of expectations for that
// method.
type methodMap map[string][]*InternalExpectation

// objectMap represents a map from mock object ID to a methodMap for that object.
type objectMap map[uintptr]methodMap

// NewController sets up a fresh controller, without any expectations set, and
// configures the controller to use the supplied error reporter.
func NewController(reporter ErrorReporter) Controller {
	return &controllerImpl{reporter, sync.RWMutex{}, objectMap{}}
}

type controllerImpl struct {
	reporter ErrorReporter

	mutex sync.RWMutex
	expectationsByObject objectMap  // Protected by mutex
}

// Return the list of registered expectations for the named method of the
// supplied object, or an empty slice if none have been registered. When this
// method returns, it is guaranteed that c.expectationsByObject has an entry
// for the object.
//
// c.mutex must be held for reading.
func (c *controllerImpl) getExpectationsLocked(
	o MockObject,
	methodName string) []*InternalExpectation {
	id := o.Oglemock_Id()

	// Look up the mock object.
	expectationsByMethod, ok := c.expectationsByObject[id]
	if !ok {
		expectationsByMethod = methodMap{}
		c.expectationsByObject[id] = expectationsByMethod
	}

	result, ok := expectationsByMethod[methodName]
	if !ok {
		return []*InternalExpectation{}
	}

	return result
}

// Add an expectation to the list registered for the named method of the
// supplied mock object.
//
// c.mutex must be held for writing.
func (c *controllerImpl) addExpectationLocked(
	o MockObject,
	methodName string,
	exp *InternalExpectation) {
	// Get the existing list.
	existing := c.getExpectationsLocked(o, methodName)

	// Store a modified list.
	id := o.Oglemock_Id()
	c.expectationsByObject[id][methodName] = append(existing, exp)
}

func (c *controllerImpl) ExpectCall(
	o MockObject,
	methodName string,
	fileName string,
	lineNumber int) PartialExpecation {
	// Find the signature for the requested method.
	ov := reflect.ValueOf(o)
	method := ov.MethodByName(methodName)
	if method.Kind() == reflect.Invalid {
		c.reporter.ReportFatalError(
			fileName,
			lineNumber,
			errors.New("Unknown method: " + methodName))
		return nil
	}

	partialAlreadyCalled := false  // Protected by c.mutex
	return func(args ...interface{}) Expectation {
		c.mutex.Lock()
		defer c.mutex.Unlock()

		// This function should only be called once.
		if partialAlreadyCalled {
			c.reporter.ReportFatalError(
				fileName,
				lineNumber,
				errors.New("Partial expectation called more than once."))
			return nil
		}

		partialAlreadyCalled = true

		// Make sure that the number of args is legal. Keep in mind that the
		// method's type has an extra receiver arg.
		if len(args) != method.Type().NumIn() {
			c.reporter.ReportFatalError(
				fileName,
				lineNumber,
				errors.New(
					fmt.Sprintf(
						"Expectation for %s given wrong number of arguments: " +
						"expected %d, got %d.",
						methodName,
						method.Type().NumIn(),
						len(args))))
			return nil
		}

		// Create an expectation and insert it into the controller's map.
		exp := InternalNewExpectation(
			c.reporter,
			method.Type(),
			args,
			fileName,
			lineNumber)

		c.addExpectationLocked(o, methodName, exp)

		// Return the expectation to the user.
		return exp
	}
}

func (c *controllerImpl) Finish() {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	// Check whether the minimum cardinality for each registered expectation has
	// been satisfied.
	for _, expectationsByMethod := range c.expectationsByObject {
		for methodName, expectations := range expectationsByMethod {
			for _, exp := range expectations {
				exp.mutex.Lock()
				defer exp.mutex.Unlock()

				minCardinality, _ := computeCardinalityLocked(exp)
				if exp.NumMatches < minCardinality {
					c.reporter.ReportError(
						exp.FileName,
						exp.LineNumber,
						errors.New(
							fmt.Sprintf(
								"Unsatisfied expectation; expected %s to be called " +
								"at least %d times; called %d times.",
								methodName,
								minCardinality,
								exp.NumMatches)))
				}
			}
		}
	}
}

// expectationMatches checks the matchers for the expectation against the
// supplied arguments.
func expectationMatches(exp *InternalExpectation, args []interface{}) bool {
	matchers := exp.ArgMatchers
	if len(args) != len(matchers) {
		panic("expectationMatches: len(args)")
	}

	// Check each matcher.
	for i, matcher := range matchers {
		if err := matcher.Matches(args[i]); err != nil {
			return false
		}
	}

	return true
}

// Return the expectation that matches the supplied arguments. If there is more
// than one such expectation, the one furthest along in the list for the method
// is returned. If there is no such expectation, nil is returned.
//
// c.mutex must be held for reading.
func (c *controllerImpl) chooseExpectationLocked(
	o MockObject,
	methodName string,
	args []interface{}) *InternalExpectation {
	// Do we have any expectations for this method?
	expectations := c.getExpectationsLocked(o, methodName)
	if len(expectations) == 0 {
		return nil
	}

	for i := len(expectations) - 1; i >= 0; i-- {
		if (expectationMatches(expectations[i], args)) {
			return expectations[i]
		}
	}

	return nil
}

// makeZeroReturnValues creates a []interface{} containing appropriate zero
// values for returning from the supplied method type.
func makeZeroReturnValues(signature reflect.Type) []interface{} {
	result := make([]interface{}, signature.NumOut())

	for i, _ := range result {
		outType := signature.Out(i)
		zeroVal := reflect.Zero(outType)
		result[i] = zeroVal.Interface()
	}

	return result
}

// computeCardinality decides on the [min, max] range of the number of expected
// matches for the supplied expectations, according to the rules documented in
// expectation.go.
//
// exp.mutex must be held for reading.
func computeCardinalityLocked(exp *InternalExpectation) (min, max uint) {
	// Explicit cardinality.
	if exp.ExpectedNumMatches >= 0 {
		min = uint(exp.ExpectedNumMatches)
		max = min
		return
	}

	// Implicit count based on one-time actions.
	if len(exp.OneTimeActions) != 0 {
		min = uint(len(exp.OneTimeActions))
		max = min

		// If there is a fallback action, this is only a lower bound.
		if exp.FallbackAction != nil {
			max = math.MaxUint32
		}

		return
	}

	// Implicit lack of restriction based on a fallback action being configured.
	if exp.FallbackAction != nil {
		min = 0
		max = math.MaxUint32
		return
	}

	// Implicit cardinality of one.
	min = 1
	max = 1
	return
}

// chooseAction returns the action that should be invoked for the i'th match to
// the supplied expectation (counting from zero). If the implicit "return zero
// values" action should be used, it returns nil.
//
// exp.mutex must be held for reading.
func chooseActionLocked(i uint, exp *InternalExpectation) Action {
	// Exhaust one-time actions first.
	if i < uint(len(exp.OneTimeActions)) {
		return exp.OneTimeActions[i]
	}

	// Fallback action (or nil if none is configured).
	return exp.FallbackAction
}

// Find an action for the method call, updating expectation match state in the
// process. Return either an action that should be invoked or a set of zero
// values to return immediately.
//
// This is split out from HandleMethodCall in order to more easily avoid
// invoking the action with locks held.
func (c *controllerImpl) chooseActionAndUpdateExpectations(
	o MockObject,
	methodName string,
	fileName string,
	lineNumber int,
	args []interface{},
) (action Action, zeroVals []interface{}) {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	// Find the signature for the requested method.
	ov := reflect.ValueOf(o)
	method := ov.MethodByName(methodName)
	if method.Kind() == reflect.Invalid {
		c.reporter.ReportFatalError(
			fileName,
			lineNumber,
			errors.New("Unknown method: " + methodName),
		)

		panic("ReportFatalError unexpectedly returned.")
	}

	// HACK(jacobsa): Make sure we got the correct number of arguments. This will
	// need to be refined when issue #5 (variadic methods) is handled.
	if len(args) != method.Type().NumIn() {
		c.reporter.ReportFatalError(
			fileName,
			lineNumber,
			errors.New(
				fmt.Sprintf(
					"Wrong number of arguments: expected %d; got %d",
					method.Type().NumIn(),
					len(args),
				),
			),
		)

		panic("ReportFatalError unexpectedly returned.")
	}

	// Find an expectation matching this call.
	expectation := c.chooseExpectationLocked(o, methodName, args)
	if expectation == nil {
		c.reporter.ReportError(
			fileName,
			lineNumber,
			errors.New(
				fmt.Sprintf("Unexpected call to %s with args: %v", methodName, args),
			),
		)

		zeroVals = makeZeroReturnValues(method.Type())
		return
	}

	expectation.mutex.Lock()
	defer expectation.mutex.Unlock()

	// Increase the number of matches recorded, and check whether we're over the
	// number expected.
	expectation.NumMatches++
	_, maxCardinality := computeCardinalityLocked(expectation)
	if expectation.NumMatches > maxCardinality {
		c.reporter.ReportError(
			expectation.FileName,
			expectation.LineNumber,
			errors.New(
				fmt.Sprintf(
					"Unexpected call to %s: " +
						"expected to be called at most %d times; called %d times.",
					methodName,
					maxCardinality,
					expectation.NumMatches,
				),
			),
		)

		zeroVals = makeZeroReturnValues(method.Type())
		return
	}

	// Choose an action to invoke. If there is none, just return zero values.
	action = chooseActionLocked(expectation.NumMatches - 1, expectation)
	if action == nil {
		zeroVals = makeZeroReturnValues(method.Type())
		return
	}

	// Let the action take over.
	return
}

func (c *controllerImpl) HandleMethodCall(
	o MockObject,
	methodName string,
	fileName string,
	lineNumber int,
	args []interface{},
) []interface{} {
	// Figure out whether to invoke an action or return zero values.
	action, zeroVals := c.chooseActionAndUpdateExpectations(
		o,
		methodName,
		fileName,
		lineNumber,
		args,
	)

	if action != nil {
		return action.Invoke(args)
	}

	return zeroVals
}
<|endoftext|>


<commit_before><commit_msg>Adding a small benchmarking tool<commit_after><|endoftext|>


<commit_before>/*
Package evidence handles all evidence storage and gossiping from detection to block proposal.
For the different types of evidence refer to the `evidence.go` file in the types package
or https://github.com/tendermint/spec/blob/master/spec/consensus/light-client/accountability.md.

## Gossiping

The core functionality begins with the evidence reactor (see reactor.
go) which operates both the sending and receiving of evidence.

The `Receive` function takes a list of evidence and does the following:
 1. breaks it down into individual evidence if it is `Composite Evidence` (see types/evidence.
go#ConflictingHeadersEvidence)
 2. checks that it does not already have the evidence stored
 3. verifies the evidence against the nodes state (see state/validation.go#VerifyEvidence)
 4. stores the evidence to a db and a concurrent list

The gossiping of evidence is initiated when a peer is added which starts a go routine to broadcast currently
uncommitted evidence at intervals of 60 seconds (set by the by broadcastEvidenceIntervalS).
It uses a concurrent list to store the evidence and before sending verifies that each evidence is still valid in the
sense that it has not exceeded the max evidence age and height. This should be set to be equal to the "trusting
period" (see types/params.go#EvidenceParams).

## Proposing

When a new block is being proposed (in state/execution.go#CreateProposalBlock),
`PendingEvidence(maxNum)` is called to send up to the maxNum number of uncommitted evidence, from the evidence store,
based on a priority that is a product of the age of the evidence and the voting power of the malicious validator.

Once the proposed evidence is submitted,
the evidence is marked as committed and is moved from the broadcasted set to the committed set (
the committed set is used to verify whether new evidence has actually already been submitted).
As a result it is also removed from the concurrent list so that it is no longer gossiped.

Last Update: 14/04/20
*/
package evidence
<commit_msg>made suggested changes<commit_after>/*
Package evidence handles all evidence storage and gossiping from detection to block proposal.
For the different types of evidence refer to the `evidence.go` file in the types package
or https://github.com/tendermint/spec/blob/master/spec/consensus/light-client/accountability.md.

Gossiping

The core functionality begins with the evidence reactor (see reactor.
go) which operates both the sending and receiving of evidence.

The `Receive` function takes a list of evidence and does the following:

1. Breaks it down into individual evidence if it is `Composite Evidence` (see types/evidence.go#ConflictingHeadersEvidence)

2. Checks that it does not already have the evidence stored

3. Verifies the evidence against the node's state (see state/validation.go#VerifyEvidence)

4. Stores the evidence to a db and a concurrent list

The gossiping of evidence is initiated when a peer is added which starts a go routine to broadcast currently
uncommitted evidence at intervals of 60 seconds (set by the by broadcastEvidenceIntervalS).
It uses a concurrent list to store the evidence and before sending verifies that each evidence is still valid in the
sense that it has not exceeded the max evidence age and height (see types/params.go#EvidenceParams).

Proposing

When a new block is being proposed (in state/execution.go#CreateProposalBlock),
`PendingEvidence(maxNum)` is called to send up to the maxNum number of uncommitted evidence, from the evidence store,
based on a priority that is a product of the age of the evidence and the voting power of the malicious validator.

Once the proposed evidence is submitted,
the evidence is marked as committed and is moved from the broadcasted set to the committed set (
the committed set is used to verify whether new evidence has actually already been submitted).
As a result it is also removed from the concurrent list so that it is no longer gossiped.
*/
package evidence
<|endoftext|>


<commit_before>package example

// Doer does things, sometimes repeatedly
type Doer interface {
	DoIt(task string, graciously bool) (int, error)
}

type Delegater struct {
	Delegate Doer
}

func (d *Delegater) DoSomething(task string) (int, error) {
	return d.Delegate.DoIt(task, false)
}
<commit_msg>Add documentation to Doer and Delegator<commit_after>package example

// Doer does things, sometimes graciously
type Doer interface {
	DoIt(task string, graciously bool) (int, error)
}

// Delegater employs a Doer to complete tasks
type Delegater struct {
	Delegate Doer
}

// DoSomething passes the work to Doer
func (d *Delegater) DoSomething(task string) (int, error) {
	return d.Delegate.DoIt(task, false)
}
<|endoftext|>


<commit_before>package sched

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"html/template"
	"io"
	"io/ioutil"
	"math"
	"net/http"
	"net/url"
	"time"

	"bosun.org/cmd/bosun/conf"
	"bosun.org/cmd/bosun/expr"
	"bosun.org/cmd/bosun/expr/parse"
	"bosun.org/opentsdb"
)

type Context struct {
	*State
	Alert *conf.Alert

	schedule    *Schedule
	runHistory  *RunHistory
	Attachments []*conf.Attachment
}

func (s *Schedule) Data(rh *RunHistory, st *State, a *conf.Alert, isEmail bool) *Context {
	c := Context{
		State:      st,
		Alert:      a,
		schedule:   s,
		runHistory: rh,
	}
	if isEmail {
		c.Attachments = make([]*conf.Attachment, 0)
	}
	return &c
}

type unknownContext struct {
	Time  time.Time
	Name  string
	Group expr.AlertKeys

	schedule *Schedule
}

func (s *Schedule) unknownData(t time.Time, name string, group expr.AlertKeys) *unknownContext {
	return &unknownContext{
		Time:     t,
		Group:    group,
		Name:     name,
		schedule: s,
	}
}

// Ack returns the URL to acknowledge an alert.
func (c *Context) Ack() string {
	return c.makeLink("/action", &url.Values{
		"type": []string{"ack"},
		"key":  []string{c.Alert.Name + c.State.Group.String()},
	})
}

// HostView returns the URL to the host view page.
func (c *Context) HostView(host string) string {
	return c.makeLink("/host", &url.Values{
		"time": []string{"1d-ago"},
		"host": []string{host},
	})
}

func (c *Context) makeLink(path string, v *url.Values) string {
	u := url.URL{
		Scheme:   "http",
		Host:     c.schedule.Conf.Hostname,
		Path:     path,
		RawQuery: v.Encode(),
	}
	return u.String()
}

func (c *Context) Expr(v string) string {
	p := url.Values{}
	p.Add("expr", base64.StdEncoding.EncodeToString([]byte(opentsdb.ReplaceTags(v, c.Group))))
	return c.makeLink("/expr", &p)
}

func (c *Context) Rule() (string, error) {
	t, err := c.schedule.Conf.AlertTemplateStrings()
	if err != nil {
		return "", err
	}
	p := url.Values{}
	adef := base64.StdEncoding.EncodeToString([]byte(t.Alerts[c.Alert.Name]))
	tdef := base64.StdEncoding.EncodeToString([]byte(t.Templates[c.Alert.Template.Name]))
	//There might be something better when we tie the notifications to evaluation time issue #395
	time := time.Now().UTC()
	p.Add("alert", adef)
	p.Add("template", tdef)
	p.Add("fromDate", time.Format("2006-01-02"))
	p.Add("fromTime", time.Format("15:04"))
	p.Add("template_group", c.Group.Tags())
	return c.makeLink("/rule", &p), nil
}

func (s *Schedule) ExecuteBody(w io.Writer, rh *RunHistory, a *conf.Alert, st *State, isEmail bool) ([]*conf.Attachment, error) {
	t := a.Template
	if t == nil || t.Body == nil {
		return nil, nil
	}
	c := s.Data(rh, st, a, isEmail)
	return c.Attachments, t.Body.Execute(w, c)
}

func (s *Schedule) ExecuteSubject(w io.Writer, rh *RunHistory, a *conf.Alert, st *State) error {
	t := a.Template
	if t == nil || t.Subject == nil {
		return nil
	}
	return t.Subject.Execute(w, s.Data(rh, st, a, false))
}

var error_body = template.Must(template.New("body_error_template").Parse(`

There was a runtime error processing alert {{.State.AlertKey}} using the {{.Alert.Template.Name}} template. The following errors occurred:

	{{if .Serr}}

Subject: {{.Serr}}

	{{end}}
	{{if .Berr}}

Body: {{.Berr}}

	{{end}}

Use this link to the rule page to correct this.

Generic Alert Information

Status: {{.Last.Status}}

Alert: {{.State.AlertKey}}

Computations

			Expression
			Value

	{{range .Computations}}

			{{.Text}}
			{{.Value}}

	{{end}}
`))

func (s *Schedule) ExecuteBadTemplate(s_err, b_err error, rh *RunHistory, a *conf.Alert, st *State) (*bytes.Buffer, *bytes.Buffer, error) {
	sub := "error: template rendering error in the "
	if s_err != nil {
		sub += "subject"
	}
	if s_err != nil && b_err != nil {
		sub += " and "
	}
	if b_err != nil {
		sub += "body"
	}
	sub += fmt.Sprintf(" for alert %v", st.AlertKey())
	c := struct {
		Serr, Berr error
		*Context
	}{
		Serr:    s_err,
		Berr:    b_err,
		Context: s.Data(rh, st, a, true),
	}
	body := new(bytes.Buffer)
	error_body.Execute(body, c)
	return bytes.NewBufferString(sub), body, nil
}

func (c *Context) eval(v interface{}, filter bool, series bool, autods int) (expr.ResultSlice, string, error) {
	var e *expr.Expr
	var err error
	switch v := v.(type) {
	case string:
		e, err = expr.New(v, c.schedule.Conf.Funcs())
	case *expr.Expr:
		e = v
	default:
		return nil, "", fmt.Errorf("expected string or expression, got %T (%v)", v, v)
	}
	if err != nil {
		return nil, "", fmt.Errorf("%v: %v", v, err)
	}
	if filter {
		e, err = expr.New(opentsdb.ReplaceTags(e.String(), c.State.Group), c.schedule.Conf.Funcs())
		if err != nil {
			return nil, "", err
		}
	}
	if series && e.Root.Return() != parse.TypeSeries {
		return nil, "", fmt.Errorf("egraph: requires an expression that returns a series")
	}
	t := time.Now().UTC()
	if c.AbnormalEvent() != nil {
		t = c.AbnormalEvent().Time
	}
	res, _, err := e.Execute(c.runHistory.Context, c.runHistory.GraphiteContext, nil, t, autods, c.Alert.UnjoinedOK, c.schedule.Search, c.schedule.Conf.AlertSquelched(c.Alert))
	if err != nil {
		return nil, "", fmt.Errorf("%s: %v", v, err)
	}
	return res.Results, e.String(), nil
}

// Lookup returns the value for a key in the lookup table for the context's tagset.
func (c *Context) Lookup(table, key string) (string, error) {
	return c.LookupAll(table, key, c.Group)
}

func (c *Context) LookupAll(table, key string, group interface{}) (string, error) {
	var t opentsdb.TagSet
	switch v := group.(type) {
	case string:
		var err error
		t, err = opentsdb.ParseTags(v)
		if err != nil {
			return "", err
		}
	case opentsdb.TagSet:
		t = v
	}
	l, ok := c.schedule.Conf.Lookups[table]
	if !ok {
		return "", fmt.Errorf("unknown lookup table %v", table)
	}
	if v, ok := l.ToExpr().Get(key, t); ok {
		return v, nil
	}
	return "", fmt.Errorf("no entry for key %v in table %v for tagset %v", key, table, c.Group)
}

// Eval executes the given expression and returns a value with corresponding
// tags to the context's tags. If no such result is found, the first result with
// nil tags is returned. If no such result is found, nil is returned.
func (c *Context) Eval(v interface{}) (interface{}, error) {
	res, _, err := c.eval(v, true, false, 0)
	if err != nil {
		return nil, err
	}
	if len(res) == 0 {
		return nil, fmt.Errorf("no results returned")
	}
	// TODO: don't choose a random result, make sure there's exactly 1
	return res[0].Value, nil
}

// EvalAll returns the executed expression.
func (c *Context) EvalAll(v interface{}) (interface{}, error) {
	res, _, err := c.eval(v, false, false, 0)
	return res, err
}

func (c *Context) IsEmail() bool {
	return c.Attachments != nil
}

func (c *Context) graph(v interface{}, filter bool) (interface{}, error) {
	res, title, err := c.eval(v, filter, true, 1000)
	if err != nil {
		return nil, err
	}
	var buf bytes.Buffer
	const width = 800
	const height = 600
	if c.IsEmail() {
		err := c.schedule.ExprPNG(nil, &buf, width, height, res, title, time.Now().UTC())
		if err != nil {
			return nil, err
		}
		name := fmt.Sprintf("%d.png", len(c.Attachments)+1)
		c.Attachments = append(c.Attachments, &conf.Attachment{
			Data:        buf.Bytes(),
			Filename:    name,
			ContentType: "image/png",
		})
		return template.HTML(fmt.Sprintf(`<img alt="%s" src="cid:%s" />`,
			template.HTMLEscapeString(fmt.Sprint(v)),
			name,
		)), nil
	}
	if err := c.schedule.ExprSVG(nil, &buf, width, height, res, title, time.Now().UTC()); err != nil {
		return nil, err
	}
	return template.HTML(buf.String()), nil
}

func (c *Context) Graph(v interface{}) (interface{}, error) {
	return c.graph(v, true)
}

func (c *Context) GraphAll(v interface{}) (interface{}, error) {
	return c.graph(v, false)
}

func (c *Context) GetMeta(metric, name string, v interface{}) (interface{}, error) {
	var t opentsdb.TagSet
	switch v := v.(type) {
	case string:
		var err error
		t, err = opentsdb.ParseTags(v)
		if err != nil {
			return t, err
		}
	case opentsdb.TagSet:
		t = v
	}
	meta := c.schedule.GetMetadata(metric, t)
	if name == "" {
		return meta, nil
	}
	for _, m := range meta {
		if m.Name == name {
			return m.Value, nil
		}
	}
	return nil, nil
}

func (c *Context) LeftJoin(q ...interface{}) (interface{}, error) {
	if len(q) < 2 {
		return nil, fmt.Errorf("need at least two expressions, got %v", len(q))
	}
	matrix := make([][]*expr.Result, 0)
	results := make([][]*expr.Result, len(q))
	for col, v := range q {
		res, _, err := c.eval(v, false, false, 0)
		if err != nil {
			return nil, err
		}
		results[col] = res
	}
	for row, first := range results[0] {
		matrix = append(matrix, make([]*expr.Result, len(q)))
		matrix[row][0] = first
		for col, res := range results[1:] {
			for _, r := range res {
				if first.Group.Subset(r.Group) {
					matrix[row][col+1] = r
					break
				}
				// Fill emtpy cells with NaN Value, so calling .Valie is not a nil pointer dereference
				matrix[row][col+1] = &expr.Result{Value: expr.Number(math.NaN())}
			}
		}
	}
	return matrix, nil
}

func (c *Context) HTTPGet(url string) string {
	resp, err := http.Get(url)
	if err != nil {
		return err.Error()
	}
	defer resp.Body.Close()
	if resp.StatusCode >= 300 {
		return fmt.Sprintf("%v: returned %v", url, resp.Status)
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err.Error()
	}
	return string(body)
}

func (c *Context) HTTPPost(url, bodyType, data string) string {
	resp, err := http.Post(url, bodyType, bytes.NewBufferString(data))
	if err != nil {
		return err.Error()
	}
	defer resp.Body.Close()
	if resp.StatusCode >= 300 {
		return fmt.Sprintf("%v: returned %v", url, resp.Status)
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err.Error()
	}
	return string(body)
}
<commit_msg>In some cases,after parsing expr's String() not equal input text,so use original input when filter is true<commit_after>package sched

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"html/template"
	"io"
	"io/ioutil"
	"math"
	"net/http"
	"net/url"
	"time"

	"bosun.org/cmd/bosun/conf"
	"bosun.org/cmd/bosun/expr"
	"bosun.org/cmd/bosun/expr/parse"
	"bosun.org/opentsdb"
)

type Context struct {
	*State
	Alert *conf.Alert

	schedule    *Schedule
	runHistory  *RunHistory
	Attachments []*conf.Attachment
}

func (s *Schedule) Data(rh *RunHistory, st *State, a *conf.Alert, isEmail bool) *Context {
	c := Context{
		State:      st,
		Alert:      a,
		schedule:   s,
		runHistory: rh,
	}
	if isEmail {
		c.Attachments = make([]*conf.Attachment, 0)
	}
	return &c
}

type unknownContext struct {
	Time  time.Time
	Name  string
	Group expr.AlertKeys

	schedule *Schedule
}

func (s *Schedule) unknownData(t time.Time, name string, group expr.AlertKeys) *unknownContext {
	return &unknownContext{
		Time:     t,
		Group:    group,
		Name:     name,
		schedule: s,
	}
}

// Ack returns the URL to acknowledge an alert.
func (c *Context) Ack() string {
	return c.makeLink("/action", &url.Values{
		"type": []string{"ack"},
		"key":  []string{c.Alert.Name + c.State.Group.String()},
	})
}

// HostView returns the URL to the host view page.
func (c *Context) HostView(host string) string {
	return c.makeLink("/host", &url.Values{
		"time": []string{"1d-ago"},
		"host": []string{host},
	})
}

func (c *Context) makeLink(path string, v *url.Values) string {
	u := url.URL{
		Scheme:   "http",
		Host:     c.schedule.Conf.Hostname,
		Path:     path,
		RawQuery: v.Encode(),
	}
	return u.String()
}

func (c *Context) Expr(v string) string {
	p := url.Values{}
	p.Add("expr", base64.StdEncoding.EncodeToString([]byte(opentsdb.ReplaceTags(v, c.Group))))
	return c.makeLink("/expr", &p)
}

func (c *Context) Rule() (string, error) {
	t, err := c.schedule.Conf.AlertTemplateStrings()
	if err != nil {
		return "", err
	}
	p := url.Values{}
	adef := base64.StdEncoding.EncodeToString([]byte(t.Alerts[c.Alert.Name]))
	tdef := base64.StdEncoding.EncodeToString([]byte(t.Templates[c.Alert.Template.Name]))
	//There might be something better when we tie the notifications to evaluation time issue #395
	time := time.Now().UTC()
	p.Add("alert", adef)
	p.Add("template", tdef)
	p.Add("fromDate", time.Format("2006-01-02"))
	p.Add("fromTime", time.Format("15:04"))
	p.Add("template_group", c.Group.Tags())
	return c.makeLink("/rule", &p), nil
}

func (s *Schedule) ExecuteBody(w io.Writer, rh *RunHistory, a *conf.Alert, st *State, isEmail bool) ([]*conf.Attachment, error) {
	t := a.Template
	if t == nil || t.Body == nil {
		return nil, nil
	}
	c := s.Data(rh, st, a, isEmail)
	return c.Attachments, t.Body.Execute(w, c)
}

func (s *Schedule) ExecuteSubject(w io.Writer, rh *RunHistory, a *conf.Alert, st *State) error {
	t := a.Template
	if t == nil || t.Subject == nil {
		return nil
	}
	return t.Subject.Execute(w, s.Data(rh, st, a, false))
}

var error_body = template.Must(template.New("body_error_template").Parse(`

There was a runtime error processing alert {{.State.AlertKey}} using the {{.Alert.Template.Name}} template. The following errors occurred:

	{{if .Serr}}

Subject: {{.Serr}}

	{{end}}
	{{if .Berr}}

Body: {{.Berr}}

	{{end}}

Use this link to the rule page to correct this.

Generic Alert Information

Status: {{.Last.Status}}

Alert: {{.State.AlertKey}}

Computations

			Expression
			Value

	{{range .Computations}}

			{{.Text}}
			{{.Value}}

	{{end}}
`))

func (s *Schedule) ExecuteBadTemplate(s_err, b_err error, rh *RunHistory, a *conf.Alert, st *State) (*bytes.Buffer, *bytes.Buffer, error) {
	sub := "error: template rendering error in the "
	if s_err != nil {
		sub += "subject"
	}
	if s_err != nil && b_err != nil {
		sub += " and "
	}
	if b_err != nil {
		sub += "body"
	}
	sub += fmt.Sprintf(" for alert %v", st.AlertKey())
	c := struct {
		Serr, Berr error
		*Context
	}{
		Serr:    s_err,
		Berr:    b_err,
		Context: s.Data(rh, st, a, true),
	}
	body := new(bytes.Buffer)
	error_body.Execute(body, c)
	return bytes.NewBufferString(sub), body, nil
}

func (c *Context) eval(v interface{}, filter bool, series bool, autods int) (expr.ResultSlice, string, error) {
	var e *expr.Expr
	var err error
	switch v := v.(type) {
	case string:
		e, err = expr.New(v, c.schedule.Conf.Funcs())
	case *expr.Expr:
		e = v
	default:
		return nil, "", fmt.Errorf("expected string or expression, got %T (%v)", v, v)
	}
	if err != nil {
		return nil, "", fmt.Errorf("%v: %v", v, err)
	}
	if filter {
		e, err = expr.New(opentsdb.ReplaceTags(e.Text, c.State.Group), c.schedule.Conf.Funcs())
		if err != nil {
			return nil, "", err
		}
	}
	if series && e.Root.Return() != parse.TypeSeries {
		return nil, "", fmt.Errorf("egraph: requires an expression that returns a series")
	}
	t := time.Now().UTC()
	if c.AbnormalEvent() != nil {
		t = c.AbnormalEvent().Time
	}
	res, _, err := e.Execute(c.runHistory.Context, c.runHistory.GraphiteContext, nil, t, autods, c.Alert.UnjoinedOK, c.schedule.Search, c.schedule.Conf.AlertSquelched(c.Alert))
	if err != nil {
		return nil, "", fmt.Errorf("%s: %v", v, err)
	}
	return res.Results, e.String(), nil
}

// Lookup returns the value for a key in the lookup table for the context's tagset.
func (c *Context) Lookup(table, key string) (string, error) {
	return c.LookupAll(table, key, c.Group)
}

func (c *Context) LookupAll(table, key string, group interface{}) (string, error) {
	var t opentsdb.TagSet
	switch v := group.(type) {
	case string:
		var err error
		t, err = opentsdb.ParseTags(v)
		if err != nil {
			return "", err
		}
	case opentsdb.TagSet:
		t = v
	}
	l, ok := c.schedule.Conf.Lookups[table]
	if !ok {
		return "", fmt.Errorf("unknown lookup table %v", table)
	}
	if v, ok := l.ToExpr().Get(key, t); ok {
		return v, nil
	}
	return "", fmt.Errorf("no entry for key %v in table %v for tagset %v", key, table, c.Group)
}

// Eval executes the given expression and returns a value with corresponding
// tags to the context's tags. If no such result is found, the first result with
// nil tags is returned. If no such result is found, nil is returned.
func (c *Context) Eval(v interface{}) (interface{}, error) {
	res, _, err := c.eval(v, true, false, 0)
	if err != nil {
		return nil, err
	}
	if len(res) == 0 {
		return nil, fmt.Errorf("no results returned")
	}
	// TODO: don't choose a random result, make sure there's exactly 1
	return res[0].Value, nil
}

// EvalAll returns the executed expression.
func (c *Context) EvalAll(v interface{}) (interface{}, error) {
	res, _, err := c.eval(v, false, false, 0)
	return res, err
}

func (c *Context) IsEmail() bool {
	return c.Attachments != nil
}

func (c *Context) graph(v interface{}, filter bool) (interface{}, error) {
	res, title, err := c.eval(v, filter, true, 1000)
	if err != nil {
		return nil, err
	}
	var buf bytes.Buffer
	const width = 800
	const height = 600
	if c.IsEmail() {
		err := c.schedule.ExprPNG(nil, &buf, width, height, res, title, time.Now().UTC())
		if err != nil {
			return nil, err
		}
		name := fmt.Sprintf("%d.png", len(c.Attachments)+1)
		c.Attachments = append(c.Attachments, &conf.Attachment{
			Data:        buf.Bytes(),
			Filename:    name,
			ContentType: "image/png",
		})
		return template.HTML(fmt.Sprintf(`<img alt="%s" src="cid:%s" />`,
			template.HTMLEscapeString(fmt.Sprint(v)),
			name,
		)), nil
	}
	if err := c.schedule.ExprSVG(nil, &buf, width, height, res, title, time.Now().UTC()); err != nil {
		return nil, err
	}
	return template.HTML(buf.String()), nil
}

func (c *Context) Graph(v interface{}) (interface{}, error) {
	return c.graph(v, true)
}

func (c *Context) GraphAll(v interface{}) (interface{}, error) {
	return c.graph(v, false)
}

func (c *Context) GetMeta(metric, name string, v interface{}) (interface{}, error) {
	var t opentsdb.TagSet
	switch v := v.(type) {
	case string:
		var err error
		t, err = opentsdb.ParseTags(v)
		if err != nil {
			return t, err
		}
	case opentsdb.TagSet:
		t = v
	}
	meta := c.schedule.GetMetadata(metric, t)
	if name == "" {
		return meta, nil
	}
	for _, m := range meta {
		if m.Name == name {
			return m.Value, nil
		}
	}
	return nil, nil
}

func (c *Context) LeftJoin(q ...interface{}) (interface{}, error) {
	if len(q) < 2 {
		return nil, fmt.Errorf("need at least two expressions, got %v", len(q))
	}
	matrix := make([][]*expr.Result, 0)
	results := make([][]*expr.Result, len(q))
	for col, v := range q {
		res, _, err := c.eval(v, false, false, 0)
		if err != nil {
			return nil, err
		}
		results[col] = res
	}
	for row, first := range results[0] {
		matrix = append(matrix, make([]*expr.Result, len(q)))
		matrix[row][0] = first
		for col, res := range results[1:] {
			for _, r := range res {
				if first.Group.Subset(r.Group) {
					matrix[row][col+1] = r
					break
				}
				// Fill emtpy cells with NaN Value, so calling .Valie is not a nil pointer dereference
				matrix[row][col+1] = &expr.Result{Value: expr.Number(math.NaN())}
			}
		}
	}
	return matrix, nil
}

func (c *Context) HTTPGet(url string) string {
	resp, err := http.Get(url)
	if err != nil {
		return err.Error()
	}
	defer resp.Body.Close()
	if resp.StatusCode >= 300 {
		return fmt.Sprintf("%v: returned %v", url, resp.Status)
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err.Error()
	}
	return string(body)
}

func (c *Context) HTTPPost(url, bodyType, data string) string {
	resp, err := http.Post(url, bodyType, bytes.NewBufferString(data))
	if err != nil {
		return err.Error()
	}
	defer resp.Body.Close()
	if resp.StatusCode >= 300 {
		return fmt.Sprintf("%v: returned %v", url, resp.Status)
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err.Error()
	}
	return string(body)
}
<|endoftext|>


<commit_before>// bulk_load_mongo loads a Mongo daemon with data from stdin.
//
// The caller is responsible for assuring that the database is empty before
// bulk load.
package main

import (
	"bufio"
	"encoding/binary"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"sync"
	"time"

	flatbuffers "github.com/google/flatbuffers/go"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"

	"github.com/influxdata/influxdb-comparisons/mongo_serialization"
	"github.com/influxdata/influxdb-comparisons/util/report"
	"strconv"
	"strings"
)

// Program option vars:
var (
	daemonUrl      string
	workers        int
	batchSize      int
	limit          int64
	doLoad         bool
	writeTimeout   time.Duration
	reportDatabase string
	reportHost     string
	reportUser     string
	reportPassword string
	reportTagsCSV  string
)

// Global vars
var (
	batchChan      chan *Batch
	inputDone      chan struct{}
	workersGroup   sync.WaitGroup
	reportTags     [][2]string
	reportHostname string
)

// Magic database constants
const (
	dbName              = "benchmark_db"
	pointCollectionName = "point_data"
)

// bufPool holds []byte instances to reduce heap churn.
var bufPool = &sync.Pool{
	New: func() interface{} {
		return make([]byte, 0, 1024)
	},
}

// Batch holds byte slices that will become mongo_serialization.Item instances.
type Batch [][]byte

func (b *Batch) ClearReferences() {
	*b = (*b)[:0]
}

// batchPool holds *Batch instances to reduce heap churn.
var batchPool = &sync.Pool{
	New: func() interface{} {
		return &Batch{}
	},
}

// Parse args:
func init() {
	flag.StringVar(&daemonUrl, "url", "localhost:27017", "Mongo URL.")

	flag.IntVar(&batchSize, "batch-size", 100, "Batch size (input items).")
	flag.IntVar(&workers, "workers", 1, "Number of parallel requests to make.")
	flag.Int64Var(&limit, "limit", -1, "Number of items to insert (default unlimited).")
	flag.DurationVar(&writeTimeout, "write-timeout", 10*time.Second, "Write timeout.")

	flag.BoolVar(&doLoad, "do-load", true, "Whether to write data. Set this flag to false to check input read speed.")

	flag.StringVar(&reportDatabase, "report-database", "database_benchmarks", "Database name where to store result metrics")
	flag.StringVar(&reportHost, "report-host", "", "Host to send result metrics")
	flag.StringVar(&reportUser, "report-user", "", "User for host to send result metrics")
	flag.StringVar(&reportPassword, "report-password", "", "User password for Host to send result metrics")
	flag.StringVar(&reportTagsCSV, "report-tags", "", "Comma separated k:v tags to send  alongside result metrics")

	flag.Parse()

	for i := 0; i < workers*batchSize; i++ {
		bufPool.Put(bufPool.New())
	}

	if reportHost != "" {
		fmt.Printf("results report destination: %v\n", reportHost)
		fmt.Printf("results report database: %v\n", reportDatabase)

		var err error
		reportHostname, err = os.Hostname()
		if err != nil {
			log.Fatalf("os.Hostname() error: %s", err.Error())
		}
		fmt.Printf("hostname for results report: %v\n", reportHostname)

		if reportTagsCSV != "" {
			pairs := strings.Split(reportTagsCSV, ",")
			for _, pair := range pairs {
				fields := strings.SplitN(pair, ":", 2)
				tagpair := [2]string{fields[0], fields[1]}
				reportTags = append(reportTags, tagpair)
			}
		}
		fmt.Printf("results report tags: %v\n", reportTags)
	}
}

func main() {
	if doLoad {
		mustCreateCollections(daemonUrl)
	}

	var session *mgo.Session

	if doLoad {
		var err error
		session, err = mgo.Dial(daemonUrl)
		if err != nil {
			log.Fatal(err)
		}

		session.SetMode(mgo.Eventual, false)

		defer session.Close()
	}

	batchChan = make(chan *Batch, workers*10)
	inputDone = make(chan struct{})

	for i := 0; i < workers; i++ {
		workersGroup.Add(1)
		go processBatches(session)
	}

	start := time.Now()
	itemsRead, bytesRead := scan(session, batchSize)

	<-inputDone
	close(batchChan)
	workersGroup.Wait()
	end := time.Now()
	took := end.Sub(start)
	itemRate := float64(itemsRead) / float64(took.Seconds())
	bytesRate := float64(bytesRead) / float64(took.Seconds())

	fmt.Printf("loaded %d values in %fsec with %d workers (mean values rate %f values/sec, %.2fMB/sec from stdin)\n", itemsRead, took.Seconds(), workers, itemRate, bytesRate/(1<<20))

	if reportHost != "" {
		//append db specific tags to custom tags
		reportTags = append(reportTags, [2]string{"write_timeout", strconv.Itoa(int(writeTimeout))})

		reportParams := &report.LoadReportParams{
			ReportParams: report.ReportParams{
				DBType:             "MongoDB",
				ReportDatabaseName: reportDatabase,
				ReportHost:         reportHost,
				ReportUser:         reportUser,
				ReportPassword:     reportPassword,
				ReportTags:         reportTags,
				Hostname:           reportHostname,
				DestinationUrl:     daemonUrl,
				Workers:            workers,
				ItemLimit:          int(limit),
			},
			IsGzip:    false,
			BatchSize: batchSize,
		}
		err := report.ReportLoadResult(reportParams, itemsRead, itemRate, bytesRate, took)

		if err != nil {
			log.Fatal(err)
		}
	}
}

// scan reads length-delimited flatbuffers items from stdin.
func scan(session *mgo.Session, itemsPerBatch int) (int64, int64) {
	var n int
	var itemsRead, bytesRead int64
	r := bufio.NewReaderSize(os.Stdin, 32<<20)

	start := time.Now()
	batch := batchPool.Get().(*Batch)
	lenBuf := make([]byte, 8)

	for {
		if itemsRead == limit {
			break
		}
		// get the serialized item length (this is the framing format)
		_, err := r.Read(lenBuf)
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatal(err.Error())
		}

		// ensure correct len of receiving buffer
		l := int(binary.LittleEndian.Uint64(lenBuf))
		itemBuf := bufPool.Get().([]byte)
		if cap(itemBuf) < l {
			itemBuf = make([]byte, l)
		}
		itemBuf = itemBuf[:l]

		// read the bytes and init the flatbuffer object
		totRead := 0
		for totRead < l {
			m, err := r.Read(itemBuf[totRead:])
			// (EOF is also fatal)
			if err != nil {
				log.Fatal(err.Error())
			}
			totRead += m
		}
		if totRead != len(itemBuf) {
			panic(fmt.Sprintf("reader/writer logic error, %d != %d", n, len(itemBuf)))
		}

		*batch = append(*batch, itemBuf)

		itemsRead++
		n++

		if n >= batchSize {
			bytesRead += int64(len(itemBuf))
			batchChan <- batch
			n = 0
			batch = batchPool.Get().(*Batch)
		}

		_ = start
		//if itemsRead > 0 && itemsRead%100000 == 0 {
		//	_ = start
		//	//took := (time.Now().UnixNano() - start.UnixNano())
		//	//if took >= 1e9 {
		//	//	tookUs := float64(took) / 1e3
		//	//	tookSec := float64(took) / 1e9
		//	//	fmt.Fprintf(os.Stderr, "itemsRead: %d, rate: %.0f/sec, lag: %.2fus/op\n",
		//	//		itemsRead, float64(itemsRead)/tookSec, tookUs/float64(itemsRead))
		//	//}
		//}
	}

	// Closing inputDone signals to the application that we've read everything and can now shut down.
	close(inputDone)

	return itemsRead, bytesRead
}

// processBatches reads byte buffers from batchChan, interprets them and writes
// them to the target server. Note that mgo forcibly incurs serialization
// overhead (it always encodes to BSON).
func processBatches(session *mgo.Session) {
	db := session.DB(dbName)

	type Tag struct {
		Key string `bson:"key"`
		Val string `bson:"val"`
	}

	type Point struct {
		// Use `string` here even though they are really `[]byte`.
		// This is so the mongo data is human-readable.
		MeasurementName string      `bson:"measurement"`
		FieldName       string      `bson:"field"`
		Timestamp       int64       `bson:"timestamp_ns"`
		Tags            []Tag       `bson:"tags"`
		Value           interface{} `bson:"value"`

		// a private union-like section
		longValue   int64
		doubleValue float64
		stringValue string
	}
	pPool := &sync.Pool{New: func() interface{} { return &Point{} }}
	pvs := []interface{}{}

	item := &mongo_serialization.Item{}
	destTag := &mongo_serialization.Tag{}
	collection := db.C(pointCollectionName)
	for batch := range batchChan {
		bulk := collection.Bulk()

		if cap(pvs) < len(*batch) {
			pvs = make([]interface{}, len(*batch))
		}
		pvs = pvs[:len(*batch)]

		for i, itemBuf := range *batch {
			// this ui could be improved on the library side:
			n := flatbuffers.GetUOffsetT(itemBuf)
			item.Init(itemBuf, n)
			x := pPool.Get().(*Point)

			x.MeasurementName = unsafeBytesToString(item.MeasurementNameBytes())
			x.FieldName = unsafeBytesToString(item.FieldNameBytes())
			x.Timestamp = item.TimestampNanos()

			tagLength := item.TagsLength()
			if cap(x.Tags) < tagLength {
				x.Tags = make([]Tag, 0, tagLength)
			}
			x.Tags = x.Tags[:tagLength]
			for i := 0; i < tagLength; i++ {
				*destTag = mongo_serialization.Tag{} // clear
				item.Tags(destTag, i)
				x.Tags[i].Key = unsafeBytesToString(destTag.KeyBytes())
				x.Tags[i].Val = unsafeBytesToString(destTag.ValBytes())
			}

			// this complexity is the result of trying to minimize
			// allocs while using an interface{} type for
			// (*Point).Value.
			switch item.ValueType() {
			case mongo_serialization.ValueTypeLong:
				x.longValue = item.LongValue()
				x.Value = &x.longValue
			case mongo_serialization.ValueTypeDouble:
				x.doubleValue = item.DoubleValue()
				x.Value = &x.doubleValue
			case mongo_serialization.ValueTypeString:
				x.stringValue = unsafeBytesToString(item.StringValueBytes())
				x.Value = &x.stringValue
			default:
				panic("logic error")
			}
			pvs[i] = x

		}
		bulk.Insert(pvs...)

		if doLoad {
			_, err := bulk.Run()
			if err != nil {
				log.Fatalf("Bulk err: %s\n", err.Error())
			}

		}

		// cleanup pvs
		for _, x := range pvs {
			p := x.(*Point)
			p.Timestamp = 0
			p.Value = nil
			p.longValue = 0
			p.doubleValue = 0
			p.Tags = p.Tags[:0]
			pPool.Put(p)
		}

		// cleanup item data
		for _, itemBuf := range *batch {
			bufPool.Put(itemBuf)
		}
		batch.ClearReferences()
		batchPool.Put(batch)
	}
	workersGroup.Done()
}

func mustCreateCollections(daemonUrl string) {
	session, err := mgo.Dial(daemonUrl)
	if err != nil {
		log.Fatal(err)
	}
	defer session.Close()

	// collection C: point data
	// from (*mgo.Collection).Create
	cmd := make(bson.D, 0, 4)
	cmd = append(cmd, bson.DocElem{"create", pointCollectionName})

	// wiredtiger settings
	cmd = append(cmd, bson.DocElem{
		"storageEngine", map[string]interface{}{
			"wiredTiger": map[string]interface{}{
				"configString": "block_compressor=snappy",
			},
		},
	})

	err = session.DB("benchmark_db").Run(cmd, nil)
	if err != nil {
		log.Fatal(err)
	}

	collection := session.DB("benchmark_db").C("point_data")
	index := mgo.Index{
		Key:        []string{"measurement", "tags", "field", "timestamp_ns"},
		Unique:     false, // Unique does not work on the entire array of tags!
		DropDups:   true,
		Background: false,
		Sparse:     false,
	}
	err = collection.EnsureIndex(index)
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>ensure writes are in journal<commit_after>// bulk_load_mongo loads a Mongo daemon with data from stdin.
//
// The caller is responsible for assuring that the database is empty before
// bulk load.
package main

import (
	"bufio"
	"encoding/binary"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"sync"
	"time"

	flatbuffers "github.com/google/flatbuffers/go"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"

	"github.com/influxdata/influxdb-comparisons/mongo_serialization"
	"github.com/influxdata/influxdb-comparisons/util/report"
	"strconv"
	"strings"
)

// Program option vars:
var (
	daemonUrl      string
	workers        int
	batchSize      int
	limit          int64
	doLoad         bool
	writeTimeout   time.Duration
	reportDatabase string
	reportHost     string
	reportUser     string
	reportPassword string
	reportTagsCSV  string
)

// Global vars
var (
	batchChan      chan *Batch
	inputDone      chan struct{}
	workersGroup   sync.WaitGroup
	reportTags     [][2]string
	reportHostname string
)

// Magic database constants
const (
	dbName              = "benchmark_db"
	pointCollectionName = "point_data"
)

// bufPool holds []byte instances to reduce heap churn.
var bufPool = &sync.Pool{
	New: func() interface{} {
		return make([]byte, 0, 1024)
	},
}

// Batch holds byte slices that will become mongo_serialization.Item instances.
type Batch [][]byte

func (b *Batch) ClearReferences() {
	*b = (*b)[:0]
}

// batchPool holds *Batch instances to reduce heap churn.
var batchPool = &sync.Pool{
	New: func() interface{} {
		return &Batch{}
	},
}

// Parse args:
func init() {
	flag.StringVar(&daemonUrl, "url", "localhost:27017", "Mongo URL.")

	flag.IntVar(&batchSize, "batch-size", 100, "Batch size (input items).")
	flag.IntVar(&workers, "workers", 1, "Number of parallel requests to make.")
	flag.Int64Var(&limit, "limit", -1, "Number of items to insert (default unlimited).")
	flag.DurationVar(&writeTimeout, "write-timeout", 10*time.Second, "Write timeout.")

	flag.BoolVar(&doLoad, "do-load", true, "Whether to write data. Set this flag to false to check input read speed.")

	flag.StringVar(&reportDatabase, "report-database", "database_benchmarks", "Database name where to store result metrics")
	flag.StringVar(&reportHost, "report-host", "", "Host to send result metrics")
	flag.StringVar(&reportUser, "report-user", "", "User for host to send result metrics")
	flag.StringVar(&reportPassword, "report-password", "", "User password for Host to send result metrics")
	flag.StringVar(&reportTagsCSV, "report-tags", "", "Comma separated k:v tags to send  alongside result metrics")

	flag.Parse()

	for i := 0; i < workers*batchSize; i++ {
		bufPool.Put(bufPool.New())
	}

	if reportHost != "" {
		fmt.Printf("results report destination: %v\n", reportHost)
		fmt.Printf("results report database: %v\n", reportDatabase)

		var err error
		reportHostname, err = os.Hostname()
		if err != nil {
			log.Fatalf("os.Hostname() error: %s", err.Error())
		}
		fmt.Printf("hostname for results report: %v\n", reportHostname)

		if reportTagsCSV != "" {
			pairs := strings.Split(reportTagsCSV, ",")
			for _, pair := range pairs {
				fields := strings.SplitN(pair, ":", 2)
				tagpair := [2]string{fields[0], fields[1]}
				reportTags = append(reportTags, tagpair)
			}
		}
		fmt.Printf("results report tags: %v\n", reportTags)
	}
}

func main() {
	if doLoad {
		mustCreateCollections(daemonUrl)
	}

	var session *mgo.Session

	if doLoad {
		var err error
		session, err = mgo.Dial(daemonUrl)
		if err != nil {
			log.Fatal(err)
		}

		session.SetMode(mgo.Eventual, false)
		session.SetSafe(&mgo.Safe{J: true})

		defer session.Close()
	}

	batchChan = make(chan *Batch, workers*10)
	inputDone = make(chan struct{})

	for i := 0; i < workers; i++ {
		workersGroup.Add(1)
		go processBatches(session)
	}

	start := time.Now()
	itemsRead, bytesRead := scan(session, batchSize)

	<-inputDone
	close(batchChan)
	workersGroup.Wait()
	end := time.Now()
	took := end.Sub(start)
	itemRate := float64(itemsRead) / float64(took.Seconds())
	bytesRate := float64(bytesRead) / float64(took.Seconds())

	fmt.Printf("loaded %d values in %fsec with %d workers (mean values rate %f values/sec, %.2fMB/sec from stdin)\n", itemsRead, took.Seconds(), workers, itemRate, bytesRate/(1<<20))

	if reportHost != "" {
		//append db specific tags to custom tags
		reportTags = append(reportTags, [2]string{"write_timeout", strconv.Itoa(int(writeTimeout))})

		reportParams := &report.LoadReportParams{
			ReportParams: report.ReportParams{
				DBType:             "MongoDB",
				ReportDatabaseName: reportDatabase,
				ReportHost:         reportHost,
				ReportUser:         reportUser,
				ReportPassword:     reportPassword,
				ReportTags:         reportTags,
				Hostname:           reportHostname,
				DestinationUrl:     daemonUrl,
				Workers:            workers,
				ItemLimit:          int(limit),
			},
			IsGzip:    false,
			BatchSize: batchSize,
		}
		err := report.ReportLoadResult(reportParams, itemsRead, itemRate, bytesRate, took)

		if err != nil {
			log.Fatal(err)
		}
	}
}

// scan reads length-delimited flatbuffers items from stdin.
func scan(session *mgo.Session, itemsPerBatch int) (int64, int64) {
	var n int
	var itemsRead, bytesRead int64
	r := bufio.NewReaderSize(os.Stdin, 32<<20)

	start := time.Now()
	batch := batchPool.Get().(*Batch)
	lenBuf := make([]byte, 8)

	for {
		if itemsRead == limit {
			break
		}
		// get the serialized item length (this is the framing format)
		_, err := r.Read(lenBuf)
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatal(err.Error())
		}

		// ensure correct len of receiving buffer
		l := int(binary.LittleEndian.Uint64(lenBuf))
		itemBuf := bufPool.Get().([]byte)
		if cap(itemBuf) < l {
			itemBuf = make([]byte, l)
		}
		itemBuf = itemBuf[:l]

		// read the bytes and init the flatbuffer object
		totRead := 0
		for totRead < l {
			m, err := r.Read(itemBuf[totRead:])
			// (EOF is also fatal)
			if err != nil {
				log.Fatal(err.Error())
			}
			totRead += m
		}
		if totRead != len(itemBuf) {
			panic(fmt.Sprintf("reader/writer logic error, %d != %d", n, len(itemBuf)))
		}

		*batch = append(*batch, itemBuf)

		itemsRead++
		n++

		if n >= batchSize {
			bytesRead += int64(len(itemBuf))
			batchChan <- batch
			n = 0
			batch = batchPool.Get().(*Batch)
		}

		_ = start
		//if itemsRead > 0 && itemsRead%100000 == 0 {
		//	_ = start
		//	//took := (time.Now().UnixNano() - start.UnixNano())
		//	//if took >= 1e9 {
		//	//	tookUs := float64(took) / 1e3
		//	//	tookSec := float64(took) / 1e9
		//	//	fmt.Fprintf(os.Stderr, "itemsRead: %d, rate: %.0f/sec, lag: %.2fus/op\n",
		//	//		itemsRead, float64(itemsRead)/tookSec, tookUs/float64(itemsRead))
		//	//}
		//}
	}

	// Closing inputDone signals to the application that we've read everything and can now shut down.
	close(inputDone)

	return itemsRead, bytesRead
}

// processBatches reads byte buffers from batchChan, interprets them and writes
// them to the target server. Note that mgo forcibly incurs serialization
// overhead (it always encodes to BSON).
func processBatches(session *mgo.Session) {
	db := session.DB(dbName)

	type Tag struct {
		Key string `bson:"key"`
		Val string `bson:"val"`
	}

	type Point struct {
		// Use `string` here even though they are really `[]byte`.
		// This is so the mongo data is human-readable.
		MeasurementName string      `bson:"measurement"`
		FieldName       string      `bson:"field"`
		Timestamp       int64       `bson:"timestamp_ns"`
		Tags            []Tag       `bson:"tags"`
		Value           interface{} `bson:"value"`

		// a private union-like section
		longValue   int64
		doubleValue float64
		stringValue string
	}
	pPool := &sync.Pool{New: func() interface{} { return &Point{} }}
	pvs := []interface{}{}

	item := &mongo_serialization.Item{}
	destTag := &mongo_serialization.Tag{}
	collection := db.C(pointCollectionName)
	for batch := range batchChan {
		bulk := collection.Bulk()

		if cap(pvs) < len(*batch) {
			pvs = make([]interface{}, len(*batch))
		}
		pvs = pvs[:len(*batch)]

		for i, itemBuf := range *batch {
			// this ui could be improved on the library side:
			n := flatbuffers.GetUOffsetT(itemBuf)
			item.Init(itemBuf, n)
			x := pPool.Get().(*Point)

			x.MeasurementName = unsafeBytesToString(item.MeasurementNameBytes())
			x.FieldName = unsafeBytesToString(item.FieldNameBytes())
			x.Timestamp = item.TimestampNanos()

			tagLength := item.TagsLength()
			if cap(x.Tags) < tagLength {
				x.Tags = make([]Tag, 0, tagLength)
			}
			x.Tags = x.Tags[:tagLength]
			for i := 0; i < tagLength; i++ {
				*destTag = mongo_serialization.Tag{} // clear
				item.Tags(destTag, i)
				x.Tags[i].Key = unsafeBytesToString(destTag.KeyBytes())
				x.Tags[i].Val = unsafeBytesToString(destTag.ValBytes())
			}

			// this complexity is the result of trying to minimize
			// allocs while using an interface{} type for
			// (*Point).Value.
			switch item.ValueType() {
			case mongo_serialization.ValueTypeLong:
				x.longValue = item.LongValue()
				x.Value = &x.longValue
			case mongo_serialization.ValueTypeDouble:
				x.doubleValue = item.DoubleValue()
				x.Value = &x.doubleValue
			case mongo_serialization.ValueTypeString:
				x.stringValue = unsafeBytesToString(item.StringValueBytes())
				x.Value = &x.stringValue
			default:
				panic("logic error")
			}
			pvs[i] = x

		}
		bulk.Insert(pvs...)

		if doLoad {
			_, err := bulk.Run()
			if err != nil {
				log.Fatalf("Bulk err: %s\n", err.Error())
			}

		}

		// cleanup pvs
		for _, x := range pvs {
			p := x.(*Point)
			p.Timestamp = 0
			p.Value = nil
			p.longValue = 0
			p.doubleValue = 0
			p.Tags = p.Tags[:0]
			pPool.Put(p)
		}

		// cleanup item data
		for _, itemBuf := range *batch {
			bufPool.Put(itemBuf)
		}
		batch.ClearReferences()
		batchPool.Put(batch)
	}
	workersGroup.Done()
}

func mustCreateCollections(daemonUrl string) {
	session, err := mgo.Dial(daemonUrl)
	if err != nil {
		log.Fatal(err)
	}
	defer session.Close()

	// collection C: point data
	// from (*mgo.Collection).Create
	cmd := make(bson.D, 0, 4)
	cmd = append(cmd, bson.DocElem{"create", pointCollectionName})

	// wiredtiger settings
	cmd = append(cmd, bson.DocElem{
		"storageEngine", map[string]interface{}{
			"wiredTiger": map[string]interface{}{
				"configString": "block_compressor=snappy",
			},
		},
	})

	err = session.DB("benchmark_db").Run(cmd, nil)
	if err != nil {
		log.Fatal(err)
	}

	collection := session.DB("benchmark_db").C("point_data")
	index := mgo.Index{
		Key:        []string{"measurement", "tags", "field", "timestamp_ns"},
		Unique:     false, // Unique does not work on the entire array of tags!
		DropDups:   true,
		Background: false,
		Sparse:     false,
	}
	err = collection.EnsureIndex(index)
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package command

import (
	"flag"
	"fmt"
	"sort"
	"strings"

	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/color"
	"github.com/funkygao/golib/gofmt"
	"github.com/ryanuber/columnize"
	gozk "github.com/samuel/go-zookeeper/zk"
)

type Consumers struct {
	Ui  cli.Ui
	Cmd string

	onlineOnly   bool
	ownerOnly    bool
	groupPattern string
	warnOnly     bool
	byHost       bool
	cleanup      bool
	confirmYes   bool
	topicPattern string
}

func (this *Consumers) Run(args []string) (exitCode int) {
	var (
		cluster string
		zone    string
	)
	cmdFlags := flag.NewFlagSet("consumers", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&zone, "z", ctx.ZkDefaultZone(), "")
	cmdFlags.StringVar(&cluster, "c", "", "")
	cmdFlags.StringVar(&this.groupPattern, "g", "", "")
	cmdFlags.BoolVar(&this.onlineOnly, "online", false, "")
	cmdFlags.BoolVar(&this.byHost, "byhost", false, "")
	cmdFlags.StringVar(&this.topicPattern, "t", "", "")
	cmdFlags.BoolVar(&this.warnOnly, "warn", false, "")
	cmdFlags.BoolVar(&this.ownerOnly, "own", false, "")
	cmdFlags.BoolVar(&this.cleanup, "cleanup", false, "")
	cmdFlags.BoolVar(&this.confirmYes, "yes", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if validateArgs(this, this.Ui).
		requireAdminRights("-cleanup").
		invalid(args) {
		return 2
	}

	if this.ownerOnly {
		this.onlineOnly = true
	}

	if zone == "" {
		forSortedZones(func(zkzone *zk.ZkZone) {
			switch {
			case this.cleanup:
				this.cleanupStaleConsumerGroups(zkzone, cluster)
			case this.byHost:
				this.printConsumersByHost(zkzone, cluster)
			default:
				this.printConsumersByGroupTable(zkzone, cluster)
			}
		})

		return
	}

	zkzone := zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
	switch {
	case this.cleanup:
		this.cleanupStaleConsumerGroups(zkzone, cluster)
	case this.byHost:
		this.printConsumersByHost(zkzone, cluster)
	default:
		this.printConsumersByGroupTable(zkzone, cluster)
	}

	return
}

func (this *Consumers) cleanupStaleConsumerGroups(zkzone *zk.ZkZone, clusterPattern string) {
	// what consumer groups are safe to delete?
	// 1. not online
	// 2. have no offsets
	this.Ui.Output(color.Blue(zkzone.Name()))

	zkzone.ForSortedClusters(func(zkcluster *zk.ZkCluster) {
		if !patternMatched(zkcluster.Name(), clusterPattern) {
			return
		}

		this.Ui.Output(strings.Repeat(" ", 4) + zkcluster.Name())
		consumerGroups := zkcluster.ConsumerGroups()
		for group, consumers := range consumerGroups {
			if len(consumers) > 0 {
				// this consumer group is online
				continue
			}

			if !patternMatched(group, this.groupPattern) {
				continue
			}

			if !strings.HasPrefix(group, "console-consumer-") {
				path := zkcluster.ConsumerGroupOffsetPath(group)
				_, _, err := zkzone.Conn().Children(path)
				if err == nil {
					this.Ui.Warn(fmt.Sprintf("%s not empty, unsafe to cleanup", path))
					continue
				}

				if err != gozk.ErrNoNode {
					// should never happen
					swallow(err)
				}
			}

			// have no offsets, safe to delete
			yes, err := this.Ui.Ask(fmt.Sprintf("confirm to remove cluster[%s] consumer group: %s? [y/N]",
				zkcluster.Name(), group))
			swallow(err)
			if this.confirmYes {
				if strings.ToLower(yes) == "n" {
					this.Ui.Info(fmt.Sprintf("%s skipped", group))
					continue
				}
			} else if strings.ToLower(yes) != "y" {
				this.Ui.Info(fmt.Sprintf("%s skipped", group))
				continue
			}

			// do delete this consumer group
			zkzone.DeleteRecursive(zkcluster.ConsumerGroupRoot(group))
			this.Ui.Info(fmt.Sprintf("%s deleted", group))
		}
	})
}

func (this *Consumers) printConsumersByHost(zkzone *zk.ZkZone, clusterPattern string) {
	outputs := make(map[string]map[string]map[string]int) // host: {cluster: {topic: count}}

	this.Ui.Output(color.Blue(zkzone.Name()))

	zkzone.ForSortedClusters(func(zkcluster *zk.ZkCluster) {
		if !patternMatched(zkcluster.Name(), clusterPattern) {
			return
		}

		consumerGroups := zkcluster.ConsumerGroups()
		for _, group := range consumerGroups {
			for _, c := range group {
				if _, present := outputs[c.Host()]; !present {
					outputs[c.Host()] = make(map[string]map[string]int)
				}

				if _, present := outputs[c.Host()][zkcluster.Name()]; !present {
					outputs[c.Host()][zkcluster.Name()] = make(map[string]int)
				}

				for topic, count := range c.Subscription {
					outputs[c.Host()][zkcluster.Name()][topic] += count
				}
			}
		}

	})

	sortedHosts := make([]string, 0, len(outputs))
	for host, _ := range outputs {
		sortedHosts = append(sortedHosts, host)
	}
	sort.Strings(sortedHosts)
	for _, host := range sortedHosts {
		tc := outputs[host]
		this.Ui.Output(fmt.Sprintf("%s %+v", color.Green("%22s", host), tc))
	}
}

func (this *Consumers) printConsumersByGroupTable(zkzone *zk.ZkZone, clusterPattern string) {
	lines := make([]string, 0)
	header := "Zone|Cluster|M|Host|ConsumerGroup|Topic/Partition|Offset|Uptime"
	lines = append(lines, header)

	zkzone.ForSortedClusters(func(zkcluster *zk.ZkCluster) {
		groupTopicsMap := make(map[string]map[string]struct{}) // group:sub topics

		if !patternMatched(zkcluster.Name(), clusterPattern) {
			return
		}

		consumerGroups := zkcluster.ConsumerGroups()
		sortedGroups := make([]string, 0, len(consumerGroups))
		for group, _ := range consumerGroups {
			if !patternMatched(group, this.groupPattern) {
				continue
			}

			sortedGroups = append(sortedGroups, group)
		}

		sort.Strings(sortedGroups)
		for _, group := range sortedGroups {
			consumers := consumerGroups[group]
			if this.onlineOnly && len(consumers) == 0 {
				continue
			}

			if len(consumers) > 0 {
				// sort by host
				sortedIds := make([]string, 0)
				consumersMap := make(map[string]*zk.ConsumerZnode)
				for _, c := range consumers {
					sortedIds = append(sortedIds, c.Id)
					consumersMap[c.Id] = c
				}
				sort.Strings(sortedIds)

				for _, consumerId := range sortedIds {
					c := consumersMap[consumerId]
					for topic, _ := range c.Subscription {
						if !patternMatched(topic, this.topicPattern) {
							continue
						}

						if len(groupTopicsMap[group]) == 0 {
							groupTopicsMap[group] = make(map[string]struct{}, 5)
						}
						groupTopicsMap[group][topic] = struct{}{}

						ownerByPartition := zkcluster.OwnersOfGroupByTopic(group, topic)

						partitionsWithOffset := make(map[string]struct{})
						for _, offset := range this.displayGroupOffsets(zkcluster, group, topic, false) {

							onlineSymbol := "◉"
							isOwner := false
							if ownerByPartition[offset.partitionId] == consumerId {
								onlineSymbol += "*" // owned by this consumer
								isOwner = true
							}
							if this.ownerOnly && !isOwner {
								continue
							}

							partitionsWithOffset[offset.partitionId] = struct{}{}

							lines = append(lines,
								fmt.Sprintf("%s|%s|%s|%s|%s|%s|%s|%s",
									zkzone.Name(), zkcluster.Name(),
									onlineSymbol,
									c.Host(),
									group+"@"+c.Id[len(c.Id)-12:],
									fmt.Sprintf("%s/%s", offset.topic, offset.partitionId),
									offset.offset,
									gofmt.PrettySince(c.Uptime())))
						}

						for partitionId, _ := range ownerByPartition {
							if _, present := partitionsWithOffset[partitionId]; !present {
								// this consumer is owner online, but has no offset
								onlineSymbol := "◉"
								isOwner := false
								if ownerByPartition[partitionId] == consumerId {
									onlineSymbol += "*"
									isOwner = true
								}
								if this.ownerOnly && !isOwner {
									continue
								}

								lines = append(lines,
									fmt.Sprintf("%s|%s|%s|%s|%s|%s|?|%s",
										zkzone.Name(), zkcluster.Name(),
										onlineSymbol,
										c.Host(),
										group+"@"+c.Id[len(c.Id)-12:],
										fmt.Sprintf("%s/%s", topic, partitionId),
										gofmt.PrettySince(c.Uptime())))
							}
						}
					}

				}
			} else {
				// offline
				for _, offset := range this.displayGroupOffsets(zkcluster, group, "", false) {
					if !patternMatched(offset.topic, this.topicPattern) {
						continue
					}

					lines = append(lines,
						fmt.Sprintf("%s|%s|%s|%s|%s|%s|%s|%s",
							zkzone.Name(), zkcluster.Name(),
							"◎",
							" ",
							group, fmt.Sprintf("%s/%s", offset.topic, offset.partitionId),
							offset.offset, " "))
				}
			}
		}

		for group, topics := range groupTopicsMap {
			if len(topics) > 1 {
				// the same consumer group is consuming more than 1 topics
				topicsLabel := make([]string, 0, len(topics))
				for t, _ := range topics {
					topicsLabel = append(topicsLabel, t)
				}
				this.Ui.Warn(fmt.Sprintf("%35s consuming: %+v", group, topicsLabel))
			}
		}
	})

	if !this.warnOnly {
		this.Ui.Output(columnize.SimpleFormat(lines))
	}

}

type consumerGroupOffset struct {
	topic, partitionId string
	offset             string // comma fmt
}

func (this *Consumers) displayGroupOffsets(zkcluster *zk.ZkCluster, group, topic string, echo bool) []consumerGroupOffset {
	offsetMap := zkcluster.ConsumerOffsetsOfGroup(group)
	sortedTopics := make([]string, 0, len(offsetMap))
	for t, _ := range offsetMap {
		sortedTopics = append(sortedTopics, t)
	}
	sort.Strings(sortedTopics)

	r := make([]consumerGroupOffset, 0)

	for _, t := range sortedTopics {
		if !patternMatched(t, this.topicPattern) || (topic != "" && t != topic) {
			continue
		}

		sortedPartitionIds := make([]string, 0, len(offsetMap[t]))
		for partitionId, _ := range offsetMap[t] {
			sortedPartitionIds = append(sortedPartitionIds, partitionId)
		}
		sort.Strings(sortedPartitionIds)

		for _, partitionId := range sortedPartitionIds {
			r = append(r, consumerGroupOffset{
				topic:       t,
				partitionId: partitionId,
				offset:      gofmt.Comma(offsetMap[t][partitionId]),
			})

			if echo {
				this.Ui.Output(fmt.Sprintf("\t\t%s/%s Offset:%s",
					t, partitionId, gofmt.Comma(offsetMap[t][partitionId])))
			}

		}
	}

	return r

}

func (*Consumers) Synopsis() string {
	return "Print high level consumer groups from Zookeeper"
}

func (this *Consumers) Help() string {
	help := fmt.Sprintf(`
Usage: %s consumers [options]

    %s

Options:

    -z zone

    -c cluster

    -g group name pattern

    -t topic pattern

    -warn
      Only show groups that consumes multiple topics.

    -online
      Only show online consumer groups.

    -own
      Only show consumer instances that owns partitions.

    -cleanup
      Cleanup the stale consumer groups after confirmation.

    -yes
      Work with -cleanup, input 'y' by default if confirm prompted.

    -byhost
      Display consumer groups by consumer hosts.

`, this.Cmd, this.Synopsis())
	return strings.TrimSpace(help)
}
<commit_msg>better prompt<commit_after>package command

import (
	"flag"
	"fmt"
	"sort"
	"strings"

	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/color"
	"github.com/funkygao/golib/gofmt"
	"github.com/ryanuber/columnize"
	gozk "github.com/samuel/go-zookeeper/zk"
)

type Consumers struct {
	Ui  cli.Ui
	Cmd string

	onlineOnly   bool
	ownerOnly    bool
	groupPattern string
	warnOnly     bool
	byHost       bool
	cleanup      bool
	confirmYes   bool
	topicPattern string
}

func (this *Consumers) Run(args []string) (exitCode int) {
	var (
		cluster string
		zone    string
	)
	cmdFlags := flag.NewFlagSet("consumers", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&zone, "z", ctx.ZkDefaultZone(), "")
	cmdFlags.StringVar(&cluster, "c", "", "")
	cmdFlags.StringVar(&this.groupPattern, "g", "", "")
	cmdFlags.BoolVar(&this.onlineOnly, "online", false, "")
	cmdFlags.BoolVar(&this.byHost, "byhost", false, "")
	cmdFlags.StringVar(&this.topicPattern, "t", "", "")
	cmdFlags.BoolVar(&this.warnOnly, "warn", false, "")
	cmdFlags.BoolVar(&this.ownerOnly, "own", false, "")
	cmdFlags.BoolVar(&this.cleanup, "cleanup", false, "")
	cmdFlags.BoolVar(&this.confirmYes, "yes", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if validateArgs(this, this.Ui).
		requireAdminRights("-cleanup").
		invalid(args) {
		return 2
	}

	if this.ownerOnly {
		this.onlineOnly = true
	}

	if zone == "" {
		forSortedZones(func(zkzone *zk.ZkZone) {
			switch {
			case this.cleanup:
				this.cleanupStaleConsumerGroups(zkzone, cluster)
			case this.byHost:
				this.printConsumersByHost(zkzone, cluster)
			default:
				this.printConsumersByGroupTable(zkzone, cluster)
			}
		})

		return
	}

	zkzone := zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
	switch {
	case this.cleanup:
		this.cleanupStaleConsumerGroups(zkzone, cluster)
	case this.byHost:
		this.printConsumersByHost(zkzone, cluster)
	default:
		this.printConsumersByGroupTable(zkzone, cluster)
	}

	return
}

func (this *Consumers) cleanupStaleConsumerGroups(zkzone *zk.ZkZone, clusterPattern string) {
	// what consumer groups are safe to delete?
	// 1. not online
	// 2. have no offsets
	this.Ui.Output(color.Blue(zkzone.Name()))

	zkzone.ForSortedClusters(func(zkcluster *zk.ZkCluster) {
		if !patternMatched(zkcluster.Name(), clusterPattern) {
			return
		}

		this.Ui.Output(strings.Repeat(" ", 4) + zkcluster.Name())
		consumerGroups := zkcluster.ConsumerGroups()
		for group, consumers := range consumerGroups {
			if len(consumers) > 0 {
				// this consumer group is online
				continue
			}

			if !patternMatched(group, this.groupPattern) {
				continue
			}

			if !strings.HasPrefix(group, "console-consumer-") {
				path := zkcluster.ConsumerGroupOffsetPath(group)
				_, _, err := zkzone.Conn().Children(path)
				if err == nil {
					this.Ui.Warn(fmt.Sprintf("%s not empty, unsafe to cleanup", path))
					continue
				}

				if err != gozk.ErrNoNode {
					// should never happen
					swallow(err)
				}
			}

			// have no offsets, safe to delete
			if this.confirmYes {
				yes, err := this.Ui.Ask(fmt.Sprintf("confirm to remove cluster[%s] consumer group: %s? [Y/n]",
					zkcluster.Name(), group))
				swallow(err)

				if strings.ToLower(yes) == "n" {
					this.Ui.Info(fmt.Sprintf("%s skipped", group))
					continue
				}
			} else if strings.ToLower(yes) != "y" {
				yes, err := this.Ui.Ask(fmt.Sprintf("confirm to remove cluster[%s] consumer group: %s? [y/N]",
					zkcluster.Name(), group))
				swallow(err)

				this.Ui.Info(fmt.Sprintf("%s skipped", group))
				continue
			}

			// do delete this consumer group
			zkzone.DeleteRecursive(zkcluster.ConsumerGroupRoot(group))
			this.Ui.Info(fmt.Sprintf("%s deleted", group))
		}
	})
}

func (this *Consumers) printConsumersByHost(zkzone *zk.ZkZone, clusterPattern string) {
	outputs := make(map[string]map[string]map[string]int) // host: {cluster: {topic: count}}

	this.Ui.Output(color.Blue(zkzone.Name()))

	zkzone.ForSortedClusters(func(zkcluster *zk.ZkCluster) {
		if !patternMatched(zkcluster.Name(), clusterPattern) {
			return
		}

		consumerGroups := zkcluster.ConsumerGroups()
		for _, group := range consumerGroups {
			for _, c := range group {
				if _, present := outputs[c.Host()]; !present {
					outputs[c.Host()] = make(map[string]map[string]int)
				}

				if _, present := outputs[c.Host()][zkcluster.Name()]; !present {
					outputs[c.Host()][zkcluster.Name()] = make(map[string]int)
				}

				for topic, count := range c.Subscription {
					outputs[c.Host()][zkcluster.Name()][topic] += count
				}
			}
		}

	})

	sortedHosts := make([]string, 0, len(outputs))
	for host, _ := range outputs {
		sortedHosts = append(sortedHosts, host)
	}
	sort.Strings(sortedHosts)
	for _, host := range sortedHosts {
		tc := outputs[host]
		this.Ui.Output(fmt.Sprintf("%s %+v", color.Green("%22s", host), tc))
	}
}

func (this *Consumers) printConsumersByGroupTable(zkzone *zk.ZkZone, clusterPattern string) {
	lines := make([]string, 0)
	header := "Zone|Cluster|M|Host|ConsumerGroup|Topic/Partition|Offset|Uptime"
	lines = append(lines, header)

	zkzone.ForSortedClusters(func(zkcluster *zk.ZkCluster) {
		groupTopicsMap := make(map[string]map[string]struct{}) // group:sub topics

		if !patternMatched(zkcluster.Name(), clusterPattern) {
			return
		}

		consumerGroups := zkcluster.ConsumerGroups()
		sortedGroups := make([]string, 0, len(consumerGroups))
		for group, _ := range consumerGroups {
			if !patternMatched(group, this.groupPattern) {
				continue
			}

			sortedGroups = append(sortedGroups, group)
		}

		sort.Strings(sortedGroups)
		for _, group := range sortedGroups {
			consumers := consumerGroups[group]
			if this.onlineOnly && len(consumers) == 0 {
				continue
			}

			if len(consumers) > 0 {
				// sort by host
				sortedIds := make([]string, 0)
				consumersMap := make(map[string]*zk.ConsumerZnode)
				for _, c := range consumers {
					sortedIds = append(sortedIds, c.Id)
					consumersMap[c.Id] = c
				}
				sort.Strings(sortedIds)

				for _, consumerId := range sortedIds {
					c := consumersMap[consumerId]
					for topic, _ := range c.Subscription {
						if !patternMatched(topic, this.topicPattern) {
							continue
						}

						if len(groupTopicsMap[group]) == 0 {
							groupTopicsMap[group] = make(map[string]struct{}, 5)
						}
						groupTopicsMap[group][topic] = struct{}{}

						ownerByPartition := zkcluster.OwnersOfGroupByTopic(group, topic)

						partitionsWithOffset := make(map[string]struct{})
						for _, offset := range this.displayGroupOffsets(zkcluster, group, topic, false) {

							onlineSymbol := "◉"
							isOwner := false
							if ownerByPartition[offset.partitionId] == consumerId {
								onlineSymbol += "*" // owned by this consumer
								isOwner = true
							}
							if this.ownerOnly && !isOwner {
								continue
							}

							partitionsWithOffset[offset.partitionId] = struct{}{}

							lines = append(lines,
								fmt.Sprintf("%s|%s|%s|%s|%s|%s|%s|%s",
									zkzone.Name(), zkcluster.Name(),
									onlineSymbol,
									c.Host(),
									group+"@"+c.Id[len(c.Id)-12:],
									fmt.Sprintf("%s/%s", offset.topic, offset.partitionId),
									offset.offset,
									gofmt.PrettySince(c.Uptime())))
						}

						for partitionId, _ := range ownerByPartition {
							if _, present := partitionsWithOffset[partitionId]; !present {
								// this consumer is owner online, but has no offset
								onlineSymbol := "◉"
								isOwner := false
								if ownerByPartition[partitionId] == consumerId {
									onlineSymbol += "*"
									isOwner = true
								}
								if this.ownerOnly && !isOwner {
									continue
								}

								lines = append(lines,
									fmt.Sprintf("%s|%s|%s|%s|%s|%s|?|%s",
										zkzone.Name(), zkcluster.Name(),
										onlineSymbol,
										c.Host(),
										group+"@"+c.Id[len(c.Id)-12:],
										fmt.Sprintf("%s/%s", topic, partitionId),
										gofmt.PrettySince(c.Uptime())))
							}
						}
					}

				}
			} else {
				// offline
				for _, offset := range this.displayGroupOffsets(zkcluster, group, "", false) {
					if !patternMatched(offset.topic, this.topicPattern) {
						continue
					}

					lines = append(lines,
						fmt.Sprintf("%s|%s|%s|%s|%s|%s|%s|%s",
							zkzone.Name(), zkcluster.Name(),
							"◎",
							" ",
							group, fmt.Sprintf("%s/%s", offset.topic, offset.partitionId),
							offset.offset, " "))
				}
			}
		}

		for group, topics := range groupTopicsMap {
			if len(topics) > 1 {
				// the same consumer group is consuming more than 1 topics
				topicsLabel := make([]string, 0, len(topics))
				for t, _ := range topics {
					topicsLabel = append(topicsLabel, t)
				}
				this.Ui.Warn(fmt.Sprintf("%35s consuming: %+v", group, topicsLabel))
			}
		}
	})

	if !this.warnOnly {
		this.Ui.Output(columnize.SimpleFormat(lines))
	}

}

type consumerGroupOffset struct {
	topic, partitionId string
	offset             string // comma fmt
}

func (this *Consumers) displayGroupOffsets(zkcluster *zk.ZkCluster, group, topic string, echo bool) []consumerGroupOffset {
	offsetMap := zkcluster.ConsumerOffsetsOfGroup(group)
	sortedTopics := make([]string, 0, len(offsetMap))
	for t, _ := range offsetMap {
		sortedTopics = append(sortedTopics, t)
	}
	sort.Strings(sortedTopics)

	r := make([]consumerGroupOffset, 0)

	for _, t := range sortedTopics {
		if !patternMatched(t, this.topicPattern) || (topic != "" && t != topic) {
			continue
		}

		sortedPartitionIds := make([]string, 0, len(offsetMap[t]))
		for partitionId, _ := range offsetMap[t] {
			sortedPartitionIds = append(sortedPartitionIds, partitionId)
		}
		sort.Strings(sortedPartitionIds)

		for _, partitionId := range sortedPartitionIds {
			r = append(r, consumerGroupOffset{
				topic:       t,
				partitionId: partitionId,
				offset:      gofmt.Comma(offsetMap[t][partitionId]),
			})

			if echo {
				this.Ui.Output(fmt.Sprintf("\t\t%s/%s Offset:%s",
					t, partitionId, gofmt.Comma(offsetMap[t][partitionId])))
			}

		}
	}

	return r

}

func (*Consumers) Synopsis() string {
	return "Print high level consumer groups from Zookeeper"
}

func (this *Consumers) Help() string {
	help := fmt.Sprintf(`
Usage: %s consumers [options]

    %s

Options:

    -z zone

    -c cluster

    -g group name pattern

    -t topic pattern

    -warn
      Only show groups that consumes multiple topics.

    -online
      Only show online consumer groups.

    -own
      Only show consumer instances that owns partitions.

    -cleanup
      Cleanup the stale consumer groups after confirmation.

    -yes
      Work with -cleanup, input 'y' by default if confirm prompted.

    -byhost
      Display consumer groups by consumer hosts.

`, this.Cmd, this.Synopsis())
	return strings.TrimSpace(help)
}
<|endoftext|>