<commit_before>package network

import (
	"fmt"
	"net"
	"strings"

	"github.com/pkg/errors"

	lxd "github.com/lxc/lxd/client"
	"github.com/lxc/lxd/lxd/cluster"
	"github.com/lxc/lxd/lxd/instance"
	"github.com/lxc/lxd/lxd/state"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	log "github.com/lxc/lxd/shared/log15"
	"github.com/lxc/lxd/shared/logger"
	"github.com/lxc/lxd/shared/logging"
)

// DHCPRange represents a range of IPs from start to end.
type DHCPRange struct {
	Start net.IP
	End   net.IP
}

// common represents a generic LXD network.
type common struct {
	logger      logger.Logger
	state       *state.State
	id          int64
	name        string
	netType     string
	description string
	config      map[string]string
}

// init initialise internal variables.
func (n *common) init(state *state.State, id int64, name string, netType string, description string, config map[string]string) {
	n.logger = logging.AddContext(logger.Log, log.Ctx{"driver": netType, "network": name})
	n.id = id
	n.name = name
	n.netType = netType
	n.config = config
	n.state = state
	n.description = description
}

// commonRules returns a map of config rules common to all drivers.
func (n *common) commonRules() map[string]func(string) error {
	return map[string]func(string) error{}
}

// validate a network config against common rules and optional driver specific rules.
func (n *common) validate(config map[string]string, driverRules map[string]func(value string) error) error {
	checkedFields := map[string]struct{}{}

	// Get rules common for all drivers.
	rules := n.commonRules()

	// Merge driver specific rules into common rules.
	for field, validator := range driverRules {
		rules[field] = validator
	}

	// Run the validator against each field.
	for k, validator := range rules {
		checkedFields[k] = struct{}{} //Mark field as checked.
		err := validator(config[k])
		if err != nil {
			return errors.Wrapf(err, "Invalid value for network %q option %q", n.name, k)
		}
	}

	// Look for any unchecked fields, as these are unknown fields and validation should fail.
	for k := range config {
		_, checked := checkedFields[k]
		if checked {
			continue
		}

		// User keys are not validated.
		if strings.HasPrefix(k, "user.") {
			continue
		}

		return fmt.Errorf("Invalid option for network %q option %q", n.name, k)
	}

	return nil
}

// Name returns the network name.
func (n *common) Name() string {
	return n.name
}

// Type returns the network type.
func (n *common) Type() string {
	return n.netType
}

// Config returns the network config.
func (n *common) Config() map[string]string {
	return n.config
}

// IsUsed returns whether the network is used by any instances.
func (n *common) IsUsed() bool {
	// Look for instances using the interface
	insts, err := instance.LoadFromAllProjects(n.state)
	if err != nil {
		return true
	}

	for _, inst := range insts {
		if IsInUseByInstance(inst, n.name) {
			return true
		}
	}

	return false
}

// HasDHCPv4 indicates whether the network has DHCPv4 enabled.
func (n *common) HasDHCPv4() bool {
	if n.config["ipv4.dhcp"] == "" || shared.IsTrue(n.config["ipv4.dhcp"]) {
		return true
	}

	return false
}

// HasDHCPv6 indicates whether the network has DHCPv6 enabled (includes stateless SLAAC router advertisement mode).
// Technically speaking stateless SLAAC RA mode isn't DHCPv6, but for consistency with LXD's config paradigm, DHCP
// here means "an ability to automatically allocate IPs and routes", rather than stateful DHCP with leases.
// To check if true stateful DHCPv6 is enabled check the "ipv6.dhcp.stateful" config key.
func (n *common) HasDHCPv6() bool {
	if n.config["ipv6.dhcp"] == "" || shared.IsTrue(n.config["ipv6.dhcp"]) {
		return true
	}

	return false
}

// DHCPv4Ranges returns a parsed set of DHCPv4 ranges for this network.
func (n *common) DHCPv4Ranges() []DHCPRange {
	dhcpRanges := make([]DHCPRange, 0)
	if n.config["ipv4.dhcp.ranges"] != "" {
		for _, r := range strings.Split(n.config["ipv4.dhcp.ranges"], ",") {
			parts := strings.SplitN(strings.TrimSpace(r), "-", 2)
			if len(parts) == 2 {
				startIP := net.ParseIP(parts[0])
				endIP := net.ParseIP(parts[1])
				dhcpRanges = append(dhcpRanges, DHCPRange{
					Start: startIP.To4(),
					End:   endIP.To4(),
				})
			}
		}
	}

	return dhcpRanges
}

// DHCPv6Ranges returns a parsed set of DHCPv6 ranges for this network.
func (n *common) DHCPv6Ranges() []DHCPRange {
	dhcpRanges := make([]DHCPRange, 0)
	if n.config["ipv6.dhcp.ranges"] != "" {
		for _, r := range strings.Split(n.config["ipv6.dhcp.ranges"], ",") {
			parts := strings.SplitN(strings.TrimSpace(r), "-", 2)
			if len(parts) == 2 {
				startIP := net.ParseIP(parts[0])
				endIP := net.ParseIP(parts[1])
				dhcpRanges = append(dhcpRanges, DHCPRange{
					Start: startIP.To16(),
					End:   endIP.To16(),
				})
			}
		}
	}

	return dhcpRanges
}

// update the internal config variables, and if not cluster notification, notifies all nodes and updates database.
func (n *common) update(applyNetwork api.NetworkPut, clusterNotification bool) error {
	// Update internal config before database has been updated (so that if update is a notification we apply
	// the config being supplied and not that in the database).
	n.description = applyNetwork.Description
	n.config = applyNetwork.Config

	// If this update isn't coming via a cluster notification itself, then notify all nodes of change and then
	// update the database.
	if !clusterNotification {
		// Notify all other nodes to update the network.
		notifier, err := cluster.NewNotifier(n.state, n.state.Endpoints.NetworkCert(), cluster.NotifyAll)
		if err != nil {
			return err
		}

		err = notifier(func(client lxd.InstanceServer) error {
			return client.UpdateNetwork(n.name, applyNetwork, "")
		})
		if err != nil {
			return err
		}

		// Update the database.
		err = n.state.Cluster.UpdateNetwork(n.name, applyNetwork.Description, applyNetwork.Config)
		if err != nil {
			return err
		}
	}

	return nil
}

// configChanged compares supplied new config with existing config. Returns a boolean indicating if differences in
// the config or description were found (and the database record needs updating), and a list of non-user config
// keys that have changed, and a copy of the current internal network config that can be used to revert if needed.
func (n *common) configChanged(newNetwork api.NetworkPut) (bool, []string, api.NetworkPut, error) {
	// Backup the current state.
	oldNetwork := api.NetworkPut{
		Description: n.description,
		Config:      map[string]string{},
	}

	err := shared.DeepCopy(&n.config, &oldNetwork.Config)
	if err != nil {
		return false, nil, oldNetwork, err
	}

	// Diff the configurations.
	changedKeys := []string{}
	dbUpdateNeeded := false

	if newNetwork.Description != n.description {
		dbUpdateNeeded = true
	}

	for k, v := range oldNetwork.Config {
		if v != newNetwork.Config[k] {
			dbUpdateNeeded = true

			// Add non-user changed key to list of changed keys.
			if !strings.HasPrefix(k, "user.") && !shared.StringInSlice(k, changedKeys) {
				changedKeys = append(changedKeys, k)
			}
		}
	}

	for k, v := range newNetwork.Config {
		if v != oldNetwork.Config[k] {
			dbUpdateNeeded = true

			// Add non-user changed key to list of changed keys.
			if !strings.HasPrefix(k, "user.") && !shared.StringInSlice(k, changedKeys) {
				changedKeys = append(changedKeys, k)
			}
		}
	}

	return dbUpdateNeeded, changedKeys, oldNetwork, nil
}
<commit_msg>lxd/network/driver/common: Removes stuttering on "common" in validation rules function<commit_after>package network

import (
	"fmt"
	"net"
	"strings"

	"github.com/pkg/errors"

	lxd "github.com/lxc/lxd/client"
	"github.com/lxc/lxd/lxd/cluster"
	"github.com/lxc/lxd/lxd/instance"
	"github.com/lxc/lxd/lxd/state"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	log "github.com/lxc/lxd/shared/log15"
	"github.com/lxc/lxd/shared/logger"
	"github.com/lxc/lxd/shared/logging"
)

// DHCPRange represents a range of IPs from start to end.
type DHCPRange struct {
	Start net.IP
	End   net.IP
}

// common represents a generic LXD network.
type common struct {
	logger      logger.Logger
	state       *state.State
	id          int64
	name        string
	netType     string
	description string
	config      map[string]string
}

// init initialise internal variables.
func (n *common) init(state *state.State, id int64, name string, netType string, description string, config map[string]string) {
	n.logger = logging.AddContext(logger.Log, log.Ctx{"driver": netType, "network": name})
	n.id = id
	n.name = name
	n.netType = netType
	n.config = config
	n.state = state
	n.description = description
}

// validationRules returns a map of config rules common to all drivers.
func (n *common) validationRules() map[string]func(string) error {
	return map[string]func(string) error{}
}

// validate a network config against common rules and optional driver specific rules.
func (n *common) validate(config map[string]string, driverRules map[string]func(value string) error) error {
	checkedFields := map[string]struct{}{}

	// Get rules common for all drivers.
	rules := n.validationRules()

	// Merge driver specific rules into common rules.
	for field, validator := range driverRules {
		rules[field] = validator
	}

	// Run the validator against each field.
	for k, validator := range rules {
		checkedFields[k] = struct{}{} //Mark field as checked.
		err := validator(config[k])
		if err != nil {
			return errors.Wrapf(err, "Invalid value for network %q option %q", n.name, k)
		}
	}

	// Look for any unchecked fields, as these are unknown fields and validation should fail.
	for k := range config {
		_, checked := checkedFields[k]
		if checked {
			continue
		}

		// User keys are not validated.
		if strings.HasPrefix(k, "user.") {
			continue
		}

		return fmt.Errorf("Invalid option for network %q option %q", n.name, k)
	}

	return nil
}

// Name returns the network name.
func (n *common) Name() string {
	return n.name
}

// Type returns the network type.
func (n *common) Type() string {
	return n.netType
}

// Config returns the network config.
func (n *common) Config() map[string]string {
	return n.config
}

// IsUsed returns whether the network is used by any instances.
func (n *common) IsUsed() bool {
	// Look for instances using the interface
	insts, err := instance.LoadFromAllProjects(n.state)
	if err != nil {
		return true
	}

	for _, inst := range insts {
		if IsInUseByInstance(inst, n.name) {
			return true
		}
	}

	return false
}

// HasDHCPv4 indicates whether the network has DHCPv4 enabled.
func (n *common) HasDHCPv4() bool {
	if n.config["ipv4.dhcp"] == "" || shared.IsTrue(n.config["ipv4.dhcp"]) {
		return true
	}

	return false
}

// HasDHCPv6 indicates whether the network has DHCPv6 enabled (includes stateless SLAAC router advertisement mode).
// Technically speaking stateless SLAAC RA mode isn't DHCPv6, but for consistency with LXD's config paradigm, DHCP
// here means "an ability to automatically allocate IPs and routes", rather than stateful DHCP with leases.
// To check if true stateful DHCPv6 is enabled check the "ipv6.dhcp.stateful" config key.
func (n *common) HasDHCPv6() bool {
	if n.config["ipv6.dhcp"] == "" || shared.IsTrue(n.config["ipv6.dhcp"]) {
		return true
	}

	return false
}

// DHCPv4Ranges returns a parsed set of DHCPv4 ranges for this network.
func (n *common) DHCPv4Ranges() []DHCPRange {
	dhcpRanges := make([]DHCPRange, 0)
	if n.config["ipv4.dhcp.ranges"] != "" {
		for _, r := range strings.Split(n.config["ipv4.dhcp.ranges"], ",") {
			parts := strings.SplitN(strings.TrimSpace(r), "-", 2)
			if len(parts) == 2 {
				startIP := net.ParseIP(parts[0])
				endIP := net.ParseIP(parts[1])
				dhcpRanges = append(dhcpRanges, DHCPRange{
					Start: startIP.To4(),
					End:   endIP.To4(),
				})
			}
		}
	}

	return dhcpRanges
}

// DHCPv6Ranges returns a parsed set of DHCPv6 ranges for this network.
func (n *common) DHCPv6Ranges() []DHCPRange {
	dhcpRanges := make([]DHCPRange, 0)
	if n.config["ipv6.dhcp.ranges"] != "" {
		for _, r := range strings.Split(n.config["ipv6.dhcp.ranges"], ",") {
			parts := strings.SplitN(strings.TrimSpace(r), "-", 2)
			if len(parts) == 2 {
				startIP := net.ParseIP(parts[0])
				endIP := net.ParseIP(parts[1])
				dhcpRanges = append(dhcpRanges, DHCPRange{
					Start: startIP.To16(),
					End:   endIP.To16(),
				})
			}
		}
	}

	return dhcpRanges
}

// update the internal config variables, and if not cluster notification, notifies all nodes and updates database.
func (n *common) update(applyNetwork api.NetworkPut, clusterNotification bool) error {
	// Update internal config before database has been updated (so that if update is a notification we apply
	// the config being supplied and not that in the database).
	n.description = applyNetwork.Description
	n.config = applyNetwork.Config

	// If this update isn't coming via a cluster notification itself, then notify all nodes of change and then
	// update the database.
	if !clusterNotification {
		// Notify all other nodes to update the network.
		notifier, err := cluster.NewNotifier(n.state, n.state.Endpoints.NetworkCert(), cluster.NotifyAll)
		if err != nil {
			return err
		}

		err = notifier(func(client lxd.InstanceServer) error {
			return client.UpdateNetwork(n.name, applyNetwork, "")
		})
		if err != nil {
			return err
		}

		// Update the database.
		err = n.state.Cluster.UpdateNetwork(n.name, applyNetwork.Description, applyNetwork.Config)
		if err != nil {
			return err
		}
	}

	return nil
}

// configChanged compares supplied new config with existing config. Returns a boolean indicating if differences in
// the config or description were found (and the database record needs updating), and a list of non-user config
// keys that have changed, and a copy of the current internal network config that can be used to revert if needed.
func (n *common) configChanged(newNetwork api.NetworkPut) (bool, []string, api.NetworkPut, error) {
	// Backup the current state.
	oldNetwork := api.NetworkPut{
		Description: n.description,
		Config:      map[string]string{},
	}

	err := shared.DeepCopy(&n.config, &oldNetwork.Config)
	if err != nil {
		return false, nil, oldNetwork, err
	}

	// Diff the configurations.
	changedKeys := []string{}
	dbUpdateNeeded := false

	if newNetwork.Description != n.description {
		dbUpdateNeeded = true
	}

	for k, v := range oldNetwork.Config {
		if v != newNetwork.Config[k] {
			dbUpdateNeeded = true

			// Add non-user changed key to list of changed keys.
			if !strings.HasPrefix(k, "user.") && !shared.StringInSlice(k, changedKeys) {
				changedKeys = append(changedKeys, k)
			}
		}
	}

	for k, v := range newNetwork.Config {
		if v != oldNetwork.Config[k] {
			dbUpdateNeeded = true

			// Add non-user changed key to list of changed keys.
			if !strings.HasPrefix(k, "user.") && !shared.StringInSlice(k, changedKeys) {
				changedKeys = append(changedKeys, k)
			}
		}
	}

	return dbUpdateNeeded, changedKeys, oldNetwork, nil
}
<|endoftext|>


<commit_before>package updater

import (
	"archive/zip"
	"github.com/sellweek/TOGY/config"
	"io"
	"net/http"
	"os"
	"strconv"
	"time"
)

func DownloadConfig(c *config.Config, destFile string) error {
	return downloadFile(c.UpdateURL+"/config/download?client="+c.Name, destFile)
}

func DownloadBroadcast(c *config.Config, ft string, destDir string) (err error) {
	srcUrl := c.UpdateURL + "/presentation/active/download?client=" + c.Name

	if ft != "zip" {
		err = downloadFile(srcUrl, destDir+string(os.PathSeparator)+"broadcast."+ft)
		return
	}

	tempFileName := os.TempDir() + string(os.PathSeparator) + "unzip-" + strconv.Itoa(int(time.Now().Unix())) + ".zip"

	err = downloadFile(srcUrl, tempFileName)
	if err != nil {
		return
	}

	err = unzip(destDir, tempFileName)

	return
}

func ColdStart(c *config.Config) (err error) {
	ui, err := GetInfo(c)
	if err != nil {
		return err
	}
	err = DownloadConfig(c, c.CentralPath)
	if err != nil {
		return err
	}
	err = DownloadBroadcast(c, ui.FileType, c.BroadcastDir)
	return
}

func unzip(dirname, fn string) (err error) {
	r, err := zip.OpenReader(fn)
	if err != nil {
		return
	}

	for _, sf := range r.File {
		var fr io.ReadCloser
		fr, err = sf.Open()
		if err != nil {
			return
		}
		var df *os.File
		df, err = os.Create(dirname + string(os.PathSeparator) + sf.Name)
		if err != nil {
			return
		}
		defer df.Close()
		_, err = io.Copy(df, fr)
		if err != nil {
			return
		}
	}
	return nil
}

func downloadFile(src, dest string) (err error) {
	resp, err := http.Get(src)
	defer resp.Body.Close()

	f, err := os.Create(dest)
	if err != nil {
		return
	}
	defer f.Close()

	_, err = io.Copy(f, resp.Body)
	return
}
<commit_msg>Added documentation to download.go<commit_after>package updater

import (
	"archive/zip"
	"github.com/sellweek/TOGY/config"
	"io"
	"net/http"
	"os"
	"strconv"
	"time"
)

//DownloadConfig ownloads a new centralConfig.json from the server
//into the given path.
func DownloadConfig(c *config.Config, destFile string) error {
	return downloadFile(c.UpdateURL+"/config/download?client="+c.Name, destFile)
}

//DownloadBroadcast downloads a new broadcast from the server into a
//given directory, unzipping it, if it has .zip extension.
func DownloadBroadcast(c *config.Config, ft string, destDir string) (err error) {
	srcUrl := c.UpdateURL + "/presentation/active/download?client=" + c.Name

	if ft != "zip" {
		err = downloadFile(srcUrl, destDir+string(os.PathSeparator)+"broadcast."+ft)
		return
	}

	tempFileName := os.TempDir() + string(os.PathSeparator) + "unzip-" + strconv.Itoa(int(time.Now().Unix())) + ".zip"

	err = downloadFile(srcUrl, tempFileName)
	if err != nil {
		return
	}

	err = unzip(destDir, tempFileName)

	return
}

//ColdStart downloads central config and the newest broadcast
//into folders specified in config, without announcing
//their downloads.
func ColdStart(c *config.Config) (err error) {
	ui, err := GetInfo(c)
	if err != nil {
		return err
	}
	err = DownloadConfig(c, c.CentralPath)
	if err != nil {
		return err
	}
	err = DownloadBroadcast(c, ui.FileType, c.BroadcastDir)
	return
}

//Unzip unzips a file into given folder.
//
//WARNING: the unzipping is not recursive therefore it doesn't support
//zip files with folders.
func unzip(dirname, fn string) (err error) {
	r, err := zip.OpenReader(fn)
	if err != nil {
		return
	}

	for _, sf := range r.File {
		var fr io.ReadCloser
		fr, err = sf.Open()
		if err != nil {
			return
		}
		var df *os.File
		df, err = os.Create(dirname + string(os.PathSeparator) + sf.Name)
		if err != nil {
			return
		}
		defer df.Close()
		_, err = io.Copy(df, fr)
		if err != nil {
			return
		}
	}
	return nil
}

//Downloads a file from given URL into given path using http.Get
func downloadFile(src, dest string) (err error) {
	resp, err := http.Get(src)
	defer resp.Body.Close()

	f, err := os.Create(dest)
	if err != nil {
		return
	}
	defer f.Close()

	_, err = io.Copy(f, resp.Body)
	return
}
<|endoftext|>


<commit_before>// ***************************************************************************
//
//  Copyright 2017 David (Dizzy) Smith, dizzyd@dizzyd.com
//
//  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 (
	"flag"
	"fmt"
	"log"
	"os"
	"sort"
	"strings"
)

var version string

var ARG_MMC bool
var ARG_VERBOSE bool
var ARG_SKIPMODS bool

type command struct {
	Fn        func() error
	Desc      string
	ArgsCount int
	Args      string
}

var gCommands = map[string]command{
	"pack.create": command{
		Fn:        cmdPackCreate,
		Desc:      "Create a new mod pack",
		ArgsCount: 2,
		Args:      "<directory> <minecraft version> [<forge version>]",
	},
	"pack.install": command{
		Fn:        cmdPackInstall,
		Desc:      "Install a mod pack, optionally using a URL to download",
		ArgsCount: 1,
		Args:      "<directory> [<url>]",
	},
	"info": command{
		Fn:        cmdInfo,
		Desc:      "Show runtime info",
		ArgsCount: 0,
	},
	"mod.list": command{
		Fn:        cmdModList,
		Desc:      "List mods matching a name and Minecraft version",
		ArgsCount: 1,
		Args:      "<mod name> [<minecraft version>]",
	},
	"mod.select": command{
		Fn:        cmdModSelect,
		Desc:      "Select a mod to include in the specified pack",
		ArgsCount: 2,
		Args:      "<directory> <mod name or URL> [<tag>]",
	},
	"mod.select.client": command{
		Fn:        cmdModSelectClient,
		Desc:      "Select a client-side only mod to include in the specified pack",
		ArgsCount: 2,
		Args:      "<directory> <mod name or URL> [<tag>]",
	},
	"server.install": command{
		Fn:        cmdServerInstall,
		Desc:      "Install a Minecraft server using an existing pack",
		ArgsCount: 1,
		Args:      "<directory>",
	},
	"db.update": command{
		Fn:        cmdDBUpdate,
		Desc:      "Update local database of available mods",
		ArgsCount: 0,
	},
	"forge.list": command{
		Fn:        cmdForgeList,
		Desc:      "List available versions of Forge",
		ArgsCount: 1,
		Args:      "<minecraft version>",
	},
}

func cmdPackCreate() error {
	dir := flag.Arg(1)
	minecraftVsn := flag.Arg(2)
	forgeVsn := flag.Arg(3)

	// If no forge version was specified, open the database and find
	// a recommended one
	if forgeVsn == "" {
		db, err := OpenDatabase()
		if err != nil {
			return err
		}

		forgeVsn, err = db.lookupForgeVsn(minecraftVsn)
		if err != nil {
			return err
		}
	}

	// Create a new pack directory
	cp, err := NewModPack(dir, false, ARG_MMC)
	if err != nil {
		return err
	}

	// Create the manifest for this new pack
	err = cp.createManifest(dir, minecraftVsn, forgeVsn)
	if err != nil {
		return err
	}

	// Create the launcher profile (and install forge if necessary)
	err = cp.createLauncherProfile()
	if err != nil {
		return err
	}

	return nil
}

func cmdPackInstall() error {
	dir := flag.Arg(1)
	url := flag.Arg(2)

	// Only require a manifest if we're not installing from a URL
	requireManifest := (url == "")

	cp, err := NewModPack(dir, requireManifest, ARG_MMC)
	if err != nil {
		return err
	}

	if url != "" {
		// Download the pack
		err = cp.download(url)
		if err != nil {
			return err
		}

		// Process manifest
		err = cp.processManifest()
		if err != nil {
			return err
		}

		// Install overrides from the modpack; this is a bit of a misnomer since
		// under usual circumstances there are no mods in the modpack file that
		// will be also be downloaded
		err = cp.installOverrides()
		if err != nil {
			return err
		}
	}

	// If the -mmc flag is provided, don't create a launcher profile; just generate
	// an instance.cfg for MultiMC to use
	if ARG_MMC == true {
		err = cp.generateMMCConfig()
		if err != nil {
			return err
		}
	} else {
		// Create launcher profile
		err = cp.createLauncherProfile()
		if err != nil {
			return err
		}
	}

	if ARG_SKIPMODS == false {
		// Install mods (include client-side only mods)
		err = cp.installMods(true)
		if err != nil {
			return err
		}
	}

	return nil
}

func cmdInfo() error {
	// Try to retrieve the latest available version info
	publishedVsn, err := getLatestVersion("release")

	if err != nil && ARG_VERBOSE {
		fmt.Printf("%s\n", err)
	}

	if err == nil && publishedVsn != "" && version != publishedVsn {
		fmt.Printf("Version: %s (%s is available for download)\n", version, publishedVsn)
	} else {
		fmt.Printf("Version: %s\n", version)
	}

	// Print the environment
	fmt.Printf("Environment:\n")
	fmt.Printf("* Minecraft dir: %s\n", env().MinecraftDir)
	fmt.Printf("* mcdex dir: %s\n", env().McdexDir)
	fmt.Printf("* Java dir: %s\n", env().JavaDir)
	return nil
}

func cmdModSelect() error {
	return _modSelect(false)
}

func cmdModSelectClient() error {
	return _modSelect(true)
}

func _modSelect(clientOnly bool) error {
	dir := flag.Arg(1)
	mod := flag.Arg(2)
	tag := flag.Arg(3)

	// Try to open the mod pack
	cp, err := NewModPack(dir, true, ARG_MMC)
	if err != nil {
		return err
	}

	// If the mod doesn't start with https://, assume it's a name and try to look it up
	if !strings.HasPrefix("https://", mod) {
		db, err := OpenDatabase()
		if err != nil {
			return err
		}

		// Lookup the URL for the mod
		mod, err = db.findModFile(mod, cp.minecraftVersion())
		if err != nil {
			return err
		}

	} else if strings.Contains(mod, "minecraft.curseforge.com") && tag == "" {
		return fmt.Errorf("Non-CurseForge URLs must include a tag argument")
	}

	err = cp.selectMod(mod, tag, clientOnly)
	if err != nil {
		return err
	}

	return nil
}

func cmdModList() error {
	name := flag.Arg(1)
	mcvsn := flag.Arg(2)

	db, err := OpenDatabase()
	if err != nil {
		return err
	}

	return db.listMods(name, mcvsn)
}

func cmdForgeList() error {
	mcvsn := flag.Arg(1)

	db, err := OpenDatabase()
	if err != nil {
		return err
	}

	return db.listForge(mcvsn, ARG_VERBOSE)
}

func cmdServerInstall() error {
	dir := flag.Arg(1)

	if ARG_MMC == true {
		return fmt.Errorf("-mmc arg not supported when installing a server")
	}

	// Open the pack; we require the manifest and any
	// config files to already be present
	cp, err := NewModPack(dir, true, false)
	if err != nil {
		return err
	}

	// Install the server jar, Forge and dependencies
	err = cp.installServer()
	if err != nil {
		return err
	}

	// Make sure all mods are installed (do NOT include client-side only)
	err = cp.installMods(false)
	if err != nil {
		return err
	}

	return nil
	// Setup the command-line
	// java -jar <forge.jar>
}

func cmdDBUpdate() error {
	return InstallDatabase()
}

func console(f string, args ...interface{}) {
	fmt.Printf(f, args...)
}

func usage() {
	console("usage: mcdex [<options>] <command> [<args>]\n")
	console(" commands:\n")

	// Sort the list of keys in gCommands
	keys := []string{}
	for k := range gCommands {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, cmd := range keys {
		console(" - %s: %s\n", cmd, gCommands[cmd].Desc)
	}
}

func main() {
	// Register
	flag.BoolVar(&ARG_MMC, "mmc", false, "Generate MultiMC instance.cfg when installing a pack")
	flag.BoolVar(&ARG_VERBOSE, "v", false, "Enable verbose logging of operations")
	flag.BoolVar(&ARG_SKIPMODS, "skipmods", false, "Skip download of mods when installing a pack")

	// Process command-line args
	flag.Parse()
	if !flag.Parsed() || flag.NArg() < 1 {
		usage()
		os.Exit(-1)
	}

	// Initialize our environment
	err := initEnv()
	if err != nil {
		log.Fatalf("Failed to initialize: %s\n", err)
	}

	commandName := flag.Arg(0)
	command, exists := gCommands[commandName]
	if !exists {
		console("ERROR: unknown command '%s'\n", commandName)
		usage()
		os.Exit(-1)
	}

	// Check that the required number of arguments is present
	if flag.NArg() < command.ArgsCount+1 {
		console("ERROR: insufficient arguments for %s\n", commandName)
		console("usage: mcdex %s %s\n", commandName, command.Args)
		os.Exit(-1)
	}

	err = command.Fn()
	if err != nil {
		log.Fatalf("%+v\n", err)
	}
}
<commit_msg>Use strings.[Contains|HasPrefix] appropriately<commit_after>// ***************************************************************************
//
//  Copyright 2017 David (Dizzy) Smith, dizzyd@dizzyd.com
//
//  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 (
	"flag"
	"fmt"
	"log"
	"os"
	"sort"
	"strings"
)

var version string

var ARG_MMC bool
var ARG_VERBOSE bool
var ARG_SKIPMODS bool

type command struct {
	Fn        func() error
	Desc      string
	ArgsCount int
	Args      string
}

var gCommands = map[string]command{
	"pack.create": command{
		Fn:        cmdPackCreate,
		Desc:      "Create a new mod pack",
		ArgsCount: 2,
		Args:      "<directory> <minecraft version> [<forge version>]",
	},
	"pack.install": command{
		Fn:        cmdPackInstall,
		Desc:      "Install a mod pack, optionally using a URL to download",
		ArgsCount: 1,
		Args:      "<directory> [<url>]",
	},
	"info": command{
		Fn:        cmdInfo,
		Desc:      "Show runtime info",
		ArgsCount: 0,
	},
	"mod.list": command{
		Fn:        cmdModList,
		Desc:      "List mods matching a name and Minecraft version",
		ArgsCount: 1,
		Args:      "<mod name> [<minecraft version>]",
	},
	"mod.select": command{
		Fn:        cmdModSelect,
		Desc:      "Select a mod to include in the specified pack",
		ArgsCount: 2,
		Args:      "<directory> <mod name or URL> [<tag>]",
	},
	"mod.select.client": command{
		Fn:        cmdModSelectClient,
		Desc:      "Select a client-side only mod to include in the specified pack",
		ArgsCount: 2,
		Args:      "<directory> <mod name or URL> [<tag>]",
	},
	"server.install": command{
		Fn:        cmdServerInstall,
		Desc:      "Install a Minecraft server using an existing pack",
		ArgsCount: 1,
		Args:      "<directory>",
	},
	"db.update": command{
		Fn:        cmdDBUpdate,
		Desc:      "Update local database of available mods",
		ArgsCount: 0,
	},
	"forge.list": command{
		Fn:        cmdForgeList,
		Desc:      "List available versions of Forge",
		ArgsCount: 1,
		Args:      "<minecraft version>",
	},
}

func cmdPackCreate() error {
	dir := flag.Arg(1)
	minecraftVsn := flag.Arg(2)
	forgeVsn := flag.Arg(3)

	// If no forge version was specified, open the database and find
	// a recommended one
	if forgeVsn == "" {
		db, err := OpenDatabase()
		if err != nil {
			return err
		}

		forgeVsn, err = db.lookupForgeVsn(minecraftVsn)
		if err != nil {
			return err
		}
	}

	// Create a new pack directory
	cp, err := NewModPack(dir, false, ARG_MMC)
	if err != nil {
		return err
	}

	// Create the manifest for this new pack
	err = cp.createManifest(dir, minecraftVsn, forgeVsn)
	if err != nil {
		return err
	}

	// Create the launcher profile (and install forge if necessary)
	err = cp.createLauncherProfile()
	if err != nil {
		return err
	}

	return nil
}

func cmdPackInstall() error {
	dir := flag.Arg(1)
	url := flag.Arg(2)

	// Only require a manifest if we're not installing from a URL
	requireManifest := (url == "")

	cp, err := NewModPack(dir, requireManifest, ARG_MMC)
	if err != nil {
		return err
	}

	if url != "" {
		// Download the pack
		err = cp.download(url)
		if err != nil {
			return err
		}

		// Process manifest
		err = cp.processManifest()
		if err != nil {
			return err
		}

		// Install overrides from the modpack; this is a bit of a misnomer since
		// under usual circumstances there are no mods in the modpack file that
		// will be also be downloaded
		err = cp.installOverrides()
		if err != nil {
			return err
		}
	}

	// If the -mmc flag is provided, don't create a launcher profile; just generate
	// an instance.cfg for MultiMC to use
	if ARG_MMC == true {
		err = cp.generateMMCConfig()
		if err != nil {
			return err
		}
	} else {
		// Create launcher profile
		err = cp.createLauncherProfile()
		if err != nil {
			return err
		}
	}

	if ARG_SKIPMODS == false {
		// Install mods (include client-side only mods)
		err = cp.installMods(true)
		if err != nil {
			return err
		}
	}

	return nil
}

func cmdInfo() error {
	// Try to retrieve the latest available version info
	publishedVsn, err := getLatestVersion("release")

	if err != nil && ARG_VERBOSE {
		fmt.Printf("%s\n", err)
	}

	if err == nil && publishedVsn != "" && version != publishedVsn {
		fmt.Printf("Version: %s (%s is available for download)\n", version, publishedVsn)
	} else {
		fmt.Printf("Version: %s\n", version)
	}

	// Print the environment
	fmt.Printf("Environment:\n")
	fmt.Printf("* Minecraft dir: %s\n", env().MinecraftDir)
	fmt.Printf("* mcdex dir: %s\n", env().McdexDir)
	fmt.Printf("* Java dir: %s\n", env().JavaDir)
	return nil
}

func cmdModSelect() error {
	return _modSelect(false)
}

func cmdModSelectClient() error {
	return _modSelect(true)
}

func _modSelect(clientOnly bool) error {
	dir := flag.Arg(1)
	mod := flag.Arg(2)
	tag := flag.Arg(3)

	// Try to open the mod pack
	cp, err := NewModPack(dir, true, ARG_MMC)
	if err != nil {
		return err
	}

	// If the mod doesn't start with https://, assume it's a name and try to look it up
	if !strings.HasPrefix(mod, "https://") {
		db, err := OpenDatabase()
		if err != nil {
			return err
		}

		// Lookup the URL for the mod
		mod, err = db.findModFile(mod, cp.minecraftVersion())
		if err != nil {
			return err
		}

	} else if !strings.Contains(mod, "minecraft.curseforge.com") && tag == "" {
		return fmt.Errorf("Non-CurseForge URLs must include a tag argument")
	}

	err = cp.selectMod(mod, tag, clientOnly)
	if err != nil {
		return err
	}

	return nil
}

func cmdModList() error {
	name := flag.Arg(1)
	mcvsn := flag.Arg(2)

	db, err := OpenDatabase()
	if err != nil {
		return err
	}

	return db.listMods(name, mcvsn)
}

func cmdForgeList() error {
	mcvsn := flag.Arg(1)

	db, err := OpenDatabase()
	if err != nil {
		return err
	}

	return db.listForge(mcvsn, ARG_VERBOSE)
}

func cmdServerInstall() error {
	dir := flag.Arg(1)

	if ARG_MMC == true {
		return fmt.Errorf("-mmc arg not supported when installing a server")
	}

	// Open the pack; we require the manifest and any
	// config files to already be present
	cp, err := NewModPack(dir, true, false)
	if err != nil {
		return err
	}

	// Install the server jar, Forge and dependencies
	err = cp.installServer()
	if err != nil {
		return err
	}

	// Make sure all mods are installed (do NOT include client-side only)
	err = cp.installMods(false)
	if err != nil {
		return err
	}

	return nil
	// Setup the command-line
	// java -jar <forge.jar>
}

func cmdDBUpdate() error {
	return InstallDatabase()
}

func console(f string, args ...interface{}) {
	fmt.Printf(f, args...)
}

func usage() {
	console("usage: mcdex [<options>] <command> [<args>]\n")
	console(" commands:\n")

	// Sort the list of keys in gCommands
	keys := []string{}
	for k := range gCommands {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, cmd := range keys {
		console(" - %s: %s\n", cmd, gCommands[cmd].Desc)
	}
}

func main() {
	// Register
	flag.BoolVar(&ARG_MMC, "mmc", false, "Generate MultiMC instance.cfg when installing a pack")
	flag.BoolVar(&ARG_VERBOSE, "v", false, "Enable verbose logging of operations")
	flag.BoolVar(&ARG_SKIPMODS, "skipmods", false, "Skip download of mods when installing a pack")

	// Process command-line args
	flag.Parse()
	if !flag.Parsed() || flag.NArg() < 1 {
		usage()
		os.Exit(-1)
	}

	// Initialize our environment
	err := initEnv()
	if err != nil {
		log.Fatalf("Failed to initialize: %s\n", err)
	}

	commandName := flag.Arg(0)
	command, exists := gCommands[commandName]
	if !exists {
		console("ERROR: unknown command '%s'\n", commandName)
		usage()
		os.Exit(-1)
	}

	// Check that the required number of arguments is present
	if flag.NArg() < command.ArgsCount+1 {
		console("ERROR: insufficient arguments for %s\n", commandName)
		console("usage: mcdex %s %s\n", commandName, command.Args)
		os.Exit(-1)
	}

	err = command.Fn()
	if err != nil {
		log.Fatalf("%+v\n", err)
	}
}
<|endoftext|>


<commit_before>package url

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/dimiro1/health"
)

func Test_NewCheckerWithTimeout(t *testing.T) {
	timeout := 2 * time.Second
	url := "http://www.google.com/"

	c := NewCheckerWithTimeout(url, timeout)

	if c.Timeout != timeout {
		t.Errorf("NewCheckerWithTimeout().Timeout == %d, wants %d", c.Timeout, timeout)
	}

	if c.URL != url {
		t.Errorf("NewCheckerWithTimeout().URL == %d, wants %d", c.URL, url)
	}
}

func Test_Checker_Check_Up(t *testing.T) {
	mux := http.NewServeMux()

	server := httptest.NewServer(mux)

	checker := NewChecker(fmt.Sprintf("%s/up/", server.URL))

	handler := health.NewHandler()
	handler.AddChecker("Up", checker)

	mux.Handle("/health/", handler)
	mux.HandleFunc("/up/", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w, "UP")
	})

	resp, _ := http.Get(fmt.Sprintf("%s/health/", server.URL))

	wants := `{"Up":{"code":200,"status":"UP"},"status":"UP"}`

	check(t, resp, wants, http.StatusOK)
}

func Test_Checker_Check_Down(t *testing.T) {
	mux := http.NewServeMux()

	server := httptest.NewServer(mux)

	checker := NewChecker(fmt.Sprintf("%s/down/", server.URL))

	handler := health.NewHandler()
	handler.AddChecker("Down", checker)

	mux.Handle("/health/", handler)
	mux.HandleFunc("/down/", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
		fmt.Fprintln(w, "Down")
	})

	resp, _ := http.Get(fmt.Sprintf("%s/health/", server.URL))

	wants := `{"Down":{"code":500,"status":"DOWN"},"status":"DOWN"}`

	check(t, resp, wants, http.StatusServiceUnavailable)
}

func Test_Checker_Check_Down_invalid(t *testing.T) {
	mux := http.NewServeMux()

	server := httptest.NewServer(mux)

	checker := NewChecker("")

	handler := health.NewHandler()
	handler.AddChecker("Down", checker)

	mux.Handle("/health/", handler)

	resp, _ := http.Get(fmt.Sprintf("%s/health/", server.URL))

	wants := `{"Down":{"code":400,"status":"DOWN"},"status":"DOWN"}`
	check(t, resp, wants, http.StatusServiceUnavailable)
}

func check(t *testing.T, resp *http.Response, wants string, code int) {
	jsonbytes, _ := ioutil.ReadAll(resp.Body)
	jsonstring := strings.TrimSpace(string(jsonbytes))

	if jsonstring != wants {
		t.Errorf("jsonstring == %s, wants %s", jsonstring, wants)
	}

	contentType := resp.Header.Get("Content-Type")
	wants = "application/json"

	if contentType != wants {
		t.Errorf("type == %s, wants %s", contentType, wants)
	}

	if resp.StatusCode != code {
		t.Errorf("resp.StatusCode == %d, wants %d", resp.StatusCode, code)
	}
}
<commit_msg>%s not %d<commit_after>package url

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/dimiro1/health"
)

func Test_NewCheckerWithTimeout(t *testing.T) {
	timeout := 2 * time.Second
	url := "http://www.google.com/"

	c := NewCheckerWithTimeout(url, timeout)

	if c.Timeout != timeout {
		t.Errorf("NewCheckerWithTimeout().Timeout == %d, wants %d", c.Timeout, timeout)
	}

	if c.URL != url {
		t.Errorf("NewCheckerWithTimeout().URL == %s, wants %s", c.URL, url)
	}
}

func Test_Checker_Check_Up(t *testing.T) {
	mux := http.NewServeMux()

	server := httptest.NewServer(mux)

	checker := NewChecker(fmt.Sprintf("%s/up/", server.URL))

	handler := health.NewHandler()
	handler.AddChecker("Up", checker)

	mux.Handle("/health/", handler)
	mux.HandleFunc("/up/", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w, "UP")
	})

	resp, _ := http.Get(fmt.Sprintf("%s/health/", server.URL))

	wants := `{"Up":{"code":200,"status":"UP"},"status":"UP"}`

	check(t, resp, wants, http.StatusOK)
}

func Test_Checker_Check_Down(t *testing.T) {
	mux := http.NewServeMux()

	server := httptest.NewServer(mux)

	checker := NewChecker(fmt.Sprintf("%s/down/", server.URL))

	handler := health.NewHandler()
	handler.AddChecker("Down", checker)

	mux.Handle("/health/", handler)
	mux.HandleFunc("/down/", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
		fmt.Fprintln(w, "Down")
	})

	resp, _ := http.Get(fmt.Sprintf("%s/health/", server.URL))

	wants := `{"Down":{"code":500,"status":"DOWN"},"status":"DOWN"}`

	check(t, resp, wants, http.StatusServiceUnavailable)
}

func Test_Checker_Check_Down_invalid(t *testing.T) {
	mux := http.NewServeMux()

	server := httptest.NewServer(mux)

	checker := NewChecker("")

	handler := health.NewHandler()
	handler.AddChecker("Down", checker)

	mux.Handle("/health/", handler)

	resp, _ := http.Get(fmt.Sprintf("%s/health/", server.URL))

	wants := `{"Down":{"code":400,"status":"DOWN"},"status":"DOWN"}`
	check(t, resp, wants, http.StatusServiceUnavailable)
}

func check(t *testing.T, resp *http.Response, wants string, code int) {
	jsonbytes, _ := ioutil.ReadAll(resp.Body)
	jsonstring := strings.TrimSpace(string(jsonbytes))

	if jsonstring != wants {
		t.Errorf("jsonstring == %s, wants %s", jsonstring, wants)
	}

	contentType := resp.Header.Get("Content-Type")
	wants = "application/json"

	if contentType != wants {
		t.Errorf("type == %s, wants %s", contentType, wants)
	}

	if resp.StatusCode != code {
		t.Errorf("resp.StatusCode == %d, wants %d", resp.StatusCode, code)
	}
}
<|endoftext|>


<commit_before>package deployment

import (
	"encoding/json"
	"time"

	"github.com/flynn/flynn/Godeps/_workspace/src/github.com/flynn/que-go"
	"github.com/flynn/flynn/Godeps/_workspace/src/gopkg.in/inconshreveable/log15.v2"
	"github.com/flynn/flynn/controller/client"
	ct "github.com/flynn/flynn/controller/types"
	"github.com/flynn/flynn/controller/worker/types"
	"github.com/flynn/flynn/pkg/attempt"
	"github.com/flynn/flynn/pkg/postgres"
)

type context struct {
	db     *postgres.DB
	client controller.Client
	logger log15.Logger
}

func JobHandler(db *postgres.DB, client controller.Client, logger log15.Logger) func(*que.Job) error {
	return (&context{db, client, logger}).HandleDeployment
}

func (c *context) HandleDeployment(job *que.Job) (e error) {
	log := c.logger.New("fn", "HandleDeployment")
	log.Info("handling deployment", "job_id", job.ID, "error_count", job.ErrorCount)

	var args ct.DeployID
	if err := json.Unmarshal(job.Args, &args); err != nil {
		log.Error("error unmarshaling job", "err", err)
		return err
	}

	log.Info("getting deployment record", "deployment_id", args.ID)
	deployment, err := c.client.GetDeployment(args.ID)
	if err != nil {
		log.Error("error getting deployment record", "deployment_id", args.ID, "err", err)
		return err
	}

	log = log.New(
		"deployment_id", deployment.ID,
		"app_id", deployment.AppID,
		"strategy", deployment.Strategy,
	)
	// for recovery purposes, fetch old formation
	log.Info("getting old formation")
	f, err := c.client.GetFormation(deployment.AppID, deployment.OldReleaseID)
	if err != nil {
		log.Error("error getting old formation", "release_id", deployment.OldReleaseID, "err", err)
		return err
	}

	events := make(chan ct.DeploymentEvent)
	defer close(events)
	go func() {
		log.Info("watching deployment events")
		for ev := range events {
			log.Info("received deployment event", "status", ev.Status, "type", ev.JobType, "state", ev.JobState)
			ev.AppID = deployment.AppID
			ev.DeploymentID = deployment.ID
			if err := c.createDeploymentEvent(ev); err != nil {
				log.Error("error creating deployment event record", "err", err)
			}
		}
		log.Info("stopped watching deployment events")
	}()
	defer func() {
		if e == worker.ErrStopped {
			return
		}
		log.Info("marking the deployment as done")
		if err := c.setDeploymentDone(deployment.ID); err != nil {
			log.Error("error marking the deployment as done", "err", err)
		}

		// rollback failed deploy
		if e != nil {
			errMsg := e.Error()
			if !IsSkipRollback(e) {
				log.Warn("rolling back deployment due to error", "err", e)
				e = c.rollback(log, deployment, f)
			}
			events <- ct.DeploymentEvent{
				ReleaseID: deployment.NewReleaseID,
				Status:    "failed",
				Error:     errMsg,
			}
		}
	}()

	j := &DeployJob{
		Deployment:      deployment,
		client:          c.client,
		deployEvents:    events,
		serviceNames:    make(map[string]string),
		jobEvents:       make(map[string]chan *JobEvent),
		useJobEvents:    make(map[string]struct{}),
		logger:          c.logger,
		oldReleaseState: make(map[string]int, len(deployment.Processes)),
		newReleaseState: make(map[string]int, len(deployment.Processes)),
		knownJobStates:  make(map[jobIDState]struct{}),
		omni:            make(map[string]struct{}),
		stop:            job.Stop,
	}

	log.Info("performing deployment")
	if err := j.Perform(); err != nil {
		log.Error("error performing deployment", "err", err)
		return err
	}
	log.Info("setting the app release")
	if err := c.client.SetAppRelease(deployment.AppID, deployment.NewReleaseID); err != nil {
		log.Error("error setting the app release", "err", err)
		return err
	}
	// signal success
	events <- ct.DeploymentEvent{
		ReleaseID: deployment.NewReleaseID,
		Status:    "complete",
	}
	log.Info("deployment complete")
	return nil
}

func (c *context) rollback(l log15.Logger, deployment *ct.Deployment, original *ct.Formation) error {
	log := l.New("fn", "rollback")

	log.Info("creating job watcher")
	jobWatcher, err := c.client.WatchJobEvents(deployment.AppID, deployment.OldReleaseID)
	if err != nil {
		log.Error("error opening job event stream", "err", err)
		return err
	}
	appJobs, err := c.client.JobList(deployment.AppID)
	if err != nil {
		log.Error("error listing app jobs", "err", err)
		return err
	}
	runningJobs := make(map[string]int)
	for _, j := range appJobs {
		if j.ReleaseID != deployment.OldReleaseID {
			continue
		}
		if j.State == ct.JobStateUp {
			runningJobs[j.Type]++
		}
	}
	expectedJobEvents := make(ct.JobEvents, len(original.Processes))
	for name, count := range original.Processes {
		count = count - runningJobs[name]
		if count > 0 {
			expectedJobEvents[name] = ct.JobUpEvents(count)
		}
	}

	log.Info("restoring the original formation", "release.id", original.ReleaseID)
	if err := c.client.PutFormation(original); err != nil {
		log.Error("error restoring the original formation", "err", err)
		return err
	}

	if len(expectedJobEvents) > 0 {
		log.Info("waiting for job events", "events", expectedJobEvents)
		callback := func(job *ct.Job) error {
			log.Info("got job event", "job.id", job.ID, "job.type", job.Type, "job.state", job.State)
			return nil
		}
		if err := jobWatcher.WaitFor(expectedJobEvents, 10*time.Second, callback); err != nil {
			log.Error("error waiting for job events", "err", err)
		}
	}

	log.Info("deleting the new formation")
	if err := c.client.DeleteFormation(deployment.AppID, deployment.NewReleaseID); err != nil {
		log.Error("error deleting the new formation:", "err", err)
		return err
	}

	log.Info("rollback complete")
	return nil
}

func (c *context) setDeploymentDone(id string) error {
	return c.execWithRetries("deployment_update_finished_at_now", id)
}

func (c *context) createDeploymentEvent(e ct.DeploymentEvent) error {
	if e.Status == "" {
		e.Status = "running"
	}
	return c.execWithRetries("event_insert", e.AppID, e.DeploymentID, string(ct.EventTypeDeployment), e)
}

var execAttempts = attempt.Strategy{
	Total: 10 * time.Second,
	Delay: 100 * time.Millisecond,
}

// Retry db queries in case postgres has been deployed
func (c *context) execWithRetries(query string, args ...interface{}) error {
	return execAttempts.Run(func() error {
		return c.db.Exec(query, args...)
	})
}
<commit_msg>worker: Don't retry deploys when ErrSkipRollback is returned<commit_after>package deployment

import (
	"encoding/json"
	"time"

	"github.com/flynn/flynn/Godeps/_workspace/src/github.com/flynn/que-go"
	"github.com/flynn/flynn/Godeps/_workspace/src/gopkg.in/inconshreveable/log15.v2"
	"github.com/flynn/flynn/controller/client"
	ct "github.com/flynn/flynn/controller/types"
	"github.com/flynn/flynn/controller/worker/types"
	"github.com/flynn/flynn/pkg/attempt"
	"github.com/flynn/flynn/pkg/postgres"
)

type context struct {
	db     *postgres.DB
	client controller.Client
	logger log15.Logger
}

func JobHandler(db *postgres.DB, client controller.Client, logger log15.Logger) func(*que.Job) error {
	return (&context{db, client, logger}).HandleDeployment
}

func (c *context) HandleDeployment(job *que.Job) (e error) {
	log := c.logger.New("fn", "HandleDeployment")
	log.Info("handling deployment", "job_id", job.ID, "error_count", job.ErrorCount)

	var args ct.DeployID
	if err := json.Unmarshal(job.Args, &args); err != nil {
		log.Error("error unmarshaling job", "err", err)
		return err
	}

	log.Info("getting deployment record", "deployment_id", args.ID)
	deployment, err := c.client.GetDeployment(args.ID)
	if err != nil {
		log.Error("error getting deployment record", "deployment_id", args.ID, "err", err)
		return err
	}

	log = log.New(
		"deployment_id", deployment.ID,
		"app_id", deployment.AppID,
		"strategy", deployment.Strategy,
	)
	// for recovery purposes, fetch old formation
	log.Info("getting old formation")
	f, err := c.client.GetFormation(deployment.AppID, deployment.OldReleaseID)
	if err != nil {
		log.Error("error getting old formation", "release_id", deployment.OldReleaseID, "err", err)
		return err
	}

	events := make(chan ct.DeploymentEvent)
	defer close(events)
	go func() {
		log.Info("watching deployment events")
		for ev := range events {
			log.Info("received deployment event", "status", ev.Status, "type", ev.JobType, "state", ev.JobState)
			ev.AppID = deployment.AppID
			ev.DeploymentID = deployment.ID
			if err := c.createDeploymentEvent(ev); err != nil {
				log.Error("error creating deployment event record", "err", err)
			}
		}
		log.Info("stopped watching deployment events")
	}()
	defer func() {
		if e == worker.ErrStopped {
			return
		}
		log.Info("marking the deployment as done")
		if err := c.setDeploymentDone(deployment.ID); err != nil {
			log.Error("error marking the deployment as done", "err", err)
		}

		// rollback failed deploy
		if e != nil {
			errMsg := e.Error()
			if IsSkipRollback(e) {
				// ErrSkipRollback indicates the deploy failed in some way
				// but no further action should be taken, so set the error
				// to nil to avoid retrying the deploy
				e = nil
			} else {
				log.Warn("rolling back deployment due to error", "err", e)
				e = c.rollback(log, deployment, f)
			}
			events <- ct.DeploymentEvent{
				ReleaseID: deployment.NewReleaseID,
				Status:    "failed",
				Error:     errMsg,
			}
		}
	}()

	j := &DeployJob{
		Deployment:      deployment,
		client:          c.client,
		deployEvents:    events,
		serviceNames:    make(map[string]string),
		jobEvents:       make(map[string]chan *JobEvent),
		useJobEvents:    make(map[string]struct{}),
		logger:          c.logger,
		oldReleaseState: make(map[string]int, len(deployment.Processes)),
		newReleaseState: make(map[string]int, len(deployment.Processes)),
		knownJobStates:  make(map[jobIDState]struct{}),
		omni:            make(map[string]struct{}),
		stop:            job.Stop,
	}

	log.Info("performing deployment")
	if err := j.Perform(); err != nil {
		log.Error("error performing deployment", "err", err)
		return err
	}
	log.Info("setting the app release")
	if err := c.client.SetAppRelease(deployment.AppID, deployment.NewReleaseID); err != nil {
		log.Error("error setting the app release", "err", err)
		return err
	}
	// signal success
	events <- ct.DeploymentEvent{
		ReleaseID: deployment.NewReleaseID,
		Status:    "complete",
	}
	log.Info("deployment complete")
	return nil
}

func (c *context) rollback(l log15.Logger, deployment *ct.Deployment, original *ct.Formation) error {
	log := l.New("fn", "rollback")

	log.Info("creating job watcher")
	jobWatcher, err := c.client.WatchJobEvents(deployment.AppID, deployment.OldReleaseID)
	if err != nil {
		log.Error("error opening job event stream", "err", err)
		return err
	}
	appJobs, err := c.client.JobList(deployment.AppID)
	if err != nil {
		log.Error("error listing app jobs", "err", err)
		return err
	}
	runningJobs := make(map[string]int)
	for _, j := range appJobs {
		if j.ReleaseID != deployment.OldReleaseID {
			continue
		}
		if j.State == ct.JobStateUp {
			runningJobs[j.Type]++
		}
	}
	expectedJobEvents := make(ct.JobEvents, len(original.Processes))
	for name, count := range original.Processes {
		count = count - runningJobs[name]
		if count > 0 {
			expectedJobEvents[name] = ct.JobUpEvents(count)
		}
	}

	log.Info("restoring the original formation", "release.id", original.ReleaseID)
	if err := c.client.PutFormation(original); err != nil {
		log.Error("error restoring the original formation", "err", err)
		return err
	}

	if len(expectedJobEvents) > 0 {
		log.Info("waiting for job events", "events", expectedJobEvents)
		callback := func(job *ct.Job) error {
			log.Info("got job event", "job.id", job.ID, "job.type", job.Type, "job.state", job.State)
			return nil
		}
		if err := jobWatcher.WaitFor(expectedJobEvents, 10*time.Second, callback); err != nil {
			log.Error("error waiting for job events", "err", err)
		}
	}

	log.Info("deleting the new formation")
	if err := c.client.DeleteFormation(deployment.AppID, deployment.NewReleaseID); err != nil {
		log.Error("error deleting the new formation:", "err", err)
		return err
	}

	log.Info("rollback complete")
	return nil
}

func (c *context) setDeploymentDone(id string) error {
	return c.execWithRetries("deployment_update_finished_at_now", id)
}

func (c *context) createDeploymentEvent(e ct.DeploymentEvent) error {
	if e.Status == "" {
		e.Status = "running"
	}
	return c.execWithRetries("event_insert", e.AppID, e.DeploymentID, string(ct.EventTypeDeployment), e)
}

var execAttempts = attempt.Strategy{
	Total: 10 * time.Second,
	Delay: 100 * time.Millisecond,
}

// Retry db queries in case postgres has been deployed
func (c *context) execWithRetries(query string, args ...interface{}) error {
	return execAttempts.Run(func() error {
		return c.db.Exec(query, args...)
	})
}
<|endoftext|>


<commit_before>package structs

import (
	"bytes"
	"fmt"
	"reflect"
	"strconv"
)

type Tag struct {
	name   string
	values map[string]string
}

func (t *Tag) Name() string {
	return t.name
}

func (t *Tag) Has(opt string) bool {
	_, ok := t.values[opt]
	return ok
}

func (t *Tag) Value(key string) string {
	return t.values[key]
}

func (t *Tag) IntValue(key string) (int, bool) {
	v := t.Value(key)
	if v != "" {
		val, err := strconv.Atoi(key)
		return val, err == nil
	}
	return 0, false
}

// Commonly used tag fields

func (t *Tag) CodecName() string {
	return t.Value("codec")
}

func (t *Tag) PipeName() string {
	return t.Value("pipe")
}

func (t *Tag) Optional() bool {
	return t.Has("optional")
}

func (t *Tag) Required() bool {
	return t.Has("required")
}

func (t *Tag) Alphanumeric() bool {
	return t.Has("alphanumeric")
}

func (t *Tag) MaxLength() (int, bool) {
	return t.IntValue("max_length")
}

func (t *Tag) MinLength() (int, bool) {
	return t.IntValue("min_length")
}

func (t *Tag) IsEmpty() bool {
	return t.name == "" && len(t.values) == 0
}

func splitFields(tag string) (string, map[string]string, error) {
	const (
		stateKey = iota
		stateValue
		stateValueQuoted
		stateEscape
	)
	hasName := false
	var name string
	var key string
	var prevState int
	state := stateKey
	var buf bytes.Buffer
	values := make(map[string]string)
	for ii, v := range []byte(tag) {
		switch v {
		case ',':
			if state != stateValueQuoted {
				if state == stateEscape {
					return "", nil, fmt.Errorf("unknown escape sequence \\%s at %d", string(v), ii)
				}
				if hasName {
					if key != "" {
						values[key] = buf.String()
						key = ""
					} else {
						values[buf.String()] = ""
					}
				} else {
					name = buf.String()
					hasName = true
				}
				buf.Reset()
				state = stateKey
			} else {
				buf.WriteByte(v)
			}
		case '\'':
			if state == stateValue {
				if buf.Len() == 0 {
					state = stateValueQuoted
				} else {
					buf.WriteByte(v)
				}
			} else if state == stateValueQuoted {
				values[key] = buf.String()
				key = ""
				buf.Reset()
				state = stateKey
			} else if state == stateEscape {
				buf.WriteByte(v)
				state = prevState
			} else {
				return "", nil, fmt.Errorf("illegal character ' in key at %d", ii)
			}
		case '\\':
			if state == stateEscape {
				buf.WriteByte(v)
			} else {
				prevState = state
				state = stateEscape
			}
		case '=':
			if state == stateKey {
				key = buf.String()
				buf.Reset()
				state = stateValue
			} else {
				buf.WriteByte(v)
			}
		default:
			if state == stateEscape {
				return "", nil, fmt.Errorf("unknown escape sequence \\%s at %d", string(v), ii)
			}
			buf.WriteByte(v)
		}
	}
	switch state {
	case stateKey:
		if k := buf.String(); k != "" {
			if hasName {
				values[k] = ""
			} else {
				name = k
			}
		}
	case stateValue:
		values[key] = buf.String()
	default:
		return "", nil, fmt.Errorf("unexpected end at %d", len(tag))
	}
	return name, values, nil
}

// ParseTag parses a Gondola style struct tag field from
// the given tag string.
func ParseTag(tag string) (*Tag, error) {
	name, values, err := splitFields(tag)
	if err != nil {
		return nil, err
	}
	return &Tag{name, values}, nil
}

// MustParseTag works like ParseTag, but panics if there's an error.
func MustParseTag(tag string) *Tag {
	t, err := ParseTag(tag)
	if err != nil {
		panic(err)
	}
	return t
}

func NewTag(field reflect.StructField, alternatives []string) *Tag {
	for _, v := range alternatives {
		t := field.Tag.Get(v)
		if t != "" {
			return MustParseTag(t)
		}
	}
	return MustParseTag("")
}

func NewTagNamed(field reflect.StructField, name string) *Tag {
	return MustParseTag(field.Tag.Get(name))
}

func NewStringTagNamed(tag string, name string) *Tag {
	t := reflect.StructTag(tag)
	return MustParseTag(t.Get(name))
}
<commit_msg>Fix IntValue()<commit_after>package structs

import (
	"bytes"
	"fmt"
	"reflect"
	"strconv"
)

type Tag struct {
	name   string
	values map[string]string
}

func (t *Tag) Name() string {
	return t.name
}

func (t *Tag) Has(opt string) bool {
	_, ok := t.values[opt]
	return ok
}

func (t *Tag) Value(key string) string {
	return t.values[key]
}

func (t *Tag) IntValue(key string) (int, bool) {
	v := t.Value(key)
	if v != "" {
		val, err := strconv.Atoi(v)
		return val, err == nil
	}
	return 0, false
}

// Commonly used tag fields

func (t *Tag) CodecName() string {
	return t.Value("codec")
}

func (t *Tag) PipeName() string {
	return t.Value("pipe")
}

func (t *Tag) Optional() bool {
	return t.Has("optional")
}

func (t *Tag) Required() bool {
	return t.Has("required")
}

func (t *Tag) Alphanumeric() bool {
	return t.Has("alphanumeric")
}

func (t *Tag) MaxLength() (int, bool) {
	return t.IntValue("max_length")
}

func (t *Tag) MinLength() (int, bool) {
	return t.IntValue("min_length")
}

func (t *Tag) IsEmpty() bool {
	return t.name == "" && len(t.values) == 0
}

func splitFields(tag string) (string, map[string]string, error) {
	const (
		stateKey = iota
		stateValue
		stateValueQuoted
		stateEscape
	)
	hasName := false
	var name string
	var key string
	var prevState int
	state := stateKey
	var buf bytes.Buffer
	values := make(map[string]string)
	for ii, v := range []byte(tag) {
		switch v {
		case ',':
			if state != stateValueQuoted {
				if state == stateEscape {
					return "", nil, fmt.Errorf("unknown escape sequence \\%s at %d", string(v), ii)
				}
				if hasName {
					if key != "" {
						values[key] = buf.String()
						key = ""
					} else {
						values[buf.String()] = ""
					}
				} else {
					name = buf.String()
					hasName = true
				}
				buf.Reset()
				state = stateKey
			} else {
				buf.WriteByte(v)
			}
		case '\'':
			if state == stateValue {
				if buf.Len() == 0 {
					state = stateValueQuoted
				} else {
					buf.WriteByte(v)
				}
			} else if state == stateValueQuoted {
				values[key] = buf.String()
				key = ""
				buf.Reset()
				state = stateKey
			} else if state == stateEscape {
				buf.WriteByte(v)
				state = prevState
			} else {
				return "", nil, fmt.Errorf("illegal character ' in key at %d", ii)
			}
		case '\\':
			if state == stateEscape {
				buf.WriteByte(v)
			} else {
				prevState = state
				state = stateEscape
			}
		case '=':
			if state == stateKey {
				key = buf.String()
				buf.Reset()
				state = stateValue
			} else {
				buf.WriteByte(v)
			}
		default:
			if state == stateEscape {
				return "", nil, fmt.Errorf("unknown escape sequence \\%s at %d", string(v), ii)
			}
			buf.WriteByte(v)
		}
	}
	switch state {
	case stateKey:
		if k := buf.String(); k != "" {
			if hasName {
				values[k] = ""
			} else {
				name = k
			}
		}
	case stateValue:
		values[key] = buf.String()
	default:
		return "", nil, fmt.Errorf("unexpected end at %d", len(tag))
	}
	return name, values, nil
}

// ParseTag parses a Gondola style struct tag field from
// the given tag string.
func ParseTag(tag string) (*Tag, error) {
	name, values, err := splitFields(tag)
	if err != nil {
		return nil, err
	}
	return &Tag{name, values}, nil
}

// MustParseTag works like ParseTag, but panics if there's an error.
func MustParseTag(tag string) *Tag {
	t, err := ParseTag(tag)
	if err != nil {
		panic(err)
	}
	return t
}

func NewTag(field reflect.StructField, alternatives []string) *Tag {
	for _, v := range alternatives {
		t := field.Tag.Get(v)
		if t != "" {
			return MustParseTag(t)
		}
	}
	return MustParseTag("")
}

func NewTagNamed(field reflect.StructField, name string) *Tag {
	return MustParseTag(field.Tag.Get(name))
}

func NewStringTagNamed(tag string, name string) *Tag {
	t := reflect.StructTag(tag)
	return MustParseTag(t.Get(name))
}
<|endoftext|>


<commit_before>/* Licensed to the Apache Software Foundation (ASF) under one or more
contributor license agreements.  See the NOTICE file distributed with
this work for additional information regarding copyright ownership.
The ASF licenses this file to You 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 framework

import (
	"fmt"
	"os"
	"os/signal"
	"strings"

	"time"

	utils "github.com/elodina/go-mesos-utils"
	"github.com/elodina/go-mesos-utils/pretty"
	"github.com/golang/protobuf/proto"
	mesos "github.com/mesos/mesos-go/mesosproto"
	util "github.com/mesos/mesos-go/mesosutil"
	"github.com/mesos/mesos-go/scheduler"
)

const (
	reconcileDelay    = 10 * time.Second
	reconcileMaxTries = 3
)

var sched *Scheduler // This is needed for HTTP server to be able to update this scheduler

type Scheduler struct {
	httpServer      *HttpServer
	cluster         *Cluster
	driver          scheduler.SchedulerDriver
	schedulerDriver *scheduler.MesosSchedulerDriver

	reconcileTime time.Time
	reconciles    int
}

func NewScheduler() *Scheduler {
	return &Scheduler{
		reconcileTime: time.Unix(0, 0),
	}
}

func (s *Scheduler) Start() error {
	Logger.Infof("Starting scheduler with configuration: \n%s", Config)
	sched = s // set this scheduler reachable for http server

	ctrlc := make(chan os.Signal, 1)
	signal.Notify(ctrlc, os.Interrupt)

	s.cluster = NewCluster()
	s.cluster.Load()

	s.httpServer = NewHttpServer(Config.Api)
	go s.httpServer.Start()

	frameworkInfo := &mesos.FrameworkInfo{
		User:            proto.String(Config.User),
		Name:            proto.String(Config.FrameworkName),
		Role:            proto.String(Config.FrameworkRole),
		FailoverTimeout: proto.Float64(float64(Config.FrameworkTimeout / 1e9)),
		Checkpoint:      proto.Bool(true),
	}

	if s.cluster.frameworkID != "" {
		frameworkInfo.Id = util.NewFrameworkID(s.cluster.frameworkID)
	}

	driverConfig := scheduler.DriverConfig{
		Scheduler: s,
		Framework: frameworkInfo,
		Master:    Config.Master,
	}

	driver, err := scheduler.NewMesosSchedulerDriver(driverConfig)
	s.schedulerDriver = driver

	if err != nil {
		return fmt.Errorf("Unable to create SchedulerDriver: %s", err)
	}

	go func() {
		if stat, err := driver.Run(); err != nil {
			Logger.Infof("Framework stopped with status %s and error: %s\n", stat.String(), err)
			panic(err)
		}
	}()

	<-ctrlc
	return nil
}

func (s *Scheduler) Registered(driver scheduler.SchedulerDriver, id *mesos.FrameworkID, master *mesos.MasterInfo) {
	Logger.Infof("[Registered] framework: %s master: %s:%d", id.GetValue(), master.GetHostname(), master.GetPort())

	s.cluster.frameworkID = id.GetValue()
	s.cluster.Save()

	s.driver = driver
	s.reconcileTasks(true)
}

func (s *Scheduler) Reregistered(driver scheduler.SchedulerDriver, master *mesos.MasterInfo) {
	Logger.Infof("[Reregistered] master: %s:%d", master.GetHostname(), master.GetPort())

	s.driver = driver
	s.reconcileTasks(true)
}

func (s *Scheduler) Disconnected(scheduler.SchedulerDriver) {
	Logger.Info("[Disconnected]")

	s.driver = nil
}

func (s *Scheduler) ResourceOffers(driver scheduler.SchedulerDriver, offers []*mesos.Offer) {
	Logger.Debugf("[ResourceOffers] %s", pretty.Offers(offers))

	for _, offer := range offers {
		declineReason := s.acceptOffer(driver, offer)
		if declineReason != "" {
			driver.DeclineOffer(offer.GetId(), &mesos.Filters{RefuseSeconds: proto.Float64(1)})
			Logger.Debugf("Declined offer: %s", declineReason)
		}
	}

	s.reconcileTasks(false)
	s.cluster.Save()
}

func (s *Scheduler) OfferRescinded(driver scheduler.SchedulerDriver, id *mesos.OfferID) {
	Logger.Infof("[OfferRescinded] %s", id.GetValue())
}

func (s *Scheduler) StatusUpdate(driver scheduler.SchedulerDriver, status *mesos.TaskStatus) {
	Logger.Infof("[StatusUpdate] %s", pretty.Status(status))

	id := s.idFromTaskId(status.GetTaskId().GetValue())

	switch status.GetState() {
	case mesos.TaskState_TASK_RUNNING:
		s.onTaskStarted(id, status)
	case mesos.TaskState_TASK_LOST, mesos.TaskState_TASK_FAILED, mesos.TaskState_TASK_ERROR:
		s.onTaskFailed(id, status)
	case mesos.TaskState_TASK_FINISHED, mesos.TaskState_TASK_KILLED:
		s.onTaskFinished(id, status)
	default:
		Logger.Warnf("Got unexpected task state %s for task %s", pretty.Status(status), id)
	}

	s.cluster.Save()
}

func (s *Scheduler) FrameworkMessage(driver scheduler.SchedulerDriver, executor *mesos.ExecutorID, slave *mesos.SlaveID, message string) {
	Logger.Infof("[FrameworkMessage] executor: %s slave: %s message: %s", executor, slave, message)
}

func (s *Scheduler) SlaveLost(driver scheduler.SchedulerDriver, slave *mesos.SlaveID) {
	Logger.Infof("[SlaveLost] %s", slave.GetValue())
}

func (s *Scheduler) ExecutorLost(driver scheduler.SchedulerDriver, executor *mesos.ExecutorID, slave *mesos.SlaveID, status int) {
	Logger.Infof("[ExecutorLost] executor: %s slave: %s status: %d", executor, slave, status)
}

func (s *Scheduler) Error(driver scheduler.SchedulerDriver, message string) {
	Logger.Errorf("[Error] %s", message)

	if s.schedulerDriver.Status() == mesos.Status_DRIVER_ABORTED {
		Logger.Errorf("Driver aborted, exiting...")
		time.Sleep(1 * time.Second) // sometimes logs do not flush so give them some time
		os.Exit(1)
	}
}

func (s *Scheduler) Shutdown(driver *scheduler.MesosSchedulerDriver) {
	Logger.Info("Shutdown triggered, stopping driver")
	driver.Stop(false)
}

func (s *Scheduler) acceptOffer(driver scheduler.SchedulerDriver, offer *mesos.Offer) string {
	declineReasons := make([]string, 0)

	tasks := s.cluster.GetTasksWithState(TaskStateStopped)
	if len(tasks) == 0 {
		return "all tasks are running"
	}

	for _, task := range tasks {
		declineReason := utils.CheckConstraints(offer, task.Data().constraints, s.cluster.GetConstrained())
		if declineReason == "" {
			declineReason = task.Matches(offer)
			if declineReason == "" {
				s.launchTask(task, offer)
				return ""
			} else {
				declineReasons = append(declineReasons, declineReason)
			}
		} else {
			declineReasons = append(declineReasons, declineReason)
		}
	}

	return strings.Join(declineReasons, ", ")
}

func (s *Scheduler) launchTask(task Task, offer *mesos.Offer) {
	taskInfo := task.NewTaskInfo(offer)
	task.Data().State = TaskStateStaging
	task.Data().Attributes = utils.OfferAttributes(offer)
	task.Data().ExecutorID = taskInfo.GetExecutor().GetExecutorId().GetValue()
	task.Data().SlaveID = taskInfo.GetSlaveId().GetValue()
	task.Data().TaskID = taskInfo.GetTaskId().GetValue()

	s.driver.LaunchTasks([]*mesos.OfferID{offer.GetId()}, []*mesos.TaskInfo{taskInfo}, &mesos.Filters{RefuseSeconds: proto.Float64(1)})
}

func (s *Scheduler) onTaskStarted(id string, status *mesos.TaskStatus) {
	if s.cluster.Exists(id) {
		task := s.cluster.Get(id)
		task.Data().State = TaskStateRunning
	} else {
		Logger.Infof("Got %s for unknown/stopped task, killing task %s", pretty.Status(status), status.GetTaskId().GetValue())
	}
}

func (s *Scheduler) onTaskFailed(id string, status *mesos.TaskStatus) {
	if s.cluster.Exists(id) {
		task := s.cluster.Get(id)
		if task.Data().State != TaskStateInactive {
			task.Data().State = TaskStateStopped
		}
	} else {
		Logger.Infof("Got %s for unknown/stopped task %s", pretty.Status(status), status.GetTaskId().GetValue())
	}
}

func (s *Scheduler) onTaskFinished(id string, status *mesos.TaskStatus) {
	if !s.cluster.Exists(id) {
		Logger.Infof("Got %s for unknown/stopped task %s", pretty.Status(status), status.GetTaskId().GetValue())
	}
}

func (s *Scheduler) stopTask(task Task) {
	if task.Data().State == TaskStateRunning || task.Data().State == TaskStateStaging {
		Logger.Infof("Stopping task %s", task.Data().TaskID)
		s.driver.KillTask(util.NewTaskID(task.Data().TaskID))
	}

	task.Data().State = TaskStateInactive
	task.Data().ResetTaskInfo()
}

func (s *Scheduler) idFromTaskId(taskId string) string {
	tokens := strings.Split(taskId, "-")
	id := tokens[1]
	Logger.Debugf("ID extracted from %s is %s", taskId, id)
	return id
}

func (s *Scheduler) reconcileTasks(force bool) {
	if time.Now().Sub(s.reconcileTime) >= reconcileDelay {
		if !s.cluster.IsReconciling() {
			s.reconciles = 0
		}
		s.reconciles++
		s.reconcileTime = time.Now()

		if s.reconciles > reconcileMaxTries {
			for _, task := range s.cluster.GetTasksWithState(TaskStateReconciling) {
				if task.Data().TaskID != "" {
					Logger.Infof("Reconciling exceeded %d tries for task %s, sending killTask for task %s", reconcileMaxTries, task.Data().ID, task.Data().TaskID)
					s.driver.KillTask(util.NewTaskID(task.Data().TaskID))

					task.Data().ResetTaskInfo()
				}
			}
		} else {
			if force {
				s.driver.ReconcileTasks(nil)
			} else {
				statuses := make([]*mesos.TaskStatus, 0)
				for _, task := range s.cluster.GetAllTasks() {
					if task.Data().TaskID != "" {
						task.Data().State = TaskStateReconciling
						Logger.Infof("Reconciling %d/%d task state for id %s, task id %s", s.reconciles, reconcileMaxTries, task.Data().ID, task.Data().TaskID)
						statuses = append(statuses, util.NewTaskStatus(util.NewTaskID(task.Data().TaskID), mesos.TaskState_TASK_STAGING))
					}
				}
				s.driver.ReconcileTasks(statuses)
			}
		}
	}
}

func getScalarResources(offer *mesos.Offer, resourceName string) float64 {
	resources := 0.0
	filteredResources := util.FilterResources(offer.Resources, func(res *mesos.Resource) bool {
		return res.GetName() == resourceName
	})
	for _, res := range filteredResources {
		resources += res.GetScalar().GetValue()
	}
	return resources
}
<commit_msg>Increased offer refuse seconds<commit_after>/* Licensed to the Apache Software Foundation (ASF) under one or more
contributor license agreements.  See the NOTICE file distributed with
this work for additional information regarding copyright ownership.
The ASF licenses this file to You 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 framework

import (
	"fmt"
	"os"
	"os/signal"
	"strings"

	"time"

	utils "github.com/elodina/go-mesos-utils"
	"github.com/elodina/go-mesos-utils/pretty"
	"github.com/golang/protobuf/proto"
	mesos "github.com/mesos/mesos-go/mesosproto"
	util "github.com/mesos/mesos-go/mesosutil"
	"github.com/mesos/mesos-go/scheduler"
)

const (
	reconcileDelay    = 10 * time.Second
	reconcileMaxTries = 3
)

var sched *Scheduler // This is needed for HTTP server to be able to update this scheduler

type Scheduler struct {
	httpServer      *HttpServer
	cluster         *Cluster
	driver          scheduler.SchedulerDriver
	schedulerDriver *scheduler.MesosSchedulerDriver

	reconcileTime time.Time
	reconciles    int
}

func NewScheduler() *Scheduler {
	return &Scheduler{
		reconcileTime: time.Unix(0, 0),
	}
}

func (s *Scheduler) Start() error {
	Logger.Infof("Starting scheduler with configuration: \n%s", Config)
	sched = s // set this scheduler reachable for http server

	ctrlc := make(chan os.Signal, 1)
	signal.Notify(ctrlc, os.Interrupt)

	s.cluster = NewCluster()
	s.cluster.Load()

	s.httpServer = NewHttpServer(Config.Api)
	go s.httpServer.Start()

	frameworkInfo := &mesos.FrameworkInfo{
		User:            proto.String(Config.User),
		Name:            proto.String(Config.FrameworkName),
		Role:            proto.String(Config.FrameworkRole),
		FailoverTimeout: proto.Float64(float64(Config.FrameworkTimeout / 1e9)),
		Checkpoint:      proto.Bool(true),
	}

	if s.cluster.frameworkID != "" {
		frameworkInfo.Id = util.NewFrameworkID(s.cluster.frameworkID)
	}

	driverConfig := scheduler.DriverConfig{
		Scheduler: s,
		Framework: frameworkInfo,
		Master:    Config.Master,
	}

	driver, err := scheduler.NewMesosSchedulerDriver(driverConfig)
	s.schedulerDriver = driver

	if err != nil {
		return fmt.Errorf("Unable to create SchedulerDriver: %s", err)
	}

	go func() {
		if stat, err := driver.Run(); err != nil {
			Logger.Infof("Framework stopped with status %s and error: %s\n", stat.String(), err)
			panic(err)
		}
	}()

	<-ctrlc
	return nil
}

func (s *Scheduler) Registered(driver scheduler.SchedulerDriver, id *mesos.FrameworkID, master *mesos.MasterInfo) {
	Logger.Infof("[Registered] framework: %s master: %s:%d", id.GetValue(), master.GetHostname(), master.GetPort())

	s.cluster.frameworkID = id.GetValue()
	s.cluster.Save()

	s.driver = driver
	s.reconcileTasks(true)
}

func (s *Scheduler) Reregistered(driver scheduler.SchedulerDriver, master *mesos.MasterInfo) {
	Logger.Infof("[Reregistered] master: %s:%d", master.GetHostname(), master.GetPort())

	s.driver = driver
	s.reconcileTasks(true)
}

func (s *Scheduler) Disconnected(scheduler.SchedulerDriver) {
	Logger.Info("[Disconnected]")

	s.driver = nil
}

func (s *Scheduler) ResourceOffers(driver scheduler.SchedulerDriver, offers []*mesos.Offer) {
	Logger.Debugf("[ResourceOffers] %s", pretty.Offers(offers))

	for _, offer := range offers {
		declineReason := s.acceptOffer(driver, offer)
		if declineReason != "" {
			driver.DeclineOffer(offer.GetId(), &mesos.Filters{RefuseSeconds: proto.Float64(10)})
			Logger.Debugf("Declined offer: %s", declineReason)
		}
	}

	s.reconcileTasks(false)
	s.cluster.Save()
}

func (s *Scheduler) OfferRescinded(driver scheduler.SchedulerDriver, id *mesos.OfferID) {
	Logger.Infof("[OfferRescinded] %s", id.GetValue())
}

func (s *Scheduler) StatusUpdate(driver scheduler.SchedulerDriver, status *mesos.TaskStatus) {
	Logger.Infof("[StatusUpdate] %s", pretty.Status(status))

	id := s.idFromTaskId(status.GetTaskId().GetValue())

	switch status.GetState() {
	case mesos.TaskState_TASK_RUNNING:
		s.onTaskStarted(id, status)
	case mesos.TaskState_TASK_LOST, mesos.TaskState_TASK_FAILED, mesos.TaskState_TASK_ERROR:
		s.onTaskFailed(id, status)
	case mesos.TaskState_TASK_FINISHED, mesos.TaskState_TASK_KILLED:
		s.onTaskFinished(id, status)
	default:
		Logger.Warnf("Got unexpected task state %s for task %s", pretty.Status(status), id)
	}

	s.cluster.Save()
}

func (s *Scheduler) FrameworkMessage(driver scheduler.SchedulerDriver, executor *mesos.ExecutorID, slave *mesos.SlaveID, message string) {
	Logger.Infof("[FrameworkMessage] executor: %s slave: %s message: %s", executor, slave, message)
}

func (s *Scheduler) SlaveLost(driver scheduler.SchedulerDriver, slave *mesos.SlaveID) {
	Logger.Infof("[SlaveLost] %s", slave.GetValue())
}

func (s *Scheduler) ExecutorLost(driver scheduler.SchedulerDriver, executor *mesos.ExecutorID, slave *mesos.SlaveID, status int) {
	Logger.Infof("[ExecutorLost] executor: %s slave: %s status: %d", executor, slave, status)
}

func (s *Scheduler) Error(driver scheduler.SchedulerDriver, message string) {
	Logger.Errorf("[Error] %s", message)

	if s.schedulerDriver.Status() == mesos.Status_DRIVER_ABORTED {
		Logger.Errorf("Driver aborted, exiting...")
		time.Sleep(1 * time.Second) // sometimes logs do not flush so give them some time
		os.Exit(1)
	}
}

func (s *Scheduler) Shutdown(driver *scheduler.MesosSchedulerDriver) {
	Logger.Info("Shutdown triggered, stopping driver")
	driver.Stop(false)
}

func (s *Scheduler) acceptOffer(driver scheduler.SchedulerDriver, offer *mesos.Offer) string {
	declineReasons := make([]string, 0)

	tasks := s.cluster.GetTasksWithState(TaskStateStopped)
	if len(tasks) == 0 {
		return "all tasks are running"
	}

	for _, task := range tasks {
		declineReason := utils.CheckConstraints(offer, task.Data().constraints, s.cluster.GetConstrained())
		if declineReason == "" {
			declineReason = task.Matches(offer)
			if declineReason == "" {
				s.launchTask(task, offer)
				return ""
			} else {
				declineReasons = append(declineReasons, declineReason)
			}
		} else {
			declineReasons = append(declineReasons, declineReason)
		}
	}

	return strings.Join(declineReasons, ", ")
}

func (s *Scheduler) launchTask(task Task, offer *mesos.Offer) {
	taskInfo := task.NewTaskInfo(offer)
	task.Data().State = TaskStateStaging
	task.Data().Attributes = utils.OfferAttributes(offer)
	task.Data().ExecutorID = taskInfo.GetExecutor().GetExecutorId().GetValue()
	task.Data().SlaveID = taskInfo.GetSlaveId().GetValue()
	task.Data().TaskID = taskInfo.GetTaskId().GetValue()

	s.driver.LaunchTasks([]*mesos.OfferID{offer.GetId()}, []*mesos.TaskInfo{taskInfo}, &mesos.Filters{RefuseSeconds: proto.Float64(10)})
}

func (s *Scheduler) onTaskStarted(id string, status *mesos.TaskStatus) {
	if s.cluster.Exists(id) {
		task := s.cluster.Get(id)
		task.Data().State = TaskStateRunning
	} else {
		Logger.Infof("Got %s for unknown/stopped task, killing task %s", pretty.Status(status), status.GetTaskId().GetValue())
	}
}

func (s *Scheduler) onTaskFailed(id string, status *mesos.TaskStatus) {
	if s.cluster.Exists(id) {
		task := s.cluster.Get(id)
		if task.Data().State != TaskStateInactive {
			task.Data().State = TaskStateStopped
		}
	} else {
		Logger.Infof("Got %s for unknown/stopped task %s", pretty.Status(status), status.GetTaskId().GetValue())
	}
}

func (s *Scheduler) onTaskFinished(id string, status *mesos.TaskStatus) {
	if !s.cluster.Exists(id) {
		Logger.Infof("Got %s for unknown/stopped task %s", pretty.Status(status), status.GetTaskId().GetValue())
	}
}

func (s *Scheduler) stopTask(task Task) {
	if task.Data().State == TaskStateRunning || task.Data().State == TaskStateStaging {
		Logger.Infof("Stopping task %s", task.Data().TaskID)
		s.driver.KillTask(util.NewTaskID(task.Data().TaskID))
	}

	task.Data().State = TaskStateInactive
	task.Data().ResetTaskInfo()
}

func (s *Scheduler) idFromTaskId(taskId string) string {
	tokens := strings.Split(taskId, "-")
	id := tokens[1]
	Logger.Debugf("ID extracted from %s is %s", taskId, id)
	return id
}

func (s *Scheduler) reconcileTasks(force bool) {
	if time.Now().Sub(s.reconcileTime) >= reconcileDelay {
		if !s.cluster.IsReconciling() {
			s.reconciles = 0
		}
		s.reconciles++
		s.reconcileTime = time.Now()

		if s.reconciles > reconcileMaxTries {
			for _, task := range s.cluster.GetTasksWithState(TaskStateReconciling) {
				if task.Data().TaskID != "" {
					Logger.Infof("Reconciling exceeded %d tries for task %s, sending killTask for task %s", reconcileMaxTries, task.Data().ID, task.Data().TaskID)
					s.driver.KillTask(util.NewTaskID(task.Data().TaskID))

					task.Data().ResetTaskInfo()
				}
			}
		} else {
			if force {
				s.driver.ReconcileTasks(nil)
			} else {
				statuses := make([]*mesos.TaskStatus, 0)
				for _, task := range s.cluster.GetAllTasks() {
					if task.Data().TaskID != "" {
						task.Data().State = TaskStateReconciling
						Logger.Infof("Reconciling %d/%d task state for id %s, task id %s", s.reconciles, reconcileMaxTries, task.Data().ID, task.Data().TaskID)
						statuses = append(statuses, util.NewTaskStatus(util.NewTaskID(task.Data().TaskID), mesos.TaskState_TASK_STAGING))
					}
				}
				s.driver.ReconcileTasks(statuses)
			}
		}
	}
}

func getScalarResources(offer *mesos.Offer, resourceName string) float64 {
	resources := 0.0
	filteredResources := util.FilterResources(offer.Resources, func(res *mesos.Resource) bool {
		return res.GetName() == resourceName
	})
	for _, res := range filteredResources {
		resources += res.GetScalar().GetValue()
	}
	return resources
}
<|endoftext|>