<commit_before>// Copyright (c) Liam Stanley <me@liamstanley.io>. All rights reserved. Use
// of this source code is governed by the MIT license that can be found in
// the LICENSE file.

package girc

import (
	"errors"
	"fmt"
)

// Commands holds a large list of useful methods to interact with the server,
// and wrappers for common events.
type Commands struct {
	c *Client
}

// Nick changes the client nickname.
func (cmd *Commands) Nick(name string) error {
	if !IsValidNick(name) {
		return &ErrInvalidTarget{Target: name}
	}

	cmd.c.Send(&Event{Command: NICK, Params: []string{name}})
	return nil
}

// Join attempts to enter a list of IRC channels, at bulk if possible to
// prevent sending extensive JOIN commands.
func (cmd *Commands) Join(channels ...string) error {
	// We can join multiple channels at once, however we need to ensure that
	// we are not exceeding the line length. (see maxLength)
	max := maxLength - len(JOIN) - 1

	var buffer string

	for i := 0; i < len(channels); i++ {
		if !IsValidChannel(channels[i]) {
			return &ErrInvalidTarget{Target: channels[i]}
		}

		if len(buffer+","+channels[i]) > max {
			cmd.c.Send(&Event{Command: JOIN, Params: []string{buffer}})
			buffer = ""
			continue
		}

		if len(buffer) == 0 {
			buffer = channels[i]
		} else {
			buffer += "," + channels[i]
		}

		if i == len(channels)-1 {
			cmd.c.Send(&Event{Command: JOIN, Params: []string{buffer}})
			return nil
		}
	}

	return nil
}

// JoinKey attempts to enter an IRC channel with a password.
func (cmd *Commands) JoinKey(channel, password string) error {
	if !IsValidChannel(channel) {
		return &ErrInvalidTarget{Target: channel}
	}

	cmd.c.Send(&Event{Command: JOIN, Params: []string{channel, password}})
	return nil
}

// Part leaves an IRC channel.
func (cmd *Commands) Part(channel, message string) error {
	if !IsValidChannel(channel) {
		return &ErrInvalidTarget{Target: channel}
	}

	cmd.c.Send(&Event{Command: JOIN, Params: []string{channel}})
	return nil
}

// PartMessage leaves an IRC channel with a specified leave message.
func (cmd *Commands) PartMessage(channel, message string) error {
	if !IsValidChannel(channel) {
		return &ErrInvalidTarget{Target: channel}
	}

	cmd.c.Send(&Event{Command: JOIN, Params: []string{channel}, Trailing: message})
	return nil
}

// SendCTCP sends a CTCP request to target. Note that this method uses
// PRIVMSG specifically.
func (cmd *Commands) SendCTCP(target, ctcpType, message string) error {
	out := encodeCTCPRaw(ctcpType, message)
	if out == "" {
		return errors.New("invalid CTCP")
	}

	return cmd.Message(target, out)
}

// SendCTCPf sends a CTCP request to target using a specific format. Note that
// this method uses PRIVMSG specifically.
func (cmd *Commands) SendCTCPf(target, ctcpType, format string, a ...interface{}) error {
	return cmd.SendCTCP(target, ctcpType, fmt.Sprintf(format, a...))
}

// SendCTCPReplyf sends a CTCP response to target using a specific format.
// Note that this method uses NOTICE specifically.
func (cmd *Commands) SendCTCPReplyf(target, ctcpType, format string, a ...interface{}) error {
	return cmd.SendCTCPReply(target, ctcpType, fmt.Sprintf(format, a...))
}

// SendCTCPReply sends a CTCP response to target. Note that this method uses
// NOTICE specifically.
func (cmd *Commands) SendCTCPReply(target, ctcpType, message string) error {
	out := encodeCTCPRaw(ctcpType, message)
	if out == "" {
		return errors.New("invalid CTCP")
	}

	return cmd.Notice(target, out)
}

// Message sends a PRIVMSG to target (either channel, service, or user).
func (cmd *Commands) Message(target, message string) error {
	if !IsValidNick(target) && !IsValidChannel(target) {
		return &ErrInvalidTarget{Target: target}
	}

	cmd.c.Send(&Event{Command: PRIVMSG, Params: []string{target}, Trailing: message})
	return nil
}

// Messagef sends a formated PRIVMSG to target (either channel, service, or
// user).
func (cmd *Commands) Messagef(target, format string, a ...interface{}) error {
	return cmd.Message(target, fmt.Sprintf(format, a...))
}

var ErrInvalidSource = errors.New("event has nil or invalid source address")

// Reply sends a reply to channel or user, based on where the supplied event
// originated from. See also ReplyTo().
func (cmd *Commands) Reply(event Event, message string) error {
	if event.Source == nil {
		return ErrInvalidSource
	}

	if len(event.Params) > 0 && IsValidChannel(event.Params[0]) {
		return cmd.Message(event.Params[0], message)
	}

	return cmd.Message(event.Source.Name, message)
}

// Replyf sends a reply to channel or user with a format string, based on
// where the supplied event originated from. See also ReplyTof().
func (cmd *Commands) Replyf(event Event, format string, a ...interface{}) error {
	return cmd.Reply(event, fmt.Sprintf(format, a...))
}

// ReplyTo sends a reply to a channel or user, based on where the supplied
// event originated from. ReplyTo(), when originating from a channel will
// default to replying with "<user>, <message>". See also Reply().
func (cmd *Commands) ReplyTo(event Event, message string) error {
	if event.Source == nil {
		return ErrInvalidSource
	}

	if len(event.Params) > 0 && IsValidChannel(event.Params[0]) {
		return cmd.Message(event.Params[0], event.Source.Name+", "+message)
	}

	return cmd.Message(event.Source.Name, message)
}

// ReplyTof sends a reply to a channel or user with a format string, based
// on where the supplied event originated from. ReplyTo(), when originating
// from a channel will default to replying with "<user>, <message>". See
// also Replyf().
func (cmd *Commands) ReplyTof(event Event, format string, a ...interface{}) error {
	return cmd.ReplyTo(event, fmt.Sprintf(format, a...))
}

// Action sends a PRIVMSG ACTION (/me) to target (either channel, service,
// or user).
func (cmd *Commands) Action(target, message string) error {
	if !IsValidNick(target) && !IsValidChannel(target) {
		return &ErrInvalidTarget{Target: target}
	}

	cmd.c.Send(&Event{
		Command:  PRIVMSG,
		Params:   []string{target},
		Trailing: fmt.Sprintf("\001ACTION %s\001", message),
	})
	return nil
}

// Actionf sends a formated PRIVMSG ACTION (/me) to target (either channel,
// service, or user).
func (cmd *Commands) Actionf(target, format string, a ...interface{}) error {
	return cmd.Action(target, fmt.Sprintf(format, a...))
}

// Notice sends a NOTICE to target (either channel, service, or user).
func (cmd *Commands) Notice(target, message string) error {
	if !IsValidNick(target) && !IsValidChannel(target) {
		return &ErrInvalidTarget{Target: target}
	}

	cmd.c.Send(&Event{Command: NOTICE, Params: []string{target}, Trailing: message})
	return nil
}

// Noticef sends a formated NOTICE to target (either channel, service, or
// user).
func (cmd *Commands) Noticef(target, format string, a ...interface{}) error {
	return cmd.Notice(target, fmt.Sprintf(format, a...))
}

// SendRaw sends a raw string back to the server, without carriage returns
// or newlines.
func (cmd *Commands) SendRaw(raw string) error {
	e := ParseEvent(raw)
	if e == nil {
		return errors.New("invalid event: " + raw)
	}

	cmd.c.Send(e)
	return nil
}

// SendRawf sends a formated string back to the server, without carriage
// returns or newlines.
func (cmd *Commands) SendRawf(format string, a ...interface{}) error {
	return cmd.SendRaw(fmt.Sprintf(format, a...))
}

// Topic sets the topic of channel to message. Does not verify the length
// of the topic.
func (cmd *Commands) Topic(channel, message string) {
	cmd.c.Send(&Event{Command: TOPIC, Params: []string{channel}, Trailing: message})
}

// Who sends a WHO query to the server, which will attempt WHOX by default.
// See http://faerion.sourceforge.net/doc/irc/whox.var for more details. This
// sends "%tcuhnr,2" per default. Do not use "1" as this will conflict with
// girc's builtin tracking functionality.
func (cmd *Commands) Who(target string) error {
	if !IsValidNick(target) && !IsValidChannel(target) && !IsValidUser(target) {
		return &ErrInvalidTarget{Target: target}
	}

	cmd.c.Send(&Event{Command: WHO, Params: []string{target, "%tcuhnr,2"}})
	return nil
}

// Whois sends a WHOIS query to the server, targeted at a specific user.
// as WHOIS is a bit slower, you may want to use WHO for brief user info.
func (cmd *Commands) Whois(nick string) error {
	if !IsValidNick(nick) {
		return &ErrInvalidTarget{Target: nick}
	}

	cmd.c.Send(&Event{Command: WHOIS, Params: []string{nick}})
	return nil
}

// Ping sends a PING query to the server, with a specific identifier that
// the server should respond with.
func (cmd *Commands) Ping(id string) {
	cmd.c.Send(&Event{Command: PING, Params: []string{id}})
}

// Pong sends a PONG query to the server, with an identifier which was
// received from a previous PING query received by the client.
func (cmd *Commands) Pong(id string) {
	cmd.c.Send(&Event{Command: PONG, Params: []string{id}})
}

// Oper sends a OPER authentication query to the server, with a username
// and password.
func (cmd *Commands) Oper(user, pass string) {
	cmd.c.Send(&Event{Command: OPER, Params: []string{user, pass}, Sensitive: true})
}

// Kick sends a KICK query to the server, attempting to kick nick from
// channel, with reason. If reason is blank, one will not be sent to the
// server.
func (cmd *Commands) Kick(channel, nick, reason string) error {
	if !IsValidChannel(channel) {
		return &ErrInvalidTarget{Target: channel}
	}

	if !IsValidNick(nick) {
		return &ErrInvalidTarget{Target: nick}
	}

	if reason != "" {
		cmd.c.Send(&Event{Command: KICK, Params: []string{channel, nick}, Trailing: reason})
		return nil
	}

	cmd.c.Send(&Event{Command: KICK, Params: []string{channel, nick}})
	return nil
}

// Invite sends a INVITE query to the server, to invite nick to channel.
func (cmd *Commands) Invite(channel, nick string) error {
	if !IsValidChannel(channel) {
		return &ErrInvalidTarget{Target: channel}
	}

	if !IsValidNick(nick) {
		return &ErrInvalidTarget{Target: nick}
	}

	cmd.c.Send(&Event{Command: INVITE, Params: []string{nick, channel}})
	return nil
}

// Away sends a AWAY query to the server, suggesting that the client is no
// longer active. If reason is blank, Client.Back() is called. Also see
// Client.Back().
func (cmd *Commands) Away(reason string) {
	if reason == "" {
		cmd.Back()
		return
	}

	cmd.c.Send(&Event{Command: AWAY, Params: []string{reason}})
}

// Back sends a AWAY query to the server, however the query is blank,
// suggesting that the client is active once again. Also see Client.Away().
func (cmd *Commands) Back() {
	cmd.c.Send(&Event{Command: AWAY})
}

// List sends a LIST query to the server, which will list channels and topics.
// Supports multiple channels at once, in hopes it will reduce extensive
// LIST queries to the server. Supply no channels to run a list against the
// entire server (warning, that may mean LOTS of channels!)
func (cmd *Commands) List(channels ...string) error {
	if len(channels) == 0 {
		cmd.c.Send(&Event{Command: LIST})
		return nil
	}

	// We can LIST multiple channels at once, however we need to ensure that
	// we are not exceeding the line length. (see maxLength)
	max := maxLength - len(JOIN) - 1

	var buffer string

	for i := 0; i < len(channels); i++ {
		if !IsValidChannel(channels[i]) {
			return &ErrInvalidTarget{Target: channels[i]}
		}

		if len(buffer+","+channels[i]) > max {
			cmd.c.Send(&Event{Command: LIST, Params: []string{buffer}})
			buffer = ""
			continue
		}

		if len(buffer) == 0 {
			buffer = channels[i]
		} else {
			buffer += "," + channels[i]
		}

		if i == len(channels)-1 {
			cmd.c.Send(&Event{Command: LIST, Params: []string{buffer}})
			return nil
		}
	}

	return nil
}

// Whowas sends a WHOWAS query to the server. amount is the amount of results
// you want back.
func (cmd *Commands) Whowas(nick string, amount int) error {
	if !IsValidNick(nick) {
		return &ErrInvalidTarget{Target: nick}
	}

	cmd.c.Send(&Event{Command: WHOWAS, Params: []string{nick, string(amount)}})
	return nil
}
<commit_msg>add missing documentation<commit_after>// Copyright (c) Liam Stanley <me@liamstanley.io>. All rights reserved. Use
// of this source code is governed by the MIT license that can be found in
// the LICENSE file.

package girc

import (
	"errors"
	"fmt"
)

// Commands holds a large list of useful methods to interact with the server,
// and wrappers for common events.
type Commands struct {
	c *Client
}

// Nick changes the client nickname.
func (cmd *Commands) Nick(name string) error {
	if !IsValidNick(name) {
		return &ErrInvalidTarget{Target: name}
	}

	cmd.c.Send(&Event{Command: NICK, Params: []string{name}})
	return nil
}

// Join attempts to enter a list of IRC channels, at bulk if possible to
// prevent sending extensive JOIN commands.
func (cmd *Commands) Join(channels ...string) error {
	// We can join multiple channels at once, however we need to ensure that
	// we are not exceeding the line length. (see maxLength)
	max := maxLength - len(JOIN) - 1

	var buffer string

	for i := 0; i < len(channels); i++ {
		if !IsValidChannel(channels[i]) {
			return &ErrInvalidTarget{Target: channels[i]}
		}

		if len(buffer+","+channels[i]) > max {
			cmd.c.Send(&Event{Command: JOIN, Params: []string{buffer}})
			buffer = ""
			continue
		}

		if len(buffer) == 0 {
			buffer = channels[i]
		} else {
			buffer += "," + channels[i]
		}

		if i == len(channels)-1 {
			cmd.c.Send(&Event{Command: JOIN, Params: []string{buffer}})
			return nil
		}
	}

	return nil
}

// JoinKey attempts to enter an IRC channel with a password.
func (cmd *Commands) JoinKey(channel, password string) error {
	if !IsValidChannel(channel) {
		return &ErrInvalidTarget{Target: channel}
	}

	cmd.c.Send(&Event{Command: JOIN, Params: []string{channel, password}})
	return nil
}

// Part leaves an IRC channel.
func (cmd *Commands) Part(channel, message string) error {
	if !IsValidChannel(channel) {
		return &ErrInvalidTarget{Target: channel}
	}

	cmd.c.Send(&Event{Command: JOIN, Params: []string{channel}})
	return nil
}

// PartMessage leaves an IRC channel with a specified leave message.
func (cmd *Commands) PartMessage(channel, message string) error {
	if !IsValidChannel(channel) {
		return &ErrInvalidTarget{Target: channel}
	}

	cmd.c.Send(&Event{Command: JOIN, Params: []string{channel}, Trailing: message})
	return nil
}

// SendCTCP sends a CTCP request to target. Note that this method uses
// PRIVMSG specifically.
func (cmd *Commands) SendCTCP(target, ctcpType, message string) error {
	out := encodeCTCPRaw(ctcpType, message)
	if out == "" {
		return errors.New("invalid CTCP")
	}

	return cmd.Message(target, out)
}

// SendCTCPf sends a CTCP request to target using a specific format. Note that
// this method uses PRIVMSG specifically.
func (cmd *Commands) SendCTCPf(target, ctcpType, format string, a ...interface{}) error {
	return cmd.SendCTCP(target, ctcpType, fmt.Sprintf(format, a...))
}

// SendCTCPReplyf sends a CTCP response to target using a specific format.
// Note that this method uses NOTICE specifically.
func (cmd *Commands) SendCTCPReplyf(target, ctcpType, format string, a ...interface{}) error {
	return cmd.SendCTCPReply(target, ctcpType, fmt.Sprintf(format, a...))
}

// SendCTCPReply sends a CTCP response to target. Note that this method uses
// NOTICE specifically.
func (cmd *Commands) SendCTCPReply(target, ctcpType, message string) error {
	out := encodeCTCPRaw(ctcpType, message)
	if out == "" {
		return errors.New("invalid CTCP")
	}

	return cmd.Notice(target, out)
}

// Message sends a PRIVMSG to target (either channel, service, or user).
func (cmd *Commands) Message(target, message string) error {
	if !IsValidNick(target) && !IsValidChannel(target) {
		return &ErrInvalidTarget{Target: target}
	}

	cmd.c.Send(&Event{Command: PRIVMSG, Params: []string{target}, Trailing: message})
	return nil
}

// Messagef sends a formated PRIVMSG to target (either channel, service, or
// user).
func (cmd *Commands) Messagef(target, format string, a ...interface{}) error {
	return cmd.Message(target, fmt.Sprintf(format, a...))
}

// ErrInvalidSource is returned when a method needs to know the origin of an
// event, however Event.Source is unknown (e.g. sent by the user, not the
// server.)
var ErrInvalidSource = errors.New("event has nil or invalid source address")

// Reply sends a reply to channel or user, based on where the supplied event
// originated from. See also ReplyTo().
func (cmd *Commands) Reply(event Event, message string) error {
	if event.Source == nil {
		return ErrInvalidSource
	}

	if len(event.Params) > 0 && IsValidChannel(event.Params[0]) {
		return cmd.Message(event.Params[0], message)
	}

	return cmd.Message(event.Source.Name, message)
}

// Replyf sends a reply to channel or user with a format string, based on
// where the supplied event originated from. See also ReplyTof().
func (cmd *Commands) Replyf(event Event, format string, a ...interface{}) error {
	return cmd.Reply(event, fmt.Sprintf(format, a...))
}

// ReplyTo sends a reply to a channel or user, based on where the supplied
// event originated from. ReplyTo(), when originating from a channel will
// default to replying with "<user>, <message>". See also Reply().
func (cmd *Commands) ReplyTo(event Event, message string) error {
	if event.Source == nil {
		return ErrInvalidSource
	}

	if len(event.Params) > 0 && IsValidChannel(event.Params[0]) {
		return cmd.Message(event.Params[0], event.Source.Name+", "+message)
	}

	return cmd.Message(event.Source.Name, message)
}

// ReplyTof sends a reply to a channel or user with a format string, based
// on where the supplied event originated from. ReplyTo(), when originating
// from a channel will default to replying with "<user>, <message>". See
// also Replyf().
func (cmd *Commands) ReplyTof(event Event, format string, a ...interface{}) error {
	return cmd.ReplyTo(event, fmt.Sprintf(format, a...))
}

// Action sends a PRIVMSG ACTION (/me) to target (either channel, service,
// or user).
func (cmd *Commands) Action(target, message string) error {
	if !IsValidNick(target) && !IsValidChannel(target) {
		return &ErrInvalidTarget{Target: target}
	}

	cmd.c.Send(&Event{
		Command:  PRIVMSG,
		Params:   []string{target},
		Trailing: fmt.Sprintf("\001ACTION %s\001", message),
	})
	return nil
}

// Actionf sends a formated PRIVMSG ACTION (/me) to target (either channel,
// service, or user).
func (cmd *Commands) Actionf(target, format string, a ...interface{}) error {
	return cmd.Action(target, fmt.Sprintf(format, a...))
}

// Notice sends a NOTICE to target (either channel, service, or user).
func (cmd *Commands) Notice(target, message string) error {
	if !IsValidNick(target) && !IsValidChannel(target) {
		return &ErrInvalidTarget{Target: target}
	}

	cmd.c.Send(&Event{Command: NOTICE, Params: []string{target}, Trailing: message})
	return nil
}

// Noticef sends a formated NOTICE to target (either channel, service, or
// user).
func (cmd *Commands) Noticef(target, format string, a ...interface{}) error {
	return cmd.Notice(target, fmt.Sprintf(format, a...))
}

// SendRaw sends a raw string back to the server, without carriage returns
// or newlines.
func (cmd *Commands) SendRaw(raw string) error {
	e := ParseEvent(raw)
	if e == nil {
		return errors.New("invalid event: " + raw)
	}

	cmd.c.Send(e)
	return nil
}

// SendRawf sends a formated string back to the server, without carriage
// returns or newlines.
func (cmd *Commands) SendRawf(format string, a ...interface{}) error {
	return cmd.SendRaw(fmt.Sprintf(format, a...))
}

// Topic sets the topic of channel to message. Does not verify the length
// of the topic.
func (cmd *Commands) Topic(channel, message string) {
	cmd.c.Send(&Event{Command: TOPIC, Params: []string{channel}, Trailing: message})
}

// Who sends a WHO query to the server, which will attempt WHOX by default.
// See http://faerion.sourceforge.net/doc/irc/whox.var for more details. This
// sends "%tcuhnr,2" per default. Do not use "1" as this will conflict with
// girc's builtin tracking functionality.
func (cmd *Commands) Who(target string) error {
	if !IsValidNick(target) && !IsValidChannel(target) && !IsValidUser(target) {
		return &ErrInvalidTarget{Target: target}
	}

	cmd.c.Send(&Event{Command: WHO, Params: []string{target, "%tcuhnr,2"}})
	return nil
}

// Whois sends a WHOIS query to the server, targeted at a specific user.
// as WHOIS is a bit slower, you may want to use WHO for brief user info.
func (cmd *Commands) Whois(nick string) error {
	if !IsValidNick(nick) {
		return &ErrInvalidTarget{Target: nick}
	}

	cmd.c.Send(&Event{Command: WHOIS, Params: []string{nick}})
	return nil
}

// Ping sends a PING query to the server, with a specific identifier that
// the server should respond with.
func (cmd *Commands) Ping(id string) {
	cmd.c.Send(&Event{Command: PING, Params: []string{id}})
}

// Pong sends a PONG query to the server, with an identifier which was
// received from a previous PING query received by the client.
func (cmd *Commands) Pong(id string) {
	cmd.c.Send(&Event{Command: PONG, Params: []string{id}})
}

// Oper sends a OPER authentication query to the server, with a username
// and password.
func (cmd *Commands) Oper(user, pass string) {
	cmd.c.Send(&Event{Command: OPER, Params: []string{user, pass}, Sensitive: true})
}

// Kick sends a KICK query to the server, attempting to kick nick from
// channel, with reason. If reason is blank, one will not be sent to the
// server.
func (cmd *Commands) Kick(channel, nick, reason string) error {
	if !IsValidChannel(channel) {
		return &ErrInvalidTarget{Target: channel}
	}

	if !IsValidNick(nick) {
		return &ErrInvalidTarget{Target: nick}
	}

	if reason != "" {
		cmd.c.Send(&Event{Command: KICK, Params: []string{channel, nick}, Trailing: reason})
		return nil
	}

	cmd.c.Send(&Event{Command: KICK, Params: []string{channel, nick}})
	return nil
}

// Invite sends a INVITE query to the server, to invite nick to channel.
func (cmd *Commands) Invite(channel, nick string) error {
	if !IsValidChannel(channel) {
		return &ErrInvalidTarget{Target: channel}
	}

	if !IsValidNick(nick) {
		return &ErrInvalidTarget{Target: nick}
	}

	cmd.c.Send(&Event{Command: INVITE, Params: []string{nick, channel}})
	return nil
}

// Away sends a AWAY query to the server, suggesting that the client is no
// longer active. If reason is blank, Client.Back() is called. Also see
// Client.Back().
func (cmd *Commands) Away(reason string) {
	if reason == "" {
		cmd.Back()
		return
	}

	cmd.c.Send(&Event{Command: AWAY, Params: []string{reason}})
}

// Back sends a AWAY query to the server, however the query is blank,
// suggesting that the client is active once again. Also see Client.Away().
func (cmd *Commands) Back() {
	cmd.c.Send(&Event{Command: AWAY})
}

// List sends a LIST query to the server, which will list channels and topics.
// Supports multiple channels at once, in hopes it will reduce extensive
// LIST queries to the server. Supply no channels to run a list against the
// entire server (warning, that may mean LOTS of channels!)
func (cmd *Commands) List(channels ...string) error {
	if len(channels) == 0 {
		cmd.c.Send(&Event{Command: LIST})
		return nil
	}

	// We can LIST multiple channels at once, however we need to ensure that
	// we are not exceeding the line length. (see maxLength)
	max := maxLength - len(JOIN) - 1

	var buffer string

	for i := 0; i < len(channels); i++ {
		if !IsValidChannel(channels[i]) {
			return &ErrInvalidTarget{Target: channels[i]}
		}

		if len(buffer+","+channels[i]) > max {
			cmd.c.Send(&Event{Command: LIST, Params: []string{buffer}})
			buffer = ""
			continue
		}

		if len(buffer) == 0 {
			buffer = channels[i]
		} else {
			buffer += "," + channels[i]
		}

		if i == len(channels)-1 {
			cmd.c.Send(&Event{Command: LIST, Params: []string{buffer}})
			return nil
		}
	}

	return nil
}

// Whowas sends a WHOWAS query to the server. amount is the amount of results
// you want back.
func (cmd *Commands) Whowas(nick string, amount int) error {
	if !IsValidNick(nick) {
		return &ErrInvalidTarget{Target: nick}
	}

	cmd.c.Send(&Event{Command: WHOWAS, Params: []string{nick, string(amount)}})
	return nil
}
<|endoftext|>


<commit_before>package rangestore

// generic interface for any Store
type Store interface {
	// lookup cluster
	ClusterLookup(string) ([]string, error)     // cluster
	KeyLookup(string, string) ([]string, error) // cluster and key

	// lookup reverse
	KeyReverseLookup(string) ([]string, error)                     // just a reverse lookup on a node
	KeyReverseLookupAttr(string, string) ([]string, error)         // reverse lookup where value and key are passed
	KeyReverseLookupHint(string, string, string) ([]string, error) // reverse lookup where value and key are passed with an hint
}
<commit_msg>changing the generic store interface args<commit_after>package rangestore

// generic interface for any Store
type Store interface {
	// lookup cluster
	ClusterLookup(string) (*[]string, error)     // cluster
	KeyLookup(string, string) (*[]string, error) // cluster and key

	// lookup reverse
	KeyReverseLookup(string) (*[]string, error)                     // just a reverse lookup on a node
	KeyReverseLookupAttr(string, string) (*[]string, error)         // reverse lookup where value and key are passed
	KeyReverseLookupHint(string, string, string) (*[]string, error) // reverse lookup where value and key are passed with an hint
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package roaring

import (
	"errors"
	"sync"

	"github.com/m3db/m3/src/m3ninx/postings"
	"github.com/m3db/m3/src/m3ninx/x"

	"github.com/RoaringBitmap/roaring"
)

var (
	errIntersectRoaringOnly  = errors.New("Intersect only supported between roaringDocId sets")
	errUnionRoaringOnly      = errors.New("Union only supported between roaringDocId sets")
	errDifferenceRoaringOnly = errors.New("Difference only supported between roaringDocId sets")
	errIteratorClosed        = errors.New("iterator has been closed")
)

// postingsList wraps a Roaring Bitmap with a mutex for thread safety.
type postingsList struct {
	sync.RWMutex
	bitmap *roaring.Bitmap
}

// NewPostingsList returns a new mutable postings list backed by a Roaring Bitmap.
func NewPostingsList() postings.MutableList {
	return &postingsList{
		bitmap: roaring.NewBitmap(),
	}
}

func (d *postingsList) Insert(i postings.ID) {
	d.Lock()
	d.bitmap.Add(uint32(i))
	d.Unlock()
}

func (d *postingsList) Intersect(other postings.List) error {
	o, ok := other.(*postingsList)
	if !ok {
		return errIntersectRoaringOnly
	}

	o.RLock()
	d.Lock()
	d.bitmap.And(o.bitmap)
	d.Unlock()
	o.RUnlock()
	return nil
}

func (d *postingsList) Difference(other postings.List) error {
	o, ok := other.(*postingsList)
	if !ok {
		return errDifferenceRoaringOnly
	}

	o.RLock()
	d.Lock()
	d.bitmap.AndNot(o.bitmap)
	d.Unlock()
	o.RUnlock()
	return nil
}

func (d *postingsList) Union(other postings.List) error {
	o, ok := other.(*postingsList)
	if !ok {
		return errUnionRoaringOnly
	}

	o.RLock()
	d.Lock()
	d.bitmap.Or(o.bitmap)
	d.Unlock()
	o.RUnlock()
	return nil
}

func (d *postingsList) AddRange(min, max postings.ID) {
	d.Lock()
	d.bitmap.AddRange(uint64(min), uint64(max))
	d.Unlock()
}

func (d *postingsList) AddIterator(iter postings.Iterator) error {
	safeIter := x.NewSafeCloser(iter)
	defer safeIter.Close()

	d.Lock()
	for iter.Next() {
		d.bitmap.Add(uint32(iter.Current()))
	}
	d.Unlock()

	if err := iter.Err(); err != nil {
		return err
	}

	return safeIter.Close()
}

func (d *postingsList) RemoveRange(min, max postings.ID) {
	d.Lock()
	d.bitmap.RemoveRange(uint64(min), uint64(max))
	d.Unlock()
}

func (d *postingsList) Reset() {
	d.Lock()
	d.bitmap.Clear()
	d.Unlock()
}

func (d *postingsList) Contains(i postings.ID) bool {
	d.RLock()
	contains := d.bitmap.Contains(uint32(i))
	d.RUnlock()
	return contains
}

func (d *postingsList) IsEmpty() bool {
	d.RLock()
	empty := d.bitmap.IsEmpty()
	d.RUnlock()
	return empty
}

func (d *postingsList) Max() (postings.ID, error) {
	d.RLock()
	if d.bitmap.IsEmpty() {
		d.RUnlock()
		return 0, postings.ErrEmptyList
	}
	max := d.bitmap.Maximum()
	d.RUnlock()
	return postings.ID(max), nil
}

func (d *postingsList) Min() (postings.ID, error) {
	d.RLock()
	if d.bitmap.IsEmpty() {
		d.RUnlock()
		return 0, postings.ErrEmptyList
	}
	min := d.bitmap.Minimum()
	d.RUnlock()
	return postings.ID(min), nil
}

func (d *postingsList) Len() int {
	d.RLock()
	l := d.bitmap.GetCardinality()
	d.RUnlock()
	return int(l)
}

func (d *postingsList) Iterator() postings.Iterator {
	return &roaringIterator{
		iter: d.bitmap.Iterator(),
	}
}

func (d *postingsList) Clone() postings.MutableList {
	d.RLock()
	// TODO: It's cheaper to Clone than to cache roaring bitmaps, see
	// `postings_list_bench_test.go`. Their internals don't allow for
	// pooling at the moment. We should address this when get a chance
	// (move to another implementation / address deficiencies).
	clone := d.bitmap.Clone()
	d.RUnlock()
	return &postingsList{
		bitmap: clone,
	}
}

func (d *postingsList) Equal(other postings.List) bool {
	if d.Len() != other.Len() {
		return false
	}

	o, ok := other.(*postingsList)
	if ok {
		return d.bitmap.Equals(o.bitmap)
	}

	iter := d.Iterator()
	otherIter := other.Iterator()

	for iter.Next() {
		if !otherIter.Next() {
			return false
		}
		if iter.Current() != otherIter.Current() {
			return false
		}
	}

	return true
}

type roaringIterator struct {
	iter    roaring.IntIterable
	current postings.ID
	closed  bool
}

func (it *roaringIterator) Current() postings.ID {
	return it.current
}

func (it *roaringIterator) Next() bool {
	if it.closed || !it.iter.HasNext() {
		return false
	}
	it.current = postings.ID(it.iter.Next())
	return true
}

func (it *roaringIterator) Err() error {
	return nil
}

func (it *roaringIterator) Close() error {
	if it.closed {
		return errIteratorClosed
	}
	it.closed = true
	return nil
}
<commit_msg>Eliminate locking from postings lists (#889)<commit_after>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package roaring

import (
	"errors"

	"github.com/m3db/m3/src/m3ninx/postings"
	"github.com/m3db/m3/src/m3ninx/x"

	"github.com/RoaringBitmap/roaring"
)

var (
	errIntersectRoaringOnly  = errors.New("Intersect only supported between roaringDocId sets")
	errUnionRoaringOnly      = errors.New("Union only supported between roaringDocId sets")
	errDifferenceRoaringOnly = errors.New("Difference only supported between roaringDocId sets")
	errIteratorClosed        = errors.New("iterator has been closed")
)

// postingsList wraps a Roaring Bitmap with the m3ninx pl API.
type postingsList struct {
	bitmap *roaring.Bitmap
}

// NewPostingsList returns a new mutable postings list backed by a Roaring Bitmap.
func NewPostingsList() postings.MutableList {
	return &postingsList{
		bitmap: roaring.NewBitmap(),
	}
}

func (d *postingsList) Insert(i postings.ID) {
	d.bitmap.Add(uint32(i))
}

func (d *postingsList) Intersect(other postings.List) error {
	o, ok := other.(*postingsList)
	if !ok {
		return errIntersectRoaringOnly
	}

	d.bitmap.And(o.bitmap)
	return nil
}

func (d *postingsList) Difference(other postings.List) error {
	o, ok := other.(*postingsList)
	if !ok {
		return errDifferenceRoaringOnly
	}

	d.bitmap.AndNot(o.bitmap)
	return nil
}

func (d *postingsList) Union(other postings.List) error {
	o, ok := other.(*postingsList)
	if !ok {
		return errUnionRoaringOnly
	}

	d.bitmap.Or(o.bitmap)
	return nil
}

func (d *postingsList) AddRange(min, max postings.ID) {
	d.bitmap.AddRange(uint64(min), uint64(max))
}

func (d *postingsList) AddIterator(iter postings.Iterator) error {
	safeIter := x.NewSafeCloser(iter)
	defer safeIter.Close()

	for iter.Next() {
		d.bitmap.Add(uint32(iter.Current()))
	}

	if err := iter.Err(); err != nil {
		return err
	}

	return safeIter.Close()
}

func (d *postingsList) RemoveRange(min, max postings.ID) {
	d.bitmap.RemoveRange(uint64(min), uint64(max))
}

func (d *postingsList) Reset() {
	d.bitmap.Clear()
}

func (d *postingsList) Contains(i postings.ID) bool {
	return d.bitmap.Contains(uint32(i))
}

func (d *postingsList) IsEmpty() bool {
	return d.bitmap.IsEmpty()
}

func (d *postingsList) Max() (postings.ID, error) {
	if d.bitmap.IsEmpty() {
		return 0, postings.ErrEmptyList
	}
	max := d.bitmap.Maximum()
	return postings.ID(max), nil
}

func (d *postingsList) Min() (postings.ID, error) {
	if d.bitmap.IsEmpty() {
		return 0, postings.ErrEmptyList
	}
	min := d.bitmap.Minimum()
	return postings.ID(min), nil
}

func (d *postingsList) Len() int {
	return int(d.bitmap.GetCardinality())
}

func (d *postingsList) Iterator() postings.Iterator {
	return &roaringIterator{
		iter: d.bitmap.Iterator(),
	}
}

func (d *postingsList) Clone() postings.MutableList {
	// TODO: It's cheaper to Clone than to cache roaring bitmaps, see
	// `postings_list_bench_test.go`. Their internals don't allow for
	// pooling at the moment. We should address this when get a chance
	// (move to another implementation / address deficiencies).
	return &postingsList{
		bitmap: d.bitmap.Clone(),
	}
}

func (d *postingsList) Equal(other postings.List) bool {
	if d.Len() != other.Len() {
		return false
	}

	o, ok := other.(*postingsList)
	if ok {
		return d.bitmap.Equals(o.bitmap)
	}

	iter := d.Iterator()
	otherIter := other.Iterator()

	for iter.Next() {
		if !otherIter.Next() {
			return false
		}
		if iter.Current() != otherIter.Current() {
			return false
		}
	}

	return true
}

type roaringIterator struct {
	iter    roaring.IntIterable
	current postings.ID
	closed  bool
}

func (it *roaringIterator) Current() postings.ID {
	return it.current
}

func (it *roaringIterator) Next() bool {
	if it.closed || !it.iter.HasNext() {
		return false
	}
	it.current = postings.ID(it.iter.Next())
	return true
}

func (it *roaringIterator) Err() error {
	return nil
}

func (it *roaringIterator) Close() error {
	if it.closed {
		return errIteratorClosed
	}
	it.closed = true
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2012 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package runtime_test

import (
	. "runtime"
	"strings"
	"sync"
	"testing"
	"time"
)

// TestStackMem measures per-thread stack segment cache behavior.
// The test consumed up to 500MB in the past.
func TestStackMem(t *testing.T) {
	const (
		BatchSize      = 32
		BatchCount     = 256
		ArraySize      = 1024
		RecursionDepth = 128
	)
	if testing.Short() {
		return
	}
	defer GOMAXPROCS(GOMAXPROCS(BatchSize))
	s0 := new(MemStats)
	ReadMemStats(s0)
	for b := 0; b < BatchCount; b++ {
		c := make(chan bool, BatchSize)
		for i := 0; i < BatchSize; i++ {
			go func() {
				var f func(k int, a [ArraySize]byte)
				f = func(k int, a [ArraySize]byte) {
					if k == 0 {
						time.Sleep(time.Millisecond)
						return
					}
					f(k-1, a)
				}
				f(RecursionDepth, [ArraySize]byte{})
				c <- true
			}()
		}
		for i := 0; i < BatchSize; i++ {
			<-c
		}

		// The goroutines have signaled via c that they are ready to exit.
		// Give them a chance to exit by sleeping. If we don't wait, we
		// might not reuse them on the next batch.
		time.Sleep(10 * time.Millisecond)
	}
	s1 := new(MemStats)
	ReadMemStats(s1)
	consumed := int64(s1.StackSys - s0.StackSys)
	t.Logf("Consumed %vMB for stack mem", consumed>>20)
	estimate := int64(8 * BatchSize * ArraySize * RecursionDepth) // 8 is to reduce flakiness.
	if consumed > estimate {
		t.Fatalf("Stack mem: want %v, got %v", estimate, consumed)
	}
	// Due to broken stack memory accounting (http://golang.org/issue/7468),
	// StackInuse can decrease during function execution, so we cast the values to int64.
	inuse := int64(s1.StackInuse) - int64(s0.StackInuse)
	t.Logf("Inuse %vMB for stack mem", inuse>>20)
	if inuse > 4<<20 {
		t.Fatalf("Stack inuse: want %v, got %v", 4<<20, inuse)
	}
}

// Test stack growing in different contexts.
func TestStackGrowth(t *testing.T) {
	switch GOARCH {
	case "386", "arm":
		t.Skipf("skipping test on %q; see issue 8083", GOARCH)
	}
	t.Parallel()
	var wg sync.WaitGroup

	// in a normal goroutine
	wg.Add(1)
	go func() {
		defer wg.Done()
		growStack()
	}()
	wg.Wait()

	// in locked goroutine
	wg.Add(1)
	go func() {
		defer wg.Done()
		LockOSThread()
		growStack()
		UnlockOSThread()
	}()
	wg.Wait()

	// in finalizer
	wg.Add(1)
	go func() {
		defer wg.Done()
		done := make(chan bool)
		go func() {
			s := new(string)
			SetFinalizer(s, func(ss *string) {
				growStack()
				done <- true
			})
			s = nil
			done <- true
		}()
		<-done
		GC()
		select {
		case <-done:
		case <-time.After(20 * time.Second):
			t.Fatal("finalizer did not run")
		}
	}()
	wg.Wait()
}

// ... and in init
//func init() {
//	growStack()
//}

func growStack() {
	n := 1 << 10
	if testing.Short() {
		n = 1 << 8
	}
	for i := 0; i < n; i++ {
		x := 0
		growStackIter(&x, i)
		if x != i+1 {
			panic("stack is corrupted")
		}
	}
	GC()
}

// This function is not an anonymous func, so that the compiler can do escape
// analysis and place x on stack (and subsequently stack growth update the pointer).
func growStackIter(p *int, n int) {
	if n == 0 {
		*p = n + 1
		GC()
		return
	}
	*p = n + 1
	x := 0
	growStackIter(&x, n-1)
	if x != n {
		panic("stack is corrupted")
	}
}

func TestStackGrowthCallback(t *testing.T) {
	t.Parallel()
	var wg sync.WaitGroup

	// test stack growth at chan op
	wg.Add(1)
	go func() {
		defer wg.Done()
		c := make(chan int, 1)
		growStackWithCallback(func() {
			c <- 1
			<-c
		})
	}()

	// test stack growth at map op
	wg.Add(1)
	go func() {
		defer wg.Done()
		m := make(map[int]int)
		growStackWithCallback(func() {
			_, _ = m[1]
			m[1] = 1
		})
	}()

	// test stack growth at goroutine creation
	wg.Add(1)
	go func() {
		defer wg.Done()
		growStackWithCallback(func() {
			done := make(chan bool)
			go func() {
				done <- true
			}()
			<-done
		})
	}()

	wg.Wait()
}

func growStackWithCallback(cb func()) {
	var f func(n int)
	f = func(n int) {
		if n == 0 {
			cb()
			return
		}
		f(n - 1)
	}
	for i := 0; i < 1<<10; i++ {
		f(i)
	}
}

// TestDeferPtrs tests the adjustment of Defer's argument pointers (p aka &y)
// during a stack copy.
func set(p *int, x int) {
	*p = x
}
func TestDeferPtrs(t *testing.T) {
	var y int

	defer func() {
		if y != 42 {
			t.Errorf("defer's stack references were not adjusted appropriately")
		}
	}()
	defer set(&y, 42)
	growStack()
}

type bigBuf [4 * 1024]byte

// TestDeferPtrsGoexit is like TestDeferPtrs but exercises the possibility that the
// stack grows as part of starting the deferred function. It calls Goexit at various
// stack depths, forcing the deferred function (with >4kB of args) to be run at
// the bottom of the stack. The goal is to find a stack depth less than 4kB from
// the end of the stack. Each trial runs in a different goroutine so that an earlier
// stack growth does not invalidate a later attempt.
func TestDeferPtrsGoexit(t *testing.T) {
	for i := 0; i < 100; i++ {
		c := make(chan int, 1)
		go testDeferPtrsGoexit(c, i)
		if n := <-c; n != 42 {
			t.Fatalf("defer's stack references were not adjusted appropriately (i=%d n=%d)", i, n)
		}
	}
}

func testDeferPtrsGoexit(c chan int, i int) {
	var y int
	defer func() {
		c <- y
	}()
	defer setBig(&y, 42, bigBuf{})
	useStackAndCall(i, Goexit)
}

func setBig(p *int, x int, b bigBuf) {
	*p = x
}

// TestDeferPtrsPanic is like TestDeferPtrsGoexit, but it's using panic instead
// of Goexit to run the Defers. Those two are different execution paths
// in the runtime.
func TestDeferPtrsPanic(t *testing.T) {
	for i := 0; i < 100; i++ {
		c := make(chan int, 1)
		go testDeferPtrsGoexit(c, i)
		if n := <-c; n != 42 {
			t.Fatalf("defer's stack references were not adjusted appropriately (i=%d n=%d)", i, n)
		}
	}
}

func testDeferPtrsPanic(c chan int, i int) {
	var y int
	defer func() {
		if recover() == nil {
			c <- -1
			return
		}
		c <- y
	}()
	defer setBig(&y, 42, bigBuf{})
	useStackAndCall(i, func() { panic(1) })
}

// TestPanicUseStack checks that a chain of Panic structs on the stack are
// updated correctly if the stack grows during the deferred execution that
// happens as a result of the panic.
func TestPanicUseStack(t *testing.T) {
	pc := make([]uintptr, 10000)
	defer func() {
		recover()
		Callers(0, pc) // force stack walk
		useStackAndCall(100, func() {
			defer func() {
				recover()
				Callers(0, pc) // force stack walk
				useStackAndCall(200, func() {
					defer func() {
						recover()
						Callers(0, pc) // force stack walk
					}()
					panic(3)
				})
			}()
			panic(2)
		})
	}()
	panic(1)
}

// use about n KB of stack and call f
func useStackAndCall(n int, f func()) {
	if n == 0 {
		f()
		return
	}
	var b [1024]byte // makes frame about 1KB
	useStackAndCall(n-1+int(b[99]), f)
}

func useStack(n int) {
	useStackAndCall(n, func() {})
}

func growing(c chan int, done chan struct{}) {
	for n := range c {
		useStack(n)
		done <- struct{}{}
	}
	done <- struct{}{}
}

func TestStackCache(t *testing.T) {
	// Allocate a bunch of goroutines and grow their stacks.
	// Repeat a few times to test the stack cache.
	const (
		R = 4
		G = 200
		S = 5
	)
	for i := 0; i < R; i++ {
		var reqchans [G]chan int
		done := make(chan struct{})
		for j := 0; j < G; j++ {
			reqchans[j] = make(chan int)
			go growing(reqchans[j], done)
		}
		for s := 0; s < S; s++ {
			for j := 0; j < G; j++ {
				reqchans[j] <- 1 << uint(s)
			}
			for j := 0; j < G; j++ {
				<-done
			}
		}
		for j := 0; j < G; j++ {
			close(reqchans[j])
		}
		for j := 0; j < G; j++ {
			<-done
		}
	}
}

func TestStackOutput(t *testing.T) {
	b := make([]byte, 1024)
	stk := string(b[:Stack(b, false)])
	if !strings.HasPrefix(stk, "goroutine ") {
		t.Errorf("Stack (len %d):\n%s", len(stk), stk)
		t.Errorf("Stack output should begin with \"goroutine \"")
	}
}

func TestStackAllOutput(t *testing.T) {
	b := make([]byte, 1024)
	stk := string(b[:Stack(b, true)])
	if !strings.HasPrefix(stk, "goroutine ") {
		t.Errorf("Stack (len %d):\n%s", len(stk), stk)
		t.Errorf("Stack output should begin with \"goroutine \"")
	}
}

func TestStackPanic(t *testing.T) {
	// Test that stack copying copies panics correctly.  This is difficult
	// to test because it is very unlikely that the stack will be copied
	// in the middle of gopanic.  But it can happen.
	// To make this test effective, edit panic.go:gopanic and uncomment
	// the GC() call just before freedefer(d).
	defer func() {
		if x := recover(); x == nil {
			t.Errorf("recover failed")
		}
	}()
	useStack(32)
	panic("test panic")
}
<commit_msg>runtime: reenable TestStackGrowth on 32-bit systems<commit_after>// Copyright 2012 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package runtime_test

import (
	. "runtime"
	"strings"
	"sync"
	"testing"
	"time"
)

// TestStackMem measures per-thread stack segment cache behavior.
// The test consumed up to 500MB in the past.
func TestStackMem(t *testing.T) {
	const (
		BatchSize      = 32
		BatchCount     = 256
		ArraySize      = 1024
		RecursionDepth = 128
	)
	if testing.Short() {
		return
	}
	defer GOMAXPROCS(GOMAXPROCS(BatchSize))
	s0 := new(MemStats)
	ReadMemStats(s0)
	for b := 0; b < BatchCount; b++ {
		c := make(chan bool, BatchSize)
		for i := 0; i < BatchSize; i++ {
			go func() {
				var f func(k int, a [ArraySize]byte)
				f = func(k int, a [ArraySize]byte) {
					if k == 0 {
						time.Sleep(time.Millisecond)
						return
					}
					f(k-1, a)
				}
				f(RecursionDepth, [ArraySize]byte{})
				c <- true
			}()
		}
		for i := 0; i < BatchSize; i++ {
			<-c
		}

		// The goroutines have signaled via c that they are ready to exit.
		// Give them a chance to exit by sleeping. If we don't wait, we
		// might not reuse them on the next batch.
		time.Sleep(10 * time.Millisecond)
	}
	s1 := new(MemStats)
	ReadMemStats(s1)
	consumed := int64(s1.StackSys - s0.StackSys)
	t.Logf("Consumed %vMB for stack mem", consumed>>20)
	estimate := int64(8 * BatchSize * ArraySize * RecursionDepth) // 8 is to reduce flakiness.
	if consumed > estimate {
		t.Fatalf("Stack mem: want %v, got %v", estimate, consumed)
	}
	// Due to broken stack memory accounting (http://golang.org/issue/7468),
	// StackInuse can decrease during function execution, so we cast the values to int64.
	inuse := int64(s1.StackInuse) - int64(s0.StackInuse)
	t.Logf("Inuse %vMB for stack mem", inuse>>20)
	if inuse > 4<<20 {
		t.Fatalf("Stack inuse: want %v, got %v", 4<<20, inuse)
	}
}

// Test stack growing in different contexts.
func TestStackGrowth(t *testing.T) {
	t.Parallel()
	var wg sync.WaitGroup

	// in a normal goroutine
	wg.Add(1)
	go func() {
		defer wg.Done()
		growStack()
	}()
	wg.Wait()

	// in locked goroutine
	wg.Add(1)
	go func() {
		defer wg.Done()
		LockOSThread()
		growStack()
		UnlockOSThread()
	}()
	wg.Wait()

	// in finalizer
	wg.Add(1)
	go func() {
		defer wg.Done()
		done := make(chan bool)
		go func() {
			s := new(string)
			SetFinalizer(s, func(ss *string) {
				growStack()
				done <- true
			})
			s = nil
			done <- true
		}()
		<-done
		GC()
		select {
		case <-done:
		case <-time.After(20 * time.Second):
			t.Fatal("finalizer did not run")
		}
	}()
	wg.Wait()
}

// ... and in init
//func init() {
//	growStack()
//}

func growStack() {
	n := 1 << 10
	if testing.Short() {
		n = 1 << 8
	}
	for i := 0; i < n; i++ {
		x := 0
		growStackIter(&x, i)
		if x != i+1 {
			panic("stack is corrupted")
		}
	}
	GC()
}

// This function is not an anonymous func, so that the compiler can do escape
// analysis and place x on stack (and subsequently stack growth update the pointer).
func growStackIter(p *int, n int) {
	if n == 0 {
		*p = n + 1
		GC()
		return
	}
	*p = n + 1
	x := 0
	growStackIter(&x, n-1)
	if x != n {
		panic("stack is corrupted")
	}
}

func TestStackGrowthCallback(t *testing.T) {
	t.Parallel()
	var wg sync.WaitGroup

	// test stack growth at chan op
	wg.Add(1)
	go func() {
		defer wg.Done()
		c := make(chan int, 1)
		growStackWithCallback(func() {
			c <- 1
			<-c
		})
	}()

	// test stack growth at map op
	wg.Add(1)
	go func() {
		defer wg.Done()
		m := make(map[int]int)
		growStackWithCallback(func() {
			_, _ = m[1]
			m[1] = 1
		})
	}()

	// test stack growth at goroutine creation
	wg.Add(1)
	go func() {
		defer wg.Done()
		growStackWithCallback(func() {
			done := make(chan bool)
			go func() {
				done <- true
			}()
			<-done
		})
	}()

	wg.Wait()
}

func growStackWithCallback(cb func()) {
	var f func(n int)
	f = func(n int) {
		if n == 0 {
			cb()
			return
		}
		f(n - 1)
	}
	for i := 0; i < 1<<10; i++ {
		f(i)
	}
}

// TestDeferPtrs tests the adjustment of Defer's argument pointers (p aka &y)
// during a stack copy.
func set(p *int, x int) {
	*p = x
}
func TestDeferPtrs(t *testing.T) {
	var y int

	defer func() {
		if y != 42 {
			t.Errorf("defer's stack references were not adjusted appropriately")
		}
	}()
	defer set(&y, 42)
	growStack()
}

type bigBuf [4 * 1024]byte

// TestDeferPtrsGoexit is like TestDeferPtrs but exercises the possibility that the
// stack grows as part of starting the deferred function. It calls Goexit at various
// stack depths, forcing the deferred function (with >4kB of args) to be run at
// the bottom of the stack. The goal is to find a stack depth less than 4kB from
// the end of the stack. Each trial runs in a different goroutine so that an earlier
// stack growth does not invalidate a later attempt.
func TestDeferPtrsGoexit(t *testing.T) {
	for i := 0; i < 100; i++ {
		c := make(chan int, 1)
		go testDeferPtrsGoexit(c, i)
		if n := <-c; n != 42 {
			t.Fatalf("defer's stack references were not adjusted appropriately (i=%d n=%d)", i, n)
		}
	}
}

func testDeferPtrsGoexit(c chan int, i int) {
	var y int
	defer func() {
		c <- y
	}()
	defer setBig(&y, 42, bigBuf{})
	useStackAndCall(i, Goexit)
}

func setBig(p *int, x int, b bigBuf) {
	*p = x
}

// TestDeferPtrsPanic is like TestDeferPtrsGoexit, but it's using panic instead
// of Goexit to run the Defers. Those two are different execution paths
// in the runtime.
func TestDeferPtrsPanic(t *testing.T) {
	for i := 0; i < 100; i++ {
		c := make(chan int, 1)
		go testDeferPtrsGoexit(c, i)
		if n := <-c; n != 42 {
			t.Fatalf("defer's stack references were not adjusted appropriately (i=%d n=%d)", i, n)
		}
	}
}

func testDeferPtrsPanic(c chan int, i int) {
	var y int
	defer func() {
		if recover() == nil {
			c <- -1
			return
		}
		c <- y
	}()
	defer setBig(&y, 42, bigBuf{})
	useStackAndCall(i, func() { panic(1) })
}

// TestPanicUseStack checks that a chain of Panic structs on the stack are
// updated correctly if the stack grows during the deferred execution that
// happens as a result of the panic.
func TestPanicUseStack(t *testing.T) {
	pc := make([]uintptr, 10000)
	defer func() {
		recover()
		Callers(0, pc) // force stack walk
		useStackAndCall(100, func() {
			defer func() {
				recover()
				Callers(0, pc) // force stack walk
				useStackAndCall(200, func() {
					defer func() {
						recover()
						Callers(0, pc) // force stack walk
					}()
					panic(3)
				})
			}()
			panic(2)
		})
	}()
	panic(1)
}

// use about n KB of stack and call f
func useStackAndCall(n int, f func()) {
	if n == 0 {
		f()
		return
	}
	var b [1024]byte // makes frame about 1KB
	useStackAndCall(n-1+int(b[99]), f)
}

func useStack(n int) {
	useStackAndCall(n, func() {})
}

func growing(c chan int, done chan struct{}) {
	for n := range c {
		useStack(n)
		done <- struct{}{}
	}
	done <- struct{}{}
}

func TestStackCache(t *testing.T) {
	// Allocate a bunch of goroutines and grow their stacks.
	// Repeat a few times to test the stack cache.
	const (
		R = 4
		G = 200
		S = 5
	)
	for i := 0; i < R; i++ {
		var reqchans [G]chan int
		done := make(chan struct{})
		for j := 0; j < G; j++ {
			reqchans[j] = make(chan int)
			go growing(reqchans[j], done)
		}
		for s := 0; s < S; s++ {
			for j := 0; j < G; j++ {
				reqchans[j] <- 1 << uint(s)
			}
			for j := 0; j < G; j++ {
				<-done
			}
		}
		for j := 0; j < G; j++ {
			close(reqchans[j])
		}
		for j := 0; j < G; j++ {
			<-done
		}
	}
}

func TestStackOutput(t *testing.T) {
	b := make([]byte, 1024)
	stk := string(b[:Stack(b, false)])
	if !strings.HasPrefix(stk, "goroutine ") {
		t.Errorf("Stack (len %d):\n%s", len(stk), stk)
		t.Errorf("Stack output should begin with \"goroutine \"")
	}
}

func TestStackAllOutput(t *testing.T) {
	b := make([]byte, 1024)
	stk := string(b[:Stack(b, true)])
	if !strings.HasPrefix(stk, "goroutine ") {
		t.Errorf("Stack (len %d):\n%s", len(stk), stk)
		t.Errorf("Stack output should begin with \"goroutine \"")
	}
}

func TestStackPanic(t *testing.T) {
	// Test that stack copying copies panics correctly.  This is difficult
	// to test because it is very unlikely that the stack will be copied
	// in the middle of gopanic.  But it can happen.
	// To make this test effective, edit panic.go:gopanic and uncomment
	// the GC() call just before freedefer(d).
	defer func() {
		if x := recover(); x == nil {
			t.Errorf("recover failed")
		}
	}()
	useStack(32)
	panic("test panic")
}
<|endoftext|>


<commit_before>package cmds

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"

	"github.com/pachyderm/pachyderm/src/client/pkg/config"
	"github.com/pachyderm/pachyderm/src/server/pkg/cmdutil"

	"github.com/gogo/protobuf/jsonpb"
	"github.com/spf13/cobra"
)

func readContext() (*config.Context, error) {
	var buf bytes.Buffer
	var decoder *json.Decoder
	var result config.Context

	contextReader := io.TeeReader(os.Stdin, &buf)
	fmt.Println("Reading from stdin.")
	decoder = json.NewDecoder(contextReader)

	if err := jsonpb.UnmarshalNext(decoder, &result); err != nil {
		if err == io.EOF {
			return nil, err
		}
		return nil, fmt.Errorf("malformed context: %s", err)
	}
	return &result, nil
}

// Cmds returns a slice containing admin commands.
func Cmds() []*cobra.Command {
	marshaller := &jsonpb.Marshaler{
		Indent:   "  ",
		OrigName: true,
	}

	var commands []*cobra.Command

	getMetrics := &cobra.Command{
		Short: "Gets whether metrics are enabled.",
		Long:  "Gets whether metrics are enabled.",
		Run: cmdutil.RunFixedArgs(0, func(args []string) (retErr error) {
			cfg, err := config.Read()
			if err != nil {
				return err
			}
			fmt.Printf("%v\n", !cfg.V2.NoMetrics)
			return nil
		}),
	}
	commands = append(commands, cmdutil.CreateAlias(getMetrics, "config get metrics"))

	setMetrics := &cobra.Command{
		Short: "Sets whether metrics are enabled.",
		Long:  "Sets whether metrics are enabled.",
		Run: cmdutil.RunFixedArgs(1, func(args []string) (retErr error) {
			metrics := true
			if args[0] == "false" {
				metrics = false
			} else if args[0] != "true" {
				return errors.New("invalid argument; use either `true` or `false`")
			}

			cfg, err := config.Read()
			if err != nil {
				return err
			}

			cfg.V2.NoMetrics = !metrics
			return cfg.Write()
		}),
	}
	commands = append(commands, cmdutil.CreateAlias(setMetrics, "config set metrics"))

	useContext := &cobra.Command{
		Short: "Sets the currently active context.",
		Long:  "Sets the currently active context.",
		Run: cmdutil.RunFixedArgs(1, func(args []string) (retErr error) {
			cfg, err := config.Read()
			if err != nil {
				return err
			}
			if _, ok := cfg.V2.Contexts[args[0]]; !ok {
				return fmt.Errorf("context does not exist: %s", args[0])
			}
			cfg.V2.ActiveContext = args[0]
			return cfg.Write()
		}),
	}
	commands = append(commands, cmdutil.CreateAlias(useContext, "config use context"))

	var active bool
	getContext := &cobra.Command{
		Short: "Gets a context.",
		Long:  "Gets the config of a context by its name, or the current active one.",
		Run: cmdutil.RunBoundedArgs(0, 1, func(args []string) (retErr error) {
			if active && len(args) == 1 {
				return errors.New("cannot get both the active context, and a specific one by its name")
			}

			cfg, err := config.Read()
			if err != nil {
				return err
			}

			var name string
			if active {
				name = cfg.V2.ActiveContext
			} else {
				name = args[0]
			}

			context, ok := cfg.V2.Contexts[name]
			if !ok {
				return fmt.Errorf("context does not exist: %s", name)
			}

			return marshaller.Marshal(os.Stdout, context)
		}),
	}
	getContext.Flags().BoolVar(&active, "active", false, "Get the active context.")
	commands = append(commands, cmdutil.CreateAlias(getContext, "config get context"))

	setContext := &cobra.Command{
		Short: "Set a context.",
		Long:  "Set a context config from a given name and JSON stdin.",
		Run: cmdutil.RunFixedArgs(1, func(args []string) (retErr error) {
			cfg, err := config.Read()
			if err != nil {
				return err
			}

			context, err := readContext()
			if err != nil {
				return err
			}

			cfg.V2.Contexts[args[0]] = context
			return cfg.Write()
		}),
	}
	commands = append(commands, cmdutil.CreateAlias(setContext, "config set context"))

	deleteContext := &cobra.Command{
		Short: "Deletes a context.",
		Long:  "Deletes a context.",
		Run: cmdutil.RunFixedArgs(1, func(args []string) (retErr error) {
			cfg, err := config.Read()
			if err != nil {
				return err
			}
			if _, ok := cfg.V2.Contexts[args[0]]; !ok {
				return fmt.Errorf("context does not exist: %s", args[0])
			}
			if cfg.V2.ActiveContext == args[0] {
				return errors.New("cannot delete an active context")
			}
			delete(cfg.V2.Contexts, args[0])
			return cfg.Write()
		}),
	}
	commands = append(commands, cmdutil.CreateAlias(deleteContext, "config delete context"))

	return commands
}
<commit_msg>Command for listing contexts<commit_after>package cmds

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"

	"github.com/pachyderm/pachyderm/src/client/pkg/config"
	"github.com/pachyderm/pachyderm/src/server/pkg/cmdutil"

	"github.com/gogo/protobuf/jsonpb"
	"github.com/spf13/cobra"
)

func readContext() (*config.Context, error) {
	var buf bytes.Buffer
	var decoder *json.Decoder
	var result config.Context

	contextReader := io.TeeReader(os.Stdin, &buf)
	fmt.Println("Reading from stdin.")
	decoder = json.NewDecoder(contextReader)

	if err := jsonpb.UnmarshalNext(decoder, &result); err != nil {
		if err == io.EOF {
			return nil, err
		}
		return nil, fmt.Errorf("malformed context: %s", err)
	}
	return &result, nil
}

// Cmds returns a slice containing admin commands.
func Cmds() []*cobra.Command {
	marshaller := &jsonpb.Marshaler{
		Indent:   "  ",
		OrigName: true,
	}

	var commands []*cobra.Command

	getMetrics := &cobra.Command{
		Short: "Gets whether metrics are enabled.",
		Long:  "Gets whether metrics are enabled.",
		Run: cmdutil.RunFixedArgs(0, func(args []string) (retErr error) {
			cfg, err := config.Read()
			if err != nil {
				return err
			}
			fmt.Printf("%v\n", !cfg.V2.NoMetrics)
			return nil
		}),
	}
	commands = append(commands, cmdutil.CreateAlias(getMetrics, "config get metrics"))

	setMetrics := &cobra.Command{
		Short: "Sets whether metrics are enabled.",
		Long:  "Sets whether metrics are enabled.",
		Run: cmdutil.RunFixedArgs(1, func(args []string) (retErr error) {
			metrics := true
			if args[0] == "false" {
				metrics = false
			} else if args[0] != "true" {
				return errors.New("invalid argument; use either `true` or `false`")
			}

			cfg, err := config.Read()
			if err != nil {
				return err
			}

			cfg.V2.NoMetrics = !metrics
			return cfg.Write()
		}),
	}
	commands = append(commands, cmdutil.CreateAlias(setMetrics, "config set metrics"))

	useContext := &cobra.Command{
		Short: "Sets the currently active context.",
		Long:  "Sets the currently active context.",
		Run: cmdutil.RunFixedArgs(1, func(args []string) (retErr error) {
			cfg, err := config.Read()
			if err != nil {
				return err
			}
			if _, ok := cfg.V2.Contexts[args[0]]; !ok {
				return fmt.Errorf("context does not exist: %s", args[0])
			}
			cfg.V2.ActiveContext = args[0]
			return cfg.Write()
		}),
	}
	commands = append(commands, cmdutil.CreateAlias(useContext, "config use context"))

	var active bool
	getContext := &cobra.Command{
		Short: "Gets a context.",
		Long:  "Gets the config of a context by its name, or the current active one.",
		Run: cmdutil.RunBoundedArgs(0, 1, func(args []string) (retErr error) {
			if active && len(args) == 1 {
				return errors.New("cannot get both the active context, and a specific one by its name")
			}

			cfg, err := config.Read()
			if err != nil {
				return err
			}

			var name string
			if active {
				name = cfg.V2.ActiveContext
			} else {
				name = args[0]
			}

			context, ok := cfg.V2.Contexts[name]
			if !ok {
				return fmt.Errorf("context does not exist: %s", name)
			}

			return marshaller.Marshal(os.Stdout, context)
		}),
	}
	getContext.Flags().BoolVar(&active, "active", false, "Get the active context.")
	commands = append(commands, cmdutil.CreateAlias(getContext, "config get context"))

	setContext := &cobra.Command{
		Short: "Set a context.",
		Long:  "Set a context config from a given name and JSON stdin.",
		Run: cmdutil.RunFixedArgs(1, func(args []string) (retErr error) {
			cfg, err := config.Read()
			if err != nil {
				return err
			}

			context, err := readContext()
			if err != nil {
				return err
			}

			cfg.V2.Contexts[args[0]] = context
			return cfg.Write()
		}),
	}
	commands = append(commands, cmdutil.CreateAlias(setContext, "config set context"))

	deleteContext := &cobra.Command{
		Short: "Deletes a context.",
		Long:  "Deletes a context.",
		Run: cmdutil.RunFixedArgs(1, func(args []string) (retErr error) {
			cfg, err := config.Read()
			if err != nil {
				return err
			}
			if _, ok := cfg.V2.Contexts[args[0]]; !ok {
				return fmt.Errorf("context does not exist: %s", args[0])
			}
			if cfg.V2.ActiveContext == args[0] {
				return errors.New("cannot delete an active context")
			}
			delete(cfg.V2.Contexts, args[0])
			return cfg.Write()
		}),
	}
	commands = append(commands, cmdutil.CreateAlias(deleteContext, "config delete context"))

	listContext := &cobra.Command{
		Short: "Lists contexts.",
		Long:  "Lists contexts.",
		Run: cmdutil.RunFixedArgs(1, func(args []string) (retErr error) {
			cfg, err := config.Read()
			if err != nil {
				return err
			}
			for name, _ := range cfg.V2.Contexts {
				if name == cfg.V2.ActiveContext {
					fmt.Printf("* %s\n", name)
				} else {
					fmt.Printf("  %s\n", name)
				}
			}
			return nil
		}),
	}
	commands = append(commands, cmdutil.CreateAlias(listContext, "config list context"))

	return commands
}
<|endoftext|>


<commit_before>package log

import (
	"encoding/json"
	"fmt"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/fatih/camelcase"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/sirupsen/logrus"
)

const (
	bucketFactor = 2.0
	bucketCount  = 20 // Which makes the max bucket 2^20 seconds or ~12 days in size
)

// This needs to be a global var, not a field on the logger, because multiple servers
// create new loggers, and the prometheus registration uses a global namespace
var reportMetricGauge prometheus.Gauge
var reportMetricsOnce sync.Once

// Logger is a helper for emitting our grpc API logs
type Logger interface {
	Log(request interface{}, response interface{}, err error, duration time.Duration)
	LogAtLevelFromDepth(request interface{}, response interface{}, err error, duration time.Duration, level logrus.Level, depth int)
}

type logger struct {
	*logrus.Entry
	histogram   map[string]*prometheus.HistogramVec
	counter     map[string]prometheus.Counter
	mutex       *sync.Mutex // synchronizes access to both histogram and counter maps
	exportStats bool
	service     string
}

// NewLogger creates a new logger
func NewLogger(service string) Logger {
	return newLogger(service, true)
}

// NewLocalLogger creates a new logger for local testing (which does not report prometheus metrics)
func NewLocalLogger(service string) Logger {
	return newLogger(service, false)
}

func newLogger(service string, exportStats bool) Logger {
	l := logrus.New()
	l.Formatter = FormatterFunc(Pretty)
	newLogger := &logger{
		l.WithFields(logrus.Fields{"service": service}),
		make(map[string]*prometheus.HistogramVec),
		make(map[string]prometheus.Counter),
		&sync.Mutex{},
		exportStats,
		service,
	}
	if exportStats {
		reportMetricsOnce.Do(func() {
			newReportMetricGauge := prometheus.NewGauge(
				prometheus.GaugeOpts{
					Namespace: "pachyderm",
					Subsystem: "pachd",
					Name:      "report_metric",
					Help:      "gauge of number of calls to ReportMetric()",
				},
			)
			if err := prometheus.Register(newReportMetricGauge); err != nil {
				entry := newLogger.WithFields(logrus.Fields{"method": "NewLogger"})
				newLogger.LogAtLevel(entry, logrus.WarnLevel, fmt.Sprintf("error registering prometheus metric: %v", newReportMetricGauge), err)
			} else {
				reportMetricGauge = newReportMetricGauge
			}
		})
	}
	return newLogger
}

// Helper function used to log requests and responses from our GRPC method
// implementations
func (l *logger) Log(request interface{}, response interface{}, err error, duration time.Duration) {
	if err != nil {
		l.LogAtLevelFromDepth(request, response, err, duration, logrus.ErrorLevel, 4)
	} else {
		l.LogAtLevelFromDepth(request, response, err, duration, logrus.InfoLevel, 4)
	}
	// We have to grab the method's name here before we
	// enter the goro's stack
	go l.ReportMetric(getMethodName(), duration, err)
}

func getMethodName() string {
	depth := 4
	pc := make([]uintptr, depth)
	runtime.Callers(depth, pc)
	split := strings.Split(runtime.FuncForPC(pc[0]).Name(), ".")
	return split[len(split)-1]
}

func (l *logger) ReportMetric(method string, duration time.Duration, err error) {
	if !l.exportStats {
		return
	}
	// Count the number of ReportMetric() goros in case we start to leak them
	if reportMetricGauge != nil {
		reportMetricGauge.Inc()
	}
	defer func() {
		if reportMetricGauge != nil {
			reportMetricGauge.Dec()
		}
	}()
	l.mutex.Lock() // for conccurent map access (histogram,counter)
	defer l.mutex.Unlock()
	state := "started"
	if err != nil {
		state = "errored"
	} else {
		if duration.Seconds() > 0 {
			state = "finished"
		}
	}
	entry := l.WithFields(logrus.Fields{"method": method})

	var tokens []string
	for _, token := range camelcase.Split(method) {
		tokens = append(tokens, strings.ToLower(token))
	}
	rootStatName := strings.Join(tokens, "_")

	// Recording the distribution of started times is meaningless
	if state != "started" {
		runTimeName := fmt.Sprintf("%v_time", rootStatName)
		runTime, ok := l.histogram[runTimeName]
		if !ok {
			runTime = prometheus.NewHistogramVec(
				prometheus.HistogramOpts{
					Namespace: "pachyderm",
					Subsystem: fmt.Sprintf("pachd_%v", topLevelService(l.service)),
					Name:      runTimeName,
					Help:      fmt.Sprintf("Run time of %v", method),
					Buckets:   prometheus.ExponentialBuckets(1.0, bucketFactor, bucketCount),
				},
				[]string{
					"state", // Since both finished and errored API calls can have run times
				},
			)
			if err := prometheus.Register(runTime); err != nil {
				l.LogAtLevel(entry, logrus.WarnLevel, fmt.Sprintf("error registering prometheus metric %v: %v", runTime, runTimeName), err)
			} else {
				l.histogram[runTimeName] = runTime
			}
		}
		if hist, err := runTime.GetMetricWithLabelValues(state); err != nil {
			l.LogAtLevel(entry, logrus.WarnLevel, "failed to get histogram w labels: state (%v) with error %v", state, err)
		} else {
			hist.Observe(duration.Seconds())
		}
	}

	secondsCountName := fmt.Sprintf("%v_seconds_count", rootStatName)
	secondsCount, ok := l.counter[secondsCountName]
	if !ok {
		secondsCount = prometheus.NewCounter(
			prometheus.CounterOpts{
				Namespace: "pachyderm",
				Subsystem: fmt.Sprintf("pachd_%v", topLevelService(l.service)),
				Name:      secondsCountName,
				Help:      fmt.Sprintf("cumulative number of seconds spent in %v", method),
			},
		)
		if err := prometheus.Register(secondsCount); err != nil {
			l.LogAtLevel(entry, logrus.WarnLevel, fmt.Sprintf("error registering prometheus metric %v: %v", secondsCount, secondsCountName), err)
		} else {
			l.counter[secondsCountName] = secondsCount
		}
	}
	secondsCount.Add(duration.Seconds())

}

func (l *logger) LogAtLevel(entry *logrus.Entry, level logrus.Level, args ...interface{}) {
	switch level {
	case logrus.PanicLevel:
		entry.Panic(args)
	case logrus.FatalLevel:
		entry.Fatal(args)
	case logrus.ErrorLevel:
		entry.Error(args)
	case logrus.WarnLevel:
		entry.Warn(args)
	case logrus.InfoLevel:
		entry.Info(args)
	case logrus.DebugLevel:
		entry.Debug(args)
	}
}

func (l *logger) LogAtLevelFromDepth(request interface{}, response interface{}, err error, duration time.Duration, level logrus.Level, depth int) {
	pc := make([]uintptr, depth)
	runtime.Callers(depth, pc)
	split := strings.Split(runtime.FuncForPC(pc[0]).Name(), ".")
	method := split[len(split)-1]

	fields := logrus.Fields{
		"method":  method,
		"request": request,
	}
	if response != nil {
		fields["response"] = response
	}
	if err != nil {
		// "err" itself might be a code or even an empty struct
		fields["error"] = err.Error()
	}
	if duration > 0 {
		fields["duration"] = duration
	}
	l.LogAtLevel(l.WithFields(fields), level)
}

func topLevelService(fullyQualifiedService string) string {
	tokens := strings.Split(fullyQualifiedService, ".")
	return tokens[0]
}

// FormatterFunc is a type alias for a function that satisfies logrus'
// `Formatter` interface
type FormatterFunc func(entry *logrus.Entry) ([]byte, error)

// Format proxies the closure in order to satisfy `logrus.Formatter`'s
// interface.
func (f FormatterFunc) Format(entry *logrus.Entry) ([]byte, error) {
	return f(entry)
}

// Pretty formats a logrus entry like so:
// ```
// 2019-02-11T16:02:02Z INFO pfs.API.InspectRepo {"request":{"repo":{"name":"images"}}} []
// ```
func Pretty(entry *logrus.Entry) ([]byte, error) {
	serialized := []byte(
		fmt.Sprintf(
			"%v %v ",
			entry.Time.Format(logrus.DefaultTimestampFormat),
			strings.ToUpper(entry.Level.String()),
		),
	)
	if entry.Data["service"] != nil {
		serialized = append(serialized, []byte(fmt.Sprintf("%v.%v ", entry.Data["service"], entry.Data["method"]))...)
	}
	if len(entry.Data) > 2 {
		delete(entry.Data, "service")
		delete(entry.Data, "method")
		if entry.Data["duration"] != nil {
			entry.Data["duration"] = entry.Data["duration"].(time.Duration).Seconds()
		}
		data, err := json.Marshal(entry.Data)
		if err != nil {
			return nil, fmt.Errorf("Failed to marshal fields to JSON, %v", err)
		}
		serialized = append(serialized, []byte(string(data))...)
		serialized = append(serialized, ' ')
	}

	serialized = append(serialized, []byte(entry.Message)...)
	serialized = append(serialized, '\n')
	return serialized, nil
}

// GRPCLogWriter proxies gRPC and etcd-produced log messages to a logrus
// logger. Because it implements `io.Writer`, it could be used anywhere where
// `io.Writer`s are used, but it has some logic specifically designed to
// handle gRPC-formatted logs.
type GRPCLogWriter struct {
	logger *logrus.Logger
	source string
}

func NewGRPCLogWriter(logger *logrus.Logger, source string) *GRPCLogWriter {
	return &GRPCLogWriter{
		logger: logger,
		source: source,
	}
}

// Write allows `GRPCInfoWriter` to implement the `io.Writer` interface. This
// will take gRPC logs, which look something like this:
// ```
// INFO: 2019/02/18 12:21:54 ClientConn switching balancer to "pick_first"
// ```
// strip out redundant content, and print the message at the appropriate log
// level in logrus. Any parse errors of the log message will be reported in
// logrus as well.
func (l *GRPCLogWriter) Write(p []byte) (int, error) {
	parts := strings.SplitN(string(p), " ", 4)
	entry := l.logger.WithField("source", l.source)

	if len(parts) == 4 {
		// parts[1] and parts[2] contain the date and time, but logrus already
		// adds this under the `time` entry field, so it's not needed (though
		// the time will presumably be marginally ahead of the original log
		// message)
		level := parts[0]
		message := strings.TrimSpace(parts[3])

		if level == "INFO:" {
			entry.Info(message)
		} else if level == "ERROR:" {
			entry.Error(message)
		} else if level == "WARNING:" {
			entry.Warning(message)
		} else if level == "FATAL:" {
			// no need to call fatal ourselves because gRPC will exit the
			// process
			entry.Error(message)
		} else {
			entry.Error(message)
			entry.Error("entry had unknown log level prefix: '%s'; this is a bug, please report it along with the previous log entry", level)
		}
	} else {
		// can't format the message -- just display the contents
		entry := l.logger.WithFields(logrus.Fields{
			"source": l.source,
		})
		entry.Error(p)
		entry.Error("entry had unexpected format; this is a bug, please report it along with the previous log entry")
	}

	return len(p), nil
}
<commit_msg>Added docstring<commit_after>package log

import (
	"encoding/json"
	"fmt"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/fatih/camelcase"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/sirupsen/logrus"
)

const (
	bucketFactor = 2.0
	bucketCount  = 20 // Which makes the max bucket 2^20 seconds or ~12 days in size
)

// This needs to be a global var, not a field on the logger, because multiple servers
// create new loggers, and the prometheus registration uses a global namespace
var reportMetricGauge prometheus.Gauge
var reportMetricsOnce sync.Once

// Logger is a helper for emitting our grpc API logs
type Logger interface {
	Log(request interface{}, response interface{}, err error, duration time.Duration)
	LogAtLevelFromDepth(request interface{}, response interface{}, err error, duration time.Duration, level logrus.Level, depth int)
}

type logger struct {
	*logrus.Entry
	histogram   map[string]*prometheus.HistogramVec
	counter     map[string]prometheus.Counter
	mutex       *sync.Mutex // synchronizes access to both histogram and counter maps
	exportStats bool
	service     string
}

// NewLogger creates a new logger
func NewLogger(service string) Logger {
	return newLogger(service, true)
}

// NewLocalLogger creates a new logger for local testing (which does not report prometheus metrics)
func NewLocalLogger(service string) Logger {
	return newLogger(service, false)
}

func newLogger(service string, exportStats bool) Logger {
	l := logrus.New()
	l.Formatter = FormatterFunc(Pretty)
	newLogger := &logger{
		l.WithFields(logrus.Fields{"service": service}),
		make(map[string]*prometheus.HistogramVec),
		make(map[string]prometheus.Counter),
		&sync.Mutex{},
		exportStats,
		service,
	}
	if exportStats {
		reportMetricsOnce.Do(func() {
			newReportMetricGauge := prometheus.NewGauge(
				prometheus.GaugeOpts{
					Namespace: "pachyderm",
					Subsystem: "pachd",
					Name:      "report_metric",
					Help:      "gauge of number of calls to ReportMetric()",
				},
			)
			if err := prometheus.Register(newReportMetricGauge); err != nil {
				entry := newLogger.WithFields(logrus.Fields{"method": "NewLogger"})
				newLogger.LogAtLevel(entry, logrus.WarnLevel, fmt.Sprintf("error registering prometheus metric: %v", newReportMetricGauge), err)
			} else {
				reportMetricGauge = newReportMetricGauge
			}
		})
	}
	return newLogger
}

// Helper function used to log requests and responses from our GRPC method
// implementations
func (l *logger) Log(request interface{}, response interface{}, err error, duration time.Duration) {
	if err != nil {
		l.LogAtLevelFromDepth(request, response, err, duration, logrus.ErrorLevel, 4)
	} else {
		l.LogAtLevelFromDepth(request, response, err, duration, logrus.InfoLevel, 4)
	}
	// We have to grab the method's name here before we
	// enter the goro's stack
	go l.ReportMetric(getMethodName(), duration, err)
}

func getMethodName() string {
	depth := 4
	pc := make([]uintptr, depth)
	runtime.Callers(depth, pc)
	split := strings.Split(runtime.FuncForPC(pc[0]).Name(), ".")
	return split[len(split)-1]
}

func (l *logger) ReportMetric(method string, duration time.Duration, err error) {
	if !l.exportStats {
		return
	}
	// Count the number of ReportMetric() goros in case we start to leak them
	if reportMetricGauge != nil {
		reportMetricGauge.Inc()
	}
	defer func() {
		if reportMetricGauge != nil {
			reportMetricGauge.Dec()
		}
	}()
	l.mutex.Lock() // for conccurent map access (histogram,counter)
	defer l.mutex.Unlock()
	state := "started"
	if err != nil {
		state = "errored"
	} else {
		if duration.Seconds() > 0 {
			state = "finished"
		}
	}
	entry := l.WithFields(logrus.Fields{"method": method})

	var tokens []string
	for _, token := range camelcase.Split(method) {
		tokens = append(tokens, strings.ToLower(token))
	}
	rootStatName := strings.Join(tokens, "_")

	// Recording the distribution of started times is meaningless
	if state != "started" {
		runTimeName := fmt.Sprintf("%v_time", rootStatName)
		runTime, ok := l.histogram[runTimeName]
		if !ok {
			runTime = prometheus.NewHistogramVec(
				prometheus.HistogramOpts{
					Namespace: "pachyderm",
					Subsystem: fmt.Sprintf("pachd_%v", topLevelService(l.service)),
					Name:      runTimeName,
					Help:      fmt.Sprintf("Run time of %v", method),
					Buckets:   prometheus.ExponentialBuckets(1.0, bucketFactor, bucketCount),
				},
				[]string{
					"state", // Since both finished and errored API calls can have run times
				},
			)
			if err := prometheus.Register(runTime); err != nil {
				l.LogAtLevel(entry, logrus.WarnLevel, fmt.Sprintf("error registering prometheus metric %v: %v", runTime, runTimeName), err)
			} else {
				l.histogram[runTimeName] = runTime
			}
		}
		if hist, err := runTime.GetMetricWithLabelValues(state); err != nil {
			l.LogAtLevel(entry, logrus.WarnLevel, "failed to get histogram w labels: state (%v) with error %v", state, err)
		} else {
			hist.Observe(duration.Seconds())
		}
	}

	secondsCountName := fmt.Sprintf("%v_seconds_count", rootStatName)
	secondsCount, ok := l.counter[secondsCountName]
	if !ok {
		secondsCount = prometheus.NewCounter(
			prometheus.CounterOpts{
				Namespace: "pachyderm",
				Subsystem: fmt.Sprintf("pachd_%v", topLevelService(l.service)),
				Name:      secondsCountName,
				Help:      fmt.Sprintf("cumulative number of seconds spent in %v", method),
			},
		)
		if err := prometheus.Register(secondsCount); err != nil {
			l.LogAtLevel(entry, logrus.WarnLevel, fmt.Sprintf("error registering prometheus metric %v: %v", secondsCount, secondsCountName), err)
		} else {
			l.counter[secondsCountName] = secondsCount
		}
	}
	secondsCount.Add(duration.Seconds())

}

func (l *logger) LogAtLevel(entry *logrus.Entry, level logrus.Level, args ...interface{}) {
	switch level {
	case logrus.PanicLevel:
		entry.Panic(args)
	case logrus.FatalLevel:
		entry.Fatal(args)
	case logrus.ErrorLevel:
		entry.Error(args)
	case logrus.WarnLevel:
		entry.Warn(args)
	case logrus.InfoLevel:
		entry.Info(args)
	case logrus.DebugLevel:
		entry.Debug(args)
	}
}

func (l *logger) LogAtLevelFromDepth(request interface{}, response interface{}, err error, duration time.Duration, level logrus.Level, depth int) {
	pc := make([]uintptr, depth)
	runtime.Callers(depth, pc)
	split := strings.Split(runtime.FuncForPC(pc[0]).Name(), ".")
	method := split[len(split)-1]

	fields := logrus.Fields{
		"method":  method,
		"request": request,
	}
	if response != nil {
		fields["response"] = response
	}
	if err != nil {
		// "err" itself might be a code or even an empty struct
		fields["error"] = err.Error()
	}
	if duration > 0 {
		fields["duration"] = duration
	}
	l.LogAtLevel(l.WithFields(fields), level)
}

func topLevelService(fullyQualifiedService string) string {
	tokens := strings.Split(fullyQualifiedService, ".")
	return tokens[0]
}

// FormatterFunc is a type alias for a function that satisfies logrus'
// `Formatter` interface
type FormatterFunc func(entry *logrus.Entry) ([]byte, error)

// Format proxies the closure in order to satisfy `logrus.Formatter`'s
// interface.
func (f FormatterFunc) Format(entry *logrus.Entry) ([]byte, error) {
	return f(entry)
}

// Pretty formats a logrus entry like so:
// ```
// 2019-02-11T16:02:02Z INFO pfs.API.InspectRepo {"request":{"repo":{"name":"images"}}} []
// ```
func Pretty(entry *logrus.Entry) ([]byte, error) {
	serialized := []byte(
		fmt.Sprintf(
			"%v %v ",
			entry.Time.Format(logrus.DefaultTimestampFormat),
			strings.ToUpper(entry.Level.String()),
		),
	)
	if entry.Data["service"] != nil {
		serialized = append(serialized, []byte(fmt.Sprintf("%v.%v ", entry.Data["service"], entry.Data["method"]))...)
	}
	if len(entry.Data) > 2 {
		delete(entry.Data, "service")
		delete(entry.Data, "method")
		if entry.Data["duration"] != nil {
			entry.Data["duration"] = entry.Data["duration"].(time.Duration).Seconds()
		}
		data, err := json.Marshal(entry.Data)
		if err != nil {
			return nil, fmt.Errorf("Failed to marshal fields to JSON, %v", err)
		}
		serialized = append(serialized, []byte(string(data))...)
		serialized = append(serialized, ' ')
	}

	serialized = append(serialized, []byte(entry.Message)...)
	serialized = append(serialized, '\n')
	return serialized, nil
}

// GRPCLogWriter proxies gRPC and etcd-produced log messages to a logrus
// logger. Because it implements `io.Writer`, it could be used anywhere where
// `io.Writer`s are used, but it has some logic specifically designed to
// handle gRPC-formatted logs.
type GRPCLogWriter struct {
	logger *logrus.Logger
	source string
}

// NewGRPCLogWriter creates a new GRPC log writer. `logger` specifies the
// underlying logger, and `source` specifies where these logs are coming from;
// it is added as a entry field for all log messages.
func NewGRPCLogWriter(logger *logrus.Logger, source string) *GRPCLogWriter {
	return &GRPCLogWriter{
		logger: logger,
		source: source,
	}
}

// Write allows `GRPCInfoWriter` to implement the `io.Writer` interface. This
// will take gRPC logs, which look something like this:
// ```
// INFO: 2019/02/18 12:21:54 ClientConn switching balancer to "pick_first"
// ```
// strip out redundant content, and print the message at the appropriate log
// level in logrus. Any parse errors of the log message will be reported in
// logrus as well.
func (l *GRPCLogWriter) Write(p []byte) (int, error) {
	parts := strings.SplitN(string(p), " ", 4)
	entry := l.logger.WithField("source", l.source)

	if len(parts) == 4 {
		// parts[1] and parts[2] contain the date and time, but logrus already
		// adds this under the `time` entry field, so it's not needed (though
		// the time will presumably be marginally ahead of the original log
		// message)
		level := parts[0]
		message := strings.TrimSpace(parts[3])

		if level == "INFO:" {
			entry.Info(message)
		} else if level == "ERROR:" {
			entry.Error(message)
		} else if level == "WARNING:" {
			entry.Warning(message)
		} else if level == "FATAL:" {
			// no need to call fatal ourselves because gRPC will exit the
			// process
			entry.Error(message)
		} else {
			entry.Error(message)
			entry.Error("entry had unknown log level prefix: '%s'; this is a bug, please report it along with the previous log entry", level)
		}
	} else {
		// can't format the message -- just display the contents
		entry := l.logger.WithFields(logrus.Fields{
			"source": l.source,
		})
		entry.Error(p)
		entry.Error("entry had unexpected format; this is a bug, please report it along with the previous log entry")
	}

	return len(p), nil
}
<|endoftext|>


<commit_before>package testing //import "go.pachyderm.com/pachyderm/src/pfs/testing"

import (
	"errors"
	"fmt"
	"os"
	"sync/atomic"
	"testing"

	"go.pedge.io/proto/test"

	"go.pachyderm.com/pachyderm/src/pfs"
	"go.pachyderm.com/pachyderm/src/pfs/drive"
	"go.pachyderm.com/pachyderm/src/pfs/drive/btrfs"
	"go.pachyderm.com/pachyderm/src/pfs/route"
	"go.pachyderm.com/pachyderm/src/pfs/server"
	"go.pachyderm.com/pachyderm/src/pkg/discovery"
	"go.pachyderm.com/pachyderm/src/pkg/grpcutil"
	"go.pachyderm.com/pachyderm/src/pkg/require"
	"google.golang.org/grpc"
)

const (
	// TODO(pedge): large numbers of shards takes forever because
	// we are doing tons of btrfs operations on init, is there anything
	// we can do about that?
	testShardsPerServer = 8
	testNumServers      = 1
	testNumReplicas     = 0
)

var (
	counter int32
)

func RunTest(
	t *testing.T,
	f func(*testing.T, pfs.ApiClient, pfs.InternalApiClient, Cluster),
) {
	discoveryClient, err := getEtcdClient()
	require.NoError(t, err)
	var cluster *cluster
	prototest.RunT(
		t,
		testNumServers,
		func(servers map[string]*grpc.Server) {
			cluster = registerFunc(t, discoveryClient, servers)
		},
		func(t *testing.T, clientConns map[string]*grpc.ClientConn) {
			var clientConn *grpc.ClientConn
			for _, c := range clientConns {
				clientConn = c
				break
			}
			go func() {
				require.Equal(t, cluster.addresser.AssignRoles(cluster.cancel), route.ErrCancelled)
			}()
			cluster.WaitForAvailability()
			f(
				t,
				pfs.NewApiClient(
					clientConn,
				),
				pfs.NewInternalApiClient(
					clientConn,
				),
				cluster,
			)
		},
	)
	cluster.Shutdown()
}

func RunBench(
	b *testing.B,
	f func(*testing.B, pfs.ApiClient),
) {
	discoveryClient, err := getEtcdClient()
	require.NoError(b, err)
	var cluster *cluster
	prototest.RunB(
		b,
		testNumServers,
		func(servers map[string]*grpc.Server) {
			cluster = registerFunc(b, discoveryClient, servers)
		},
		func(b *testing.B, clientConns map[string]*grpc.ClientConn) {
			var clientConn *grpc.ClientConn
			for _, c := range clientConns {
				clientConn = c
				break
			}
			go func() {
				require.Equal(b, cluster.addresser.AssignRoles(cluster.cancel), route.ErrCancelled)
			}()
			cluster.WaitForAvailability()
			f(
				b,
				pfs.NewApiClient(
					clientConn,
				),
			)
		},
	)
	cluster.Shutdown()
}

type Cluster interface {
	WaitForAvailability()
	Kill(server int)
	Restart(server int)
	KillRoleAssigner()
	RestartRoleAssigner()
	Shutdown()
}

type cluster struct {
	addresses       []string
	servers         map[string]server.APIServer
	internalServers map[string]server.InternalAPIServer
	cancels         map[string]chan bool
	internalCancels map[string]chan bool
	cancel          chan bool
	addresser       route.TestAddresser
	sharder         route.Sharder
	tb              testing.TB
}

func (c *cluster) WaitForAvailability() {
	// We use address as the id for servers
	var frontendIds []string
	var serverIds []string
	for _, address := range c.addresses {
		if _, ok := c.cancels[address]; ok {
			frontendIds = append(frontendIds, address)
		}
		if _, ok := c.internalCancels[address]; ok {
			serverIds = append(serverIds, address)
		}
	}
	require.NoError(c.tb, c.addresser.WaitForAvailability(frontendIds, serverIds))
}

func (c *cluster) Kill(index int) {
	close(c.cancels[c.addresses[index]])
	delete(c.cancels, c.addresses[index])
	delete(c.internalServers, c.addresses[index])
}

func (c *cluster) Restart(index int) {
	address := c.addresses[index]
	c.cancels[address] = make(chan bool)
	internalAPIServer := server.NewInternalAPIServer(
		c.sharder,
		route.NewRouter(
			c.addresser,
			grpcutil.NewDialer(
				grpc.WithInsecure(),
			),
			address,
		),
		getDriver(c.tb, address),
	)
	c.internalServers[address] = internalAPIServer
	go func() {
		require.Equal(c.tb, c.addresser.Register(c.cancels[address], address, address, c.internalServers[address]), route.ErrCancelled)
	}()
}

func (c *cluster) KillRoleAssigner() {
	close(c.cancel)
}

func (c *cluster) RestartRoleAssigner() {
	c.cancel = make(chan bool)
	go func() {
		require.Equal(c.tb, c.addresser.AssignRoles(c.cancel), route.ErrCancelled)
	}()
}

func (c *cluster) Shutdown() {
	close(c.cancel)
	for _, cancel := range c.cancels {
		close(cancel)
	}
}

func newCluster(tb testing.TB, discoveryClient discovery.Client, servers map[string]*grpc.Server) *cluster {
	sharder := route.NewSharder(
		testShardsPerServer*testNumServers,
		testNumReplicas,
	)
	addresser := route.NewDiscoveryTestAddresser(discoveryClient, sharder, testNamespace())
	cluster := cluster{
		servers:         make(map[string]server.APIServer),
		internalServers: make(map[string]server.InternalAPIServer),
		cancels:         make(map[string]chan bool),
		internalCancels: make(map[string]chan bool),
		cancel:          make(chan bool),
		addresser:       addresser,
		sharder:         sharder,
		tb:              tb,
	}
	for address, s := range servers {
		cluster.addresses = append(cluster.addresses, address)
		apiServer := server.NewAPIServer(
			cluster.sharder,
			route.NewRouter(
				cluster.addresser,
				grpcutil.NewDialer(
					grpc.WithInsecure(),
				),
				address,
			),
		)
		cluster.servers[address] = apiServer
		cluster.cancels[address] = make(chan bool)
		go func(address string) {
			require.Equal(tb, cluster.addresser.RegisterFrontend(cluster.cancels[address], address, cluster.servers[address]), route.ErrCancelled)
		}(address)
		pfs.RegisterApiServer(s, apiServer)
		internalAPIServer := server.NewInternalAPIServer(
			cluster.sharder,
			route.NewRouter(
				cluster.addresser,
				grpcutil.NewDialer(
					grpc.WithInsecure(),
				),
				address,
			),
			getDriver(tb, address),
		)
		pfs.RegisterInternalApiServer(s, internalAPIServer)
		cluster.internalServers[address] = internalAPIServer
		cluster.internalCancels[address] = make(chan bool)
		go func(address string) {
			require.Equal(tb, cluster.addresser.Register(cluster.internalCancels[address], address, address, cluster.internalServers[address]), route.ErrCancelled)
		}(address)
	}
	return &cluster
}

func registerFunc(tb testing.TB, discoveryClient discovery.Client, servers map[string]*grpc.Server) *cluster {
	return newCluster(tb, discoveryClient, servers)
}

func getDriver(tb testing.TB, namespace string) drive.Driver {
	driver, err := btrfs.NewDriver(getBtrfsRootDir(tb), namespace)
	require.NoError(tb, err)
	return driver
}

func getBtrfsRootDir(tb testing.TB) string {
	// TODO(pedge)
	rootDir := os.Getenv("PFS_DRIVER_ROOT")
	if rootDir == "" {
		tb.Fatal("PFS_DRIVER_ROOT not set")
	}
	return rootDir
}

func getEtcdClient() (discovery.Client, error) {
	etcdAddress, err := getEtcdAddress()
	if err != nil {
		return nil, err
	}
	return discovery.NewEtcdClient(etcdAddress), nil
}

func getEtcdAddress() (string, error) {
	etcdAddr := os.Getenv("ETCD_PORT_2379_TCP_ADDR")
	if etcdAddr == "" {
		return "", errors.New("ETCD_PORT_2379_TCP_ADDR not set")
	}
	return fmt.Sprintf("http://%s:2379", etcdAddr), nil
}

func testNamespace() string {
	return fmt.Sprintf("test-%d", atomic.AddInt32(&counter, 1))
}
<commit_msg>Make tests harder again.<commit_after>package testing //import "go.pachyderm.com/pachyderm/src/pfs/testing"

import (
	"errors"
	"fmt"
	"os"
	"sync/atomic"
	"testing"

	"go.pedge.io/proto/test"

	"go.pachyderm.com/pachyderm/src/pfs"
	"go.pachyderm.com/pachyderm/src/pfs/drive"
	"go.pachyderm.com/pachyderm/src/pfs/drive/btrfs"
	"go.pachyderm.com/pachyderm/src/pfs/route"
	"go.pachyderm.com/pachyderm/src/pfs/server"
	"go.pachyderm.com/pachyderm/src/pkg/discovery"
	"go.pachyderm.com/pachyderm/src/pkg/grpcutil"
	"go.pachyderm.com/pachyderm/src/pkg/require"
	"google.golang.org/grpc"
)

const (
	// TODO(pedge): large numbers of shards takes forever because
	// we are doing tons of btrfs operations on init, is there anything
	// we can do about that?
	testShardsPerServer = 8
	testNumServers      = 8
	testNumReplicas     = 1
)

var (
	counter int32
)

func RunTest(
	t *testing.T,
	f func(*testing.T, pfs.ApiClient, pfs.InternalApiClient, Cluster),
) {
	discoveryClient, err := getEtcdClient()
	require.NoError(t, err)
	var cluster *cluster
	prototest.RunT(
		t,
		testNumServers,
		func(servers map[string]*grpc.Server) {
			cluster = registerFunc(t, discoveryClient, servers)
		},
		func(t *testing.T, clientConns map[string]*grpc.ClientConn) {
			var clientConn *grpc.ClientConn
			for _, c := range clientConns {
				clientConn = c
				break
			}
			go func() {
				require.Equal(t, cluster.addresser.AssignRoles(cluster.cancel), route.ErrCancelled)
			}()
			cluster.WaitForAvailability()
			f(
				t,
				pfs.NewApiClient(
					clientConn,
				),
				pfs.NewInternalApiClient(
					clientConn,
				),
				cluster,
			)
		},
	)
	cluster.Shutdown()
}

func RunBench(
	b *testing.B,
	f func(*testing.B, pfs.ApiClient),
) {
	discoveryClient, err := getEtcdClient()
	require.NoError(b, err)
	var cluster *cluster
	prototest.RunB(
		b,
		testNumServers,
		func(servers map[string]*grpc.Server) {
			cluster = registerFunc(b, discoveryClient, servers)
		},
		func(b *testing.B, clientConns map[string]*grpc.ClientConn) {
			var clientConn *grpc.ClientConn
			for _, c := range clientConns {
				clientConn = c
				break
			}
			go func() {
				require.Equal(b, cluster.addresser.AssignRoles(cluster.cancel), route.ErrCancelled)
			}()
			cluster.WaitForAvailability()
			f(
				b,
				pfs.NewApiClient(
					clientConn,
				),
			)
		},
	)
	cluster.Shutdown()
}

type Cluster interface {
	WaitForAvailability()
	Kill(server int)
	Restart(server int)
	KillRoleAssigner()
	RestartRoleAssigner()
	Shutdown()
}

type cluster struct {
	addresses       []string
	servers         map[string]server.APIServer
	internalServers map[string]server.InternalAPIServer
	cancels         map[string]chan bool
	internalCancels map[string]chan bool
	cancel          chan bool
	addresser       route.TestAddresser
	sharder         route.Sharder
	tb              testing.TB
}

func (c *cluster) WaitForAvailability() {
	// We use address as the id for servers
	var frontendIds []string
	var serverIds []string
	for _, address := range c.addresses {
		if _, ok := c.cancels[address]; ok {
			frontendIds = append(frontendIds, address)
		}
		if _, ok := c.internalCancels[address]; ok {
			serverIds = append(serverIds, address)
		}
	}
	require.NoError(c.tb, c.addresser.WaitForAvailability(frontendIds, serverIds))
}

func (c *cluster) Kill(index int) {
	close(c.cancels[c.addresses[index]])
	delete(c.cancels, c.addresses[index])
	delete(c.internalServers, c.addresses[index])
}

func (c *cluster) Restart(index int) {
	address := c.addresses[index]
	c.cancels[address] = make(chan bool)
	internalAPIServer := server.NewInternalAPIServer(
		c.sharder,
		route.NewRouter(
			c.addresser,
			grpcutil.NewDialer(
				grpc.WithInsecure(),
			),
			address,
		),
		getDriver(c.tb, address),
	)
	c.internalServers[address] = internalAPIServer
	go func() {
		require.Equal(c.tb, c.addresser.Register(c.cancels[address], address, address, c.internalServers[address]), route.ErrCancelled)
	}()
}

func (c *cluster) KillRoleAssigner() {
	close(c.cancel)
}

func (c *cluster) RestartRoleAssigner() {
	c.cancel = make(chan bool)
	go func() {
		require.Equal(c.tb, c.addresser.AssignRoles(c.cancel), route.ErrCancelled)
	}()
}

func (c *cluster) Shutdown() {
	close(c.cancel)
	for _, cancel := range c.cancels {
		close(cancel)
	}
}

func newCluster(tb testing.TB, discoveryClient discovery.Client, servers map[string]*grpc.Server) *cluster {
	sharder := route.NewSharder(
		testShardsPerServer*testNumServers,
		testNumReplicas,
	)
	addresser := route.NewDiscoveryTestAddresser(discoveryClient, sharder, testNamespace())
	cluster := cluster{
		servers:         make(map[string]server.APIServer),
		internalServers: make(map[string]server.InternalAPIServer),
		cancels:         make(map[string]chan bool),
		internalCancels: make(map[string]chan bool),
		cancel:          make(chan bool),
		addresser:       addresser,
		sharder:         sharder,
		tb:              tb,
	}
	for address, s := range servers {
		cluster.addresses = append(cluster.addresses, address)
		apiServer := server.NewAPIServer(
			cluster.sharder,
			route.NewRouter(
				cluster.addresser,
				grpcutil.NewDialer(
					grpc.WithInsecure(),
				),
				address,
			),
		)
		cluster.servers[address] = apiServer
		cluster.cancels[address] = make(chan bool)
		go func(address string) {
			require.Equal(tb, cluster.addresser.RegisterFrontend(cluster.cancels[address], address, cluster.servers[address]), route.ErrCancelled)
		}(address)
		pfs.RegisterApiServer(s, apiServer)
		internalAPIServer := server.NewInternalAPIServer(
			cluster.sharder,
			route.NewRouter(
				cluster.addresser,
				grpcutil.NewDialer(
					grpc.WithInsecure(),
				),
				address,
			),
			getDriver(tb, address),
		)
		pfs.RegisterInternalApiServer(s, internalAPIServer)
		cluster.internalServers[address] = internalAPIServer
		cluster.internalCancels[address] = make(chan bool)
		go func(address string) {
			require.Equal(tb, cluster.addresser.Register(cluster.internalCancels[address], address, address, cluster.internalServers[address]), route.ErrCancelled)
		}(address)
	}
	return &cluster
}

func registerFunc(tb testing.TB, discoveryClient discovery.Client, servers map[string]*grpc.Server) *cluster {
	return newCluster(tb, discoveryClient, servers)
}

func getDriver(tb testing.TB, namespace string) drive.Driver {
	driver, err := btrfs.NewDriver(getBtrfsRootDir(tb), namespace)
	require.NoError(tb, err)
	return driver
}

func getBtrfsRootDir(tb testing.TB) string {
	// TODO(pedge)
	rootDir := os.Getenv("PFS_DRIVER_ROOT")
	if rootDir == "" {
		tb.Fatal("PFS_DRIVER_ROOT not set")
	}
	return rootDir
}

func getEtcdClient() (discovery.Client, error) {
	etcdAddress, err := getEtcdAddress()
	if err != nil {
		return nil, err
	}
	return discovery.NewEtcdClient(etcdAddress), nil
}

func getEtcdAddress() (string, error) {
	etcdAddr := os.Getenv("ETCD_PORT_2379_TCP_ADDR")
	if etcdAddr == "" {
		return "", errors.New("ETCD_PORT_2379_TCP_ADDR not set")
	}
	return fmt.Sprintf("http://%s:2379", etcdAddr), nil
}

func testNamespace() string {
	return fmt.Sprintf("test-%d", atomic.AddInt32(&counter, 1))
}
<|endoftext|>