<commit_before>// Package router is a network/router selector
package router

import (
	"context"
	"net"
	"os"
	"sort"
	"strconv"
	"sync"

	"github.com/micro/go-micro/client"
	"github.com/micro/go-micro/client/selector"
	"github.com/micro/go-micro/network/router"
	pb "github.com/micro/go-micro/network/router/proto"
	"github.com/micro/go-micro/registry"
)

type routerSelector struct {
	opts selector.Options

	// the router
	r router.Router

	// the client for the remote router
	c pb.RouterService

	// address of the remote router
	addr string

	// whether to use the remote router
	remote bool
}

type clientKey struct{}
type routerKey struct{}

// getRoutes returns the routes whether they are remote or local
func (r *routerSelector) getRoutes(service string) ([]router.Route, error) {
	if !r.remote {
		// lookup router for routes for the service
		return r.r.Table().Lookup(router.NewQuery(
			router.QueryDestination(service),
		))
	}

	// lookup the remote router

	var clientOpts []client.CallOption

	// set the remote address if specified
	if len(r.addr) > 0 {
		clientOpts = append(clientOpts, client.WithAddress(r.addr))
	}

	// call the router
	pbRoutes, err := r.c.Lookup(context.Background(), &pb.LookupRequest{
		Query: &pb.Query{
			Destination: service,
		},
	}, clientOpts...)
	if err != nil {
		return nil, err
	}

	var routes []router.Route

	// convert from pb to []*router.Route
	for _, r := range pbRoutes.Routes {
		routes = append(routes, router.Route{
			Destination: r.Destination,
			Gateway:     r.Gateway,
			Router:      r.Router,
			Network:     r.Network,
			Metric:      int(r.Metric),
		})
	}

	return routes, nil
}

func (r *routerSelector) Init(opts ...selector.Option) error {
	// no op
	return nil
}

func (r *routerSelector) Options() selector.Options {
	return r.opts
}

func (r *routerSelector) Select(service string, opts ...selector.SelectOption) (selector.Next, error) {
	// TODO: pull routes asynchronously and cache
	routes, err := r.getRoutes(service)
	if err != nil {
		return nil, err
	}

	// no routes return not found error
	if len(routes) == 0 {
		return nil, selector.ErrNotFound
	}

	// TODO: apply filters by pseudo constructing service

	// sort the routes based on metric
	sort.Slice(routes, func(i, j int) bool {
		return routes[i].Metric < routes[j].Metric
	})

	// roundrobin assuming routes are in metric preference order
	var i int
	var mtx sync.Mutex

	return func() (*registry.Node, error) {
		// get index and increment counter with every call to next
		mtx.Lock()
		idx := i
		i++
		mtx.Unlock()

		// get route based on idx
		route := routes[idx%len(routes)]

		// defaults to gateway and no port
		address := route.Gateway
		port := 0

		// check if its host:port
		host, pr, err := net.SplitHostPort(address)
		if err == nil {
			pp, _ := strconv.Atoi(pr)
			// set port
			port = pp
			// set address
			address = host
		}

		// return as a node
		return &registry.Node{
			// TODO: add id and metadata if we can
			Address: address,
			Port:    port,
		}, nil
	}, nil
}

func (r *routerSelector) Mark(service string, node *registry.Node, err error) {
	// TODO: pass back metrics or information to the router
	return
}

func (r *routerSelector) Reset(service string) {
	// TODO: reset the metrics or information at the router
	return
}

func (r *routerSelector) Close() error {
	// stop the router advertisements
	return r.r.Stop()
}

func (r *routerSelector) String() string {
	return "router"
}

// NewSelector returns a new router based selector
func NewSelector(opts ...selector.Option) selector.Selector {
	options := selector.Options{
		Context: context.Background(),
	}

	for _, o := range opts {
		o(&options)
	}

	// set default registry if not set
	if options.Registry == nil {
		options.Registry = registry.DefaultRegistry
	}

	// try get router from the context
	r, ok := options.Context.Value(routerKey{}).(router.Router)
	if !ok {
		// TODO: Use router.DefaultRouter?
		r = router.NewRouter(
			router.Registry(options.Registry),
		)
	}

	// try get client from the context
	c, ok := options.Context.Value(clientKey{}).(client.Client)
	if !ok {
		c = client.DefaultClient
	}

	// get the router from env vars if its a remote service
	remote := true
	routerName := os.Getenv("MICRO_ROUTER")
	routerAddress := os.Getenv("MICRO_ROUTER_ADDRESS")

	// start the router advertisements if we're running it locally
	if len(routerName) == 0 && len(routerAddress) == 0 {
		go r.Advertise()
		remote = false
	}

	return &routerSelector{
		opts: options,
		// set the internal router
		r: r,
		// set the client
		c: pb.NewRouterService(routerName, c),
		// address of router
		addr: routerAddress,
		// let ourselves know to use the remote router
		remote: remote,
	}
}

// WithClient sets the client for the request
func WithClient(c client.Client) selector.Option {
	return func(o *selector.Options) {
		if o.Context == nil {
			o.Context = context.Background()
		}
		o.Context = context.WithValue(o.Context, clientKey{}, c)
	}
}

// WithRouter sets the router as an option
func WithRouter(r router.Router) selector.Option {
	return func(o *selector.Options) {
		if o.Context == nil {
			o.Context = context.Background()
		}
		o.Context = context.WithValue(o.Context, routerKey{}, r)
	}
}
<commit_msg>Lookup router via registry<commit_after>// Package router is a network/router selector
package router

import (
	"context"
	"fmt"
	"net"
	"os"
	"sort"
	"strconv"
	"sync"

	"github.com/micro/go-micro/client"
	"github.com/micro/go-micro/client/selector"
	"github.com/micro/go-micro/network/router"
	pb "github.com/micro/go-micro/network/router/proto"
	"github.com/micro/go-micro/registry"
)

type routerSelector struct {
	opts selector.Options

	// the router
	r router.Router

	// the client we have
	c client.Client

	// the client for the remote router
	rs pb.RouterService

	// name of the router
	name string

	// address of the remote router
	addr string

	// whether to use the remote router
	remote bool
}

type clientKey struct{}
type routerKey struct{}

// getRoutes returns the routes whether they are remote or local
func (r *routerSelector) getRoutes(service string) ([]router.Route, error) {
	if !r.remote {
		// lookup router for routes for the service
		return r.r.Table().Lookup(router.NewQuery(
			router.QueryDestination(service),
		))
	}

	// lookup the remote router

	var addrs []string

	// set the remote address if specified
	if len(r.addr) > 0 {
		addrs = append(addrs, r.addr)
	} else {
		// we have a name so we need to check the registry
		services, err := r.c.Options().Registry.GetService(r.name)
		if err != nil {
			return nil, err
		}

		for _, service := range services {
			for _, node := range service.Nodes {
				addr := node.Address
				if node.Port > 0 {
					addr = fmt.Sprintf("%s:%d", node.Address, node.Port)
				}
				addrs = append(addrs, addr)
			}
		}
	}

	// no router addresses available
	if len(addrs) == 0 {
		return nil, selector.ErrNoneAvailable
	}

	var pbRoutes *pb.LookupResponse
	var err error

	// TODO: implement backoff and retries
	for _, addr := range addrs {
		// call the router
		pbRoutes, err = r.rs.Lookup(context.Background(), &pb.LookupRequest{
			Query: &pb.Query{
				Destination: service,
			},
		}, client.WithAddress(addr))
		if err != nil {
			continue
		}
		break
	}

	// errored out
	if err != nil {
		return nil, err
	}

	// no routes
	if pbRoutes == nil {
		return nil, selector.ErrNoneAvailable
	}

	var routes []router.Route

	// convert from pb to []*router.Route
	for _, r := range pbRoutes.Routes {
		routes = append(routes, router.Route{
			Destination: r.Destination,
			Gateway:     r.Gateway,
			Router:      r.Router,
			Network:     r.Network,
			Metric:      int(r.Metric),
		})
	}

	return routes, nil
}

func (r *routerSelector) Init(opts ...selector.Option) error {
	// no op
	return nil
}

func (r *routerSelector) Options() selector.Options {
	return r.opts
}

func (r *routerSelector) Select(service string, opts ...selector.SelectOption) (selector.Next, error) {
	// TODO: pull routes asynchronously and cache
	routes, err := r.getRoutes(service)
	if err != nil {
		return nil, err
	}

	// no routes return not found error
	if len(routes) == 0 {
		return nil, selector.ErrNotFound
	}

	// TODO: apply filters by pseudo constructing service

	// sort the routes based on metric
	sort.Slice(routes, func(i, j int) bool {
		return routes[i].Metric < routes[j].Metric
	})

	// roundrobin assuming routes are in metric preference order
	var i int
	var mtx sync.Mutex

	return func() (*registry.Node, error) {
		// get index and increment counter with every call to next
		mtx.Lock()
		idx := i
		i++
		mtx.Unlock()

		// get route based on idx
		route := routes[idx%len(routes)]

		// defaults to gateway and no port
		address := route.Gateway
		port := 0

		// check if its host:port
		host, pr, err := net.SplitHostPort(address)
		if err == nil {
			pp, _ := strconv.Atoi(pr)
			// set port
			port = pp
			// set address
			address = host
		}

		// return as a node
		return &registry.Node{
			// TODO: add id and metadata if we can
			Address: address,
			Port:    port,
		}, nil
	}, nil
}

func (r *routerSelector) Mark(service string, node *registry.Node, err error) {
	// TODO: pass back metrics or information to the router
	return
}

func (r *routerSelector) Reset(service string) {
	// TODO: reset the metrics or information at the router
	return
}

func (r *routerSelector) Close() error {
	// stop the router advertisements
	return r.r.Stop()
}

func (r *routerSelector) String() string {
	return "router"
}

// NewSelector returns a new router based selector
func NewSelector(opts ...selector.Option) selector.Selector {
	options := selector.Options{
		Context: context.Background(),
	}

	for _, o := range opts {
		o(&options)
	}

	// set default registry if not set
	if options.Registry == nil {
		options.Registry = registry.DefaultRegistry
	}

	// try get router from the context
	r, ok := options.Context.Value(routerKey{}).(router.Router)
	if !ok {
		// TODO: Use router.DefaultRouter?
		r = router.NewRouter(
			router.Registry(options.Registry),
		)
	}

	// try get client from the context
	c, ok := options.Context.Value(clientKey{}).(client.Client)
	if !ok {
		c = client.DefaultClient
	}

	// get the router from env vars if its a remote service
	remote := true
	routerName := os.Getenv("MICRO_ROUTER")
	routerAddress := os.Getenv("MICRO_ROUTER_ADDRESS")

	// start the router advertisements if we're running it locally
	if len(routerName) == 0 && len(routerAddress) == 0 {
		go r.Advertise()
		remote = false
	}

	return &routerSelector{
		opts: options,
		// set the internal router
		r: r,
		// set the client
		c: c,
		// set the router client
		rs: pb.NewRouterService(routerName, c),
		// name of the router
		name: routerName,
		// address of router
		addr: routerAddress,
		// let ourselves know to use the remote router
		remote: remote,
	}
}

// WithClient sets the client for the request
func WithClient(c client.Client) selector.Option {
	return func(o *selector.Options) {
		if o.Context == nil {
			o.Context = context.Background()
		}
		o.Context = context.WithValue(o.Context, clientKey{}, c)
	}
}

// WithRouter sets the router as an option
func WithRouter(r router.Router) selector.Option {
	return func(o *selector.Options) {
		if o.Context == nil {
			o.Context = context.Background()
		}
		o.Context = context.WithValue(o.Context, routerKey{}, r)
	}
}
<|endoftext|>


<commit_before>package client

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"text/template"

	services "github.com/maximilien/softlayer-go/services"
	softlayer "github.com/maximilien/softlayer-go/softlayer"
)

const (
	SOFTLAYER_API_URL  = "api.softlayer.com/rest/v3"
	TEMPLATE_ROOT_PATH = "templates"
)

type softLayerClient struct {
	username string
	apiKey   string

	templatePath string

	httpClient *http.Client

	softLayerServices map[string]softlayer.Service
}

func NewSoftLayerClient(username, apiKey string) *softLayerClient {
	pwd, _ := os.Getwd()
	slc := &softLayerClient{
		username: username,
		apiKey:   apiKey,

		templatePath: filepath.Join(pwd, TEMPLATE_ROOT_PATH),

		httpClient: &http.Client{
			Transport: &http.Transport{
				Proxy: http.ProxyFromEnvironment,
			},
		},

		softLayerServices: map[string]softlayer.Service{},
	}

	slc.initSoftLayerServices()

	return slc
}

//softlayer.Client interface methods

func (slc *softLayerClient) GetService(serviceName string) (softlayer.Service, error) {
	slService, ok := slc.softLayerServices[serviceName]
	if !ok {
		return nil, errors.New(fmt.Sprintf("softlayer-go does not support service '%s'", serviceName))
	}

	return slService, nil
}

func (slc *softLayerClient) GetSoftLayer_Account_Service() (softlayer.SoftLayer_Account_Service, error) {
	slService, err := slc.GetService("SoftLayer_Account")
	if err != nil {
		return nil, err
	}

	return slService.(softlayer.SoftLayer_Account_Service), nil
}

func (slc *softLayerClient) GetSoftLayer_Virtual_Guest_Service() (softlayer.SoftLayer_Virtual_Guest_Service, error) {
	slService, err := slc.GetService("SoftLayer_Virtual_Guest")
	if err != nil {
		return nil, err
	}

	return slService.(softlayer.SoftLayer_Virtual_Guest_Service), nil
}

func (slc *softLayerClient) GetSoftLayer_Virtual_Disk_Image_Service() (softlayer.SoftLayer_Virtual_Disk_Image_Service, error) {
	slService, err := slc.GetService("SoftLayer_Virtual_Disk_Image")
	if err != nil {
		return nil, err
	}

	return slService.(softlayer.SoftLayer_Virtual_Disk_Image_Service), nil
}

func (slc *softLayerClient) GetSoftLayer_Security_Ssh_Key_Service() (softlayer.SoftLayer_Security_Ssh_Key_Service, error) {
	slService, err := slc.GetService("SoftLayer_Security_Ssh_Key")
	if err != nil {
		return nil, err
	}

	return slService.(softlayer.SoftLayer_Security_Ssh_Key_Service), nil
}

func (slc *softLayerClient) GetSoftLayer_Product_Package_Service() (softlayer.SoftLayer_Product_Package_Service, error) {
	slService, err := slc.GetService("SoftLayer_Product_Package")
	if err != nil {
		return nil, err
	}

	return slService.(softlayer.SoftLayer_Product_Package_Service), nil
}

func (slc *softLayerClient) GetSoftLayer_Virtual_Guest_Block_Device_Template_Group_Service() (softlayer.SoftLayer_Virtual_Guest_Block_Device_Template_Group_Service, error) {
	slService, err := slc.GetService("SoftLayer_Virtual_Guest_Block_Device_Template_Group_Service")
	if err != nil {
		return nil, err
	}

	return slService.(softlayer.SoftLayer_Virtual_Guest_Block_Device_Template_Group_Service), nil
}

func (slc *softLayerClient) GetSoftLayer_Network_Storage_Service() (softlayer.SoftLayer_Network_Storage_Service, error) {
	slService, err := slc.GetService("SoftLayer_Network_Storage")
	if err != nil {
		return nil, err
	}

	return slService.(softlayer.SoftLayer_Network_Storage_Service), nil
}

func (slc *softLayerClient) GetSoftLayer_Product_Order_Service() (softlayer.SoftLayer_Product_Order_Service, error) {
	slService, err := slc.GetService("SoftLayer_Product_Order")
	if err != nil {
		return nil, err
	}

	return slService.(softlayer.SoftLayer_Product_Order_Service), nil
}

func (slc *softLayerClient) GetSoftLayer_Billing_Item_Cancellation_Request_Service() (softlayer.SoftLayer_Billing_Item_Cancellation_Request_Service, error) {
	slService, err := slc.GetService("SoftLayer_Billing_Item_Cancellation_Request")
	if err != nil {
		return nil, err
	}

	return slService.(softlayer.SoftLayer_Billing_Item_Cancellation_Request_Service), nil
}

func (slc *softLayerClient) GetSoftLayer_Hardware_Service() (softlayer.SoftLayer_Hardware_Service, error) {
	slService, err := slc.GetService("SoftLayer_Hardware")
	if err != nil {
		return nil, err
	}

	return slService.(softlayer.SoftLayer_Hardware_Service), nil
}

//Public methods

func (slc *softLayerClient) DoRawHttpRequestWithObjectMask(path string, masks []string, requestType string, requestBody *bytes.Buffer) ([]byte, error) {
	url := fmt.Sprintf("https://%s:%s@%s/%s", slc.username, slc.apiKey, SOFTLAYER_API_URL, path)

	url += "?objectMask="
	for i := 0; i < len(masks); i++ {
		url += masks[i]
		if i != len(masks)-1 {
			url += ";"
		}
	}

	return slc.makeHttpRequest(url, requestType, requestBody)
}

func (slc *softLayerClient) DoRawHttpRequest(path string, requestType string, requestBody *bytes.Buffer) ([]byte, error) {
	url := fmt.Sprintf("https://%s:%s@%s/%s", slc.username, slc.apiKey, SOFTLAYER_API_URL, path)
	return slc.makeHttpRequest(url, requestType, requestBody)
}

func (slc *softLayerClient) GenerateRequestBody(templateData interface{}) (*bytes.Buffer, error) {
	cwd, err := os.Getwd()
	if err != nil {
		return nil, err
	}

	bodyTemplate := template.Must(template.ParseFiles(filepath.Join(cwd, slc.templatePath)))
	body := new(bytes.Buffer)
	bodyTemplate.Execute(body, templateData)

	return body, nil
}

func (slc *softLayerClient) HasErrors(body map[string]interface{}) error {
	if errString, ok := body["error"]; !ok {
		return nil
	} else {
		return errors.New(errString.(string))
	}
}

func (slc *softLayerClient) CheckForHttpResponseErrors(data []byte) error {
	var decodedResponse map[string]interface{}
	err := json.Unmarshal(data, &decodedResponse)
	if err != nil {
		return err
	}

	if err := slc.HasErrors(decodedResponse); err != nil {
		return err
	}

	return nil
}

//Private methods

func (slc *softLayerClient) initSoftLayerServices() {
	slc.softLayerServices["SoftLayer_Account"] = services.NewSoftLayer_Account_Service(slc)
	slc.softLayerServices["SoftLayer_Virtual_Guest"] = services.NewSoftLayer_Virtual_Guest_Service(slc)
	slc.softLayerServices["SoftLayer_Virtual_Disk_Image"] = services.NewSoftLayer_Virtual_Disk_Image_Service(slc)
	slc.softLayerServices["SoftLayer_Security_Ssh_Key"] = services.NewSoftLayer_Security_Ssh_Key_Service(slc)
	slc.softLayerServices["SoftLayer_Product_Package"] = services.NewSoftLayer_Product_Package_Service(slc)
	slc.softLayerServices["SoftLayer_Network_Storage"] = services.NewSoftLayer_Network_Storage_Service(slc)
	slc.softLayerServices["SoftLayer_Product_Order"] = services.NewSoftLayer_Product_Order_Service(slc)
	slc.softLayerServices["SoftLayer_Billing_Item_Cancellation_Request"] = services.NewSoftLayer_Billing_Item_Cancellation_Request_Service(slc)
	slc.softLayerServices["SoftLayer_Virtual_Guest_Block_Device_Template_Group"] = services.NewSoftLayer_Virtual_Guest_Block_Device_Template_Group_Service(slc)
	slc.softLayerServices["SoftLayer_Hardware"] = services.NewSoftLayer_Hardware_Service(slc)
}

func (slc *softLayerClient) makeHttpRequest(url string, requestType string, requestBody *bytes.Buffer) ([]byte, error) {
	var lastResponse http.Response
	switch requestType {
	case "POST", "DELETE":
		req, err := http.NewRequest(requestType, url, requestBody)

		if err != nil {
			return nil, err
		}
		resp, err := slc.httpClient.Do(req)

		if err != nil {
			return nil, err
		} else {
			lastResponse = *resp
		}
	case "GET":
		resp, err := http.Get(url)

		if err != nil {
			return nil, err
		} else {
			lastResponse = *resp
		}
	default:
		return nil, errors.New(fmt.Sprintf("Undefined request type '%s', only GET/POST/DELETE are available!", requestType))
	}

	responseBody, err := ioutil.ReadAll(lastResponse.Body)
	lastResponse.Body.Close()
	if err != nil {
		return nil, err
	}

	return responseBody, nil
}
<commit_msg>log HTTP req/responses<commit_after>package client

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httputil"
	"os"
	"path/filepath"
	"text/template"

	services "github.com/maximilien/softlayer-go/services"
	softlayer "github.com/maximilien/softlayer-go/softlayer"
)

const (
	SOFTLAYER_API_URL  = "api.softlayer.com/rest/v3"
	TEMPLATE_ROOT_PATH = "templates"
)

type softLayerClient struct {
	username string
	apiKey   string

	templatePath string

	httpClient *http.Client

	softLayerServices map[string]softlayer.Service
}

func NewSoftLayerClient(username, apiKey string) *softLayerClient {
	pwd, _ := os.Getwd()
	slc := &softLayerClient{
		username: username,
		apiKey:   apiKey,

		templatePath: filepath.Join(pwd, TEMPLATE_ROOT_PATH),

		httpClient: &http.Client{
			Transport: &http.Transport{
				Proxy: http.ProxyFromEnvironment,
			},
		},

		softLayerServices: map[string]softlayer.Service{},
	}

	slc.initSoftLayerServices()

	return slc
}

//softlayer.Client interface methods

func (slc *softLayerClient) GetService(serviceName string) (softlayer.Service, error) {
	slService, ok := slc.softLayerServices[serviceName]
	if !ok {
		return nil, errors.New(fmt.Sprintf("softlayer-go does not support service '%s'", serviceName))
	}

	return slService, nil
}

func (slc *softLayerClient) GetSoftLayer_Account_Service() (softlayer.SoftLayer_Account_Service, error) {
	slService, err := slc.GetService("SoftLayer_Account")
	if err != nil {
		return nil, err
	}

	return slService.(softlayer.SoftLayer_Account_Service), nil
}

func (slc *softLayerClient) GetSoftLayer_Virtual_Guest_Service() (softlayer.SoftLayer_Virtual_Guest_Service, error) {
	slService, err := slc.GetService("SoftLayer_Virtual_Guest")
	if err != nil {
		return nil, err
	}

	return slService.(softlayer.SoftLayer_Virtual_Guest_Service), nil
}

func (slc *softLayerClient) GetSoftLayer_Virtual_Disk_Image_Service() (softlayer.SoftLayer_Virtual_Disk_Image_Service, error) {
	slService, err := slc.GetService("SoftLayer_Virtual_Disk_Image")
	if err != nil {
		return nil, err
	}

	return slService.(softlayer.SoftLayer_Virtual_Disk_Image_Service), nil
}

func (slc *softLayerClient) GetSoftLayer_Security_Ssh_Key_Service() (softlayer.SoftLayer_Security_Ssh_Key_Service, error) {
	slService, err := slc.GetService("SoftLayer_Security_Ssh_Key")
	if err != nil {
		return nil, err
	}

	return slService.(softlayer.SoftLayer_Security_Ssh_Key_Service), nil
}

func (slc *softLayerClient) GetSoftLayer_Product_Package_Service() (softlayer.SoftLayer_Product_Package_Service, error) {
	slService, err := slc.GetService("SoftLayer_Product_Package")
	if err != nil {
		return nil, err
	}

	return slService.(softlayer.SoftLayer_Product_Package_Service), nil
}

func (slc *softLayerClient) GetSoftLayer_Virtual_Guest_Block_Device_Template_Group_Service() (softlayer.SoftLayer_Virtual_Guest_Block_Device_Template_Group_Service, error) {
	slService, err := slc.GetService("SoftLayer_Virtual_Guest_Block_Device_Template_Group_Service")
	if err != nil {
		return nil, err
	}

	return slService.(softlayer.SoftLayer_Virtual_Guest_Block_Device_Template_Group_Service), nil
}

func (slc *softLayerClient) GetSoftLayer_Network_Storage_Service() (softlayer.SoftLayer_Network_Storage_Service, error) {
	slService, err := slc.GetService("SoftLayer_Network_Storage")
	if err != nil {
		return nil, err
	}

	return slService.(softlayer.SoftLayer_Network_Storage_Service), nil
}

func (slc *softLayerClient) GetSoftLayer_Product_Order_Service() (softlayer.SoftLayer_Product_Order_Service, error) {
	slService, err := slc.GetService("SoftLayer_Product_Order")
	if err != nil {
		return nil, err
	}

	return slService.(softlayer.SoftLayer_Product_Order_Service), nil
}

func (slc *softLayerClient) GetSoftLayer_Billing_Item_Cancellation_Request_Service() (softlayer.SoftLayer_Billing_Item_Cancellation_Request_Service, error) {
	slService, err := slc.GetService("SoftLayer_Billing_Item_Cancellation_Request")
	if err != nil {
		return nil, err
	}

	return slService.(softlayer.SoftLayer_Billing_Item_Cancellation_Request_Service), nil
}

func (slc *softLayerClient) GetSoftLayer_Hardware_Service() (softlayer.SoftLayer_Hardware_Service, error) {
	slService, err := slc.GetService("SoftLayer_Hardware")
	if err != nil {
		return nil, err
	}

	return slService.(softlayer.SoftLayer_Hardware_Service), nil
}

//Public methods

func (slc *softLayerClient) DoRawHttpRequestWithObjectMask(path string, masks []string, requestType string, requestBody *bytes.Buffer) ([]byte, error) {
	url := fmt.Sprintf("https://%s:%s@%s/%s", slc.username, slc.apiKey, SOFTLAYER_API_URL, path)

	url += "?objectMask="
	for i := 0; i < len(masks); i++ {
		url += masks[i]
		if i != len(masks)-1 {
			url += ";"
		}
	}

	return slc.makeHttpRequest(url, requestType, requestBody)
}

func (slc *softLayerClient) DoRawHttpRequest(path string, requestType string, requestBody *bytes.Buffer) ([]byte, error) {
	url := fmt.Sprintf("https://%s:%s@%s/%s", slc.username, slc.apiKey, SOFTLAYER_API_URL, path)
	return slc.makeHttpRequest(url, requestType, requestBody)
}

func (slc *softLayerClient) GenerateRequestBody(templateData interface{}) (*bytes.Buffer, error) {
	cwd, err := os.Getwd()
	if err != nil {
		return nil, err
	}

	bodyTemplate := template.Must(template.ParseFiles(filepath.Join(cwd, slc.templatePath)))
	body := new(bytes.Buffer)
	bodyTemplate.Execute(body, templateData)

	return body, nil
}

func (slc *softLayerClient) HasErrors(body map[string]interface{}) error {
	if errString, ok := body["error"]; !ok {
		return nil
	} else {
		return errors.New(errString.(string))
	}
}

func (slc *softLayerClient) CheckForHttpResponseErrors(data []byte) error {
	var decodedResponse map[string]interface{}
	err := json.Unmarshal(data, &decodedResponse)
	if err != nil {
		return err
	}

	if err := slc.HasErrors(decodedResponse); err != nil {
		return err
	}

	return nil
}

//Private methods

func (slc *softLayerClient) initSoftLayerServices() {
	slc.softLayerServices["SoftLayer_Account"] = services.NewSoftLayer_Account_Service(slc)
	slc.softLayerServices["SoftLayer_Virtual_Guest"] = services.NewSoftLayer_Virtual_Guest_Service(slc)
	slc.softLayerServices["SoftLayer_Virtual_Disk_Image"] = services.NewSoftLayer_Virtual_Disk_Image_Service(slc)
	slc.softLayerServices["SoftLayer_Security_Ssh_Key"] = services.NewSoftLayer_Security_Ssh_Key_Service(slc)
	slc.softLayerServices["SoftLayer_Product_Package"] = services.NewSoftLayer_Product_Package_Service(slc)
	slc.softLayerServices["SoftLayer_Network_Storage"] = services.NewSoftLayer_Network_Storage_Service(slc)
	slc.softLayerServices["SoftLayer_Product_Order"] = services.NewSoftLayer_Product_Order_Service(slc)
	slc.softLayerServices["SoftLayer_Billing_Item_Cancellation_Request"] = services.NewSoftLayer_Billing_Item_Cancellation_Request_Service(slc)
	slc.softLayerServices["SoftLayer_Virtual_Guest_Block_Device_Template_Group"] = services.NewSoftLayer_Virtual_Guest_Block_Device_Template_Group_Service(slc)
	slc.softLayerServices["SoftLayer_Hardware"] = services.NewSoftLayer_Hardware_Service(slc)
}

func (slc *softLayerClient) makeHttpRequest(url string, requestType string, requestBody *bytes.Buffer) ([]byte, error) {
	req, err := http.NewRequest(requestType, url, requestBody)
	if err != nil {
		return nil, err
	}

	bs, err := httputil.DumpRequest(req, true)
	if err != nil {
		return nil, err
	}

	fmt.Fprintf(os.Stderr, "\n---\n[softlayer-go] Request:\n%s\n", string(bs))

	resp, err := slc.httpClient.Do(req)
	if err != nil {
		return nil, err
	}

	defer resp.Body.Close()

	bs, err = httputil.DumpResponse(resp, true)
	if err != nil {
		return nil, err
	}

	fmt.Fprintf(os.Stderr, "[softlayer-go] Response:\n%s\n", string(bs))

	responseBody, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	return responseBody, nil
}
<|endoftext|>


<commit_before>package model

import (
	"log"
	"sort"
	"sync"
	"time"
)

type Monitor interface {
	FileBegins(<-chan content) error
	FileDone() error
}

type FileQueue struct {
	files        queuedFileList
	sorted       bool
	fmut         sync.Mutex // protects files and sorted
	availability map[string][]string
	amut         sync.Mutex // protects availability
	queued       map[string]bool
}

type queuedFile struct {
	name         string
	blocks       []Block
	activeBlocks []bool
	given        int
	remaining    int
	channel      chan content
	nodes        []string
	nodesChecked time.Time
	monitor      Monitor
}

type content struct {
	offset int64
	data   []byte
}

type queuedFileList []queuedFile

func (l queuedFileList) Len() int { return len(l) }

func (l queuedFileList) Swap(a, b int) { l[a], l[b] = l[b], l[a] }

func (l queuedFileList) Less(a, b int) bool {
	// Sort by most blocks already given out, then alphabetically
	if l[a].given != l[b].given {
		return l[a].given > l[b].given
	}
	return l[a].name < l[b].name
}

type queuedBlock struct {
	name  string
	block Block
	index int
}

func NewFileQueue() *FileQueue {
	return &FileQueue{
		availability: make(map[string][]string),
		queued:       make(map[string]bool),
	}
}

func (q *FileQueue) Add(name string, blocks []Block, monitor Monitor) {
	q.fmut.Lock()
	defer q.fmut.Unlock()

	if q.queued[name] {
		return
	}

	q.files = append(q.files, queuedFile{
		name:         name,
		blocks:       blocks,
		activeBlocks: make([]bool, len(blocks)),
		remaining:    len(blocks),
		channel:      make(chan content),
		monitor:      monitor,
	})
	q.queued[name] = true
	q.sorted = false
}

func (q *FileQueue) Len() int {
	q.fmut.Lock()
	defer q.fmut.Unlock()

	return len(q.files)
}

func (q *FileQueue) Get(nodeID string) (queuedBlock, bool) {
	q.fmut.Lock()
	defer q.fmut.Unlock()

	if !q.sorted {
		sort.Sort(q.files)
		q.sorted = true
	}

	for i := range q.files {
		qf := &q.files[i]

		q.amut.Lock()
		av := q.availability[qf.name]
		q.amut.Unlock()

		if len(av) == 0 {
			// Noone has the file we want; abort.
			if qf.remaining != len(qf.blocks) {
				// We have already started on this file; close it down
				close(qf.channel)
				if mon := qf.monitor; mon != nil {
					mon.FileDone()
				}
			}
			delete(q.queued, qf.name)
			q.deleteAt(i)
			return queuedBlock{}, false
		}

		for _, ni := range av {
			// Find and return the next block in the queue
			if ni == nodeID {
				for j, b := range qf.blocks {
					if !qf.activeBlocks[j] {
						qf.activeBlocks[j] = true
						qf.given++
						return queuedBlock{
							name:  qf.name,
							block: b,
							index: j,
						}, true
					}
				}
				break
			}
		}
	}

	// We found nothing to do
	return queuedBlock{}, false
}

func (q *FileQueue) Done(file string, offset int64, data []byte) {
	q.fmut.Lock()
	defer q.fmut.Unlock()

	c := content{
		offset: offset,
		data:   data,
	}
	for i := range q.files {
		qf := &q.files[i]

		if qf.name == file {
			if qf.monitor != nil && qf.remaining == len(qf.blocks) {
				err := qf.monitor.FileBegins(qf.channel)
				if err != nil {
					log.Printf("WARNING: %s: %v (not synced)", qf.name, err)
					delete(q.queued, qf.name)
					q.deleteAt(i)
					return
				}
			}

			qf.channel <- c
			qf.remaining--

			if qf.remaining == 0 {
				close(qf.channel)
				if qf.monitor != nil {
					err := qf.monitor.FileDone()
					if err != nil {
						log.Printf("WARNING: %s: %v", qf.name, err)
					}
				}
				delete(q.queued, qf.name)
				q.deleteAt(i)
			}
			return
		}
	}

	// We found nothing, might have errored out already
}

func (q *FileQueue) QueuedFiles() (files []string) {
	q.fmut.Lock()
	defer q.fmut.Unlock()

	for _, qf := range q.files {
		files = append(files, qf.name)
	}
	return
}

func (q *FileQueue) deleteAt(i int) {
	q.files = append(q.files[:i], q.files[i+1:]...)
}

func (q *FileQueue) deleteFile(n string) {
	for i, file := range q.files {
		if n == file.name {
			q.deleteAt(i)
			delete(q.queued, file.name)
			return
		}
	}
}

func (q *FileQueue) SetAvailable(file string, nodes []string) {
	q.amut.Lock()
	defer q.amut.Unlock()

	q.availability[file] = nodes
}

func (q *FileQueue) RemoveAvailable(toRemove string) {
	q.amut.Lock()
	defer q.amut.Unlock()

	for file, nodes := range q.availability {
		for i, node := range nodes {
			if node == toRemove {
				q.availability[file] = nodes[:i+copy(nodes[i:], nodes[i+1:])]
				if len(q.availability[file]) == 0 {
					q.fmut.Lock()
					q.deleteFile(file)
					q.fmut.Unlock()
				}
			}
			break
		}
	}
}
<commit_msg>(Re)Fix locking around deleteFile (fixes #64)<commit_after>package model

import (
	"log"
	"sort"
	"sync"
	"time"
)

type Monitor interface {
	FileBegins(<-chan content) error
	FileDone() error
}

type FileQueue struct {
	files        queuedFileList
	sorted       bool
	fmut         sync.Mutex // protects files and sorted
	availability map[string][]string
	amut         sync.Mutex // protects availability
	queued       map[string]bool
}

type queuedFile struct {
	name         string
	blocks       []Block
	activeBlocks []bool
	given        int
	remaining    int
	channel      chan content
	nodes        []string
	nodesChecked time.Time
	monitor      Monitor
}

type content struct {
	offset int64
	data   []byte
}

type queuedFileList []queuedFile

func (l queuedFileList) Len() int { return len(l) }

func (l queuedFileList) Swap(a, b int) { l[a], l[b] = l[b], l[a] }

func (l queuedFileList) Less(a, b int) bool {
	// Sort by most blocks already given out, then alphabetically
	if l[a].given != l[b].given {
		return l[a].given > l[b].given
	}
	return l[a].name < l[b].name
}

type queuedBlock struct {
	name  string
	block Block
	index int
}

func NewFileQueue() *FileQueue {
	return &FileQueue{
		availability: make(map[string][]string),
		queued:       make(map[string]bool),
	}
}

func (q *FileQueue) Add(name string, blocks []Block, monitor Monitor) {
	q.fmut.Lock()
	defer q.fmut.Unlock()

	if q.queued[name] {
		return
	}

	q.files = append(q.files, queuedFile{
		name:         name,
		blocks:       blocks,
		activeBlocks: make([]bool, len(blocks)),
		remaining:    len(blocks),
		channel:      make(chan content),
		monitor:      monitor,
	})
	q.queued[name] = true
	q.sorted = false
}

func (q *FileQueue) Len() int {
	q.fmut.Lock()
	defer q.fmut.Unlock()

	return len(q.files)
}

func (q *FileQueue) Get(nodeID string) (queuedBlock, bool) {
	q.fmut.Lock()
	defer q.fmut.Unlock()

	if !q.sorted {
		sort.Sort(q.files)
		q.sorted = true
	}

	for i := range q.files {
		qf := &q.files[i]

		q.amut.Lock()
		av := q.availability[qf.name]
		q.amut.Unlock()

		if len(av) == 0 {
			// Noone has the file we want; abort.
			if qf.remaining != len(qf.blocks) {
				// We have already started on this file; close it down
				close(qf.channel)
				if mon := qf.monitor; mon != nil {
					mon.FileDone()
				}
			}
			delete(q.queued, qf.name)
			q.deleteAt(i)
			return queuedBlock{}, false
		}

		for _, ni := range av {
			// Find and return the next block in the queue
			if ni == nodeID {
				for j, b := range qf.blocks {
					if !qf.activeBlocks[j] {
						qf.activeBlocks[j] = true
						qf.given++
						return queuedBlock{
							name:  qf.name,
							block: b,
							index: j,
						}, true
					}
				}
				break
			}
		}
	}

	// We found nothing to do
	return queuedBlock{}, false
}

func (q *FileQueue) Done(file string, offset int64, data []byte) {
	q.fmut.Lock()
	defer q.fmut.Unlock()

	c := content{
		offset: offset,
		data:   data,
	}
	for i := range q.files {
		qf := &q.files[i]

		if qf.name == file {
			if qf.monitor != nil && qf.remaining == len(qf.blocks) {
				err := qf.monitor.FileBegins(qf.channel)
				if err != nil {
					log.Printf("WARNING: %s: %v (not synced)", qf.name, err)
					delete(q.queued, qf.name)
					q.deleteAt(i)
					return
				}
			}

			qf.channel <- c
			qf.remaining--

			if qf.remaining == 0 {
				close(qf.channel)
				if qf.monitor != nil {
					err := qf.monitor.FileDone()
					if err != nil {
						log.Printf("WARNING: %s: %v", qf.name, err)
					}
				}
				delete(q.queued, qf.name)
				q.deleteAt(i)
			}
			return
		}
	}

	// We found nothing, might have errored out already
}

func (q *FileQueue) QueuedFiles() (files []string) {
	q.fmut.Lock()
	defer q.fmut.Unlock()

	for _, qf := range q.files {
		files = append(files, qf.name)
	}
	return
}

func (q *FileQueue) deleteAt(i int) {
	q.files = append(q.files[:i], q.files[i+1:]...)
}

func (q *FileQueue) deleteFile(n string) {
	for i, file := range q.files {
		if n == file.name {
			q.deleteAt(i)
			delete(q.queued, file.name)
			return
		}
	}
}

func (q *FileQueue) SetAvailable(file string, nodes []string) {
	q.amut.Lock()
	defer q.amut.Unlock()

	q.availability[file] = nodes
}

func (q *FileQueue) RemoveAvailable(toRemove string) {
	q.amut.Lock()
	q.fmut.Lock()
	defer q.fmut.Unlock()
	defer q.amut.Unlock()

	for file, nodes := range q.availability {
		for i, node := range nodes {
			if node == toRemove {
				q.availability[file] = nodes[:i+copy(nodes[i:], nodes[i+1:])]
				if len(q.availability[file]) == 0 {
					q.deleteFile(file)
				}
			}
			break
		}
	}
}
<|endoftext|>


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

package zip

import (
	"compress/flate"
	"io"
	"io/ioutil"
	"sync"
)

// A Compressor returns a compressing writer, writing to the
// provided writer. On Close, any pending data should be flushed.
type Compressor func(io.Writer) (io.WriteCloser, error)

// Decompressor is a function that wraps a Reader with a decompressing Reader.
// The decompressed ReadCloser is returned to callers who open files from
// within the archive.  These callers are responsible for closing this reader
// when they're finished reading.
type Decompressor func(io.Reader) io.ReadCloser

var (
	mu sync.RWMutex // guards compressor and decompressor maps

	compressors = map[uint16]Compressor{
		Store:   func(w io.Writer) (io.WriteCloser, error) { return &nopCloser{w}, nil },
		Deflate: func(w io.Writer) (io.WriteCloser, error) { return flate.NewWriter(w, 5) },
	}

	decompressors = map[uint16]Decompressor{
		Store:   ioutil.NopCloser,
		Deflate: flate.NewReader,
	}
)

// RegisterDecompressor allows custom decompressors for a specified method ID.
func RegisterDecompressor(method uint16, d Decompressor) {
	mu.Lock()
	defer mu.Unlock()

	if _, ok := decompressors[method]; ok {
		panic("decompressor already registered")
	}
	decompressors[method] = d
}

// RegisterCompressor registers custom compressors for a specified method ID.
// The common methods Store and Deflate are built in.
func RegisterCompressor(method uint16, comp Compressor) {
	mu.Lock()
	defer mu.Unlock()

	if _, ok := compressors[method]; ok {
		panic("compressor already registered")
	}
	compressors[method] = comp
}

func compressor(method uint16) Compressor {
	mu.RLock()
	defer mu.RUnlock()
	return compressors[method]
}

func decompressor(method uint16) Decompressor {
	mu.RLock()
	defer mu.RUnlock()
	return decompressors[method]
}
<commit_msg>archive/zip: re-use flate.Writers when writing compressed files<commit_after>// Copyright 2010 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package zip

import (
	"compress/flate"
	"errors"
	"io"
	"io/ioutil"
	"sync"
)

// A Compressor returns a compressing writer, writing to the
// provided writer. On Close, any pending data should be flushed.
type Compressor func(io.Writer) (io.WriteCloser, error)

// Decompressor is a function that wraps a Reader with a decompressing Reader.
// The decompressed ReadCloser is returned to callers who open files from
// within the archive.  These callers are responsible for closing this reader
// when they're finished reading.
type Decompressor func(io.Reader) io.ReadCloser

var flateWriterPool sync.Pool

func newFlateWriter(w io.Writer) io.WriteCloser {
	fw, ok := flateWriterPool.Get().(*flate.Writer)
	if ok {
		fw.Reset(w)
	} else {
		fw, _ = flate.NewWriter(w, 5)
	}
	return &pooledFlateWriter{fw: fw}
}

type pooledFlateWriter struct {
	mu sync.Mutex // guards Close and Write
	fw *flate.Writer
}

func (w *pooledFlateWriter) Write(p []byte) (n int, err error) {
	w.mu.Lock()
	defer w.mu.Unlock()
	if w.fw == nil {
		return 0, errors.New("Write after Close")
	}
	return w.fw.Write(p)
}

func (w *pooledFlateWriter) Close() error {
	w.mu.Lock()
	defer w.mu.Unlock()
	var err error
	if w.fw != nil {
		err = w.fw.Close()
		flateWriterPool.Put(w.fw)
		w.fw = nil
	}
	return err
}

var (
	mu sync.RWMutex // guards compressor and decompressor maps

	compressors = map[uint16]Compressor{
		Store:   func(w io.Writer) (io.WriteCloser, error) { return &nopCloser{w}, nil },
		Deflate: func(w io.Writer) (io.WriteCloser, error) { return newFlateWriter(w), nil },
	}

	decompressors = map[uint16]Decompressor{
		Store:   ioutil.NopCloser,
		Deflate: flate.NewReader,
	}
)

// RegisterDecompressor allows custom decompressors for a specified method ID.
func RegisterDecompressor(method uint16, d Decompressor) {
	mu.Lock()
	defer mu.Unlock()

	if _, ok := decompressors[method]; ok {
		panic("decompressor already registered")
	}
	decompressors[method] = d
}

// RegisterCompressor registers custom compressors for a specified method ID.
// The common methods Store and Deflate are built in.
func RegisterCompressor(method uint16, comp Compressor) {
	mu.Lock()
	defer mu.Unlock()

	if _, ok := compressors[method]; ok {
		panic("compressor already registered")
	}
	compressors[method] = comp
}

func compressor(method uint16) Compressor {
	mu.RLock()
	defer mu.RUnlock()
	return compressors[method]
}

func decompressor(method uint16) Decompressor {
	mu.RLock()
	defer mu.RUnlock()
	return decompressors[method]
}
<|endoftext|>


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

// Tests that involve both reading and writing.

package zip

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"strings"
	"testing"
	"time"
)

func TestOver65kFiles(t *testing.T) {
	buf := new(bytes.Buffer)
	w := NewWriter(buf)
	const nFiles = (1 << 16) + 42
	for i := 0; i < nFiles; i++ {
		_, err := w.CreateHeader(&FileHeader{
			Name:   fmt.Sprintf("%d.dat", i),
			Method: Store, // avoid Issue 6136 and Issue 6138
		})
		if err != nil {
			t.Fatalf("creating file %d: %v", i, err)
		}
	}
	if err := w.Close(); err != nil {
		t.Fatalf("Writer.Close: %v", err)
	}
	s := buf.String()
	zr, err := NewReader(strings.NewReader(s), int64(len(s)))
	if err != nil {
		t.Fatalf("NewReader: %v", err)
	}
	if got := len(zr.File); got != nFiles {
		t.Fatalf("File contains %d files, want %d", got, nFiles)
	}
	for i := 0; i < nFiles; i++ {
		want := fmt.Sprintf("%d.dat", i)
		if zr.File[i].Name != want {
			t.Fatalf("File(%d) = %q, want %q", i, zr.File[i].Name, want)
		}
	}
}

func TestModTime(t *testing.T) {
	var testTime = time.Date(2009, time.November, 10, 23, 45, 58, 0, time.UTC)
	fh := new(FileHeader)
	fh.SetModTime(testTime)
	outTime := fh.ModTime()
	if !outTime.Equal(testTime) {
		t.Errorf("times don't match: got %s, want %s", outTime, testTime)
	}
}

func testHeaderRoundTrip(fh *FileHeader, wantUncompressedSize uint32, wantUncompressedSize64 uint64, t *testing.T) {
	fi := fh.FileInfo()
	fh2, err := FileInfoHeader(fi)
	if err != nil {
		t.Fatal(err)
	}
	if got, want := fh2.Name, fh.Name; got != want {
		t.Errorf("Name: got %s, want %s\n", got, want)
	}
	if got, want := fh2.UncompressedSize, wantUncompressedSize; got != want {
		t.Errorf("UncompressedSize: got %d, want %d\n", got, want)
	}
	if got, want := fh2.UncompressedSize64, wantUncompressedSize64; got != want {
		t.Errorf("UncompressedSize64: got %d, want %d\n", got, want)
	}
	if got, want := fh2.ModifiedTime, fh.ModifiedTime; got != want {
		t.Errorf("ModifiedTime: got %d, want %d\n", got, want)
	}
	if got, want := fh2.ModifiedDate, fh.ModifiedDate; got != want {
		t.Errorf("ModifiedDate: got %d, want %d\n", got, want)
	}

	if sysfh, ok := fi.Sys().(*FileHeader); !ok && sysfh != fh {
		t.Errorf("Sys didn't return original *FileHeader")
	}
}

func TestFileHeaderRoundTrip(t *testing.T) {
	fh := &FileHeader{
		Name:             "foo.txt",
		UncompressedSize: 987654321,
		ModifiedTime:     1234,
		ModifiedDate:     5678,
	}
	testHeaderRoundTrip(fh, fh.UncompressedSize, uint64(fh.UncompressedSize), t)
}

func TestFileHeaderRoundTrip64(t *testing.T) {
	fh := &FileHeader{
		Name:               "foo.txt",
		UncompressedSize64: 9876543210,
		ModifiedTime:       1234,
		ModifiedDate:       5678,
	}
	testHeaderRoundTrip(fh, uint32max, fh.UncompressedSize64, t)
}

func TestZip64(t *testing.T) {
	if testing.Short() {
		t.Skip("slow test; skipping")
	}
	// write 2^32 bytes plus "END\n" to a zip file
	buf := new(bytes.Buffer)
	w := NewWriter(buf)
	f, err := w.Create("huge.txt")
	if err != nil {
		t.Fatal(err)
	}
	chunk := make([]byte, 1024)
	for i := range chunk {
		chunk[i] = '.'
	}
	chunk[len(chunk)-1] = '\n'
	end := []byte("END\n")
	for i := 0; i < (1<<32)/1024; i++ {
		_, err := f.Write(chunk)
		if err != nil {
			t.Fatal("write chunk:", err)
		}
	}
	_, err = f.Write(end)
	if err != nil {
		t.Fatal("write end:", err)
	}
	if err := w.Close(); err != nil {
		t.Fatal(err)
	}

	// read back zip file and check that we get to the end of it
	r, err := NewReader(bytes.NewReader(buf.Bytes()), int64(buf.Len()))
	if err != nil {
		t.Fatal("reader:", err)
	}
	f0 := r.File[0]
	rc, err := f0.Open()
	if err != nil {
		t.Fatal("opening:", err)
	}
	for i := 0; i < (1<<32)/1024; i++ {
		_, err := io.ReadFull(rc, chunk)
		if err != nil {
			t.Fatal("read:", err)
		}
	}
	gotEnd, err := ioutil.ReadAll(rc)
	if err != nil {
		t.Fatal("read end:", err)
	}
	if !bytes.Equal(gotEnd, end) {
		t.Errorf("End of zip64 archive %q, want %q", gotEnd, end)
	}
	err = rc.Close()
	if err != nil {
		t.Fatal("closing:", err)
	}
	if got, want := f0.UncompressedSize, uint32(uint32max); got != want {
		t.Errorf("UncompressedSize %d, want %d", got, want)
	}

	if got, want := f0.UncompressedSize64, (1<<32)+uint64(len(end)); got != want {
		t.Errorf("UncompressedSize64 %d, want %d", got, want)
	}
}

func testInvalidHeader(h *FileHeader, t *testing.T) {
	var buf bytes.Buffer
	z := NewWriter(&buf)

	f, err := z.CreateHeader(h)
	if err != nil {
		t.Fatalf("error creating header: %v", err)
	}
	if _, err := f.Write([]byte("hi")); err != nil {
		t.Fatalf("error writing content: %v", err)
	}
	if err := z.Close(); err != nil {
		t.Fatalf("error closing zip writer: %v", err)
	}

	b := buf.Bytes()
	if _, err = NewReader(bytes.NewReader(b), int64(len(b))); err != ErrFormat {
		t.Fatalf("got %v, expected ErrFormat", err)
	}
}

func testValidHeader(h *FileHeader, t *testing.T) {
	var buf bytes.Buffer
	z := NewWriter(&buf)

	f, err := z.CreateHeader(h)
	if err != nil {
		t.Fatalf("error creating header: %v", err)
	}
	if _, err := f.Write([]byte("hi")); err != nil {
		t.Fatalf("error writing content: %v", err)
	}
	if err := z.Close(); err != nil {
		t.Fatalf("error closing zip writer: %v", err)
	}

	b := buf.Bytes()
	if _, err = NewReader(bytes.NewReader(b), int64(len(b))); err != nil {
		t.Fatalf("got %v, expected nil", err)
	}
}

// Issue 4302.
func TestHeaderInvalidTagAndSize(t *testing.T) {
	const timeFormat = "20060102T150405.000.txt"

	ts := time.Now()
	filename := ts.Format(timeFormat)

	h := FileHeader{
		Name:   filename,
		Method: Deflate,
		Extra:  []byte(ts.Format(time.RFC3339Nano)), // missing tag and len
	}
	h.SetModTime(ts)

	testInvalidHeader(&h, t)
}

func TestHeaderTooShort(t *testing.T) {
	h := FileHeader{
		Name:   "foo.txt",
		Method: Deflate,
		Extra:  []byte{zip64ExtraId}, // missing size
	}
	testInvalidHeader(&h, t)
}

// Issue 4393. It is valid to have an extra data header
// which contains no body.
func TestZeroLengthHeader(t *testing.T) {
	h := FileHeader{
		Name:   "extadata.txt",
		Method: Deflate,
		Extra: []byte{
			85, 84, 5, 0, 3, 154, 144, 195, 77, // tag 21589 size 5
			85, 120, 0, 0, // tag 30805 size 0
		},
	}
	testValidHeader(&h, t)
}
<commit_msg>archive/zip: speed up Zip64 test<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Tests that involve both reading and writing.

package zip

import (
	"bytes"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"sort"
	"strings"
	"testing"
	"time"
)

func TestOver65kFiles(t *testing.T) {
	buf := new(bytes.Buffer)
	w := NewWriter(buf)
	const nFiles = (1 << 16) + 42
	for i := 0; i < nFiles; i++ {
		_, err := w.CreateHeader(&FileHeader{
			Name:   fmt.Sprintf("%d.dat", i),
			Method: Store, // avoid Issue 6136 and Issue 6138
		})
		if err != nil {
			t.Fatalf("creating file %d: %v", i, err)
		}
	}
	if err := w.Close(); err != nil {
		t.Fatalf("Writer.Close: %v", err)
	}
	s := buf.String()
	zr, err := NewReader(strings.NewReader(s), int64(len(s)))
	if err != nil {
		t.Fatalf("NewReader: %v", err)
	}
	if got := len(zr.File); got != nFiles {
		t.Fatalf("File contains %d files, want %d", got, nFiles)
	}
	for i := 0; i < nFiles; i++ {
		want := fmt.Sprintf("%d.dat", i)
		if zr.File[i].Name != want {
			t.Fatalf("File(%d) = %q, want %q", i, zr.File[i].Name, want)
		}
	}
}

func TestModTime(t *testing.T) {
	var testTime = time.Date(2009, time.November, 10, 23, 45, 58, 0, time.UTC)
	fh := new(FileHeader)
	fh.SetModTime(testTime)
	outTime := fh.ModTime()
	if !outTime.Equal(testTime) {
		t.Errorf("times don't match: got %s, want %s", outTime, testTime)
	}
}

func testHeaderRoundTrip(fh *FileHeader, wantUncompressedSize uint32, wantUncompressedSize64 uint64, t *testing.T) {
	fi := fh.FileInfo()
	fh2, err := FileInfoHeader(fi)
	if err != nil {
		t.Fatal(err)
	}
	if got, want := fh2.Name, fh.Name; got != want {
		t.Errorf("Name: got %s, want %s\n", got, want)
	}
	if got, want := fh2.UncompressedSize, wantUncompressedSize; got != want {
		t.Errorf("UncompressedSize: got %d, want %d\n", got, want)
	}
	if got, want := fh2.UncompressedSize64, wantUncompressedSize64; got != want {
		t.Errorf("UncompressedSize64: got %d, want %d\n", got, want)
	}
	if got, want := fh2.ModifiedTime, fh.ModifiedTime; got != want {
		t.Errorf("ModifiedTime: got %d, want %d\n", got, want)
	}
	if got, want := fh2.ModifiedDate, fh.ModifiedDate; got != want {
		t.Errorf("ModifiedDate: got %d, want %d\n", got, want)
	}

	if sysfh, ok := fi.Sys().(*FileHeader); !ok && sysfh != fh {
		t.Errorf("Sys didn't return original *FileHeader")
	}
}

func TestFileHeaderRoundTrip(t *testing.T) {
	fh := &FileHeader{
		Name:             "foo.txt",
		UncompressedSize: 987654321,
		ModifiedTime:     1234,
		ModifiedDate:     5678,
	}
	testHeaderRoundTrip(fh, fh.UncompressedSize, uint64(fh.UncompressedSize), t)
}

func TestFileHeaderRoundTrip64(t *testing.T) {
	fh := &FileHeader{
		Name:               "foo.txt",
		UncompressedSize64: 9876543210,
		ModifiedTime:       1234,
		ModifiedDate:       5678,
	}
	testHeaderRoundTrip(fh, uint32max, fh.UncompressedSize64, t)
}

type repeatedByte struct {
	off int64
	b   byte
	n   int64
}

// rleBuffer is a run-length-encoded byte buffer.
// It's an io.Writer (like a bytes.Buffer) and also an io.ReaderAt,
// allowing random-access reads.
type rleBuffer struct {
	buf []repeatedByte
}

func (r *rleBuffer) Size() int64 {
	if len(r.buf) == 0 {
		return 0
	}
	last := &r.buf[len(r.buf)-1]
	return last.off + last.n
}

func (r *rleBuffer) Write(p []byte) (n int, err error) {
	var rp *repeatedByte
	if len(r.buf) > 0 {
		rp = &r.buf[len(r.buf)-1]
		// Fast path, if p is entirely the same byte repeated.
		if lastByte := rp.b; len(p) > 0 && p[0] == lastByte {
			all := true
			for _, b := range p {
				if b != lastByte {
					all = false
					break
				}
			}
			if all {
				rp.n += int64(len(p))
				return len(p), nil
			}
		}
	}

	for _, b := range p {
		if rp == nil || rp.b != b {
			r.buf = append(r.buf, repeatedByte{r.Size(), b, 1})
			rp = &r.buf[len(r.buf)-1]
		} else {
			rp.n++
		}
	}
	return len(p), nil
}

func (r *rleBuffer) ReadAt(p []byte, off int64) (n int, err error) {
	if len(p) == 0 {
		return
	}
	skipParts := sort.Search(len(r.buf), func(i int) bool {
		part := &r.buf[i]
		return part.off+part.n > off
	})
	parts := r.buf[skipParts:]
	if len(parts) > 0 {
		skipBytes := off - parts[0].off
		for len(parts) > 0 {
			part := parts[0]
			for i := skipBytes; i < part.n; i++ {
				if n == len(p) {
					return
				}
				p[n] = part.b
				n++
			}
			parts = parts[1:]
			skipBytes = 0
		}
	}
	if n != len(p) {
		err = io.ErrUnexpectedEOF
	}
	return
}

// Just testing the rleBuffer used in the Zip64 test above. Not used by the zip code.
func TestRLEBuffer(t *testing.T) {
	b := new(rleBuffer)
	var all []byte
	writes := []string{"abcdeee", "eeeeeee", "eeeefghaaiii"}
	for _, w := range writes {
		b.Write([]byte(w))
		all = append(all, w...)
	}
	if len(b.buf) != 10 {
		t.Fatalf("len(b.buf) = %d; want 10", len(b.buf))
	}

	for i := 0; i < len(all); i++ {
		for j := 0; j < len(all)-i; j++ {
			buf := make([]byte, j)
			n, err := b.ReadAt(buf, int64(i))
			if err != nil || n != len(buf) {
				t.Errorf("ReadAt(%d, %d) = %d, %v; want %d, nil", i, j, n, err, len(buf))
			}
			if !bytes.Equal(buf, all[i:i+j]) {
				t.Errorf("ReadAt(%d, %d) = %q; want %q", i, j, buf, all[i:i+j])
			}
		}
	}
}

// fakeHash32 is a dummy Hash32 that always returns 0.
type fakeHash32 struct {
	hash.Hash32
}

func (fakeHash32) Write(p []byte) (int, error) { return len(p), nil }
func (fakeHash32) Sum32() uint32               { return 0 }

func TestZip64(t *testing.T) {
	if testing.Short() {
		t.Skip("slow test; skipping")
	}
	const size = 1 << 32 // before the "END\n" part
	testZip64(t, size)
}

func testZip64(t testing.TB, size int64) {
	const chunkSize = 1024
	chunks := int(size / chunkSize)
	// write 2^32 bytes plus "END\n" to a zip file
	buf := new(rleBuffer)
	w := NewWriter(buf)
	f, err := w.CreateHeader(&FileHeader{
		Name:   "huge.txt",
		Method: Store,
	})
	if err != nil {
		t.Fatal(err)
	}
	f.(*fileWriter).crc32 = fakeHash32{}
	chunk := make([]byte, chunkSize)
	for i := range chunk {
		chunk[i] = '.'
	}
	for i := 0; i < chunks; i++ {
		_, err := f.Write(chunk)
		if err != nil {
			t.Fatal("write chunk:", err)
		}
	}
	end := []byte("END\n")
	_, err = f.Write(end)
	if err != nil {
		t.Fatal("write end:", err)
	}
	if err := w.Close(); err != nil {
		t.Fatal(err)
	}

	// read back zip file and check that we get to the end of it
	r, err := NewReader(buf, int64(buf.Size()))
	if err != nil {
		t.Fatal("reader:", err)
	}
	f0 := r.File[0]
	rc, err := f0.Open()
	if err != nil {
		t.Fatal("opening:", err)
	}
	rc.(*checksumReader).hash = fakeHash32{}
	for i := 0; i < chunks; i++ {
		_, err := io.ReadFull(rc, chunk)
		if err != nil {
			t.Fatal("read:", err)
		}
	}
	gotEnd, err := ioutil.ReadAll(rc)
	if err != nil {
		t.Fatal("read end:", err)
	}
	if !bytes.Equal(gotEnd, end) {
		t.Errorf("End of zip64 archive %q, want %q", gotEnd, end)
	}
	err = rc.Close()
	if err != nil {
		t.Fatal("closing:", err)
	}
	if size == 1<<32 {
		if got, want := f0.UncompressedSize, uint32(uint32max); got != want {
			t.Errorf("UncompressedSize %d, want %d", got, want)
		}
	}

	if got, want := f0.UncompressedSize64, uint64(size)+uint64(len(end)); got != want {
		t.Errorf("UncompressedSize64 %d, want %d", got, want)
	}
}

func testInvalidHeader(h *FileHeader, t *testing.T) {
	var buf bytes.Buffer
	z := NewWriter(&buf)

	f, err := z.CreateHeader(h)
	if err != nil {
		t.Fatalf("error creating header: %v", err)
	}
	if _, err := f.Write([]byte("hi")); err != nil {
		t.Fatalf("error writing content: %v", err)
	}
	if err := z.Close(); err != nil {
		t.Fatalf("error closing zip writer: %v", err)
	}

	b := buf.Bytes()
	if _, err = NewReader(bytes.NewReader(b), int64(len(b))); err != ErrFormat {
		t.Fatalf("got %v, expected ErrFormat", err)
	}
}

func testValidHeader(h *FileHeader, t *testing.T) {
	var buf bytes.Buffer
	z := NewWriter(&buf)

	f, err := z.CreateHeader(h)
	if err != nil {
		t.Fatalf("error creating header: %v", err)
	}
	if _, err := f.Write([]byte("hi")); err != nil {
		t.Fatalf("error writing content: %v", err)
	}
	if err := z.Close(); err != nil {
		t.Fatalf("error closing zip writer: %v", err)
	}

	b := buf.Bytes()
	if _, err = NewReader(bytes.NewReader(b), int64(len(b))); err != nil {
		t.Fatalf("got %v, expected nil", err)
	}
}

// Issue 4302.
func TestHeaderInvalidTagAndSize(t *testing.T) {
	const timeFormat = "20060102T150405.000.txt"

	ts := time.Now()
	filename := ts.Format(timeFormat)

	h := FileHeader{
		Name:   filename,
		Method: Deflate,
		Extra:  []byte(ts.Format(time.RFC3339Nano)), // missing tag and len
	}
	h.SetModTime(ts)

	testInvalidHeader(&h, t)
}

func TestHeaderTooShort(t *testing.T) {
	h := FileHeader{
		Name:   "foo.txt",
		Method: Deflate,
		Extra:  []byte{zip64ExtraId}, // missing size
	}
	testInvalidHeader(&h, t)
}

// Issue 4393. It is valid to have an extra data header
// which contains no body.
func TestZeroLengthHeader(t *testing.T) {
	h := FileHeader{
		Name:   "extadata.txt",
		Method: Deflate,
		Extra: []byte{
			85, 84, 5, 0, 3, 154, 144, 195, 77, // tag 21589 size 5
			85, 120, 0, 0, // tag 30805 size 0
		},
	}
	testValidHeader(&h, t)
}

// Just benchmarking how fast the Zip64 test above is. Not related to
// our zip performance, since the test above disabled CRC32 and flate.
func BenchmarkZip64Test(b *testing.B) {
	for i := 0; i < b.N; i++ {
		testZip64(b, 1<<26)
	}
}
<|endoftext|>


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

package strings_test

import (
	"fmt"
	"strings"
	"unicode"
)

func ExampleFields() {
	fmt.Printf("Fields are: %q", strings.Fields("  foo bar  baz   "))
	// Output: Fields are: ["foo" "bar" "baz"]
}

func ExampleFieldsFunc() {
	f := func(c rune) bool {
		return !unicode.IsLetter(c) && !unicode.IsNumber(c)
	}
	fmt.Printf("Fields are: %q", strings.FieldsFunc("  foo1;bar2,baz3...", f))
	// Output: Fields are: ["foo1" "bar2" "baz3"]
}

func ExampleContains() {
	fmt.Println(strings.Contains("seafood", "foo"))
	fmt.Println(strings.Contains("seafood", "bar"))
	fmt.Println(strings.Contains("seafood", ""))
	fmt.Println(strings.Contains("", ""))
	// Output:
	// true
	// false
	// true
	// true
}

func ExampleContainsAny() {
	fmt.Println(strings.ContainsAny("team", "i"))
	fmt.Println(strings.ContainsAny("failure", "u & i"))
	fmt.Println(strings.ContainsAny("foo", ""))
	fmt.Println(strings.ContainsAny("", ""))
	// Output:
	// false
	// true
	// false
	// false
}

func ExampleCount() {
	fmt.Println(strings.Count("cheese", "e"))
	fmt.Println(strings.Count("five", "")) // before & after each rune
	// Output:
	// 3
	// 5
}

func ExampleEqualFold() {
	fmt.Println(strings.EqualFold("Go", "go"))
	// Output: true
}

func ExampleIndex() {
	fmt.Println(strings.Index("chicken", "ken"))
	fmt.Println(strings.Index("chicken", "dmr"))
	// Output:
	// 4
	// -1
}

func ExampleIndexRune() {
	fmt.Println(strings.IndexRune("chicken", 'k'))
	fmt.Println(strings.IndexRune("chicken", 'd'))
	// Output:
	// 4
	// -1
}

func ExampleLastIndex() {
	fmt.Println(strings.Index("go gopher", "go"))
	fmt.Println(strings.LastIndex("go gopher", "go"))
	fmt.Println(strings.LastIndex("go gopher", "rodent"))
	// Output:
	// 0
	// 3
	// -1
}

func ExampleJoin() {
	s := []string{"foo", "bar", "baz"}
	fmt.Println(strings.Join(s, ", "))
	// Output: foo, bar, baz
}

func ExampleRepeat() {
	fmt.Println("ba" + strings.Repeat("na", 2))
	// Output: banana
}

func ExampleReplace() {
	fmt.Println(strings.Replace("oink oink oink", "k", "ky", 2))
	fmt.Println(strings.Replace("oink oink oink", "oink", "moo", -1))
	// Output:
	// oinky oinky oink
	// moo moo moo
}

func ExampleSplit() {
	fmt.Printf("%q\n", strings.Split("a,b,c", ","))
	fmt.Printf("%q\n", strings.Split("a man a plan a canal panama", "a "))
	fmt.Printf("%q\n", strings.Split(" xyz ", ""))
	fmt.Printf("%q\n", strings.Split("", "Bernardo O'Higgins"))
	// Output:
	// ["a" "b" "c"]
	// ["" "man " "plan " "canal panama"]
	// [" " "x" "y" "z" " "]
	// [""]
}

func ExampleSplitN() {
	fmt.Printf("%q\n", strings.SplitN("a,b,c", ",", 2))
	z := strings.SplitN("a,b,c", ",", 0)
	fmt.Printf("%q (nil = %v)\n", z, z == nil)
	// Output:
	// ["a" "b,c"]
	// [] (nil = true)
}

func ExampleSplitAfter() {
	fmt.Printf("%q\n", strings.SplitAfter("a,b,c", ","))
	// Output: ["a," "b," "c"]
}

func ExampleSplitAfterN() {
	fmt.Printf("%q\n", strings.SplitAfterN("a,b,c", ",", 2))
	// Output: ["a," "b,c"]
}

func ExampleTitle() {
	fmt.Println(strings.Title("her royal highness"))
	// Output: Her Royal Highness
}

func ExampleToTitle() {
	fmt.Println(strings.ToTitle("loud noises"))
	fmt.Println(strings.ToTitle("хлеб"))
	// Output:
	// LOUD NOISES
	// ХЛЕБ
}

func ExampleTrim() {
	fmt.Printf("[%q]", strings.Trim(" !!! Achtung! Achtung! !!! ", "! "))
	// Output: ["Achtung! Achtung"]
}

func ExampleMap() {
	rot13 := func(r rune) rune {
		switch {
		case r >= 'A' && r <= 'Z':
			return 'A' + (r-'A'+13)%26
		case r >= 'a' && r <= 'z':
			return 'a' + (r-'a'+13)%26
		}
		return r
	}
	fmt.Println(strings.Map(rot13, "'Twas brillig and the slithy gopher..."))
	// Output: 'Gjnf oevyyvt naq gur fyvgul tbcure...
}

func ExampleTrimSpace() {
	fmt.Println(strings.TrimSpace(" \t\n a lone gopher \n\t\r\n"))
	// Output: a lone gopher
}

func ExampleNewReplacer() {
	r := strings.NewReplacer("<", "<", ">", ">")
	fmt.Println(r.Replace("This is HTML!"))
	// Output: This is <b>HTML</b>!
}

func ExampleToUpper() {
	fmt.Println(strings.ToUpper("Gopher"))
	// Output: GOPHER
}

func ExampleToLower() {
	fmt.Println(strings.ToLower("Gopher"))
	// Output: gopher
}

func ExampleTrimSuffix() {
	var s = "Hello, goodbye, etc!"
	s = strings.TrimSuffix(s, "goodbye, etc!")
	s = strings.TrimSuffix(s, "planet")
	fmt.Print(s, "world!")
	// Output: Hello, world!
}

func ExampleTrimPrefix() {
	var s = "Goodbye,, world!"
	s = strings.TrimPrefix(s, "Goodbye,")
	s = strings.TrimPrefix(s, "Howdy,")
	fmt.Print("Hello" + s)
	// Output: Hello, world!
}
<commit_msg>strings: Add IndexFunc example<commit_after>// Copyright 2012 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package strings_test

import (
	"fmt"
	"strings"
	"unicode"
)

func ExampleFields() {
	fmt.Printf("Fields are: %q", strings.Fields("  foo bar  baz   "))
	// Output: Fields are: ["foo" "bar" "baz"]
}

func ExampleFieldsFunc() {
	f := func(c rune) bool {
		return !unicode.IsLetter(c) && !unicode.IsNumber(c)
	}
	fmt.Printf("Fields are: %q", strings.FieldsFunc("  foo1;bar2,baz3...", f))
	// Output: Fields are: ["foo1" "bar2" "baz3"]
}

func ExampleContains() {
	fmt.Println(strings.Contains("seafood", "foo"))
	fmt.Println(strings.Contains("seafood", "bar"))
	fmt.Println(strings.Contains("seafood", ""))
	fmt.Println(strings.Contains("", ""))
	// Output:
	// true
	// false
	// true
	// true
}

func ExampleContainsAny() {
	fmt.Println(strings.ContainsAny("team", "i"))
	fmt.Println(strings.ContainsAny("failure", "u & i"))
	fmt.Println(strings.ContainsAny("foo", ""))
	fmt.Println(strings.ContainsAny("", ""))
	// Output:
	// false
	// true
	// false
	// false
}

func ExampleCount() {
	fmt.Println(strings.Count("cheese", "e"))
	fmt.Println(strings.Count("five", "")) // before & after each rune
	// Output:
	// 3
	// 5
}

func ExampleEqualFold() {
	fmt.Println(strings.EqualFold("Go", "go"))
	// Output: true
}

func ExampleIndex() {
	fmt.Println(strings.Index("chicken", "ken"))
	fmt.Println(strings.Index("chicken", "dmr"))
	// Output:
	// 4
	// -1
}

func ExampleIndexFunc() {
	f := func(c rune) bool {
		return unicode.Is(unicode.Han, c)
	}
	fmt.Println(strings.IndexFunc("Hello, 世界", f))
	fmt.Println(strings.IndexFunc("Hello, world", f))
	// Output:
	// 7
	// -1
}

func ExampleIndexRune() {
	fmt.Println(strings.IndexRune("chicken", 'k'))
	fmt.Println(strings.IndexRune("chicken", 'd'))
	// Output:
	// 4
	// -1
}

func ExampleLastIndex() {
	fmt.Println(strings.Index("go gopher", "go"))
	fmt.Println(strings.LastIndex("go gopher", "go"))
	fmt.Println(strings.LastIndex("go gopher", "rodent"))
	// Output:
	// 0
	// 3
	// -1
}

func ExampleJoin() {
	s := []string{"foo", "bar", "baz"}
	fmt.Println(strings.Join(s, ", "))
	// Output: foo, bar, baz
}

func ExampleRepeat() {
	fmt.Println("ba" + strings.Repeat("na", 2))
	// Output: banana
}

func ExampleReplace() {
	fmt.Println(strings.Replace("oink oink oink", "k", "ky", 2))
	fmt.Println(strings.Replace("oink oink oink", "oink", "moo", -1))
	// Output:
	// oinky oinky oink
	// moo moo moo
}

func ExampleSplit() {
	fmt.Printf("%q\n", strings.Split("a,b,c", ","))
	fmt.Printf("%q\n", strings.Split("a man a plan a canal panama", "a "))
	fmt.Printf("%q\n", strings.Split(" xyz ", ""))
	fmt.Printf("%q\n", strings.Split("", "Bernardo O'Higgins"))
	// Output:
	// ["a" "b" "c"]
	// ["" "man " "plan " "canal panama"]
	// [" " "x" "y" "z" " "]
	// [""]
}

func ExampleSplitN() {
	fmt.Printf("%q\n", strings.SplitN("a,b,c", ",", 2))
	z := strings.SplitN("a,b,c", ",", 0)
	fmt.Printf("%q (nil = %v)\n", z, z == nil)
	// Output:
	// ["a" "b,c"]
	// [] (nil = true)
}

func ExampleSplitAfter() {
	fmt.Printf("%q\n", strings.SplitAfter("a,b,c", ","))
	// Output: ["a," "b," "c"]
}

func ExampleSplitAfterN() {
	fmt.Printf("%q\n", strings.SplitAfterN("a,b,c", ",", 2))
	// Output: ["a," "b,c"]
}

func ExampleTitle() {
	fmt.Println(strings.Title("her royal highness"))
	// Output: Her Royal Highness
}

func ExampleToTitle() {
	fmt.Println(strings.ToTitle("loud noises"))
	fmt.Println(strings.ToTitle("хлеб"))
	// Output:
	// LOUD NOISES
	// ХЛЕБ
}

func ExampleTrim() {
	fmt.Printf("[%q]", strings.Trim(" !!! Achtung! Achtung! !!! ", "! "))
	// Output: ["Achtung! Achtung"]
}

func ExampleMap() {
	rot13 := func(r rune) rune {
		switch {
		case r >= 'A' && r <= 'Z':
			return 'A' + (r-'A'+13)%26
		case r >= 'a' && r <= 'z':
			return 'a' + (r-'a'+13)%26
		}
		return r
	}
	fmt.Println(strings.Map(rot13, "'Twas brillig and the slithy gopher..."))
	// Output: 'Gjnf oevyyvt naq gur fyvgul tbcure...
}

func ExampleTrimSpace() {
	fmt.Println(strings.TrimSpace(" \t\n a lone gopher \n\t\r\n"))
	// Output: a lone gopher
}

func ExampleNewReplacer() {
	r := strings.NewReplacer("<", "<", ">", ">")
	fmt.Println(r.Replace("This is HTML!"))
	// Output: This is <b>HTML</b>!
}

func ExampleToUpper() {
	fmt.Println(strings.ToUpper("Gopher"))
	// Output: GOPHER
}

func ExampleToLower() {
	fmt.Println(strings.ToLower("Gopher"))
	// Output: gopher
}

func ExampleTrimSuffix() {
	var s = "Hello, goodbye, etc!"
	s = strings.TrimSuffix(s, "goodbye, etc!")
	s = strings.TrimSuffix(s, "planet")
	fmt.Print(s, "world!")
	// Output: Hello, world!
}

func ExampleTrimPrefix() {
	var s = "Goodbye,, world!"
	s = strings.TrimPrefix(s, "Goodbye,")
	s = strings.TrimPrefix(s, "Howdy,")
	fmt.Print("Hello" + s)
	// Output: Hello, world!
}
<|endoftext|>


<commit_before>package cloudconfig

const (
	decryptTLSAssetsScriptTemplate = `#!/bin/bash -e

rkt run \
  --volume=ssl,kind=host,source=/etc/kubernetes/ssl,readOnly=false \
  --mount=volume=ssl,target=/etc/kubernetes/ssl \
  --uuid-file-save=/var/run/coreos/decrypt-tls-assets.uuid \
  --volume=dns,kind=host,source=/etc/resolv.conf,readOnly=true --mount volume=dns,target=/etc/resolv.conf \
  --net=host \
  --trust-keys-from-https \
  quay.io/coreos/awscli:025a357f05242fdad6a81e8a6b520098aa65a600 --exec=/bin/bash -- \
    -ec \
    'echo decrypting tls assets
    shopt -s nullglob
    for encKey in $(find /etc/kubernetes/ssl -name "*.pem.enc"); do
      echo decrypting $encKey
      f=$(mktemp $encKey.XXXXXXXX)
      /usr/bin/aws \
        --region {{.AWS.Region}} kms decrypt \
        --ciphertext-blob fileb://$encKey \
        --output text \
        --query Plaintext \
      | base64 -d > $f
      mv -f $f ${encKey%.enc}
    done;'

rkt rm --uuid-file=/var/run/coreos/decrypt-tls-assets.uuid || :

chown -R etcd:etcd /etc/kubernetes/ssl/etcd`

	decryptTLSAssetsServiceTemplate = `
[Unit]
Description=Decrypt TLS certificates

[Service]
Type=oneshot
ExecStart=/opt/bin/decrypt-tls-assets

[Install]
WantedBy=multi-user.target
`

	decryptKeysAssetsScriptTemplate = `#!/bin/bash -e

rkt run \
  --volume=keys,kind=host,source=/etc/kubernetes/encryption,readOnly=false \
  --mount=volume=keys,target=/etc/kubernetes/encryption \
  --uuid-file-save=/var/run/coreos/decrypt-keys-assets.uuid \
  --volume=dns,kind=host,source=/etc/resolv.conf,readOnly=true --mount volume=dns,target=/etc/resolv.conf \
  --net=host \
  --trust-keys-from-https \
  quay.io/coreos/awscli:025a357f05242fdad6a81e8a6b520098aa65a600 --exec=/bin/bash -- \
    -ec \
    'echo decrypting keys assets
    shopt -s nullglob
    for encKey in $(find /etc/kubernetes/encryption -name "*.enc"); do
      echo decrypting $encKey
      f=$(mktemp $encKey.XXXXXXXX)
      /usr/bin/aws \
        --region {{.AWS.Region}} kms decrypt \
        --ciphertext-blob fileb://$encKey \
        --output text \
        --query Plaintext \
      | base64 -d > $f
      mv -f $f ${encKey%.enc}
    done;
    echo done.'

rkt rm --uuid-file=/var/run/coreos/decrypt-keys-assets.uuid || :
`

	decryptKeysServiceTemplate = `
[Unit]
Description=Decrypt Secret Keys

[Service]
Type=oneshot
ExecStart=/opt/bin/decrypt-keys-assets

[Install]
WantedBy=multi-user.target
`

	masterFormatVarLibDockerServiceTemplate = `
[Unit]
Description=Format /var/lib/docker to XFS
Before=docker.service var-lib-docker.mount
ConditionPathExists=!/var/lib/docker

[Service]
Type=oneshot
ExecStart=/usr/sbin/mkfs.xfs -f /dev/xvdb

[Install]
WantedBy=multi-user.target
`

	workerFormatVarLibDockerServiceTemplate = `
[Unit]
Description=Format /var/lib/docker to XFS
Before=docker.service var-lib-docker.mount
ConditionPathExists=!/var/lib/docker

[Service]
Type=oneshot
ExecStart=/usr/sbin/mkfs.xfs -f /dev/xvdh

[Install]
WantedBy=multi-user.target
`

	ephemeralVarLibDockerMountTemplate = `
[Unit]
Description=Mount ephemeral volume on /var/lib/docker

[Mount]
What=/dev/xvdb
Where=/var/lib/docker
Type=xfs

[Install]
RequiredBy=local-fs.target
`
	persistentVarLibDockerMountTemplate = `
[Unit]
Description=Mount persistent volume on /var/lib/docker

[Mount]
What=/dev/xvdh
Where=/var/lib/docker
Type=xfs

[Install]
RequiredBy=local-fs.target
`

	waitDockerConfTemplate = `
[Unit]
After=var-lib-docker.mount
Requires=var-lib-docker.mount
`

	instanceStorageTemplate = `
storage:
  filesystems:
    - name: ephemeral1
      mount:
        device: /dev/xvdb
        format: xfs
        create:
          force: true
`

	instanceStorageClassTemplate = `
write_files:
- path: /srv/default-storage-class.yaml
  owner: root
  permissions: 644
  content: |
    apiVersion: storage.k8s.io/v1beta1
    kind: StorageClass
    metadata:
      name: gp2
      annotations:
        storageclass.beta.kubernetes.io/is-default-class: "true"
      labels:
        kubernetes.io/cluster-service: "true"
        addonmanager.kubernetes.io/mode: EnsureExists
    provisioner: kubernetes.io/aws-ebs
    parameters:
      type: gp2
      encrypted: "true"
`

	ingressControllerConfigMapTemplate = `kind: ConfigMap
apiVersion: v1
metadata:
  name: ingress-nginx
  namespace: kube-system
  labels:
    k8s-addon: ingress-nginx.addons.k8s.io
data:
  server-name-hash-bucket-size: "1024"
  server-name-hash-max-size: "1024"
  use-proxy-protocol: "true"
`
)
<commit_msg>fix-var-lib-mount3 (#505)<commit_after>package cloudconfig

const (
	decryptTLSAssetsScriptTemplate = `#!/bin/bash -e

rkt run \
  --volume=ssl,kind=host,source=/etc/kubernetes/ssl,readOnly=false \
  --mount=volume=ssl,target=/etc/kubernetes/ssl \
  --uuid-file-save=/var/run/coreos/decrypt-tls-assets.uuid \
  --volume=dns,kind=host,source=/etc/resolv.conf,readOnly=true --mount volume=dns,target=/etc/resolv.conf \
  --net=host \
  --trust-keys-from-https \
  quay.io/coreos/awscli:025a357f05242fdad6a81e8a6b520098aa65a600 --exec=/bin/bash -- \
    -ec \
    'echo decrypting tls assets
    shopt -s nullglob
    for encKey in $(find /etc/kubernetes/ssl -name "*.pem.enc"); do
      echo decrypting $encKey
      f=$(mktemp $encKey.XXXXXXXX)
      /usr/bin/aws \
        --region {{.AWS.Region}} kms decrypt \
        --ciphertext-blob fileb://$encKey \
        --output text \
        --query Plaintext \
      | base64 -d > $f
      mv -f $f ${encKey%.enc}
    done;'

rkt rm --uuid-file=/var/run/coreos/decrypt-tls-assets.uuid || :

chown -R etcd:etcd /etc/kubernetes/ssl/etcd`

	decryptTLSAssetsServiceTemplate = `
[Unit]
Description=Decrypt TLS certificates

[Service]
Type=oneshot
ExecStart=/opt/bin/decrypt-tls-assets

[Install]
WantedBy=multi-user.target
`

	decryptKeysAssetsScriptTemplate = `#!/bin/bash -e

rkt run \
  --volume=keys,kind=host,source=/etc/kubernetes/encryption,readOnly=false \
  --mount=volume=keys,target=/etc/kubernetes/encryption \
  --uuid-file-save=/var/run/coreos/decrypt-keys-assets.uuid \
  --volume=dns,kind=host,source=/etc/resolv.conf,readOnly=true --mount volume=dns,target=/etc/resolv.conf \
  --net=host \
  --trust-keys-from-https \
  quay.io/coreos/awscli:025a357f05242fdad6a81e8a6b520098aa65a600 --exec=/bin/bash -- \
    -ec \
    'echo decrypting keys assets
    shopt -s nullglob
    for encKey in $(find /etc/kubernetes/encryption -name "*.enc"); do
      echo decrypting $encKey
      f=$(mktemp $encKey.XXXXXXXX)
      /usr/bin/aws \
        --region {{.AWS.Region}} kms decrypt \
        --ciphertext-blob fileb://$encKey \
        --output text \
        --query Plaintext \
      | base64 -d > $f
      mv -f $f ${encKey%.enc}
    done;
    echo done.'

rkt rm --uuid-file=/var/run/coreos/decrypt-keys-assets.uuid || :
`

	decryptKeysServiceTemplate = `
[Unit]
Description=Decrypt Secret Keys

[Service]
Type=oneshot
ExecStart=/opt/bin/decrypt-keys-assets

[Install]
WantedBy=multi-user.target
`

	masterFormatVarLibDockerServiceTemplate = `
[Unit]
Description=Format /var/lib/docker to XFS
Before=docker.service var-lib-docker.mount

[Service]
Type=oneshot
ExecStart=/usr/sbin/mkfs.xfs -f /dev/xvdb

[Install]
WantedBy=multi-user.target
`

	workerFormatVarLibDockerServiceTemplate = `
[Unit]
Description=Format /var/lib/docker to XFS
Before=docker.service var-lib-docker.mount
ConditionPathExists=!/var/lib/docker

[Service]
Type=oneshot
ExecStart=/usr/sbin/mkfs.xfs -f /dev/xvdh

[Install]
WantedBy=multi-user.target
`

	ephemeralVarLibDockerMountTemplate = `
[Unit]
Description=Mount ephemeral volume on /var/lib/docker

[Mount]
What=/dev/xvdb
Where=/var/lib/docker
Type=xfs

[Install]
RequiredBy=local-fs.target
`
	persistentVarLibDockerMountTemplate = `
[Unit]
Description=Mount persistent volume on /var/lib/docker

[Mount]
What=/dev/xvdh
Where=/var/lib/docker
Type=xfs

[Install]
RequiredBy=local-fs.target
`

	waitDockerConfTemplate = `
[Unit]
After=var-lib-docker.mount
Requires=var-lib-docker.mount
`

	instanceStorageTemplate = `
storage:
  filesystems:
    - name: ephemeral1
      mount:
        device: /dev/xvdb
        format: xfs
        create:
          force: true
`

	instanceStorageClassTemplate = `
write_files:
- path: /srv/default-storage-class.yaml
  owner: root
  permissions: 644
  content: |
    apiVersion: storage.k8s.io/v1beta1
    kind: StorageClass
    metadata:
      name: gp2
      annotations:
        storageclass.beta.kubernetes.io/is-default-class: "true"
      labels:
        kubernetes.io/cluster-service: "true"
        addonmanager.kubernetes.io/mode: EnsureExists
    provisioner: kubernetes.io/aws-ebs
    parameters:
      type: gp2
      encrypted: "true"
`

	ingressControllerConfigMapTemplate = `kind: ConfigMap
apiVersion: v1
metadata:
  name: ingress-nginx
  namespace: kube-system
  labels:
    k8s-addon: ingress-nginx.addons.k8s.io
data:
  server-name-hash-bucket-size: "1024"
  server-name-hash-max-size: "1024"
  use-proxy-protocol: "true"
`
)
<|endoftext|>


<commit_before>package servicebuilder

import (
	"log"
	"path/filepath"
	"strings"

	"github.com/bndr/gopencils"
	"github.com/essentier/nomockutil"
	"github.com/essentier/spickspan/config"
	"github.com/go-errors/errors"
)

func createServiceBuilder(serviceConfig config.Service, providerUrl string, token string) *serviceBuilder {
	nomockApi := gopencils.Api(providerUrl)
	return &serviceBuilder{nomockApi: nomockApi, token: token, serviceConfig: serviceConfig}
}

type serviceBuilder struct {
	nomockApi     *gopencils.Resource
	serviceConfig config.Service
	token         string
}

func (p *serviceBuilder) buildService() error {
	if !p.serviceConfig.IsSourceProject() {
		return errors.New("Could not build a service that is not a source project.")
	}

	err := gitPush(p.serviceConfig.ServiceName, p.serviceConfig.ProjectSrcRoot, p.nomockApi.Api.BaseUrl.String()+"/nomockbuilder", p.token)
	if err != nil {
		return err
	}

	log.Printf("going to build code")
	return p.buildServiceOnNomockBuilder(p.serviceConfig.ServiceName)
}

func (p *serviceBuilder) buildServiceOnNomockBuilder(serviceName string) error {
	log.Printf("building service %v on nomock builder", serviceName)
	var result interface{}
	builderResource := p.nomockApi.Res("nomockbuilder/build/"+serviceName, &result)
	builderResource.SetHeader("Authorization", "Bearer "+p.token)
	_, err := builderResource.Get()
	return err
}

func gitPush(serviceName string, projectDir string, builderUrl string, token string) error {
	log.Printf("gitpusher.Push project dir: %v", projectDir)
	gitDir := filepath.Join(projectDir, ".git")
	if !nomockutil.Exists(gitDir) {
		return errors.Errorf("Project %v is not initialized with git. Use 'git init' to initialize the project.", projectDir)
	}

	remoteUrl := getEssentierGitRemote(serviceName, builderUrl, token)
	git := &gitProject{projectDir: projectDir, err: nil}
	originalBranch := git.getCurrentBranch()
	git.stashAll()
	if git.err == nil {
		defer git.deferredPopStashed()
	}

	git.branch("nomock")
	if git.err == nil {
		defer git.deferredDeleteBranch("nomock")
	}

	git.checkout("nomock")
	if git.err == nil {
		defer git.deferredCheckout(originalBranch)
	}

	git.pull(remoteUrl, "nomock")
	git.applyStash()
	git.addAll()
	git.commit("'done by nomock'")
	git.push(remoteUrl, "nomock")
	return git.err
}

func getEssentierGitRemote(serviceName string, builderUrl string, token string) string {
	remoteUrl := ""
	if strings.HasPrefix(builderUrl, "git") {
		remoteUrl = builderUrl + ":" + serviceName
	} else if strings.HasPrefix(builderUrl, "https://") {
		if strings.HasSuffix(builderUrl, "/") {
			remoteUrl = "https://" + token + ":@" + builderUrl[8:] + serviceName
		} else {
			remoteUrl = "https://" + token + ":@" + builderUrl[8:] + "/" + serviceName
		}
	}
	return remoteUrl
}
<commit_msg>use essentier gopencils<commit_after>package servicebuilder

import (
	"log"
	"path/filepath"
	"strings"

	"github.com/essentier/gopencils"
	"github.com/essentier/nomockutil"
	"github.com/essentier/spickspan/config"
	"github.com/go-errors/errors"
)

func createServiceBuilder(serviceConfig config.Service, providerUrl string, token string) *serviceBuilder {
	nomockApi := gopencils.Api(providerUrl)
	return &serviceBuilder{nomockApi: nomockApi, token: token, serviceConfig: serviceConfig}
}

type serviceBuilder struct {
	nomockApi     *gopencils.Resource
	serviceConfig config.Service
	token         string
}

func (p *serviceBuilder) buildService() error {
	if !p.serviceConfig.IsSourceProject() {
		return errors.New("Could not build a service that is not a source project.")
	}

	err := gitPush(p.serviceConfig.ServiceName, p.serviceConfig.ProjectSrcRoot, p.nomockApi.Api.BaseUrl.String()+"/nomockbuilder", p.token)
	if err != nil {
		return err
	}

	log.Printf("going to build code")
	return p.buildServiceOnNomockBuilder(p.serviceConfig.ServiceName)
}

func (p *serviceBuilder) buildServiceOnNomockBuilder(serviceName string) error {
	log.Printf("building service %v on nomock builder", serviceName)
	var result interface{}
	builderResource := p.nomockApi.Res("nomockbuilder/build/"+serviceName, &result)
	builderResource.SetHeader("Authorization", "Bearer "+p.token)
	_, err := builderResource.Get()
	return err
}

func gitPush(serviceName string, projectDir string, builderUrl string, token string) error {
	log.Printf("gitpusher.Push project dir: %v", projectDir)
	gitDir := filepath.Join(projectDir, ".git")
	if !nomockutil.Exists(gitDir) {
		return errors.Errorf("Project %v is not initialized with git. Use 'git init' to initialize the project.", projectDir)
	}

	remoteUrl := getEssentierGitRemote(serviceName, builderUrl, token)
	git := &gitProject{projectDir: projectDir, err: nil}
	originalBranch := git.getCurrentBranch()
	git.stashAll()
	if git.err == nil {
		defer git.deferredPopStashed()
	}

	git.branch("nomock")
	if git.err == nil {
		defer git.deferredDeleteBranch("nomock")
	}

	git.checkout("nomock")
	if git.err == nil {
		defer git.deferredCheckout(originalBranch)
	}

	git.pull(remoteUrl, "nomock")
	git.applyStash()
	git.addAll()
	git.commit("'done by nomock'")
	git.push(remoteUrl, "nomock")
	return git.err
}

func getEssentierGitRemote(serviceName string, builderUrl string, token string) string {
	remoteUrl := ""
	if strings.HasPrefix(builderUrl, "git") {
		remoteUrl = builderUrl + ":" + serviceName
	} else if strings.HasPrefix(builderUrl, "https://") {
		if strings.HasSuffix(builderUrl, "/") {
			remoteUrl = "https://" + token + ":@" + builderUrl[8:] + serviceName
		} else {
			remoteUrl = "https://" + token + ":@" + builderUrl[8:] + "/" + serviceName
		}
	}
	return remoteUrl
}
<|endoftext|>


<commit_before>package client

import (
	"log"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/hashicorp/nomad/client/driver"
	"github.com/hashicorp/nomad/nomad/mock"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/nomad/testutil"

	ctestutil "github.com/hashicorp/nomad/client/testutil"
)

func testLogger() *log.Logger {
	return log.New(os.Stderr, "", log.LstdFlags)
}

type MockTaskStateUpdater struct {
	Count       int
	Name        []string
	Status      []string
	Description []string
}

func (m *MockTaskStateUpdater) Update(name, status, desc string) {
	m.Count += 1
	m.Name = append(m.Name, name)
	m.Status = append(m.Status, status)
	m.Description = append(m.Description, desc)
}

func testTaskRunner() (*MockTaskStateUpdater, *TaskRunner) {
	logger := testLogger()
	conf := DefaultConfig()
	conf.StateDir = "/tmp"
	upd := &MockTaskStateUpdater{}
	ctx := driver.NewExecContext()
	alloc := mock.Alloc()
	task := alloc.Job.TaskGroups[0].Tasks[0]
	tr := NewTaskRunner(logger, conf, upd.Update, ctx, alloc.ID, task)
	return upd, tr
}

func testTaskRunnerPostOffer(*MockTaskStateUpdater, *TaskRunner) {
	state, runner := testTaskRunner()
	// testTaskRunner gives us a job in pre-offer state. We need it in post-
	// offer state some some other things are initialized. We'll init them here.
	runner.task.Resources.Networks[0].ReservedPorts[0] = 80
	return state, runner
}

func TestTaskRunner_SimpleRun(t *testing.T) {
	ctestutil.ExecCompatible(t)
	upd, tr := testTaskRunnerPostOffer()
	go tr.Run()
	defer tr.Destroy()

	select {
	case <-tr.WaitCh():
	case <-time.After(2 * time.Second):
		t.Fatalf("timeout")
	}

	if upd.Count != 2 {
		t.Fatalf("should have 2 updates: %#v", upd)
	}
	if upd.Name[0] != tr.task.Name {
		t.Fatalf("bad: %#v", upd.Name)
	}
	if upd.Status[0] != structs.AllocClientStatusRunning {
		t.Fatalf("bad: %#v", upd.Status)
	}
	if upd.Description[0] != "task started" {
		t.Fatalf("bad: %#v", upd.Description)
	}

	if upd.Name[1] != tr.task.Name {
		t.Fatalf("bad: %#v", upd.Name)
	}
	if upd.Status[1] != structs.AllocClientStatusDead {
		t.Fatalf("bad: %#v", upd.Status)
	}
	if upd.Description[1] != "task completed" {
		t.Fatalf("bad: %#v", upd.Description)
	}
}

func TestTaskRunner_Destroy(t *testing.T) {
	ctestutil.ExecCompatible(t)
	upd, tr := testTaskRunner()

	// Change command to ensure we run for a bit
	tr.task.Config["command"] = "/bin/sleep"
	tr.task.Config["args"] = "10"
	go tr.Run()

	// Begin the tear down
	go func() {
		time.Sleep(100 * time.Millisecond)
		tr.Destroy()
	}()

	select {
	case <-tr.WaitCh():
	case <-time.After(2 * time.Second):
		t.Fatalf("timeout")
	}

	if upd.Count != 2 {
		t.Fatalf("should have 2 updates: %#v", upd)
	}
	if upd.Status[0] != structs.AllocClientStatusRunning {
		t.Fatalf("bad: %#v", upd.Status)
	}
	if upd.Status[1] != structs.AllocClientStatusDead {
		t.Fatalf("bad: %#v", upd.Status)
	}
	if !strings.Contains(upd.Description[1], "task failed") {
		t.Fatalf("bad: %#v", upd.Description)
	}
}

func TestTaskRunner_Update(t *testing.T) {
	ctestutil.ExecCompatible(t)
	_, tr := testTaskRunner()

	// Change command to ensure we run for a bit
	tr.task.Config["command"] = "/bin/sleep"
	tr.task.Config["args"] = "10"
	go tr.Run()
	defer tr.Destroy()

	// Update the task definition
	newTask := new(structs.Task)
	*newTask = *tr.task
	newTask.Driver = "foobar"
	tr.Update(newTask)

	// Wait for update to take place
	testutil.WaitForResult(func() (bool, error) {
		return tr.task == newTask, nil
	}, func(err error) {
		t.Fatalf("err: %v", err)
	})
}

/*
TODO: This test is disabled til a follow-up api changes the restore state interface.
The driver/executor interface will be changed from Open to Cleanup, in which
clean-up tears down previous allocs.

func TestTaskRunner_SaveRestoreState(t *testing.T) {
	upd, tr := testTaskRunner()

	// Change command to ensure we run for a bit
	tr.task.Config["command"] = "/bin/sleep"
	tr.task.Config["args"] = "10"
	go tr.Run()
	defer tr.Destroy()

	// Snapshot state
	time.Sleep(200 * time.Millisecond)
	err := tr.SaveState()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create a new task runner
	tr2 := NewTaskRunner(tr.logger, tr.config, upd.Update,
		tr.ctx, tr.allocID, &structs.Task{Name: tr.task.Name})
	err = tr2.RestoreState()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	go tr2.Run()
	defer tr2.Destroy()

	// Destroy and wait
	tr2.Destroy()

	select {
	case <-tr.WaitCh():
	case <-time.After(2 * time.Second):
		t.Fatalf("timeout")
	}
}
*/
<commit_msg>Missed a set of parens<commit_after>package client

import (
	"log"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/hashicorp/nomad/client/driver"
	"github.com/hashicorp/nomad/nomad/mock"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/nomad/testutil"

	ctestutil "github.com/hashicorp/nomad/client/testutil"
)

func testLogger() *log.Logger {
	return log.New(os.Stderr, "", log.LstdFlags)
}

type MockTaskStateUpdater struct {
	Count       int
	Name        []string
	Status      []string
	Description []string
}

func (m *MockTaskStateUpdater) Update(name, status, desc string) {
	m.Count += 1
	m.Name = append(m.Name, name)
	m.Status = append(m.Status, status)
	m.Description = append(m.Description, desc)
}

func testTaskRunner() (*MockTaskStateUpdater, *TaskRunner) {
	logger := testLogger()
	conf := DefaultConfig()
	conf.StateDir = "/tmp"
	upd := &MockTaskStateUpdater{}
	ctx := driver.NewExecContext()
	alloc := mock.Alloc()
	task := alloc.Job.TaskGroups[0].Tasks[0]
	tr := NewTaskRunner(logger, conf, upd.Update, ctx, alloc.ID, task)
	return upd, tr
}

func testTaskRunnerPostOffer() (*MockTaskStateUpdater, *TaskRunner) {
	state, runner := testTaskRunner()
	// testTaskRunner gives us a job in pre-offer state. We need it in post-
	// offer state some some other things are initialized. We'll init them here.
	runner.task.Resources.Networks[0].ReservedPorts[0] = 80
	return state, runner
}

func TestTaskRunner_SimpleRun(t *testing.T) {
	ctestutil.ExecCompatible(t)
	upd, tr := testTaskRunnerPostOffer()
	go tr.Run()
	defer tr.Destroy()

	select {
	case <-tr.WaitCh():
	case <-time.After(2 * time.Second):
		t.Fatalf("timeout")
	}

	if upd.Count != 2 {
		t.Fatalf("should have 2 updates: %#v", upd)
	}
	if upd.Name[0] != tr.task.Name {
		t.Fatalf("bad: %#v", upd.Name)
	}
	if upd.Status[0] != structs.AllocClientStatusRunning {
		t.Fatalf("bad: %#v", upd.Status)
	}
	if upd.Description[0] != "task started" {
		t.Fatalf("bad: %#v", upd.Description)
	}

	if upd.Name[1] != tr.task.Name {
		t.Fatalf("bad: %#v", upd.Name)
	}
	if upd.Status[1] != structs.AllocClientStatusDead {
		t.Fatalf("bad: %#v", upd.Status)
	}
	if upd.Description[1] != "task completed" {
		t.Fatalf("bad: %#v", upd.Description)
	}
}

func TestTaskRunner_Destroy(t *testing.T) {
	ctestutil.ExecCompatible(t)
	upd, tr := testTaskRunner()

	// Change command to ensure we run for a bit
	tr.task.Config["command"] = "/bin/sleep"
	tr.task.Config["args"] = "10"
	go tr.Run()

	// Begin the tear down
	go func() {
		time.Sleep(100 * time.Millisecond)
		tr.Destroy()
	}()

	select {
	case <-tr.WaitCh():
	case <-time.After(2 * time.Second):
		t.Fatalf("timeout")
	}

	if upd.Count != 2 {
		t.Fatalf("should have 2 updates: %#v", upd)
	}
	if upd.Status[0] != structs.AllocClientStatusRunning {
		t.Fatalf("bad: %#v", upd.Status)
	}
	if upd.Status[1] != structs.AllocClientStatusDead {
		t.Fatalf("bad: %#v", upd.Status)
	}
	if !strings.Contains(upd.Description[1], "task failed") {
		t.Fatalf("bad: %#v", upd.Description)
	}
}

func TestTaskRunner_Update(t *testing.T) {
	ctestutil.ExecCompatible(t)
	_, tr := testTaskRunner()

	// Change command to ensure we run for a bit
	tr.task.Config["command"] = "/bin/sleep"
	tr.task.Config["args"] = "10"
	go tr.Run()
	defer tr.Destroy()

	// Update the task definition
	newTask := new(structs.Task)
	*newTask = *tr.task
	newTask.Driver = "foobar"
	tr.Update(newTask)

	// Wait for update to take place
	testutil.WaitForResult(func() (bool, error) {
		return tr.task == newTask, nil
	}, func(err error) {
		t.Fatalf("err: %v", err)
	})
}

/*
TODO: This test is disabled til a follow-up api changes the restore state interface.
The driver/executor interface will be changed from Open to Cleanup, in which
clean-up tears down previous allocs.

func TestTaskRunner_SaveRestoreState(t *testing.T) {
	upd, tr := testTaskRunner()

	// Change command to ensure we run for a bit
	tr.task.Config["command"] = "/bin/sleep"
	tr.task.Config["args"] = "10"
	go tr.Run()
	defer tr.Destroy()

	// Snapshot state
	time.Sleep(200 * time.Millisecond)
	err := tr.SaveState()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create a new task runner
	tr2 := NewTaskRunner(tr.logger, tr.config, upd.Update,
		tr.ctx, tr.allocID, &structs.Task{Name: tr.task.Name})
	err = tr2.RestoreState()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	go tr2.Run()
	defer tr2.Destroy()

	// Destroy and wait
	tr2.Destroy()

	select {
	case <-tr.WaitCh():
	case <-time.After(2 * time.Second):
		t.Fatalf("timeout")
	}
}
*/
<|endoftext|>