<commit_before>package clt

import (
	"bufio"
	"fmt"
	"io"

	"os"
	"strings"

	"golang.org/x/crypto/ssh/terminal"
)

type getOpt int

const (
	noColon getOpt = iota
)

// InteractiveSession creates a system for collecting user input
// in response to questions and choices
type InteractiveSession struct {
	Prompt  string
	Default string
	ValHint string

	response string
	input    *bufio.Reader
	output   io.Writer
}

// NewInteractiveSession returns a new InteractiveSession outputting to Stdout
// and reading from Stdin by default, but other inputs and outputs may be specified
// with SessionOptions
func NewInteractiveSession(opts ...SessionOption) *InteractiveSession {
	i := &InteractiveSession{
		input:  bufio.NewReader(os.Stdin),
		output: os.Stdout,
	}
	for _, opt := range opts {
		opt(i)
	}
	return i
}

// SessionOption optionally configures aspects of the interactive session
type SessionOption func(i *InteractiveSession)

// WithInput uses an input other than os.Stdin
func WithInput(r io.Reader) SessionOption {
	return func(i *InteractiveSession) {
		i.input = bufio.NewReader(r)
	}
}

// WithOutput uses an output other than os.Stdout
func WithOutput(w io.Writer) SessionOption {
	return func(i *InteractiveSession) {
		i.output = w
	}
}

// Reset allows reuse of the same interactive session by reseting its state and keeping
// its current input and output
func (i *InteractiveSession) Reset() {
	i.Prompt = ""
	i.Default = ""
	i.ValHint = ""
	i.response = ""
}

// Say is a short form of fmt.Fprintf but allows you to chain additional terminators to
// the interactive session to collect user input
func (i *InteractiveSession) Say(format string, args ...interface{}) *InteractiveSession {
	fmt.Fprintf(i.output, fmt.Sprintf("\n%s\n", format), args...)
	return i
}

// Pause is a terminator that will render long-form text added via the another method
// that returns *InteractiveSession and will wait for the user to press enter to continue.
// It is useful for long-form content or paging.
func (i *InteractiveSession) Pause() {
	i.Prompt = "\nPress [Enter] to continue."
	i.get(noColon)
}

// PauseWithPrompt is a terminator that will render long-form text added via the another method
// that returns *InteractiveSession and will wait for the user to press enter to continue.
// This will use the custom prompt specified by format and args.
func (i *InteractiveSession) PauseWithPrompt(format string, args ...interface{}) {
	i.Prompt = fmt.Sprintf(format, args...)
	i.get(noColon)
}

func (i *InteractiveSession) get(opts ...getOpt) (err error) {
	contains := func(wanted getOpt) bool {
		for _, opt := range opts {
			if opt == wanted {
				return true
			}
		}
		return false
	}

	if i.output == nil {
		i.output = bufio.NewWriter(os.Stdout)
	}
	if i.input == nil {
		i.input = bufio.NewReader(os.Stdin)
	}

	switch {
	case len(i.Default) > 0:
		fmt.Fprintf(i.output, "%s  [%s]: ", i.Prompt, i.Default)
	case len(i.ValHint) > 0:
		fmt.Fprintf(i.output, "%s (%s): ", i.Prompt, i.ValHint)
	case contains(noColon):
		fmt.Fprintf(i.output, "%s", i.Prompt)
	case len(i.Prompt) > 0:
		fmt.Fprintf(i.output, "%s: ", i.Prompt)
	default:
	}

	i.response, err = i.input.ReadString('\n')
	if err != nil {
		return err
	}
	i.response = strings.TrimRight(i.response, " \n")
	if len(i.Default) > 0 && len(i.response) == 0 {
		switch i.Default {
		case "y/N":
			i.response = "n"
		case "Y/n":
			i.response = "y"
		default:
			i.response = i.Default
		}
	}

	return nil
}

// Warn adds an informational warning message to the user in format
// Warning: <user defined string>
func (i *InteractiveSession) Warn(format string, args ...interface{}) *InteractiveSession {
	fmt.Fprintf(i.output, "\n%s: %s\n", Styled(Yellow).ApplyTo("Warning"), fmt.Sprintf(format, args...))
	return i
}

// Error is a terminator that gives an informational error message to the user in format
// Error: <user defined string>.  Exits the program returning status code 1
func (i *InteractiveSession) Error(format string, args ...interface{}) {
	fmt.Fprintf(i.output, "\n\n%s: %s\n", Styled(Red).ApplyTo("Error:"), fmt.Sprintf(format, args...))
	os.Exit(1)
}

// Ask is a terminator for an interactive session that results in returning the user's
// input.  Validators can optionally be applied to ensure that acceptable input is returned
// or the question will be asked again.
func (i *InteractiveSession) Ask(prompt string, validators ...ValidationFunc) string {
	return i.ask(prompt, "", "", validators...)
}

// AskWithDefault is like ask, but sets a default choice that the user can select by pressing enter.
func (i *InteractiveSession) AskWithDefault(prompt string, defaultChoice string, validators ...ValidationFunc) string {
	return i.ask(prompt, defaultChoice, "", validators...)
}

// AskWithHint is like ask, but gives a hint about the proper format of the response.  This is useful
// combined with a validation function to get input in the right format.
func (i *InteractiveSession) AskWithHint(prompt string, hint string, validators ...ValidationFunc) string {
	return i.ask(prompt, "", hint, validators...)
}

func (i *InteractiveSession) ask(prompt string, def string, hint string, validators ...ValidationFunc) string {
	i.Prompt = prompt
	i.Default = def
	i.ValHint = hint
	i.get()
	for _, validator := range validators {
		if ok, err := validator(i.response); !ok {
			i.Say("\nError: %s\n\n", err)
			i.ask(prompt, def, hint, validators...)
		}
	}
	return i.response
}

// AskPassword is a terminator that asks for a password and does not echo input
// to the terminal.
func (i *InteractiveSession) AskPassword(validators ...ValidationFunc) string {
	return askPassword(i, "Password: ", validators...)
}

// AskPasswordPrompt is a terminator that asks for a password with a custom prompt
func (i *InteractiveSession) AskPasswordPrompt(prompt string, validators ...ValidationFunc) string {
	return askPassword(i, prompt, validators...)
}

func askPassword(i *InteractiveSession, prompt string, validators ...ValidationFunc) string {
	fmt.Fprintf(i.output, "Password: ")
	pw, err := terminal.ReadPassword(0)
	if err != nil {
		i.Error("\n%s\n", err)
	}

	pwS := strings.TrimSpace(string(pw))
	for _, validator := range validators {
		if ok, err := validator(pwS); !ok {
			i.Say("\nError: %s\n\n", err)
			i.AskPassword(validators...)
		}
	}
	return pwS
}

// AskYesNo asks the user a yes or no question with a default value.  Defaults of `y` or `yes` will
// set the default to yes.  Anything else will default to no.  You can use IsYes or IsNo to act on the response
// without worrying about what version of y, Y, YES, yes, etc. that the user entered.
func (i *InteractiveSession) AskYesNo(prompt string, defaultChoice string) string {
	switch def := strings.ToLower(defaultChoice); def {
	case "y", "yes":
		i.Default = "Y/n"
	default:
		i.Default = "y/N"
	}
	return i.ask(prompt, i.Default, "", ValidateYesNo())
}

// AskFromTable creates a table to select choices from.  It has a built-in validation function that will
// ensure that only the options listed are valid choices.
func (i *InteractiveSession) AskFromTable(prompt string, choices map[string]string, def string) string {
	t := NewTable(2).
		ColumnHeaders("Option", "")
	var allKeys []string
	for key, choice := range choices {
		t.AddRow(key, choice)
		allKeys = append(allKeys, key)
	}
	tAsString := t.AsString()

	i.Prompt = fmt.Sprintf("\n%s%s\nChoice", prompt, tAsString)
	i.Default = def
	i.get()
	if ok, err := AllowedOptions(allKeys)(i.response); !ok {
		i.Say("\nError: %s\n\n", err)
		i.AskFromTable(prompt, choices, def)
	}

	return strings.TrimSpace(i.response)
}
<commit_msg>stablizes sort order for asking from a table of choices<commit_after>package clt

import (
	"bufio"
	"fmt"
	"io"
	"sort"

	"os"
	"strings"

	"golang.org/x/crypto/ssh/terminal"
)

type getOpt int

const (
	noColon getOpt = iota
)

// InteractiveSession creates a system for collecting user input
// in response to questions and choices
type InteractiveSession struct {
	Prompt  string
	Default string
	ValHint string

	response string
	input    *bufio.Reader
	output   io.Writer
}

// NewInteractiveSession returns a new InteractiveSession outputting to Stdout
// and reading from Stdin by default, but other inputs and outputs may be specified
// with SessionOptions
func NewInteractiveSession(opts ...SessionOption) *InteractiveSession {
	i := &InteractiveSession{
		input:  bufio.NewReader(os.Stdin),
		output: os.Stdout,
	}
	for _, opt := range opts {
		opt(i)
	}
	return i
}

// SessionOption optionally configures aspects of the interactive session
type SessionOption func(i *InteractiveSession)

// WithInput uses an input other than os.Stdin
func WithInput(r io.Reader) SessionOption {
	return func(i *InteractiveSession) {
		i.input = bufio.NewReader(r)
	}
}

// WithOutput uses an output other than os.Stdout
func WithOutput(w io.Writer) SessionOption {
	return func(i *InteractiveSession) {
		i.output = w
	}
}

// Reset allows reuse of the same interactive session by reseting its state and keeping
// its current input and output
func (i *InteractiveSession) Reset() {
	i.Prompt = ""
	i.Default = ""
	i.ValHint = ""
	i.response = ""
}

// Say is a short form of fmt.Fprintf but allows you to chain additional terminators to
// the interactive session to collect user input
func (i *InteractiveSession) Say(format string, args ...interface{}) *InteractiveSession {
	fmt.Fprintf(i.output, fmt.Sprintf("\n%s\n", format), args...)
	return i
}

// Pause is a terminator that will render long-form text added via the another method
// that returns *InteractiveSession and will wait for the user to press enter to continue.
// It is useful for long-form content or paging.
func (i *InteractiveSession) Pause() {
	i.Prompt = "\nPress [Enter] to continue."
	i.get(noColon)
}

// PauseWithPrompt is a terminator that will render long-form text added via the another method
// that returns *InteractiveSession and will wait for the user to press enter to continue.
// This will use the custom prompt specified by format and args.
func (i *InteractiveSession) PauseWithPrompt(format string, args ...interface{}) {
	i.Prompt = fmt.Sprintf(format, args...)
	i.get(noColon)
}

func (i *InteractiveSession) get(opts ...getOpt) (err error) {
	contains := func(wanted getOpt) bool {
		for _, opt := range opts {
			if opt == wanted {
				return true
			}
		}
		return false
	}

	if i.output == nil {
		i.output = bufio.NewWriter(os.Stdout)
	}
	if i.input == nil {
		i.input = bufio.NewReader(os.Stdin)
	}

	switch {
	case len(i.Default) > 0:
		fmt.Fprintf(i.output, "%s  [%s]: ", i.Prompt, i.Default)
	case len(i.ValHint) > 0:
		fmt.Fprintf(i.output, "%s (%s): ", i.Prompt, i.ValHint)
	case contains(noColon):
		fmt.Fprintf(i.output, "%s", i.Prompt)
	case len(i.Prompt) > 0:
		fmt.Fprintf(i.output, "%s: ", i.Prompt)
	default:
	}

	i.response, err = i.input.ReadString('\n')
	if err != nil {
		return err
	}
	i.response = strings.TrimRight(i.response, " \n")
	if len(i.Default) > 0 && len(i.response) == 0 {
		switch i.Default {
		case "y/N":
			i.response = "n"
		case "Y/n":
			i.response = "y"
		default:
			i.response = i.Default
		}
	}

	return nil
}

// Warn adds an informational warning message to the user in format
// Warning: <user defined string>
func (i *InteractiveSession) Warn(format string, args ...interface{}) *InteractiveSession {
	fmt.Fprintf(i.output, "\n%s: %s\n", Styled(Yellow).ApplyTo("Warning"), fmt.Sprintf(format, args...))
	return i
}

// Error is a terminator that gives an informational error message to the user in format
// Error: <user defined string>.  Exits the program returning status code 1
func (i *InteractiveSession) Error(format string, args ...interface{}) {
	fmt.Fprintf(i.output, "\n\n%s: %s\n", Styled(Red).ApplyTo("Error:"), fmt.Sprintf(format, args...))
	os.Exit(1)
}

// Ask is a terminator for an interactive session that results in returning the user's
// input.  Validators can optionally be applied to ensure that acceptable input is returned
// or the question will be asked again.
func (i *InteractiveSession) Ask(prompt string, validators ...ValidationFunc) string {
	return i.ask(prompt, "", "", validators...)
}

// AskWithDefault is like ask, but sets a default choice that the user can select by pressing enter.
func (i *InteractiveSession) AskWithDefault(prompt string, defaultChoice string, validators ...ValidationFunc) string {
	return i.ask(prompt, defaultChoice, "", validators...)
}

// AskWithHint is like ask, but gives a hint about the proper format of the response.  This is useful
// combined with a validation function to get input in the right format.
func (i *InteractiveSession) AskWithHint(prompt string, hint string, validators ...ValidationFunc) string {
	return i.ask(prompt, "", hint, validators...)
}

func (i *InteractiveSession) ask(prompt string, def string, hint string, validators ...ValidationFunc) string {
	i.Prompt = prompt
	i.Default = def
	i.ValHint = hint
	i.get()
	for _, validator := range validators {
		if ok, err := validator(i.response); !ok {
			i.Say("\nError: %s\n\n", err)
			i.ask(prompt, def, hint, validators...)
		}
	}
	return i.response
}

// AskPassword is a terminator that asks for a password and does not echo input
// to the terminal.
func (i *InteractiveSession) AskPassword(validators ...ValidationFunc) string {
	return askPassword(i, "Password: ", validators...)
}

// AskPasswordPrompt is a terminator that asks for a password with a custom prompt
func (i *InteractiveSession) AskPasswordPrompt(prompt string, validators ...ValidationFunc) string {
	return askPassword(i, prompt, validators...)
}

func askPassword(i *InteractiveSession, prompt string, validators ...ValidationFunc) string {
	fmt.Fprintf(i.output, "Password: ")
	pw, err := terminal.ReadPassword(0)
	if err != nil {
		i.Error("\n%s\n", err)
	}

	pwS := strings.TrimSpace(string(pw))
	for _, validator := range validators {
		if ok, err := validator(pwS); !ok {
			i.Say("\nError: %s\n\n", err)
			i.AskPassword(validators...)
		}
	}
	return pwS
}

// AskYesNo asks the user a yes or no question with a default value.  Defaults of `y` or `yes` will
// set the default to yes.  Anything else will default to no.  You can use IsYes or IsNo to act on the response
// without worrying about what version of y, Y, YES, yes, etc. that the user entered.
func (i *InteractiveSession) AskYesNo(prompt string, defaultChoice string) string {
	switch def := strings.ToLower(defaultChoice); def {
	case "y", "yes":
		i.Default = "Y/n"
	default:
		i.Default = "y/N"
	}
	return i.ask(prompt, i.Default, "", ValidateYesNo())
}

// AskFromTable creates a table to select choices from.  It has a built-in validation function that will
// ensure that only the options listed are valid choices.
func (i *InteractiveSession) AskFromTable(prompt string, choices map[string]string, def string) string {
	t := NewTable(2).
		ColumnHeaders("Option", "")
	var allKeys []string
	for key := range choices {
		allKeys = append(allKeys, key)
	}
	sort.Strings(allKeys)

	for _, key := range allKeys {
		t.AddRow(key, choices[key])
	}
	tAsString := t.AsString()

	i.Prompt = fmt.Sprintf("\n%s%s\nChoice", prompt, tAsString)
	i.Default = def
	i.get()
	if ok, err := AllowedOptions(allKeys)(i.response); !ok {
		i.Say("\nError: %s\n\n", err)
		i.AskFromTable(prompt, choices, def)
	}

	return strings.TrimSpace(i.response)
}
<|endoftext|>


<commit_before>// Copyright 2015 Netflix, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package common contains all protocol-agnostic types and interfaces, as well as some variables
// used by multiple other packages in the application. Metrics are useful in many places for
// tracking bytes written and read. Errors exist here to be the neutral representation of either
// protocol errors or other application errors that happen all over. Other interfaces and types used
// by multiple protocols are here so any protocol can be used on the front end as well as the back
// with the same intermediate representation.
package common

import (
	"errors"

	"github.com/netflix/rend/metrics"
)

const VersionString = "Rend 0.1"

// Common metrics used across packages
var (
	MetricBytesReadRemote     = metrics.AddCounter("bytes_read_remote")
	MetricBytesReadLocal      = metrics.AddCounter("bytes_read_local")
	MetricBytesReadLocalL1    = metrics.AddCounter("bytes_read_local_l1")
	MetricBytesReadLocalL2    = metrics.AddCounter("bytes_read_local_l2")
	MetricBytesWrittenRemote  = metrics.AddCounter("bytes_written_remote")
	MetricBytesWrittenLocal   = metrics.AddCounter("bytes_written_local")
	MetricBytesWrittenLocalL1 = metrics.AddCounter("bytes_written_local_l1")
	MetricBytesWrittenLocalL2 = metrics.AddCounter("bytes_written_local_l2")

	// Errors used across the application
	ErrBadRequest = errors.New("CLIENT_ERROR bad request")
	ErrBadLength  = errors.New("CLIENT_ERROR length is not a valid integer")
	ErrBadFlags   = errors.New("CLIENT_ERROR flags is not a valid integer")
	ErrBadExptime = errors.New("CLIENT_ERROR exptime is not a valid integer")

	ErrNoError        = errors.New("Success")
	ErrKeyNotFound    = errors.New("ERROR Key not found")
	ErrKeyExists      = errors.New("ERROR Key already exists")
	ErrValueTooBig    = errors.New("ERROR Value too big")
	ErrInvalidArgs    = errors.New("ERROR Invalid arguments")
	ErrItemNotStored  = errors.New("ERROR Item not stored (CAS didn't match)")
	ErrBadIncDecValue = errors.New("ERROR Bad increment/decrement value")
	ErrAuth           = errors.New("ERROR Authentication error")
	ErrUnknownCmd     = errors.New("ERROR Unknown command")
	ErrNoMem          = errors.New("ERROR Out of memory")
	ErrNotSupported   = errors.New("ERROR Not supported")
	ErrInternal       = errors.New("ERROR Internal error")
	ErrBusy           = errors.New("ERROR Busy")
	ErrTempFailure    = errors.New("ERROR Temporary error")
)

// IsAppError differentiates between protocol-defined errors that are relatively benign and other
// fatal errors like an IO error because of some socket problem or network issue.
// Make sure to keep this list in sync with the one above. It should contain all Err* that could
// come back from memcached itself
func IsAppError(err error) bool {
	return err == ErrKeyNotFound ||
		err == ErrKeyExists ||
		err == ErrValueTooBig ||
		err == ErrInvalidArgs ||
		err == ErrItemNotStored ||
		err == ErrBadIncDecValue ||
		err == ErrAuth ||
		err == ErrUnknownCmd ||
		err == ErrNoMem ||
		err == ErrNotSupported ||
		err == ErrInternal ||
		err == ErrBusy ||
		err == ErrTempFailure
}

// RequestType is the protocol-agnostic identifier for the command
type RequestType int

const (
	// RequestUnknown means the parser doesn't know what the request represents. Valid protocol
	// parsing but invalid values.
	RequestUnknown = iota

	// RequestGet represents both a single get and a multi-get, which take differen forms in
	// different protocols. This means it can also be the accumulation of many GETQ commands.
	RequestGet

	// RequestGat is a get-and-touch operation that retrieves the information while updating the TTL
	RequestGat

	// RequestGetE is a custom get which returns the TTL remaining with the data
	RequestGetE

	// RequestSet is to insert a new piece of data unconditionally. What that means is different
	// depending on L1 / L2 handling.
	RequestSet

	// RequestAdd will perform the same operations as set, but only if the key does not exist
	RequestAdd

	// RequestReplace will perform the same operations as set, but only if the key already exists
	RequestReplace

	// RequestDelete deletes a piece of data from all levels of cache
	RequestDelete

	// RequestTouch updates the TTL for the item specified to a new TTL
	RequestTouch

	// RequestNoop does nothing
	RequestNoop

	// RequestQuit closes the connection
	RequestQuit

	// RequestVersion replies with a string designating the current software version
	RequestVersion
)

// RequestParser represents an interface to parse incoming requests. Each protocol provides its own
// implementation. The return value is an interface{}, but not all hope is lost. The return result
// is guaranteed by implementations to be castable to the type that matches the RequestType returned.
type RequestParser interface {
	Parse() (Request, RequestType, error)
}

// Responder is the interface for a protocol to respond to different commands. It responds in
// whatever way is appropriate, including doing nothing or panic()-ing for unsupported interactions.
// Unsupported interactions are OK to panic() on because they should never be returned from the
// corresponding RequestParser.
type Responder interface {
	Set(opaque uint32, quiet bool) error
	Add(opaque uint32, quiet bool) error
	Replace(opaque uint32, quiet bool) error
	Get(response GetResponse) error
	GetEnd(opaque uint32, noopEnd bool) error
	GetE(response GetEResponse) error
	GAT(response GetResponse) error
	Delete(opaque uint32) error
	Touch(opaque uint32) error
	Noop(opaque uint32) error
	Quit(opaque uint32, quiet bool) error
	Version(opaque uint32) error
	Error(opaque uint32, reqType RequestType, err error) error
}

type Request interface {
	Opq() uint32
}

// SetRequest corresponds to common.RequestSet. It contains all the information required to fulfill
// a set request.
type SetRequest struct {
	Key     []byte
	Data    []byte
	Flags   uint32
	Exptime uint32
	Opaque  uint32
	Quiet   bool
}

func (r SetRequest) Opq() uint32 {
	return r.Opaque
}

// GetRequest corresponds to common.RequestGet. It contains all the information required to fulfill
// a get requestGets are batch by default, so single gets and batch gets are both represented by the
// same type.
type GetRequest struct {
	Keys       [][]byte
	Opaques    []uint32
	Quiet      []bool
	NoopOpaque uint32
	NoopEnd    bool
}

func (r GetRequest) Opq() uint32 {
	// TODO: better implementation?
	// It's nonsensical but the best way to react in this case since it's a bad situation already.
	// Typically if this method was needed we're already in a fatal error sitation.
	return 0
}

// DeleteRequest corresponds to common.RequestDelete. It contains all the information required to
// fulfill a delete request.
type DeleteRequest struct {
	Key    []byte
	Opaque uint32
}

func (r DeleteRequest) Opq() uint32 {
	return r.Opaque
}

// TouchRequest corresponds to common.RequestTouch. It contains all the information required to
// fulfill a touch request.
type TouchRequest struct {
	Key     []byte
	Exptime uint32
	Opaque  uint32
}

func (r TouchRequest) Opq() uint32 {
	return r.Opaque
}

// GATRequest corresponds to common.RequestGat. It contains all the information required to fulfill
// a get-and-touch request.
type GATRequest struct {
	Key     []byte
	Exptime uint32
	Opaque  uint32
}

func (r GATRequest) Opq() uint32 {
	return r.Opaque
}

// QuitRequest corresponds to common.RequestQuit. It contains all the information required to
// fulfill a quit request.
type QuitRequest struct {
	Opaque uint32
	Quiet  bool
}

func (r QuitRequest) Opq() uint32 {
	return r.Opaque
}

// NoopRequest corresponds to common.RequestNoop. It contains all the information required to
// fulfill a version request.
type NoopRequest struct {
	Opaque uint32
}

func (r NoopRequest) Opq() uint32 {
	return r.Opaque
}

// VersionRequest corresponds to common.RequestVersion. It contains all the information required to
// fulfill a version request.
type VersionRequest struct {
	Opaque uint32
}

func (r VersionRequest) Opq() uint32 {
	return r.Opaque
}

// GetResponse is used in both RequestGet and RequestGat handling. Both respond in the same manner
// but with different opcodes. It is binary-protocol specific, but is still a part of the interface
// of responder to make the handling code more protocol-agnostic.
type GetResponse struct {
	Key    []byte
	Data   []byte
	Opaque uint32
	Flags  uint32
	Miss   bool
	Quiet  bool
}

// GetEResponse is used in the GetE protocol extension
type GetEResponse struct {
	Key     []byte
	Data    []byte
	Opaque  uint32
	Flags   uint32
	Exptime uint32
	Miss    bool
	Quiet   bool
}
<commit_msg>Adding append and prepend in the common package<commit_after>// Copyright 2015 Netflix, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package common contains all protocol-agnostic types and interfaces, as well as some variables
// used by multiple other packages in the application. Metrics are useful in many places for
// tracking bytes written and read. Errors exist here to be the neutral representation of either
// protocol errors or other application errors that happen all over. Other interfaces and types used
// by multiple protocols are here so any protocol can be used on the front end as well as the back
// with the same intermediate representation.
package common

import (
	"errors"

	"github.com/netflix/rend/metrics"
)

const VersionString = "Rend 0.1"

// Common metrics used across packages
var (
	MetricBytesReadRemote     = metrics.AddCounter("bytes_read_remote")
	MetricBytesReadLocal      = metrics.AddCounter("bytes_read_local")
	MetricBytesReadLocalL1    = metrics.AddCounter("bytes_read_local_l1")
	MetricBytesReadLocalL2    = metrics.AddCounter("bytes_read_local_l2")
	MetricBytesWrittenRemote  = metrics.AddCounter("bytes_written_remote")
	MetricBytesWrittenLocal   = metrics.AddCounter("bytes_written_local")
	MetricBytesWrittenLocalL1 = metrics.AddCounter("bytes_written_local_l1")
	MetricBytesWrittenLocalL2 = metrics.AddCounter("bytes_written_local_l2")

	// Errors used across the application
	ErrBadRequest = errors.New("CLIENT_ERROR bad request")
	ErrBadLength  = errors.New("CLIENT_ERROR length is not a valid integer")
	ErrBadFlags   = errors.New("CLIENT_ERROR flags is not a valid integer")
	ErrBadExptime = errors.New("CLIENT_ERROR exptime is not a valid integer")

	ErrNoError        = errors.New("Success")
	ErrKeyNotFound    = errors.New("ERROR Key not found")
	ErrKeyExists      = errors.New("ERROR Key already exists")
	ErrValueTooBig    = errors.New("ERROR Value too big")
	ErrInvalidArgs    = errors.New("ERROR Invalid arguments")
	ErrItemNotStored  = errors.New("ERROR Item not stored (CAS didn't match)")
	ErrBadIncDecValue = errors.New("ERROR Bad increment/decrement value")
	ErrAuth           = errors.New("ERROR Authentication error")
	ErrUnknownCmd     = errors.New("ERROR Unknown command")
	ErrNoMem          = errors.New("ERROR Out of memory")
	ErrNotSupported   = errors.New("ERROR Not supported")
	ErrInternal       = errors.New("ERROR Internal error")
	ErrBusy           = errors.New("ERROR Busy")
	ErrTempFailure    = errors.New("ERROR Temporary error")
)

// IsAppError differentiates between protocol-defined errors that are relatively benign and other
// fatal errors like an IO error because of some socket problem or network issue.
// Make sure to keep this list in sync with the one above. It should contain all Err* that could
// come back from memcached itself
func IsAppError(err error) bool {
	return err == ErrKeyNotFound ||
		err == ErrKeyExists ||
		err == ErrValueTooBig ||
		err == ErrInvalidArgs ||
		err == ErrItemNotStored ||
		err == ErrBadIncDecValue ||
		err == ErrAuth ||
		err == ErrUnknownCmd ||
		err == ErrNoMem ||
		err == ErrNotSupported ||
		err == ErrInternal ||
		err == ErrBusy ||
		err == ErrTempFailure
}

// RequestType is the protocol-agnostic identifier for the command
type RequestType int

const (
	// RequestUnknown means the parser doesn't know what the request represents. Valid protocol
	// parsing but invalid values.
	RequestUnknown = iota

	// RequestGet represents both a single get and a multi-get, which take differen forms in
	// different protocols. This means it can also be the accumulation of many GETQ commands.
	RequestGet

	// RequestGat is a get-and-touch operation that retrieves the information while updating the TTL
	RequestGat

	// RequestGetE is a custom get which returns the TTL remaining with the data
	RequestGetE

	// RequestSet is to insert a new piece of data unconditionally. What that means is different
	// depending on L1 / L2 handling.
	RequestSet

	// RequestAdd will perform the same operations as set, but only if the key does not exist
	RequestAdd

	// RequestReplace will perform the same operations as set, but only if the key already exists
	RequestReplace

	// RequestAppend appends data to the end of the already existing data for a given key. Does not
	// change the flags or TTL values even if they are given.
	RequestAppend

	// RequestPrepend appends data to the end of the already existing data for a given key. Does not
	// change the flags or TTL values even if they are given.
	RequestPrepend

	// RequestDelete deletes a piece of data from all levels of cache
	RequestDelete

	// RequestTouch updates the TTL for the item specified to a new TTL
	RequestTouch

	// RequestNoop does nothing
	RequestNoop

	// RequestQuit closes the connection
	RequestQuit

	// RequestVersion replies with a string designating the current software version
	RequestVersion
)

// RequestParser represents an interface to parse incoming requests. Each protocol provides its own
// implementation. The return value is an interface{}, but not all hope is lost. The return result
// is guaranteed by implementations to be castable to the type that matches the RequestType returned.
type RequestParser interface {
	Parse() (Request, RequestType, error)
}

// Responder is the interface for a protocol to respond to different commands. It responds in
// whatever way is appropriate, including doing nothing or panic()-ing for unsupported interactions.
// Unsupported interactions are OK to panic() on because they should never be returned from the
// corresponding RequestParser.
type Responder interface {
	Set(opaque uint32, quiet bool) error
	Add(opaque uint32, quiet bool) error
	Replace(opaque uint32, quiet bool) error
	Append(opaque uint32, quiet bool) error
	Prepend(opaque uint32, quiet bool) error
	Get(response GetResponse) error
	GetEnd(opaque uint32, noopEnd bool) error
	GetE(response GetEResponse) error
	GAT(response GetResponse) error
	Delete(opaque uint32) error
	Touch(opaque uint32) error
	Noop(opaque uint32) error
	Quit(opaque uint32, quiet bool) error
	Version(opaque uint32) error
	Error(opaque uint32, reqType RequestType, err error) error
}

type Request interface {
	Opq() uint32
}

// SetRequest corresponds to common.RequestSet. It contains all the information required to fulfill
// a set request.
type SetRequest struct {
	Key     []byte
	Data    []byte
	Flags   uint32
	Exptime uint32
	Opaque  uint32
	Quiet   bool
}

func (r SetRequest) Opq() uint32 {
	return r.Opaque
}

// GetRequest corresponds to common.RequestGet. It contains all the information required to fulfill
// a get requestGets are batch by default, so single gets and batch gets are both represented by the
// same type.
type GetRequest struct {
	Keys       [][]byte
	Opaques    []uint32
	Quiet      []bool
	NoopOpaque uint32
	NoopEnd    bool
}

func (r GetRequest) Opq() uint32 {
	// TODO: better implementation?
	// It's nonsensical but the best way to react in this case since it's a bad situation already.
	// Typically if this method was needed we're already in a fatal error sitation.
	return 0
}

// DeleteRequest corresponds to common.RequestDelete. It contains all the information required to
// fulfill a delete request.
type DeleteRequest struct {
	Key    []byte
	Opaque uint32
}

func (r DeleteRequest) Opq() uint32 {
	return r.Opaque
}

// TouchRequest corresponds to common.RequestTouch. It contains all the information required to
// fulfill a touch request.
type TouchRequest struct {
	Key     []byte
	Exptime uint32
	Opaque  uint32
}

func (r TouchRequest) Opq() uint32 {
	return r.Opaque
}

// GATRequest corresponds to common.RequestGat. It contains all the information required to fulfill
// a get-and-touch request.
type GATRequest struct {
	Key     []byte
	Exptime uint32
	Opaque  uint32
}

func (r GATRequest) Opq() uint32 {
	return r.Opaque
}

// QuitRequest corresponds to common.RequestQuit. It contains all the information required to
// fulfill a quit request.
type QuitRequest struct {
	Opaque uint32
	Quiet  bool
}

func (r QuitRequest) Opq() uint32 {
	return r.Opaque
}

// NoopRequest corresponds to common.RequestNoop. It contains all the information required to
// fulfill a version request.
type NoopRequest struct {
	Opaque uint32
}

func (r NoopRequest) Opq() uint32 {
	return r.Opaque
}

// VersionRequest corresponds to common.RequestVersion. It contains all the information required to
// fulfill a version request.
type VersionRequest struct {
	Opaque uint32
}

func (r VersionRequest) Opq() uint32 {
	return r.Opaque
}

// GetResponse is used in both RequestGet and RequestGat handling. Both respond in the same manner
// but with different opcodes. It is binary-protocol specific, but is still a part of the interface
// of responder to make the handling code more protocol-agnostic.
type GetResponse struct {
	Key    []byte
	Data   []byte
	Opaque uint32
	Flags  uint32
	Miss   bool
	Quiet  bool
}

// GetEResponse is used in the GetE protocol extension
type GetEResponse struct {
	Key     []byte
	Data    []byte
	Opaque  uint32
	Flags   uint32
	Exptime uint32
	Miss    bool
	Quiet   bool
}
<|endoftext|>


<commit_before>// The amazonebs package contains a packer.Builder implementation that
// builds AMIs for Amazon EC2.
//
// In general, there are two types of AMIs that can be created: ebs-backed or
// instance-store. This builder _only_ builds ebs-backed images.
package amazonebs

import (
	"github.com/mitchellh/goamz/aws"
	"github.com/mitchellh/goamz/ec2"
	"github.com/mitchellh/mapstructure"
	"github.com/mitchellh/packer/packer"
	"log"
)

// The unique ID for this builder
const BuilderId = "mitchellh.amazonebs"

type config struct {
	// Access information
	AccessKey string `mapstructure:"access_key"`
	SecretKey string `mapstructure:"secret_key"`

	// Information for the source instance
	Region       string
	SourceAmi    string `mapstructure:"source_ami"`
	InstanceType string `mapstructure:"instance_type"`
	SSHUsername  string `mapstructure:"ssh_username"`
	SSHPort      int    `mapstructure:"ssh_port"`

	// Configuration of the resulting AMI
	AMIName string `mapstructure:"ami_name"`
}

type Builder struct {
	config config
}

func (b *Builder) Prepare(raw interface{}) (err error) {
	err = mapstructure.Decode(raw, &b.config)
	if err != nil {
		return
	}

	log.Printf("Config: %+v", b.config)

	// TODO: config validation and asking for fields:
	// * access key
	// * secret key
	// * region
	// * source ami
	// * instance type
	// * SSH username
	// * SSH port? (or default to 22?)

	return
}

func (b *Builder) Run(ui packer.Ui, hook packer.Hook) packer.Artifact {
	auth := aws.Auth{b.config.AccessKey, b.config.SecretKey}
	region := aws.Regions[b.config.Region]
	ec2conn := ec2.New(auth, region)

	// Setup the state bag and initial state for the steps
	state := make(map[string]interface{})
	state["config"] = b.config
	state["ec2"] = ec2conn
	state["hook"] = hook
	state["ui"] = ui

	// Build the steps
	steps := []Step{
		&stepKeyPair{},
		&stepRunSourceInstance{},
		&stepConnectSSH{},
		&stepProvision{},
		&stepStopInstance{},
		&stepCreateAMI{},
	}

	// Run!
	RunSteps(state, steps)

	// Build the artifact and return it
	return &artifact{state["amis"].(map[string]string)}
}
<commit_msg>builder/amazonebs: panics in edge cases<commit_after>// The amazonebs package contains a packer.Builder implementation that
// builds AMIs for Amazon EC2.
//
// In general, there are two types of AMIs that can be created: ebs-backed or
// instance-store. This builder _only_ builds ebs-backed images.
package amazonebs

import (
	"github.com/mitchellh/goamz/aws"
	"github.com/mitchellh/goamz/ec2"
	"github.com/mitchellh/mapstructure"
	"github.com/mitchellh/packer/packer"
	"log"
)

// The unique ID for this builder
const BuilderId = "mitchellh.amazonebs"

type config struct {
	// Access information
	AccessKey string `mapstructure:"access_key"`
	SecretKey string `mapstructure:"secret_key"`

	// Information for the source instance
	Region       string
	SourceAmi    string `mapstructure:"source_ami"`
	InstanceType string `mapstructure:"instance_type"`
	SSHUsername  string `mapstructure:"ssh_username"`
	SSHPort      int    `mapstructure:"ssh_port"`

	// Configuration of the resulting AMI
	AMIName string `mapstructure:"ami_name"`
}

type Builder struct {
	config config
}

func (b *Builder) Prepare(raw interface{}) (err error) {
	err = mapstructure.Decode(raw, &b.config)
	if err != nil {
		return
	}

	log.Printf("Config: %+v", b.config)

	// TODO: config validation and asking for fields:
	// * access key
	// * secret key
	// * region (exists and valid)
	// * source ami
	// * instance type
	// * SSH username
	// * SSH port? (or default to 22?)

	return
}

func (b *Builder) Run(ui packer.Ui, hook packer.Hook) packer.Artifact {
	// Basic sanity checks. These are panics now because the Prepare
	// method should verify these exist and such.
	if b.config.AccessKey == "" {
		panic("access key not filled in")
	}

	if b.config.SecretKey == "" {
		panic("secret key not filled in")
	}

	if b.config.Region == "" {
		panic("region not filled in")
	}

	region, ok := aws.Regions[b.config.Region]
	if !ok {
		panic("region not found")
	}

	auth := aws.Auth{b.config.AccessKey, b.config.SecretKey}
	ec2conn := ec2.New(auth, region)

	// Setup the state bag and initial state for the steps
	state := make(map[string]interface{})
	state["config"] = b.config
	state["ec2"] = ec2conn
	state["hook"] = hook
	state["ui"] = ui

	// Build the steps
	steps := []Step{
		&stepKeyPair{},
		&stepRunSourceInstance{},
		&stepConnectSSH{},
		&stepProvision{},
		&stepStopInstance{},
		&stepCreateAMI{},
	}

	// Run!
	RunSteps(state, steps)

	// Build the artifact and return it
	return &artifact{state["amis"].(map[string]string)}
}
<|endoftext|>


<commit_before>package chainntnfs

import (
	"fmt"

	"github.com/roasbeef/btcd/chaincfg/chainhash"
	"github.com/roasbeef/btcutil"
)

// ConfNtfn represents a notifier client's request to receive a notification
// once the target transaction gets sufficient confirmations. The client is
// asynchronously notified via the ConfirmationEvent channels.
type ConfNtfn struct {
	// TxID is the hash of the transaction for which confirmation notifications
	// are requested.
	TxID *chainhash.Hash

	// NumConfirmations is the number of confirmations after which the
	// notification is to be sent.
	NumConfirmations uint32

	// Event contains references to the channels that the notifications are to
	// be sent over.
	Event *ConfirmationEvent

	// details describes the transaction's position is the blockchain. May be
	// nil for unconfirmed transactions.
	details *TxConfirmation

	// dispatched is false if the confirmed notification has not been sent yet.
	dispatched bool
}

// NewConfirmationEvent constructs a new ConfirmationEvent with newly opened
// channels.
func NewConfirmationEvent(numConfs uint32) *ConfirmationEvent {
	return &ConfirmationEvent{
		Confirmed:    make(chan *TxConfirmation, 1),
		Updates:      make(chan uint32, numConfs),
		NegativeConf: make(chan int32, 1),
	}
}

// TxConfNotifier is used to register transaction confirmation notifications and
// dispatch them as the transactions confirm. A client can request to be
// notified when a particular transaction has sufficient on-chain confirmations
// (or be notified immediately if the tx already does), and the TxConfNotifier
// will watch changes to the blockchain in order to satisfy these requests.
type TxConfNotifier struct {
	// currentHeight is the height of the tracked blockchain. It is used to
	// determine the number of confirmations a tx has and ensure blocks are
	// connected and disconnected in order.
	currentHeight uint32

	// reorgSafetyLimit is the chain depth beyond which it is assumed a block
	// will not be reorganized out of the chain. This is used to determine when
	// to prune old confirmation requests so that reorgs are handled correctly.
	// The coinbase maturity period is a reasonable value to use.
	reorgSafetyLimit uint32

	// reorgDepth is the depth of a chain organization that this system is being
	// informed of. This is incremented as long as a sequence of blocks are
	// disconnected without being interrupted by a new block.
	reorgDepth uint32

	// confNotifications is an index of notification requests by transaction
	// hash.
	confNotifications map[chainhash.Hash][]*ConfNtfn

	// confTxsByInitialHeight is an index of watched transactions by the height
	// that they are included at in the blockchain. This is tracked so that
	// incorrect notifications are not sent if a transaction is reorganized out
	// of the chain and so that negative confirmations can be recognized.
	confTxsByInitialHeight map[uint32][]*chainhash.Hash

	// ntfnsByConfirmHeight is an index of notification requests by the height
	// at which the transaction will have sufficient confirmations.
	ntfnsByConfirmHeight map[uint32]map[*ConfNtfn]struct{}

	// quit is closed in order to signal that the notifier is gracefully
	// exiting.
	quit chan struct{}
}

// NewTxConfNotifier creates a TxConfNotifier. The current height of the
// blockchain is accepted as a parameter.
func NewTxConfNotifier(startHeight uint32, reorgSafetyLimit uint32) *TxConfNotifier {
	return &TxConfNotifier{
		currentHeight:          startHeight,
		reorgSafetyLimit:       reorgSafetyLimit,
		confNotifications:      make(map[chainhash.Hash][]*ConfNtfn),
		confTxsByInitialHeight: make(map[uint32][]*chainhash.Hash),
		ntfnsByConfirmHeight:   make(map[uint32]map[*ConfNtfn]struct{}),
		quit:                   make(chan struct{}),
	}
}

// Register handles a new notification request. The client will be notified when
// the transaction gets a sufficient number of confirmations on the blockchain.
// If the transaction has already been included in a block on the chain, the
// confirmation details must be given as the txConf argument, otherwise it
// should be nil. If the transaction already has the sufficient number of
// confirmations, this dispatches the notification immediately.
func (tcn *TxConfNotifier) Register(ntfn *ConfNtfn, txConf *TxConfirmation) error {
	select {
	case <-tcn.quit:
		return fmt.Errorf("TxConfNotifier is exiting")
	default:
	}

	if txConf == nil || txConf.BlockHeight > tcn.currentHeight {
		// Transaction is unconfirmed.
		tcn.confNotifications[*ntfn.TxID] =
			append(tcn.confNotifications[*ntfn.TxID], ntfn)
		return nil
	}

	// If the transaction already has the required confirmations, dispatch
	// notification immediately, otherwise record along with the height at
	// which to notify.
	confHeight := txConf.BlockHeight + ntfn.NumConfirmations - 1
	if confHeight <= tcn.currentHeight {
		Log.Infof("Dispatching %v conf notification for %v",
			ntfn.NumConfirmations, ntfn.TxID)
		select {
		case <-tcn.quit:
			return fmt.Errorf("TxConfNotifier is exiting")
		case ntfn.Event.Confirmed <- txConf:
			ntfn.dispatched = true
		}
	} else {
		ntfn.details = txConf
		ntfnSet, exists := tcn.ntfnsByConfirmHeight[confHeight]
		if !exists {
			ntfnSet = make(map[*ConfNtfn]struct{})
			tcn.ntfnsByConfirmHeight[confHeight] = ntfnSet
		}
		ntfnSet[ntfn] = struct{}{}
	}

	// Unless the transaction is finalized, include transaction information in
	// confNotifications and confTxsByInitialHeight in case the tx gets
	// reorganized out of the chain.
	if txConf.BlockHeight+tcn.reorgSafetyLimit > tcn.currentHeight {
		tcn.confNotifications[*ntfn.TxID] =
			append(tcn.confNotifications[*ntfn.TxID], ntfn)
		tcn.confTxsByInitialHeight[txConf.BlockHeight] =
			append(tcn.confTxsByInitialHeight[txConf.BlockHeight], ntfn.TxID)
	}

	return nil
}

// ConnectTip handles a new block extending the current chain. This checks each
// transaction in the block to see if any watched transactions are included.
// Also, if any watched transactions now have the required number of
// confirmations as a result of this block being connected, this dispatches
// notifications.
func (tcn *TxConfNotifier) ConnectTip(blockHash *chainhash.Hash,
	blockHeight uint32, txns []*btcutil.Tx) error {

	select {
	case <-tcn.quit:
		return fmt.Errorf("TxConfNotifier is exiting")
	default:
	}

	if blockHeight != tcn.currentHeight+1 {
		return fmt.Errorf("Received blocks out of order: "+
			"current height=%d, new height=%d",
			tcn.currentHeight, blockHeight)
	}
	tcn.currentHeight++
	tcn.reorgDepth = 0

	// Record any newly confirmed transactions in ntfnsByConfirmHeight so that
	// notifications get dispatched when the tx gets sufficient confirmations.
	// Also record txs in confTxsByInitialHeight so reorgs can be handled
	// correctly.
	for _, tx := range txns {
		txHash := tx.Hash()
		for _, ntfn := range tcn.confNotifications[*txHash] {
			ntfn.details = &TxConfirmation{
				BlockHash:   blockHash,
				BlockHeight: blockHeight,
				TxIndex:     uint32(tx.Index()),
			}

			confHeight := blockHeight + ntfn.NumConfirmations - 1
			ntfnSet, exists := tcn.ntfnsByConfirmHeight[confHeight]
			if !exists {
				ntfnSet = make(map[*ConfNtfn]struct{})
				tcn.ntfnsByConfirmHeight[confHeight] = ntfnSet
			}
			ntfnSet[ntfn] = struct{}{}

			tcn.confTxsByInitialHeight[blockHeight] =
				append(tcn.confTxsByInitialHeight[blockHeight], tx.Hash())
		}
	}

	// Dispatch notifications for all transactions that are considered confirmed
	// at this new block height.
	for ntfn := range tcn.ntfnsByConfirmHeight[tcn.currentHeight] {
		Log.Infof("Dispatching %v conf notification for %v",
			ntfn.NumConfirmations, ntfn.TxID)
		select {
		case ntfn.Event.Confirmed <- ntfn.details:
			ntfn.dispatched = true
		case <-tcn.quit:
			return fmt.Errorf("TxConfNotifier is exiting")
		}
	}
	delete(tcn.ntfnsByConfirmHeight, tcn.currentHeight)

	// Clear entries from confNotifications and confTxsByInitialHeight. We
	// assume that reorgs deeper than the reorg safety limit do not happen, so
	// we can clear out entries for the block that is now mature.
	if tcn.currentHeight >= tcn.reorgSafetyLimit {
		matureBlockHeight := tcn.currentHeight - tcn.reorgSafetyLimit
		for _, txHash := range tcn.confTxsByInitialHeight[matureBlockHeight] {
			delete(tcn.confNotifications, *txHash)
		}
		delete(tcn.confTxsByInitialHeight, matureBlockHeight)
	}

	return nil
}

// DisconnectTip handles the tip of the current chain being disconnected during
// a chain reorganization. If any watched transactions were included in this
// block, internal structures are updated to ensure a confirmation notification
// is not sent unless the transaction is included in the new chain.
func (tcn *TxConfNotifier) DisconnectTip(blockHeight uint32) error {
	select {
	case <-tcn.quit:
		return fmt.Errorf("TxConfNotifier is exiting")
	default:
	}

	if blockHeight != tcn.currentHeight {
		return fmt.Errorf("Received blocks out of order: "+
			"current height=%d, disconnected height=%d",
			tcn.currentHeight, blockHeight)
	}
	tcn.currentHeight--
	tcn.reorgDepth++

	for _, txHash := range tcn.confTxsByInitialHeight[blockHeight] {
		for _, ntfn := range tcn.confNotifications[*txHash] {
			// If notification has been dispatched with sufficient
			// confirmations, notify of the reversal.
			if ntfn.dispatched {
				select {
				case <-ntfn.Event.Confirmed:
					// Drain confirmation notification instead of sending
					// negative conf if the receiver has not processed it yet.
					// This ensures sends to the Confirmed channel are always
					// non-blocking.
				case ntfn.Event.NegativeConf <- int32(tcn.reorgDepth):
				case <-tcn.quit:
					return fmt.Errorf("TxConfNotifier is exiting")
				}
				ntfn.dispatched = false
				continue
			}

			confHeight := blockHeight + ntfn.NumConfirmations - 1
			ntfnSet, exists := tcn.ntfnsByConfirmHeight[confHeight]
			if !exists {
				continue
			}
			delete(ntfnSet, ntfn)
		}
	}
	delete(tcn.confTxsByInitialHeight, blockHeight)

	return nil
}

// TearDown is to be called when the owner of the TxConfNotifier is exiting.
// This closes the event channels of all registered notifications that have
// not been dispatched yet.
func (tcn *TxConfNotifier) TearDown() {
	close(tcn.quit)

	for _, ntfns := range tcn.confNotifications {
		for _, ntfn := range ntfns {
			if ntfn.dispatched {
				continue
			}

			close(ntfn.Event.Confirmed)
			close(ntfn.Event.Updates)
			close(ntfn.Event.NegativeConf)
		}
	}
}
<commit_msg>chainntnfs: modify watched transactions by height to use a set<commit_after>package chainntnfs

import (
	"fmt"

	"github.com/roasbeef/btcd/chaincfg/chainhash"
	"github.com/roasbeef/btcutil"
)

// ConfNtfn represents a notifier client's request to receive a notification
// once the target transaction gets sufficient confirmations. The client is
// asynchronously notified via the ConfirmationEvent channels.
type ConfNtfn struct {
	// TxID is the hash of the transaction for which confirmation notifications
	// are requested.
	TxID *chainhash.Hash

	// NumConfirmations is the number of confirmations after which the
	// notification is to be sent.
	NumConfirmations uint32

	// Event contains references to the channels that the notifications are to
	// be sent over.
	Event *ConfirmationEvent

	// details describes the transaction's position is the blockchain. May be
	// nil for unconfirmed transactions.
	details *TxConfirmation

	// dispatched is false if the confirmed notification has not been sent yet.
	dispatched bool
}

// NewConfirmationEvent constructs a new ConfirmationEvent with newly opened
// channels.
func NewConfirmationEvent(numConfs uint32) *ConfirmationEvent {
	return &ConfirmationEvent{
		Confirmed:    make(chan *TxConfirmation, 1),
		Updates:      make(chan uint32, numConfs),
		NegativeConf: make(chan int32, 1),
	}
}

// TxConfNotifier is used to register transaction confirmation notifications and
// dispatch them as the transactions confirm. A client can request to be
// notified when a particular transaction has sufficient on-chain confirmations
// (or be notified immediately if the tx already does), and the TxConfNotifier
// will watch changes to the blockchain in order to satisfy these requests.
type TxConfNotifier struct {
	// currentHeight is the height of the tracked blockchain. It is used to
	// determine the number of confirmations a tx has and ensure blocks are
	// connected and disconnected in order.
	currentHeight uint32

	// reorgSafetyLimit is the chain depth beyond which it is assumed a block
	// will not be reorganized out of the chain. This is used to determine when
	// to prune old confirmation requests so that reorgs are handled correctly.
	// The coinbase maturity period is a reasonable value to use.
	reorgSafetyLimit uint32

	// reorgDepth is the depth of a chain organization that this system is being
	// informed of. This is incremented as long as a sequence of blocks are
	// disconnected without being interrupted by a new block.
	reorgDepth uint32

	// confNotifications is an index of notification requests by transaction
	// hash.
	confNotifications map[chainhash.Hash][]*ConfNtfn

	// txsByInitialHeight is an index of watched transactions by the height
	// that they are included at in the blockchain. This is tracked so that
	// incorrect notifications are not sent if a transaction is reorganized
	// out of the chain and so that negative confirmations can be recognized.
	txsByInitialHeight map[uint32]map[chainhash.Hash]struct{}

	// ntfnsByConfirmHeight is an index of notification requests by the height
	// at which the transaction will have sufficient confirmations.
	ntfnsByConfirmHeight map[uint32]map[*ConfNtfn]struct{}

	// quit is closed in order to signal that the notifier is gracefully
	// exiting.
	quit chan struct{}
}

// NewTxConfNotifier creates a TxConfNotifier. The current height of the
// blockchain is accepted as a parameter.
func NewTxConfNotifier(startHeight uint32, reorgSafetyLimit uint32) *TxConfNotifier {
	return &TxConfNotifier{
		currentHeight:        startHeight,
		reorgSafetyLimit:     reorgSafetyLimit,
		confNotifications:    make(map[chainhash.Hash][]*ConfNtfn),
		txsByInitialHeight:   make(map[uint32]map[chainhash.Hash]struct{}),
		ntfnsByConfirmHeight: make(map[uint32]map[*ConfNtfn]struct{}),
		quit:                 make(chan struct{}),
	}
}

// Register handles a new notification request. The client will be notified when
// the transaction gets a sufficient number of confirmations on the blockchain.
// If the transaction has already been included in a block on the chain, the
// confirmation details must be given as the txConf argument, otherwise it
// should be nil. If the transaction already has the sufficient number of
// confirmations, this dispatches the notification immediately.
func (tcn *TxConfNotifier) Register(ntfn *ConfNtfn, txConf *TxConfirmation) error {
	select {
	case <-tcn.quit:
		return fmt.Errorf("TxConfNotifier is exiting")
	default:
	}

	if txConf == nil || txConf.BlockHeight > tcn.currentHeight {
		// Transaction is unconfirmed.
		tcn.confNotifications[*ntfn.TxID] =
			append(tcn.confNotifications[*ntfn.TxID], ntfn)
		return nil
	}

	// If the transaction already has the required confirmations, dispatch
	// notification immediately, otherwise record along with the height at
	// which to notify.
	confHeight := txConf.BlockHeight + ntfn.NumConfirmations - 1
	if confHeight <= tcn.currentHeight {
		Log.Infof("Dispatching %v conf notification for %v",
			ntfn.NumConfirmations, ntfn.TxID)
		select {
		case <-tcn.quit:
			return fmt.Errorf("TxConfNotifier is exiting")
		case ntfn.Event.Confirmed <- txConf:
			ntfn.dispatched = true
		}
	} else {
		ntfn.details = txConf
		ntfnSet, exists := tcn.ntfnsByConfirmHeight[confHeight]
		if !exists {
			ntfnSet = make(map[*ConfNtfn]struct{})
			tcn.ntfnsByConfirmHeight[confHeight] = ntfnSet
		}
		ntfnSet[ntfn] = struct{}{}
	}

	// As a final check, we'll also watch the transaction if it's still
	// possible for it to get reorganized out of the chain.
	if txConf.BlockHeight+tcn.reorgSafetyLimit > tcn.currentHeight {
		tcn.confNotifications[*ntfn.TxID] =
			append(tcn.confNotifications[*ntfn.TxID], ntfn)

		txSet, exists := tcn.txsByInitialHeight[txConf.BlockHeight]
		if !exists {
			txSet = make(map[chainhash.Hash]struct{})
			tcn.txsByInitialHeight[txConf.BlockHeight] = txSet
		}
		txSet[*ntfn.TxID] = struct{}{}
	}

	return nil
}

// ConnectTip handles a new block extending the current chain. This checks each
// transaction in the block to see if any watched transactions are included.
// Also, if any watched transactions now have the required number of
// confirmations as a result of this block being connected, this dispatches
// notifications.
func (tcn *TxConfNotifier) ConnectTip(blockHash *chainhash.Hash,
	blockHeight uint32, txns []*btcutil.Tx) error {

	select {
	case <-tcn.quit:
		return fmt.Errorf("TxConfNotifier is exiting")
	default:
	}

	if blockHeight != tcn.currentHeight+1 {
		return fmt.Errorf("Received blocks out of order: "+
			"current height=%d, new height=%d",
			tcn.currentHeight, blockHeight)
	}
	tcn.currentHeight++
	tcn.reorgDepth = 0

	// Record any newly confirmed transactions in ntfnsByConfirmHeight so that
	// notifications get dispatched when the tx gets sufficient confirmations.
	// Also record txs in confTxsByInitialHeight so reorgs can be handled
	// correctly.
	for _, tx := range txns {
		txHash := tx.Hash()
		for _, ntfn := range tcn.confNotifications[*txHash] {
			ntfn.details = &TxConfirmation{
				BlockHash:   blockHash,
				BlockHeight: blockHeight,
				TxIndex:     uint32(tx.Index()),
			}

			confHeight := blockHeight + ntfn.NumConfirmations - 1
			ntfnSet, exists := tcn.ntfnsByConfirmHeight[confHeight]
			if !exists {
				ntfnSet = make(map[*ConfNtfn]struct{})
				tcn.ntfnsByConfirmHeight[confHeight] = ntfnSet
			}
			ntfnSet[ntfn] = struct{}{}

			txSet, exists := tcn.txsByInitialHeight[blockHeight]
			if !exists {
				txSet = make(map[chainhash.Hash]struct{})
				tcn.txsByInitialHeight[blockHeight] = txSet
			}
			txSet[*txHash] = struct{}{}
		}
	}

	// Dispatch notifications for all transactions that are considered confirmed
	// at this new block height.
	for ntfn := range tcn.ntfnsByConfirmHeight[tcn.currentHeight] {
		Log.Infof("Dispatching %v conf notification for %v",
			ntfn.NumConfirmations, ntfn.TxID)
		select {
		case ntfn.Event.Confirmed <- ntfn.details:
			ntfn.dispatched = true
		case <-tcn.quit:
			return fmt.Errorf("TxConfNotifier is exiting")
		}
	}
	delete(tcn.ntfnsByConfirmHeight, tcn.currentHeight)

	// Clear entries from confNotifications and confTxsByInitialHeight. We
	// assume that reorgs deeper than the reorg safety limit do not happen,
	// so we can clear out entries for the block that is now mature.
	if tcn.currentHeight >= tcn.reorgSafetyLimit {
		matureBlockHeight := tcn.currentHeight - tcn.reorgSafetyLimit
		for txHash := range tcn.txsByInitialHeight[matureBlockHeight] {
			delete(tcn.confNotifications, txHash)
		}
		delete(tcn.txsByInitialHeight, matureBlockHeight)
	}

	return nil
}

// DisconnectTip handles the tip of the current chain being disconnected during
// a chain reorganization. If any watched transactions were included in this
// block, internal structures are updated to ensure a confirmation notification
// is not sent unless the transaction is included in the new chain.
func (tcn *TxConfNotifier) DisconnectTip(blockHeight uint32) error {
	select {
	case <-tcn.quit:
		return fmt.Errorf("TxConfNotifier is exiting")
	default:
	}

	if blockHeight != tcn.currentHeight {
		return fmt.Errorf("Received blocks out of order: "+
			"current height=%d, disconnected height=%d",
			tcn.currentHeight, blockHeight)
	}
	tcn.currentHeight--
	tcn.reorgDepth++

	for _, txHash := range tcn.confTxsByInitialHeight[blockHeight] {
		for _, ntfn := range tcn.confNotifications[*txHash] {
			// If notification has been dispatched with sufficient
			// confirmations, notify of the reversal.
			if ntfn.dispatched {
				select {
				case <-ntfn.Event.Confirmed:
					// Drain confirmation notification instead of sending
					// negative conf if the receiver has not processed it yet.
					// This ensures sends to the Confirmed channel are always
					// non-blocking.
				case ntfn.Event.NegativeConf <- int32(tcn.reorgDepth):
				case <-tcn.quit:
					return fmt.Errorf("TxConfNotifier is exiting")
				}
				ntfn.dispatched = false
				continue
			}

			confHeight := blockHeight + ntfn.NumConfirmations - 1
			ntfnSet, exists := tcn.ntfnsByConfirmHeight[confHeight]
			if !exists {
				continue
			}
			delete(ntfnSet, ntfn)
		}
	}

	// Finally, we can remove the transactions we're currently watching that
	// were included in this block height.
	delete(tcn.txsByInitialHeight, blockHeight)

	return nil
}

// TearDown is to be called when the owner of the TxConfNotifier is exiting.
// This closes the event channels of all registered notifications that have
// not been dispatched yet.
func (tcn *TxConfNotifier) TearDown() {
	close(tcn.quit)

	for _, ntfns := range tcn.confNotifications {
		for _, ntfn := range ntfns {
			if ntfn.dispatched {
				continue
			}

			close(ntfn.Event.Confirmed)
			close(ntfn.Event.Updates)
			close(ntfn.Event.NegativeConf)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/Shopify/sarama"

	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"strings"
	"time"
)

var (
	addr    = flag.String("addr", ":8080", "The address to bind to")
	brokers = flag.String("brokers", os.Getenv("KAFKA_PEERS"), "The Kafka brokers to connect to, as a comma separated list")
	verbose = flag.Bool("verbose", false, "Turn on Sarama logging")
)

func main() {
	flag.Parse()

	if *verbose {
		sarama.Logger = log.New(os.Stdout, "[sarama] ", log.LstdFlags)
	}

	if *brokers == "" {
		flag.PrintDefaults()
		os.Exit(1)
	}

	brokerList := strings.Split(*brokers, ",")
	log.Printf("Kafka brokers: %s", strings.Join(brokerList, ", "))

	server := &Server{
		DataCollector:     newDataCollector(brokerList),
		AccessLogProducer: newAccessLogProducer(brokerList),
	}
	defer func() {
		if err := server.Close(); err != nil {
			log.Println("Failed to close server", err)
		}
	}()

	log.Fatal(server.Run(*addr))
}

type Server struct {
	DataCollector     sarama.SyncProducer
	AccessLogProducer sarama.AsyncProducer
}

func (s *Server) Close() error {
	if err := s.DataCollector.Close(); err != nil {
		log.Println("Failed to shut down data collector cleanly", err)
	}

	if err := s.AccessLogProducer.Close(); err != nil {
		log.Println("Failed to shut down access log producer cleanly", err)
	}

	return nil
}

func (s *Server) Handler() http.Handler {
	return s.withAccessLog(s.collectQueryStringData())
}

func (s *Server) Run(addr string) error {
	httpServer := &http.Server{
		Addr:    addr,
		Handler: s.Handler(),
	}

	log.Printf("Listening for requests on %s...\n", addr)
	return httpServer.ListenAndServe()
}

func (s *Server) collectQueryStringData() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/" {
			http.NotFound(w, r)
			return
		}

		// We are not setting a message key, which means that all messages will
		// be distributed randomly over the different partitions.
		partition, offset, err := s.DataCollector.SendMessage(&sarama.ProducerMessage{
			Topic: "important",
			Value: sarama.StringEncoder(r.URL.RawQuery),
		})

		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "Failed to store your data:, %s", err)
		} else {
			// The tuple (topic, partition, offset) can be used as a unique identifier
			// for a message in a Kafka cluster.
			fmt.Fprintf(w, "Your data is stored with unique identifier important/%d/%d", partition, offset)
		}
	})
}

type accessLogEntry struct {
	Method       string  `json:"method"`
	Host         string  `json:"host"`
	Path         string  `json:"path"`
	IP           string  `json:"ip"`
	ResponseTime float64 `json:"response_time"`

	encoded []byte
	err     error
}

func (ale *accessLogEntry) ensureEncoded() {
	if ale.encoded == nil && ale.err == nil {
		ale.encoded, ale.err = json.Marshal(ale)
	}
}

func (ale *accessLogEntry) Length() int {
	ale.ensureEncoded()
	return len(ale.encoded)
}

func (ale *accessLogEntry) Encode() ([]byte, error) {
	ale.ensureEncoded()
	return ale.encoded, ale.err
}

func (s *Server) withAccessLog(next http.Handler) http.Handler {

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		started := time.Now()

		next.ServeHTTP(w, r)

		entry := &accessLogEntry{
			Method:       r.Method,
			Host:         r.Host,
			Path:         r.RequestURI,
			IP:           r.RemoteAddr,
			ResponseTime: float64(time.Since(started)) / float64(time.Second),
		}

		// We will use the client's IP address as key. This will cause
		// all the access log entries of the same IP address to end up
		// on the same partition.
		s.AccessLogProducer.Input() <- &sarama.ProducerMessage{
			Topic: "access_log",
			Key:   sarama.StringEncoder(r.RemoteAddr),
			Value: entry,
		}
	})
}

func newDataCollector(brokerList []string) sarama.SyncProducer {

	// For the data collector, we are looking for strong consistency semantics.
	// Because we don't change the flush settings, sarama will try to produce messages
	// as fast as possible to keep latency low.
	config := sarama.NewConfig()
	config.Producer.RequiredAcks = sarama.WaitForAll // Wait for all in-sync replicas to ack the message
	config.Producer.Retry.Max = 10                   // Retry up to 10 times to produce the message

	// On the broker side, you may want to change the following settings to get
	// stronger consistency guarantees:
	// - For your broker, set `unclean.leader.election.enable` to false
	// - For the topic, you could increase `min.insync.replicas`.

	producer, err := sarama.NewSyncProducer(brokerList, config)
	if err != nil {
		log.Fatalln("Failed to start Sarama producer:", err)
	}

	return producer
}

func newAccessLogProducer(brokerList []string) sarama.AsyncProducer {

	// For the access log, we are looking for AP semantics, with high throughput.
	// By creating batches of compressed messages, we reduce network I/O at a cost of more latency.
	config := sarama.NewConfig()
	config.Producer.RequiredAcks = sarama.WaitForLocal       // Only wait for the leader to ack
	config.Producer.Compression = sarama.CompressionSnappy   // Compress messages
	config.Producer.Flush.Frequency = 500 * time.Millisecond // Flush batches every 500ms

	producer, err := sarama.NewAsyncProducer(brokerList, config)
	if err != nil {
		log.Fatalln("Failed to start Sarama producer:", err)
	}

	// We will just log to STDOUT if we're not able to produce messages.
	// Note: messages will only be returned here after all retry attempts are exhausted.
	go func() {
		for err := range producer.Errors() {
			log.Println("Failed to write access log entry:", err)
		}
	}()

	return producer
}
<commit_msg>add the client key, certificate and ca file to the http_server test program<commit_after>package main

import (
	"github.com/Shopify/sarama"

	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
	"time"
)

var (
	addr       = flag.String("addr", ":8080", "The address to bind to")
	brokers    = flag.String("brokers", os.Getenv("KAFKA_PEERS"), "The Kafka brokers to connect to, as a comma separated list")
	verbose    = flag.Bool("verbose", false, "Turn on Sarama logging")
	cert_file  = flag.String("certificate", "", "The certificate file for client authentication")
	key_file   = flag.String("key", "", "The key file for client authentication")
	ca_file    = flag.String("ca", "", "The certificate authority file for client authentication")
	tls_config *tls.Config
)

func main() {
	flag.Parse()

	if *cert_file != "" && *key_file != "" && *ca_file != "" {
		cert, err := tls.LoadX509KeyPair(*cert_file, *key_file)
		if err != nil {
			log.Fatal(err)
		}

		caCert, err := ioutil.ReadFile(*ca_file)
		if err != nil {
			log.Fatal(err)
		}

		caCertPool := x509.NewCertPool()
		caCertPool.AppendCertsFromPEM(caCert)

		tls_config = &tls.Config{
			Certificates: []tls.Certificate{cert},
			//		RootCAs:            caCertPool,
			InsecureSkipVerify: true,
		}
	}

	if *verbose {
		sarama.Logger = log.New(os.Stdout, "[sarama] ", log.LstdFlags)
	}

	if *brokers == "" {
		flag.PrintDefaults()
		os.Exit(1)
	}

	brokerList := strings.Split(*brokers, ",")
	log.Printf("Kafka brokers: %s", strings.Join(brokerList, ", "))

	server := &Server{
		DataCollector:     newDataCollector(brokerList),
		AccessLogProducer: newAccessLogProducer(brokerList),
	}
	defer func() {
		if err := server.Close(); err != nil {
			log.Println("Failed to close server", err)
		}
	}()

	log.Fatal(server.Run(*addr))
}

type Server struct {
	DataCollector     sarama.SyncProducer
	AccessLogProducer sarama.AsyncProducer
}

func (s *Server) Close() error {
	if err := s.DataCollector.Close(); err != nil {
		log.Println("Failed to shut down data collector cleanly", err)
	}

	if err := s.AccessLogProducer.Close(); err != nil {
		log.Println("Failed to shut down access log producer cleanly", err)
	}

	return nil
}

func (s *Server) Handler() http.Handler {
	return s.withAccessLog(s.collectQueryStringData())
}

func (s *Server) Run(addr string) error {
	httpServer := &http.Server{
		Addr:    addr,
		Handler: s.Handler(),
	}

	log.Printf("Listening for requests on %s...\n", addr)
	return httpServer.ListenAndServe()
}

func (s *Server) collectQueryStringData() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/" {
			http.NotFound(w, r)
			return
		}

		// We are not setting a message key, which means that all messages will
		// be distributed randomly over the different partitions.
		partition, offset, err := s.DataCollector.SendMessage(&sarama.ProducerMessage{
			Topic: "important",
			Value: sarama.StringEncoder(r.URL.RawQuery),
		})

		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "Failed to store your data:, %s", err)
		} else {
			// The tuple (topic, partition, offset) can be used as a unique identifier
			// for a message in a Kafka cluster.
			fmt.Fprintf(w, "Your data is stored with unique identifier important/%d/%d", partition, offset)
		}
	})
}

type accessLogEntry struct {
	Method       string  `json:"method"`
	Host         string  `json:"host"`
	Path         string  `json:"path"`
	IP           string  `json:"ip"`
	ResponseTime float64 `json:"response_time"`

	encoded []byte
	err     error
}

func (ale *accessLogEntry) ensureEncoded() {
	if ale.encoded == nil && ale.err == nil {
		ale.encoded, ale.err = json.Marshal(ale)
	}
}

func (ale *accessLogEntry) Length() int {
	ale.ensureEncoded()
	return len(ale.encoded)
}

func (ale *accessLogEntry) Encode() ([]byte, error) {
	ale.ensureEncoded()
	return ale.encoded, ale.err
}

func (s *Server) withAccessLog(next http.Handler) http.Handler {

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		started := time.Now()

		next.ServeHTTP(w, r)

		entry := &accessLogEntry{
			Method:       r.Method,
			Host:         r.Host,
			Path:         r.RequestURI,
			IP:           r.RemoteAddr,
			ResponseTime: float64(time.Since(started)) / float64(time.Second),
		}

		// We will use the client's IP address as key. This will cause
		// all the access log entries of the same IP address to end up
		// on the same partition.
		s.AccessLogProducer.Input() <- &sarama.ProducerMessage{
			Topic: "access_log",
			Key:   sarama.StringEncoder(r.RemoteAddr),
			Value: entry,
		}
	})
}

func newDataCollector(brokerList []string) sarama.SyncProducer {

	// For the data collector, we are looking for strong consistency semantics.
	// Because we don't change the flush settings, sarama will try to produce messages
	// as fast as possible to keep latency low.
	config := sarama.NewConfig()
	config.Producer.RequiredAcks = sarama.WaitForAll // Wait for all in-sync replicas to ack the message
	config.Producer.Retry.Max = 10                   // Retry up to 10 times to produce the message
	if tls_config != nil {
		config.Net.TLS.Config = tls_config
		config.Net.TLS.Enable = true
	}

	// On the broker side, you may want to change the following settings to get
	// stronger consistency guarantees:
	// - For your broker, set `unclean.leader.election.enable` to false
	// - For the topic, you could increase `min.insync.replicas`.

	producer, err := sarama.NewSyncProducer(brokerList, config)
	if err != nil {
		log.Fatalln("Failed to start Sarama producer:", err)
	}

	return producer
}

func newAccessLogProducer(brokerList []string) sarama.AsyncProducer {

	// For the access log, we are looking for AP semantics, with high throughput.
	// By creating batches of compressed messages, we reduce network I/O at a cost of more latency.
	config := sarama.NewConfig()
	if tls_config != nil {
		config.Net.TLS.Enable = true
		config.Net.TLS.Config = tls_config
	}
	config.Producer.RequiredAcks = sarama.WaitForLocal       // Only wait for the leader to ack
	config.Producer.Compression = sarama.CompressionSnappy   // Compress messages
	config.Producer.Flush.Frequency = 500 * time.Millisecond // Flush batches every 500ms

	producer, err := sarama.NewAsyncProducer(brokerList, config)
	if err != nil {
		log.Fatalln("Failed to start Sarama producer:", err)
	}

	// We will just log to STDOUT if we're not able to produce messages.
	// Note: messages will only be returned here after all retry attempts are exhausted.
	go func() {
		for err := range producer.Errors() {
			log.Println("Failed to write access log entry:", err)
		}
	}()

	return producer
}
<|endoftext|>


<commit_before>package service

import (
	"fmt"
	"strconv"
	"strings"

	"golang.org/x/net/context"

	"github.com/docker/cli/cli"
	"github.com/docker/cli/cli/command"
	"github.com/docker/docker/api/types"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
)

type scaleOptions struct {
	detach bool
}

func newScaleCommand(dockerCli *command.DockerCli) *cobra.Command {
	options := &scaleOptions{}

	cmd := &cobra.Command{
		Use:   "scale SERVICE=REPLICAS [SERVICE=REPLICAS...]",
		Short: "Scale one or multiple replicated services",
		Args:  scaleArgs,
		RunE: func(cmd *cobra.Command, args []string) error {
			return runScale(dockerCli, cmd.Flags(), options, args)
		},
	}

	flags := cmd.Flags()
	addDetachFlag(flags, &options.detach)
	return cmd
}

func scaleArgs(cmd *cobra.Command, args []string) error {
	if err := cli.RequiresMinArgs(1)(cmd, args); err != nil {
		return err
	}
	for _, arg := range args {
		if parts := strings.SplitN(arg, "=", 2); len(parts) != 2 {
			return errors.Errorf(
				"Invalid scale specifier '%s'.\nSee '%s --help'.\n\nUsage:  %s\n\n%s",
				arg,
				cmd.CommandPath(),
				cmd.UseLine(),
				cmd.Short,
			)
		}
	}
	return nil
}

func runScale(dockerCli *command.DockerCli, flags *pflag.FlagSet, options *scaleOptions, args []string) error {
	var errs []string
	ctx := context.Background()

	for _, arg := range args {
		parts := strings.SplitN(arg, "=", 2)
		serviceID, scaleStr := parts[0], parts[1]

		// validate input arg scale number
		scale, err := strconv.ParseUint(scaleStr, 10, 64)
		if err != nil {
			errs = append(errs, fmt.Sprintf("%s: invalid replicas value %s: %v", serviceID, scaleStr, err))
			continue
		}

		if err := runServiceScale(ctx, dockerCli, serviceID, scale); err != nil {
			errs = append(errs, fmt.Sprintf("%s: %v", serviceID, err))
		}

		if options.detach {
			continue
		}

		if err := waitOnService(ctx, dockerCli, serviceID, false); err != nil {
			errs = append(errs, fmt.Sprintf("%s: %v", serviceID, err))
		}
	}

	if options.detach {
		warnDetachDefault(dockerCli.Err(), dockerCli.Client().ClientVersion(), flags, "scaled")
	}

	if len(errs) == 0 {
		return nil
	}
	return errors.Errorf(strings.Join(errs, "\n"))
}

func runServiceScale(ctx context.Context, dockerCli *command.DockerCli, serviceID string, scale uint64) error {
	client := dockerCli.Client()

	service, _, err := client.ServiceInspectWithRaw(ctx, serviceID, types.ServiceInspectOptions{})
	if err != nil {
		return err
	}

	serviceMode := &service.Spec.Mode
	if serviceMode.Replicated == nil {
		return errors.Errorf("scale can only be used with replicated mode")
	}

	serviceMode.Replicated.Replicas = &scale

	response, err := client.ServiceUpdate(ctx, service.ID, service.Version, service.Spec, types.ServiceUpdateOptions{})
	if err != nil {
		return err
	}

	for _, warning := range response.Warnings {
		fmt.Fprintln(dockerCli.Err(), warning)
	}

	fmt.Fprintf(dockerCli.Out(), "%s scaled to %d\n", serviceID, scale)
	return nil
}
<commit_msg>scale all services at once<commit_after>package service

import (
	"fmt"
	"strconv"
	"strings"

	"golang.org/x/net/context"

	"github.com/docker/cli/cli"
	"github.com/docker/cli/cli/command"
	"github.com/docker/docker/api/types"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
)

type scaleOptions struct {
	detach bool
}

func newScaleCommand(dockerCli *command.DockerCli) *cobra.Command {
	options := &scaleOptions{}

	cmd := &cobra.Command{
		Use:   "scale SERVICE=REPLICAS [SERVICE=REPLICAS...]",
		Short: "Scale one or multiple replicated services",
		Args:  scaleArgs,
		RunE: func(cmd *cobra.Command, args []string) error {
			return runScale(dockerCli, cmd.Flags(), options, args)
		},
	}

	flags := cmd.Flags()
	addDetachFlag(flags, &options.detach)
	return cmd
}

func scaleArgs(cmd *cobra.Command, args []string) error {
	if err := cli.RequiresMinArgs(1)(cmd, args); err != nil {
		return err
	}
	for _, arg := range args {
		if parts := strings.SplitN(arg, "=", 2); len(parts) != 2 {
			return errors.Errorf(
				"Invalid scale specifier '%s'.\nSee '%s --help'.\n\nUsage:  %s\n\n%s",
				arg,
				cmd.CommandPath(),
				cmd.UseLine(),
				cmd.Short,
			)
		}
	}
	return nil
}

func runScale(dockerCli *command.DockerCli, flags *pflag.FlagSet, options *scaleOptions, args []string) error {
	var errs []string
	var serviceIDs []string
	ctx := context.Background()

	for _, arg := range args {
		parts := strings.SplitN(arg, "=", 2)
		serviceID, scaleStr := parts[0], parts[1]

		// validate input arg scale number
		scale, err := strconv.ParseUint(scaleStr, 10, 64)
		if err != nil {
			errs = append(errs, fmt.Sprintf("%s: invalid replicas value %s: %v", serviceID, scaleStr, err))
			continue
		}

		if err := runServiceScale(ctx, dockerCli, serviceID, scale); err != nil {
			errs = append(errs, fmt.Sprintf("%s: %v", serviceID, err))
		} else {
			serviceIDs = append(serviceIDs, serviceID)
		}

	}

	if len(serviceIDs) > 0 {
		if options.detach {
			warnDetachDefault(dockerCli.Err(), dockerCli.Client().ClientVersion(), flags, "scaled")
		} else {
			for _, serviceID := range serviceIDs {
				if err := waitOnService(ctx, dockerCli, serviceID, false); err != nil {
					errs = append(errs, fmt.Sprintf("%s: %v", serviceID, err))
				}
			}
		}
	}

	if len(errs) == 0 {
		return nil
	}
	return errors.Errorf(strings.Join(errs, "\n"))
}

func runServiceScale(ctx context.Context, dockerCli *command.DockerCli, serviceID string, scale uint64) error {
	client := dockerCli.Client()

	service, _, err := client.ServiceInspectWithRaw(ctx, serviceID, types.ServiceInspectOptions{})
	if err != nil {
		return err
	}

	serviceMode := &service.Spec.Mode
	if serviceMode.Replicated == nil {
		return errors.Errorf("scale can only be used with replicated mode")
	}

	serviceMode.Replicated.Replicas = &scale

	response, err := client.ServiceUpdate(ctx, service.ID, service.Version, service.Spec, types.ServiceUpdateOptions{})
	if err != nil {
		return err
	}

	for _, warning := range response.Warnings {
		fmt.Fprintln(dockerCli.Err(), warning)
	}

	fmt.Fprintf(dockerCli.Out(), "%s scaled to %d\n", serviceID, scale)
	return nil
}
<|endoftext|>


<commit_before>package driver

import (
	"log"
	"os"
	"path/filepath"
	"reflect"
	"testing"

	"github.com/hashicorp/nomad/client/allocdir"
	"github.com/hashicorp/nomad/client/config"
	"github.com/hashicorp/nomad/nomad/structs"
)

var basicResources = &structs.Resources{
	CPU:      250,
	MemoryMB: 256,
	Networks: []*structs.NetworkResource{
		&structs.NetworkResource{
			IP:            "1.2.3.4",
			ReservedPorts: []int{12345},
			DynamicPorts:  []string{"HTTP"},
		},
	},
}

func testLogger() *log.Logger {
	return log.New(os.Stderr, "", log.LstdFlags)
}

func testConfig() *config.Config {
	conf := &config.Config{}
	conf.StateDir = os.TempDir()
	conf.AllocDir = os.TempDir()
	return conf
}

func testDriverContext(task string) *DriverContext {
	cfg := testConfig()
	return NewDriverContext(task, cfg, cfg.Node, testLogger())
}

func testDriverExecContext(task *structs.Task, driverCtx *DriverContext) *ExecContext {
	allocDir := allocdir.NewAllocDir(filepath.Join(driverCtx.config.AllocDir, structs.GenerateUUID()))
	allocDir.Build([]*structs.Task{task})
	ctx := NewExecContext(allocDir, "dummyAllocId")
	return ctx
}

func TestDriver_TaskEnvironmentVariables(t *testing.T) {
	ctx := &ExecContext{}
	task := &structs.Task{
		Env: map[string]string{
			"HELLO": "world",
			"lorem": "ipsum",
		},
		Resources: &structs.Resources{
			CPU:      1000,
			MemoryMB: 500,
			Networks: []*structs.NetworkResource{
				&structs.NetworkResource{
					IP:            "1.2.3.4",
					ReservedPorts: []int{80, 443, 8080, 12345},
					DynamicPorts:  []string{"admin", "5000"},
				},
			},
		},
		Meta: map[string]string{
			"chocolate":  "cake",
			"strawberry": "icecream",
		},
	}

	env := TaskEnvironmentVariables(ctx, task)
	exp := map[string]string{
		"NOMAD_CPU_LIMIT":       "1000",
		"NOMAD_MEMORY_LIMIT":    "500",
		"NOMAD_IP":              "1.2.3.4",
		"NOMAD_PORT_admin":      "8080",
		"NOMAD_PORT_5000":       "12345",
		"NOMAD_META_CHOCOLATE":  "cake",
		"NOMAD_META_STRAWBERRY": "icecream",
		"HELLO":                 "world",
		"lorem":                 "ipsum",
	}

	act := env.Map()
	if !reflect.DeepEqual(act, exp) {
		t.Fatalf("TaskEnvironmentVariables(%#v, %#v) returned %#v; want %#v", ctx, task, act, exp)
	}
}
<commit_msg>Change this to a valid IP address so docker can bind to it<commit_after>package driver

import (
	"log"
	"os"
	"path/filepath"
	"reflect"
	"testing"

	"github.com/hashicorp/nomad/client/allocdir"
	"github.com/hashicorp/nomad/client/config"
	"github.com/hashicorp/nomad/nomad/structs"
)

var basicResources = &structs.Resources{
	CPU:      250,
	MemoryMB: 256,
	Networks: []*structs.NetworkResource{
		&structs.NetworkResource{
			IP:            "0.0.0.0",
			ReservedPorts: []int{12345},
			DynamicPorts:  []string{"HTTP"},
		},
	},
}

func testLogger() *log.Logger {
	return log.New(os.Stderr, "", log.LstdFlags)
}

func testConfig() *config.Config {
	conf := &config.Config{}
	conf.StateDir = os.TempDir()
	conf.AllocDir = os.TempDir()
	return conf
}

func testDriverContext(task string) *DriverContext {
	cfg := testConfig()
	return NewDriverContext(task, cfg, cfg.Node, testLogger())
}

func testDriverExecContext(task *structs.Task, driverCtx *DriverContext) *ExecContext {
	allocDir := allocdir.NewAllocDir(filepath.Join(driverCtx.config.AllocDir, structs.GenerateUUID()))
	allocDir.Build([]*structs.Task{task})
	ctx := NewExecContext(allocDir, "dummyAllocId")
	return ctx
}

func TestDriver_TaskEnvironmentVariables(t *testing.T) {
	ctx := &ExecContext{}
	task := &structs.Task{
		Env: map[string]string{
			"HELLO": "world",
			"lorem": "ipsum",
		},
		Resources: &structs.Resources{
			CPU:      1000,
			MemoryMB: 500,
			Networks: []*structs.NetworkResource{
				&structs.NetworkResource{
					IP:            "1.2.3.4",
					ReservedPorts: []int{80, 443, 8080, 12345},
					DynamicPorts:  []string{"admin", "5000"},
				},
			},
		},
		Meta: map[string]string{
			"chocolate":  "cake",
			"strawberry": "icecream",
		},
	}

	env := TaskEnvironmentVariables(ctx, task)
	exp := map[string]string{
		"NOMAD_CPU_LIMIT":       "1000",
		"NOMAD_MEMORY_LIMIT":    "500",
		"NOMAD_IP":              "1.2.3.4",
		"NOMAD_PORT_admin":      "8080",
		"NOMAD_PORT_5000":       "12345",
		"NOMAD_META_CHOCOLATE":  "cake",
		"NOMAD_META_STRAWBERRY": "icecream",
		"HELLO":                 "world",
		"lorem":                 "ipsum",
	}

	act := env.Map()
	if !reflect.DeepEqual(act, exp) {
		t.Fatalf("TaskEnvironmentVariables(%#v, %#v) returned %#v; want %#v", ctx, task, act, exp)
	}
}
<|endoftext|>


<commit_before>package parser

import (
	"plaid/lexer"
	"testing"
)

var tok = lexer.Token{}

func TestProgram(t *testing.T) {
	(Program{}).isNode()

	prog := Program{[]Stmt{
		DeclarationStmt{tok, IdentExpr{tok, "a"}, NumberExpr{tok, 123}},
		DeclarationStmt{tok, IdentExpr{tok, "b"}, NumberExpr{tok, 456}},
	}}
	expectString(t, prog, "(let a 123)\n(let b 456)")

	prog = Program{[]Stmt{}}
	expectString(t, prog, "")
}

func TestStmtBlock(t *testing.T) {
	(StmtBlock{}).isNode()

	block := StmtBlock{tok, []Stmt{
		DeclarationStmt{tok, IdentExpr{tok, "a"}, NumberExpr{tok, 123}},
		DeclarationStmt{tok, IdentExpr{tok, "b"}, NumberExpr{tok, 456}},
	}, tok}
	expectString(t, block, "{\n  (let a 123)\n  (let b 456)}")
}

func TestDeclarationStmt(t *testing.T) {
	(DeclarationStmt{}).isNode()
	(DeclarationStmt{}).isStmt()

	expectString(t, DeclarationStmt{tok, IdentExpr{tok, "a"}, NumberExpr{tok, 123}}, "(let a 123)")
}

func TestTypeIdent(t *testing.T) {
	(TypeIdent{}).isNode()
	(TypeIdent{}).isType()

	expectString(t, TypeIdent{tok, "Int"}, "Int")
}

func TestTypeList(t *testing.T) {
	(TypeList{}).isNode()
	(TypeList{}).isType()

	expectString(t, TypeList{tok, TypeIdent{tok, "Int"}}, "[Int]")
}

func TestTypeOptional(t *testing.T) {
	(TypeOptional{}).isNode()
	(TypeOptional{}).isType()

	expectString(t, TypeOptional{tok, TypeIdent{tok, "Int"}}, "Int?")
}

func TestFunctionExpr(t *testing.T) {
	(FunctionExpr{}).isNode()
	(FunctionExpr{}).isExpr()

	params := []FunctionParam{
		FunctionParam{IdentExpr{tok, "x"}, TypeIdent{tok, "Int"}},
		FunctionParam{IdentExpr{tok, "y"}, nil},
	}
	ret := TypeIdent{tok, "Str"}
	block := StmtBlock{tok, []Stmt{
		DeclarationStmt{tok, IdentExpr{tok, "z"}, NumberExpr{tok, 123}},
	}, tok}

	expectString(t, FunctionExpr{tok, params, ret, block}, "(fn (x:Int y:Bool):Str {\n  (let z 123)})")
}

func TestBinaryExpr(t *testing.T) {
	(BinaryExpr{}).isNode()
	(BinaryExpr{}).isExpr()

	expectString(t, BinaryExpr{"+", tok, NumberExpr{tok, 123}, NumberExpr{tok, 456}}, "(+ 123 456)")
}

func TestUnaryExpr(t *testing.T) {
	(UnaryExpr{}).isNode()
	(UnaryExpr{}).isExpr()

	expectString(t, UnaryExpr{"+", tok, NumberExpr{tok, 123}}, "(+ 123)")
}

func TestIdentExpr(t *testing.T) {
	(IdentExpr{}).isNode()
	(IdentExpr{}).isExpr()

	expectString(t, IdentExpr{tok, "abc"}, "abc")
}

func TestStringExpr(t *testing.T) {
	(StringExpr{}).isNode()
	(StringExpr{}).isExpr()

	expectString(t, StringExpr{tok, "abc"}, "\"abc\"")
}

func TestNumberExpr(t *testing.T) {
	(NumberExpr{}).isNode()
	(NumberExpr{}).isExpr()

	expectString(t, NumberExpr{tok, 123}, "123")
}

func TestIndentBlock(t *testing.T) {
	source := "foo\nbar\n  baz"
	exp := "...foo\n...bar\n...  baz"
	got := indentBlock("...", source)

	if exp != got {
		t.Errorf("Expected '%s', got '%s'\n", exp, got)
	}
}

func expectString(t *testing.T, node Node, exp string) {
	got := node.String()

	if exp != got {
		t.Errorf("Expected %s, got %s\n", exp, got)
	}
}
<commit_msg>test FunctionParam 'String()' with no type signature<commit_after>package parser

import (
	"plaid/lexer"
	"testing"
)

var tok = lexer.Token{}

func TestProgram(t *testing.T) {
	(Program{}).isNode()

	prog := Program{[]Stmt{
		DeclarationStmt{tok, IdentExpr{tok, "a"}, NumberExpr{tok, 123}},
		DeclarationStmt{tok, IdentExpr{tok, "b"}, NumberExpr{tok, 456}},
	}}
	expectString(t, prog, "(let a 123)\n(let b 456)")

	prog = Program{[]Stmt{}}
	expectString(t, prog, "")
}

func TestStmtBlock(t *testing.T) {
	(StmtBlock{}).isNode()

	block := StmtBlock{tok, []Stmt{
		DeclarationStmt{tok, IdentExpr{tok, "a"}, NumberExpr{tok, 123}},
		DeclarationStmt{tok, IdentExpr{tok, "b"}, NumberExpr{tok, 456}},
	}, tok}
	expectString(t, block, "{\n  (let a 123)\n  (let b 456)}")
}

func TestDeclarationStmt(t *testing.T) {
	(DeclarationStmt{}).isNode()
	(DeclarationStmt{}).isStmt()

	expectString(t, DeclarationStmt{tok, IdentExpr{tok, "a"}, NumberExpr{tok, 123}}, "(let a 123)")
}

func TestTypeIdent(t *testing.T) {
	(TypeIdent{}).isNode()
	(TypeIdent{}).isType()

	expectString(t, TypeIdent{tok, "Int"}, "Int")
}

func TestTypeList(t *testing.T) {
	(TypeList{}).isNode()
	(TypeList{}).isType()

	expectString(t, TypeList{tok, TypeIdent{tok, "Int"}}, "[Int]")
}

func TestTypeOptional(t *testing.T) {
	(TypeOptional{}).isNode()
	(TypeOptional{}).isType()

	expectString(t, TypeOptional{tok, TypeIdent{tok, "Int"}}, "Int?")
}

func TestFunctionExpr(t *testing.T) {
	(FunctionExpr{}).isNode()
	(FunctionExpr{}).isExpr()

	params := []FunctionParam{
		FunctionParam{IdentExpr{tok, "x"}, TypeIdent{tok, "Int"}},
		FunctionParam{IdentExpr{tok, "y"}, nil},
	}
	ret := TypeIdent{tok, "Str"}
	block := StmtBlock{tok, []Stmt{
		DeclarationStmt{tok, IdentExpr{tok, "z"}, NumberExpr{tok, 123}},
	}, tok}

	expectString(t, FunctionExpr{tok, params, ret, block}, "(fn (x:Int y):Str {\n  (let z 123)})")
}

func TestBinaryExpr(t *testing.T) {
	(BinaryExpr{}).isNode()
	(BinaryExpr{}).isExpr()

	expectString(t, BinaryExpr{"+", tok, NumberExpr{tok, 123}, NumberExpr{tok, 456}}, "(+ 123 456)")
}

func TestUnaryExpr(t *testing.T) {
	(UnaryExpr{}).isNode()
	(UnaryExpr{}).isExpr()

	expectString(t, UnaryExpr{"+", tok, NumberExpr{tok, 123}}, "(+ 123)")
}

func TestIdentExpr(t *testing.T) {
	(IdentExpr{}).isNode()
	(IdentExpr{}).isExpr()

	expectString(t, IdentExpr{tok, "abc"}, "abc")
}

func TestStringExpr(t *testing.T) {
	(StringExpr{}).isNode()
	(StringExpr{}).isExpr()

	expectString(t, StringExpr{tok, "abc"}, "\"abc\"")
}

func TestNumberExpr(t *testing.T) {
	(NumberExpr{}).isNode()
	(NumberExpr{}).isExpr()

	expectString(t, NumberExpr{tok, 123}, "123")
}

func TestIndentBlock(t *testing.T) {
	source := "foo\nbar\n  baz"
	exp := "...foo\n...bar\n...  baz"
	got := indentBlock("...", source)

	if exp != got {
		t.Errorf("Expected '%s', got '%s'\n", exp, got)
	}
}

func expectString(t *testing.T, node Node, exp string) {
	got := node.String()

	if exp != got {
		t.Errorf("Expected %s, got %s\n", exp, got)
	}
}
<|endoftext|>