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

package main

import (
	"fmt"
	"launchpad.net/gnuflag"
	"launchpad.net/juju-core/charm"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/log"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/api"
	"launchpad.net/juju-core/state/api/machineagent"
	"launchpad.net/juju-core/state/api/params"
	"launchpad.net/juju-core/state/apiserver"
	"launchpad.net/juju-core/worker"
	"launchpad.net/juju-core/worker/firewaller"
	"launchpad.net/juju-core/worker/machiner"
	"launchpad.net/juju-core/worker/provisioner"
	"launchpad.net/tomb"
	"path/filepath"
	"time"
)

var retryDelay = 3 * time.Second

// MachineAgent is a cmd.Command responsible for running a machine agent.
type MachineAgent struct {
	cmd.CommandBase
	tomb      tomb.Tomb
	Conf      AgentConf
	MachineId string
	runner    *worker.Runner
}

// Info returns usage information for the command.
func (a *MachineAgent) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "machine",
		Purpose: "run a juju machine agent",
	}
}

func (a *MachineAgent) SetFlags(f *gnuflag.FlagSet) {
	a.Conf.addFlags(f)
	f.StringVar(&a.MachineId, "machine-id", "", "id of the machine to run")
}

// Init initializes the command for running.
func (a *MachineAgent) Init(args []string) error {
	if !state.IsMachineId(a.MachineId) {
		return fmt.Errorf("--machine-id option must be set, and expects a non-negative integer")
	}
	if err := a.Conf.checkArgs(args); err != nil {
		return err
	}
	a.runner = worker.NewRunner(isFatal, moreImportant)
	return nil
}

// Wait waits for the machine agent to finish.
func (a *MachineAgent) Wait() error {
	return a.tomb.Wait()
}

// Stop stops the machine agent.
func (a *MachineAgent) Stop() error {
	a.runner.Kill()
	return a.tomb.Wait()
}

// Run runs a machine agent.
func (a *MachineAgent) Run(_ *cmd.Context) error {
	defer a.tomb.Done()
	log.Infof("machine agent %v start", a.Tag())
	if err := a.Conf.read(a.Tag()); err != nil {
		return err
	}
	charm.CacheDir = filepath.Join(a.Conf.DataDir, "charmcache")

	// ensureStateWorker ensures that there is a worker that
	// connects to the state that runs within itself all the workers
	// that need a state connection Unless we're bootstrapping, we
	// need to connect to the API server to find out if we need to
	// call this, so we make the APIWorker call it when necessary if
	// the machine requires it.  Note that startStateWorker can be
	// called many times - StartWorker does nothing if there is
	// already a worker started with the given name.
	ensureStateWorker := func() {
		a.runner.StartWorker("state", func() (worker.Worker, error) {
			// TODO(rog) go1.1: use method expression
			return a.StateWorker()
		})
	}
	if a.MachineId == "0" {
		// If we're bootstrapping, we don't have an API
		// server to connect to, so start the state worker regardless.

		// TODO(rog) When we have HA, we only want to do this
		// when we really are bootstrapping - once other
		// instances of the API server have been started, we
		// should follow the normal course of things and ignore
		// the fact that this was once the bootstrap machine.
		ensureStateWorker()
	}
	a.runner.StartWorker("api", func() (worker.Worker, error) {
		// TODO(rog) go1.1: use method expression
		return a.APIWorker(ensureStateWorker)
	})
	err := agentDone(a.runner.Wait())
	a.tomb.Kill(err)
	return err
}

func allFatal(error) bool {
	return true
}

var stateJobs = map[params.MachineJob]bool{
	params.JobHostUnits:     true,
	params.JobManageEnviron: true,
	params.JobManageState:   true,
}

// APIWorker returns a Worker that connects to the API and starts any
// workers that need an API connection.
//
// If a state worker is necessary, APIWorker calls startStateWorker.
func (a *MachineAgent) APIWorker(ensureStateWorker func()) (worker.Worker, error) {
	st, entity, err := openAPIState(a.Conf.Conf, a)
	if err != nil {
		return nil, err
	}
	m := entity.(*machineagent.Machine)
	needsStateWorker := false
	for _, job := range m.Jobs() {
		needsStateWorker = needsStateWorker || stateJobs[job]
	}
	if needsStateWorker {
		ensureStateWorker()
	}
	runner := worker.NewRunner(allFatal, moreImportant)
	// No agents currently connect to the API, so just
	// return the runner running nothing.
	return newCloseWorker(runner, st), nil // Note: a worker.Runner is itself a worker.Worker.
}

// StateJobs returns a worker running all the workers that require
// a *state.State connection.
func (a *MachineAgent) StateWorker() (worker.Worker, error) {
	st, entity, err := openState(a.Conf.Conf, a)
	if err != nil {
		return nil, err
	}
	m := entity.(*state.Machine)
	// TODO(rog) use more discriminating test for errors
	// rather than taking everything down indiscriminately.
	runner := worker.NewRunner(allFatal, moreImportant)
	runner.StartWorker("upgrader", func() (worker.Worker, error) {
		// TODO(rog) use id instead of *Machine (or introduce Clone method)
		return NewUpgrader(st, m, a.Conf.DataDir), nil
	})
	runner.StartWorker("machiner", func() (worker.Worker, error) {
		return machiner.NewMachiner(st, m.Id()), nil
	})
	for _, job := range m.Jobs() {
		switch job {
		case state.JobHostUnits:
			runner.StartWorker("deployer", func() (worker.Worker, error) {
				return newDeployer(st, m.WatchPrincipalUnits(), a.Conf.DataDir), nil
			})
		case state.JobManageEnviron:
			runner.StartWorker("provisioner", func() (worker.Worker, error) {
				return provisioner.NewProvisioner(st, a.MachineId), nil
			})
			runner.StartWorker("firewaller", func() (worker.Worker, error) {
				return firewaller.NewFirewaller(st), nil
			})
		case state.JobManageState:
			runner.StartWorker("apiserver", func() (worker.Worker, error) {
				// If the configuration does not have the required information,
				// it is currently not a recoverable error, so we kill the whole
				// agent, potentially enabling human intervention to fix
				// the agent's configuration file. In the future, we may retrieve
				// the state server certificate and key from the state, and
				// this should then change.
				if len(a.Conf.StateServerCert) == 0 || len(a.Conf.StateServerKey) == 0 {
					return nil, &fatalError{"configuration does not have state server cert/key"}
				}
				return apiserver.NewServer(st, fmt.Sprintf(":%d", a.Conf.APIPort), a.Conf.StateServerCert, a.Conf.StateServerKey)
			})
		default:
			log.Warningf("ignoring unknown job %q", job)
		}
	}
	return newCloseWorker(runner, st), nil
}

func (a *MachineAgent) Entity(st *state.State) (AgentState, error) {
	m, err := st.Machine(a.MachineId)
	if err != nil {
		return nil, err
	}
	// Check the machine nonce as provisioned matches the agent.Conf value.
	if !m.CheckProvisioned(a.Conf.MachineNonce) {
		// The agent is running on a different machine to the one it
		// should be according to state. It must stop immediately.
		log.Errorf("running machine %v agent on inappropriate instance", m)
		return nil, worker.ErrTerminateAgent
	}
	return m, nil
}

func (a *MachineAgent) APIEntity(st *api.State) (AgentAPIState, error) {
	m, err := st.MachineAgent().Machine(a.MachineId)
	if err != nil {
		return nil, err
	}
	// TODO(rog) move the CheckProvisioned test into
	// this method when it's implemented in the API
	return m, nil
}

func (a *MachineAgent) Tag() string {
	return state.MachineTag(a.MachineId)
}
<commit_msg>cmd/jujud: fix comments<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package main

import (
	"fmt"
	"launchpad.net/gnuflag"
	"launchpad.net/juju-core/charm"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/log"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/api"
	"launchpad.net/juju-core/state/api/machineagent"
	"launchpad.net/juju-core/state/api/params"
	"launchpad.net/juju-core/state/apiserver"
	"launchpad.net/juju-core/worker"
	"launchpad.net/juju-core/worker/firewaller"
	"launchpad.net/juju-core/worker/machiner"
	"launchpad.net/juju-core/worker/provisioner"
	"launchpad.net/tomb"
	"path/filepath"
	"time"
)

var retryDelay = 3 * time.Second

// MachineAgent is a cmd.Command responsible for running a machine agent.
type MachineAgent struct {
	cmd.CommandBase
	tomb      tomb.Tomb
	Conf      AgentConf
	MachineId string
	runner    *worker.Runner
}

// Info returns usage information for the command.
func (a *MachineAgent) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "machine",
		Purpose: "run a juju machine agent",
	}
}

func (a *MachineAgent) SetFlags(f *gnuflag.FlagSet) {
	a.Conf.addFlags(f)
	f.StringVar(&a.MachineId, "machine-id", "", "id of the machine to run")
}

// Init initializes the command for running.
func (a *MachineAgent) Init(args []string) error {
	if !state.IsMachineId(a.MachineId) {
		return fmt.Errorf("--machine-id option must be set, and expects a non-negative integer")
	}
	if err := a.Conf.checkArgs(args); err != nil {
		return err
	}
	a.runner = worker.NewRunner(isFatal, moreImportant)
	return nil
}

// Wait waits for the machine agent to finish.
func (a *MachineAgent) Wait() error {
	return a.tomb.Wait()
}

// Stop stops the machine agent.
func (a *MachineAgent) Stop() error {
	a.runner.Kill()
	return a.tomb.Wait()
}

// Run runs a machine agent.
func (a *MachineAgent) Run(_ *cmd.Context) error {
	defer a.tomb.Done()
	log.Infof("machine agent %v start", a.Tag())
	if err := a.Conf.read(a.Tag()); err != nil {
		return err
	}
	charm.CacheDir = filepath.Join(a.Conf.DataDir, "charmcache")

	// ensureStateWorker ensures that there is a worker that
	// connects to the state that runs within itself all the workers
	// that need a state connection Unless we're bootstrapping, we
	// need to connect to the API server to find out if we need to
	// call this, so we make the APIWorker call it when necessary if
	// the machine requires it.  Note that ensureStateWorker can be
	// called many times - StartWorker does nothing if there is
	// already a worker started with the given name.
	ensureStateWorker := func() {
		a.runner.StartWorker("state", func() (worker.Worker, error) {
			// TODO(rog) go1.1: use method expression
			return a.StateWorker()
		})
	}
	if a.MachineId == "0" {
		// If we're bootstrapping, we don't have an API
		// server to connect to, so start the state worker regardless.

		// TODO(rog) When we have HA, we only want to do this
		// when we really are bootstrapping - once other
		// instances of the API server have been started, we
		// should follow the normal course of things and ignore
		// the fact that this was once the bootstrap machine.
		ensureStateWorker()
	}
	a.runner.StartWorker("api", func() (worker.Worker, error) {
		// TODO(rog) go1.1: use method expression
		return a.APIWorker(ensureStateWorker)
	})
	err := agentDone(a.runner.Wait())
	a.tomb.Kill(err)
	return err
}

func allFatal(error) bool {
	return true
}

var stateJobs = map[params.MachineJob]bool{
	params.JobHostUnits:     true,
	params.JobManageEnviron: true,
	params.JobManageState:   true,
}

// APIWorker returns a Worker that connects to the API and starts any
// workers that need an API connection.
//
// If a state worker is necessary, APIWorker calls ensureStateWorker.
func (a *MachineAgent) APIWorker(ensureStateWorker func()) (worker.Worker, error) {
	st, entity, err := openAPIState(a.Conf.Conf, a)
	if err != nil {
		return nil, err
	}
	m := entity.(*machineagent.Machine)
	needsStateWorker := false
	for _, job := range m.Jobs() {
		needsStateWorker = needsStateWorker || stateJobs[job]
	}
	if needsStateWorker {
		ensureStateWorker()
	}
	runner := worker.NewRunner(allFatal, moreImportant)
	// No agents currently connect to the API, so just
	// return the runner running nothing.
	return newCloseWorker(runner, st), nil // Note: a worker.Runner is itself a worker.Worker.
}

// StateJobs returns a worker running all the workers that require
// a *state.State connection.
func (a *MachineAgent) StateWorker() (worker.Worker, error) {
	st, entity, err := openState(a.Conf.Conf, a)
	if err != nil {
		return nil, err
	}
	m := entity.(*state.Machine)
	// TODO(rog) use more discriminating test for errors
	// rather than taking everything down indiscriminately.
	runner := worker.NewRunner(allFatal, moreImportant)
	runner.StartWorker("upgrader", func() (worker.Worker, error) {
		// TODO(rog) use id instead of *Machine (or introduce Clone method)
		return NewUpgrader(st, m, a.Conf.DataDir), nil
	})
	runner.StartWorker("machiner", func() (worker.Worker, error) {
		return machiner.NewMachiner(st, m.Id()), nil
	})
	for _, job := range m.Jobs() {
		switch job {
		case state.JobHostUnits:
			runner.StartWorker("deployer", func() (worker.Worker, error) {
				return newDeployer(st, m.WatchPrincipalUnits(), a.Conf.DataDir), nil
			})
		case state.JobManageEnviron:
			runner.StartWorker("provisioner", func() (worker.Worker, error) {
				return provisioner.NewProvisioner(st, a.MachineId), nil
			})
			runner.StartWorker("firewaller", func() (worker.Worker, error) {
				return firewaller.NewFirewaller(st), nil
			})
		case state.JobManageState:
			runner.StartWorker("apiserver", func() (worker.Worker, error) {
				// If the configuration does not have the required information,
				// it is currently not a recoverable error, so we kill the whole
				// agent, potentially enabling human intervention to fix
				// the agent's configuration file. In the future, we may retrieve
				// the state server certificate and key from the state, and
				// this should then change.
				if len(a.Conf.StateServerCert) == 0 || len(a.Conf.StateServerKey) == 0 {
					return nil, &fatalError{"configuration does not have state server cert/key"}
				}
				return apiserver.NewServer(st, fmt.Sprintf(":%d", a.Conf.APIPort), a.Conf.StateServerCert, a.Conf.StateServerKey)
			})
		default:
			log.Warningf("ignoring unknown job %q", job)
		}
	}
	return newCloseWorker(runner, st), nil
}

func (a *MachineAgent) Entity(st *state.State) (AgentState, error) {
	m, err := st.Machine(a.MachineId)
	if err != nil {
		return nil, err
	}
	// Check the machine nonce as provisioned matches the agent.Conf value.
	if !m.CheckProvisioned(a.Conf.MachineNonce) {
		// The agent is running on a different machine to the one it
		// should be according to state. It must stop immediately.
		log.Errorf("running machine %v agent on inappropriate instance", m)
		return nil, worker.ErrTerminateAgent
	}
	return m, nil
}

func (a *MachineAgent) APIEntity(st *api.State) (AgentAPIState, error) {
	m, err := st.MachineAgent().Machine(a.MachineId)
	if err != nil {
		return nil, err
	}
	// TODO(rog) move the CheckProvisioned test into
	// this method when it's implemented in the API
	return m, nil
}

func (a *MachineAgent) Tag() string {
	return state.MachineTag(a.MachineId)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"time"

	"github.com/juju/errors"
	"github.com/juju/names"
	"github.com/juju/utils"

	"github.com/juju/juju/agent"
	"github.com/juju/juju/environs"
	"github.com/juju/juju/mongo"
	"github.com/juju/juju/state"
	"github.com/juju/juju/state/api"
	"github.com/juju/juju/state/api/params"
	"github.com/juju/juju/upgrades"
	"github.com/juju/juju/version"
	"github.com/juju/juju/worker"
)

type upgradingMachineAgent interface {
	ensureMongoServer(agent.Config) error
	setMachineStatus(*api.State, params.Status, string) error
	CurrentConfig() agent.Config
	ChangeConfig(AgentConfigMutator) error
}

var upgradesPerformUpgrade = upgrades.PerformUpgrade // Allow patching for tests

func NewUpgradeWorkerContext() *upgradeWorkerContext {
	return &upgradeWorkerContext{
		UpgradeComplete: make(chan struct{}),
	}
}

type upgradeWorkerContext struct {
	UpgradeComplete chan struct{}
	agent           upgradingMachineAgent
	apiState        *api.State
	jobs            []params.MachineJob
	agentConfig     agent.Config
	isStateServer   bool
	st              *state.State
}

// InitialiseUsingAgent sets up a upgradeWorkerContext from a machine agent instance.
// It may update the agent's configuration.
func (c *upgradeWorkerContext) InitializeUsingAgent(a upgradingMachineAgent) error {
	return a.ChangeConfig(func(agentConfig agent.ConfigSetter) error {
		if !upgrades.AreUpgradesDefined(agentConfig.UpgradedToVersion()) {
			logger.Infof("no upgrade steps required or upgrade steps for %v "+
				"have already been run.", version.Current.Number)
			close(c.UpgradeComplete)

			// Even if no upgrade is required the version number in
			// the agent's config still needs to be bumped.
			agentConfig.SetUpgradedToVersion(version.Current.Number)
		}
		return nil
	})
}

func (c *upgradeWorkerContext) Worker(
	agent upgradingMachineAgent,
	apiState *api.State,
	jobs []params.MachineJob,
) worker.Worker {
	c.agent = agent
	c.apiState = apiState
	c.jobs = jobs
	return worker.NewSimpleWorker(c.run)
}

func (c *upgradeWorkerContext) IsUpgradeRunning() bool {
	select {
	case <-c.UpgradeComplete:
		return false
	default:
		return true
	}
}

type apiLostDuringUpgrade struct {
	err error
}

func (e *apiLostDuringUpgrade) Error() string {
	return fmt.Sprintf("API connection lost during upgrade: %v", e.err)
}

func isAPILostDuringUpgrade(err error) bool {
	_, ok := err.(*apiLostDuringUpgrade)
	return ok
}

func (c *upgradeWorkerContext) run(stop <-chan struct{}) error {
	select {
	case <-c.UpgradeComplete:
		// Our work is already done (we're probably being restarted
		// because the API connection has gone down), so do nothing.
		return nil
	default:
	}

	c.agentConfig = c.agent.CurrentConfig()

	// If the machine agent is a state server, flag that state
	// needs to be opened before running upgrade steps
	for _, job := range c.jobs {
		if job == params.JobManageEnviron {
			c.isStateServer = true
		}
	}
	// We need a *state.State for upgrades. We open it independently
	// of StateWorker, because we have no guarantees about when
	// and how often StateWorker might run.
	if c.isStateServer {
		var err error
		c.st, err = openStateForUpgrade(c.agent, c.agentConfig)
		if err != nil {
			return err
		}
		defer c.st.Close()
	}
	if err := c.runUpgrades(); err != nil {
		// Only return an error from the worker if the connection to
		// state went away (possible mongo master change). Returning
		// an error when the connection is lost will cause the agent
		// to restart.
		//
		// For other errors, the error is not returned because we want
		// the machine agent to stay running in an error state waiting
		// for user intervention.
		if isAPILostDuringUpgrade(err) {
			return err
		}
	} else {
		// Upgrade succeeded - signal that the upgrade is complete.
		close(c.UpgradeComplete)
	}
	return nil
}

// runUpgrades runs the upgrade operations for each job type and
// updates the updatedToVersion on success.
func (c *upgradeWorkerContext) runUpgrades() error {
	from := version.Current
	from.Number = c.agentConfig.UpgradedToVersion()
	if from == version.Current {
		logger.Infof("upgrade to %v already completed.", version.Current)
		return nil
	}

	a := c.agent
	tag := c.agentConfig.Tag().(names.MachineTag)

	isMaster, err := isMachineMaster(c.st, tag)
	if err != nil {
		return errors.Trace(err)
	}

	if c.isStateServer {
		// State servers need to wait for other state servers to be
		// ready to run the upgrade.
		if err := waitForOtherStateServers(c.st, isMaster); err != nil {
			logger.Errorf(`other state servers failed to come up for upgrade `+
				`to %s - aborting: %v`, version.Current, err)
			a.setMachineStatus(c.apiState, params.StatusError,
				fmt.Sprintf("upgrade to %v aborted while waiting for other "+
					"state servers: %v", version.Current, err))
			// If master, trigger a rollback to the previous agent version.
			if isMaster {
				logger.Errorf("downgrading environment agent version to %v due to aborted upgrade",
					from.Number)
				if rollbackErr := c.st.SetEnvironAgentVersion(from.Number); rollbackErr != nil {
					return errors.Annotate(rollbackErr, "failed to roll back desired agent version")
				}
			}
			return err
		}
	}

	err = a.ChangeConfig(func(agentConfig agent.ConfigSetter) error {
		var upgradeErr error
		a.setMachineStatus(c.apiState, params.StatusStarted,
			fmt.Sprintf("upgrading to %v", version.Current))

		context := upgrades.NewContext(agentConfig, c.apiState, c.st)
		for _, job := range c.jobs {
			target := upgradeTarget(job, isMaster)
			if target == "" {
				continue
			}
			logger.Infof("starting upgrade from %v to %v for %v %q",
				from, version.Current, target, tag)

			attempts := getUpgradeRetryStrategy()
			for attempt := attempts.Start(); attempt.Next(); {
				upgradeErr = upgradesPerformUpgrade(from.Number, target, context)
				if upgradeErr == nil {
					break
				}
				if connectionIsDead(c.apiState) {
					// API connection has gone away - abort!
					return &apiLostDuringUpgrade{upgradeErr}
				}
				retryText := "will retry"
				if !attempt.HasNext() {
					retryText = "giving up"
				}
				logger.Errorf("upgrade from %v to %v for %v %q failed (%s): %v",
					from, version.Current, target, tag, retryText, upgradeErr)
				a.setMachineStatus(c.apiState, params.StatusError,
					fmt.Sprintf("upgrade to %v failed (%s): %v",
						version.Current, retryText, upgradeErr))
			}
		}
		if upgradeErr != nil {
			return upgradeErr
		}
		agentConfig.SetUpgradedToVersion(version.Current.Number)
		return nil
	})
	if err != nil {
		logger.Errorf("upgrade to %v failed: %v", version.Current, err)
		return err
	}

	logger.Infof("upgrade to %v completed successfully.", version.Current)
	a.setMachineStatus(c.apiState, params.StatusStarted, "")
	return nil
}

var openStateForUpgrade = func(
	agent upgradingMachineAgent,
	agentConfig agent.Config,
) (*state.State, error) {
	if err := agent.ensureMongoServer(agentConfig); err != nil {
		return nil, err
	}
	var err error
	info, ok := agentConfig.MongoInfo()
	if !ok {
		return nil, fmt.Errorf("no state info available")
	}
	st, err := state.Open(info, mongo.DefaultDialOpts(), environs.NewStatePolicy())
	if err != nil {
		return nil, err
	}
	return st, nil
}

var isMachineMaster = func(st *state.State, tag names.MachineTag) (bool, error) {
	if st == nil {
		// If there is no state, we aren't a master.
		return false, nil
	}
	// Not calling the agent openState method as it does other checks
	// we really don't care about here.  All we need here is the machine
	// so we can determine if we are the master or not.
	machine, err := st.Machine(tag.Id())
	if err != nil {
		// This shouldn't happen, and if it does, the state worker will have
		// found out before us, and already errored, or is likely to error out
		// very shortly.  All we do here is return the error. The state worker
		// returns an error that will cause the agent to be terminated.
		return false, errors.Trace(err)
	}
	isMaster, err := mongo.IsMaster(st.MongoSession(), machine)
	if err != nil {
		return false, errors.Trace(err)
	}
	return isMaster, nil
}

var waitForOtherStateServers = func(st *state.State, isMaster bool) error {
	// TODO(mjs) - for now, assume that the other state servers are
	// ready. This function will be fleshed out once the UpgradeInfo
	// work is done.
	return nil
}

var getUpgradeRetryStrategy = func() utils.AttemptStrategy {
	return utils.AttemptStrategy{
		Delay: 2 * time.Minute,
		Min:   5,
	}
}

func upgradeTarget(job params.MachineJob, isMaster bool) upgrades.Target {
	switch job {
	case params.JobManageEnviron:
		if isMaster {
			return upgrades.DatabaseMaster
		}
		return upgrades.StateServer
	case params.JobHostUnits:
		return upgrades.HostMachine
	}
	return ""
}
<commit_msg>cmd/jujud: added wrench for state server upgrade wait<commit_after>package main

import (
	"fmt"
	"time"

	"github.com/juju/errors"
	"github.com/juju/names"
	"github.com/juju/utils"

	"github.com/juju/juju/agent"
	"github.com/juju/juju/environs"
	"github.com/juju/juju/mongo"
	"github.com/juju/juju/state"
	"github.com/juju/juju/state/api"
	"github.com/juju/juju/state/api/params"
	"github.com/juju/juju/upgrades"
	"github.com/juju/juju/version"
	"github.com/juju/juju/worker"
	"github.com/juju/juju/wrench"
)

type upgradingMachineAgent interface {
	ensureMongoServer(agent.Config) error
	setMachineStatus(*api.State, params.Status, string) error
	CurrentConfig() agent.Config
	ChangeConfig(AgentConfigMutator) error
}

var upgradesPerformUpgrade = upgrades.PerformUpgrade // Allow patching for tests

func NewUpgradeWorkerContext() *upgradeWorkerContext {
	return &upgradeWorkerContext{
		UpgradeComplete: make(chan struct{}),
	}
}

type upgradeWorkerContext struct {
	UpgradeComplete chan struct{}
	agent           upgradingMachineAgent
	apiState        *api.State
	jobs            []params.MachineJob
	agentConfig     agent.Config
	isStateServer   bool
	st              *state.State
}

// InitialiseUsingAgent sets up a upgradeWorkerContext from a machine agent instance.
// It may update the agent's configuration.
func (c *upgradeWorkerContext) InitializeUsingAgent(a upgradingMachineAgent) error {
	return a.ChangeConfig(func(agentConfig agent.ConfigSetter) error {
		if !upgrades.AreUpgradesDefined(agentConfig.UpgradedToVersion()) {
			logger.Infof("no upgrade steps required or upgrade steps for %v "+
				"have already been run.", version.Current.Number)
			close(c.UpgradeComplete)

			// Even if no upgrade is required the version number in
			// the agent's config still needs to be bumped.
			agentConfig.SetUpgradedToVersion(version.Current.Number)
		}
		return nil
	})
}

func (c *upgradeWorkerContext) Worker(
	agent upgradingMachineAgent,
	apiState *api.State,
	jobs []params.MachineJob,
) worker.Worker {
	c.agent = agent
	c.apiState = apiState
	c.jobs = jobs
	return worker.NewSimpleWorker(c.run)
}

func (c *upgradeWorkerContext) IsUpgradeRunning() bool {
	select {
	case <-c.UpgradeComplete:
		return false
	default:
		return true
	}
}

type apiLostDuringUpgrade struct {
	err error
}

func (e *apiLostDuringUpgrade) Error() string {
	return fmt.Sprintf("API connection lost during upgrade: %v", e.err)
}

func isAPILostDuringUpgrade(err error) bool {
	_, ok := err.(*apiLostDuringUpgrade)
	return ok
}

func (c *upgradeWorkerContext) run(stop <-chan struct{}) error {
	select {
	case <-c.UpgradeComplete:
		// Our work is already done (we're probably being restarted
		// because the API connection has gone down), so do nothing.
		return nil
	default:
	}

	c.agentConfig = c.agent.CurrentConfig()

	// If the machine agent is a state server, flag that state
	// needs to be opened before running upgrade steps
	for _, job := range c.jobs {
		if job == params.JobManageEnviron {
			c.isStateServer = true
		}
	}
	// We need a *state.State for upgrades. We open it independently
	// of StateWorker, because we have no guarantees about when
	// and how often StateWorker might run.
	if c.isStateServer {
		var err error
		c.st, err = openStateForUpgrade(c.agent, c.agentConfig)
		if err != nil {
			return err
		}
		defer c.st.Close()
	}
	if err := c.runUpgrades(); err != nil {
		// Only return an error from the worker if the connection to
		// state went away (possible mongo master change). Returning
		// an error when the connection is lost will cause the agent
		// to restart.
		//
		// For other errors, the error is not returned because we want
		// the machine agent to stay running in an error state waiting
		// for user intervention.
		if isAPILostDuringUpgrade(err) {
			return err
		}
	} else {
		// Upgrade succeeded - signal that the upgrade is complete.
		close(c.UpgradeComplete)
	}
	return nil
}

// runUpgrades runs the upgrade operations for each job type and
// updates the updatedToVersion on success.
func (c *upgradeWorkerContext) runUpgrades() error {
	from := version.Current
	from.Number = c.agentConfig.UpgradedToVersion()
	if from == version.Current {
		logger.Infof("upgrade to %v already completed.", version.Current)
		return nil
	}

	a := c.agent
	tag := c.agentConfig.Tag().(names.MachineTag)

	isMaster, err := isMachineMaster(c.st, tag)
	if err != nil {
		return errors.Trace(err)
	}

	if c.isStateServer {
		// State servers need to wait for other state servers to be
		// ready to run the upgrade.
		if err := waitForOtherStateServers(c.st, isMaster); err != nil {
			logger.Errorf(`other state servers failed to come up for upgrade `+
				`to %s - aborting: %v`, version.Current, err)
			a.setMachineStatus(c.apiState, params.StatusError,
				fmt.Sprintf("upgrade to %v aborted while waiting for other "+
					"state servers: %v", version.Current, err))
			// If master, trigger a rollback to the previous agent version.
			if isMaster {
				logger.Errorf("downgrading environment agent version to %v due to aborted upgrade",
					from.Number)
				if rollbackErr := c.st.SetEnvironAgentVersion(from.Number); rollbackErr != nil {
					return errors.Annotate(rollbackErr, "failed to roll back desired agent version")
				}
			}
			return err
		}
	}

	err = a.ChangeConfig(func(agentConfig agent.ConfigSetter) error {
		var upgradeErr error
		a.setMachineStatus(c.apiState, params.StatusStarted,
			fmt.Sprintf("upgrading to %v", version.Current))

		context := upgrades.NewContext(agentConfig, c.apiState, c.st)
		for _, job := range c.jobs {
			target := upgradeTarget(job, isMaster)
			if target == "" {
				continue
			}
			logger.Infof("starting upgrade from %v to %v for %v %q",
				from, version.Current, target, tag)

			attempts := getUpgradeRetryStrategy()
			for attempt := attempts.Start(); attempt.Next(); {
				upgradeErr = upgradesPerformUpgrade(from.Number, target, context)
				if upgradeErr == nil {
					break
				}
				if connectionIsDead(c.apiState) {
					// API connection has gone away - abort!
					return &apiLostDuringUpgrade{upgradeErr}
				}
				retryText := "will retry"
				if !attempt.HasNext() {
					retryText = "giving up"
				}
				logger.Errorf("upgrade from %v to %v for %v %q failed (%s): %v",
					from, version.Current, target, tag, retryText, upgradeErr)
				a.setMachineStatus(c.apiState, params.StatusError,
					fmt.Sprintf("upgrade to %v failed (%s): %v",
						version.Current, retryText, upgradeErr))
			}
		}
		if upgradeErr != nil {
			return upgradeErr
		}
		agentConfig.SetUpgradedToVersion(version.Current.Number)
		return nil
	})
	if err != nil {
		logger.Errorf("upgrade to %v failed: %v", version.Current, err)
		return err
	}

	logger.Infof("upgrade to %v completed successfully.", version.Current)
	a.setMachineStatus(c.apiState, params.StatusStarted, "")
	return nil
}

var openStateForUpgrade = func(
	agent upgradingMachineAgent,
	agentConfig agent.Config,
) (*state.State, error) {
	if err := agent.ensureMongoServer(agentConfig); err != nil {
		return nil, err
	}
	var err error
	info, ok := agentConfig.MongoInfo()
	if !ok {
		return nil, fmt.Errorf("no state info available")
	}
	st, err := state.Open(info, mongo.DefaultDialOpts(), environs.NewStatePolicy())
	if err != nil {
		return nil, err
	}
	return st, nil
}

var isMachineMaster = func(st *state.State, tag names.MachineTag) (bool, error) {
	if st == nil {
		// If there is no state, we aren't a master.
		return false, nil
	}
	// Not calling the agent openState method as it does other checks
	// we really don't care about here.  All we need here is the machine
	// so we can determine if we are the master or not.
	machine, err := st.Machine(tag.Id())
	if err != nil {
		// This shouldn't happen, and if it does, the state worker will have
		// found out before us, and already errored, or is likely to error out
		// very shortly.  All we do here is return the error. The state worker
		// returns an error that will cause the agent to be terminated.
		return false, errors.Trace(err)
	}
	isMaster, err := mongo.IsMaster(st.MongoSession(), machine)
	if err != nil {
		return false, errors.Trace(err)
	}
	return isMaster, nil
}

var waitForOtherStateServers = func(st *state.State, isMaster bool) error {
	if wrench.IsActive("machine-agent", "fail-state-server-upgrade-wait") {
		return errors.New("failing other state servers check due to wrench")
	}
	// TODO(mjs) - for now, assume that the other state servers are
	// ready. This function will be fleshed out once the UpgradeInfo
	// work is done.
	return nil
}

var getUpgradeRetryStrategy = func() utils.AttemptStrategy {
	return utils.AttemptStrategy{
		Delay: 2 * time.Minute,
		Min:   5,
	}
}

func upgradeTarget(job params.MachineJob, isMaster bool) upgrades.Target {
	switch job {
	case params.JobManageEnviron:
		if isMaster {
			return upgrades.DatabaseMaster
		}
		return upgrades.StateServer
	case params.JobHostUnits:
		return upgrades.HostMachine
	}
	return ""
}
<|endoftext|>


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

import (
	"fmt"
	"strings"

	"github.com/coreos/go-iptables/iptables"
)

type rules struct {
	host          string
	kiamPort      int
	hostInterface string
}

const (
	metadataAddress = "169.254.169.254"
)

func newIPTablesRules(host string, kiamPort int, hostInterface string) *rules {
	return &rules{host: host, kiamPort: kiamPort, hostInterface: hostInterface}
}

func (r *rules) Add() error {
	ipt, err := iptables.New()
	if err != nil {
		return err
	}

	return ipt.AppendUnique("nat", "PREROUTING", r.ruleSpec()...)
}

func (r *rules) ruleSpec() []string {
	rules := []string{
		"-p", "tcp",
		"-d", metadataAddress,
		"--dport", "80",
		"-j", "DNAT",
		"--to-destination", r.kiamAddress(),
	}
	if strings.HasPrefix(r.hostInterface, "!") {
		rules = append(rules, "!")
	}
	rules = append(rules, "-i", strings.TrimPrefix(r.hostInterface, "!"))

	return rules
}

func (r *rules) Remove() error {
	ipt, err := iptables.New()
	if err != nil {
		return err
	}

	return ipt.Delete("nat", "PREROUTING", r.ruleSpec()...)
}

func (r *rules) kiamAddress() string {
	return fmt.Sprintf("%s:%d", r.host, r.kiamPort)
}
<commit_msg>add retries around iptable rule removal (#402)<commit_after>// Copyright 2017 uSwitch
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package main

import (
	"fmt"
	"strings"
	"time"

	"github.com/coreos/go-iptables/iptables"
	log "github.com/sirupsen/logrus"
)

type rules struct {
	host          string
	kiamPort      int
	hostInterface string
}

const (
	metadataAddress = "169.254.169.254"
)

func newIPTablesRules(host string, kiamPort int, hostInterface string) *rules {
	return &rules{host: host, kiamPort: kiamPort, hostInterface: hostInterface}
}

func (r *rules) Add() error {
	ipt, err := iptables.New()
	if err != nil {
		return err
	}

	return ipt.AppendUnique("nat", "PREROUTING", r.ruleSpec()...)
}

func (r *rules) ruleSpec() []string {
	rules := []string{
		"-p", "tcp",
		"-d", metadataAddress,
		"--dport", "80",
		"-j", "DNAT",
		"--to-destination", r.kiamAddress(),
	}
	if strings.HasPrefix(r.hostInterface, "!") {
		rules = append(rules, "!")
	}
	rules = append(rules, "-i", strings.TrimPrefix(r.hostInterface, "!"))

	return rules
}

var (
	retryInterval = time.Millisecond * 500
	maxAttempts   = 30
)

func (r *rules) Remove() error {
	ipt, err := iptables.New()
	if err != nil {
		return err
	}

	var attempt int
	for {
		if attempt >= maxAttempts {
			log.Errorf("failed to remove iptables rule, retries exhausted: %s", err.Error())
			break
		}
		if err := ipt.Delete("nat", "PREROUTING", r.ruleSpec()...); err == nil {
			log.Info("iptables rule was successfully removed")
			break
		}
		log.Warnf("failed to remove iptables rule, will retry: %s", err.Error())
		time.Sleep(retryInterval)
		attempt++
	}
	return nil
}

func (r *rules) kiamAddress() string {
	return fmt.Sprintf("%s:%d", r.host, r.kiamPort)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/gob"
	"fmt"
	"io"
	"os"
	"strconv"
	"strings"

	"github.com/Symantec/Dominator/lib/errors"
	"github.com/Symantec/Dominator/lib/log"
	"github.com/Symantec/Dominator/lib/srpc"
	proto "github.com/Symantec/Dominator/proto/logger"
)

func watchSubcommand(clients []*srpc.Client, addrs, args []string,
	logger log.Logger) {
	level, err := strconv.ParseInt(args[0], 10, 16)
	if err != nil {
		logger.Fatalf("Error parsing level: %s\n", err)
	}
	if err != nil {
		logger.Fatalf("Error dialing: %s\n", err)
	}
	if err := watchAll(clients, addrs, int16(level)); err != nil {
		logger.Fatalf("Error watching: %s\n", err)
	}
	os.Exit(0)
}

func watchAll(clients []*srpc.Client, addrs []string, level int16) error {
	if len(clients) == 1 {
		return watchOne(clients[0], level, "")
	}
	maxWidth := 0
	for _, addr := range addrs {
		if len(addr) > maxWidth {
			maxWidth = len(addr)
		}
	}
	errors := make(chan error, 1)
	for index, client := range clients {
		prefix := addrs[index]
		if len(prefix) < maxWidth {
			prefix += strings.Repeat(" ", maxWidth-len(prefix))
		}
		go func(client *srpc.Client, level int16, prefix string) {
			errors <- watchOne(client, level, prefix)
		}(client, level, prefix)
	}
	for range clients {
		if err := <-errors; err != nil {
			return err
		}
	}
	return nil
}

func watchOne(client *srpc.Client, level int16, prefix string) error {
	request := proto.WatchRequest{
		ExcludeRegex: *excludeRegex,
		IncludeRegex: *includeRegex,
		Name:         *loggerName,
		DebugLevel:   level,
	}
	if conn, err := client.Call("Logger.Watch"); err != nil {
		return err
	} else {
		defer conn.Close()
		encoder := gob.NewEncoder(conn)
		if err := encoder.Encode(request); err != nil {
			return err
		}
		if err := conn.Flush(); err != nil {
			return err
		}
		decoder := gob.NewDecoder(conn)
		var response proto.WatchResponse
		if err := decoder.Decode(&response); err != nil {
			return fmt.Errorf("error decoding: %s", err)
		}
		if response.Error != "" {
			return errors.New(response.Error)
		}
		if prefix == "" {
			_, err := io.Copy(os.Stdout, conn)
			return err
		}
		for {
			line, err := conn.ReadString('\n')
			if len(line) > 0 {
				if prefix != "" {
					line = prefix + " " + line
				}
				if _, err := os.Stdout.Write([]byte(line)); err != nil {
					return err
				}
			}
			if err != nil {
				if err == io.EOF {
					return nil
				}
				return err
			}
		}
	}
}
<commit_msg>Switch logtool command to use connection Decode and Encode methods.<commit_after>package main

import (
	"fmt"
	"io"
	"os"
	"strconv"
	"strings"

	"github.com/Symantec/Dominator/lib/errors"
	"github.com/Symantec/Dominator/lib/log"
	"github.com/Symantec/Dominator/lib/srpc"
	proto "github.com/Symantec/Dominator/proto/logger"
)

func watchSubcommand(clients []*srpc.Client, addrs, args []string,
	logger log.Logger) {
	level, err := strconv.ParseInt(args[0], 10, 16)
	if err != nil {
		logger.Fatalf("Error parsing level: %s\n", err)
	}
	if err != nil {
		logger.Fatalf("Error dialing: %s\n", err)
	}
	if err := watchAll(clients, addrs, int16(level)); err != nil {
		logger.Fatalf("Error watching: %s\n", err)
	}
	os.Exit(0)
}

func watchAll(clients []*srpc.Client, addrs []string, level int16) error {
	if len(clients) == 1 {
		return watchOne(clients[0], level, "")
	}
	maxWidth := 0
	for _, addr := range addrs {
		if len(addr) > maxWidth {
			maxWidth = len(addr)
		}
	}
	errors := make(chan error, 1)
	for index, client := range clients {
		prefix := addrs[index]
		if len(prefix) < maxWidth {
			prefix += strings.Repeat(" ", maxWidth-len(prefix))
		}
		go func(client *srpc.Client, level int16, prefix string) {
			errors <- watchOne(client, level, prefix)
		}(client, level, prefix)
	}
	for range clients {
		if err := <-errors; err != nil {
			return err
		}
	}
	return nil
}

func watchOne(client *srpc.Client, level int16, prefix string) error {
	request := proto.WatchRequest{
		ExcludeRegex: *excludeRegex,
		IncludeRegex: *includeRegex,
		Name:         *loggerName,
		DebugLevel:   level,
	}
	if conn, err := client.Call("Logger.Watch"); err != nil {
		return err
	} else {
		defer conn.Close()
		if err := conn.Encode(request); err != nil {
			return err
		}
		if err := conn.Flush(); err != nil {
			return err
		}
		var response proto.WatchResponse
		if err := conn.Decode(&response); err != nil {
			return fmt.Errorf("error decoding: %s", err)
		}
		if response.Error != "" {
			return errors.New(response.Error)
		}
		if prefix == "" {
			_, err := io.Copy(os.Stdout, conn)
			return err
		}
		for {
			line, err := conn.ReadString('\n')
			if len(line) > 0 {
				if prefix != "" {
					line = prefix + " " + line
				}
				if _, err := os.Stdout.Write([]byte(line)); err != nil {
					return err
				}
			}
			if err != nil {
				if err == io.EOF {
					return nil
				}
				return err
			}
		}
	}
}
<|endoftext|>


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

package main

import (
	"context"
	"flag"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"cloud.google.com/go/compute/metadata"
	"github.com/oklog/run"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"go.uber.org/zap/zapcore"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/log/zap"
	ctrlmetrics "sigs.k8s.io/controller-runtime/pkg/metrics"

	// Blank import required to register GCP auth handlers to talk to GKE clusters.
	_ "k8s.io/client-go/plugin/pkg/client/auth/gcp"

	"github.com/GoogleCloudPlatform/prometheus-engine/pkg/operator"
)

func unstableFlagHelp(help string) string {
	return help + " (Setting this flag voids any guarantees of proper behavior of the operator.)"
}

func main() {
	var (
		defaultProjectID string
		defaultCluster   string
		defaultLocation  string
	)
	if metadata.OnGCE() {
		defaultProjectID, _ = metadata.ProjectID()
		defaultCluster, _ = metadata.InstanceAttributeValue("cluster-name")
		defaultLocation, _ = metadata.InstanceAttributeValue("cluster-location")
	}
	var (
		logVerbosity      = flag.Int("v", 0, "Logging verbosity")
		projectID         = flag.String("project-id", defaultProjectID, "Project ID of the cluster. May be left empty on GKE.")
		location          = flag.String("location", defaultLocation, "GCP location of the cluster. Maybe be left empty on GKE.")
		cluster           = flag.String("cluster", defaultCluster, "Name of the cluster the operator acts on. May be left empty on GKE.")
		operatorNamespace = flag.String("operator-namespace", operator.DefaultOperatorNamespace,
			"Namespace in which the operator manages its resources.")
		publicNamespace = flag.String("public-namespace", operator.DefaultPublicNamespace,
			"Namespace in which the operator reads user-provided resources.")

		imageCollector = flag.String("image-collector", operator.ImageCollector,
			unstableFlagHelp("Override for the container image of the collector."))
		imageConfigReloader = flag.String("image-config-reloader", operator.ImageConfigReloader,
			unstableFlagHelp("Override for the container image of the config reloader."))
		imageRuleEvaluator = flag.String("image-rule-evaluator", operator.ImageRuleEvaluator,
			unstableFlagHelp("Override for the container image of the rule evaluator."))

		hostNetwork = flag.Bool("host-network", true,
			"Whether pods are deployed with hostNetwork enabled. If true, GKE clusters with Workload Identity will not require additional permission for the components deployed by the operator. Must be false on GKE Autopilot clusters.")
		priorityClass = flag.String("priority-class", "",
			"Priority class at which the collector pods are run.")
		gcmEndpoint = flag.String("cloud-monitoring-endpoint", "",
			"Override for the Cloud Monitoring endpoint to use for all collectors.")
		tlsCert     = flag.String("tls-cert-base64", "", "The base64-encoded TLS certificate.")
		tlsKey      = flag.String("tls-key-base64", "", "The base64-encoded TLS key.")
		caCert      = flag.String("ca-cert-base64", "", "The base64-encoded certificate authority.")
		webhookAddr = flag.String("webhook-addr", ":8443",
			"Address to listen to for incoming kube admission webhook connections.")
		metricsAddr = flag.String("metrics-addr", ":18080", "Address to emit metrics on.")

		collectorMemoryResource = flag.Int64("collector-memory-resource", 200, "The Memory Resource of collector pod, in mega bytes")
		collectorMemoryLimit    = flag.Int64("collector-memory-limit", 3000, "The Memory Limit of collector pod, in mega bytes.")
		collectorCPUResource    = flag.Int64("collector-cpu-resource", 100, "The CPU Resource of collector pod, in milli cpu.")
		evaluatorMemoryResource = flag.Int64("evaluator-memory-resource", 200, "The Memory Resource of evaluator pod, in mega bytes.")
		evaluatorMemoryLimit    = flag.Int64("evaluator-memory-limit", 1000, "The Memory Limit of evaluator pod, in mega bytesv.")
		evaluatorCPUResource    = flag.Int64("evaluator-cpu-resource", 100, "The CPU Resource of evaluator pod, in milli cpu.")
		mode                    = flag.String("mode", "kubectl", "how managed collection was provisioned.")
	)
	flag.Parse()

	logger := zap.New(zap.Level(zapcore.Level(-*logVerbosity)))
	ctrl.SetLogger(logger)

	cfg, err := ctrl.GetConfig()
	if err != nil {
		logger.Error(err, "loading kubeconfig failed")
		os.Exit(1)
	}
	switch *mode {
	// repo manifest always defaults to "kubectl".
	case "kubectl":
	case "gcloud":
	case "gcloud-auto":
	default:
		logger.Error(err, "--mode must be one of {'kubectl', 'gcloud', 'gcloud-auto'}")
		os.Exit(1)
	}

	// controller-runtime creates a registry against which its metrics are registered globally.
	// Using it as our non-global registry is the easiest way to combine metrics into a single
	// /metrics endpoint.
	// It already has the GoCollector and ProcessCollector metrics installed.
	metrics := ctrlmetrics.Registry

	op, err := operator.New(logger, cfg, metrics, operator.Options{
		ProjectID:               *projectID,
		Location:                *location,
		Cluster:                 *cluster,
		OperatorNamespace:       *operatorNamespace,
		PublicNamespace:         *publicNamespace,
		ImageCollector:          *imageCollector,
		ImageConfigReloader:     *imageConfigReloader,
		ImageRuleEvaluator:      *imageRuleEvaluator,
		HostNetwork:             *hostNetwork,
		PriorityClass:           *priorityClass,
		CloudMonitoringEndpoint: *gcmEndpoint,
		TLSCert:                 *tlsCert,
		TLSKey:                  *tlsKey,
		CACert:                  *caCert,
		ListenAddr:              *webhookAddr,
		CollectorMemoryResource: *collectorMemoryResource,
		CollectorMemoryLimit:    *collectorMemoryLimit,
		CollectorCPUResource:    *collectorCPUResource,
		EvaluatorCPUResource:    *evaluatorCPUResource,
		EvaluatorMemoryResource: *evaluatorMemoryResource,
		EvaluatorMemoryLimit:    *evaluatorMemoryLimit,
		Mode:                    *mode,
	})
	if err != nil {
		logger.Error(err, "instantiating operator failed")
		os.Exit(1)
	}

	var g run.Group
	// Termination handler.
	{
		term := make(chan os.Signal, 1)
		cancel := make(chan struct{})
		signal.Notify(term, os.Interrupt, syscall.SIGTERM)

		g.Add(
			func() error {
				select {
				case <-term:
					logger.Info("received SIGTERM, exiting gracefully...")
				case <-cancel:
				}
				return nil
			},
			func(err error) {
				close(cancel)
			},
		)
	}
	// Operator monitoring.
	{
		server := &http.Server{Addr: *metricsAddr}
		http.Handle("/metrics", promhttp.HandlerFor(metrics, promhttp.HandlerOpts{Registry: metrics}))
		g.Add(func() error {
			return server.ListenAndServe()
		}, func(err error) {
			ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
			server.Shutdown(ctx)
			cancel()
		})
	}
	// Main operator loop.
	{
		ctx, cancel := context.WithCancel(context.Background())
		g.Add(func() error {
			return op.Run(ctx)
		}, func(err error) {
			cancel()
		})
	}
	if err := g.Run(); err != nil {
		logger.Error(err, "exit with error")
		os.Exit(1)
	}
}
<commit_msg>change gcloud to gke<commit_after>// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"context"
	"flag"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"cloud.google.com/go/compute/metadata"
	"github.com/oklog/run"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"go.uber.org/zap/zapcore"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/log/zap"
	ctrlmetrics "sigs.k8s.io/controller-runtime/pkg/metrics"

	// Blank import required to register GCP auth handlers to talk to GKE clusters.
	_ "k8s.io/client-go/plugin/pkg/client/auth/gcp"

	"github.com/GoogleCloudPlatform/prometheus-engine/pkg/operator"
)

func unstableFlagHelp(help string) string {
	return help + " (Setting this flag voids any guarantees of proper behavior of the operator.)"
}

func main() {
	var (
		defaultProjectID string
		defaultCluster   string
		defaultLocation  string
	)
	if metadata.OnGCE() {
		defaultProjectID, _ = metadata.ProjectID()
		defaultCluster, _ = metadata.InstanceAttributeValue("cluster-name")
		defaultLocation, _ = metadata.InstanceAttributeValue("cluster-location")
	}
	var (
		logVerbosity      = flag.Int("v", 0, "Logging verbosity")
		projectID         = flag.String("project-id", defaultProjectID, "Project ID of the cluster. May be left empty on GKE.")
		location          = flag.String("location", defaultLocation, "GCP location of the cluster. Maybe be left empty on GKE.")
		cluster           = flag.String("cluster", defaultCluster, "Name of the cluster the operator acts on. May be left empty on GKE.")
		operatorNamespace = flag.String("operator-namespace", operator.DefaultOperatorNamespace,
			"Namespace in which the operator manages its resources.")
		publicNamespace = flag.String("public-namespace", operator.DefaultPublicNamespace,
			"Namespace in which the operator reads user-provided resources.")

		imageCollector = flag.String("image-collector", operator.ImageCollector,
			unstableFlagHelp("Override for the container image of the collector."))
		imageConfigReloader = flag.String("image-config-reloader", operator.ImageConfigReloader,
			unstableFlagHelp("Override for the container image of the config reloader."))
		imageRuleEvaluator = flag.String("image-rule-evaluator", operator.ImageRuleEvaluator,
			unstableFlagHelp("Override for the container image of the rule evaluator."))

		hostNetwork = flag.Bool("host-network", true,
			"Whether pods are deployed with hostNetwork enabled. If true, GKE clusters with Workload Identity will not require additional permission for the components deployed by the operator. Must be false on GKE Autopilot clusters.")
		priorityClass = flag.String("priority-class", "",
			"Priority class at which the collector pods are run.")
		gcmEndpoint = flag.String("cloud-monitoring-endpoint", "",
			"Override for the Cloud Monitoring endpoint to use for all collectors.")
		tlsCert     = flag.String("tls-cert-base64", "", "The base64-encoded TLS certificate.")
		tlsKey      = flag.String("tls-key-base64", "", "The base64-encoded TLS key.")
		caCert      = flag.String("ca-cert-base64", "", "The base64-encoded certificate authority.")
		webhookAddr = flag.String("webhook-addr", ":8443",
			"Address to listen to for incoming kube admission webhook connections.")
		metricsAddr = flag.String("metrics-addr", ":18080", "Address to emit metrics on.")

		collectorMemoryResource = flag.Int64("collector-memory-resource", 200, "The Memory Resource of collector pod, in mega bytes")
		collectorMemoryLimit    = flag.Int64("collector-memory-limit", 3000, "The Memory Limit of collector pod, in mega bytes.")
		collectorCPUResource    = flag.Int64("collector-cpu-resource", 100, "The CPU Resource of collector pod, in milli cpu.")
		evaluatorMemoryResource = flag.Int64("evaluator-memory-resource", 200, "The Memory Resource of evaluator pod, in mega bytes.")
		evaluatorMemoryLimit    = flag.Int64("evaluator-memory-limit", 1000, "The Memory Limit of evaluator pod, in mega bytesv.")
		evaluatorCPUResource    = flag.Int64("evaluator-cpu-resource", 100, "The CPU Resource of evaluator pod, in milli cpu.")
		mode                    = flag.String("mode", "kubectl", "how managed collection was provisioned.")
	)
	flag.Parse()

	logger := zap.New(zap.Level(zapcore.Level(-*logVerbosity)))
	ctrl.SetLogger(logger)

	cfg, err := ctrl.GetConfig()
	if err != nil {
		logger.Error(err, "loading kubeconfig failed")
		os.Exit(1)
	}
	switch *mode {
	// repo manifest always defaults to "kubectl".
	case "kubectl":
	case "gke":
	case "gke-auto":
	default:
		logger.Error(err, "--mode must be one of {'kubectl', 'gcloud', 'gcloud-auto'}")
		os.Exit(1)
	}

	// controller-runtime creates a registry against which its metrics are registered globally.
	// Using it as our non-global registry is the easiest way to combine metrics into a single
	// /metrics endpoint.
	// It already has the GoCollector and ProcessCollector metrics installed.
	metrics := ctrlmetrics.Registry

	op, err := operator.New(logger, cfg, metrics, operator.Options{
		ProjectID:               *projectID,
		Location:                *location,
		Cluster:                 *cluster,
		OperatorNamespace:       *operatorNamespace,
		PublicNamespace:         *publicNamespace,
		ImageCollector:          *imageCollector,
		ImageConfigReloader:     *imageConfigReloader,
		ImageRuleEvaluator:      *imageRuleEvaluator,
		HostNetwork:             *hostNetwork,
		PriorityClass:           *priorityClass,
		CloudMonitoringEndpoint: *gcmEndpoint,
		TLSCert:                 *tlsCert,
		TLSKey:                  *tlsKey,
		CACert:                  *caCert,
		ListenAddr:              *webhookAddr,
		CollectorMemoryResource: *collectorMemoryResource,
		CollectorMemoryLimit:    *collectorMemoryLimit,
		CollectorCPUResource:    *collectorCPUResource,
		EvaluatorCPUResource:    *evaluatorCPUResource,
		EvaluatorMemoryResource: *evaluatorMemoryResource,
		EvaluatorMemoryLimit:    *evaluatorMemoryLimit,
		Mode:                    *mode,
	})
	if err != nil {
		logger.Error(err, "instantiating operator failed")
		os.Exit(1)
	}

	var g run.Group
	// Termination handler.
	{
		term := make(chan os.Signal, 1)
		cancel := make(chan struct{})
		signal.Notify(term, os.Interrupt, syscall.SIGTERM)

		g.Add(
			func() error {
				select {
				case <-term:
					logger.Info("received SIGTERM, exiting gracefully...")
				case <-cancel:
				}
				return nil
			},
			func(err error) {
				close(cancel)
			},
		)
	}
	// Operator monitoring.
	{
		server := &http.Server{Addr: *metricsAddr}
		http.Handle("/metrics", promhttp.HandlerFor(metrics, promhttp.HandlerOpts{Registry: metrics}))
		g.Add(func() error {
			return server.ListenAndServe()
		}, func(err error) {
			ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
			server.Shutdown(ctx)
			cancel()
		})
	}
	// Main operator loop.
	{
		ctx, cancel := context.WithCancel(context.Background())
		g.Add(func() error {
			return op.Run(ctx)
		}, func(err error) {
			cancel()
		})
	}
	if err := g.Run(); err != nil {
		logger.Error(err, "exit with error")
		os.Exit(1)
	}
}
<|endoftext|>


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

package main

import (
	"fmt"
	"net/http"
	"regexp"
	"strings"
	"time"

	"github.com/coreos/go-semver/semver"
	"github.com/spf13/cobra"
	"golang.org/x/net/context"
	"google.golang.org/api/compute/v1"
	gs "google.golang.org/api/storage/v1"

	"github.com/coreos/mantle/auth"
	"github.com/coreos/mantle/platform/api/aws"
	"github.com/coreos/mantle/platform/api/azure"
	"github.com/coreos/mantle/platform/api/gcloud"
	"github.com/coreos/mantle/storage"
	"github.com/coreos/mantle/storage/index"
)

var (
	releaseDryRun bool
	cmdRelease    = &cobra.Command{
		Use:   "release [options]",
		Short: "Publish a new CoreOS release.",
		Run:   runRelease,
		Long:  `Publish a new CoreOS release.`,
	}
)

func init() {
	cmdRelease.Flags().StringVar(&awsCredentialsFile, "aws-credentials", "", "AWS credentials file")
	cmdRelease.Flags().StringVar(&azureProfile, "azure-profile", "", "Azure Profile json file")
	cmdRelease.Flags().BoolVarP(&releaseDryRun, "dry-run", "n", false,
		"perform a trial run, do not make changes")
	AddSpecFlags(cmdRelease.Flags())
	root.AddCommand(cmdRelease)
}

func runRelease(cmd *cobra.Command, args []string) {
	if len(args) > 0 {
		plog.Fatal("No args accepted")
	}

	spec := ChannelSpec()
	ctx := context.Background()
	client, err := auth.GoogleClient()
	if err != nil {
		plog.Fatalf("Authentication failed: %v", err)
	}

	src, err := storage.NewBucket(client, spec.SourceURL())
	if err != nil {
		plog.Fatal(err)
	}
	src.WriteDryRun(releaseDryRun)

	if err := src.Fetch(ctx); err != nil {
		plog.Fatal(err)
	}

	// Sanity check!
	if vertxt := src.Object(src.Prefix() + "version.txt"); vertxt == nil {
		verurl := src.URL().String() + "version.txt"
		plog.Fatalf("File not found: %s", verurl)
	}

	// Register GCE image if needed.
	doGCE(ctx, client, src, &spec)

	// Make Azure images public.
	doAzure(ctx, client, src, &spec)

	// Make AWS images public.
	doAWS(ctx, client, src, &spec)

	for _, dSpec := range spec.Destinations {
		dst, err := storage.NewBucket(client, dSpec.BaseURL)
		if err != nil {
			plog.Fatal(err)
		}
		dst.WriteDryRun(releaseDryRun)

		// Fetch parent directories non-recursively to re-index it later.
		for _, prefix := range dSpec.ParentPrefixes() {
			if err := dst.FetchPrefix(ctx, prefix, false); err != nil {
				plog.Fatal(err)
			}
		}

		// Fetch and sync each destination directory.
		for _, prefix := range dSpec.FinalPrefixes() {
			if err := dst.FetchPrefix(ctx, prefix, true); err != nil {
				plog.Fatal(err)
			}

			sync := index.NewSyncIndexJob(src, dst)
			sync.DestinationPrefix(prefix)
			sync.DirectoryHTML(dSpec.DirectoryHTML)
			sync.IndexHTML(dSpec.IndexHTML)
			sync.Delete(true)
			if dSpec.Title != "" {
				sync.Name(dSpec.Title)
			}
			if err := sync.Do(ctx); err != nil {
				plog.Fatal(err)
			}
		}

		// Now refresh the parent directory indexes.
		for _, prefix := range dSpec.ParentPrefixes() {
			parent := index.NewIndexJob(dst)
			parent.Prefix(prefix)
			parent.DirectoryHTML(dSpec.DirectoryHTML)
			parent.IndexHTML(dSpec.IndexHTML)
			parent.Recursive(false)
			parent.Delete(true)
			if dSpec.Title != "" {
				parent.Name(dSpec.Title)
			}
			if err := parent.Do(ctx); err != nil {
				plog.Fatal(err)
			}
		}
	}
}

func sanitizeVersion() string {
	v := strings.Replace(specVersion, ".", "-", -1)
	return strings.Replace(v, "+", "-", -1)
}

func doGCE(ctx context.Context, client *http.Client, src *storage.Bucket, spec *channelSpec) {
	if spec.GCE.Project == "" || spec.GCE.Image == "" {
		plog.Notice("GCE image creation disabled.")
		return
	}

	api, err := gcloud.New(&gcloud.Options{
		Project: spec.GCE.Project,
	})
	if err != nil {
		plog.Fatalf("GCE client failed: %v", err)
	}

	publishImage := func(image string) {
		if spec.GCE.Publish == "" {
			plog.Notice("GCE image name publishing disabled.")
			return
		}
		obj := gs.Object{
			Name:        src.Prefix() + spec.GCE.Publish,
			ContentType: "text/plain",
		}
		media := strings.NewReader(
			fmt.Sprintf("projects/%s/global/images/%s\n",
				spec.GCE.Project, image))
		if err := src.Upload(ctx, &obj, media); err != nil {
			plog.Fatal(err)
		}
	}

	nameVer := fmt.Sprintf("%s-%s-v", spec.GCE.Family, sanitizeVersion())
	date := time.Now().UTC()
	name := nameVer + date.Format("20060102")
	desc := fmt.Sprintf("%s, %s, %s published on %s", spec.GCE.Description,
		specVersion, specBoard, date.Format("2006-01-02"))

	images, err := api.ListImages(ctx, spec.GCE.Family+"-")
	if err != nil {
		plog.Fatal(err)
	}

	var conflicting []string
	for _, image := range images {
		if strings.HasPrefix(image.Name, nameVer) {
			conflicting = append(conflicting, image.Name)
		}
	}

	// Check for any with the same version but possibly different dates.
	if len(conflicting) > 1 {
		plog.Fatalf("Duplicate GCE images found: %v", conflicting)
	} else if len(conflicting) == 1 {
		plog.Noticef("GCE image already exists: %s", conflicting[0])
		publishImage(conflicting[0])
		return
	}

	if spec.GCE.Limit > 0 && len(images) > spec.GCE.Limit {
		plog.Noticef("Pruning %d GCE images.", len(images)-spec.GCE.Limit)
		plog.Notice("NOPE! JUST KIDDING, TODO")
	}

	obj := src.Object(src.Prefix() + spec.GCE.Image)
	if obj == nil {
		plog.Fatalf("GCE image not found %s%s", src.URL(), spec.GCE.Image)
	}

	if releaseDryRun {
		plog.Noticef("Would create GCE image %s", name)
		return
	}

	plog.Noticef("Creating GCE image %s", name)
	parsedVersion, err := semver.NewVersion(specVersion)
	if err != nil {
		plog.Fatalf("couldn't parse version %s: %v", specVersion, err)
	}
	disableMultiqueue := false
	if parsedVersion.LessThan(semver.Version{Major: 1409}) {
		disableMultiqueue = true
		plog.Noticef("Not enabling multiqueue for version %v", specVersion)
	}
	op, pending, err := api.CreateImage(&gcloud.ImageSpec{
		SourceImage:           obj.MediaLink,
		Family:                spec.GCE.Family,
		Name:                  name,
		Description:           desc,
		Licenses:              spec.GCE.Licenses,
		DisableSCSIMultiqueue: disableMultiqueue,
	}, false)
	if err != nil {
		plog.Fatalf("GCE image creation failed: %v", err)
	}

	plog.Infof("Waiting for image creation to finish...")
	pending.Interval = 3 * time.Second
	pending.Progress = func(_ string, _ time.Duration, op *compute.Operation) error {
		status := strings.ToLower(op.Status)
		if op.Progress != 0 {
			plog.Infof("Image creation is %s: %s % 2d%%", status, op.StatusMessage, op.Progress)
		} else {
			plog.Infof("Image creation is %s. %s", status, op.StatusMessage)
		}
		return nil
	}
	if err := pending.Wait(); err != nil {
		plog.Fatal(err)
	}
	plog.Info("Success!")

	publishImage(name)

	var pendings []*gcloud.Pending
	for _, old := range images {
		if old.Deprecated != nil && old.Deprecated.State != "" {
			continue
		}
		plog.Noticef("Deprecating old image %s", old.Name)
		pending, err := api.DeprecateImage(old.Name, gcloud.DeprecationStateDeprecated, op.TargetLink)
		if err != nil {
			plog.Fatal(err)
		}
		pending.Interval = 1 * time.Second
		pending.Timeout = 0
		pendings = append(pendings, pending)
	}

	plog.Infof("Waiting on %d operations.", len(pendings))
	for _, pending := range pendings {
		err := pending.Wait()
		if err != nil {
			plog.Fatal(err)
		}
	}
}

func doAzure(ctx context.Context, client *http.Client, src *storage.Bucket, spec *channelSpec) {
	if spec.Azure.StorageAccount == "" {
		plog.Notice("Azure image creation disabled.")
		return
	}

	prof, err := auth.ReadAzureProfile(azureProfile)
	if err != nil {
		plog.Fatalf("failed reading Azure profile: %v", err)
	}

	// channel name should be caps for azure image
	imageName := fmt.Sprintf("%s-%s-%s", spec.Azure.Offer, strings.Title(specChannel), specVersion)

	for _, environment := range spec.Azure.Environments {
		opt := prof.SubscriptionOptions(environment.SubscriptionName)
		if opt == nil {
			plog.Fatalf("couldn't find subscription %q", environment.SubscriptionName)
		}

		api, err := azure.New(opt)
		if err != nil {
			plog.Fatalf("failed to create Azure API: %v", err)
		}

		if releaseDryRun {
			// TODO(bgilbert): check that the image exists
			plog.Printf("Would share %q on %v", imageName, environment.SubscriptionName)
			continue
		} else {
			plog.Printf("Sharing %q on %v...", imageName, environment.SubscriptionName)
		}

		if err := api.ShareImage(imageName, "public"); err != nil {
			plog.Fatalf("failed to share image %q: %v", imageName, err)
		}
	}
}

func doAWS(ctx context.Context, client *http.Client, src *storage.Bucket, spec *channelSpec) {
	if spec.AWS.Image == "" {
		plog.Notice("AWS image creation disabled.")
		return
	}

	imageName := fmt.Sprintf("%v-%v-%v", spec.AWS.BaseName, specChannel, specVersion)
	imageName = regexp.MustCompile(`[^A-Za-z0-9()\\./_-]`).ReplaceAllLiteralString(imageName, "_")

	for _, part := range spec.AWS.Partitions {
		for _, region := range part.Regions {
			if releaseDryRun {
				plog.Printf("Checking for images in %v %v...", part.Name, region)
			} else {
				plog.Printf("Publishing images in %v %v...", part.Name, region)
			}

			api, err := aws.New(&aws.Options{
				CredentialsFile: awsCredentialsFile,
				Profile:         part.Profile,
				Region:          region,
			})
			if err != nil {
				plog.Fatalf("creating client for %v %v: %v", part.Name, region, err)
			}

			publish := func(imageName string) {
				imageID, err := api.FindImage(imageName)
				if err != nil {
					plog.Fatalf("couldn't find image %q in %v %v: %v", imageName, part.Name, region, err)
				}

				if !releaseDryRun {
					err := api.PublishImage(imageID)
					if err != nil {
						plog.Fatalf("couldn't publish image in %v %v: %v", part.Name, region, err)
					}
				}
			}
			publish(imageName)
			publish(imageName + "-hvm")
		}
	}
}
<commit_msg>cmd/plume: Only pretend to prune GCE images at the end<commit_after>// Copyright 2016 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"fmt"
	"net/http"
	"regexp"
	"strings"
	"time"

	"github.com/coreos/go-semver/semver"
	"github.com/spf13/cobra"
	"golang.org/x/net/context"
	"google.golang.org/api/compute/v1"
	gs "google.golang.org/api/storage/v1"

	"github.com/coreos/mantle/auth"
	"github.com/coreos/mantle/platform/api/aws"
	"github.com/coreos/mantle/platform/api/azure"
	"github.com/coreos/mantle/platform/api/gcloud"
	"github.com/coreos/mantle/storage"
	"github.com/coreos/mantle/storage/index"
)

var (
	releaseDryRun bool
	cmdRelease    = &cobra.Command{
		Use:   "release [options]",
		Short: "Publish a new CoreOS release.",
		Run:   runRelease,
		Long:  `Publish a new CoreOS release.`,
	}
)

func init() {
	cmdRelease.Flags().StringVar(&awsCredentialsFile, "aws-credentials", "", "AWS credentials file")
	cmdRelease.Flags().StringVar(&azureProfile, "azure-profile", "", "Azure Profile json file")
	cmdRelease.Flags().BoolVarP(&releaseDryRun, "dry-run", "n", false,
		"perform a trial run, do not make changes")
	AddSpecFlags(cmdRelease.Flags())
	root.AddCommand(cmdRelease)
}

func runRelease(cmd *cobra.Command, args []string) {
	if len(args) > 0 {
		plog.Fatal("No args accepted")
	}

	spec := ChannelSpec()
	ctx := context.Background()
	client, err := auth.GoogleClient()
	if err != nil {
		plog.Fatalf("Authentication failed: %v", err)
	}

	src, err := storage.NewBucket(client, spec.SourceURL())
	if err != nil {
		plog.Fatal(err)
	}
	src.WriteDryRun(releaseDryRun)

	if err := src.Fetch(ctx); err != nil {
		plog.Fatal(err)
	}

	// Sanity check!
	if vertxt := src.Object(src.Prefix() + "version.txt"); vertxt == nil {
		verurl := src.URL().String() + "version.txt"
		plog.Fatalf("File not found: %s", verurl)
	}

	// Register GCE image if needed.
	doGCE(ctx, client, src, &spec)

	// Make Azure images public.
	doAzure(ctx, client, src, &spec)

	// Make AWS images public.
	doAWS(ctx, client, src, &spec)

	for _, dSpec := range spec.Destinations {
		dst, err := storage.NewBucket(client, dSpec.BaseURL)
		if err != nil {
			plog.Fatal(err)
		}
		dst.WriteDryRun(releaseDryRun)

		// Fetch parent directories non-recursively to re-index it later.
		for _, prefix := range dSpec.ParentPrefixes() {
			if err := dst.FetchPrefix(ctx, prefix, false); err != nil {
				plog.Fatal(err)
			}
		}

		// Fetch and sync each destination directory.
		for _, prefix := range dSpec.FinalPrefixes() {
			if err := dst.FetchPrefix(ctx, prefix, true); err != nil {
				plog.Fatal(err)
			}

			sync := index.NewSyncIndexJob(src, dst)
			sync.DestinationPrefix(prefix)
			sync.DirectoryHTML(dSpec.DirectoryHTML)
			sync.IndexHTML(dSpec.IndexHTML)
			sync.Delete(true)
			if dSpec.Title != "" {
				sync.Name(dSpec.Title)
			}
			if err := sync.Do(ctx); err != nil {
				plog.Fatal(err)
			}
		}

		// Now refresh the parent directory indexes.
		for _, prefix := range dSpec.ParentPrefixes() {
			parent := index.NewIndexJob(dst)
			parent.Prefix(prefix)
			parent.DirectoryHTML(dSpec.DirectoryHTML)
			parent.IndexHTML(dSpec.IndexHTML)
			parent.Recursive(false)
			parent.Delete(true)
			if dSpec.Title != "" {
				parent.Name(dSpec.Title)
			}
			if err := parent.Do(ctx); err != nil {
				plog.Fatal(err)
			}
		}
	}
}

func sanitizeVersion() string {
	v := strings.Replace(specVersion, ".", "-", -1)
	return strings.Replace(v, "+", "-", -1)
}

func doGCE(ctx context.Context, client *http.Client, src *storage.Bucket, spec *channelSpec) {
	if spec.GCE.Project == "" || spec.GCE.Image == "" {
		plog.Notice("GCE image creation disabled.")
		return
	}

	api, err := gcloud.New(&gcloud.Options{
		Project: spec.GCE.Project,
	})
	if err != nil {
		plog.Fatalf("GCE client failed: %v", err)
	}

	publishImage := func(image string) {
		if spec.GCE.Publish == "" {
			plog.Notice("GCE image name publishing disabled.")
			return
		}
		obj := gs.Object{
			Name:        src.Prefix() + spec.GCE.Publish,
			ContentType: "text/plain",
		}
		media := strings.NewReader(
			fmt.Sprintf("projects/%s/global/images/%s\n",
				spec.GCE.Project, image))
		if err := src.Upload(ctx, &obj, media); err != nil {
			plog.Fatal(err)
		}
	}

	nameVer := fmt.Sprintf("%s-%s-v", spec.GCE.Family, sanitizeVersion())
	date := time.Now().UTC()
	name := nameVer + date.Format("20060102")
	desc := fmt.Sprintf("%s, %s, %s published on %s", spec.GCE.Description,
		specVersion, specBoard, date.Format("2006-01-02"))

	images, err := api.ListImages(ctx, spec.GCE.Family+"-")
	if err != nil {
		plog.Fatal(err)
	}

	var conflicting []string
	for _, image := range images {
		if strings.HasPrefix(image.Name, nameVer) {
			conflicting = append(conflicting, image.Name)
		}
	}

	// Check for any with the same version but possibly different dates.
	if len(conflicting) > 1 {
		plog.Fatalf("Duplicate GCE images found: %v", conflicting)
	} else if len(conflicting) == 1 {
		plog.Noticef("GCE image already exists: %s", conflicting[0])
		publishImage(conflicting[0])
		return
	}

	obj := src.Object(src.Prefix() + spec.GCE.Image)
	if obj == nil {
		plog.Fatalf("GCE image not found %s%s", src.URL(), spec.GCE.Image)
	}

	if releaseDryRun {
		plog.Noticef("Would create GCE image %s", name)
		return
	}

	plog.Noticef("Creating GCE image %s", name)
	parsedVersion, err := semver.NewVersion(specVersion)
	if err != nil {
		plog.Fatalf("couldn't parse version %s: %v", specVersion, err)
	}
	disableMultiqueue := false
	if parsedVersion.LessThan(semver.Version{Major: 1409}) {
		disableMultiqueue = true
		plog.Noticef("Not enabling multiqueue for version %v", specVersion)
	}
	op, pending, err := api.CreateImage(&gcloud.ImageSpec{
		SourceImage:           obj.MediaLink,
		Family:                spec.GCE.Family,
		Name:                  name,
		Description:           desc,
		Licenses:              spec.GCE.Licenses,
		DisableSCSIMultiqueue: disableMultiqueue,
	}, false)
	if err != nil {
		plog.Fatalf("GCE image creation failed: %v", err)
	}

	plog.Infof("Waiting for image creation to finish...")
	pending.Interval = 3 * time.Second
	pending.Progress = func(_ string, _ time.Duration, op *compute.Operation) error {
		status := strings.ToLower(op.Status)
		if op.Progress != 0 {
			plog.Infof("Image creation is %s: %s % 2d%%", status, op.StatusMessage, op.Progress)
		} else {
			plog.Infof("Image creation is %s. %s", status, op.StatusMessage)
		}
		return nil
	}
	if err := pending.Wait(); err != nil {
		plog.Fatal(err)
	}
	plog.Info("Success!")

	publishImage(name)

	var pendings []*gcloud.Pending
	for _, old := range images {
		if old.Deprecated != nil && old.Deprecated.State != "" {
			continue
		}
		plog.Noticef("Deprecating old image %s", old.Name)
		pending, err := api.DeprecateImage(old.Name, gcloud.DeprecationStateDeprecated, op.TargetLink)
		if err != nil {
			plog.Fatal(err)
		}
		pending.Interval = 1 * time.Second
		pending.Timeout = 0
		pendings = append(pendings, pending)
	}

	if spec.GCE.Limit > 0 && len(images) > spec.GCE.Limit {
		plog.Noticef("Pruning %d GCE images.", len(images)-spec.GCE.Limit)
		plog.Notice("NOPE! JUST KIDDING, TODO")
	}

	plog.Infof("Waiting on %d operations.", len(pendings))
	for _, pending := range pendings {
		err := pending.Wait()
		if err != nil {
			plog.Fatal(err)
		}
	}
}

func doAzure(ctx context.Context, client *http.Client, src *storage.Bucket, spec *channelSpec) {
	if spec.Azure.StorageAccount == "" {
		plog.Notice("Azure image creation disabled.")
		return
	}

	prof, err := auth.ReadAzureProfile(azureProfile)
	if err != nil {
		plog.Fatalf("failed reading Azure profile: %v", err)
	}

	// channel name should be caps for azure image
	imageName := fmt.Sprintf("%s-%s-%s", spec.Azure.Offer, strings.Title(specChannel), specVersion)

	for _, environment := range spec.Azure.Environments {
		opt := prof.SubscriptionOptions(environment.SubscriptionName)
		if opt == nil {
			plog.Fatalf("couldn't find subscription %q", environment.SubscriptionName)
		}

		api, err := azure.New(opt)
		if err != nil {
			plog.Fatalf("failed to create Azure API: %v", err)
		}

		if releaseDryRun {
			// TODO(bgilbert): check that the image exists
			plog.Printf("Would share %q on %v", imageName, environment.SubscriptionName)
			continue
		} else {
			plog.Printf("Sharing %q on %v...", imageName, environment.SubscriptionName)
		}

		if err := api.ShareImage(imageName, "public"); err != nil {
			plog.Fatalf("failed to share image %q: %v", imageName, err)
		}
	}
}

func doAWS(ctx context.Context, client *http.Client, src *storage.Bucket, spec *channelSpec) {
	if spec.AWS.Image == "" {
		plog.Notice("AWS image creation disabled.")
		return
	}

	imageName := fmt.Sprintf("%v-%v-%v", spec.AWS.BaseName, specChannel, specVersion)
	imageName = regexp.MustCompile(`[^A-Za-z0-9()\\./_-]`).ReplaceAllLiteralString(imageName, "_")

	for _, part := range spec.AWS.Partitions {
		for _, region := range part.Regions {
			if releaseDryRun {
				plog.Printf("Checking for images in %v %v...", part.Name, region)
			} else {
				plog.Printf("Publishing images in %v %v...", part.Name, region)
			}

			api, err := aws.New(&aws.Options{
				CredentialsFile: awsCredentialsFile,
				Profile:         part.Profile,
				Region:          region,
			})
			if err != nil {
				plog.Fatalf("creating client for %v %v: %v", part.Name, region, err)
			}

			publish := func(imageName string) {
				imageID, err := api.FindImage(imageName)
				if err != nil {
					plog.Fatalf("couldn't find image %q in %v %v: %v", imageName, part.Name, region, err)
				}

				if !releaseDryRun {
					err := api.PublishImage(imageID)
					if err != nil {
						plog.Fatalf("couldn't publish image in %v %v: %v", part.Name, region, err)
					}
				}
			}
			publish(imageName)
			publish(imageName + "-hvm")
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"reflect"

	"github.com/cortexproject/cortex/pkg/util"
	"github.com/go-kit/kit/log/level"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/version"
	"github.com/weaveworks/common/logging"

	"github.com/grafana/loki/pkg/cfg"
	"github.com/grafana/loki/pkg/logentry/stages"
	"github.com/grafana/loki/pkg/promtail"
	"github.com/grafana/loki/pkg/promtail/config"
)

func init() {
	prometheus.MustRegister(version.NewCollector("promtail"))
}

func main() {
	printVersion := flag.Bool("version", false, "Print this builds version information")

	var config config.Config
	if err := cfg.Parse(&config); err != nil {
		level.Error(util.Logger).Log("msg", "parsing config", "error", err)
		os.Exit(1)
	}
	if *printVersion {
		fmt.Print(version.Print("promtail"))
		os.Exit(0)
	}

	// Init the logger which will honor the log level set in cfg.Server
	if reflect.DeepEqual(&config.ServerConfig.Config.LogLevel, &logging.Level{}) {
		level.Error(util.Logger).Log("msg", "invalid log level")
		os.Exit(1)
	}
	util.InitLogger(&config.ServerConfig.Config)

	// Set the global debug variable in the stages package which is used to conditionally log
	// debug messages which otherwise cause huge allocations processing log lines for log messages never printed
	if config.ServerConfig.Config.LogLevel.String() == "debug" {
		stages.Debug = true
	}

	p, err := promtail.New(config)
	if err != nil {
		level.Error(util.Logger).Log("msg", "error creating promtail", "error", err)
		os.Exit(1)
	}

	level.Info(util.Logger).Log("msg", "Starting Promtail", "version", version.Info())

	if err := p.Run(); err != nil {
		level.Error(util.Logger).Log("msg", "error starting promtail", "error", err)
		os.Exit(1)
	}

	p.Shutdown()
}
<commit_msg>promtail: fix error message displaying on invalid config file<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"reflect"

	"github.com/cortexproject/cortex/pkg/util"
	"github.com/go-kit/kit/log/level"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/version"
	"github.com/weaveworks/common/logging"

	"github.com/grafana/loki/pkg/cfg"
	"github.com/grafana/loki/pkg/logentry/stages"
	"github.com/grafana/loki/pkg/promtail"
	"github.com/grafana/loki/pkg/promtail/config"
)

func init() {
	prometheus.MustRegister(version.NewCollector("promtail"))
}

func main() {
	printVersion := flag.Bool("version", false, "Print this builds version information")

	// Load config, merging config file and CLI flags
	var config config.Config
	if err := cfg.Parse(&config); err != nil {
		fmt.Println("Unable to parse config:", err)
		os.Exit(1)
	}

	// Handle -version CLI flag
	if *printVersion {
		fmt.Println(version.Print("promtail"))
		os.Exit(0)
	}

	// Init the logger which will honor the log level set in cfg.Server
	if reflect.DeepEqual(&config.ServerConfig.Config.LogLevel, &logging.Level{}) {
		fmt.Println("Invalid log level")
		os.Exit(1)
	}
	util.InitLogger(&config.ServerConfig.Config)

	// Set the global debug variable in the stages package which is used to conditionally log
	// debug messages which otherwise cause huge allocations processing log lines for log messages never printed
	if config.ServerConfig.Config.LogLevel.String() == "debug" {
		stages.Debug = true
	}

	p, err := promtail.New(config)
	if err != nil {
		level.Error(util.Logger).Log("msg", "error creating promtail", "error", err)
		os.Exit(1)
	}

	level.Info(util.Logger).Log("msg", "Starting Promtail", "version", version.Info())

	if err := p.Run(); err != nil {
		level.Error(util.Logger).Log("msg", "error starting promtail", "error", err)
		os.Exit(1)
	}

	p.Shutdown()
}
<|endoftext|>