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

package ec2

import (
	"fmt"

	"github.com/juju/errors"
	"github.com/juju/juju/network"
	"github.com/juju/utils/set"
	"gopkg.in/amz.v3/ec2"
)

const (
	activeState           = "active"
	availableState        = "available"
	attachedState         = "attached"
	localRouteGatewayID   = "local"
	defaultRouteCIDRBlock = "0.0.0.0/0"
	defaultVPCIDNone      = "none"
)

// vpcAPIClient defines a subset of the goamz API calls needed to validate a VPC.
type vpcAPIClient interface {
	AccountAttributes(attributeNames ...string) (*ec2.AccountAttributesResp, error)
	VPCs(ids []string, filter *ec2.Filter) (*ec2.VPCsResp, error)
	Subnets(ids []string, filter *ec2.Filter) (*ec2.SubnetsResp, error)
	InternetGateways(ids []string, filter *ec2.Filter) (*ec2.InternetGatewaysResp, error)
	RouteTables(ids []string, filter *ec2.Filter) (*ec2.RouteTablesResp, error)
}

// validateVPC requires both arguments to be set and validates that vpcID refers
// to an existing AWS VPC (default or non-default) for the current region.
// Returns an error satifying errors.IsNotFound() when the VPC with the given
// vpcID cannot be found, or when the VPC exists but contains no subnets.
// Returns an error satisfying errors.IsNotValid() in the following cases:
//
// 1. The VPC's state is not "available".
// 2. The VPC does not have an Internet Gateway (IGW) attached.
// 3. A main route table is not associated with the VPC.
// 4. The main route table lacks both a default route via the IGW and a local
//    route matching the VPC's CIDR block.
// 5. One or more of the VPC's subnets are not associated with the main route
//    table of the VPC.
// 6. None of the the VPC's subnets have the MapPublicIPOnLaunch attribute set.
//
// With the force-vpc-id config setting set to true, the provider can ignore a
// NotValidError. NotFoundError cannot be ignored, while unexpected API
// responses and errors could be retried.
func validateVPC(apiClient vpcAPIClient, vpcID string) error {
	if vpcID == "" || apiClient == nil {
		return errors.Errorf("invalid arguments: empty VPC ID or nil client")
	}

	vpc, err := getVPCByID(apiClient, vpcID)
	if err != nil {
		return errors.Trace(err)
	}

	if err := checkVPCIsAvailable(vpc); err != nil {
		return errors.Trace(err)
	}

	subnets, err := getVPCSubnets(apiClient, vpc)
	if err != nil {
		return errors.Trace(err)
	}

	publicSubnet, err := findFirstPublicSubnet(subnets)
	if err != nil {
		return errors.Trace(err)
	}
	logger.Infof(
		"found subnet %q (%s) in AZ %q, suitable for a Juju controller instance",
		publicSubnet.Id, publicSubnet.CIDRBlock, publicSubnet.AvailZone,
	)

	gateway, err := getVPCInternetGateway(apiClient, vpc)
	if err != nil {
		return errors.Trace(err)
	}

	if err := checkInternetGatewayIsAttached(gateway); err != nil {
		return errors.Trace(err)
	}

	routeTables, err := getVPCRouteTables(apiClient, vpc)
	if err != nil {
		return errors.Trace(err)
	}

	mainRouteTable, err := findVPCMainRouteTable(vpc, routeTables)
	if err != nil {
		return errors.Trace(err)
	}

	if err := checkVPCRouteTableRoutes(vpc, mainRouteTable, gateway); err != nil {
		return errors.Trace(err)
	}

	logger.Infof("VPC %q is suitable for Juju controllers and expose-able workloads", vpc.Id)
	return nil
}

func getVPCByID(apiClient vpcAPIClient, vpcID string) (*ec2.VPC, error) {
	response, err := apiClient.VPCs([]string{vpcID}, nil)
	if isVPCNotFoundError(err) {
		return nil, errors.NewNotFound(err, "")
	} else if err != nil {
		return nil, errors.Annotatef(err, "unexpected AWS response getting VPC %q", vpcID)
	}

	if numResults := len(response.VPCs); numResults == 0 {
		return nil, errors.NotFoundf("VPC %q", vpcID)
	} else if numResults > 1 {
		logger.Debugf("VPCs() returned %#v", response)
		return nil, errors.Errorf("expected 1 result from AWS, got %d", numResults)
	}

	vpc := response.VPCs[0]
	return &vpc, nil
}

func isVPCNotFoundError(err error) bool {
	return err != nil && ec2ErrCode(err) == "InvalidVpcID.NotFound"
}

func checkVPCIsAvailable(vpc *ec2.VPC) error {
	if vpc.State != availableState {
		return errors.NotValidf("VPC with unexpected state %q", vpc.State)
	}

	if vpc.IsDefault {
		logger.Infof("VPC %q is the default VPC for the region", vpc.Id)
	}

	return nil
}

func getVPCSubnets(apiClient vpcAPIClient, vpc *ec2.VPC) ([]ec2.Subnet, error) {
	filter := ec2.NewFilter()
	filter.Add("vpc-id", vpc.Id)
	response, err := apiClient.Subnets(nil, filter)
	if err != nil {
		return nil, errors.Annotate(err, "unexpected AWS response getting VPC subnets")
	}

	if len(response.Subnets) == 0 {
		message := fmt.Sprintf("no subnets found for VPC %q", vpc.Id)
		return nil, errors.NewNotFound(nil, message)
	}

	return response.Subnets, nil
}

func findFirstPublicSubnet(subnets []ec2.Subnet) (*ec2.Subnet, error) {
	for _, subnet := range subnets {
		if subnet.MapPublicIPOnLaunch {
			return &subnet, nil
		}

	}
	return nil, errors.NotValidf("VPC without any public subnets")
}

func getVPCInternetGateway(apiClient vpcAPIClient, vpc *ec2.VPC) (*ec2.InternetGateway, error) {
	filter := ec2.NewFilter()
	filter.Add("vpc-id", vpc.Id)
	response, err := apiClient.InternetGateways(nil, filter)
	if err != nil {
		return nil, errors.Annotatef(err, "unexpected AWS response getting Internet Gateway of VPC %q", vpc.Id)
	}

	if numResults := len(response.InternetGateways); numResults == 0 {
		return nil, errors.NotValidf("VPC without Internet Gateway")
	} else if numResults > 1 {
		logger.Debugf("InternetGateways() returned %#v", response)
		return nil, errors.Errorf("expected 1 result from AWS, got %d", numResults)
	}

	gateway := response.InternetGateways[0]
	return &gateway, nil
}

func checkInternetGatewayIsAttached(gateway *ec2.InternetGateway) error {
	if state := gateway.AttachmentState; state != attachedState {
		errors.NotValidf("VPC Internet Gateway %q in unexpected state %q", gateway.Id, state)
	}

	return nil
}

func getVPCRouteTables(apiClient vpcAPIClient, vpc *ec2.VPC) ([]ec2.RouteTable, error) {
	filter := ec2.NewFilter()
	filter.Add("vpc-id", vpc.Id)
	response, err := apiClient.RouteTables(nil, filter)
	if err != nil {
		return nil, errors.Annotatef(err, "unexpected AWS response getting VPC %q route tables", vpc.Id)
	}

	if len(response.Tables) == 0 {
		return nil, errors.NotValidf("VPC without any route tables")
	}
	logger.Tracef("RouteTables() returned %#v", response)

	return response.Tables, nil
}

func findVPCMainRouteTable(vpc *ec2.VPC, routeTables []ec2.RouteTable) (*ec2.RouteTable, error) {
	var mainTable *ec2.RouteTable
	for _, table := range routeTables {
		if len(table.Associations) < 1 {
			logger.Tracef("ignoring VPC %q route table %q with no associations", vpc.Id, table.Id)
			continue
		}

		for _, association := range table.Associations {
			if subnetID := association.SubnetId; subnetID != "" {
				return nil, errors.NotValidf("subnet %q not associated with VPC main route table", subnetID)
			}
			if association.IsMain && mainTable == nil {
				logger.Tracef("main route table of VPC %q has ID %q", vpc.Id, table.Id)
				mainTable = &table
			}
		}
	}

	if mainTable == nil {
		return nil, errors.NotValidf("VPC without associated main route table")
	}

	return mainTable, nil
}

func checkVPCRouteTableRoutes(vpc *ec2.VPC, routeTable *ec2.RouteTable, gateway *ec2.InternetGateway) error {
	hasDefaultRoute := false
	hasLocalRoute := false

	logger.Tracef("checking route table %+v routes", routeTable)
	for _, route := range routeTable.Routes {
		if route.State != activeState {
			logger.Tracef("skipping inactive route %+v", route)
			continue
		}

		switch route.DestinationCIDRBlock {
		case defaultRouteCIDRBlock:
			if route.GatewayId == gateway.Id {
				logger.Tracef("default route uses expected gateway %q", gateway.Id)
				hasDefaultRoute = true
			}
		case vpc.CIDRBlock:
			if route.GatewayId == localRouteGatewayID {
				logger.Tracef("local route uses expected CIDR %q", vpc.CIDRBlock)
				hasLocalRoute = true
			}
		default:
			logger.Tracef("route %+v is neither local nor default (skipping)", route)
		}
	}

	if hasDefaultRoute && hasLocalRoute {
		return nil
	}

	if !hasDefaultRoute {
		return errors.NotValidf("missing default route via gateway %q", gateway.Id)
	}
	return errors.NotValidf("missing local route with destination %q", vpc.CIDRBlock)
}

func findDefaultVPCID(apiClient vpcAPIClient) (string, error) {
	response, err := apiClient.AccountAttributes("default-vpc")
	if err != nil {
		return "", errors.Annotate(err, "unexpected AWS response getting default-vpc")
	}

	if len(response.Attributes) == 0 || len(response.Attributes[0].Values) == 0 {
		// No value for the requested "default-vpc" attribute, all bets are off.
		return "", errors.NotFoundf("default-vpc account attribute")
	}

	firstAttributeValue := response.Attributes[0].Values[0]
	if firstAttributeValue == defaultVPCIDNone {
		return "", errors.NotFoundf("default VPC")
	}

	return firstAttributeValue, nil
}

func getVPCSubnetIDsForAvailabilityZone(apiClient vpcAPIClient, vpcID, zoneName string) ([]string, error) {
	vpc := &ec2.VPC{Id: vpcID}
	subnets, err := getVPCSubnets(apiClient, vpc)
	if err != nil {
		return nil, errors.Annotatef(err, "cannot get VPC %q subnets", vpcID)
	}

	matchingSubnetIDs := set.NewStrings()
	for _, subnet := range subnets {
		if subnet.AvailZone == zoneName {
			matchingSubnetIDs.Add(subnet.Id)
		}
	}

	if matchingSubnetIDs.IsEmpty() {
		return nil, errors.NotFoundf("VPC %q subnet in AZ %q", vpcID, zoneName)
	}

	return matchingSubnetIDs.SortedValues(), nil
}

func findSubnetIDsForAvailabilityZone(zoneName string, subnetsToZones map[network.Id][]string) ([]string, error) {
	matchingSubnetIDs := set.NewStrings()
	for subnetID, zones := range subnetsToZones {
		zonesSet := set.NewStrings(zones...)
		if zonesSet.Contains(zoneName) {
			matchingSubnetIDs.Add(string(subnetID))
		}
	}

	if matchingSubnetIDs.IsEmpty() {
		return nil, errors.NotFoundf("subnets in AZ %q", zoneName)
	}

	return matchingSubnetIDs.SortedValues(), nil
}
<commit_msg>Fixed imports ordering and an issue discovered while live testing<commit_after>// Copyright 2016 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package ec2

import (
	"fmt"

	"github.com/juju/errors"
	"github.com/juju/utils/set"
	"gopkg.in/amz.v3/ec2"

	"github.com/juju/juju/network"
)

const (
	activeState           = "active"
	availableState        = "available"
	attachedState         = "attached"
	localRouteGatewayID   = "local"
	defaultRouteCIDRBlock = "0.0.0.0/0"
	defaultVPCIDNone      = "none"
)

// vpcAPIClient defines a subset of the goamz API calls needed to validate a VPC.
type vpcAPIClient interface {
	AccountAttributes(attributeNames ...string) (*ec2.AccountAttributesResp, error)
	VPCs(ids []string, filter *ec2.Filter) (*ec2.VPCsResp, error)
	Subnets(ids []string, filter *ec2.Filter) (*ec2.SubnetsResp, error)
	InternetGateways(ids []string, filter *ec2.Filter) (*ec2.InternetGatewaysResp, error)
	RouteTables(ids []string, filter *ec2.Filter) (*ec2.RouteTablesResp, error)
}

// validateVPC requires both arguments to be set and validates that vpcID refers
// to an existing AWS VPC (default or non-default) for the current region.
// Returns an error satifying errors.IsNotFound() when the VPC with the given
// vpcID cannot be found, or when the VPC exists but contains no subnets.
// Returns an error satisfying errors.IsNotValid() in the following cases:
//
// 1. The VPC's state is not "available".
// 2. The VPC does not have an Internet Gateway (IGW) attached.
// 3. A main route table is not associated with the VPC.
// 4. The main route table lacks both a default route via the IGW and a local
//    route matching the VPC's CIDR block.
// 5. One or more of the VPC's subnets are not associated with the main route
//    table of the VPC.
// 6. None of the the VPC's subnets have the MapPublicIPOnLaunch attribute set.
//
// With the force-vpc-id config setting set to true, the provider can ignore a
// NotValidError. NotFoundError cannot be ignored, while unexpected API
// responses and errors could be retried.
func validateVPC(apiClient vpcAPIClient, vpcID string) error {
	if vpcID == "" || apiClient == nil {
		return errors.Errorf("invalid arguments: empty VPC ID or nil client")
	}

	vpc, err := getVPCByID(apiClient, vpcID)
	if err != nil {
		return errors.Trace(err)
	}

	if err := checkVPCIsAvailable(vpc); err != nil {
		return errors.Trace(err)
	}

	subnets, err := getVPCSubnets(apiClient, vpc)
	if err != nil {
		return errors.Trace(err)
	}

	publicSubnet, err := findFirstPublicSubnet(subnets)
	if err != nil {
		return errors.Trace(err)
	}
	logger.Infof(
		"found subnet %q (%s) in AZ %q, suitable for a Juju controller instance",
		publicSubnet.Id, publicSubnet.CIDRBlock, publicSubnet.AvailZone,
	)

	gateway, err := getVPCInternetGateway(apiClient, vpc)
	if err != nil {
		return errors.Trace(err)
	}

	if err := checkInternetGatewayIsAttached(gateway); err != nil {
		return errors.Trace(err)
	}

	routeTables, err := getVPCRouteTables(apiClient, vpc)
	if err != nil {
		return errors.Trace(err)
	}

	mainRouteTable, err := findVPCMainRouteTable(vpc, routeTables)
	if err != nil {
		return errors.Trace(err)
	}

	if err := checkVPCRouteTableRoutes(vpc, mainRouteTable, gateway); err != nil {
		return errors.Trace(err)
	}

	logger.Infof("VPC %q is suitable for Juju controllers and expose-able workloads", vpc.Id)
	return nil
}

func getVPCByID(apiClient vpcAPIClient, vpcID string) (*ec2.VPC, error) {
	response, err := apiClient.VPCs([]string{vpcID}, nil)
	if isVPCNotFoundError(err) {
		return nil, errors.NewNotFound(err, "")
	} else if err != nil {
		return nil, errors.Annotatef(err, "unexpected AWS response getting VPC %q", vpcID)
	}

	if numResults := len(response.VPCs); numResults == 0 {
		return nil, errors.NotFoundf("VPC %q", vpcID)
	} else if numResults > 1 {
		logger.Debugf("VPCs() returned %#v", response)
		return nil, errors.Errorf("expected 1 result from AWS, got %d", numResults)
	}

	vpc := response.VPCs[0]
	return &vpc, nil
}

func isVPCNotFoundError(err error) bool {
	return err != nil && ec2ErrCode(err) == "InvalidVpcID.NotFound"
}

func checkVPCIsAvailable(vpc *ec2.VPC) error {
	if vpc.State != availableState {
		return errors.NotValidf("VPC with unexpected state %q", vpc.State)
	}

	if vpc.IsDefault {
		logger.Infof("VPC %q is the default VPC for the region", vpc.Id)
	}

	return nil
}

func getVPCSubnets(apiClient vpcAPIClient, vpc *ec2.VPC) ([]ec2.Subnet, error) {
	filter := ec2.NewFilter()
	filter.Add("vpc-id", vpc.Id)
	response, err := apiClient.Subnets(nil, filter)
	if err != nil {
		return nil, errors.Annotate(err, "unexpected AWS response getting VPC subnets")
	}

	if len(response.Subnets) == 0 {
		message := fmt.Sprintf("no subnets found for VPC %q", vpc.Id)
		return nil, errors.NewNotFound(nil, message)
	}

	return response.Subnets, nil
}

func findFirstPublicSubnet(subnets []ec2.Subnet) (*ec2.Subnet, error) {
	for _, subnet := range subnets {
		if subnet.MapPublicIPOnLaunch {
			return &subnet, nil
		}

	}
	return nil, errors.NotValidf("VPC without any public subnets")
}

func getVPCInternetGateway(apiClient vpcAPIClient, vpc *ec2.VPC) (*ec2.InternetGateway, error) {
	filter := ec2.NewFilter()
	filter.Add("attachment.vpc-id", vpc.Id)
	response, err := apiClient.InternetGateways(nil, filter)
	if err != nil {
		return nil, errors.Annotatef(err, "unexpected AWS response getting Internet Gateway of VPC %q", vpc.Id)
	}

	if numResults := len(response.InternetGateways); numResults == 0 {
		return nil, errors.NotValidf("VPC without Internet Gateway")
	} else if numResults > 1 {
		logger.Debugf("InternetGateways() returned %#v", response)
		return nil, errors.Errorf("expected 1 result from AWS, got %d", numResults)
	}

	gateway := response.InternetGateways[0]
	return &gateway, nil
}

func checkInternetGatewayIsAttached(gateway *ec2.InternetGateway) error {
	if state := gateway.AttachmentState; state != attachedState {
		errors.NotValidf("VPC Internet Gateway %q in unexpected state %q", gateway.Id, state)
	}

	return nil
}

func getVPCRouteTables(apiClient vpcAPIClient, vpc *ec2.VPC) ([]ec2.RouteTable, error) {
	filter := ec2.NewFilter()
	filter.Add("vpc-id", vpc.Id)
	response, err := apiClient.RouteTables(nil, filter)
	if err != nil {
		return nil, errors.Annotatef(err, "unexpected AWS response getting VPC %q route tables", vpc.Id)
	}

	if len(response.Tables) == 0 {
		return nil, errors.NotValidf("VPC without any route tables")
	}
	logger.Tracef("RouteTables() returned %#v", response)

	return response.Tables, nil
}

func findVPCMainRouteTable(vpc *ec2.VPC, routeTables []ec2.RouteTable) (*ec2.RouteTable, error) {
	var mainTable *ec2.RouteTable
	for _, table := range routeTables {
		if len(table.Associations) < 1 {
			logger.Tracef("ignoring VPC %q route table %q with no associations", vpc.Id, table.Id)
			continue
		}

		for _, association := range table.Associations {
			if subnetID := association.SubnetId; subnetID != "" {
				return nil, errors.NotValidf("subnet %q not associated with VPC main route table", subnetID)
			}
			if association.IsMain && mainTable == nil {
				logger.Tracef("main route table of VPC %q has ID %q", vpc.Id, table.Id)
				mainTable = &table
			}
		}
	}

	if mainTable == nil {
		return nil, errors.NotValidf("VPC without associated main route table")
	}

	return mainTable, nil
}

func checkVPCRouteTableRoutes(vpc *ec2.VPC, routeTable *ec2.RouteTable, gateway *ec2.InternetGateway) error {
	hasDefaultRoute := false
	hasLocalRoute := false

	logger.Tracef("checking route table %+v routes", routeTable)
	for _, route := range routeTable.Routes {
		if route.State != activeState {
			logger.Tracef("skipping inactive route %+v", route)
			continue
		}

		switch route.DestinationCIDRBlock {
		case defaultRouteCIDRBlock:
			if route.GatewayId == gateway.Id {
				logger.Tracef("default route uses expected gateway %q", gateway.Id)
				hasDefaultRoute = true
			}
		case vpc.CIDRBlock:
			if route.GatewayId == localRouteGatewayID {
				logger.Tracef("local route uses expected CIDR %q", vpc.CIDRBlock)
				hasLocalRoute = true
			}
		default:
			logger.Tracef("route %+v is neither local nor default (skipping)", route)
		}
	}

	if hasDefaultRoute && hasLocalRoute {
		return nil
	}

	if !hasDefaultRoute {
		return errors.NotValidf("missing default route via gateway %q", gateway.Id)
	}
	return errors.NotValidf("missing local route with destination %q", vpc.CIDRBlock)
}

func findDefaultVPCID(apiClient vpcAPIClient) (string, error) {
	response, err := apiClient.AccountAttributes("default-vpc")
	if err != nil {
		return "", errors.Annotate(err, "unexpected AWS response getting default-vpc")
	}

	if len(response.Attributes) == 0 || len(response.Attributes[0].Values) == 0 {
		// No value for the requested "default-vpc" attribute, all bets are off.
		return "", errors.NotFoundf("default-vpc account attribute")
	}

	firstAttributeValue := response.Attributes[0].Values[0]
	if firstAttributeValue == defaultVPCIDNone {
		return "", errors.NotFoundf("default VPC")
	}

	return firstAttributeValue, nil
}

func getVPCSubnetIDsForAvailabilityZone(apiClient vpcAPIClient, vpcID, zoneName string) ([]string, error) {
	vpc := &ec2.VPC{Id: vpcID}
	subnets, err := getVPCSubnets(apiClient, vpc)
	if err != nil {
		return nil, errors.Annotatef(err, "cannot get VPC %q subnets", vpcID)
	}

	matchingSubnetIDs := set.NewStrings()
	for _, subnet := range subnets {
		if subnet.AvailZone == zoneName {
			matchingSubnetIDs.Add(subnet.Id)
		}
	}

	if matchingSubnetIDs.IsEmpty() {
		return nil, errors.NotFoundf("VPC %q subnet in AZ %q", vpcID, zoneName)
	}

	return matchingSubnetIDs.SortedValues(), nil
}

func findSubnetIDsForAvailabilityZone(zoneName string, subnetsToZones map[network.Id][]string) ([]string, error) {
	matchingSubnetIDs := set.NewStrings()
	for subnetID, zones := range subnetsToZones {
		zonesSet := set.NewStrings(zones...)
		if zonesSet.Contains(zoneName) {
			matchingSubnetIDs.Add(string(subnetID))
		}
	}

	if matchingSubnetIDs.IsEmpty() {
		return nil, errors.NotFoundf("subnets in AZ %q", zoneName)
	}

	return matchingSubnetIDs.SortedValues(), nil
}
<|endoftext|>


<commit_before>//go:generate ./e.sh

package main

import "encoding/json"
import "flag"
import "fmt"
import "io/ioutil"
import "os"
import "path/filepath"
import "sort"
import "strconv"
import "strings"
import "text/tabwriter"

import "github.com/appc/spec/schema"
import "github.com/appc/spec/schema/types"
import "github.com/juju/errors"

import "lib/jetpack"
import "lib/run"

var Host *jetpack.Host

func die(err error) {
	if err != nil {
		fmt.Fprintln(os.Stderr, errors.ErrorStack(err))
		os.Exit(1)
	}
}

func show(obj ...interface{}) {
	die(Show("", obj...))
}

func subcommand(def string, args []string) (string, []string) {
	if len(args) == 0 {
		return def, args
	}
	return args[0], args[1:]
}

func image(name string) *jetpack.Image {
	img, err := Host.FindImage(name)
	if err == jetpack.ErrNotFound {
		die(errors.Errorf("No such image: %#v", name))
	}
	die(err)
	return img
}

func getRuntimeApp(name string) (*schema.RuntimeApp, error) {
	if img, err := Host.FindImage(name); err != nil {
		return nil, err
	} else {
		rta := img.RuntimeApp()
		return &rta, nil
	}
}

func main() {
	configPath := jetpack.DefaultConfigPath
	help := false

	if cfg := os.Getenv("JETPACK_CONF"); cfg != "" {
		configPath = cfg
	}

	flag.StringVar(&configPath, "config", configPath, "Configuration file")
	flag.BoolVar(&help, "h", false, "Show help")
	flag.BoolVar(&help, "help", false, "Show help")

	flag.Parse()
	args := flag.Args()

	if help || len(args) == 0 || args[0] == "help" {
		fmt.Fprintf(os.Stderr, `Usage: %s [OPTIONS] COMMAND...
Options:
  -config=PATH  Configuration file (%s)
  -help, -h     Display this help screen
Commands:
  help                                    Display this help screen
  init                                    Initialize host
  info                                    Show global information
  test                                    Run integration tests
  trust [-l|-list]                        List trusted ACI signign keys
  trust [-prefix PREFIX] [-root] [KEY]    Trust ACI signing key
  trust -d FINGERPRINT                    Untrust ACI signing key
  fetch NAME...                           Fetch ACI
  image list [QUERY]                      List images
  image IMAGE build [OPTIONS] COMMAND...  Build new image from an existing one
                    -dir=.                Location on build directory on host
                    -cp=PATH...           Copy additional files from host
  image IMAGE show                        Display image details
  image IMAGE export [PATH]               Export image to an AMI file
                                          Output to stdout if no PATH given
  image IMAGE destroy                     Destroy image
  pod list                                List pods
  pod create [FLAGS] IMAGE [IMAGE FLAGS] [IMAGE [IMAGE FLAGS] ...]
                                          Create new pod from image
             -help                        Show detailed help
  pod POD show                            Display pod details
  pod POD run [APP]                       Run pod's application
  pod POD console [APP]                   Open console inside the pod
  pod POD ps|top|killall [OPTIONS...]
                                          Manage pod's processes
  pod POD kill                            Kill running pod
  pod POD destroy                         Destroy pod
Needs Explanation:
  ARCHIVE, MANIFEST  May be filesystem paths or URLs.
            cp=PATH  This option can be given multiple times
              QUERY  Is an expression that looks like this:
                      - NAME[,LABEL=VALUE[,LABEL=VALUE[,...]]]
                      - NAME:VERSION (alias for NAME:version=VERSION)
              IMAGE  Can be:
                      - an UUID (XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXX),
                      - a checksum (sha512-...), or
                      - a QUERY (which can't be ambiguous).
          POD  Has to be an UUID for now
Helpful Aliases:
  i|img ... -- image ...
  p ... -- pod ...
  image, images -- image list
  pod, pods -- pod list
  image build|show|export|destroy IMAGE ... -- image IMAGE build|show|... ...
`,
			filepath.Base(os.Args[0]), configPath)
		return
	}

	command := args[0]
	args = args[1:]

	if host, err := jetpack.NewHost(configPath); err != nil {
		die(err)
	} else {
		Host = host
	}

	if command == "init" {
		// Init is special: it doesn't need an initialized host
		die(Host.Initialize())
		show(Host)
		return
	}

	if Host.Dataset == nil {
		die(errors.New("Host is not initialized"))
	}

	switch command {
	case "info":
		show(Host)
	case "test":
		die(run.Command(filepath.Join(jetpack.LibexecPath, "test.integration"),
			append(args, "dataset="+Host.Dataset.Name)...).Run())
	case "trust":
		die(runTrust(args))
	case "fetch":
		die(runFetch(args))
	case "images":
		command = "image"
		args = append([]string{"list"}, args...)
		fallthrough
	case "image", "img", "i":
		switch command, args := subcommand("list", args); command {
		case "list":
			var machineFriendly, showHash, idOnly bool
			fl := flag.NewFlagSet("image list", flag.ExitOnError)
			fl.BoolVar(&machineFriendly, "H", false, "Machine-friendly output")
			fl.BoolVar(&showHash, "hash", false, "Show image hash instead of UUID")
			fl.BoolVar(&idOnly, "q", false, "Show only ID")
			die(fl.Parse(args))

			images := Host.Images()

			if idOnly {
				for _, img := range images {
					if showHash {
						fmt.Println(img.Hash)
					} else {
						fmt.Println(img.UUID)
					}
				}
			} else if len(images) == 0 {
				if !machineFriendly {
					show("No images")
				}
			} else {
				lines := make([]string, len(images))
				for i, img := range images {
					labels := make([]string, len(img.Manifest.Labels))
					for j, label := range img.Manifest.Labels {
						labels[j] = fmt.Sprintf("%v=%#v", label.Name, label.Value)
					}
					sort.Strings(labels)
					first := img.UUID.String()
					if showHash {
						first = img.Hash.String()
					}
					lines[i] = fmt.Sprintf("%v\t%v\t%v",
						first,
						img.Manifest.Name,
						strings.Join(labels, ","))
				}
				sort.Strings(lines)
				output := strings.Join(lines, "\n")

				if machineFriendly {
					fmt.Println(output)
				} else {
					first := "UUID"
					if showHash {
						first = "HASH"
					}
					w := tabwriter.NewWriter(os.Stdout, 2, 8, 2, ' ', 0)
					fmt.Fprintf(w, "%v\tNAME\tLABELS\n%v\n", first, output)
					die(w.Flush())
				}
			}
		case "build", "show", "export", "destroy":
			// be nice to people who prefer to type UUID after command
			command, args[0] = args[0], command
			fallthrough
		default:
			img := image(command)

			switch command, args := subcommand("show", args); command {
			case "build":
				var copyFiles sliceFlag
				var buildDir string

				fs := flag.NewFlagSet("build", flag.ExitOnError)
				fs.Var(&copyFiles, "cp", "")
				fs.StringVar(&buildDir, "dir", ".", "")
				die(fs.Parse(args))

				newImage, err := img.Build(buildDir, copyFiles, fs.Args())
				die(err)
				show(newImage)
			case "show":
				show(img)
			case "export":
				path := "-"
				if len(args) > 0 {
					path = args[0]
				}
				if hash, err := img.SaveAMI(path, 0644); err != nil {
					die(err)
				} else {
					fmt.Fprintln(os.Stderr, hash)
				}
			case "destroy":
				die(img.Destroy())
			default:
				die(errors.Errorf("Unknown command %#v", command))
			}
		}
	case "pods":
		command = "pod"
		args = append([]string{"list"}, args...)
		fallthrough
	case "pod", "p":
		switch command, args := subcommand("list", args); command {
		case "create":
			var dryRun, doRun, doDestroy bool
			var saveId string
			fl := flag.NewFlagSet("jetpack pod create", flag.ContinueOnError)
			fl.BoolVar(&dryRun, "n", false, "Dry run (don't actually create pod, just show manifest)")
			fl.BoolVar(&doRun, "run", false, "Run pod immediately")
			fl.BoolVar(&doDestroy, "destroy", false, "Destroy pod after running (meaningless without -run)")
			fl.StringVar(&saveId, "saveid", "", "Save pod UUID to file")

			if pm, err := ConstructPod(args, fl, getRuntimeApp); err == flag.ErrHelp {
				// It's all right. Help has been shown.
			} else if err != nil {
				panic(err)
			} else if dryRun {
				if jb, err := json.MarshalIndent(pm, "", "  "); err != nil {
					panic(err)
				} else {
					fmt.Println(string(jb))
				}
			} else {
				pod, err := Host.CreatePod(pm)
				die(err)
				if saveId != "" {
					die(ioutil.WriteFile(saveId, []byte(pod.UUID.String()), 0644))
				}
				if doRun {
					if len(pod.Manifest.Apps) > 1 {
						die(errors.New("Pod has multiple apps, cannot run"))
					}
					err := pod.RunApp(pod.Manifest.Apps[0].Name)
					if doDestroy {
						err1 := pod.Destroy()
						if err == nil {
							err = err1
						}
					}
					die(err)
				} else {
					show(pod)
				}
			}
		case "list":
			var machineFriendly, idOnly bool
			fl := flag.NewFlagSet("pod list", flag.ExitOnError)
			fl.BoolVar(&machineFriendly, "H", false, "Machine-friendly output")
			fl.BoolVar(&idOnly, "q", false, "Show only ID")
			fl.Parse(args)

			pods := Host.Pods()

			if idOnly {
				for _, pod := range pods {
					fmt.Println(pod.UUID)
				}
			} else if len(pods) == 0 {
				if !machineFriendly {
					show("No pods")
				}
			} else {
				lines := make([]string, len(pods))
				for i, pod := range pods {
					apps := make([]string, len(pod.Manifest.Apps))
					for j, app := range pod.Manifest.Apps {
						apps[j] = app.Name.String()
					}
					ipAddress, _ := pod.Manifest.Annotations.Get("ip-address")
					lines[i] = fmt.Sprintf("%v\t%v\t%v\t%v",
						pod.UUID,
						pod.Status().String(),
						ipAddress,
						strings.Join(apps, " "))
				}
				sort.Strings(lines)
				output := strings.Join(lines, "\n")

				if machineFriendly {
					fmt.Println(output)
				} else {
					w := tabwriter.NewWriter(os.Stdout, 2, 8, 2, ' ', 0)
					fmt.Fprintf(w, "UUID\tSTATUS\tIP\tAPPS\n%v\n", output)
					die(w.Flush())
				}
			}
		case "show", "run", "ps", "top", "killall", "kill", "destroy":
			// be nice to people who prefer to type UUID after command
			command, args[0] = args[0], command
			fallthrough
		default:
			pod, err := Host.FindPod(command)
			if err == jetpack.ErrNotFound {
				die(errors.Errorf("No such pod: %#v", command))
			}
			die(err)
			switch command, args := subcommand("show", args); command {
			case "show":
				show(pod)
			case "run":
				switch len(args) {
				case 0:
					if len(pod.Manifest.Apps) > 1 {
						die(errors.New("Pod has multiple apps, you need to specify one"))
					}
					die(pod.RunApp(pod.Manifest.Apps[0].Name))
				case 1:
					die(pod.RunApp(types.ACName(args[0])))
				default:
					die(errors.New("Command `run' takes at most one argument"))
				}
			case "console":
				switch len(args) {
				case 0:
					if len(pod.Manifest.Apps) > 1 {
						die(errors.New("Pod has multiple apps, you need to specify one"))
					}
					die(pod.Console(pod.Manifest.Apps[0].Name, "root"))
				case 1:
					die(pod.Console(types.ACName(args[0]), "root"))
				default:
					die(errors.New("Command `console' takes at most one argument"))
				}
			case "ps", "top", "killall":
				jid := pod.Jid()
				if jid == 0 {
					die(errors.New("Pod is not running"))
				}

				flag := "-J"
				if command == "killall" {
					flag = "-j"
				}

				die(run.Command(command, append([]string{flag, strconv.Itoa(jid)}, args...)...).Run())
			case "kill":
				die(pod.Kill())
			case "destroy":
				die(pod.Destroy())
			default:
				die(errors.Errorf("Unknown command %#v", command))
			}
		}
	default:
		die(errors.Errorf("Unknown command %#v", command))
	}
}
<commit_msg>Delete broken `jetpack test`<commit_after>//go:generate ./e.sh

package main

import "encoding/json"
import "flag"
import "fmt"
import "io/ioutil"
import "os"
import "path/filepath"
import "sort"
import "strconv"
import "strings"
import "text/tabwriter"

import "github.com/appc/spec/schema"
import "github.com/appc/spec/schema/types"
import "github.com/juju/errors"

import "lib/jetpack"
import "lib/run"

var Host *jetpack.Host

func die(err error) {
	if err != nil {
		fmt.Fprintln(os.Stderr, errors.ErrorStack(err))
		os.Exit(1)
	}
}

func show(obj ...interface{}) {
	die(Show("", obj...))
}

func subcommand(def string, args []string) (string, []string) {
	if len(args) == 0 {
		return def, args
	}
	return args[0], args[1:]
}

func image(name string) *jetpack.Image {
	img, err := Host.FindImage(name)
	if err == jetpack.ErrNotFound {
		die(errors.Errorf("No such image: %#v", name))
	}
	die(err)
	return img
}

func getRuntimeApp(name string) (*schema.RuntimeApp, error) {
	if img, err := Host.FindImage(name); err != nil {
		return nil, err
	} else {
		rta := img.RuntimeApp()
		return &rta, nil
	}
}

func main() {
	configPath := jetpack.DefaultConfigPath
	help := false

	if cfg := os.Getenv("JETPACK_CONF"); cfg != "" {
		configPath = cfg
	}

	flag.StringVar(&configPath, "config", configPath, "Configuration file")
	flag.BoolVar(&help, "h", false, "Show help")
	flag.BoolVar(&help, "help", false, "Show help")

	flag.Parse()
	args := flag.Args()

	if help || len(args) == 0 || args[0] == "help" {
		fmt.Fprintf(os.Stderr, `Usage: %s [OPTIONS] COMMAND...
Options:
  -config=PATH  Configuration file (%s)
  -help, -h     Display this help screen
Commands:
  help                                    Display this help screen
  init                                    Initialize host
  info                                    Show global information
  trust [-l|-list]                        List trusted ACI signign keys
  trust [-prefix PREFIX] [-root] [KEY]    Trust ACI signing key
  trust -d FINGERPRINT                    Untrust ACI signing key
  fetch NAME...                           Fetch ACI
  image list [QUERY]                      List images
  image IMAGE build [OPTIONS] COMMAND...  Build new image from an existing one
                    -dir=.                Location on build directory on host
                    -cp=PATH...           Copy additional files from host
  image IMAGE show                        Display image details
  image IMAGE export [PATH]               Export image to an AMI file
                                          Output to stdout if no PATH given
  image IMAGE destroy                     Destroy image
  pod list                                List pods
  pod create [FLAGS] IMAGE [IMAGE FLAGS] [IMAGE [IMAGE FLAGS] ...]
                                          Create new pod from image
             -help                        Show detailed help
  pod POD show                            Display pod details
  pod POD run [APP]                       Run pod's application
  pod POD console [APP]                   Open console inside the pod
  pod POD ps|top|killall [OPTIONS...]
                                          Manage pod's processes
  pod POD kill                            Kill running pod
  pod POD destroy                         Destroy pod
Needs Explanation:
  ARCHIVE, MANIFEST  May be filesystem paths or URLs.
            cp=PATH  This option can be given multiple times
              QUERY  Is an expression that looks like this:
                      - NAME[,LABEL=VALUE[,LABEL=VALUE[,...]]]
                      - NAME:VERSION (alias for NAME:version=VERSION)
              IMAGE  Can be:
                      - an UUID (XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXX),
                      - a checksum (sha512-...), or
                      - a QUERY (which can't be ambiguous).
          POD  Has to be an UUID for now
Helpful Aliases:
  i|img ... -- image ...
  p ... -- pod ...
  image, images -- image list
  pod, pods -- pod list
  image build|show|export|destroy IMAGE ... -- image IMAGE build|show|... ...
`,
			filepath.Base(os.Args[0]), configPath)
		return
	}

	command := args[0]
	args = args[1:]

	if host, err := jetpack.NewHost(configPath); err != nil {
		die(err)
	} else {
		Host = host
	}

	if command == "init" {
		// Init is special: it doesn't need an initialized host
		die(Host.Initialize())
		show(Host)
		return
	}

	if Host.Dataset == nil {
		die(errors.New("Host is not initialized"))
	}

	switch command {
	case "info":
		show(Host)
	case "trust":
		die(runTrust(args))
	case "fetch":
		die(runFetch(args))
	case "images":
		command = "image"
		args = append([]string{"list"}, args...)
		fallthrough
	case "image", "img", "i":
		switch command, args := subcommand("list", args); command {
		case "list":
			var machineFriendly, showHash, idOnly bool
			fl := flag.NewFlagSet("image list", flag.ExitOnError)
			fl.BoolVar(&machineFriendly, "H", false, "Machine-friendly output")
			fl.BoolVar(&showHash, "hash", false, "Show image hash instead of UUID")
			fl.BoolVar(&idOnly, "q", false, "Show only ID")
			die(fl.Parse(args))

			images := Host.Images()

			if idOnly {
				for _, img := range images {
					if showHash {
						fmt.Println(img.Hash)
					} else {
						fmt.Println(img.UUID)
					}
				}
			} else if len(images) == 0 {
				if !machineFriendly {
					show("No images")
				}
			} else {
				lines := make([]string, len(images))
				for i, img := range images {
					labels := make([]string, len(img.Manifest.Labels))
					for j, label := range img.Manifest.Labels {
						labels[j] = fmt.Sprintf("%v=%#v", label.Name, label.Value)
					}
					sort.Strings(labels)
					first := img.UUID.String()
					if showHash {
						first = img.Hash.String()
					}
					lines[i] = fmt.Sprintf("%v\t%v\t%v",
						first,
						img.Manifest.Name,
						strings.Join(labels, ","))
				}
				sort.Strings(lines)
				output := strings.Join(lines, "\n")

				if machineFriendly {
					fmt.Println(output)
				} else {
					first := "UUID"
					if showHash {
						first = "HASH"
					}
					w := tabwriter.NewWriter(os.Stdout, 2, 8, 2, ' ', 0)
					fmt.Fprintf(w, "%v\tNAME\tLABELS\n%v\n", first, output)
					die(w.Flush())
				}
			}
		case "build", "show", "export", "destroy":
			// be nice to people who prefer to type UUID after command
			command, args[0] = args[0], command
			fallthrough
		default:
			img := image(command)

			switch command, args := subcommand("show", args); command {
			case "build":
				var copyFiles sliceFlag
				var buildDir string

				fs := flag.NewFlagSet("build", flag.ExitOnError)
				fs.Var(&copyFiles, "cp", "")
				fs.StringVar(&buildDir, "dir", ".", "")
				die(fs.Parse(args))

				newImage, err := img.Build(buildDir, copyFiles, fs.Args())
				die(err)
				show(newImage)
			case "show":
				show(img)
			case "export":
				path := "-"
				if len(args) > 0 {
					path = args[0]
				}
				if hash, err := img.SaveAMI(path, 0644); err != nil {
					die(err)
				} else {
					fmt.Fprintln(os.Stderr, hash)
				}
			case "destroy":
				die(img.Destroy())
			default:
				die(errors.Errorf("Unknown command %#v", command))
			}
		}
	case "pods":
		command = "pod"
		args = append([]string{"list"}, args...)
		fallthrough
	case "pod", "p":
		switch command, args := subcommand("list", args); command {
		case "create":
			var dryRun, doRun, doDestroy bool
			var saveId string
			fl := flag.NewFlagSet("jetpack pod create", flag.ContinueOnError)
			fl.BoolVar(&dryRun, "n", false, "Dry run (don't actually create pod, just show manifest)")
			fl.BoolVar(&doRun, "run", false, "Run pod immediately")
			fl.BoolVar(&doDestroy, "destroy", false, "Destroy pod after running (meaningless without -run)")
			fl.StringVar(&saveId, "saveid", "", "Save pod UUID to file")

			if pm, err := ConstructPod(args, fl, getRuntimeApp); err == flag.ErrHelp {
				// It's all right. Help has been shown.
			} else if err != nil {
				panic(err)
			} else if dryRun {
				if jb, err := json.MarshalIndent(pm, "", "  "); err != nil {
					panic(err)
				} else {
					fmt.Println(string(jb))
				}
			} else {
				pod, err := Host.CreatePod(pm)
				die(err)
				if saveId != "" {
					die(ioutil.WriteFile(saveId, []byte(pod.UUID.String()), 0644))
				}
				if doRun {
					if len(pod.Manifest.Apps) > 1 {
						die(errors.New("Pod has multiple apps, cannot run"))
					}
					err := pod.RunApp(pod.Manifest.Apps[0].Name)
					if doDestroy {
						err1 := pod.Destroy()
						if err == nil {
							err = err1
						}
					}
					die(err)
				} else {
					show(pod)
				}
			}
		case "list":
			var machineFriendly, idOnly bool
			fl := flag.NewFlagSet("pod list", flag.ExitOnError)
			fl.BoolVar(&machineFriendly, "H", false, "Machine-friendly output")
			fl.BoolVar(&idOnly, "q", false, "Show only ID")
			fl.Parse(args)

			pods := Host.Pods()

			if idOnly {
				for _, pod := range pods {
					fmt.Println(pod.UUID)
				}
			} else if len(pods) == 0 {
				if !machineFriendly {
					show("No pods")
				}
			} else {
				lines := make([]string, len(pods))
				for i, pod := range pods {
					apps := make([]string, len(pod.Manifest.Apps))
					for j, app := range pod.Manifest.Apps {
						apps[j] = app.Name.String()
					}
					ipAddress, _ := pod.Manifest.Annotations.Get("ip-address")
					lines[i] = fmt.Sprintf("%v\t%v\t%v\t%v",
						pod.UUID,
						pod.Status().String(),
						ipAddress,
						strings.Join(apps, " "))
				}
				sort.Strings(lines)
				output := strings.Join(lines, "\n")

				if machineFriendly {
					fmt.Println(output)
				} else {
					w := tabwriter.NewWriter(os.Stdout, 2, 8, 2, ' ', 0)
					fmt.Fprintf(w, "UUID\tSTATUS\tIP\tAPPS\n%v\n", output)
					die(w.Flush())
				}
			}
		case "show", "run", "ps", "top", "killall", "kill", "destroy":
			// be nice to people who prefer to type UUID after command
			command, args[0] = args[0], command
			fallthrough
		default:
			pod, err := Host.FindPod(command)
			if err == jetpack.ErrNotFound {
				die(errors.Errorf("No such pod: %#v", command))
			}
			die(err)
			switch command, args := subcommand("show", args); command {
			case "show":
				show(pod)
			case "run":
				switch len(args) {
				case 0:
					if len(pod.Manifest.Apps) > 1 {
						die(errors.New("Pod has multiple apps, you need to specify one"))
					}
					die(pod.RunApp(pod.Manifest.Apps[0].Name))
				case 1:
					die(pod.RunApp(types.ACName(args[0])))
				default:
					die(errors.New("Command `run' takes at most one argument"))
				}
			case "console":
				switch len(args) {
				case 0:
					if len(pod.Manifest.Apps) > 1 {
						die(errors.New("Pod has multiple apps, you need to specify one"))
					}
					die(pod.Console(pod.Manifest.Apps[0].Name, "root"))
				case 1:
					die(pod.Console(types.ACName(args[0]), "root"))
				default:
					die(errors.New("Command `console' takes at most one argument"))
				}
			case "ps", "top", "killall":
				jid := pod.Jid()
				if jid == 0 {
					die(errors.New("Pod is not running"))
				}

				flag := "-J"
				if command == "killall" {
					flag = "-j"
				}

				die(run.Command(command, append([]string{flag, strconv.Itoa(jid)}, args...)...).Run())
			case "kill":
				die(pod.Kill())
			case "destroy":
				die(pod.Destroy())
			default:
				die(errors.Errorf("Unknown command %#v", command))
			}
		}
	default:
		die(errors.Errorf("Unknown command %#v", command))
	}
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/Comcast/webpa-common/xmetrics"
)

const (
	ErrorRequestBodyCounter       = "error_request_body_count"
	EmptyRequestBodyCounter       = "empty_request_body_count"
	DeliveryCounter               = "delivery_count"
	DeliveryRetryCounter          = "delivery_retry_count"
	SlowConsumerDroppedMsgCounter = "slow_consumer_dropped_message_count"
	SlowConsumerCounter           = "slow_consumer_cut_off_count"
	IncomingQueueDepth            = "incoming_queue_depth"
	IncomingContentTypeCounter    = "incoming_content_type_count"
	IncomingEventTypeCounter      = "incoming_event_type_count"
	DropsDueToInvalidPayload      = "drops_due_to_invalid_payload"
	OutgoingQueueDepth            = "outgoing_queue_depths"
)

func Metrics() []xmetrics.Metric {
	return []xmetrics.Metric{
		{
			Name: IncomingQueueDepth,
			Help: "The depth of the queue behind the incoming handlers.",
			Type: "gauge",
		},
		{
			Name: ErrorRequestBodyCounter,
			Help: "Count of the number of errors encountered reading the body.",
			Type: "counter",
		},
		{
			Name: EmptyRequestBodyCounter,
			Help: "Count of the number of times the request is an empty body.",
			Type: "counter",
		},
		{
			Name: DropsDueToInvalidPayload,
			Help: "Dropped messages due to invalid payloads.",
			Type: "counter",
		},
		{
			Name:       IncomingContentTypeCounter,
			Help:       "Count of the content type processed.",
			Type:       "counter",
			LabelNames: []string{"content"},
		},
		{
			Name:       DeliveryRetryCounter,
			Help:       "Number of delivery retries made",
			Type:       "counter",
			LabelNames: []string{"url", "event"},
		},
		{
			Name:       DeliveryCounter,
			Help:       "Count of delivered messages to a url with a status code",
			Type:       "counter",
			LabelNames: []string{"url", "code", "event"},
		},
		{
			Name:       SlowConsumerDroppedMsgCounter,
			Help:       "Count of dropped messages due to a slow consumer",
			Type:       "counter",
			LabelNames: []string{"url", "reason"},
		},
		{
			Name:       SlowConsumerCounter,
			Help:       "Count of the number of times a consumer has been deemed too slow and is cut off.",
			Type:       "counter",
			LabelNames: []string{"url"},
		},
		{
			Name:       OutgoingQueueDepth,
			Help:       "The depth of the queue per outgoing url.",
			Type:       "gauge",
			LabelNames: []string{"url"},
		},
		{
			Name:       IncomingEventTypeCounter,
			Help:       "Incoming count of events by event type",
			Type:       "counter",
			LabelNames: []string{"event"},
		},
	}
}

func CreateOutbounderMetrics(m CaduceusMetricsRegistry, c *CaduceusOutboundSender) {
	c.deliveryCounter = m.NewCounter(DeliveryCounter)
	c.deliveryRetryCounter = m.NewCounter(DeliveryRetryCounter)
	c.cutOffCounter = m.NewCounter(SlowConsumerCounter).With("url", c.id)
	c.droppedQueueFullCounter = m.NewCounter(SlowConsumerDroppedMsgCounter).With("url", c.id, "reason", "queue_full")
	c.droppedExpiredCounter = m.NewCounter(SlowConsumerDroppedMsgCounter).With("url", c.id, "reason", "expired")
	c.droppedCutoffCounter = m.NewCounter(SlowConsumerDroppedMsgCounter).With("url", c.id, "reason", "cut_off")
	c.droppedInvalidConfig = m.NewCounter(SlowConsumerDroppedMsgCounter).With("url", c.id, "reason", "invalid_config")
	c.droppedNetworkErrCounter = m.NewCounter(SlowConsumerDroppedMsgCounter).With("url", c.id, "reason", "network_err")
	c.queueDepthGauge = m.NewGauge(OutgoingQueueDepth).With("url", c.id)
	c.contentTypeCounter = m.NewCounter(IncomingContentTypeCounter)
}
<commit_msg>caduceus: revert to content_type<commit_after>package main

import (
	"github.com/Comcast/webpa-common/xmetrics"
)

const (
	ErrorRequestBodyCounter       = "error_request_body_count"
	EmptyRequestBodyCounter       = "empty_request_body_count"
	DeliveryCounter               = "delivery_count"
	DeliveryRetryCounter          = "delivery_retry_count"
	SlowConsumerDroppedMsgCounter = "slow_consumer_dropped_message_count"
	SlowConsumerCounter           = "slow_consumer_cut_off_count"
	IncomingQueueDepth            = "incoming_queue_depth"
	IncomingContentTypeCounter    = "incoming_content_type_count"
	IncomingEventTypeCounter      = "incoming_event_type_count"
	DropsDueToInvalidPayload      = "drops_due_to_invalid_payload"
	OutgoingQueueDepth            = "outgoing_queue_depths"
)

func Metrics() []xmetrics.Metric {
	return []xmetrics.Metric{
		{
			Name: IncomingQueueDepth,
			Help: "The depth of the queue behind the incoming handlers.",
			Type: "gauge",
		},
		{
			Name: ErrorRequestBodyCounter,
			Help: "Count of the number of errors encountered reading the body.",
			Type: "counter",
		},
		{
			Name: EmptyRequestBodyCounter,
			Help: "Count of the number of times the request is an empty body.",
			Type: "counter",
		},
		{
			Name: DropsDueToInvalidPayload,
			Help: "Dropped messages due to invalid payloads.",
			Type: "counter",
		},
		{
			Name:       IncomingContentTypeCounter,
			Help:       "Count of the content type processed.",
			Type:       "counter",
			LabelNames: []string{"content_type"},
		},
		{
			Name:       DeliveryRetryCounter,
			Help:       "Number of delivery retries made",
			Type:       "counter",
			LabelNames: []string{"url", "event"},
		},
		{
			Name:       DeliveryCounter,
			Help:       "Count of delivered messages to a url with a status code",
			Type:       "counter",
			LabelNames: []string{"url", "code", "event"},
		},
		{
			Name:       SlowConsumerDroppedMsgCounter,
			Help:       "Count of dropped messages due to a slow consumer",
			Type:       "counter",
			LabelNames: []string{"url", "reason"},
		},
		{
			Name:       SlowConsumerCounter,
			Help:       "Count of the number of times a consumer has been deemed too slow and is cut off.",
			Type:       "counter",
			LabelNames: []string{"url"},
		},
		{
			Name:       OutgoingQueueDepth,
			Help:       "The depth of the queue per outgoing url.",
			Type:       "gauge",
			LabelNames: []string{"url"},
		},
		{
			Name:       IncomingEventTypeCounter,
			Help:       "Incoming count of events by event type",
			Type:       "counter",
			LabelNames: []string{"event"},
		},
	}
}

func CreateOutbounderMetrics(m CaduceusMetricsRegistry, c *CaduceusOutboundSender) {
	c.deliveryCounter = m.NewCounter(DeliveryCounter)
	c.deliveryRetryCounter = m.NewCounter(DeliveryRetryCounter)
	c.cutOffCounter = m.NewCounter(SlowConsumerCounter).With("url", c.id)
	c.droppedQueueFullCounter = m.NewCounter(SlowConsumerDroppedMsgCounter).With("url", c.id, "reason", "queue_full")
	c.droppedExpiredCounter = m.NewCounter(SlowConsumerDroppedMsgCounter).With("url", c.id, "reason", "expired")
	c.droppedCutoffCounter = m.NewCounter(SlowConsumerDroppedMsgCounter).With("url", c.id, "reason", "cut_off")
	c.droppedInvalidConfig = m.NewCounter(SlowConsumerDroppedMsgCounter).With("url", c.id, "reason", "invalid_config")
	c.droppedNetworkErrCounter = m.NewCounter(SlowConsumerDroppedMsgCounter).With("url", c.id, "reason", "network_err")
	c.queueDepthGauge = m.NewGauge(OutgoingQueueDepth).With("url", c.id)
	c.contentTypeCounter = m.NewCounter(IncomingContentTypeCounter)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"errors"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"log"
	"os"
	"path/filepath"
	"reflect"
	"strconv"
	"strings"
	"unicode"
	"unicode/utf8"

	"github.com/lestrrat/go-pdebug"
)

var ErrAnnotatedStructNotFound = errors.New("annotated struct was not found")

func snakeCase(s string) string {
	ret := []rune{}
	wasLower := false
	for len(s) > 0 {
		r, n := utf8.DecodeRuneInString(s)
		if r == utf8.RuneError {
			panic("yikes")
		}

		s = s[n:]
		if unicode.IsUpper(r) {
			if wasLower {
				ret = append(ret, '_')
			}
			wasLower = false
		} else {
			wasLower = true
		}

		ret = append(ret, unicode.ToLower(r))
	}
	return string(ret)
}

type Processor struct {
	Types []string
	Dir    string
}

func skipGenerated(fi os.FileInfo) bool {
	switch {
	case strings.HasSuffix(fi.Name(), "gendb.go"):
		return false
	case strings.HasSuffix(fi.Name(), "_gen.go"):
		return false
	}
	return true
}

func (p *Processor) Do() error {
	fset := token.NewFileSet()
	pkgs, err := parser.ParseDir(fset, p.Dir, skipGenerated, parser.ParseComments)
	if err != nil {
		return err
	}

	if len(pkgs) == 0 {
		return errors.New("no packages to process...")
	}

	for _, pkg := range pkgs {
		if err := p.ProcessPkg(pkg); err != nil {
			return err
		}
	}

	return nil
}

func (p *Processor) ShouldProceed(s Struct) bool {
	if len(p.Types) == 0 {
		return true
	}

	for _, t := range p.Types {
		if t == s.Name {
			return true
		}
	}
	return false
}

func (p *Processor) ProcessPkg(pkg *ast.Package) error {
	if pdebug.Enabled {
		g := pdebug.Marker("ProcessPkg %s", pkg.Name)
		defer g.End()
	}
	for fn, f := range pkg.Files {
		pdebug.Printf("Checking file %s", fn)
		for _, s := range p.ExtractStructs(pkg, f) {
			if pdebug.Enabled {
				pdebug.Printf("Checking struct %s", s.Name)
			}
			if !p.ShouldProceed(s) {
				if pdebug.Enabled {
					pdebug.Printf("Skipping struct %s", s.Name)
				}
				continue
			}

			if err := p.ProcessStruct(s); err != nil {
				return err
			}
		}
	}
	return nil
}

func (p *Processor) ProcessStruct(s Struct) error {
	if pdebug.Enabled {
		g := pdebug.Marker("ProcessStruct %s", s.Name)
		defer g.End()
	}

	buf := bytes.Buffer{}

	buf.WriteString("// Automatically generated by gendb utility. DO NOT EDIT!\n")
	buf.WriteString("package ")
	buf.WriteString(s.PackageName)
	buf.WriteString("\n\n")
	buf.WriteString("\nimport (")
	buf.WriteString("\n" + `"errors"`)
	buf.WriteString("\n)")
	varname := unicode.ToLower(rune(s.Name[0]))

	scols := bytes.Buffer{}
	icols := bytes.Buffer{}
	sfields := bytes.Buffer{} // Scan fields
	ifields := bytes.Buffer{} // Scan fields (everything except for OID)
	ipholders := bytes.Buffer{} // Insert place holders (everything except for OID)

	sansOidFields := make([]StructField, 0, len(s.Fields)-1)
	for _, f := range s.Fields {
		if f.Name == "OID" {
			continue
		}
		sansOidFields = append(sansOidFields, f)
	}

	for i, f := range sansOidFields {
		icols.WriteString(f.ColumnName)
		ipholders.WriteRune('?')
		ifields.WriteRune(varname)
		ifields.WriteRune('.')
		ifields.WriteString(f.Name)
		if i < len(sansOidFields) - 1 {
			icols.WriteString(", ")
			ipholders.WriteString(", ")
			ifields.WriteString(", ")
		}
	}

	for i, f := range s.Fields {
		scols.WriteString(f.ColumnName)
		sfields.WriteRune('&')
		sfields.WriteRune(varname)
		sfields.WriteRune('.')
		sfields.WriteString(f.Name)
		if i < len(s.Fields) - 1 {
			scols.WriteString(", ")
			sfields.WriteString(", ")
		}
	}

	fmt.Fprintf(&buf, "\nvar %sTable = %s\n", s.Name, strconv.Quote(s.Tablename))
	fmt.Fprintf(&buf, "\nfunc (%c *%s) Scan(scanner interface { Scan(...interface{}) error }) error {", varname, s.Name)
	fmt.Fprintf(&buf, "\nreturn scanner.Scan(%s)", sfields.String())
	buf.WriteString("\n}\n")

	fmt.Fprintf(&buf, "\nfunc (%c *%s) LoadByEID(tx *Tx, eid string) error {", varname, s.Name)
	fmt.Fprintf(&buf, "\nrow := tx.QueryRow(`SELECT %s FROM ` + SessionTable + ` WHERE eid = ?`, eid)", scols.String())
	fmt.Fprintf(&buf, "\nif err := %c.Scan(row); err != nil {", varname)
	buf.WriteString("\nreturn err")
	buf.WriteString("\n}")
	buf.WriteString("\nreturn nil")
	buf.WriteString("\n}\n")

	fmt.Fprintf(&buf, "\nfunc (%c *%s) Create(tx *Tx) error {", varname, s.Name)
	fmt.Fprintf(&buf, "\nif %c.EID == " + `""` + " {", varname)
	buf.WriteString("\n" + `return errors.New("create: non-empty EID required")`)
	buf.WriteString("\n}\n\n")
	fmt.Fprintf(&buf, "\nresult, err := tx.Exec(`INSERT INTO %s VALUES (%s)`, %s)", icols.String(), ipholders.String(), ifields.String())
	buf.WriteString("\nif err != nil {")
	buf.WriteString("\nreturn err")
	buf.WriteString("\n}\n")
	buf.WriteString("\nlii, err := result.LastInsertId()")
	buf.WriteString("\nif err != nil {")
	buf.WriteString("\nreturn err")
	buf.WriteString("\n}\n")
	fmt.Fprintf(&buf, "\n%c.OID = lii", varname)
	buf.WriteString("\nreturn nil")
	buf.WriteString("\n}\n")

	fsrc, err := format.Source(buf.Bytes())
	if err != nil {
		log.Printf("%s", buf.Bytes())
		return err
	}

	fn := filepath.Join(p.Dir, snakeCase(s.Name)+"_gen.go")
	if pdebug.Enabled {
		pdebug.Printf("Generating file %s", fn)
	}
	fi, err := os.Create(fn)
	if err != nil {
		return err
	}
	defer fi.Close()

	if _, err := fi.Write(fsrc); err != nil {
		return err
	}

	return nil
}

func (p *Processor) ExtractStructs(pkg *ast.Package, f *ast.File) []Struct {
	ctx := &InspectionCtx{
		Marker:  "+DB",
		Package: pkg.Name,
	}

	ast.Inspect(f, ctx.ExtractStructs)
	return ctx.Structs
}

func (c *InspectionCtx) ExtractStructs(n ast.Node) bool {
	var decl *ast.GenDecl
	var ok bool
	var err error

	if decl, ok = n.(*ast.GenDecl); !ok {
		return true
	}

	if decl.Tok != token.TYPE {
		return true
	}

	if err = c.ExtractStructsFromDecl(decl); err != nil {
		return true
	}

	return false
}

func (ctx *InspectionCtx) ExtractStructsFromDecl(decl *ast.GenDecl) error {
	marked := false
	cacheEnabled := true
	noScanner := false
	tablename := ""
	preCreate := ""
	postCreate := ""
	cacheExpires := "1800"
	if cg := decl.Doc; cg != nil {
		for _, c := range cg.List {
			txt := strings.TrimSpace(strings.TrimLeft(c.Text, "//"))
			if strings.HasPrefix(txt, ctx.Marker) {
				marked = true
				tag := reflect.StructTag(strings.TrimPrefix(txt, ctx.Marker))
				tablename = tag.Get("tablename")
				preCreate = tag.Get("pre_create")
				postCreate = tag.Get("post_create")
				if tag.Get("cache") == "false" {
					cacheEnabled = false
				}
				if ce := tag.Get("cache_expires"); ce != "" {
					cacheExpires = ce
				}
				if tag.Get("scanner") == "false" {
					noScanner = true
				}
				break
			}
		}
	}

	if !marked {
		return ErrAnnotatedStructNotFound
	}

	for _, spec := range decl.Specs {
		var t *ast.TypeSpec
		var s *ast.StructType
		var ok bool
		var ident *ast.Ident

		if t, ok = spec.(*ast.TypeSpec); !ok {
			return ErrAnnotatedStructNotFound
		}

		if s, ok = t.Type.(*ast.StructType); !ok {
			return ErrAnnotatedStructNotFound
		}

		if tablename == "" {
			tablename = fmt.Sprintf("%s_%s",
				ctx.Package,
				snakeCase(t.Name.Name),
			)
		}

		st := Struct{
			PackageName:  ctx.Package,
			CacheEnabled: cacheEnabled,
			CacheExpires: cacheExpires,
			Fields:       make([]StructField, 0, len(s.Fields.List)),
			Name:         t.Name.Name,
			NoScanner:    noScanner,
			PreCreate:    preCreate,
			PostCreate:   postCreate,
			Tablename:    tablename,
		}

	LoopFields:
		for _, f := range s.Fields.List {
			autoIncrement := false
			primaryKey := false
			unique := false
			converter := ""
			columnName := ""

			if f.Tag != nil {
				v := f.Tag.Value
				if len(v) >= 2 {
					if v[0] == '`' {
						v = v[1:]
					}
					if v[len(v)-1] == '`' {
						v = v[:len(v)-1]
					}
				}

				tags := reflect.StructTag(v).Get("db")
				for _, tag := range strings.Split(tags, ",") {
					switch tag {
					case "ignore":
						continue LoopFields
					case "auto_increment":
						autoIncrement = true
					case "pk":
						primaryKey = true
					case "unique":
						unique = true
					default:
						switch {
						case strings.HasPrefix(tag, "converter="):
							converter = tag[10:]
						case strings.HasPrefix(tag, "column="):
							columnName = tag[7:]
						}
					}
				}
			}

			if columnName == "" {
				columnName = snakeCase(f.Names[0].Name)
			}

			switch f.Type.(type) {
			case *ast.Ident:
				ident = f.Type.(*ast.Ident)
			case *ast.SelectorExpr:
				ident = f.Type.(*ast.SelectorExpr).Sel
			case *ast.StarExpr:
				ident = f.Type.(*ast.StarExpr).X.(*ast.SelectorExpr).Sel
			default:
				fmt.Printf("%#v\n", f.Type)
				panic("field type not supported")
			}

			field := StructField{
				AutoIncrement: autoIncrement,
				Converter:     converter,
				ColumnName:    columnName,
				Name:          f.Names[0].Name,
				Type:          ident.Name,
				Unique:        unique,
			}

			if autoIncrement {
				st.AutoIncrementField = &field
			}

			if primaryKey {
				st.PrimaryKey = &field
			}

			// if auto_increment, this key does not need to
			// be included in the st.Fields list
			st.Fields = append(st.Fields, field)
		}
		ctx.Structs = append(ctx.Structs, st)
	}

	return nil
}
<commit_msg>fix generated methods<commit_after>package main

import (
	"bytes"
	"errors"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"log"
	"os"
	"path/filepath"
	"reflect"
	"strconv"
	"strings"
	"unicode"
	"unicode/utf8"

	"github.com/lestrrat/go-pdebug"
)

var ErrAnnotatedStructNotFound = errors.New("annotated struct was not found")

func snakeCase(s string) string {
	ret := []rune{}
	wasLower := false
	for len(s) > 0 {
		r, n := utf8.DecodeRuneInString(s)
		if r == utf8.RuneError {
			panic("yikes")
		}

		s = s[n:]
		if unicode.IsUpper(r) {
			if wasLower {
				ret = append(ret, '_')
			}
			wasLower = false
		} else {
			wasLower = true
		}

		ret = append(ret, unicode.ToLower(r))
	}
	return string(ret)
}

type Processor struct {
	Types []string
	Dir    string
}

func skipGenerated(fi os.FileInfo) bool {
	switch {
	case strings.HasSuffix(fi.Name(), "gendb.go"):
		return false
	case strings.HasSuffix(fi.Name(), "_gen.go"):
		return false
	}
	return true
}

func (p *Processor) Do() error {
	fset := token.NewFileSet()
	pkgs, err := parser.ParseDir(fset, p.Dir, skipGenerated, parser.ParseComments)
	if err != nil {
		return err
	}

	if len(pkgs) == 0 {
		return errors.New("no packages to process...")
	}

	for _, pkg := range pkgs {
		if err := p.ProcessPkg(pkg); err != nil {
			return err
		}
	}

	return nil
}

func (p *Processor) ShouldProceed(s Struct) bool {
	if len(p.Types) == 0 {
		return true
	}

	for _, t := range p.Types {
		if t == s.Name {
			return true
		}
	}
	return false
}

func (p *Processor) ProcessPkg(pkg *ast.Package) error {
	if pdebug.Enabled {
		g := pdebug.Marker("ProcessPkg %s", pkg.Name)
		defer g.End()
	}
	for fn, f := range pkg.Files {
		pdebug.Printf("Checking file %s", fn)
		for _, s := range p.ExtractStructs(pkg, f) {
			if pdebug.Enabled {
				pdebug.Printf("Checking struct %s", s.Name)
			}
			if !p.ShouldProceed(s) {
				if pdebug.Enabled {
					pdebug.Printf("Skipping struct %s", s.Name)
				}
				continue
			}

			if err := p.ProcessStruct(s); err != nil {
				return err
			}
		}
	}
	return nil
}

func (p *Processor) ProcessStruct(s Struct) error {
	if pdebug.Enabled {
		g := pdebug.Marker("ProcessStruct %s", s.Name)
		defer g.End()
	}

	buf := bytes.Buffer{}

	buf.WriteString("// Automatically generated by gendb utility. DO NOT EDIT!\n")
	buf.WriteString("package ")
	buf.WriteString(s.PackageName)
	buf.WriteString("\n\n")
	buf.WriteString("\nimport (")
	buf.WriteString("\n" + `"errors"`)
	buf.WriteString("\n)")
	varname := unicode.ToLower(rune(s.Name[0]))

	scols := bytes.Buffer{}
	icols := bytes.Buffer{}
	sfields := bytes.Buffer{} // Scan fields
	ifields := bytes.Buffer{} // Scan fields (everything except for OID)
	ipholders := bytes.Buffer{} // Insert place holders (everything except for OID)

	sansOidFields := make([]StructField, 0, len(s.Fields)-1)
	for _, f := range s.Fields {
		if f.Name == "OID" {
			continue
		}
		sansOidFields = append(sansOidFields, f)
	}

	for i, f := range sansOidFields {
		icols.WriteString(f.ColumnName)
		ipholders.WriteRune('?')
		ifields.WriteRune(varname)
		ifields.WriteRune('.')
		ifields.WriteString(f.Name)
		if i < len(sansOidFields) - 1 {
			icols.WriteString(", ")
			ipholders.WriteString(", ")
			ifields.WriteString(", ")
		}
	}

	for i, f := range s.Fields {
		scols.WriteString(f.ColumnName)
		sfields.WriteRune('&')
		sfields.WriteRune(varname)
		sfields.WriteRune('.')
		sfields.WriteString(f.Name)
		if i < len(s.Fields) - 1 {
			scols.WriteString(", ")
			sfields.WriteString(", ")
		}
	}

	fmt.Fprintf(&buf, "\nvar %sTable = %s\n", s.Name, strconv.Quote(s.Tablename))
	fmt.Fprintf(&buf, "\nfunc (%c *%s) Scan(scanner interface { Scan(...interface{}) error }) error {", varname, s.Name)
	fmt.Fprintf(&buf, "\nreturn scanner.Scan(%s)", sfields.String())
	buf.WriteString("\n}\n")

	fmt.Fprintf(&buf, "\nfunc (%c *%s) LoadByEID(tx *Tx, eid string) error {", varname, s.Name)
	fmt.Fprintf(&buf, "\nrow := tx.QueryRow(`SELECT %s FROM ` + %sTable + ` WHERE eid = ?`, eid)", scols.String(), s.Name)
	fmt.Fprintf(&buf, "\nif err := %c.Scan(row); err != nil {", varname)
	buf.WriteString("\nreturn err")
	buf.WriteString("\n}")
	buf.WriteString("\nreturn nil")
	buf.WriteString("\n}\n")

	fmt.Fprintf(&buf, "\nfunc (%c *%s) Create(tx *Tx) error {", varname, s.Name)
	fmt.Fprintf(&buf, "\nif %c.EID == " + `""` + " {", varname)
	buf.WriteString("\n" + `return errors.New("create: non-empty EID required")`)
	buf.WriteString("\n}\n\n")
	fmt.Fprintf(&buf, "\nresult, err := tx.Exec(`INSERT INTO ` + %sTable + ` (%s) VALUES (%s)`, %s)", s.Name, icols.String(), ipholders.String(), ifields.String())
	buf.WriteString("\nif err != nil {")
	buf.WriteString("\nreturn err")
	buf.WriteString("\n}\n")
	buf.WriteString("\nlii, err := result.LastInsertId()")
	buf.WriteString("\nif err != nil {")
	buf.WriteString("\nreturn err")
	buf.WriteString("\n}\n")
	fmt.Fprintf(&buf, "\n%c.OID = lii", varname)
	buf.WriteString("\nreturn nil")
	buf.WriteString("\n}\n")

	fmt.Fprintf(&buf, "\nfunc (%c %s) Delete(tx *Tx) error {", varname, s.Name)
	fmt.Fprintf(&buf, "\nif %c.OID != 0 {", varname)
	fmt.Fprintf(&buf, "\n_, err := tx.Exec(`DELETE FROM ` + %sTable + `WHERE oid = ?`, %c.OID)", s.Name, varname)
	buf.WriteString("\nreturn err")
	buf.WriteString("\n}\n")
	fmt.Fprintf(&buf, "\nif %c.EID != %s {", varname, `""`)
	fmt.Fprintf(&buf, "\n_, err := tx.Exec(`DELETE FROM ` + %sTable + `WHERE eid = ?`, %c.EID)", s.Name, varname)
	buf.WriteString("\nreturn err")
	buf.WriteString("\n}\n")
	fmt.Fprintf(&buf, "\nreturn errors.New(%s)", strconv.Quote("either OID/EID mustbe filled"))
	buf.WriteString("\n}\n")

	fsrc, err := format.Source(buf.Bytes())
	if err != nil {
		log.Printf("%s", buf.Bytes())
		return err
	}

	fn := filepath.Join(p.Dir, snakeCase(s.Name)+"_gen.go")
	if pdebug.Enabled {
		pdebug.Printf("Generating file %s", fn)
	}
	fi, err := os.Create(fn)
	if err != nil {
		return err
	}
	defer fi.Close()

	if _, err := fi.Write(fsrc); err != nil {
		return err
	}

	return nil
}

func (p *Processor) ExtractStructs(pkg *ast.Package, f *ast.File) []Struct {
	ctx := &InspectionCtx{
		Marker:  "+DB",
		Package: pkg.Name,
	}

	ast.Inspect(f, ctx.ExtractStructs)
	return ctx.Structs
}

func (c *InspectionCtx) ExtractStructs(n ast.Node) bool {
	var decl *ast.GenDecl
	var ok bool
	var err error

	if decl, ok = n.(*ast.GenDecl); !ok {
		return true
	}

	if decl.Tok != token.TYPE {
		return true
	}

	if err = c.ExtractStructsFromDecl(decl); err != nil {
		return true
	}

	return false
}

func (ctx *InspectionCtx) ExtractStructsFromDecl(decl *ast.GenDecl) error {
	marked := false
	cacheEnabled := true
	noScanner := false
	tablename := ""
	preCreate := ""
	postCreate := ""
	cacheExpires := "1800"
	if cg := decl.Doc; cg != nil {
		for _, c := range cg.List {
			txt := strings.TrimSpace(strings.TrimLeft(c.Text, "//"))
			if strings.HasPrefix(txt, ctx.Marker) {
				marked = true
				tag := reflect.StructTag(strings.TrimPrefix(txt, ctx.Marker))
				tablename = tag.Get("tablename")
				preCreate = tag.Get("pre_create")
				postCreate = tag.Get("post_create")
				if tag.Get("cache") == "false" {
					cacheEnabled = false
				}
				if ce := tag.Get("cache_expires"); ce != "" {
					cacheExpires = ce
				}
				if tag.Get("scanner") == "false" {
					noScanner = true
				}
				break
			}
		}
	}

	if !marked {
		return ErrAnnotatedStructNotFound
	}

	for _, spec := range decl.Specs {
		var t *ast.TypeSpec
		var s *ast.StructType
		var ok bool
		var ident *ast.Ident

		if t, ok = spec.(*ast.TypeSpec); !ok {
			return ErrAnnotatedStructNotFound
		}

		if s, ok = t.Type.(*ast.StructType); !ok {
			return ErrAnnotatedStructNotFound
		}

		if tablename == "" {
			tablename = fmt.Sprintf("%s_%s",
				ctx.Package,
				snakeCase(t.Name.Name),
			)
		}

		st := Struct{
			PackageName:  ctx.Package,
			CacheEnabled: cacheEnabled,
			CacheExpires: cacheExpires,
			Fields:       make([]StructField, 0, len(s.Fields.List)),
			Name:         t.Name.Name,
			NoScanner:    noScanner,
			PreCreate:    preCreate,
			PostCreate:   postCreate,
			Tablename:    tablename,
		}

	LoopFields:
		for _, f := range s.Fields.List {
			autoIncrement := false
			primaryKey := false
			unique := false
			converter := ""
			columnName := ""

			if f.Tag != nil {
				v := f.Tag.Value
				if len(v) >= 2 {
					if v[0] == '`' {
						v = v[1:]
					}
					if v[len(v)-1] == '`' {
						v = v[:len(v)-1]
					}
				}

				tags := reflect.StructTag(v).Get("db")
				for _, tag := range strings.Split(tags, ",") {
					switch tag {
					case "ignore":
						continue LoopFields
					case "auto_increment":
						autoIncrement = true
					case "pk":
						primaryKey = true
					case "unique":
						unique = true
					default:
						switch {
						case strings.HasPrefix(tag, "converter="):
							converter = tag[10:]
						case strings.HasPrefix(tag, "column="):
							columnName = tag[7:]
						}
					}
				}
			}

			if columnName == "" {
				columnName = snakeCase(f.Names[0].Name)
			}

			switch f.Type.(type) {
			case *ast.Ident:
				ident = f.Type.(*ast.Ident)
			case *ast.SelectorExpr:
				ident = f.Type.(*ast.SelectorExpr).Sel
			case *ast.StarExpr:
				ident = f.Type.(*ast.StarExpr).X.(*ast.SelectorExpr).Sel
			default:
				fmt.Printf("%#v\n", f.Type)
				panic("field type not supported")
			}

			field := StructField{
				AutoIncrement: autoIncrement,
				Converter:     converter,
				ColumnName:    columnName,
				Name:          f.Names[0].Name,
				Type:          ident.Name,
				Unique:        unique,
			}

			if autoIncrement {
				st.AutoIncrementField = &field
			}

			if primaryKey {
				st.PrimaryKey = &field
			}

			// if auto_increment, this key does not need to
			// be included in the st.Fields list
			st.Fields = append(st.Fields, field)
		}
		ctx.Structs = append(ctx.Structs, st)
	}

	return nil
}
<|endoftext|>


<commit_before>// Package service contains all the services necessary for interacting
// with all the collections(tables) in the database.
//
// Each service file should be written in it's own file and should represent only one dbmodels
package service

import (
	"errors"
	"gost/config"
	"log"

	"gopkg.in/mgo.v2"
)

var mongoDbSession *mgo.Session

var (
	// ErrNoIDSpecified says that no ID was specified for a fetch operation
	ErrNoIDSpecified = errors.New("No Id was specified for the entity")
)

// InitDbService initializes the connection (known as session) parameters to the database
func InitDbService() {
	if config.DbConnectionString == "" {
		log.Fatal("Database error: No connection string provided")
	}

	if mongoDbSession == nil {
		var err error
		mongoDbSession, err = mgo.Dial(config.DbConnectionString)
		mongoDbSession.SetMode(mgo.Monotonic, true)

		if err != nil {
			log.Fatalf("Can't connect to mongo, go error: %v\n", err)
		}
	}
}

// CloseDbService closes the current mongodb session
func CloseDbService() {
	mongoDbSession.Close()
}

// Connect creates the connection to the database.
// To avoid forgetting to close the session, always write: *defer session.Close()* right after getting it through this method
func Connect(collectionName string) (*mgo.Session, *mgo.Collection) {
	return mongoDbSession.Copy(), mongoDbSession.DB(config.DbName).C(collectionName)
}
<commit_msg>fixed session desync for when initializing a new connection to mongo db<commit_after>// Package service contains all the services necessary for interacting
// with all the collections(tables) in the database.
//
// Each service file should be written in it's own file and should represent only one dbmodels
package service

import (
	"errors"
	"gost/config"
	"log"

	"gopkg.in/mgo.v2"
)

var mongoDbSession *mgo.Session

var (
	// ErrNoIDSpecified says that no ID was specified for a fetch operation
	ErrNoIDSpecified = errors.New("No Id was specified for the entity")
)

// InitDbService initializes the connection (known as session) parameters to the database
func InitDbService() {
	if config.DbConnectionString == "" {
		log.Fatal("Database error: No connection string provided")
	}

	if mongoDbSession == nil {
		var err error
		mongoDbSession, err = mgo.Dial(config.DbConnectionString)
		mongoDbSession.SetMode(mgo.Monotonic, true)

		if err != nil {
			log.Fatalf("Can't connect to mongo, go error: %v\n", err)
		}
	}
}

// CloseDbService closes the current mongodb session
func CloseDbService() {
	mongoDbSession.Close()
}

// Connect creates the connection to the database.
// To avoid forgetting to close the session, always write: *defer session.Close()* right after getting it through this method
func Connect(collectionName string) (*mgo.Session, *mgo.Collection) {
	var session = mongoDbSession.Copy()

	return session, session.DB(config.DbName).C(collectionName)
}
<|endoftext|>


<commit_before>package hyperone

import (
	"context"
	"fmt"
	"strings"

	"github.com/hashicorp/packer/helper/multistep"
	"github.com/hashicorp/packer/packer"
	openapi "github.com/hyperonecom/h1-client-go"
)

type stepCreateVM struct {
	vmID string
}

const (
	chrootDiskName = "packer-chroot-disk"
)

func (s *stepCreateVM) Run(ctx context.Context, state multistep.StateBag) multistep.StepAction {
	client := state.Get("client").(*openapi.APIClient)
	ui := state.Get("ui").(packer.Ui)
	config := state.Get("config").(*Config)
	sshKey := state.Get("ssh_public_key").(string)

	ui.Say("Creating VM...")

	netAdapter := pickNetAdapter(config)

	var sshKeys = []string{sshKey}
	sshKeys = append(sshKeys, config.SSHKeys...)

	disks := []openapi.VmCreateDisk{
		{
			Service: config.DiskType,
			Size:    config.DiskSize,
		},
	}

	if config.ChrootDisk {
		disks = append(disks, openapi.VmCreateDisk{
			Service: config.ChrootDiskType,
			Size:    config.ChrootDiskSize,
			Name:    chrootDiskName,
		})
	}

	options := openapi.VmCreate{
		Name:         config.VmName,
		Image:        config.SourceImage,
		Service:      config.VmType,
		SshKeys:      sshKeys,
		Disk:         disks,
		Netadp:       []openapi.VmCreateNetadp{netAdapter},
		UserMetadata: config.UserData,
		Tag:          config.VmTags,
	}

	vm, _, err := client.VmApi.VmCreate(ctx, options)
	if err != nil {
		err := fmt.Errorf("error creating VM: %s", formatOpenAPIError(err))
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	s.vmID = vm.Id
	state.Put("vm_id", vm.Id)
	// instance_id is the generic term used so that users can have access to the
	// instance id inside of the provisioners, used in step_provision.
	state.Put("instance_id", vm.Id)

	hdds, _, err := client.VmApi.VmListHdd(ctx, vm.Id)
	if err != nil {
		err := fmt.Errorf("error listing hdd: %s", formatOpenAPIError(err))
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	for _, hdd := range hdds {
		if hdd.Disk.Name == chrootDiskName {
			state.Put("chroot_disk_id", hdd.Disk.Id)
			controllerNumber := strings.ToLower(strings.Trim(hdd.ControllerNumber, "{}"))
			state.Put("chroot_controller_number", controllerNumber)
			state.Put("chroot_controller_location", int(hdd.ControllerLocation))
			break
		}
	}

	netadp, _, err := client.VmApi.VmListNetadp(ctx, vm.Id)
	if err != nil {
		err := fmt.Errorf("error listing netadp: %s", formatOpenAPIError(err))
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	if len(netadp) < 1 {
		err := fmt.Errorf("no network adapters found")
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	publicIP, err := associatePublicIP(ctx, config, client, netadp[0])
	if err != nil {
		err := fmt.Errorf("error associating IP: %s", formatOpenAPIError(err))
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	state.Put("public_ip", publicIP)

	return multistep.ActionContinue
}

func pickNetAdapter(config *Config) openapi.VmCreateNetadp {
	if config.Network == "" {
		if config.PublicIP != "" {
			return openapi.VmCreateNetadp{
				Service: config.PublicNetAdpService,
				Ip:      []string{config.PublicIP},
			}
		}
	} else {
		var privateIPs []string

		if config.PrivateIP == "" {
			privateIPs = nil
		} else {
			privateIPs = []string{config.PrivateIP}
		}

		return openapi.VmCreateNetadp{
			Service: "private",
			Network: config.Network,
			Ip:      privateIPs,
		}
	}

	return openapi.VmCreateNetadp{
		Service: config.PublicNetAdpService,
	}
}

func associatePublicIP(ctx context.Context, config *Config, client *openapi.APIClient, netadp openapi.Netadp) (string, error) {
	if config.Network == "" || config.PublicIP == "" {
		// Public IP belongs to attached net adapter
		return netadp.Ip[0].Address, nil
	}

	var privateIP string
	if config.PrivateIP == "" {
		privateIP = netadp.Ip[0].Id
	} else {
		privateIP = config.PrivateIP
	}

	ip, _, err := client.IpApi.IpActionAssociate(ctx, config.PublicIP, openapi.IpActionAssociate{Ip: privateIP})
	if err != nil {
		return "", err
	}

	return ip.Address, nil
}

func (s *stepCreateVM) Cleanup(state multistep.StateBag) {
	if s.vmID == "" {
		return
	}

	ui := state.Get("ui").(packer.Ui)

	ui.Say(fmt.Sprintf("Deleting VM %s...", s.vmID))
	err := deleteVMWithDisks(s.vmID, state)
	if err != nil {
		ui.Error(err.Error())
	}
}

func deleteVMWithDisks(vmID string, state multistep.StateBag) error {
	client := state.Get("client").(*openapi.APIClient)
	hdds, _, err := client.VmApi.VmListHdd(context.TODO(), vmID)
	if err != nil {
		return fmt.Errorf("error listing hdd: %s", formatOpenAPIError(err))
	}

	deleteOptions := openapi.VmDelete{}
	for _, hdd := range hdds {
		deleteOptions.RemoveDisks = append(deleteOptions.RemoveDisks, hdd.Disk.Id)
	}

	_, err = client.VmApi.VmDelete(context.TODO(), vmID, deleteOptions)
	if err != nil {
		return fmt.Errorf("Error deleting server '%s' - please delete it manually: %s", vmID, formatOpenAPIError(err))
	}

	return nil
}
<commit_msg>builder/hyperone: Add support for custom username in vm create<commit_after>package hyperone

import (
	"context"
	"fmt"
	"strings"

	"github.com/hashicorp/packer/helper/multistep"
	"github.com/hashicorp/packer/packer"
	openapi "github.com/hyperonecom/h1-client-go"
)

type stepCreateVM struct {
	vmID string
}

const (
	chrootDiskName = "packer-chroot-disk"
)

func (s *stepCreateVM) Run(ctx context.Context, state multistep.StateBag) multistep.StepAction {
	client := state.Get("client").(*openapi.APIClient)
	ui := state.Get("ui").(packer.Ui)
	config := state.Get("config").(*Config)
	sshKey := state.Get("ssh_public_key").(string)

	ui.Say("Creating VM...")

	netAdapter := pickNetAdapter(config)

	var sshKeys = []string{sshKey}
	sshKeys = append(sshKeys, config.SSHKeys...)

	disks := []openapi.VmCreateDisk{
		{
			Service: config.DiskType,
			Size:    config.DiskSize,
		},
	}

	if config.ChrootDisk {
		disks = append(disks, openapi.VmCreateDisk{
			Service: config.ChrootDiskType,
			Size:    config.ChrootDiskSize,
			Name:    chrootDiskName,
		})
	}

	options := openapi.VmCreate{
		Name:         config.VmName,
		Image:        config.SourceImage,
		Service:      config.VmType,
		SshKeys:      sshKeys,
		Disk:         disks,
		Netadp:       []openapi.VmCreateNetadp{netAdapter},
		UserMetadata: config.UserData,
		Tag:          config.VmTags,
		Username:     config.Comm.SSHUsername,
	}

	vm, _, err := client.VmApi.VmCreate(ctx, options)
	if err != nil {
		err := fmt.Errorf("error creating VM: %s", formatOpenAPIError(err))
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	s.vmID = vm.Id
	state.Put("vm_id", vm.Id)
	// instance_id is the generic term used so that users can have access to the
	// instance id inside of the provisioners, used in step_provision.
	state.Put("instance_id", vm.Id)

	hdds, _, err := client.VmApi.VmListHdd(ctx, vm.Id)
	if err != nil {
		err := fmt.Errorf("error listing hdd: %s", formatOpenAPIError(err))
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	for _, hdd := range hdds {
		if hdd.Disk.Name == chrootDiskName {
			state.Put("chroot_disk_id", hdd.Disk.Id)
			controllerNumber := strings.ToLower(strings.Trim(hdd.ControllerNumber, "{}"))
			state.Put("chroot_controller_number", controllerNumber)
			state.Put("chroot_controller_location", int(hdd.ControllerLocation))
			break
		}
	}

	netadp, _, err := client.VmApi.VmListNetadp(ctx, vm.Id)
	if err != nil {
		err := fmt.Errorf("error listing netadp: %s", formatOpenAPIError(err))
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	if len(netadp) < 1 {
		err := fmt.Errorf("no network adapters found")
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	publicIP, err := associatePublicIP(ctx, config, client, netadp[0])
	if err != nil {
		err := fmt.Errorf("error associating IP: %s", formatOpenAPIError(err))
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	state.Put("public_ip", publicIP)

	return multistep.ActionContinue
}

func pickNetAdapter(config *Config) openapi.VmCreateNetadp {
	if config.Network == "" {
		if config.PublicIP != "" {
			return openapi.VmCreateNetadp{
				Service: config.PublicNetAdpService,
				Ip:      []string{config.PublicIP},
			}
		}
	} else {
		var privateIPs []string

		if config.PrivateIP == "" {
			privateIPs = nil
		} else {
			privateIPs = []string{config.PrivateIP}
		}

		return openapi.VmCreateNetadp{
			Service: "private",
			Network: config.Network,
			Ip:      privateIPs,
		}
	}

	return openapi.VmCreateNetadp{
		Service: config.PublicNetAdpService,
	}
}

func associatePublicIP(ctx context.Context, config *Config, client *openapi.APIClient, netadp openapi.Netadp) (string, error) {
	if config.Network == "" || config.PublicIP == "" {
		// Public IP belongs to attached net adapter
		return netadp.Ip[0].Address, nil
	}

	var privateIP string
	if config.PrivateIP == "" {
		privateIP = netadp.Ip[0].Id
	} else {
		privateIP = config.PrivateIP
	}

	ip, _, err := client.IpApi.IpActionAssociate(ctx, config.PublicIP, openapi.IpActionAssociate{Ip: privateIP})
	if err != nil {
		return "", err
	}

	return ip.Address, nil
}

func (s *stepCreateVM) Cleanup(state multistep.StateBag) {
	if s.vmID == "" {
		return
	}

	ui := state.Get("ui").(packer.Ui)

	ui.Say(fmt.Sprintf("Deleting VM %s...", s.vmID))
	err := deleteVMWithDisks(s.vmID, state)
	if err != nil {
		ui.Error(err.Error())
	}
}

func deleteVMWithDisks(vmID string, state multistep.StateBag) error {
	client := state.Get("client").(*openapi.APIClient)
	hdds, _, err := client.VmApi.VmListHdd(context.TODO(), vmID)
	if err != nil {
		return fmt.Errorf("error listing hdd: %s", formatOpenAPIError(err))
	}

	deleteOptions := openapi.VmDelete{}
	for _, hdd := range hdds {
		deleteOptions.RemoveDisks = append(deleteOptions.RemoveDisks, hdd.Disk.Id)
	}

	_, err = client.VmApi.VmDelete(context.TODO(), vmID, deleteOptions)
	if err != nil {
		return fmt.Errorf("Error deleting server '%s' - please delete it manually: %s", vmID, formatOpenAPIError(err))
	}

	return nil
}
<|endoftext|>


<commit_before>// The sentinel package provides a convenient interface with a redis sentinel
// which will automatically handle pooling connections and automatic failover.
//
// This package only gaurantees that when GetMaster is called the returned
// connection will be a connection to the master as of the moment that method is
// called. It is still possible that there is a failover as that connection is
// being used by the application, or before the PutMaster. Because of this,
// always check errors and never PutMaster on a connection which has returned an
// error.
//
// As a final note, a Client can be interacted with from multiple routines at
// once safely, except for the Close method. To safely Close, ensure that only
// one routine ever makes the call and that once the call is made no other
// methods are ever called by any routines.
package sentinel

import (
	"errors"
	"github.com/fzzy/radix/redis"
	"strings"

	"github.com/fzzy/radix/extra/pool"
	"github.com/fzzy/radix/extra/pubsub"
)

// An error wrapper returned by operations in this package. It implements the
// error interface and can therefore be passed around as a normal error.
type ClientError struct {
	err error

	// If this is true the error is due to a problem with the sentinel
	// connection, either it being closed or otherwise unavailable. If false the
	// error is due to some other circumstances. This is useful if you want to
	// implement some kind of reconnecting to sentinel on an error.
	SentinelErr bool
}

func (ce *ClientError) Error() string {
	return ce.err.Error()
}

type getReqRet struct {
	conn *redis.Client
	err  *ClientError
}

type getReq struct {
	name  string
	retCh chan *getReqRet
}

type putReq struct {
	name string
	conn *redis.Client
}

type switchMaster struct {
	name string
	addr string
}

type Client struct {
	poolSize    int
	masterPools map[string]*pool.Pool
	subClient   *pubsub.SubClient

	getCh   chan *getReq
	putCh   chan *putReq
	closeCh chan struct{}

	alwaysErr      *ClientError
	alwaysErrCh    chan *ClientError
	switchMasterCh chan *switchMaster
}

// Creates a sentinel client. Connects to the given sentinel instance, pulls the
// information for the masters of the given names, and creates an intial pool of
// connections for each master. The client will automatically replace the pool
// for any master should sentinel decide to fail the master over. The returned
// error is a *ClientError.
func NewClient(
	network, address string, poolSize int, names ...string,
) (
	*Client, error,
) {

	// We use this to fetch initial details about masters before we upgrade it
	// to a pubsub client
	client, err := redis.Dial(network, address)
	if err != nil {
		return nil, &ClientError{err: err}
	}

	masterPools := map[string]*pool.Pool{}
	for _, name := range names {
		r := client.Cmd("SENTINEL", "MASTER", name)
		l, err := r.List()
		if err != nil {
			return nil, &ClientError{err: err, SentinelErr: true}
		}
		addr := l[3] + ":" + l[5]
		pool, err := pool.NewPool("tcp", addr, poolSize)
		if err != nil {
			return nil, &ClientError{err: err}
		}
		masterPools[name] = pool
	}

	subClient := pubsub.NewSubClient(client)
	r := subClient.Subscribe("+switch-master")
	if r.Err != nil {
		return nil, &ClientError{err: r.Err, SentinelErr: true}
	}

	c := &Client{
		poolSize:       poolSize,
		masterPools:    masterPools,
		subClient:      subClient,
		getCh:          make(chan *getReq),
		putCh:          make(chan *putReq),
		closeCh:        make(chan struct{}),
		alwaysErrCh:    make(chan *ClientError),
		switchMasterCh: make(chan *switchMaster),
	}

	go c.subSpin()
	go c.spin()
	return c, nil
}

func (c *Client) subSpin() {
	for {
		r := c.subClient.Receive()
		if r.Timeout() {
			continue
		}
		if r.Err != nil {
			select {
			case c.alwaysErrCh <- &ClientError{err: r.Err, SentinelErr: true}:
			case <-c.closeCh:
			}
			return
		}
		sMsg := strings.Split(r.Message, " ")
		name := sMsg[0]
		newAddr := sMsg[3] + ":" + sMsg[4]
		select {
		case c.switchMasterCh <- &switchMaster{name, newAddr}:
		case <-c.closeCh:
			return
		}
	}
}

func (c *Client) spin() {
	for {
		select {
		case req := <-c.getCh:
			if c.alwaysErr != nil {
				req.retCh <- &getReqRet{nil, c.alwaysErr}
				continue
			}
			pool, ok := c.masterPools[req.name]
			if !ok {
				err := errors.New("unknown name: " + req.name)
				req.retCh <- &getReqRet{nil, &ClientError{err: err}}
				continue
			}
			conn, err := pool.Get()
			if err != nil {
				req.retCh <- &getReqRet{nil, &ClientError{err: err}}
				continue
			}
			req.retCh <- &getReqRet{conn, nil}

		case req := <-c.putCh:
			if pool, ok := c.masterPools[req.name]; ok {
				pool.Put(req.conn)
			}

		case err := <-c.alwaysErrCh:
			c.alwaysErr = err

		case sm := <-c.switchMasterCh:
			if p, ok := c.masterPools[sm.name]; ok {
				p.Empty()
				p = pool.NewOrEmptyPool("tcp", sm.addr, c.poolSize)
				c.masterPools[sm.name] = p
			}

		case <-c.closeCh:
			for name := range c.masterPools {
				c.masterPools[name].Empty()
			}
			c.subClient.Client.Close()
			close(c.getCh)
			close(c.putCh)
			return
		}
	}
}

// Retrieves a connection for the master of the given name. If sentinel has
// become unreachable this will always return an error. Close should be called
// in that case. The returned error is a *ClientError.
func (c *Client) GetMaster(name string) (*redis.Client, error) {
	req := getReq{name, make(chan *getReqRet)}
	c.getCh <- &req
	ret := <-req.retCh
	if ret.err != nil {
		return nil, ret.err
	}
	return ret.conn, nil
}

// Return a connection for a master of a given name. As with the pool package,
// do not return a connection which is having connectivity issues, or which is
// otherwise unable to perform requests.
func (c *Client) PutMaster(name string, client *redis.Client) {
	c.putCh <- &putReq{name, client}
}

// A useful helper method, analagous to the pool package's CarefullyPut method.
// Since we don't want to Put a connection which is having connectivity
// issues, this can be defered inside a function to make sure we only put back a
// connection when we should. It should be used like the following:
//
//	func doSomeThings(c *Client) error {
//		conn, redisErr := c.GetMaster("bucket0")
//		if redisErr != nil {
//			return redisErr
//		}
//		defer c.CarefullyPutMaster("bucket0", conn, &redisErr)
//
//		var i int
//		i, redisErr = conn.Cmd("GET", "foo").Int()
//		if redisErr != nil {
//			return redisErr
//		}
//
//		redisErr = conn.Cmd("SET", "foo", i * 3).Err
//		return redisErr
//	}
func (c *Client) CarefullyPutMaster(
	name string, client *redis.Client, potentialErr *error,
) {
	if potentialErr != nil && *potentialErr != nil {
		// If the client sent back that it's READONLY then we don't want to keep
		// this connection around. Otherwise, we don't care about command errors
		if cerr, ok := (*potentialErr).(*redis.CmdError); !ok || cerr.Readonly() {
			client.Close()
			return
		}
	}
	c.PutMaster(name, client)
}

// Closes all connection pools as well as the connection to sentinel.
func (c *Client) Close() {
	close(c.closeCh)
}
<commit_msg>update sentinel documentation<commit_after>// The sentinel package provides a convenient interface with a redis sentinel
// which will automatically handle pooling connections and automatic failover.
//
// Here's an example of creating a sentinel client and then using it to perform
// some commands
//
//	func example() error {
//		// If there exists sentinel masters "bucket0" and "bucket1", and we want out
//		// client to create pools for both:
//		client, err := sentinel.NewClient("tcp", "localhost:6379", 100, "bucket0", "bucket1")
//		if err != nil {
//			return err
//		}
//
//		if err := exampleCmd(client); err != nil {
//			return err
//		}
//
//		return nil
//	}
//
//	func exampleCmd(client *sentinel.Client) error {
//		conn, redisErr := client.GetMaster("bucket0")
//		if redisErr != nil {
//			return redisErr
//		}
//		// We use CarefullyPutMaster to conditionally put the connection back in the
//		// pool depending on the last error seen
//		defer client.CarefullyPutMaster("bucket0", conn, &redisErr)
//
//		var i int
//		if i, redisErr = conn.Cmd("GET", "foo").Int(); redisErr != nil {
//			return redisErr
//		}
//
//		if redisErr = conn.Cmd("SET", "foo", i+1); redisErr != nil {
//			return redisErr
//		}
//
//		return nil
//	}
//
// This package only gaurantees that when GetMaster is called the returned
// connection will be a connection to the master as of the moment that method is
// called. It is still possible that there is a failover as that connection is
// being used by the application. The Readonly() method on CmdError will be
// helpful if you want to gracefully handle this case.
//
// As a final note, a Client can be interacted with from multiple routines at
// once safely, except for the Close method. To safely Close, ensure that only
// one routine ever makes the call and that once the call is made no other
// methods are ever called by any routines.
package sentinel

import (
	"errors"
	"github.com/fzzy/radix/redis"
	"strings"

	"github.com/fzzy/radix/extra/pool"
	"github.com/fzzy/radix/extra/pubsub"
)

// An error wrapper returned by operations in this package. It implements the
// error interface and can therefore be passed around as a normal error.
type ClientError struct {
	err error

	// If this is true the error is due to a problem with the sentinel
	// connection, either it being closed or otherwise unavailable. If false the
	// error is due to some other circumstances. This is useful if you want to
	// implement some kind of reconnecting to sentinel on an error.
	SentinelErr bool
}

func (ce *ClientError) Error() string {
	return ce.err.Error()
}

type getReqRet struct {
	conn *redis.Client
	err  *ClientError
}

type getReq struct {
	name  string
	retCh chan *getReqRet
}

type putReq struct {
	name string
	conn *redis.Client
}

type switchMaster struct {
	name string
	addr string
}

type Client struct {
	poolSize    int
	masterPools map[string]*pool.Pool
	subClient   *pubsub.SubClient

	getCh   chan *getReq
	putCh   chan *putReq
	closeCh chan struct{}

	alwaysErr      *ClientError
	alwaysErrCh    chan *ClientError
	switchMasterCh chan *switchMaster
}

// Creates a sentinel client. Connects to the given sentinel instance, pulls the
// information for the masters of the given names, and creates an intial pool of
// connections for each master. The client will automatically replace the pool
// for any master should sentinel decide to fail the master over. The returned
// error is a *ClientError.
func NewClient(
	network, address string, poolSize int, names ...string,
) (
	*Client, error,
) {

	// We use this to fetch initial details about masters before we upgrade it
	// to a pubsub client
	client, err := redis.Dial(network, address)
	if err != nil {
		return nil, &ClientError{err: err}
	}

	masterPools := map[string]*pool.Pool{}
	for _, name := range names {
		r := client.Cmd("SENTINEL", "MASTER", name)
		l, err := r.List()
		if err != nil {
			return nil, &ClientError{err: err, SentinelErr: true}
		}
		addr := l[3] + ":" + l[5]
		pool, err := pool.NewPool("tcp", addr, poolSize)
		if err != nil {
			return nil, &ClientError{err: err}
		}
		masterPools[name] = pool
	}

	subClient := pubsub.NewSubClient(client)
	r := subClient.Subscribe("+switch-master")
	if r.Err != nil {
		return nil, &ClientError{err: r.Err, SentinelErr: true}
	}

	c := &Client{
		poolSize:       poolSize,
		masterPools:    masterPools,
		subClient:      subClient,
		getCh:          make(chan *getReq),
		putCh:          make(chan *putReq),
		closeCh:        make(chan struct{}),
		alwaysErrCh:    make(chan *ClientError),
		switchMasterCh: make(chan *switchMaster),
	}

	go c.subSpin()
	go c.spin()
	return c, nil
}

func (c *Client) subSpin() {
	for {
		r := c.subClient.Receive()
		if r.Timeout() {
			continue
		}
		if r.Err != nil {
			select {
			case c.alwaysErrCh <- &ClientError{err: r.Err, SentinelErr: true}:
			case <-c.closeCh:
			}
			return
		}
		sMsg := strings.Split(r.Message, " ")
		name := sMsg[0]
		newAddr := sMsg[3] + ":" + sMsg[4]
		select {
		case c.switchMasterCh <- &switchMaster{name, newAddr}:
		case <-c.closeCh:
			return
		}
	}
}

func (c *Client) spin() {
	for {
		select {
		case req := <-c.getCh:
			if c.alwaysErr != nil {
				req.retCh <- &getReqRet{nil, c.alwaysErr}
				continue
			}
			pool, ok := c.masterPools[req.name]
			if !ok {
				err := errors.New("unknown name: " + req.name)
				req.retCh <- &getReqRet{nil, &ClientError{err: err}}
				continue
			}
			conn, err := pool.Get()
			if err != nil {
				req.retCh <- &getReqRet{nil, &ClientError{err: err}}
				continue
			}
			req.retCh <- &getReqRet{conn, nil}

		case req := <-c.putCh:
			if pool, ok := c.masterPools[req.name]; ok {
				pool.Put(req.conn)
			}

		case err := <-c.alwaysErrCh:
			c.alwaysErr = err

		case sm := <-c.switchMasterCh:
			if p, ok := c.masterPools[sm.name]; ok {
				p.Empty()
				p = pool.NewOrEmptyPool("tcp", sm.addr, c.poolSize)
				c.masterPools[sm.name] = p
			}

		case <-c.closeCh:
			for name := range c.masterPools {
				c.masterPools[name].Empty()
			}
			c.subClient.Client.Close()
			close(c.getCh)
			close(c.putCh)
			return
		}
	}
}

// Retrieves a connection for the master of the given name. If sentinel has
// become unreachable this will always return an error. Close should be called
// in that case. The returned error is a *ClientError.
func (c *Client) GetMaster(name string) (*redis.Client, error) {
	req := getReq{name, make(chan *getReqRet)}
	c.getCh <- &req
	ret := <-req.retCh
	if ret.err != nil {
		return nil, ret.err
	}
	return ret.conn, nil
}

// Return a connection for a master of a given name. As with the pool package,
// do not return a connection which is having connectivity issues, or which is
// otherwise unable to perform requests.
func (c *Client) PutMaster(name string, client *redis.Client) {
	c.putCh <- &putReq{name, client}
}

// A useful helper method, analagous to the pool package's CarefullyPut method.
// Since we don't want to Put a connection which is having connectivity
// issues, this can be defered inside a function to make sure we only put back a
// connection when we should. It should be used like the following:
//
//	func doSomeThings(c *Client) error {
//		conn, redisErr := c.GetMaster("bucket0")
//		if redisErr != nil {
//			return redisErr
//		}
//		defer c.CarefullyPutMaster("bucket0", conn, &redisErr)
//
//		var i int
//		i, redisErr = conn.Cmd("GET", "foo").Int()
//		if redisErr != nil {
//			return redisErr
//		}
//
//		redisErr = conn.Cmd("SET", "foo", i * 3).Err
//		return redisErr
//	}
func (c *Client) CarefullyPutMaster(
	name string, client *redis.Client, potentialErr *error,
) {
	if potentialErr != nil && *potentialErr != nil {
		// If the client sent back that it's READONLY then we don't want to keep
		// this connection around. Otherwise, we don't care about command errors
		if cerr, ok := (*potentialErr).(*redis.CmdError); !ok || cerr.Readonly() {
			client.Close()
			return
		}
	}
	c.PutMaster(name, client)
}

// Closes all connection pools as well as the connection to sentinel.
func (c *Client) Close() {
	close(c.closeCh)
}
<|endoftext|>