<commit_before>package internal

import (
	"errors"
	"fmt"
	"math"
	"sync"

	"github.com/ready-steady/adapt"
)

type Target adapt.Target

func NewTarget(problem *Problem) (Target, error) {
	config := problem.Config.Target

	nj, nf, ni := len(config.Rejection), len(config.Refinement), len(config.Importance)
	if nj == 0 || nj != nf || nf != ni {
		return nil, errors.New("the rejection, refinement, and importance " +
			"should not be empty and should have the same number of elements")
	}

	switch config.Name {
	case "end-to-end-delay":
		return newDelayTarget(problem, &config), nil
	case "total-energy":
		return newEnergyTarget(problem, &config), nil
	case "temperature-profile":
		return newProfileTarget(problem, &config)
	default:
		return nil, errors.New("the target is unknown")
	}
}

func String(target Target) string {
	ni, no := target.Dimensions()
	return fmt.Sprintf(`{"inputs": %d, "outputs": %d}`, ni, no)
}

func Monitor(target Target, progress *adapt.Progress) {
	if progress.Iteration == 0 {
		fmt.Printf("%10s %15s %15s %15s\n", "Iteration",
			"Accepted Nodes", "Rejected Nodes", "Current Nodes")
	}
	fmt.Printf("%10d %15d %15d %15d\n", progress.Iteration,
		progress.Accepted, progress.Rejected, progress.Current)
}

func Score(target Target, config *TargetConfig,
	location *adapt.Location, progress *adapt.Progress) float64 {

	rejection := config.Rejection
	refinement := config.Refinement

	_, no := target.Dimensions()
	nj := uint(len(rejection))

	score, reject := 0.0, true
	for i := uint(0); i < no; i++ {
		j := i % nj
		ε := math.Abs(location.Surplus[i])
		if ε >= rejection[j] {
			reject = false
		}
		if ε > refinement[j] {
			score += ε
		}
	}

	if reject {
		score = -1
	}

	return score
}

func Invoke(target Target, points []float64, nw uint) []float64 {
	ni, no := target.Dimensions()
	np := uint(len(points)) / ni

	values := make([]float64, np*no)
	jobs := make(chan uint, np)
	group := sync.WaitGroup{}
	group.Add(int(np))

	for i := uint(0); i < nw; i++ {
		go func() {
			for j := range jobs {
				target.Compute(points[j*ni:(j+1)*ni], values[j*no:(j+1)*no])
				group.Done()
			}
		}()
	}

	for i := uint(0); i < np; i++ {
		jobs <- i
	}

	group.Wait()
	close(jobs)

	return values
}
<commit_msg>Fix a leftover<commit_after>package internal

import (
	"errors"
	"fmt"
	"math"
	"sync"

	"github.com/ready-steady/adapt"
)

type Target adapt.Target

func NewTarget(problem *Problem) (Target, error) {
	config := problem.Config.Target

	nj, nf := len(config.Rejection), len(config.Refinement)
	if nj == 0 || nj != nf {
		return nil, errors.New("the rejection and refinement " +
			"should not be empty and should have the same number of elements")
	}

	switch config.Name {
	case "end-to-end-delay":
		return newDelayTarget(problem, &config), nil
	case "total-energy":
		return newEnergyTarget(problem, &config), nil
	case "temperature-profile":
		return newProfileTarget(problem, &config)
	default:
		return nil, errors.New("the target is unknown")
	}
}

func String(target Target) string {
	ni, no := target.Dimensions()
	return fmt.Sprintf(`{"inputs": %d, "outputs": %d}`, ni, no)
}

func Monitor(target Target, progress *adapt.Progress) {
	if progress.Iteration == 0 {
		fmt.Printf("%10s %15s %15s %15s\n", "Iteration",
			"Accepted Nodes", "Rejected Nodes", "Current Nodes")
	}
	fmt.Printf("%10d %15d %15d %15d\n", progress.Iteration,
		progress.Accepted, progress.Rejected, progress.Current)
}

func Score(target Target, config *TargetConfig,
	location *adapt.Location, progress *adapt.Progress) float64 {

	rejection := config.Rejection
	refinement := config.Refinement

	_, no := target.Dimensions()
	nj := uint(len(rejection))

	score, reject := 0.0, true
	for i := uint(0); i < no; i++ {
		j := i % nj
		ε := math.Abs(location.Surplus[i])
		if ε >= rejection[j] {
			reject = false
		}
		if ε > refinement[j] {
			score += ε
		}
	}

	if reject {
		score = -1
	}

	return score
}

func Invoke(target Target, points []float64, nw uint) []float64 {
	ni, no := target.Dimensions()
	np := uint(len(points)) / ni

	values := make([]float64, np*no)
	jobs := make(chan uint, np)
	group := sync.WaitGroup{}
	group.Add(int(np))

	for i := uint(0); i < nw; i++ {
		go func() {
			for j := range jobs {
				target.Compute(points[j*ni:(j+1)*ni], values[j*no:(j+1)*no])
				group.Done()
			}
		}()
	}

	for i := uint(0); i < np; i++ {
		jobs <- i
	}

	group.Wait()
	close(jobs)

	return values
}
<|endoftext|>


<commit_before>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package main

import (
	"encoding/base64"
	"fmt"
	"net"

	"launchpad.net/gnuflag"
	"launchpad.net/goyaml"

	"launchpad.net/juju-core/agent"
	"launchpad.net/juju-core/agent/mongo"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/constraints"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/api/params"
	"launchpad.net/juju-core/worker/peergrouper"
)

type BootstrapCommand struct {
	cmd.CommandBase
	AgentConf
	EnvConfig   map[string]interface{}
	Constraints constraints.Value
	Hardware    instance.HardwareCharacteristics
	InstanceId  string
}

// Info returns a decription of the command.
func (c *BootstrapCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "bootstrap-state",
		Purpose: "initialize juju state",
	}
}

func (c *BootstrapCommand) SetFlags(f *gnuflag.FlagSet) {
	c.AgentConf.AddFlags(f)
	yamlBase64Var(f, &c.EnvConfig, "env-config", "", "initial environment configuration (yaml, base64 encoded)")
	f.Var(constraints.ConstraintsValue{Target: &c.Constraints}, "constraints", "initial environment constraints (space-separated strings)")
	f.Var(&c.Hardware, "hardware", "hardware characteristics (space-separated strings)")
	f.StringVar(&c.InstanceId, "instance-id", "", "unique instance-id for bootstrap machine")
}

// Init initializes the command for running.
func (c *BootstrapCommand) Init(args []string) error {
	if len(c.EnvConfig) == 0 {
		return requiredError("env-config")
	}
	if c.InstanceId == "" {
		return requiredError("instance-id")
	}
	return c.AgentConf.CheckArgs(args)
}

// Run initializes state for an environment.
func (c *BootstrapCommand) Run(_ *cmd.Context) error {
	envCfg, err := config.New(config.NoDefaults, c.EnvConfig)
	if err != nil {
		return err
	}
	err = c.ReadConfig("machine-0")
	if err != nil {
		return err
	}
	agentConfig := c.CurrentConfig()

	// agent.Jobs is an optional field in the agent config, and was
	// introduced after 1.17.2. We default to allowing units on
	// machine-0 if missing.
	jobs := agentConfig.Jobs()
	if len(jobs) == 0 {
		jobs = []params.MachineJob{
			params.JobManageEnviron,
			params.JobHostUnits,
		}
	}

	// Get the bootstrap machine's addresses from the provider.
	env, err := environs.New(envCfg)
	if err != nil {
		return err
	}
	instanceId := instance.Id(c.InstanceId)
	instances, err := env.Instances([]instance.Id{instanceId})
	if err != nil {
		return err
	}
	addrs, err := instances[0].Addresses()
	if err != nil {
		return err
	}

	// Generate a shared secret for the Mongo replica set, and write it out.
	sharedSecret, err := mongo.GenerateSharedSecret()
	if err != nil {
		return err
	}
	info, ok := agentConfig.StateServingInfo()
	if !ok {
		return fmt.Errorf("bootstrap machine config has no state serving info")
	}
	info.SharedSecret = sharedSecret
	if err := c.ChangeConfig(func(agentConfig agent.ConfigSetter) {
		agentConfig.SetStateServingInfo(info)
	}); err != nil {
		return fmt.Errorf("cannot write agent config: %v", err)
	}
	agentConfig = c.CurrentConfig()

	if err := c.startMongo(addrs, agentConfig); err != nil {
		return err
	}

	logger.Debugf("started mongo")
	// Initialise state, and store any agent config (e.g. password) changes.
	var st *state.State
	err = nil
	writeErr := c.ChangeConfig(func(agentConfig agent.ConfigSetter) {
		st, _, err = agent.InitializeState(
			agentConfig,
			envCfg,
			agent.BootstrapMachineConfig{
				Addresses:       addrs,
				Constraints:     c.Constraints,
				Jobs:            jobs,
				InstanceId:      instanceId,
				Characteristics: c.Hardware,
				SharedSecret:    sharedSecret,
			},
			state.DefaultDialOpts(),
			environs.NewStatePolicy(),
		)
	})
	if writeErr != nil {
		return fmt.Errorf("cannot write initial configuration: %v", err)
	}
	if err != nil {
		return err
	}
	st.Close()
	return nil
}

func (c *BootstrapCommand) startMongo(addrs []instance.Address, agentConfig agent.Config) error {
	logger.Debugf("starting mongo")

	dialInfo, err := state.DialInfo(agentConfig.StateInfo(), state.DefaultDialOpts())
	if err != nil {
		return err
	}
	servingInfo, ok := agentConfig.StateServingInfo()
	if !ok {
		return fmt.Errorf("agent config has no state serving info")
	}
	// Use localhost to dial the mongo server, because it's running in
	// auth mode and will refuse to perform any operations unless
	// we dial that address.
	dialInfo.Addrs = []string{
		net.JoinHostPort("127.0.0.1", fmt.Sprint(servingInfo.StatePort)),
	}
	logger.Infof("calling ensureMongoServer")
	err = ensureMongoServer(
		agentConfig.DataDir(),
		agentConfig.Value(agent.Namespace),
		servingInfo)
	if err != nil {
		return err
	}

	peerAddr := mongo.SelectPeerAddress(addrs)
	if peerAddr == "" {
		return fmt.Errorf("no appropriate peer address found in %q", addrs)
	}
	peerHostPort := net.JoinHostPort(peerAddr, fmt.Sprint(servingInfo.StatePort))

	return peergrouper.MaybeInitiateMongoServer(peergrouper.InitiateMongoParams{
		DialInfo:       dialInfo,
		MemberHostPort: peerHostPort,
	})
}

// yamlBase64Value implements gnuflag.Value on a map[string]interface{}.
type yamlBase64Value map[string]interface{}

// Set decodes the base64 value into yaml then expands that into a map.
func (v *yamlBase64Value) Set(value string) error {
	decoded, err := base64.StdEncoding.DecodeString(value)
	if err != nil {
		return err
	}
	return goyaml.Unmarshal(decoded, v)
}

func (v *yamlBase64Value) String() string {
	return fmt.Sprintf("%v", *v)
}

// yamlBase64Var sets up a gnuflag flag analogous to the FlagSet.*Var methods.
func yamlBase64Var(fs *gnuflag.FlagSet, target *map[string]interface{}, name string, value string, usage string) {
	fs.Var((*yamlBase64Value)(target), name, usage)
}
<commit_msg>cmd/jujud: use mock MaybeInitiateMongoServer<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package main

import (
	"encoding/base64"
	"fmt"
	"net"

	"launchpad.net/gnuflag"
	"launchpad.net/goyaml"

	"launchpad.net/juju-core/agent"
	"launchpad.net/juju-core/agent/mongo"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/constraints"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/api/params"
	"launchpad.net/juju-core/worker/peergrouper"
)

type BootstrapCommand struct {
	cmd.CommandBase
	AgentConf
	EnvConfig   map[string]interface{}
	Constraints constraints.Value
	Hardware    instance.HardwareCharacteristics
	InstanceId  string
}

// Info returns a decription of the command.
func (c *BootstrapCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "bootstrap-state",
		Purpose: "initialize juju state",
	}
}

func (c *BootstrapCommand) SetFlags(f *gnuflag.FlagSet) {
	c.AgentConf.AddFlags(f)
	yamlBase64Var(f, &c.EnvConfig, "env-config", "", "initial environment configuration (yaml, base64 encoded)")
	f.Var(constraints.ConstraintsValue{Target: &c.Constraints}, "constraints", "initial environment constraints (space-separated strings)")
	f.Var(&c.Hardware, "hardware", "hardware characteristics (space-separated strings)")
	f.StringVar(&c.InstanceId, "instance-id", "", "unique instance-id for bootstrap machine")
}

// Init initializes the command for running.
func (c *BootstrapCommand) Init(args []string) error {
	if len(c.EnvConfig) == 0 {
		return requiredError("env-config")
	}
	if c.InstanceId == "" {
		return requiredError("instance-id")
	}
	return c.AgentConf.CheckArgs(args)
}

// Run initializes state for an environment.
func (c *BootstrapCommand) Run(_ *cmd.Context) error {
	envCfg, err := config.New(config.NoDefaults, c.EnvConfig)
	if err != nil {
		return err
	}
	err = c.ReadConfig("machine-0")
	if err != nil {
		return err
	}
	agentConfig := c.CurrentConfig()

	// agent.Jobs is an optional field in the agent config, and was
	// introduced after 1.17.2. We default to allowing units on
	// machine-0 if missing.
	jobs := agentConfig.Jobs()
	if len(jobs) == 0 {
		jobs = []params.MachineJob{
			params.JobManageEnviron,
			params.JobHostUnits,
		}
	}

	// Get the bootstrap machine's addresses from the provider.
	env, err := environs.New(envCfg)
	if err != nil {
		return err
	}
	instanceId := instance.Id(c.InstanceId)
	instances, err := env.Instances([]instance.Id{instanceId})
	if err != nil {
		return err
	}
	addrs, err := instances[0].Addresses()
	if err != nil {
		return err
	}

	// Generate a shared secret for the Mongo replica set, and write it out.
	sharedSecret, err := mongo.GenerateSharedSecret()
	if err != nil {
		return err
	}
	info, ok := agentConfig.StateServingInfo()
	if !ok {
		return fmt.Errorf("bootstrap machine config has no state serving info")
	}
	info.SharedSecret = sharedSecret
	if err := c.ChangeConfig(func(agentConfig agent.ConfigSetter) {
		agentConfig.SetStateServingInfo(info)
	}); err != nil {
		return fmt.Errorf("cannot write agent config: %v", err)
	}
	agentConfig = c.CurrentConfig()

	if err := c.startMongo(addrs, agentConfig); err != nil {
		return err
	}

	logger.Debugf("started mongo")
	// Initialise state, and store any agent config (e.g. password) changes.
	var st *state.State
	err = nil
	writeErr := c.ChangeConfig(func(agentConfig agent.ConfigSetter) {
		st, _, err = agent.InitializeState(
			agentConfig,
			envCfg,
			agent.BootstrapMachineConfig{
				Addresses:       addrs,
				Constraints:     c.Constraints,
				Jobs:            jobs,
				InstanceId:      instanceId,
				Characteristics: c.Hardware,
				SharedSecret:    sharedSecret,
			},
			state.DefaultDialOpts(),
			environs.NewStatePolicy(),
		)
	})
	if writeErr != nil {
		return fmt.Errorf("cannot write initial configuration: %v", err)
	}
	if err != nil {
		return err
	}
	st.Close()
	return nil
}

func (c *BootstrapCommand) startMongo(addrs []instance.Address, agentConfig agent.Config) error {
	logger.Debugf("starting mongo")

	dialInfo, err := state.DialInfo(agentConfig.StateInfo(), state.DefaultDialOpts())
	if err != nil {
		return err
	}
	servingInfo, ok := agentConfig.StateServingInfo()
	if !ok {
		return fmt.Errorf("agent config has no state serving info")
	}
	// Use localhost to dial the mongo server, because it's running in
	// auth mode and will refuse to perform any operations unless
	// we dial that address.
	dialInfo.Addrs = []string{
		net.JoinHostPort("127.0.0.1", fmt.Sprint(servingInfo.StatePort)),
	}
	logger.Infof("calling ensureMongoServer")
	err = ensureMongoServer(
		agentConfig.DataDir(),
		agentConfig.Value(agent.Namespace),
		servingInfo)
	if err != nil {
		return err
	}

	peerAddr := mongo.SelectPeerAddress(addrs)
	if peerAddr == "" {
		return fmt.Errorf("no appropriate peer address found in %q", addrs)
	}
	peerHostPort := net.JoinHostPort(peerAddr, fmt.Sprint(servingInfo.StatePort))

	return maybeInitiateMongoServer(peergrouper.InitiateMongoParams{
		DialInfo:       dialInfo,
		MemberHostPort: peerHostPort,
	})
}

// yamlBase64Value implements gnuflag.Value on a map[string]interface{}.
type yamlBase64Value map[string]interface{}

// Set decodes the base64 value into yaml then expands that into a map.
func (v *yamlBase64Value) Set(value string) error {
	decoded, err := base64.StdEncoding.DecodeString(value)
	if err != nil {
		return err
	}
	return goyaml.Unmarshal(decoded, v)
}

func (v *yamlBase64Value) String() string {
	return fmt.Sprintf("%v", *v)
}

// yamlBase64Var sets up a gnuflag flag analogous to the FlagSet.*Var methods.
func yamlBase64Var(fs *gnuflag.FlagSet, target *map[string]interface{}, name string, value string, usage string) {
	fs.Var((*yamlBase64Value)(target), name, usage)
}
<|endoftext|>


<commit_before>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package main

import (
	"encoding/base64"
	"fmt"
	"net"

	"launchpad.net/gnuflag"
	"launchpad.net/goyaml"

	"launchpad.net/juju-core/agent"
	"launchpad.net/juju-core/agent/mongo"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/constraints"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/api/params"
	"launchpad.net/juju-core/worker/peergrouper"
)

type BootstrapCommand struct {
	cmd.CommandBase
	AgentConf
	EnvConfig   map[string]interface{}
	Constraints constraints.Value
	Hardware    instance.HardwareCharacteristics
	InstanceId  string
}

// Info returns a decription of the command.
func (c *BootstrapCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "bootstrap-state",
		Purpose: "initialize juju state",
	}
}

func (c *BootstrapCommand) SetFlags(f *gnuflag.FlagSet) {
	c.AgentConf.AddFlags(f)
	yamlBase64Var(f, &c.EnvConfig, "env-config", "", "initial environment configuration (yaml, base64 encoded)")
	f.Var(constraints.ConstraintsValue{Target: &c.Constraints}, "constraints", "initial environment constraints (space-separated strings)")
	f.Var(&c.Hardware, "hardware", "hardware characteristics (space-separated strings)")
	f.StringVar(&c.InstanceId, "instance-id", "", "unique instance-id for bootstrap machine")
}

// Init initializes the command for running.
func (c *BootstrapCommand) Init(args []string) error {
	if len(c.EnvConfig) == 0 {
		return requiredError("env-config")
	}
	if c.InstanceId == "" {
		return requiredError("instance-id")
	}
	return c.AgentConf.CheckArgs(args)
}

// Run initializes state for an environment.
func (c *BootstrapCommand) Run(_ *cmd.Context) error {
	envCfg, err := config.New(config.NoDefaults, c.EnvConfig)
	if err != nil {
		return err
	}
	err = c.ReadConfig("machine-0")
	if err != nil {
		return err
	}
	agentConfig := c.CurrentConfig()

	// agent.Jobs is an optional field in the agent config, and was
	// introduced after 1.17.2. We default to allowing units on
	// machine-0 if missing.
	jobs := agentConfig.Jobs()
	if len(jobs) == 0 {
		jobs = []params.MachineJob{
			params.JobManageEnviron,
			params.JobHostUnits,
		}
	}

	// Get the bootstrap machine's addresses from the provider.
	env, err := environs.New(envCfg)
	if err != nil {
		return err
	}
	instanceId := instance.Id(c.InstanceId)
	instances, err := env.Instances([]instance.Id{instanceId})
	if err != nil {
		return err
	}
	addrs, err := instances[0].Addresses()
	if err != nil {
		return err
	}

	// Create system-identity file
	if err := agent.WriteSystemIdentityFile(agentConfig); err != nil {
		return err
	}

	// Generate a shared secret for the Mongo replica set, and write it out.
	sharedSecret, err := mongo.GenerateSharedSecret()
	if err != nil {
		return err
	}
	info, ok := agentConfig.StateServingInfo()
	if !ok {
		return fmt.Errorf("bootstrap machine config has no state serving info")
	}
	info.SharedSecret = sharedSecret
	err = c.ChangeConfig(func(agentConfig agent.ConfigSetter) {
		agentConfig.SetStateServingInfo(info)
	})
	if err != nil {
		return fmt.Errorf("cannot write agent config: %v", err)
	}
	agentConfig = c.CurrentConfig()

	if err := c.startMongo(addrs, agentConfig); err != nil {
		return err
	}

	logger.Infof("started mongo")
	// Initialise state, and store any agent config (e.g. password) changes.
	var st *state.State
	var m *state.Machine
	err = nil
	writeErr := c.ChangeConfig(func(agentConfig agent.ConfigSetter) {
		st, m, err = agent.InitializeState(
			agentConfig,
			envCfg,
			agent.BootstrapMachineConfig{
				Addresses:       addrs,
				Constraints:     c.Constraints,
				Jobs:            jobs,
				InstanceId:      instanceId,
				Characteristics: c.Hardware,
				SharedSecret:    sharedSecret,
			},
			state.DefaultDialOpts(),
			environs.NewStatePolicy(),
		)
	})
	if writeErr != nil {
		return fmt.Errorf("cannot write initial configuration: %v", err)
	}
	if err != nil {
		return err
	}
	defer st.Close()

	// bootstrap machine always gets the vote
	return m.SetHasVote(true)
}

func (c *BootstrapCommand) startMongo(addrs []instance.Address, agentConfig agent.Config) error {
	logger.Debugf("starting mongo")

	info, ok := agentConfig.StateInfo()
	if !ok {
		return fmt.Errorf("no state info available")
	}
	dialInfo, err := state.DialInfo(info, state.DefaultDialOpts())
	if err != nil {
		return err
	}
	servingInfo, ok := agentConfig.StateServingInfo()
	if !ok {
		return fmt.Errorf("agent config has no state serving info")
	}
	// Use localhost to dial the mongo server, because it's running in
	// auth mode and will refuse to perform any operations unless
	// we dial that address.
	dialInfo.Addrs = []string{
		net.JoinHostPort("127.0.0.1", fmt.Sprint(servingInfo.StatePort)),
	}

	logger.Debugf("calling ensureMongoServer")
	withHA := shouldEnableHA(agentConfig)
	err = ensureMongoServer(
		agentConfig.DataDir(),
		agentConfig.Value(agent.Namespace),
		servingInfo,
		withHA,
	)
	if err != nil {
		return err
	}
	// If we are not doing HA, there is no need to set up replica set.
	if !withHA {
		return nil
	}

	peerAddr := mongo.SelectPeerAddress(addrs)
	if peerAddr == "" {
		return fmt.Errorf("no appropriate peer address found in %q", addrs)
	}
	peerHostPort := net.JoinHostPort(peerAddr, fmt.Sprint(servingInfo.StatePort))

	return maybeInitiateMongoServer(peergrouper.InitiateMongoParams{
		DialInfo:       dialInfo,
		MemberHostPort: peerHostPort,
	})
}

// yamlBase64Value implements gnuflag.Value on a map[string]interface{}.
type yamlBase64Value map[string]interface{}

// Set decodes the base64 value into yaml then expands that into a map.
func (v *yamlBase64Value) Set(value string) error {
	decoded, err := base64.StdEncoding.DecodeString(value)
	if err != nil {
		return err
	}
	return goyaml.Unmarshal(decoded, v)
}

func (v *yamlBase64Value) String() string {
	return fmt.Sprintf("%v", *v)
}

// yamlBase64Var sets up a gnuflag flag analogous to the FlagSet.*Var methods.
func yamlBase64Var(fs *gnuflag.FlagSet, target *map[string]interface{}, name string, value string, usage string) {
	fs.Var((*yamlBase64Value)(target), name, usage)
}
<commit_msg>Increase dial timeout when contacting mongo at bootstrap time<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package main

import (
	"encoding/base64"
	"fmt"
	"time"
	"net"

	"launchpad.net/gnuflag"
	"launchpad.net/goyaml"

	"launchpad.net/juju-core/agent"
	"launchpad.net/juju-core/agent/mongo"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/constraints"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/api/params"
	"launchpad.net/juju-core/worker/peergrouper"
)

type BootstrapCommand struct {
	cmd.CommandBase
	AgentConf
	EnvConfig   map[string]interface{}
	Constraints constraints.Value
	Hardware    instance.HardwareCharacteristics
	InstanceId  string
}

// Info returns a decription of the command.
func (c *BootstrapCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "bootstrap-state",
		Purpose: "initialize juju state",
	}
}

func (c *BootstrapCommand) SetFlags(f *gnuflag.FlagSet) {
	c.AgentConf.AddFlags(f)
	yamlBase64Var(f, &c.EnvConfig, "env-config", "", "initial environment configuration (yaml, base64 encoded)")
	f.Var(constraints.ConstraintsValue{Target: &c.Constraints}, "constraints", "initial environment constraints (space-separated strings)")
	f.Var(&c.Hardware, "hardware", "hardware characteristics (space-separated strings)")
	f.StringVar(&c.InstanceId, "instance-id", "", "unique instance-id for bootstrap machine")
}

// Init initializes the command for running.
func (c *BootstrapCommand) Init(args []string) error {
	if len(c.EnvConfig) == 0 {
		return requiredError("env-config")
	}
	if c.InstanceId == "" {
		return requiredError("instance-id")
	}
	return c.AgentConf.CheckArgs(args)
}

// Run initializes state for an environment.
func (c *BootstrapCommand) Run(_ *cmd.Context) error {
	envCfg, err := config.New(config.NoDefaults, c.EnvConfig)
	if err != nil {
		return err
	}
	err = c.ReadConfig("machine-0")
	if err != nil {
		return err
	}
	agentConfig := c.CurrentConfig()

	// agent.Jobs is an optional field in the agent config, and was
	// introduced after 1.17.2. We default to allowing units on
	// machine-0 if missing.
	jobs := agentConfig.Jobs()
	if len(jobs) == 0 {
		jobs = []params.MachineJob{
			params.JobManageEnviron,
			params.JobHostUnits,
		}
	}

	// Get the bootstrap machine's addresses from the provider.
	env, err := environs.New(envCfg)
	if err != nil {
		return err
	}
	instanceId := instance.Id(c.InstanceId)
	instances, err := env.Instances([]instance.Id{instanceId})
	if err != nil {
		return err
	}
	addrs, err := instances[0].Addresses()
	if err != nil {
		return err
	}

	// Create system-identity file
	if err := agent.WriteSystemIdentityFile(agentConfig); err != nil {
		return err
	}

	// Generate a shared secret for the Mongo replica set, and write it out.
	sharedSecret, err := mongo.GenerateSharedSecret()
	if err != nil {
		return err
	}
	info, ok := agentConfig.StateServingInfo()
	if !ok {
		return fmt.Errorf("bootstrap machine config has no state serving info")
	}
	info.SharedSecret = sharedSecret
	err = c.ChangeConfig(func(agentConfig agent.ConfigSetter) {
		agentConfig.SetStateServingInfo(info)
	})
	if err != nil {
		return fmt.Errorf("cannot write agent config: %v", err)
	}
	agentConfig = c.CurrentConfig()

	if err := c.startMongo(addrs, agentConfig); err != nil {
		return err
	}

	logger.Infof("started mongo")
	// Initialise state, and store any agent config (e.g. password) changes.
	var st *state.State
	var m *state.Machine
	err = nil
	writeErr := c.ChangeConfig(func(agentConfig agent.ConfigSetter) {
		st, m, err = agent.InitializeState(
			agentConfig,
			envCfg,
			agent.BootstrapMachineConfig{
				Addresses:       addrs,
				Constraints:     c.Constraints,
				Jobs:            jobs,
				InstanceId:      instanceId,
				Characteristics: c.Hardware,
				SharedSecret:    sharedSecret,
			},
			state.DefaultDialOpts(),
			environs.NewStatePolicy(),
		)
	})
	if writeErr != nil {
		return fmt.Errorf("cannot write initial configuration: %v", err)
	}
	if err != nil {
		return err
	}
	defer st.Close()

	// bootstrap machine always gets the vote
	return m.SetHasVote(true)
}

func (c *BootstrapCommand) startMongo(addrs []instance.Address, agentConfig agent.Config) error {
	logger.Debugf("starting mongo")

	info, ok := agentConfig.StateInfo()
	if !ok {
		return fmt.Errorf("no state info available")
	}
	// When bootstrapping, we need to allow enough time for mongo
	// to start as there's no retry loop in place.
	// 5 minutes should suffice.
	bootstrapDialOpts := DialOpts{Timeout:  5 * time.Minute}
	dialInfo, err := state.DialInfo(info, bootstrapDialOpts)
	if err != nil {
		return err
	}
	servingInfo, ok := agentConfig.StateServingInfo()
	if !ok {
		return fmt.Errorf("agent config has no state serving info")
	}
	// Use localhost to dial the mongo server, because it's running in
	// auth mode and will refuse to perform any operations unless
	// we dial that address.
	dialInfo.Addrs = []string{
		net.JoinHostPort("127.0.0.1", fmt.Sprint(servingInfo.StatePort)),
	}

	logger.Debugf("calling ensureMongoServer")
	withHA := shouldEnableHA(agentConfig)
	err = ensureMongoServer(
		agentConfig.DataDir(),
		agentConfig.Value(agent.Namespace),
		servingInfo,
		withHA,
	)
	if err != nil {
		return err
	}
	// If we are not doing HA, there is no need to set up replica set.
	if !withHA {
		return nil
	}

	peerAddr := mongo.SelectPeerAddress(addrs)
	if peerAddr == "" {
		return fmt.Errorf("no appropriate peer address found in %q", addrs)
	}
	peerHostPort := net.JoinHostPort(peerAddr, fmt.Sprint(servingInfo.StatePort))

	return maybeInitiateMongoServer(peergrouper.InitiateMongoParams{
		DialInfo:       dialInfo,
		MemberHostPort: peerHostPort,
	})
}

// yamlBase64Value implements gnuflag.Value on a map[string]interface{}.
type yamlBase64Value map[string]interface{}

// Set decodes the base64 value into yaml then expands that into a map.
func (v *yamlBase64Value) Set(value string) error {
	decoded, err := base64.StdEncoding.DecodeString(value)
	if err != nil {
		return err
	}
	return goyaml.Unmarshal(decoded, v)
}

func (v *yamlBase64Value) String() string {
	return fmt.Sprintf("%v", *v)
}

// yamlBase64Var sets up a gnuflag flag analogous to the FlagSet.*Var methods.
func yamlBase64Var(fs *gnuflag.FlagSet, target *map[string]interface{}, name string, value string, usage string) {
	fs.Var((*yamlBase64Value)(target), name, usage)
}
<|endoftext|>


<commit_before>package main

import (
	"sync"
	"time"

	"github.com/Shopify/sarama"
	log "github.com/funkygao/log4go"
	"github.com/wvanbergen/kafka/consumergroup"
)

type subPool struct {
	gw *Gateway

	// {topic: {group: {client: consumerGroup}}}
	cgs         map[string]map[string]map[string]*consumergroup.ConsumerGroup
	cgsLock     sync.RWMutex
	rebalancing bool
}

func newSubPool(gw *Gateway) *subPool {
	return &subPool{
		gw:  gw,
		cgs: make(map[string]map[string]map[string]*consumergroup.ConsumerGroup),
	}
}

func (this *subPool) PickConsumerGroup(ver, topic, group,
	client string) (cg *consumergroup.ConsumerGroup, err error) {
	this.cgsLock.Lock()
	defer this.cgsLock.Unlock()

	var present bool
	if _, present = this.cgs[topic]; !present {
		this.cgs[topic] = make(map[string]map[string]*consumergroup.ConsumerGroup)
	}
	if _, present = this.cgs[topic][group]; !present {
		this.cgs[topic][group] = make(map[string]*consumergroup.ConsumerGroup)
	}
	cg, present = this.cgs[topic][group][client]
	if present {
		return
	}

	if this.rebalancing {
		err = ErrRebalancing
		return
	}

	// FIXME 2 partition, if 3 client concurrently connects, got 3 consumer
	if this.gw.metaStore.OnlineConsumersCount(topic, group) >= len(this.gw.metaStore.Partitions(topic)) {
		err = ErrTooManyConsumers
		return
	}

	// create the consumer group for this client
	cf := consumergroup.NewConfig()
	cf.ChannelBufferSize = 0
	cf.Offsets.Initial = sarama.OffsetOldest
	cf.Offsets.CommitInterval = options.offsetCommitInterval
	// time to wait for all the offsets for a partition to be processed after stopping to consume from it.
	cf.Offsets.ProcessingTimeout = time.Second * 10 // TODO
	cf.Zookeeper.Chroot = this.gw.metaStore.ZkChroot()
	for i := 0; i < 3; i++ {
		// join group will async register zk owners znodes
		// so, if many client concurrently connects to kateway, will not
		// strictly throw ErrTooManyConsumers
		cg, err = consumergroup.JoinConsumerGroup(group, []string{topic},
			this.gw.metaStore.ZkAddrs(), cf)
		if err == nil {
			this.cgs[topic][group][client] = cg
			break
		}
	}

	return
}

func (this *subPool) KillClient(topic, group, client string) {
	// TODO golang keep-alive max idle defaults 60s
	this.rebalancing = true
	this.cgs[topic][group][client].Close() // will flush offset

	this.cgsLock.Lock()
	delete(this.cgs[topic][group], client)
	this.cgsLock.Unlock()

	this.rebalancing = false

	log.Info("consumer %s{topic:%s, group:%s} closed, rebalanced ok", client, topic, group)
}

func (this *subPool) Start() {
	this.gw.wg.Add(1)
	defer this.gw.wg.Done()

	ever := true
	for ever {
		select {
		case <-this.gw.shutdownCh:
			log.Info("sub pool shutdown")
			this.Stop()
			ever = false

		case remoteAddr := <-this.gw.closedConnCh: // TODO
			log.Info("sub client %s closed", remoteAddr)
		}
	}

}

func (this *subPool) Stop() {
	this.cgsLock.Lock()
	defer this.cgsLock.Unlock()

	var wg sync.WaitGroup
	for topic, ts := range this.cgs {
		for group, gs := range ts {
			for client, c := range gs {
				wg.Add(1)
				go func() {
					defer wg.Done()

					if err := c.Close(); err != nil {
						// will commit the offset
						log.Error("{topic:%s, group:%s, client:%s}: %v", topic,
							group, client, err)
					}
				}()
			}
		}
	}

	// wait for all consumers commit offset
	wg.Wait()

	// reinit the vars
	this.cgs = make(map[string]map[string]map[string]*consumergroup.ConsumerGroup)
}
<commit_msg>sub group should backoff between retries<commit_after>package main

import (
	"sync"
	"time"

	"github.com/Shopify/sarama"
	log "github.com/funkygao/log4go"
	"github.com/wvanbergen/kafka/consumergroup"
)

type subPool struct {
	gw *Gateway

	// {topic: {group: {client: consumerGroup}}}
	cgs         map[string]map[string]map[string]*consumergroup.ConsumerGroup
	cgsLock     sync.RWMutex
	rebalancing bool
}

func newSubPool(gw *Gateway) *subPool {
	return &subPool{
		gw:  gw,
		cgs: make(map[string]map[string]map[string]*consumergroup.ConsumerGroup),
	}
}

func (this *subPool) PickConsumerGroup(ver, topic, group,
	client string) (cg *consumergroup.ConsumerGroup, err error) {
	this.cgsLock.Lock()
	defer this.cgsLock.Unlock()

	var present bool
	if _, present = this.cgs[topic]; !present {
		this.cgs[topic] = make(map[string]map[string]*consumergroup.ConsumerGroup)
	}
	if _, present = this.cgs[topic][group]; !present {
		this.cgs[topic][group] = make(map[string]*consumergroup.ConsumerGroup)
	}
	cg, present = this.cgs[topic][group][client]
	if present {
		return
	}

	if this.rebalancing {
		err = ErrRebalancing
		return
	}

	// FIXME 2 partition, if 3 client concurrently connects, got 3 consumer
	if this.gw.metaStore.OnlineConsumersCount(topic, group) >= len(this.gw.metaStore.Partitions(topic)) {
		err = ErrTooManyConsumers
		return
	}

	// create the consumer group for this client
	cf := consumergroup.NewConfig()
	cf.ChannelBufferSize = 0
	cf.Offsets.Initial = sarama.OffsetOldest
	cf.Offsets.CommitInterval = options.offsetCommitInterval
	// time to wait for all the offsets for a partition to be processed after stopping to consume from it.
	cf.Offsets.ProcessingTimeout = time.Second * 10 // TODO
	cf.Zookeeper.Chroot = this.gw.metaStore.ZkChroot()
	for i := 0; i < 3; i++ {
		// join group will async register zk owners znodes
		// so, if many client concurrently connects to kateway, will not
		// strictly throw ErrTooManyConsumers
		cg, err = consumergroup.JoinConsumerGroup(group, []string{topic},
			this.gw.metaStore.ZkAddrs(), cf)
		if err == nil {
			this.cgs[topic][group][client] = cg
			break
		}

		// backoff
		time.Sleep(time.Millisecond * 100)
	}

	return
}

func (this *subPool) KillClient(topic, group, client string) {
	// TODO golang keep-alive max idle defaults 60s
	this.rebalancing = true
	this.cgs[topic][group][client].Close() // will flush offset

	this.cgsLock.Lock()
	delete(this.cgs[topic][group], client)
	this.cgsLock.Unlock()

	this.rebalancing = false

	log.Info("consumer %s{topic:%s, group:%s} closed, rebalanced ok", client, topic, group)
}

func (this *subPool) Start() {
	this.gw.wg.Add(1)
	defer this.gw.wg.Done()

	ever := true
	for ever {
		select {
		case <-this.gw.shutdownCh:
			log.Info("sub pool shutdown")
			this.Stop()
			ever = false

		case remoteAddr := <-this.gw.closedConnCh: // TODO
			log.Info("sub client %s closed", remoteAddr)
		}
	}

}

func (this *subPool) Stop() {
	this.cgsLock.Lock()
	defer this.cgsLock.Unlock()

	var wg sync.WaitGroup
	for topic, ts := range this.cgs {
		for group, gs := range ts {
			for client, c := range gs {
				wg.Add(1)
				go func() {
					defer wg.Done()

					if err := c.Close(); err != nil {
						// will commit the offset
						log.Error("{topic:%s, group:%s, client:%s}: %v", topic,
							group, client, err)
					}
				}()
			}
		}
	}

	// wait for all consumers commit offset
	wg.Wait()

	// reinit the vars
	this.cgs = make(map[string]map[string]map[string]*consumergroup.ConsumerGroup)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"strings"

	"github.com/vpn-kill-switch/killswitch"
)

// PadRight add spaces for aligning the output
func PadRight(str, pad string, length int) string {
	for {
		str += pad
		if len(str) > length {
			return str[0:length]
		}
	}
}

func exit1(err error) {
	fmt.Println(err)
	os.Exit(1)
}

var version string

func main() {

	var (
		ip       = flag.String("ip", "", "VPN peer `IPv4`, killswitch tries to find this automatically")
		d        = flag.Bool("d", false, "`Disable` load /etc/pf.conf rules")
		e        = flag.Bool("e", false, "`Enable` load the pf rules")
		p        = flag.Bool("p", false, "`Print` the pf rules")
		v        = flag.Bool("v", false, fmt.Sprintf("Print version: %s", version))
		paranoid = flag.Bool("paranoid", false, "block icmp (ping) and any non VPN DNS requests")
	)

	flag.Parse()

	if *v {
		fmt.Printf("%s\n", version)
		os.Exit(0)
	}

	if *d {
		exec.Command("pfctl", "-e").CombinedOutput()
		fmt.Printf("# %s\n", strings.Repeat("-", 62))
		fmt.Println("# Loading /etc/pf.conf rules")
		fmt.Printf("# %s\n", strings.Repeat("-", 62))
		out, _ := exec.Command("pfctl",
			"-Fa",
			"-f",
			"/etc/pf.conf").CombinedOutput()
		fmt.Printf("%s\n", out)
		out, _ = exec.Command("pfctl", "-sr").CombinedOutput()
		fmt.Printf("%s\n", out)
		return
	}

	ks, err := killswitch.New(*ip)
	if err != nil {
		exit1(err)
	}

	err = ks.GetActive()
	if err != nil {
		exit1(err)
	}

	if len(ks.UpInterfaces) == 0 {
		exit1(fmt.Errorf("No active interfaces found, verify network settings, use (\"%s -h\") for help.\n", os.Args[0]))
	}

	fmt.Println("Interface  MAC address         IP")
	for k, v := range ks.UpInterfaces {
		fmt.Printf("%s %s   %s\n", PadRight(k, " ", 10), v[0], v[1])
	}
	for k, v := range ks.P2PInterfaces {
		fmt.Printf("%s %s   %s\n", PadRight(k, " ", 10), PadRight(v[0], " ", 17), v[1])
	}
	// check for DNS leaks
	if ipDNS, err := killswitch.WhoamiDNS(); err == nil {
		if ipWWW, err := killswitch.WhoamiWWW(); err == nil {
			if ipDNS != ipWWW {
				fmt.Printf("\n%s:\n", killswitch.Red("DNS leaking"))
				fmt.Printf("Public IP address (DNS): %s\n", killswitch.Red(ipDNS))
				fmt.Printf("Public IP address (WWW): %s\n", killswitch.Red(ipWWW))
			} else {
				fmt.Printf("\nPublic IP address: %s\n", killswitch.Red(ipDNS))
			}
		}
	}

	// add some space
	println()

	if len(ks.P2PInterfaces) == 0 {
		exit1(fmt.Errorf("No VPN interface found, verify VPN is connected"))
	}

	fmt.Printf("PEER IP address:   %s\n", killswitch.Yellow(ks.PeerIP))

	if *ip != "" {
		if ipv4 := net.ParseIP(*ip); ipv4.To4() == nil {
			exit1(fmt.Errorf("%s is not a valid IPv4 address, use (\"%s -h\") for help.\n", *ip, os.Args[0]))
		}
	}

	ks.CreatePF(*paranoid)

	fmt.Printf("\n%s: %s\n", "To enable the kill switch run", killswitch.Green("sudo killswitch -e"))
	fmt.Printf("%s: %s\n\n", "To disable", killswitch.Yellow("sudo killswitch -d"))

	if *p {
		fmt.Printf("PF rules to be loaded:\n")
		fmt.Println(ks.PFRules.String())
	}

	if err = ioutil.WriteFile("/tmp/killswitch.pf.conf",
		ks.PFRules.Bytes(),
		0644,
	); err != nil {
		exit1(err)
	}

	if *e {
		fmt.Printf("# %s\n", strings.Repeat("-", 62))
		fmt.Println("# Loading rules")
		fmt.Printf("# %s\n", strings.Repeat("-", 62))
		out, _ := exec.Command("pfctl", "-e").CombinedOutput()
		fmt.Printf("%s\n", out)
		out, _ = exec.Command("pfctl",
			"-Fa",
			"-f",
			"/tmp/killswitch.pf.conf").CombinedOutput()
		fmt.Printf("%s\n", out)
		out, _ = exec.Command("pfctl", "-sr").CombinedOutput()
		fmt.Printf("%s\n", out)
	}
}
<commit_msg>update flag msg<commit_after>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"strings"

	"github.com/vpn-kill-switch/killswitch"
)

// PadRight add spaces for aligning the output
func PadRight(str, pad string, length int) string {
	for {
		str += pad
		if len(str) > length {
			return str[0:length]
		}
	}
}

func exit1(err error) {
	fmt.Println(err)
	os.Exit(1)
}

var version string

func main() {

	var (
		ip       = flag.String("ip", "", "VPN peer `IPv4`, killswitch tries to find this automatically")
		d        = flag.Bool("d", false, "`Disable` load /etc/pf.conf rules")
		e        = flag.Bool("e", false, "`Enable` load the pf rules")
		p        = flag.Bool("p", false, "`Print` the pf rules")
		v        = flag.Bool("v", false, fmt.Sprintf("Print version: %s", version))
		paranoid = flag.Bool("paranoid", false, "ICMP (ping) and DNS requests only through the VPN")
	)

	flag.Parse()

	if *v {
		fmt.Printf("%s\n", version)
		os.Exit(0)
	}

	if *d {
		exec.Command("pfctl", "-e").CombinedOutput()
		fmt.Printf("# %s\n", strings.Repeat("-", 62))
		fmt.Println("# Loading /etc/pf.conf rules")
		fmt.Printf("# %s\n", strings.Repeat("-", 62))
		out, _ := exec.Command("pfctl",
			"-Fa",
			"-f",
			"/etc/pf.conf").CombinedOutput()
		fmt.Printf("%s\n", out)
		out, _ = exec.Command("pfctl", "-sr").CombinedOutput()
		fmt.Printf("%s\n", out)
		return
	}

	ks, err := killswitch.New(*ip)
	if err != nil {
		exit1(err)
	}

	err = ks.GetActive()
	if err != nil {
		exit1(err)
	}

	if len(ks.UpInterfaces) == 0 {
		exit1(fmt.Errorf("No active interfaces found, verify network settings, use (\"%s -h\") for help.\n", os.Args[0]))
	}

	fmt.Println("Interface  MAC address         IP")
	for k, v := range ks.UpInterfaces {
		fmt.Printf("%s %s   %s\n", PadRight(k, " ", 10), v[0], v[1])
	}
	for k, v := range ks.P2PInterfaces {
		fmt.Printf("%s %s   %s\n", PadRight(k, " ", 10), PadRight(v[0], " ", 17), v[1])
	}
	// check for DNS leaks
	if ipDNS, err := killswitch.WhoamiDNS(); err == nil {
		if ipWWW, err := killswitch.WhoamiWWW(); err == nil {
			if ipDNS != ipWWW {
				fmt.Printf("\n%s:\n", killswitch.Red("DNS leaking"))
				fmt.Printf("Public IP address (DNS): %s\n", killswitch.Red(ipDNS))
				fmt.Printf("Public IP address (WWW): %s\n", killswitch.Red(ipWWW))
			} else {
				fmt.Printf("\nPublic IP address: %s\n", killswitch.Red(ipDNS))
			}
		}
	}

	// add some space
	println()

	if len(ks.P2PInterfaces) == 0 {
		exit1(fmt.Errorf("No VPN interface found, verify VPN is connected"))
	}

	fmt.Printf("PEER IP address:   %s\n", killswitch.Yellow(ks.PeerIP))

	if *ip != "" {
		if ipv4 := net.ParseIP(*ip); ipv4.To4() == nil {
			exit1(fmt.Errorf("%s is not a valid IPv4 address, use (\"%s -h\") for help.\n", *ip, os.Args[0]))
		}
	}

	ks.CreatePF(*paranoid)

	fmt.Printf("\n%s: %s\n", "To enable the kill switch run", killswitch.Green("sudo killswitch -e"))
	fmt.Printf("%s: %s\n\n", "To disable", killswitch.Yellow("sudo killswitch -d"))

	if *p {
		fmt.Printf("PF rules to be loaded:\n")
		fmt.Println(ks.PFRules.String())
	}

	if err = ioutil.WriteFile("/tmp/killswitch.pf.conf",
		ks.PFRules.Bytes(),
		0644,
	); err != nil {
		exit1(err)
	}

	if *e {
		fmt.Printf("# %s\n", strings.Repeat("-", 62))
		fmt.Println("# Loading rules")
		fmt.Printf("# %s\n", strings.Repeat("-", 62))
		out, _ := exec.Command("pfctl", "-e").CombinedOutput()
		fmt.Printf("%s\n", out)
		out, _ = exec.Command("pfctl",
			"-Fa",
			"-f",
			"/tmp/killswitch.pf.conf").CombinedOutput()
		fmt.Printf("%s\n", out)
		out, _ = exec.Command("pfctl", "-sr").CombinedOutput()
		fmt.Printf("%s\n", out)
	}
}
<|endoftext|>


<commit_before>// Copyright 2017-2019, Square, Inc.

// Package spinc provides a framework for integration with other programs.
package spinc

import (
	"fmt"
	"net/http"
	"os"
	"time"

	rm "github.com/square/spincycle/request-manager"
	"github.com/square/spincycle/spinc/app"
	"github.com/square/spincycle/spinc/cmd"
	"github.com/square/spincycle/spinc/config"
)

// Run runs spinc and exits when done. When using a standard spinc bin, Run is
// called by spinc/bin/main.go. When spinc is wrapped by custom code, that code
// imports this pkg then call spinc.Run() with its custom factories. If a factory
// is not set (nil), then the default/standard factory is used.
func Run(ctx app.Context) error {
	// //////////////////////////////////////////////////////////////////////
	// Config and command line
	// //////////////////////////////////////////////////////////////////////

	// Options are set in this order: config -> env var -> cmd line option.
	// So first we must apply config files, then do cmd line parsing which
	// will apply env vars and cmd line options.

	// Parse cmd line to get --config files
	cmdLine := config.ParseCommandLine(config.Options{})

	// --config files override defaults if given
	configFiles := config.DEFAULT_CONFIG_FILES
	if cmdLine.Config != "" {
		configFiles = cmdLine.Config
	}

	// Parse default options from config files
	def := config.ParseConfigFiles(configFiles, cmdLine.Debug)

	// Parse env vars and cmd line options, override default config
	cmdLine = config.ParseCommandLine(def)

	// Final options and commands
	var o config.Options = cmdLine.Options
	var c config.Command = cmdLine.Command

	// Apply defaults
	if o.Timeout == 0 {
		o.Timeout = config.DEFAULT_TIMEOUT
	}
	if o.Addr == "" {
		o.Addr = config.DEFAULT_ADDR
	}

	// Let hook modify options, if set
	if ctx.Hooks.AfterParseOptions != nil {
		if o.Debug {
			app.Debug("calling hook AfterParseOptions")
		}
		ctx.Hooks.AfterParseOptions(&o)
	}

	// This is a little hack to make spinc -> quick help work, i.e. print
	// quick help when there is no command. We can't check os.Args because
	// it'll be >0 if any flag, like --debug, is specified but we ignore
	// flags. And we can't check c.Cmd == "" because we set c.Cmd = "help".
	ctx.Nargs = len(c.Args) + 1
	if c.Cmd == "" {
		ctx.Nargs -= 1
	}

	// spinc with no args or --help = spinc help
	if len(os.Args) == 1 || o.Help || c.Cmd == "" {
		c.Cmd = "help"
	}

	// --version = spinc version
	if o.Version {
		c.Cmd = "version"
	}

	ctx.Options = o
	ctx.Command = c
	if o.Debug {
		app.Debug("command: %#v\n", c)
		app.Debug("options: %#v\n", o)
	}

	// Use default, built-in command factory if not set by user
	if ctx.Factories.Command == nil {
		ctx.Factories.Command = &cmd.DefaultFactory{}
	}

	// //////////////////////////////////////////////////////////////////////
	// Request Manager Client
	// //////////////////////////////////////////////////////////////////////
	var err error
	ctx.RMClient, err = makeRMC(ctx)
	if err != nil {
		if o.Debug {
			app.Debug("error making RM client: %s", err)
		}
		// All cmds except help and version require an RM client
		if c.Cmd != "help" && c.Cmd != "version" {
			return err
		}
	}

	// //////////////////////////////////////////////////////////////////////
	// Commands
	// //////////////////////////////////////////////////////////////////////

	spincCmd, err := ctx.Factories.Command.Make(c.Cmd, ctx)
	if err != nil {
		switch err {
		case cmd.ErrNotExist:
			return fmt.Errorf("Unknown command: %s. Run 'spinc help' to list commands.", c.Cmd)
		default:
			return fmt.Errorf("Command factory error: %s", err)
		}
	}

	// Let command prepare to run. The start command makes heavy use of this.
	if err := spincCmd.Prepare(); err != nil {
		if o.Debug {
			app.Debug("%s Prepare error: %s", c.Cmd, err)
		}
		switch err {
		case app.ErrUnknownRequest:
			reqName := c.Args[0]
			return fmt.Errorf("Unknown request: %s. Run spinc (no arguments) to list all requests.", reqName)
		default:
			return err
		}
	}

	err = spincCmd.Run()
	if o.Debug {
		app.Debug("%s Run error: %s", c.Cmd, err)
	}
	return err
}

func makeRMC(ctx app.Context) (rm.Client, error) {
	if ctx.Options.Addr == "" {
		return nil, fmt.Errorf("Request Manager API address is not set."+
			" It is best to specify addr in a config file (%s). Or, specify"+
			" --addr on the command line option or set the ADDR environment"+
			" variable. Use --ping to test addr when set.", config.DEFAULT_CONFIG_FILES)
	}
	if ctx.Options.Debug {
		app.Debug("addr: %s", ctx.Options.Addr)
	}
	var httpClient *http.Client
	var err error
	if ctx.Factories.HTTPClient != nil {
		httpClient, err = ctx.Factories.HTTPClient.Make(ctx)
	} else {
		httpClient = &http.Client{
			Timeout: time.Duration(ctx.Options.Timeout) * time.Millisecond,
		}
	}
	if err != nil {
		return nil, fmt.Errorf("Error making http.Client: %s", err)
	}
	rmc := rm.NewClient(httpClient, ctx.Options.Addr)
	return rmc, nil
}
<commit_msg>Move defaults after options hook so --addr can't be empty. Remove a mention to --ping.<commit_after>// Copyright 2017-2019, Square, Inc.

// Package spinc provides a framework for integration with other programs.
package spinc

import (
	"fmt"
	"net/http"
	"os"
	"time"

	rm "github.com/square/spincycle/request-manager"
	"github.com/square/spincycle/spinc/app"
	"github.com/square/spincycle/spinc/cmd"
	"github.com/square/spincycle/spinc/config"
)

// Run runs spinc and exits when done. When using a standard spinc bin, Run is
// called by spinc/bin/main.go. When spinc is wrapped by custom code, that code
// imports this pkg then call spinc.Run() with its custom factories. If a factory
// is not set (nil), then the default/standard factory is used.
func Run(ctx app.Context) error {
	// //////////////////////////////////////////////////////////////////////
	// Config and command line
	// //////////////////////////////////////////////////////////////////////

	// Options are set in this order: config -> env var -> cmd line option.
	// So first we must apply config files, then do cmd line parsing which
	// will apply env vars and cmd line options.

	// Parse cmd line to get --config files
	cmdLine := config.ParseCommandLine(config.Options{})

	// --config files override defaults if given
	configFiles := config.DEFAULT_CONFIG_FILES
	if cmdLine.Config != "" {
		configFiles = cmdLine.Config
	}

	// Parse default options from config files
	def := config.ParseConfigFiles(configFiles, cmdLine.Debug)

	// Parse env vars and cmd line options, override default config
	cmdLine = config.ParseCommandLine(def)

	// Final options and commands
	var o config.Options = cmdLine.Options
	var c config.Command = cmdLine.Command

	// Let hook modify options, if set
	if ctx.Hooks.AfterParseOptions != nil {
		if o.Debug {
			app.Debug("calling hook AfterParseOptions")
		}
		ctx.Hooks.AfterParseOptions(&o)
	}

	// Apply defaults
	if o.Timeout == 0 {
		o.Timeout = config.DEFAULT_TIMEOUT
	}
	if o.Addr == "" {
		o.Addr = config.DEFAULT_ADDR
	}

	// This is a little hack to make spinc -> quick help work, i.e. print
	// quick help when there is no command. We can't check os.Args because
	// it'll be >0 if any flag, like --debug, is specified but we ignore
	// flags. And we can't check c.Cmd == "" because we set c.Cmd = "help".
	ctx.Nargs = len(c.Args) + 1
	if c.Cmd == "" {
		ctx.Nargs -= 1
	}

	// spinc with no args or --help = spinc help
	if len(os.Args) == 1 || o.Help || c.Cmd == "" {
		c.Cmd = "help"
	}

	// --version = spinc version
	if o.Version {
		c.Cmd = "version"
	}

	ctx.Options = o
	ctx.Command = c
	if o.Debug {
		app.Debug("command: %#v\n", c)
		app.Debug("options: %#v\n", o)
	}

	// Use default, built-in command factory if not set by user
	if ctx.Factories.Command == nil {
		ctx.Factories.Command = &cmd.DefaultFactory{}
	}

	// //////////////////////////////////////////////////////////////////////
	// Request Manager Client
	// //////////////////////////////////////////////////////////////////////
	var err error
	ctx.RMClient, err = makeRMC(ctx)
	if err != nil {
		if o.Debug {
			app.Debug("error making RM client: %s", err)
		}
		// All cmds except help and version require an RM client
		if c.Cmd != "help" && c.Cmd != "version" {
			return err
		}
	}

	// //////////////////////////////////////////////////////////////////////
	// Commands
	// //////////////////////////////////////////////////////////////////////

	spincCmd, err := ctx.Factories.Command.Make(c.Cmd, ctx)
	if err != nil {
		switch err {
		case cmd.ErrNotExist:
			return fmt.Errorf("Unknown command: %s. Run 'spinc help' to list commands.", c.Cmd)
		default:
			return fmt.Errorf("Command factory error: %s", err)
		}
	}

	// Let command prepare to run. The start command makes heavy use of this.
	if err := spincCmd.Prepare(); err != nil {
		if o.Debug {
			app.Debug("%s Prepare error: %s", c.Cmd, err)
		}
		switch err {
		case app.ErrUnknownRequest:
			reqName := c.Args[0]
			return fmt.Errorf("Unknown request: %s. Run spinc (no arguments) to list all requests.", reqName)
		default:
			return err
		}
	}

	err = spincCmd.Run()
	if o.Debug {
		app.Debug("%s Run error: %s", c.Cmd, err)
	}
	return err
}

func makeRMC(ctx app.Context) (rm.Client, error) {
	if ctx.Options.Debug {
		app.Debug("addr: %s", ctx.Options.Addr)
	}
	var httpClient *http.Client
	var err error
	if ctx.Factories.HTTPClient != nil {
		httpClient, err = ctx.Factories.HTTPClient.Make(ctx)
	} else {
		httpClient = &http.Client{
			Timeout: time.Duration(ctx.Options.Timeout) * time.Millisecond,
		}
	}
	if err != nil {
		return nil, fmt.Errorf("Error making http.Client: %s", err)
	}
	rmc := rm.NewClient(httpClient, ctx.Options.Addr)
	return rmc, nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2015-2021 MinIO, Inc.
//
// This file is part of MinIO Object Storage stack
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package cmd

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"sync"
	"time"

	"github.com/dustin/go-humanize"
	xhttp "github.com/minio/minio/internal/http"
	"github.com/minio/minio/internal/logger"
)

var printEndpointError = func() func(Endpoint, error, bool) {
	var mutex sync.Mutex
	printOnce := make(map[Endpoint]map[string]int)

	return func(endpoint Endpoint, err error, once bool) {
		reqInfo := (&logger.ReqInfo{}).AppendTags("endpoint", endpoint.String())
		ctx := logger.SetReqInfo(GlobalContext, reqInfo)
		mutex.Lock()
		defer mutex.Unlock()

		m, ok := printOnce[endpoint]
		if !ok {
			m = make(map[string]int)
			m[err.Error()]++
			printOnce[endpoint] = m
			if once {
				logger.LogAlwaysIf(ctx, err)
				return
			}
		}
		// Once is set and we are here means error was already
		// printed once.
		if once {
			return
		}
		// once not set, check if same error occurred 3 times in
		// a row, then make sure we print it to call attention.
		if m[err.Error()] > 2 {
			logger.LogAlwaysIf(ctx, fmt.Errorf("Following error has been printed %d times.. %w", m[err.Error()], err))
			// Reduce the count to introduce further delay in printing
			// but let it again print after the 2th attempt
			m[err.Error()]--
			m[err.Error()]--
		}
		m[err.Error()]++
	}
}()

// Cleans up tmp directory of the local disk.
func formatErasureCleanupTmp(diskPath string) {
	// Need to move temporary objects left behind from previous run of minio
	// server to a unique directory under `minioMetaTmpBucket-old` to clean
	// up `minioMetaTmpBucket` for the current run.
	//
	// /disk1/.minio.sys/tmp-old/
	//  |__ 33a58b40-aecc-4c9f-a22f-ff17bfa33b62
	//  |__ e870a2c1-d09c-450c-a69c-6eaa54a89b3e
	//
	// In this example, `33a58b40-aecc-4c9f-a22f-ff17bfa33b62` directory contains
	// temporary objects from one of the previous runs of minio server.
	tmpID := mustGetUUID()
	tmpOld := pathJoin(diskPath, minioMetaTmpBucket+"-old", tmpID)
	if err := renameAll(pathJoin(diskPath, minioMetaTmpBucket),
		tmpOld); err != nil && !errors.Is(err, errFileNotFound) {
		logger.LogIf(GlobalContext, fmt.Errorf("unable to rename (%s -> %s) %w, drive may be faulty please investigate",
			pathJoin(diskPath, minioMetaTmpBucket),
			tmpOld,
			osErrToFileErr(err)))
	}

	if err := mkdirAll(pathJoin(diskPath, minioMetaTmpDeletedBucket), 0o777); err != nil {
		logger.LogIf(GlobalContext, fmt.Errorf("unable to create (%s) %w, drive may be faulty please investigate",
			pathJoin(diskPath, minioMetaTmpBucket),
			err))
	}

	if err := renameAll(tmpOld, pathJoin(diskPath, minioMetaTmpDeletedBucket, tmpID)); err != nil && !errors.Is(err, errFileNotFound) {
		logger.LogIf(GlobalContext, fmt.Errorf("unable to rename (%s -> %s) %w, drive may be faulty please investigate",
			pathJoin(diskPath, minioMetaTmpBucket),
			tmpOld,
			osErrToFileErr(err)))
	}

	// Renames and schedules for purging all bucket metacache.
	renameAllBucketMetacache(diskPath)
}

// Following error message is added to fix a regression in release
// RELEASE.2018-03-16T22-52-12Z after migrating v1 to v2 to v3. This
// migration failed to capture '.This' field properly which indicates
// the disk UUID association. Below error message is returned when
// we see this situation in format.json, for more info refer
// https://github.com/minio/minio/issues/5667
var errErasureV3ThisEmpty = fmt.Errorf("Erasure format version 3 has This field empty")

// isServerResolvable - checks if the endpoint is resolvable
// by sending a naked HTTP request with liveness checks.
func isServerResolvable(endpoint Endpoint, timeout time.Duration) error {
	serverURL := &url.URL{
		Scheme: endpoint.Scheme,
		Host:   endpoint.Host,
		Path:   pathJoin(healthCheckPathPrefix, healthCheckLivenessPath),
	}

	var tlsConfig *tls.Config
	if globalIsTLS {
		tlsConfig = &tls.Config{
			RootCAs: globalRootCAs,
		}
	}

	httpClient := &http.Client{
		Transport:
		// For more details about various values used here refer
		// https://golang.org/pkg/net/http/#Transport documentation
		&http.Transport{
			Proxy:                 http.ProxyFromEnvironment,
			DialContext:           xhttp.NewCustomDialContext(3 * time.Second),
			ResponseHeaderTimeout: 3 * time.Second,
			TLSHandshakeTimeout:   3 * time.Second,
			ExpectContinueTimeout: 3 * time.Second,
			TLSClientConfig:       tlsConfig,
			// Go net/http automatically unzip if content-type is
			// gzip disable this feature, as we are always interested
			// in raw stream.
			DisableCompression: true,
		},
	}
	defer httpClient.CloseIdleConnections()

	ctx, cancel := context.WithTimeout(GlobalContext, timeout)

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, serverURL.String(), nil)
	if err != nil {
		cancel()
		return err
	}

	resp, err := httpClient.Do(req)
	cancel()
	if err != nil {
		return err
	}
	xhttp.DrainBody(resp.Body)

	return nil
}

// connect to list of endpoints and load all Erasure disk formats, validate the formats are correct
// and are in quorum, if no formats are found attempt to initialize all of them for the first
// time. additionally make sure to close all the disks used in this attempt.
func connectLoadInitFormats(retryCount int, firstDisk bool, endpoints Endpoints, poolCount, setCount, setDriveCount int, deploymentID, distributionAlgo string) (storageDisks []StorageAPI, format *formatErasureV3, err error) {
	// Initialize all storage disks
	storageDisks, errs := initStorageDisksWithErrors(endpoints)

	defer func(storageDisks []StorageAPI) {
		if err != nil {
			closeStorageDisks(storageDisks)
		}
	}(storageDisks)

	for i, err := range errs {
		if err != nil {
			if err == errDiskNotFound && retryCount >= 10 {
				logger.Error("Unable to connect to %s: %v", endpoints[i], isServerResolvable(endpoints[i], time.Second))
			} else {
				logger.Error("Unable to use the drive %s: %v", endpoints[i], err)
			}
		}
	}

	if err := checkDiskFatalErrs(errs); err != nil {
		return nil, nil, err
	}

	// Attempt to load all `format.json` from all disks.
	formatConfigs, sErrs := loadFormatErasureAll(storageDisks, false)
	// Check if we have
	for i, sErr := range sErrs {
		// print the error, nonetheless, which is perhaps unhandled
		if sErr != errUnformattedDisk && sErr != errDiskNotFound && retryCount >= 10 {
			if sErr != nil {
				logger.Error("Unable to read 'format.json' from %s: %v\n", endpoints[i], sErr)
			}
		}
	}

	// Pre-emptively check if one of the formatted disks
	// is invalid. This function returns success for the
	// most part unless one of the formats is not consistent
	// with expected Erasure format. For example if a user is
	// trying to pool FS backend into an Erasure set.
	if err = checkFormatErasureValues(formatConfigs, storageDisks, setDriveCount); err != nil {
		return nil, nil, err
	}

	// All disks report unformatted we should initialized everyone.
	if shouldInitErasureDisks(sErrs) && firstDisk {
		logger.Info("Formatting %s pool, %v set(s), %v drives per set.",
			humanize.Ordinal(poolCount), setCount, setDriveCount)

		// Initialize erasure code format on disks
		format, err = initFormatErasure(GlobalContext, storageDisks, setCount, setDriveCount, deploymentID, distributionAlgo, sErrs)
		if err != nil {
			return nil, nil, err
		}

		// Assign globalDeploymentID on first run for the
		// minio server managing the first disk
		globalDeploymentID = format.ID
		return storageDisks, format, nil
	}

	// Return error when quorum unformatted disks - indicating we are
	// waiting for first server to be online.
	unformattedDisks := quorumUnformattedDisks(sErrs)
	if unformattedDisks && !firstDisk {
		return nil, nil, errNotFirstDisk
	}

	// Return error when quorum unformatted disks but waiting for rest
	// of the servers to be online.
	if unformattedDisks && firstDisk {
		return nil, nil, errFirstDiskWait
	}

	// Mark all root disks down
	markRootDisksAsDown(storageDisks, sErrs)

	// Following function is added to fix a regressions which was introduced
	// in release RELEASE.2018-03-16T22-52-12Z after migrating v1 to v2 to v3.
	// This migration failed to capture '.This' field properly which indicates
	// the disk UUID association. Below function is called to handle and fix
	// this regression, for more info refer https://github.com/minio/minio/issues/5667
	if err = fixFormatErasureV3(storageDisks, endpoints, formatConfigs); err != nil {
		logger.LogIf(GlobalContext, err)
		return nil, nil, err
	}

	// If any of the .This field is still empty, we return error.
	if formatErasureV3ThisEmpty(formatConfigs) {
		return nil, nil, errErasureV3ThisEmpty
	}

	format, err = getFormatErasureInQuorum(formatConfigs)
	if err != nil {
		logger.LogIf(GlobalContext, err)
		return nil, nil, err
	}

	if format.ID == "" {
		// Not a first disk, wait until first disk fixes deploymentID
		if !firstDisk {
			return nil, nil, errNotFirstDisk
		}
		if err = formatErasureFixDeploymentID(endpoints, storageDisks, format); err != nil {
			logger.LogIf(GlobalContext, err)
			return nil, nil, err
		}
	}

	globalDeploymentID = format.ID

	if err = formatErasureFixLocalDeploymentID(endpoints, storageDisks, format); err != nil {
		logger.LogIf(GlobalContext, err)
		return nil, nil, err
	}

	return storageDisks, format, nil
}

// Format disks before initialization of object layer.
func waitForFormatErasure(firstDisk bool, endpoints Endpoints, poolCount, setCount, setDriveCount int, deploymentID, distributionAlgo string) ([]StorageAPI, *formatErasureV3, error) {
	if len(endpoints) == 0 || setCount == 0 || setDriveCount == 0 {
		return nil, nil, errInvalidArgument
	}

	// prepare getElapsedTime() to calculate elapsed time since we started trying formatting disks.
	// All times are rounded to avoid showing milli, micro and nano seconds
	formatStartTime := time.Now().Round(time.Second)
	getElapsedTime := func() string {
		return time.Now().Round(time.Second).Sub(formatStartTime).String()
	}

	var tries int
	storageDisks, format, err := connectLoadInitFormats(tries, firstDisk, endpoints, poolCount, setCount, setDriveCount, deploymentID, distributionAlgo)
	if err == nil {
		return storageDisks, format, nil
	}

	tries++ // tried already once

	// Wait on each try for an update.
	ticker := time.NewTicker(150 * time.Millisecond)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			if tries == 10 {
				// Reset the tries count such that we log only for every 10 retries.
				tries = 1
			}

			storageDisks, format, err := connectLoadInitFormats(tries, firstDisk, endpoints, poolCount, setCount, setDriveCount, deploymentID, distributionAlgo)
			if err != nil {
				tries++
				switch err {
				case errNotFirstDisk:
					// Fresh setup, wait for first server to be up.
					logger.Info("Waiting for the first server to format the disks (elapsed %s)\n", getElapsedTime())
					continue
				case errFirstDiskWait:
					// Fresh setup, wait for other servers to come up.
					logger.Info("Waiting for all other servers to be online to format the disks (elapses %s)\n", getElapsedTime())
					continue
				case errErasureReadQuorum:
					// no quorum available continue to wait for minimum number of servers.
					logger.Info("Waiting for a minimum of %d disks to come online (elapsed %s)\n",
						len(endpoints)/2, getElapsedTime())
					continue
				case errErasureWriteQuorum:
					// no quorum available continue to wait for minimum number of servers.
					logger.Info("Waiting for a minimum of %d disks to come online (elapsed %s)\n",
						(len(endpoints)/2)+1, getElapsedTime())
					continue
				case errErasureV3ThisEmpty:
					// need to wait for this error to be healed, so continue.
					continue
				default:
					// For all other unhandled errors we exit and fail.
					return nil, nil, err
				}
			}
			return storageDisks, format, nil
		case <-globalOSSignalCh:
			return nil, nil, fmt.Errorf("Initializing data volumes gracefully stopped")
		}
	}
}
<commit_msg> fix the error logs have been omitted because of retryCount never exceed 10 (#14268)<commit_after>// Copyright (c) 2015-2021 MinIO, Inc.
//
// This file is part of MinIO Object Storage stack
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package cmd

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"sync"
	"time"

	"github.com/dustin/go-humanize"
	xhttp "github.com/minio/minio/internal/http"
	"github.com/minio/minio/internal/logger"
)

var printEndpointError = func() func(Endpoint, error, bool) {
	var mutex sync.Mutex
	printOnce := make(map[Endpoint]map[string]int)

	return func(endpoint Endpoint, err error, once bool) {
		reqInfo := (&logger.ReqInfo{}).AppendTags("endpoint", endpoint.String())
		ctx := logger.SetReqInfo(GlobalContext, reqInfo)
		mutex.Lock()
		defer mutex.Unlock()

		m, ok := printOnce[endpoint]
		if !ok {
			m = make(map[string]int)
			m[err.Error()]++
			printOnce[endpoint] = m
			if once {
				logger.LogAlwaysIf(ctx, err)
				return
			}
		}
		// Once is set and we are here means error was already
		// printed once.
		if once {
			return
		}
		// once not set, check if same error occurred 3 times in
		// a row, then make sure we print it to call attention.
		if m[err.Error()] > 2 {
			logger.LogAlwaysIf(ctx, fmt.Errorf("Following error has been printed %d times.. %w", m[err.Error()], err))
			// Reduce the count to introduce further delay in printing
			// but let it again print after the 2th attempt
			m[err.Error()]--
			m[err.Error()]--
		}
		m[err.Error()]++
	}
}()

// Cleans up tmp directory of the local disk.
func formatErasureCleanupTmp(diskPath string) {
	// Need to move temporary objects left behind from previous run of minio
	// server to a unique directory under `minioMetaTmpBucket-old` to clean
	// up `minioMetaTmpBucket` for the current run.
	//
	// /disk1/.minio.sys/tmp-old/
	//  |__ 33a58b40-aecc-4c9f-a22f-ff17bfa33b62
	//  |__ e870a2c1-d09c-450c-a69c-6eaa54a89b3e
	//
	// In this example, `33a58b40-aecc-4c9f-a22f-ff17bfa33b62` directory contains
	// temporary objects from one of the previous runs of minio server.
	tmpID := mustGetUUID()
	tmpOld := pathJoin(diskPath, minioMetaTmpBucket+"-old", tmpID)
	if err := renameAll(pathJoin(diskPath, minioMetaTmpBucket),
		tmpOld); err != nil && !errors.Is(err, errFileNotFound) {
		logger.LogIf(GlobalContext, fmt.Errorf("unable to rename (%s -> %s) %w, drive may be faulty please investigate",
			pathJoin(diskPath, minioMetaTmpBucket),
			tmpOld,
			osErrToFileErr(err)))
	}

	if err := mkdirAll(pathJoin(diskPath, minioMetaTmpDeletedBucket), 0o777); err != nil {
		logger.LogIf(GlobalContext, fmt.Errorf("unable to create (%s) %w, drive may be faulty please investigate",
			pathJoin(diskPath, minioMetaTmpBucket),
			err))
	}

	if err := renameAll(tmpOld, pathJoin(diskPath, minioMetaTmpDeletedBucket, tmpID)); err != nil && !errors.Is(err, errFileNotFound) {
		logger.LogIf(GlobalContext, fmt.Errorf("unable to rename (%s -> %s) %w, drive may be faulty please investigate",
			pathJoin(diskPath, minioMetaTmpBucket),
			tmpOld,
			osErrToFileErr(err)))
	}

	// Renames and schedules for purging all bucket metacache.
	renameAllBucketMetacache(diskPath)
}

// Following error message is added to fix a regression in release
// RELEASE.2018-03-16T22-52-12Z after migrating v1 to v2 to v3. This
// migration failed to capture '.This' field properly which indicates
// the disk UUID association. Below error message is returned when
// we see this situation in format.json, for more info refer
// https://github.com/minio/minio/issues/5667
var errErasureV3ThisEmpty = fmt.Errorf("Erasure format version 3 has This field empty")

// isServerResolvable - checks if the endpoint is resolvable
// by sending a naked HTTP request with liveness checks.
func isServerResolvable(endpoint Endpoint, timeout time.Duration) error {
	serverURL := &url.URL{
		Scheme: endpoint.Scheme,
		Host:   endpoint.Host,
		Path:   pathJoin(healthCheckPathPrefix, healthCheckLivenessPath),
	}

	var tlsConfig *tls.Config
	if globalIsTLS {
		tlsConfig = &tls.Config{
			RootCAs: globalRootCAs,
		}
	}

	httpClient := &http.Client{
		Transport:
		// For more details about various values used here refer
		// https://golang.org/pkg/net/http/#Transport documentation
		&http.Transport{
			Proxy:                 http.ProxyFromEnvironment,
			DialContext:           xhttp.NewCustomDialContext(3 * time.Second),
			ResponseHeaderTimeout: 3 * time.Second,
			TLSHandshakeTimeout:   3 * time.Second,
			ExpectContinueTimeout: 3 * time.Second,
			TLSClientConfig:       tlsConfig,
			// Go net/http automatically unzip if content-type is
			// gzip disable this feature, as we are always interested
			// in raw stream.
			DisableCompression: true,
		},
	}
	defer httpClient.CloseIdleConnections()

	ctx, cancel := context.WithTimeout(GlobalContext, timeout)

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, serverURL.String(), nil)
	if err != nil {
		cancel()
		return err
	}

	resp, err := httpClient.Do(req)
	cancel()
	if err != nil {
		return err
	}
	xhttp.DrainBody(resp.Body)

	return nil
}

// connect to list of endpoints and load all Erasure disk formats, validate the formats are correct
// and are in quorum, if no formats are found attempt to initialize all of them for the first
// time. additionally make sure to close all the disks used in this attempt.
func connectLoadInitFormats(verboseLogging bool, firstDisk bool, endpoints Endpoints, poolCount, setCount, setDriveCount int, deploymentID, distributionAlgo string) (storageDisks []StorageAPI, format *formatErasureV3, err error) {
	// Initialize all storage disks
	storageDisks, errs := initStorageDisksWithErrors(endpoints)

	defer func(storageDisks []StorageAPI) {
		if err != nil {
			closeStorageDisks(storageDisks)
		}
	}(storageDisks)

	for i, err := range errs {
		if err != nil {
			if err == errDiskNotFound && verboseLogging {
				logger.Error("Unable to connect to %s: %v", endpoints[i], isServerResolvable(endpoints[i], time.Second))
			} else {
				logger.Error("Unable to use the drive %s: %v", endpoints[i], err)
			}
		}
	}

	if err := checkDiskFatalErrs(errs); err != nil {
		return nil, nil, err
	}

	// Attempt to load all `format.json` from all disks.
	formatConfigs, sErrs := loadFormatErasureAll(storageDisks, false)
	// Check if we have
	for i, sErr := range sErrs {
		// print the error, nonetheless, which is perhaps unhandled
		if sErr != errUnformattedDisk && sErr != errDiskNotFound && verboseLogging {
			if sErr != nil {
				logger.Error("Unable to read 'format.json' from %s: %v\n", endpoints[i], sErr)
			}
		}
	}

	// Pre-emptively check if one of the formatted disks
	// is invalid. This function returns success for the
	// most part unless one of the formats is not consistent
	// with expected Erasure format. For example if a user is
	// trying to pool FS backend into an Erasure set.
	if err = checkFormatErasureValues(formatConfigs, storageDisks, setDriveCount); err != nil {
		return nil, nil, err
	}

	// All disks report unformatted we should initialized everyone.
	if shouldInitErasureDisks(sErrs) && firstDisk {
		logger.Info("Formatting %s pool, %v set(s), %v drives per set.",
			humanize.Ordinal(poolCount), setCount, setDriveCount)

		// Initialize erasure code format on disks
		format, err = initFormatErasure(GlobalContext, storageDisks, setCount, setDriveCount, deploymentID, distributionAlgo, sErrs)
		if err != nil {
			return nil, nil, err
		}

		// Assign globalDeploymentID on first run for the
		// minio server managing the first disk
		globalDeploymentID = format.ID
		return storageDisks, format, nil
	}

	// Return error when quorum unformatted disks - indicating we are
	// waiting for first server to be online.
	unformattedDisks := quorumUnformattedDisks(sErrs)
	if unformattedDisks && !firstDisk {
		return nil, nil, errNotFirstDisk
	}

	// Return error when quorum unformatted disks but waiting for rest
	// of the servers to be online.
	if unformattedDisks && firstDisk {
		return nil, nil, errFirstDiskWait
	}

	// Mark all root disks down
	markRootDisksAsDown(storageDisks, sErrs)

	// Following function is added to fix a regressions which was introduced
	// in release RELEASE.2018-03-16T22-52-12Z after migrating v1 to v2 to v3.
	// This migration failed to capture '.This' field properly which indicates
	// the disk UUID association. Below function is called to handle and fix
	// this regression, for more info refer https://github.com/minio/minio/issues/5667
	if err = fixFormatErasureV3(storageDisks, endpoints, formatConfigs); err != nil {
		logger.LogIf(GlobalContext, err)
		return nil, nil, err
	}

	// If any of the .This field is still empty, we return error.
	if formatErasureV3ThisEmpty(formatConfigs) {
		return nil, nil, errErasureV3ThisEmpty
	}

	format, err = getFormatErasureInQuorum(formatConfigs)
	if err != nil {
		logger.LogIf(GlobalContext, err)
		return nil, nil, err
	}

	if format.ID == "" {
		// Not a first disk, wait until first disk fixes deploymentID
		if !firstDisk {
			return nil, nil, errNotFirstDisk
		}
		if err = formatErasureFixDeploymentID(endpoints, storageDisks, format); err != nil {
			logger.LogIf(GlobalContext, err)
			return nil, nil, err
		}
	}

	globalDeploymentID = format.ID

	if err = formatErasureFixLocalDeploymentID(endpoints, storageDisks, format); err != nil {
		logger.LogIf(GlobalContext, err)
		return nil, nil, err
	}

	return storageDisks, format, nil
}

// Format disks before initialization of object layer.
func waitForFormatErasure(firstDisk bool, endpoints Endpoints, poolCount, setCount, setDriveCount int, deploymentID, distributionAlgo string) ([]StorageAPI, *formatErasureV3, error) {
	if len(endpoints) == 0 || setCount == 0 || setDriveCount == 0 {
		return nil, nil, errInvalidArgument
	}

	// prepare getElapsedTime() to calculate elapsed time since we started trying formatting disks.
	// All times are rounded to avoid showing milli, micro and nano seconds
	formatStartTime := time.Now().Round(time.Second)
	getElapsedTime := func() string {
		return time.Now().Round(time.Second).Sub(formatStartTime).String()
	}

	var tries int
	var verboseLogging bool
	storageDisks, format, err := connectLoadInitFormats(verboseLogging, firstDisk, endpoints, poolCount, setCount, setDriveCount, deploymentID, distributionAlgo)
	if err == nil {
		return storageDisks, format, nil
	}

	tries++ // tried already once

	// Wait on each try for an update.
	ticker := time.NewTicker(150 * time.Millisecond)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			// Only log once every 10 iterations, then reset the tries count.
			verboseLogging = tries >= 10
			if verboseLogging {
				tries = 1
			}

			storageDisks, format, err := connectLoadInitFormats(verboseLogging, firstDisk, endpoints, poolCount, setCount, setDriveCount, deploymentID, distributionAlgo)
			if err != nil {
				tries++
				switch err {
				case errNotFirstDisk:
					// Fresh setup, wait for first server to be up.
					logger.Info("Waiting for the first server to format the disks (elapsed %s)\n", getElapsedTime())
					continue
				case errFirstDiskWait:
					// Fresh setup, wait for other servers to come up.
					logger.Info("Waiting for all other servers to be online to format the disks (elapses %s)\n", getElapsedTime())
					continue
				case errErasureReadQuorum:
					// no quorum available continue to wait for minimum number of servers.
					logger.Info("Waiting for a minimum of %d disks to come online (elapsed %s)\n",
						len(endpoints)/2, getElapsedTime())
					continue
				case errErasureWriteQuorum:
					// no quorum available continue to wait for minimum number of servers.
					logger.Info("Waiting for a minimum of %d disks to come online (elapsed %s)\n",
						(len(endpoints)/2)+1, getElapsedTime())
					continue
				case errErasureV3ThisEmpty:
					// need to wait for this error to be healed, so continue.
					continue
				default:
					// For all other unhandled errors we exit and fail.
					return nil, nil, err
				}
			}
			return storageDisks, format, nil
		case <-globalOSSignalCh:
			return nil, nil, fmt.Errorf("Initializing data volumes gracefully stopped")
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"log"
	"net/http"
	"runtime"
	"sync/atomic"
	"time"
)

var rc *rateCalculator

func statusService(addr string) {
	rc = newRateCalculator(360, 10*time.Second, &bytesProxied)

	http.HandleFunc("/status", getStatus)
	if err := http.ListenAndServe(addr, nil); err != nil {
		log.Fatal(err)
	}
}

func getStatus(w http.ResponseWriter, r *http.Request) {
	status := make(map[string]interface{})

	sessionMut.Lock()
	// This can potentially be double the number of pending sessions, as each session has two keys, one for each side.
	status["startTime"] = rc.startTime
	status["uptimeSeconds"] = time.Since(rc.startTime) / time.Second
	status["numPendingSessionKeys"] = len(pendingSessions)
	status["numActiveSessions"] = len(activeSessions)
	sessionMut.Unlock()
	status["numConnections"] = atomic.LoadInt64(&numConnections)
	status["numProxies"] = atomic.LoadInt64(&numProxies)
	status["bytesProxied"] = atomic.LoadInt64(&bytesProxied)
	status["goVersion"] = runtime.Version()
	status["goOS"] = runtime.GOOS
	status["goAarch"] = runtime.GOARCH
	status["goMaxProcs"] = runtime.GOMAXPROCS(-1)
	status["kbps10s1m5m15m30m60m"] = []int64{
		rc.rate(10/10) * 8 / 1000,
		rc.rate(60/10) * 8 / 1000,
		rc.rate(5*60/10) * 8 / 1000,
		rc.rate(15*60/10) * 8 / 1000,
		rc.rate(30*60/10) * 8 / 1000,
		rc.rate(60*60/10) * 8 / 1000,
	}
	status["options"] = map[string]interface{}{
		"network-timeout":  networkTimeout / time.Second,
		"ping-interval":    pingInterval / time.Second,
		"message-timeout":  messageTimeout / time.Second,
		"per-session-rate": sessionLimitBps,
		"global-rate":      globalLimitBps,
		"pools":            pools,
		"provided-by":      providedBy,
	}

	bs, err := json.MarshalIndent(status, "", "    ")
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.Write(bs)
}

type rateCalculator struct {
	rates     []int64
	prev      int64
	counter   *int64
	startTime time.Time
}

func newRateCalculator(keepIntervals int, interval time.Duration, counter *int64) *rateCalculator {
	r := &rateCalculator{
		rates:     make([]int64, keepIntervals),
		counter:   counter,
		startTime: time.Now(),
	}

	go r.updateRates(interval)

	return r
}

func (r *rateCalculator) updateRates(interval time.Duration) {
	for {
		now := time.Now()
		next := now.Truncate(interval).Add(interval)
		time.Sleep(next.Sub(now))

		cur := atomic.LoadInt64(r.counter)
		rate := int64(float64(cur-r.prev) / interval.Seconds())
		copy(r.rates[1:], r.rates)
		r.rates[0] = rate
		r.prev = cur
	}
}

func (r *rateCalculator) rate(periods int) int64 {
	var tot int64
	for i := 0; i < periods; i++ {
		tot += r.rates[i]
	}
	return tot / int64(periods)
}
<commit_msg>Add CORS headers<commit_after>package main

import (
	"encoding/json"
	"log"
	"net/http"
	"runtime"
	"sync/atomic"
	"time"
)

var rc *rateCalculator

func statusService(addr string) {
	rc = newRateCalculator(360, 10*time.Second, &bytesProxied)

	http.HandleFunc("/status", getStatus)
	if err := http.ListenAndServe(addr, nil); err != nil {
		log.Fatal(err)
	}
}

func getStatus(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Access-Control-Allow-Origin", "*")
	status := make(map[string]interface{})

	sessionMut.Lock()
	// This can potentially be double the number of pending sessions, as each session has two keys, one for each side.
	status["startTime"] = rc.startTime
	status["uptimeSeconds"] = time.Since(rc.startTime) / time.Second
	status["numPendingSessionKeys"] = len(pendingSessions)
	status["numActiveSessions"] = len(activeSessions)
	sessionMut.Unlock()
	status["numConnections"] = atomic.LoadInt64(&numConnections)
	status["numProxies"] = atomic.LoadInt64(&numProxies)
	status["bytesProxied"] = atomic.LoadInt64(&bytesProxied)
	status["goVersion"] = runtime.Version()
	status["goOS"] = runtime.GOOS
	status["goAarch"] = runtime.GOARCH
	status["goMaxProcs"] = runtime.GOMAXPROCS(-1)
	status["kbps10s1m5m15m30m60m"] = []int64{
		rc.rate(10/10) * 8 / 1000,
		rc.rate(60/10) * 8 / 1000,
		rc.rate(5*60/10) * 8 / 1000,
		rc.rate(15*60/10) * 8 / 1000,
		rc.rate(30*60/10) * 8 / 1000,
		rc.rate(60*60/10) * 8 / 1000,
	}
	status["options"] = map[string]interface{}{
		"network-timeout":  networkTimeout / time.Second,
		"ping-interval":    pingInterval / time.Second,
		"message-timeout":  messageTimeout / time.Second,
		"per-session-rate": sessionLimitBps,
		"global-rate":      globalLimitBps,
		"pools":            pools,
		"provided-by":      providedBy,
	}

	bs, err := json.MarshalIndent(status, "", "    ")
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.Write(bs)
}

type rateCalculator struct {
	rates     []int64
	prev      int64
	counter   *int64
	startTime time.Time
}

func newRateCalculator(keepIntervals int, interval time.Duration, counter *int64) *rateCalculator {
	r := &rateCalculator{
		rates:     make([]int64, keepIntervals),
		counter:   counter,
		startTime: time.Now(),
	}

	go r.updateRates(interval)

	return r
}

func (r *rateCalculator) updateRates(interval time.Duration) {
	for {
		now := time.Now()
		next := now.Truncate(interval).Add(interval)
		time.Sleep(next.Sub(now))

		cur := atomic.LoadInt64(r.counter)
		rate := int64(float64(cur-r.prev) / interval.Seconds())
		copy(r.rates[1:], r.rates)
		r.rates[0] = rate
		r.prev = cur
	}
}

func (r *rateCalculator) rate(periods int) int64 {
	var tot int64
	for i := 0; i < periods; i++ {
		tot += r.rates[i]
	}
	return tot / int64(periods)
}
<|endoftext|>


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

// The skylark command interprets a Skylark file.
//
// With no arguments, it starts a read-eval-print loop (REPL).
// If an input line can be parsed as an expression,
// the REPL parses and evaluates it and prints its result.
// Otherwise the REPL reads lines until a blank line,
// then tries again to parse the multi-line input as an
// expression. If the input still cannot be parsed as an expression,
// the REPL parses and executes it as a file (a list of statements),
// for side effects.
package main

// TODO(adonovan):
//
// - Distinguish expressions from statements more precisely.
//   Otherwise e.g. 1 is parsed as an expression but
//   1000000000000000000000000000 is parsed as a file
//   because the scanner fails to convert it to an int64.
//   The spec should clarify limits on numeric literals.
//
// - Unparenthesized tuples are not parsed as a single expression:
//     >>> (1, 2)
//     (1, 2)
//     >>> 1, 2
//     ...
//     >>>
//   This is not necessarily a bug.

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"runtime/pprof"
	"sort"
	"strings"

	"github.com/google/skylark"
	"github.com/google/skylark/resolve"
	"github.com/google/skylark/syntax"
)

// flags
var (
	cpuprofile = flag.String("cpuprofile", "", "gather CPU profile in this file")
	showenv    = flag.Bool("showenv", false, "on success, print final global environment")
)

// non-standard dialect flags
func init() {
	flag.BoolVar(&resolve.AllowFloat, "fp", resolve.AllowFloat, "allow floating-point numbers")
	flag.BoolVar(&resolve.AllowFreeze, "freeze", resolve.AllowFreeze, "add freeze built-in function")
	flag.BoolVar(&resolve.AllowSet, "set", resolve.AllowSet, "allow set data type")
	flag.BoolVar(&resolve.AllowLambda, "lambda", resolve.AllowLambda, "allow lambda expressions")
	flag.BoolVar(&resolve.AllowNestedDef, "nesteddef", resolve.AllowNestedDef, "allow nested def statements")
}

func main() {
	log.SetPrefix("skylark: ")
	log.SetFlags(0)
	flag.Parse()

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		if err := pprof.StartCPUProfile(f); err != nil {
			log.Fatal(err)
		}
		defer pprof.StopCPUProfile()
	}

	switch len(flag.Args()) {
	case 0:
		repl()
	case 1:
		execfile(flag.Args()[0])
	default:
		log.Fatal("want at most one Skylark file name")
	}
}

func execfile(filename string) {
	thread := &skylark.Thread{Load: load}
	globals := make(skylark.StringDict)
	if err := skylark.ExecFile(thread, filename, nil, globals); err != nil {
		printError(err)
		os.Exit(1)
	}

	// Print the global environment.
	if *showenv {
		var names []string
		for name := range globals {
			if !strings.HasPrefix(name, "_") {
				names = append(names, name)
			}
		}
		sort.Strings(names)
		for _, name := range names {
			fmt.Fprintf(os.Stderr, "%s = %s\n", name, globals[name])
		}
	}
}

func repl() {
	thread := &skylark.Thread{Load: load}
	globals := make(skylark.StringDict)

	sc := bufio.NewScanner(os.Stdin)
outer:
	for {
		io.WriteString(os.Stderr, ">>> ")
		if !sc.Scan() {
			break
		}
		line := sc.Text()
		if l := strings.TrimSpace(line); l == "" || l[0] == '#' {
			continue // blank or comment
		}

		// If the line contains a well-formed
		// expression, evaluate it.
		if expr, err := syntax.ParseExpr("<stdin>", line); err == nil && !isLoad(expr) {
			if v, err := skylark.Eval(thread, "<stdin>", line, globals); err != nil {
				printError(err)
			} else if v != skylark.None {
				fmt.Println(v)
			}
			continue
		}

		// Otherwise assume it is the first of several
		// comprising a file, followed by a blank line.
		var buf bytes.Buffer
		fmt.Fprintln(&buf, line)
		for {
			io.WriteString(os.Stderr, "... ")
			if !sc.Scan() {
				break outer
			}
			line := sc.Text()
			if l := strings.TrimSpace(line); l == "" {
				break // blank
			}
			fmt.Fprintln(&buf, line)
		}
		text := buf.Bytes()

		// Try parsing it once more as an expression,
		// such as a call spread over several lines:
		//   f(
		//     1,
		//     2
		//   )
		if expr, err := syntax.ParseExpr("<stdin>", text); err == nil && !isLoad(expr) {
			if v, err := skylark.Eval(thread, "<stdin>", text, globals); err != nil {
				printError(err)
			} else if v != skylark.None {
				fmt.Println(v)
			}
			continue
		}

		// Execute it as a file.
		if err := execFileNoFreeze(thread, text, globals); err != nil {
			printError(err)
		}
	}
	fmt.Println()
}

// execFileNoFreeze is skylark.ExecFile without globals.Freeze().
func execFileNoFreeze(thread *skylark.Thread, src interface{}, globals skylark.StringDict) error {
	// parse
	f, err := syntax.Parse("<stdin>", src)
	if err != nil {
		return err
	}

	// resolve
	if err := resolve.File(f, globals.Has, skylark.Universe.Has); err != nil {
		return err

	}

	// execute
	fr := thread.Push(globals, len(f.Locals))
	defer thread.Pop()
	return fr.ExecStmts(f.Stmts)
}

type entry struct {
	globals skylark.StringDict
	err     error
}

var cache = make(map[string]*entry)

// load is a simple sequential implementation of module loading.
func load(thread *skylark.Thread, module string) (skylark.StringDict, error) {
	e, ok := cache[module]
	if e == nil {
		if ok {
			// request for package whose loading is in progress
			return nil, fmt.Errorf("cycle in load graph")
		}

		// Add a placeholder to indicate "load in progress".
		cache[module] = nil

		// Load it.
		thread := &skylark.Thread{Load: load}
		globals := make(skylark.StringDict)
		err := skylark.ExecFile(thread, module, nil, globals)
		e = &entry{globals, err}

		// Update the cache.
		cache[module] = e
	}
	return e.globals, e.err
}

func printError(err error) {
	if evalErr, ok := err.(*skylark.EvalError); ok {
		fmt.Fprintln(os.Stderr, evalErr.Backtrace())
	} else {
		fmt.Fprintln(os.Stderr, err)
	}
}

// isLoad reports whether e is a load(...) function call.
// If so, we must parse it again as a file, not an expression,
// so that it is is converted to a load statement.
// ("load" should really be a reserved word.)
func isLoad(e syntax.Expr) bool {
	if call, ok := e.(*syntax.CallExpr); ok {
		if id, ok := call.Fn.(*syntax.Ident); ok && id.Name == "load" {
			return true
		}
	}
	return false
}
<commit_msg>cmd/skylark: process a single-line load statement without delay (#13)<commit_after>// Copyright 2017 The Bazel Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// The skylark command interprets a Skylark file.
//
// With no arguments, it starts a read-eval-print loop (REPL).
// If an input line can be parsed as an expression,
// the REPL parses and evaluates it and prints its result.
// Otherwise the REPL reads lines until a blank line,
// then tries again to parse the multi-line input as an
// expression. If the input still cannot be parsed as an expression,
// the REPL parses and executes it as a file (a list of statements),
// for side effects.
package main

// TODO(adonovan):
//
// - Distinguish expressions from statements more precisely.
//   Otherwise e.g. 1 is parsed as an expression but
//   1000000000000000000000000000 is parsed as a file
//   because the scanner fails to convert it to an int64.
//   The spec should clarify limits on numeric literals.
//
// - Unparenthesized tuples are not parsed as a single expression:
//     >>> (1, 2)
//     (1, 2)
//     >>> 1, 2
//     ...
//     >>>
//   This is not necessarily a bug.

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"runtime/pprof"
	"sort"
	"strings"

	"github.com/google/skylark"
	"github.com/google/skylark/resolve"
	"github.com/google/skylark/syntax"
)

// flags
var (
	cpuprofile = flag.String("cpuprofile", "", "gather CPU profile in this file")
	showenv    = flag.Bool("showenv", false, "on success, print final global environment")
)

// non-standard dialect flags
func init() {
	flag.BoolVar(&resolve.AllowFloat, "fp", resolve.AllowFloat, "allow floating-point numbers")
	flag.BoolVar(&resolve.AllowFreeze, "freeze", resolve.AllowFreeze, "add freeze built-in function")
	flag.BoolVar(&resolve.AllowSet, "set", resolve.AllowSet, "allow set data type")
	flag.BoolVar(&resolve.AllowLambda, "lambda", resolve.AllowLambda, "allow lambda expressions")
	flag.BoolVar(&resolve.AllowNestedDef, "nesteddef", resolve.AllowNestedDef, "allow nested def statements")
}

func main() {
	log.SetPrefix("skylark: ")
	log.SetFlags(0)
	flag.Parse()

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		if err := pprof.StartCPUProfile(f); err != nil {
			log.Fatal(err)
		}
		defer pprof.StopCPUProfile()
	}

	switch len(flag.Args()) {
	case 0:
		repl()
	case 1:
		execfile(flag.Args()[0])
	default:
		log.Fatal("want at most one Skylark file name")
	}
}

func execfile(filename string) {
	thread := &skylark.Thread{Load: load}
	globals := make(skylark.StringDict)
	if err := skylark.ExecFile(thread, filename, nil, globals); err != nil {
		printError(err)
		os.Exit(1)
	}

	// Print the global environment.
	if *showenv {
		var names []string
		for name := range globals {
			if !strings.HasPrefix(name, "_") {
				names = append(names, name)
			}
		}
		sort.Strings(names)
		for _, name := range names {
			fmt.Fprintf(os.Stderr, "%s = %s\n", name, globals[name])
		}
	}
}

func repl() {
	thread := &skylark.Thread{Load: load}
	globals := make(skylark.StringDict)

	sc := bufio.NewScanner(os.Stdin)
outer:
	for {
		io.WriteString(os.Stderr, ">>> ")
		if !sc.Scan() {
			break
		}
		line := sc.Text()
		if l := strings.TrimSpace(line); l == "" || l[0] == '#' {
			continue // blank or comment
		}

		// If the line contains a well-formed
		// expression, evaluate it.
		if expr, err := syntax.ParseExpr("<stdin>", line); err == nil {
			if isLoad(expr) {
				if err := execFileNoFreeze(thread, line, globals); err != nil {
					printError(err)
				}
			} else if v, err := skylark.Eval(thread, "<stdin>", line, globals); err != nil {
				printError(err)
			} else if v != skylark.None {
				fmt.Println(v)
			}
			continue
		}

		// Otherwise assume it is the first of several
		// comprising a file, followed by a blank line.
		var buf bytes.Buffer
		fmt.Fprintln(&buf, line)
		for {
			io.WriteString(os.Stderr, "... ")
			if !sc.Scan() {
				break outer
			}
			line := sc.Text()
			if l := strings.TrimSpace(line); l == "" {
				break // blank
			}
			fmt.Fprintln(&buf, line)
		}
		text := buf.Bytes()

		// Try parsing it once more as an expression,
		// such as a call spread over several lines:
		//   f(
		//     1,
		//     2
		//   )
		if expr, err := syntax.ParseExpr("<stdin>", text); err == nil && !isLoad(expr) {
			if v, err := skylark.Eval(thread, "<stdin>", text, globals); err != nil {
				printError(err)
			} else if v != skylark.None {
				fmt.Println(v)
			}
			continue
		}

		// Execute it as a file.
		if err := execFileNoFreeze(thread, text, globals); err != nil {
			printError(err)
		}
	}
	fmt.Println()
}

// execFileNoFreeze is skylark.ExecFile without globals.Freeze().
func execFileNoFreeze(thread *skylark.Thread, src interface{}, globals skylark.StringDict) error {
	// parse
	f, err := syntax.Parse("<stdin>", src)
	if err != nil {
		return err
	}

	// resolve
	if err := resolve.File(f, globals.Has, skylark.Universe.Has); err != nil {
		return err

	}

	// execute
	fr := thread.Push(globals, len(f.Locals))
	defer thread.Pop()
	return fr.ExecStmts(f.Stmts)
}

type entry struct {
	globals skylark.StringDict
	err     error
}

var cache = make(map[string]*entry)

// load is a simple sequential implementation of module loading.
func load(thread *skylark.Thread, module string) (skylark.StringDict, error) {
	e, ok := cache[module]
	if e == nil {
		if ok {
			// request for package whose loading is in progress
			return nil, fmt.Errorf("cycle in load graph")
		}

		// Add a placeholder to indicate "load in progress".
		cache[module] = nil

		// Load it.
		thread := &skylark.Thread{Load: load}
		globals := make(skylark.StringDict)
		err := skylark.ExecFile(thread, module, nil, globals)
		e = &entry{globals, err}

		// Update the cache.
		cache[module] = e
	}
	return e.globals, e.err
}

func printError(err error) {
	if evalErr, ok := err.(*skylark.EvalError); ok {
		fmt.Fprintln(os.Stderr, evalErr.Backtrace())
	} else {
		fmt.Fprintln(os.Stderr, err)
	}
}

// isLoad reports whether e is a load(...) function call.
// If so, we must parse it again as a file, not an expression,
// so that it is is converted to a load statement.
// ("load" should really be a reserved word.)
func isLoad(e syntax.Expr) bool {
	if call, ok := e.(*syntax.CallExpr); ok {
		if id, ok := call.Fn.(*syntax.Ident); ok && id.Name == "load" {
			return true
		}
	}
	return false
}
<|endoftext|>