<commit_before>// Copyright 2013

/*
EVE Online API downloader/cacher.

Begin by initializing the default client.
	NewClient(NilCache)

Create a new request for the API page needed.
	req := NewRequest("eve/ConquerableStationList.xml.aspx")

Set any options you may need, including keyid/vcode.
	req.Set("keyid", "1234")
	req.Set("vcode", "abcd")
	req.Set("charactername", "innominate")
	req.Set("characterid", fmt.Sprintf("%d", 123))

Get your response.
	resp, err := req.Do()
	xml := resp.Data
*/
package apicache

import (
	"crypto/rand"
	"encoding/xml"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"sync"
	"time"
)

const sqlDateTime = "2006-01-02 15:04:05"

var DebugLog = log.New(ioutil.Discard, "apicache", log.Ldate|log.Ltime)

//
//  API Error Type, hopefully contains the error code and some useful information.
//  This is CCP so nothing is guaranteed.
//
type APIError struct {
	ErrorCode int    `xml:"code,attr"`
	ErrorText string `xml:",innerxml"`
}

// Error String Generator
func (e APIError) Error() string {
	return fmt.Sprintf("API Error %d: %s", e.ErrorCode, e.ErrorText)
}

// API Client structure.  Must be created using NewClient() function.
type Client struct {
	// Base URL, defaults to CCPs api but can be changed to a proxy
	BaseURL string

	// Default three retries, can be changed at will.
	Retries int

	timeout      time.Duration
	maxIdleConns int

	cacher     Cacher
	httpClient *http.Client

	panicUntil  time.Time
	panicCode   int
	panicReason string

	sync.RWMutex
}

// Default client, first client created lives here as well.
var client *Client

// Default BaseURL for new clients, change as necessary.  Must contain
// trailing slash or bad things happen.
var DefaultBaseURL = "https://api.eveonline.com/"

// Return the default client for
func GetDefaultClient() *Client {
	if client == nil {
		panic("Tried to get nonexistent client, must be initialized first.")
	}

	return client
}

// Create a new API Client.  The first time this is called will become
// the default client.  Requires a cacher.
func NewClient(cacher Cacher) *Client {
	var newClient Client

	newClient.BaseURL = DefaultBaseURL
	newClient.Retries = 5
	newClient.cacher = cacher
	newClient.maxIdleConns = 2

	// Also sets up our initial http client
	newClient.SetTimeout(60 * time.Second)

	if client == nil {
		client = &newClient
	}
	return &newClient
}

func (c *Client) SetMaxIdleConns(maxIdleConns int) {
	// Enforce some sanity.
	if maxIdleConns <= 0 || maxIdleConns >= 64 {
		maxIdleConns = 2
	}
	c.maxIdleConns = maxIdleConns
	c.newHttpClient()
}

// Set timeout for each API request.
func (c *Client) SetTimeout(timeout time.Duration) {
	if timeout.Seconds() <= 0 || timeout.Seconds() > 3600 {
		timeout = 60 * time.Second
	}

	c.timeout = timeout
	c.newHttpClient()
}

// Set max idle conns for the default client
func SetMaxIdleConns(maxIdleConns int) {
	client.SetMaxIdleConns(maxIdleConns)
}

// Set timeout for default client.
func SetTimeout(timeout time.Duration) {
	client.SetTimeout(timeout)
}

func (c *Client) newHttpClient() {
	c.httpClient = &http.Client{
		Transport: &http.Transport{
			Dial: func(netw, addr string) (net.Conn, error) {
				c, err := net.DialTimeout(netw, addr, c.timeout)
				if err != nil {
					return nil, err
				}
				//c.SetDeadline(deadline)
				return c, nil
			},
			ResponseHeaderTimeout: c.timeout,
			MaxIdleConnsPerHost:   c.maxIdleConns,
		},
	}
}

// Create a new request.
func (c *Client) NewRequest(url string) *Request {
	var r Request
	r.params = make(map[string]string)

	if url[0] == '/' {
		url = url[1:]
	}
	r.url = url
	r.client = c

	return &r
}

// Create a new request using default client.
func NewRequest(url string) *Request {
	return client.NewRequest(url)
}

var (
	ErrCannotConnect = fmt.Errorf("Error connecting to API.")
	ErrNetwork       = fmt.Errorf("Network error.")
	ErrHTTP          = fmt.Errorf("HTTP error.")
	ErrForbidden     = fmt.Errorf("HTTP Forbidden, invalid API provided.")
	ErrUnknown       = fmt.Errorf("Unknown Error.")
	ErrXML           = fmt.Errorf("Malformed XML Detected")
	ErrTime          = fmt.Errorf("Malformed cache directive.")
)

type cacheResp struct {
	Error       APIError `xml:"error"`
	CachedUntil string   `xml:"cachedUntil"`
}

// Data structure returned from the API.
type Response struct {
	// Raw XML data
	Data []byte

	// Signals if we used the cache or not
	FromCache bool

	// Data expiration time
	Expires time.Time

	// true if an error occured due to invalid API rather than server problems
	Invalidate bool

	// Contains API Error if one occured
	Error APIError

	//Pass on CCP's HTTP code because why not?
	HTTPCode int
}

func (c *Client) GetCached(r *Request) (retresp *Response, reterr error) {
	resp := &Response{}

	// Check for cached version
	cacheTag := r.cacheTag()
	httpCode, data, expires, err := c.cacher.Get(cacheTag)
	if err == nil && !r.Force && !r.NoCache {
		resp.Data = data
		resp.FromCache = true
		resp.Expires = expires
		resp.HTTPCode = httpCode

		var cR cacheResp
		err = xml.Unmarshal(data, &cR)
		if err != nil {
			DebugLog.Printf("XML Error: %s", err)
			return resp, ErrXML
		}

		// Handle extended expiration requests.
		if r.Expires.After(expires) {
			resp.Expires = r.Expires
		} else {
			resp.Expires = expires
		}

		// Pass on any API errors
		resp.Error = cR.Error

		return resp, nil
	}
	return resp, err
}

func MakeID() string {
	buf := make([]byte, 5)
	io.ReadFull(rand.Reader, buf)
	return fmt.Sprintf("%x", buf)
}

// Perform a request, usually called by the request itself.
// User friendly error is enclosed in the response, returned error should be
// for internal use only.
func (c *Client) Do(r *Request) (retresp *Response, reterr error) {
	var data []byte
	var expires time.Time

	cacheTag := r.cacheTag()

	// Check for cached version
	resp, err := c.GetCached(r)
	if err == nil {
		return resp, nil
	}

	// If we're panicking, bail out early and spit back a fake error
	c.RLock()
	if c.panicUntil.After(time.Now()) {
		DebugLog.Printf("Got Request, but we're currently panicing until %s", c.panicUntil.Format(sqlDateTime))
		data := SynthesizeAPIError(c.panicCode, c.panicReason, c.panicUntil.Sub(time.Now()))
		c.RUnlock()

		resp.Data = data
		resp.FromCache = true
		resp.Expires = c.panicUntil
		resp.HTTPCode = 418
		resp.Error = APIError{c.panicCode, c.panicReason}

		return resp, nil
	}
	c.RUnlock()

	// Build parameter list
	formValues := make(url.Values)
	for k, v := range r.params {
		formValues.Set(k, v)
	}

	// Use defer to cache so we can synthesize error pages if necessary
	defer func() {
		if reterr != nil {
			resp.HTTPCode = 504
			resp.Data = SynthesizeAPIError(500, "APIProxy Error: "+reterr.Error(), 5*time.Minute)
		} else if resp.Data == nil {
			resp.HTTPCode = 504
			resp.Data = SynthesizeAPIError(900, "This shouldn't happen.", 15*time.Minute)
		}
		if !r.NoCache {
			err := c.cacher.Store(cacheTag, resp.HTTPCode, resp.Data, resp.Expires)
			if err != nil {
				log.Printf("Cache Error: %s", err)
			}
		}
	}()

	//Post the shit, retry if necessary.

	tries := 0
	var httpResp *http.Response
	for tries < c.Retries {
		tries++

		httpResp, err = c.httpClient.PostForm(c.BaseURL+r.url, formValues)
		if err != nil {
			DebugLog.Printf("Error Connecting to API, retrying: %s", err)
			time.Sleep(3 * time.Second)
			continue
		}
		defer httpResp.Body.Close()

		resp.HTTPCode = httpResp.StatusCode

		// We're going to do this asynchronously so we can time it out, AAAAAAA
		type ioRead struct {
			body []byte
			err  error
		}

		readBodyChan := make(chan ioRead)
		go func() {
			bytes, err := ioutil.ReadAll(httpResp.Body)
			readBodyChan <- ioRead{bytes, err}
			close(readBodyChan)
		}()

		select {
		case readBody := <-readBodyChan:
			err = readBody.err
			data = readBody.body
		case <-time.After(c.timeout):
			data = nil
			err = fmt.Errorf("read timed out after %f seconds", c.timeout.Seconds())

			// if ioutil ever does come back, let's handle it.
			go func() {
				id := MakeID()
				DebugLog.Printf("zombie body read %s: %s ? %s", id, r.url, formValues)
				rb := <-readBodyChan
				DebugLog.Printf("zombie read completed %s: %s - %s ? %s\n%s", id, rb.err, r.url, formValues, rb.body)
			}()
		}
		if err != nil {
			DebugLog.Printf("Error Reading from API(%s), retrying...", err)
			time.Sleep(3 * time.Second)
			continue
		}

		break
		log.Printf("WARNING MAJOR REGRESSION: This should NEVER appear.")
	}
	if err != nil {
		DebugLog.Printf("Failed to access api, giving up: %s - %#v", c.BaseURL+r.url, formValues)
		return resp, ErrNetwork
	}

	//data = SynthesizeAPIError(904, "Your IP address has been temporarily blocked because it is causing too many errors. See the cacheUntil timestamp for when it will be opened again. IPs that continuall    y cause a lot of errors in the API will be permanently banned, please take measures to minimize problematic API calls from your application.", time.Second*30)

	// Get cache directive, bail with an error if anything is wrong with XML or
	// time format.  If these produce an error the rest of the data should be
	// considered worthless.
	var cR cacheResp
	err = xml.Unmarshal(data, &cR)
	if err != nil {
		DebugLog.Printf("XML Error: %s", err)
		return resp, ErrXML
	}

	// Get expiration
	expires, err = time.Parse(sqlDateTime, cR.CachedUntil)
	if err != nil {
		return resp, ErrTime
	}

	// Handle extended expiration requests.
	if r.Expires.After(expires) {
		resp.Expires = r.Expires
	} else {
		resp.Expires = expires
	}

	// Pass on any API errors
	resp.Error = cR.Error

	code := cR.Error.ErrorCode
	if code >= 901 && code <= 905 {
		log.Printf("Major API Error: %d - %s for %s %+v", cR.Error.ErrorCode, cR.Error.ErrorText, r.url, r.params)
		log.Printf("Pausing all API actions until %s...", resp.Expires.Format(sqlDateTime))
		c.Lock()
		c.panicUntil = resp.Expires
		c.panicCode = code
		c.panicReason = cR.Error.ErrorText
		c.Unlock()
	}
	if resp.HTTPCode == 403 || (code >= 100 && code <= 299) {
		resp.Invalidate = true
	}

	resp.Data = data
	return resp, nil
}
<commit_msg>don't return bogus responses<commit_after>// Copyright 2013

/*
EVE Online API downloader/cacher.

Begin by initializing the default client.
	NewClient(NilCache)

Create a new request for the API page needed.
	req := NewRequest("eve/ConquerableStationList.xml.aspx")

Set any options you may need, including keyid/vcode.
	req.Set("keyid", "1234")
	req.Set("vcode", "abcd")
	req.Set("charactername", "innominate")
	req.Set("characterid", fmt.Sprintf("%d", 123))

Get your response.
	resp, err := req.Do()
	xml := resp.Data
*/
package apicache

import (
	"crypto/rand"
	"encoding/xml"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"sync"
	"time"
)

const sqlDateTime = "2006-01-02 15:04:05"

var DebugLog = log.New(ioutil.Discard, "apicache", log.Ldate|log.Ltime)

//
//  API Error Type, hopefully contains the error code and some useful information.
//  This is CCP so nothing is guaranteed.
//
type APIError struct {
	ErrorCode int    `xml:"code,attr"`
	ErrorText string `xml:",innerxml"`
}

// Error String Generator
func (e APIError) Error() string {
	return fmt.Sprintf("API Error %d: %s", e.ErrorCode, e.ErrorText)
}

// API Client structure.  Must be created using NewClient() function.
type Client struct {
	// Base URL, defaults to CCPs api but can be changed to a proxy
	BaseURL string

	// Default three retries, can be changed at will.
	Retries int

	timeout      time.Duration
	maxIdleConns int

	cacher     Cacher
	httpClient *http.Client

	panicUntil  time.Time
	panicCode   int
	panicReason string

	sync.RWMutex
}

// Default client, first client created lives here as well.
var client *Client

// Default BaseURL for new clients, change as necessary.  Must contain
// trailing slash or bad things happen.
var DefaultBaseURL = "https://api.eveonline.com/"

// Return the default client for
func GetDefaultClient() *Client {
	if client == nil {
		panic("Tried to get nonexistent client, must be initialized first.")
	}

	return client
}

// Create a new API Client.  The first time this is called will become
// the default client.  Requires a cacher.
func NewClient(cacher Cacher) *Client {
	var newClient Client

	newClient.BaseURL = DefaultBaseURL
	newClient.Retries = 5
	newClient.cacher = cacher
	newClient.maxIdleConns = 2

	// Also sets up our initial http client
	newClient.SetTimeout(60 * time.Second)

	if client == nil {
		client = &newClient
	}
	return &newClient
}

func (c *Client) SetMaxIdleConns(maxIdleConns int) {
	// Enforce some sanity.
	if maxIdleConns <= 0 || maxIdleConns >= 64 {
		maxIdleConns = 2
	}
	c.maxIdleConns = maxIdleConns
	c.newHttpClient()
}

// Set timeout for each API request.
func (c *Client) SetTimeout(timeout time.Duration) {
	if timeout.Seconds() <= 0 || timeout.Seconds() > 3600 {
		timeout = 60 * time.Second
	}

	c.timeout = timeout
	c.newHttpClient()
}

// Set max idle conns for the default client
func SetMaxIdleConns(maxIdleConns int) {
	client.SetMaxIdleConns(maxIdleConns)
}

// Set timeout for default client.
func SetTimeout(timeout time.Duration) {
	client.SetTimeout(timeout)
}

func (c *Client) newHttpClient() {
	c.httpClient = &http.Client{
		Transport: &http.Transport{
			Dial: func(netw, addr string) (net.Conn, error) {
				c, err := net.DialTimeout(netw, addr, c.timeout)
				if err != nil {
					return nil, err
				}
				//c.SetDeadline(deadline)
				return c, nil
			},
			ResponseHeaderTimeout: c.timeout,
			MaxIdleConnsPerHost:   c.maxIdleConns,
		},
	}
}

// Create a new request.
func (c *Client) NewRequest(url string) *Request {
	var r Request
	r.params = make(map[string]string)

	if url[0] == '/' {
		url = url[1:]
	}
	r.url = url
	r.client = c

	return &r
}

// Create a new request using default client.
func NewRequest(url string) *Request {
	return client.NewRequest(url)
}

var (
	ErrCannotConnect = fmt.Errorf("Error connecting to API.")
	ErrNetwork       = fmt.Errorf("Network error.")
	ErrHTTP          = fmt.Errorf("HTTP error.")
	ErrForbidden     = fmt.Errorf("HTTP Forbidden, invalid API provided.")
	ErrUnknown       = fmt.Errorf("Unknown Error.")
	ErrXML           = fmt.Errorf("Malformed XML Detected")
	ErrTime          = fmt.Errorf("Malformed cache directive.")
)

type cacheResp struct {
	Error       APIError `xml:"error"`
	CachedUntil string   `xml:"cachedUntil"`
}

// Data structure returned from the API.
type Response struct {
	// Raw XML data
	Data []byte

	// Signals if we used the cache or not
	FromCache bool

	// Data expiration time
	Expires time.Time

	// true if an error occured due to invalid API rather than server problems
	Invalidate bool

	// Contains API Error if one occured
	Error APIError

	//Pass on CCP's HTTP code because why not?
	HTTPCode int
}

func (c *Client) GetCached(r *Request) (retresp *Response, reterr error) {
	resp := &Response{}

	// Check for cached version
	cacheTag := r.cacheTag()
	httpCode, data, expires, err := c.cacher.Get(cacheTag)
	if err == nil && !r.Force && !r.NoCache {
		resp.Data = data
		resp.FromCache = true
		resp.Expires = expires
		resp.HTTPCode = httpCode

		var cR cacheResp
		err = xml.Unmarshal(data, &cR)
		if err != nil {
			DebugLog.Printf("XML Error: %s", err)
			return resp, ErrXML
		}

		// Handle extended expiration requests.
		if r.Expires.After(expires) {
			resp.Expires = r.Expires
		} else {
			resp.Expires = expires
		}

		// Pass on any API errors
		resp.Error = cR.Error

		return resp, nil
	}
	return nil, err
}

func MakeID() string {
	buf := make([]byte, 5)
	io.ReadFull(rand.Reader, buf)
	return fmt.Sprintf("%x", buf)
}

// Perform a request, usually called by the request itself.
// User friendly error is enclosed in the response, returned error should be
// for internal use only.
func (c *Client) Do(r *Request) (retresp *Response, reterr error) {
	var data []byte
	var expires time.Time

	cacheTag := r.cacheTag()

	// Check for cached version
	resp, err := c.GetCached(r)
	if err == nil {
		return resp, nil
	}

	// If we're panicking, bail out early and spit back a fake error
	c.RLock()
	if c.panicUntil.After(time.Now()) {
		DebugLog.Printf("Got Request, but we're currently panicing until %s", c.panicUntil.Format(sqlDateTime))
		data := SynthesizeAPIError(c.panicCode, c.panicReason, c.panicUntil.Sub(time.Now()))
		c.RUnlock()

		resp.Data = data
		resp.FromCache = true
		resp.Expires = c.panicUntil
		resp.HTTPCode = 418
		resp.Error = APIError{c.panicCode, c.panicReason}

		return resp, nil
	}
	c.RUnlock()

	// Build parameter list
	formValues := make(url.Values)
	for k, v := range r.params {
		formValues.Set(k, v)
	}

	// Use defer to cache so we can synthesize error pages if necessary
	defer func() {
		if reterr != nil {
			resp.HTTPCode = 504
			resp.Data = SynthesizeAPIError(500, "APIProxy Error: "+reterr.Error(), 5*time.Minute)
		} else if resp.Data == nil {
			resp.HTTPCode = 504
			resp.Data = SynthesizeAPIError(900, "This shouldn't happen.", 15*time.Minute)
		}
		if !r.NoCache {
			err := c.cacher.Store(cacheTag, resp.HTTPCode, resp.Data, resp.Expires)
			if err != nil {
				log.Printf("Cache Error: %s", err)
			}
		}
	}()

	//Post the shit, retry if necessary.

	tries := 0
	var httpResp *http.Response
	for tries < c.Retries {
		tries++

		httpResp, err = c.httpClient.PostForm(c.BaseURL+r.url, formValues)
		if err != nil {
			DebugLog.Printf("Error Connecting to API, retrying: %s", err)
			time.Sleep(3 * time.Second)
			continue
		}
		defer httpResp.Body.Close()

		resp.HTTPCode = httpResp.StatusCode

		// We're going to do this asynchronously so we can time it out, AAAAAAA
		type ioRead struct {
			body []byte
			err  error
		}

		readBodyChan := make(chan ioRead)
		go func() {
			bytes, err := ioutil.ReadAll(httpResp.Body)
			readBodyChan <- ioRead{bytes, err}
			close(readBodyChan)
		}()

		select {
		case readBody := <-readBodyChan:
			err = readBody.err
			data = readBody.body
		case <-time.After(c.timeout):
			data = nil
			err = fmt.Errorf("read timed out after %f seconds", c.timeout.Seconds())

			// if ioutil ever does come back, let's handle it.
			go func() {
				id := MakeID()
				DebugLog.Printf("zombie body read %s: %s ? %s", id, r.url, formValues)
				rb := <-readBodyChan
				DebugLog.Printf("zombie read completed %s: %s - %s ? %s\n%s", id, rb.err, r.url, formValues, rb.body)
			}()
		}
		if err != nil {
			DebugLog.Printf("Error Reading from API(%s), retrying...", err)
			time.Sleep(3 * time.Second)
			continue
		}

		break
		log.Printf("WARNING MAJOR REGRESSION: This should NEVER appear.")
	}
	if err != nil {
		DebugLog.Printf("Failed to access api, giving up: %s - %#v", c.BaseURL+r.url, formValues)
		return resp, ErrNetwork
	}

	//data = SynthesizeAPIError(904, "Your IP address has been temporarily blocked because it is causing too many errors. See the cacheUntil timestamp for when it will be opened again. IPs that continuall    y cause a lot of errors in the API will be permanently banned, please take measures to minimize problematic API calls from your application.", time.Second*30)

	// Get cache directive, bail with an error if anything is wrong with XML or
	// time format.  If these produce an error the rest of the data should be
	// considered worthless.
	var cR cacheResp
	err = xml.Unmarshal(data, &cR)
	if err != nil {
		DebugLog.Printf("XML Error: %s", err)
		return resp, ErrXML
	}

	// Get expiration
	expires, err = time.Parse(sqlDateTime, cR.CachedUntil)
	if err != nil {
		return resp, ErrTime
	}

	// Handle extended expiration requests.
	if r.Expires.After(expires) {
		resp.Expires = r.Expires
	} else {
		resp.Expires = expires
	}

	// Pass on any API errors
	resp.Error = cR.Error

	code := cR.Error.ErrorCode
	if code >= 901 && code <= 905 {
		log.Printf("Major API Error: %d - %s for %s %+v", cR.Error.ErrorCode, cR.Error.ErrorText, r.url, r.params)
		log.Printf("Pausing all API actions until %s...", resp.Expires.Format(sqlDateTime))
		c.Lock()
		c.panicUntil = resp.Expires
		c.panicCode = code
		c.panicReason = cR.Error.ErrorText
		c.Unlock()
	}
	if resp.HTTPCode == 403 || (code >= 100 && code <= 299) {
		resp.Invalidate = true
	}

	resp.Data = data
	return resp, nil
}
<|endoftext|>


<commit_before>package azure

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/hashicorp/go-retryablehttp"
	"github.com/mitchellh/mapstructure"
)

type Request struct {
	URI      *string
	location *string
	tags     *map[string]*string
	etag     *string
	Command  APICall

	client *Client
}

func readLocation(req interface{}) (string, bool) {
	var value reflect.Value
	if reflect.ValueOf(req).Kind() == reflect.Ptr {
		value = reflect.ValueOf(req).Elem()
	} else {
		value = reflect.ValueOf(req)
	}

	for i := 0; i < value.NumField(); i++ { // iterates through every struct type field
		tag := value.Type().Field(i).Tag // returns the tag string
		if tag.Get("riviera") == "location" {
			return value.Field(i).String(), true
		}
	}
	return "", false
}

func readTags(req interface{}) (map[string]*string, bool) {
	var value reflect.Value
	if reflect.ValueOf(req).Kind() == reflect.Ptr {
		value = reflect.ValueOf(req).Elem()
	} else {
		value = reflect.ValueOf(req)
	}

	for i := 0; i < value.NumField(); i++ { // iterates through every struct type field
		tag := value.Type().Field(i).Tag // returns the tag string
		if tag.Get("riviera") == "tags" {
			tags := value.Field(i)
			return tags.Interface().(map[string]*string), true
		}
	}
	return make(map[string]*string), false
}

func (request *Request) pollForAsynchronousResponse(acceptedResponse *http.Response) (*http.Response, error) {
	var resp *http.Response = acceptedResponse

	for {
		if resp.StatusCode != http.StatusAccepted {
			return resp, nil
		}

		if retryAfter := resp.Header.Get("Retry-After"); retryAfter != "" {
			retryTime, err := strconv.Atoi(strings.TrimSpace(retryAfter))
			if err != nil {
				return nil, err
			}

			request.client.logger.Printf("[INFO] Polling pausing for %d seconds as per Retry-After header", retryTime)
			time.Sleep(time.Duration(retryTime) * time.Second)
		}

		pollLocation, err := resp.Location()
		if err != nil {
			return nil, err
		}

		request.client.logger.Printf("[INFO] Polling %q for operation completion", pollLocation.String())
		req, err := retryablehttp.NewRequest("GET", pollLocation.String(), bytes.NewReader([]byte{}))
		if err != nil {
			return nil, err
		}

		err = request.client.tokenRequester.addAuthorizationToRequest(req)
		if err != nil {
			return nil, err
		}

		resp, err := request.client.httpClient.Do(req)
		if err != nil {
			return nil, err
		}
		if resp.StatusCode == http.StatusAccepted {
			continue
		}

		return resp, err
	}
}

func (request *Request) Execute() (*Response, error) {
	apiInfo := request.Command.APIInfo()

	var urlString string

	// Base URL should already be validated by now so Parse is safe without error handling
	urlObj, _ := url.Parse(request.client.BaseURL)

	// Determine whether to use the URLPathFunc or the URI explicitly set in the request
	if request.URI == nil {
		urlObj.Path = fmt.Sprintf("/subscriptions/%s/%s", request.client.subscriptionID, strings.TrimPrefix(apiInfo.URLPathFunc(), "/"))
		urlString = urlObj.String()
	} else {
		urlObj.Path = *request.URI
		urlString = urlObj.String()
	}

	// Encode the request body if necessary
	body := bytes.NewReader([]byte{})
	if apiInfo.HasBody() {
		bodyStruct := struct {
			Location   *string             `json:"location,omitempty"`
			Tags       *map[string]*string `json:"tags,omitempty"`
			Properties interface{}         `json:"properties"`
		}{
			Properties: request.Command,
		}

		if location, hasLocation := readLocation(request.Command); hasLocation {
			bodyStruct.Location = &location
		}
		if tags, hasTags := readTags(request.Command); hasTags {
			if len(tags) > 0 {
				bodyStruct.Tags = &tags
			}
		}

		jsonEncodedRequest, err := json.Marshal(bodyStruct)
		if err != nil {
			return nil, err
		}
		body = bytes.NewReader(jsonEncodedRequest)
	}

	// Create an HTTP request
	req, err := retryablehttp.NewRequest(apiInfo.Method, urlString, body)
	if err != nil {
		return nil, err
	}

	query := req.URL.Query()
	query.Set("api-version", apiInfo.APIVersion)
	req.URL.RawQuery = query.Encode()

	if apiInfo.HasBody() {
		req.Header.Add("Content-Type", "application/json")
	}

	err = request.client.tokenRequester.addAuthorizationToRequest(req)
	if err != nil {
		return nil, err
	}

	httpResponse, err := request.client.httpClient.Do(req)
	if err != nil {
		return nil, err
	}

	// This is safe to use for every request: we check for it being http.StatusAccepted
	httpResponse, err = request.pollForAsynchronousResponse(httpResponse)
	if err != nil {
		return nil, err
	}

	var responseObj interface{}
	var errorObj *Error

	if isSuccessCode(httpResponse.StatusCode) {
		responseObj = apiInfo.ResponseTypeFunc()
		// The response factory func returns nil as a signal that there is no body
		if responseObj != nil {
			responseMap, err := unmarshalFlattenPropertiesAndClose(httpResponse)
			if err != nil {
				return nil, err
			}

			err = mapstructure.WeakDecode(responseMap, responseObj)
			if err != nil {
				return nil, err
			}
		}
	} else {
		responseMap, err := unmarshalFlattenErrorAndClose(httpResponse)

		err = mapstructure.WeakDecode(responseMap, &errorObj)
		if err != nil {
			return nil, err
		}

		errorObj.StatusCode = httpResponse.StatusCode
	}

	return &Response{
		HTTP:   httpResponse,
		Parsed: responseObj,
		Error:  errorObj,
	}, nil
}
<commit_msg>core: Pull serialization out of request execution<commit_after>package azure

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/hashicorp/go-retryablehttp"
	"github.com/mitchellh/mapstructure"
)

type Request struct {
	URI      *string
	location *string
	tags     *map[string]*string
	etag     *string
	Command  APICall

	client *Client
}

func readLocation(req interface{}) (string, bool) {
	var value reflect.Value
	if reflect.ValueOf(req).Kind() == reflect.Ptr {
		value = reflect.ValueOf(req).Elem()
	} else {
		value = reflect.ValueOf(req)
	}

	for i := 0; i < value.NumField(); i++ { // iterates through every struct type field
		tag := value.Type().Field(i).Tag // returns the tag string
		if tag.Get("riviera") == "location" {
			return value.Field(i).String(), true
		}
	}
	return "", false
}

func readTags(req interface{}) (map[string]*string, bool) {
	var value reflect.Value
	if reflect.ValueOf(req).Kind() == reflect.Ptr {
		value = reflect.ValueOf(req).Elem()
	} else {
		value = reflect.ValueOf(req)
	}

	for i := 0; i < value.NumField(); i++ { // iterates through every struct type field
		tag := value.Type().Field(i).Tag // returns the tag string
		if tag.Get("riviera") == "tags" {
			tags := value.Field(i)
			return tags.Interface().(map[string]*string), true
		}
	}
	return make(map[string]*string), false
}

func (request *Request) pollForAsynchronousResponse(acceptedResponse *http.Response) (*http.Response, error) {
	var resp *http.Response = acceptedResponse

	for {
		if resp.StatusCode != http.StatusAccepted {
			return resp, nil
		}

		if retryAfter := resp.Header.Get("Retry-After"); retryAfter != "" {
			retryTime, err := strconv.Atoi(strings.TrimSpace(retryAfter))
			if err != nil {
				return nil, err
			}

			request.client.logger.Printf("[INFO] Polling pausing for %d seconds as per Retry-After header", retryTime)
			time.Sleep(time.Duration(retryTime) * time.Second)
		}

		pollLocation, err := resp.Location()
		if err != nil {
			return nil, err
		}

		request.client.logger.Printf("[INFO] Polling %q for operation completion", pollLocation.String())
		req, err := retryablehttp.NewRequest("GET", pollLocation.String(), bytes.NewReader([]byte{}))
		if err != nil {
			return nil, err
		}

		err = request.client.tokenRequester.addAuthorizationToRequest(req)
		if err != nil {
			return nil, err
		}

		resp, err := request.client.httpClient.Do(req)
		if err != nil {
			return nil, err
		}
		if resp.StatusCode == http.StatusAccepted {
			continue
		}

		return resp, err
	}
}

func defaultARMRequestStruct(request *Request) interface{} {
	bodyStruct := struct {
		Location   *string             `json:"location,omitempty"`
		Tags       *map[string]*string `json:"tags,omitempty"`
		Properties interface{}         `json:"properties"`
	}{
		Properties: request.Command,
	}

	if location, hasLocation := readLocation(request.Command); hasLocation {
		bodyStruct.Location = &location
	}
	if tags, hasTags := readTags(request.Command); hasTags {
		if len(tags) > 0 {
			bodyStruct.Tags = &tags
		}
	}

	return bodyStruct
}

func defaultARMRequestSerialize(body interface{}) (io.ReadSeeker, error) {
	jsonEncodedRequest, err := json.Marshal(body)
	if err != nil {
		return nil, err
	}
	return bytes.NewReader(jsonEncodedRequest), nil
}

func (request *Request) Execute() (*Response, error) {
	apiInfo := request.Command.APIInfo()

	var urlString string

	// Base URL should already be validated by now so Parse is safe without error handling
	urlObj, _ := url.Parse(request.client.BaseURL)

	// Determine whether to use the URLPathFunc or the URI explicitly set in the request
	if request.URI == nil {
		urlObj.Path = fmt.Sprintf("/subscriptions/%s/%s", request.client.subscriptionID, strings.TrimPrefix(apiInfo.URLPathFunc(), "/"))
		urlString = urlObj.String()
	} else {
		urlObj.Path = *request.URI
		urlString = urlObj.String()
	}

	// Encode the request body if necessary
	var body io.ReadSeeker
	if apiInfo.HasBody() {
		bodyStruct := defaultARMRequestStruct(request)
		serialized, err := defaultARMRequestSerialize(bodyStruct)
		if err != nil {
			return nil, err
		}

		body = serialized
	} else {

		body = bytes.NewReader([]byte{})
	}

	// Create an HTTP request
	req, err := retryablehttp.NewRequest(apiInfo.Method, urlString, body)
	if err != nil {
		return nil, err
	}

	query := req.URL.Query()
	query.Set("api-version", apiInfo.APIVersion)
	req.URL.RawQuery = query.Encode()

	if apiInfo.HasBody() {
		req.Header.Add("Content-Type", "application/json")
	}

	err = request.client.tokenRequester.addAuthorizationToRequest(req)
	if err != nil {
		return nil, err
	}

	httpResponse, err := request.client.httpClient.Do(req)
	if err != nil {
		return nil, err
	}

	// This is safe to use for every request: we check for it being http.StatusAccepted
	httpResponse, err = request.pollForAsynchronousResponse(httpResponse)
	if err != nil {
		return nil, err
	}

	var responseObj interface{}
	var errorObj *Error

	if isSuccessCode(httpResponse.StatusCode) {
		responseObj = apiInfo.ResponseTypeFunc()
		// The response factory func returns nil as a signal that there is no body
		if responseObj != nil {
			responseMap, err := unmarshalFlattenPropertiesAndClose(httpResponse)
			if err != nil {
				return nil, err
			}

			err = mapstructure.WeakDecode(responseMap, responseObj)
			if err != nil {
				return nil, err
			}
		}
	} else {
		responseMap, err := unmarshalFlattenErrorAndClose(httpResponse)

		err = mapstructure.WeakDecode(responseMap, &errorObj)
		if err != nil {
			return nil, err
		}

		errorObj.StatusCode = httpResponse.StatusCode
	}

	return &Response{
		HTTP:   httpResponse,
		Parsed: responseObj,
		Error:  errorObj,
	}, nil
}
<|endoftext|>


<commit_before>package app

import (
	"errors"
	"io"
	"log"
	"strconv"
	"time"

	"github.com/mesos/mesos-go"
	"github.com/mesos/mesos-go/backoff"
	xmetrics "github.com/mesos/mesos-go/extras/metrics"
	"github.com/mesos/mesos-go/extras/scheduler/controller"
	"github.com/mesos/mesos-go/scheduler"
	"github.com/mesos/mesos-go/scheduler/calls"
	"github.com/mesos/mesos-go/scheduler/events"
)

var (
	RegistrationMinBackoff = 1 * time.Second
	RegistrationMaxBackoff = 15 * time.Second
)

func Run(cfg Config) error {
	log.Printf("scheduler running with configuration: %+v", cfg)
	shutdown := make(chan struct{})
	defer close(shutdown)

	state, err := newInternalState(cfg)
	if err != nil {
		return err
	}
	err = controller.New().Run(buildControllerConfig(state, shutdown))
	if state.err != nil {
		err = state.err
	}
	return err
}

func buildControllerConfig(state *internalState, shutdown <-chan struct{}) controller.Config {
	controlContext := &controller.ContextAdapter{
		DoneFunc:        func() bool { return state.done },
		FrameworkIDFunc: func() string { return state.frameworkID },
		ErrorFunc: func(err error) {
			if err != nil && err != io.EOF {
				log.Println(err)
			} else {
				log.Println("disconnected")
			}
		},
	}

	state.cli = calls.Decorators{
		callMetrics(state.metricsAPI, time.Now, state.config.summaryMetrics),
		logCalls(map[scheduler.Call_Type]string{scheduler.Call_SUBSCRIBE: "connecting..."}),
	}.Apply(state.cli)

	return controller.Config{
		Context:            controlContext,
		Framework:          buildFrameworkInfo(state.config),
		Caller:             state.cli,
		RegistrationTokens: backoff.Notifier(RegistrationMinBackoff, RegistrationMaxBackoff, shutdown),

		Handler: events.Decorators{
			eventMetrics(state.metricsAPI, time.Now, state.config.summaryMetrics),
			events.Decorator(logAllEvents).If(state.config.verbose),
		}.Apply(buildEventHandler(state)),
	}
}

// buildEventHandler generates and returns a handler to process events received from the subscription.
func buildEventHandler(state *internalState) events.Handler {
	// TODO(jdef) would be nice to merge this ack handler with the status update handler below; need to
	// figure out appropriate error propagation among chained handlers.
	ack := events.AcknowledgeUpdates(func() calls.Caller { return state.cli })
	return events.NewMux(
		events.DefaultHandler(events.HandlerFunc(controller.DefaultHandler)),
		events.MapFuncs(map[scheduler.Event_Type]events.HandlerFunc{
			scheduler.Event_FAILURE: func(e *scheduler.Event) error {
				log.Println("received a FAILURE event")
				f := e.GetFailure()
				failure(f.ExecutorID, f.AgentID, f.Status)
				return nil
			},
			scheduler.Event_OFFERS: func(e *scheduler.Event) error {
				if state.config.verbose {
					log.Println("received an OFFERS event")
				}
				offers := e.GetOffers().GetOffers()
				state.metricsAPI.offersReceived.Int(len(offers))
				resourceOffers(state, offers)
				return nil
			},
			scheduler.Event_UPDATE: func(e *scheduler.Event) error {
				if err := ack.HandleEvent(e); err != nil {
					log.Printf("failed to ack status update for task: %+v", err)
					// TODO(jdef) we don't return the error because that would cause the subscription
					// to terminate; is that the right thing to do?
				}
				statusUpdate(state, e.GetUpdate().GetStatus())
				return nil
			},
			scheduler.Event_SUBSCRIBED: func(e *scheduler.Event) (err error) {
				log.Println("received a SUBSCRIBED event")
				if state.frameworkID == "" {
					state.frameworkID = e.GetSubscribed().GetFrameworkID().GetValue()
					if state.frameworkID == "" {
						// sanity check
						err = errors.New("mesos gave us an empty frameworkID")
					} else {
						// automatically set the frameworkID for all outgoing calls
						state.cli = calls.FrameworkCaller(state.frameworkID).Apply(state.cli)
					}
				}
				return
			},
		}),
	)
}

func failure(eid *mesos.ExecutorID, aid *mesos.AgentID, stat *int32) {
	if eid != nil {
		// executor failed..
		msg := "executor '" + eid.Value + "' terminated"
		if aid != nil {
			msg += " on agent '" + aid.Value + "'"
		}
		if stat != nil {
			msg += " with status=" + strconv.Itoa(int(*stat))
		}
		log.Println(msg)
	} else if aid != nil {
		// agent failed..
		log.Println("agent '" + aid.Value + "' terminated")
	}
}

func resourceOffers(state *internalState, offers []mesos.Offer) {
	callOption := calls.RefuseSecondsWithJitter(state.random, state.config.maxRefuseSeconds)
	tasksLaunchedThisCycle := 0
	offersDeclined := 0
	for i := range offers {
		var (
			remaining = mesos.Resources(offers[i].Resources)
			tasks     = []mesos.TaskInfo{}
		)

		if state.config.verbose {
			log.Println("received offer id '" + offers[i].ID.Value + "' with resources " + remaining.String())
		}

		var wantsExecutorResources mesos.Resources
		if len(offers[i].ExecutorIDs) == 0 {
			wantsExecutorResources = mesos.Resources(state.executor.Resources)
		}

		flattened := remaining.Flatten()

		// avoid the expense of computing these if we can...
		if state.config.summaryMetrics && state.config.resourceTypeMetrics {
			for name, restype := range flattened.Types() {
				if restype == mesos.SCALAR {
					sum := flattened.SumScalars(mesos.NamedResources(name))
					state.metricsAPI.offeredResources(sum.GetValue(), name)
				}
			}
		}

		taskWantsResources := state.wantsTaskResources.Plus(wantsExecutorResources...)
		for state.tasksLaunched < state.totalTasks && flattened.ContainsAll(taskWantsResources) {
			state.tasksLaunched++
			taskID := state.tasksLaunched

			if state.config.verbose {
				log.Println("launching task " + strconv.Itoa(taskID) + " using offer " + offers[i].ID.Value)
			}

			task := mesos.TaskInfo{
				TaskID:    mesos.TaskID{Value: strconv.Itoa(taskID)},
				AgentID:   offers[i].AgentID,
				Executor:  state.executor,
				Resources: remaining.Find(state.wantsTaskResources.Flatten(mesos.Role(state.role).Assign())),
			}
			task.Name = "Task " + task.TaskID.Value

			remaining.Subtract(task.Resources...)
			tasks = append(tasks, task)

			flattened = remaining.Flatten()
		}

		// build Accept call to launch all of the tasks we've assembled
		accept := calls.Accept(
			calls.OfferWithOperations(
				offers[i].ID,
				calls.OpLaunch(tasks...),
			),
		).With(callOption)

		// send Accept call to mesos
		err := calls.CallNoData(state.cli, accept)
		if err != nil {
			log.Printf("failed to launch tasks: %+v", err)
		} else {
			if n := len(tasks); n > 0 {
				tasksLaunchedThisCycle += n
			} else {
				offersDeclined++
			}
		}
	}
	state.metricsAPI.offersDeclined.Int(offersDeclined)
	state.metricsAPI.tasksLaunched.Int(tasksLaunchedThisCycle)
	if state.config.summaryMetrics {
		state.metricsAPI.launchesPerOfferCycle(float64(tasksLaunchedThisCycle))
	}
}

func statusUpdate(state *internalState, s mesos.TaskStatus) {
	if state.config.verbose {
		msg := "Task " + s.TaskID.Value + " is in state " + s.GetState().String()
		if m := s.GetMessage(); m != "" {
			msg += " with message '" + m + "'"
		}
		log.Println(msg)
	}

	switch st := s.GetState(); st {
	case mesos.TASK_FINISHED:
		state.tasksFinished++
		state.metricsAPI.tasksFinished()

		if state.tasksFinished == state.totalTasks {
			log.Println("mission accomplished, terminating")
			state.done = true
		} else {
			tryReviveOffers(state)
		}

	case mesos.TASK_LOST, mesos.TASK_KILLED, mesos.TASK_FAILED, mesos.TASK_ERROR:
		state.err = errors.New("Exiting because task " + s.GetTaskID().Value +
			" is in an unexpected state " + st.String() +
			" with reason " + s.GetReason().String() +
			" from source " + s.GetSource().String() +
			" with message '" + s.GetMessage() + "'")
		state.done = true
	}
}

func tryReviveOffers(state *internalState) {
	// limit the rate at which we request offer revival
	select {
	case <-state.reviveTokens:
		// not done yet, revive offers!
		err := calls.CallNoData(state.cli, calls.Revive())
		if err != nil {
			log.Printf("failed to revive offers: %+v", err)
			return
		}
	default:
		// noop
	}
}

// logAllEvents logs every observed event; this is somewhat expensive to do
func logAllEvents(h events.Handler) events.Handler {
	return events.HandlerFunc(func(e *scheduler.Event) error {
		log.Printf("%+v\n", *e)
		return h.HandleEvent(e)
	})
}

// eventMetrics logs metrics for every processed API event
func eventMetrics(metricsAPI *metricsAPI, clock func() time.Time, timingMetrics bool) events.Decorator {
	timed := metricsAPI.eventReceivedLatency
	if !timingMetrics {
		timed = nil
	}
	harness := xmetrics.NewHarness(metricsAPI.eventReceivedCount, metricsAPI.eventErrorCount, timed, clock)
	return events.Metrics(harness)
}

// callMetrics logs metrics for every outgoing Mesos call
func callMetrics(metricsAPI *metricsAPI, clock func() time.Time, timingMetrics bool) calls.Decorator {
	timed := metricsAPI.callLatency
	if !timingMetrics {
		timed = nil
	}
	harness := xmetrics.NewHarness(metricsAPI.callCount, metricsAPI.callErrorCount, timed, clock)
	return calls.CallerMetrics(harness)
}

// logCalls logs a specific message string when a particular call-type is observed
func logCalls(messages map[scheduler.Call_Type]string) calls.Decorator {
	return func(caller calls.Caller) calls.Caller {
		return calls.CallerFunc(func(c *scheduler.Call) (mesos.Response, error) {
			if message, ok := messages[c.GetType()]; ok {
				log.Println(message)
			}
			return caller.Call(c)
		})
	}
}
<commit_msg>added TODO for handling repeated timeouts SUBSCRIBE calls<commit_after>package app

import (
	"errors"
	"io"
	"log"
	"strconv"
	"time"

	"github.com/mesos/mesos-go"
	"github.com/mesos/mesos-go/backoff"
	xmetrics "github.com/mesos/mesos-go/extras/metrics"
	"github.com/mesos/mesos-go/extras/scheduler/controller"
	"github.com/mesos/mesos-go/scheduler"
	"github.com/mesos/mesos-go/scheduler/calls"
	"github.com/mesos/mesos-go/scheduler/events"
)

var (
	RegistrationMinBackoff = 1 * time.Second
	RegistrationMaxBackoff = 15 * time.Second
)

func Run(cfg Config) error {
	log.Printf("scheduler running with configuration: %+v", cfg)
	shutdown := make(chan struct{})
	defer close(shutdown)

	state, err := newInternalState(cfg)
	if err != nil {
		return err
	}

	// TODO(jdef) how to track/handle timeout errors that occur for SUBSCRIBE calls? we should
	// probably tolerate X number of subsequent subscribe failures before bailing. we'll need
	// to track the lastCallAttempted along with subsequentSubscribeTimeouts.
	err = controller.New().Run(buildControllerConfig(state, shutdown))
	if state.err != nil {
		err = state.err
	}
	return err
}

func buildControllerConfig(state *internalState, shutdown <-chan struct{}) controller.Config {
	controlContext := &controller.ContextAdapter{
		DoneFunc:        func() bool { return state.done },
		FrameworkIDFunc: func() string { return state.frameworkID },
		ErrorFunc: func(err error) {
			if err != nil && err != io.EOF {
				log.Println(err)
			} else {
				log.Println("disconnected")
			}
		},
	}

	state.cli = calls.Decorators{
		callMetrics(state.metricsAPI, time.Now, state.config.summaryMetrics),
		logCalls(map[scheduler.Call_Type]string{scheduler.Call_SUBSCRIBE: "connecting..."}),
	}.Apply(state.cli)

	return controller.Config{
		Context:            controlContext,
		Framework:          buildFrameworkInfo(state.config),
		Caller:             state.cli,
		RegistrationTokens: backoff.Notifier(RegistrationMinBackoff, RegistrationMaxBackoff, shutdown),

		Handler: events.Decorators{
			eventMetrics(state.metricsAPI, time.Now, state.config.summaryMetrics),
			events.Decorator(logAllEvents).If(state.config.verbose),
		}.Apply(buildEventHandler(state)),
	}
}

// buildEventHandler generates and returns a handler to process events received from the subscription.
func buildEventHandler(state *internalState) events.Handler {
	// TODO(jdef) would be nice to merge this ack handler with the status update handler below; need to
	// figure out appropriate error propagation among chained handlers.
	ack := events.AcknowledgeUpdates(func() calls.Caller { return state.cli })
	return events.NewMux(
		events.DefaultHandler(events.HandlerFunc(controller.DefaultHandler)),
		events.MapFuncs(map[scheduler.Event_Type]events.HandlerFunc{
			scheduler.Event_FAILURE: func(e *scheduler.Event) error {
				log.Println("received a FAILURE event")
				f := e.GetFailure()
				failure(f.ExecutorID, f.AgentID, f.Status)
				return nil
			},
			scheduler.Event_OFFERS: func(e *scheduler.Event) error {
				if state.config.verbose {
					log.Println("received an OFFERS event")
				}
				offers := e.GetOffers().GetOffers()
				state.metricsAPI.offersReceived.Int(len(offers))
				resourceOffers(state, offers)
				return nil
			},
			scheduler.Event_UPDATE: func(e *scheduler.Event) error {
				if err := ack.HandleEvent(e); err != nil {
					log.Printf("failed to ack status update for task: %+v", err)
					// TODO(jdef) we don't return the error because that would cause the subscription
					// to terminate; is that the right thing to do?
				}
				statusUpdate(state, e.GetUpdate().GetStatus())
				return nil
			},
			scheduler.Event_SUBSCRIBED: func(e *scheduler.Event) (err error) {
				log.Println("received a SUBSCRIBED event")
				if state.frameworkID == "" {
					state.frameworkID = e.GetSubscribed().GetFrameworkID().GetValue()
					if state.frameworkID == "" {
						// sanity check
						err = errors.New("mesos gave us an empty frameworkID")
					} else {
						// automatically set the frameworkID for all outgoing calls
						state.cli = calls.FrameworkCaller(state.frameworkID).Apply(state.cli)
					}
				}
				return
			},
		}),
	)
}

func failure(eid *mesos.ExecutorID, aid *mesos.AgentID, stat *int32) {
	if eid != nil {
		// executor failed..
		msg := "executor '" + eid.Value + "' terminated"
		if aid != nil {
			msg += " on agent '" + aid.Value + "'"
		}
		if stat != nil {
			msg += " with status=" + strconv.Itoa(int(*stat))
		}
		log.Println(msg)
	} else if aid != nil {
		// agent failed..
		log.Println("agent '" + aid.Value + "' terminated")
	}
}

func resourceOffers(state *internalState, offers []mesos.Offer) {
	callOption := calls.RefuseSecondsWithJitter(state.random, state.config.maxRefuseSeconds)
	tasksLaunchedThisCycle := 0
	offersDeclined := 0
	for i := range offers {
		var (
			remaining = mesos.Resources(offers[i].Resources)
			tasks     = []mesos.TaskInfo{}
		)

		if state.config.verbose {
			log.Println("received offer id '" + offers[i].ID.Value + "' with resources " + remaining.String())
		}

		var wantsExecutorResources mesos.Resources
		if len(offers[i].ExecutorIDs) == 0 {
			wantsExecutorResources = mesos.Resources(state.executor.Resources)
		}

		flattened := remaining.Flatten()

		// avoid the expense of computing these if we can...
		if state.config.summaryMetrics && state.config.resourceTypeMetrics {
			for name, restype := range flattened.Types() {
				if restype == mesos.SCALAR {
					sum := flattened.SumScalars(mesos.NamedResources(name))
					state.metricsAPI.offeredResources(sum.GetValue(), name)
				}
			}
		}

		taskWantsResources := state.wantsTaskResources.Plus(wantsExecutorResources...)
		for state.tasksLaunched < state.totalTasks && flattened.ContainsAll(taskWantsResources) {
			state.tasksLaunched++
			taskID := state.tasksLaunched

			if state.config.verbose {
				log.Println("launching task " + strconv.Itoa(taskID) + " using offer " + offers[i].ID.Value)
			}

			task := mesos.TaskInfo{
				TaskID:    mesos.TaskID{Value: strconv.Itoa(taskID)},
				AgentID:   offers[i].AgentID,
				Executor:  state.executor,
				Resources: remaining.Find(state.wantsTaskResources.Flatten(mesos.Role(state.role).Assign())),
			}
			task.Name = "Task " + task.TaskID.Value

			remaining.Subtract(task.Resources...)
			tasks = append(tasks, task)

			flattened = remaining.Flatten()
		}

		// build Accept call to launch all of the tasks we've assembled
		accept := calls.Accept(
			calls.OfferWithOperations(
				offers[i].ID,
				calls.OpLaunch(tasks...),
			),
		).With(callOption)

		// send Accept call to mesos
		err := calls.CallNoData(state.cli, accept)
		if err != nil {
			log.Printf("failed to launch tasks: %+v", err)
		} else {
			if n := len(tasks); n > 0 {
				tasksLaunchedThisCycle += n
			} else {
				offersDeclined++
			}
		}
	}
	state.metricsAPI.offersDeclined.Int(offersDeclined)
	state.metricsAPI.tasksLaunched.Int(tasksLaunchedThisCycle)
	if state.config.summaryMetrics {
		state.metricsAPI.launchesPerOfferCycle(float64(tasksLaunchedThisCycle))
	}
}

func statusUpdate(state *internalState, s mesos.TaskStatus) {
	if state.config.verbose {
		msg := "Task " + s.TaskID.Value + " is in state " + s.GetState().String()
		if m := s.GetMessage(); m != "" {
			msg += " with message '" + m + "'"
		}
		log.Println(msg)
	}

	switch st := s.GetState(); st {
	case mesos.TASK_FINISHED:
		state.tasksFinished++
		state.metricsAPI.tasksFinished()

		if state.tasksFinished == state.totalTasks {
			log.Println("mission accomplished, terminating")
			state.done = true
		} else {
			tryReviveOffers(state)
		}

	case mesos.TASK_LOST, mesos.TASK_KILLED, mesos.TASK_FAILED, mesos.TASK_ERROR:
		state.err = errors.New("Exiting because task " + s.GetTaskID().Value +
			" is in an unexpected state " + st.String() +
			" with reason " + s.GetReason().String() +
			" from source " + s.GetSource().String() +
			" with message '" + s.GetMessage() + "'")
		state.done = true
	}
}

func tryReviveOffers(state *internalState) {
	// limit the rate at which we request offer revival
	select {
	case <-state.reviveTokens:
		// not done yet, revive offers!
		err := calls.CallNoData(state.cli, calls.Revive())
		if err != nil {
			log.Printf("failed to revive offers: %+v", err)
			return
		}
	default:
		// noop
	}
}

// logAllEvents logs every observed event; this is somewhat expensive to do
func logAllEvents(h events.Handler) events.Handler {
	return events.HandlerFunc(func(e *scheduler.Event) error {
		log.Printf("%+v\n", *e)
		return h.HandleEvent(e)
	})
}

// eventMetrics logs metrics for every processed API event
func eventMetrics(metricsAPI *metricsAPI, clock func() time.Time, timingMetrics bool) events.Decorator {
	timed := metricsAPI.eventReceivedLatency
	if !timingMetrics {
		timed = nil
	}
	harness := xmetrics.NewHarness(metricsAPI.eventReceivedCount, metricsAPI.eventErrorCount, timed, clock)
	return events.Metrics(harness)
}

// callMetrics logs metrics for every outgoing Mesos call
func callMetrics(metricsAPI *metricsAPI, clock func() time.Time, timingMetrics bool) calls.Decorator {
	timed := metricsAPI.callLatency
	if !timingMetrics {
		timed = nil
	}
	harness := xmetrics.NewHarness(metricsAPI.callCount, metricsAPI.callErrorCount, timed, clock)
	return calls.CallerMetrics(harness)
}

// logCalls logs a specific message string when a particular call-type is observed
func logCalls(messages map[scheduler.Call_Type]string) calls.Decorator {
	return func(caller calls.Caller) calls.Caller {
		return calls.CallerFunc(func(c *scheduler.Call) (mesos.Response, error) {
			if message, ok := messages[c.GetType()]; ok {
				log.Println(message)
			}
			return caller.Call(c)
		})
	}
}
<|endoftext|>


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

package environment_test

import (
	"github.com/juju/cmd"
	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/apiserver/params"
	"github.com/juju/juju/cmd/envcmd"
	"github.com/juju/juju/cmd/juju/environment"
	"github.com/juju/juju/testing"
)

type SetSuite struct {
	fakeEnvSuite
}

var _ = gc.Suite(&SetSuite{})

func (s *SetSuite) run(c *gc.C, args ...string) (*cmd.Context, error) {
	command := environment.NewSetCommand(s.fake)
	return testing.RunCommand(c, envcmd.Wrap(command), args...)
}

func (s *SetSuite) TestInit(c *gc.C) {
	for i, test := range []struct {
		args       []string
		errorMatch string
	}{
		{
			errorMatch: "no key, value pairs specified",
		}, {
			args:       []string{"special"},
			errorMatch: `expected "key=value", got "special"`,
		}, {
			args:       []string{"special=extra", "special=other"},
			errorMatch: `key "special" specified more than once`,
		}, {
			args:       []string{"agent-version=2.0.0"},
			errorMatch: "agent-version must be set via upgrade-juju",
		},
	} {
		c.Log("test %d", i)
		setCmd := &environment.SetCommand{}
		err := testing.InitCommand(setCmd, test.args)
		c.Check(err, gc.ErrorMatches, test.errorMatch)
	}
}

func (s *SetSuite) TestPassesValues(c *gc.C) {
	_, err := s.run(c, "special=extra", "unknown=foo")
	c.Assert(err, jc.ErrorIsNil)
	expected := map[string]interface{}{
		"special": "extra",
		"unknown": "foo",
	}
	c.Assert(s.fake.values, jc.DeepEquals, expected)
}

func (s *SetSuite) TestSettingKnownValue(c *gc.C) {
	_, err := s.run(c, "special=extra", "unknown=foo")
	c.Assert(err, jc.ErrorIsNil)
	// Command succeeds, but warning logged.
	expected := `key "unknown" is not defined in the current environment configuration: possible misspelling`
	c.Check(c.GetTestLog(), jc.Contains, expected)
}

func (s *SetSuite) TestBlockedError(c *gc.C) {
	s.fake.err = &params.Error{Code: params.CodeOperationBlocked}
	_, err := s.run(c, "special=extra")
	c.Assert(err, gc.Equals, cmd.ErrSilent)
	// msg is logged
	c.Check(c.GetTestLog(), jc.Contains, "To unblock changes")
}
<commit_msg>Typo fix<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package environment_test

import (
	"github.com/juju/cmd"
	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/apiserver/params"
	"github.com/juju/juju/cmd/envcmd"
	"github.com/juju/juju/cmd/juju/environment"
	"github.com/juju/juju/testing"
)

type SetSuite struct {
	fakeEnvSuite
}

var _ = gc.Suite(&SetSuite{})

func (s *SetSuite) run(c *gc.C, args ...string) (*cmd.Context, error) {
	command := environment.NewSetCommand(s.fake)
	return testing.RunCommand(c, envcmd.Wrap(command), args...)
}

func (s *SetSuite) TestInit(c *gc.C) {
	for i, test := range []struct {
		args       []string
		errorMatch string
	}{
		{
			errorMatch: "no key, value pairs specified",
		}, {
			args:       []string{"special"},
			errorMatch: `expected "key=value", got "special"`,
		}, {
			args:       []string{"special=extra", "special=other"},
			errorMatch: `key "special" specified more than once`,
		}, {
			args:       []string{"agent-version=2.0.0"},
			errorMatch: "agent-version must be set via upgrade-juju",
		},
	} {
		c.Logf("test %d", i)
		setCmd := &environment.SetCommand{}
		err := testing.InitCommand(setCmd, test.args)
		c.Check(err, gc.ErrorMatches, test.errorMatch)
	}
}

func (s *SetSuite) TestPassesValues(c *gc.C) {
	_, err := s.run(c, "special=extra", "unknown=foo")
	c.Assert(err, jc.ErrorIsNil)
	expected := map[string]interface{}{
		"special": "extra",
		"unknown": "foo",
	}
	c.Assert(s.fake.values, jc.DeepEquals, expected)
}

func (s *SetSuite) TestSettingKnownValue(c *gc.C) {
	_, err := s.run(c, "special=extra", "unknown=foo")
	c.Assert(err, jc.ErrorIsNil)
	// Command succeeds, but warning logged.
	expected := `key "unknown" is not defined in the current environment configuration: possible misspelling`
	c.Check(c.GetTestLog(), jc.Contains, expected)
}

func (s *SetSuite) TestBlockedError(c *gc.C) {
	s.fake.err = &params.Error{Code: params.CodeOperationBlocked}
	_, err := s.run(c, "special=extra")
	c.Assert(err, gc.Equals, cmd.ErrSilent)
	// msg is logged
	c.Check(c.GetTestLog(), jc.Contains, "To unblock changes")
}
<|endoftext|>


<commit_before>package compose

import (
	"errors"
	"fmt"
	"regexp"
	"strings"
)

// shorthandFileFlagMatcher matches shorthand flag specifications containing
// file specifications at their end or as the next argument. This requires
// supporting other shorthand flags which don't take an argument and which may
// precede the file flag specification. Such flags are specified in the bracket
// expression of the regular expression and may need to be updated if Docker
// Compose's command line interface evolves.
var shorthandFileFlagMatcher = regexp.MustCompile(`^-[hv]*f`)

// arguments stores the components of the (partially) parsed Docker Compose
// command line arguments and flags.
type arguments struct {
	// files are any configuration files explicitly specified using the -f or
	// --file flags, in order of specification.
	files []string
	// preCommandArguments are those arguments and flags preceeding the command
	// specification (if any). It will have file specifications filtered out. If
	// a shorthand file specification is combined with other nullary shorthand
	// flags, then the file specification will be removed from that clustering
	// of flags but the remainder of the cluster will be included. No other
	// flags are filtered out, including those whose value may be extracted.
	preCommandArguments []string
	// command is the command specification (if any).
	command *string
	// postCommandArguments are those arguments and flags following the command
	// specification (if any). If a command was not specified, then this slice
	// will be empty.
	postCommandArguments []string
	// projectDirectory is the project directory specified by the
	// --project-directory flag, if any.
	projectDirectory *string
	// environmentFile is the environment variables file specified by the
	// --env-file flag, if any.
	environmentFile *string
}

// reconstitute recombines arguments into a single slice suitable for using with
// os/exec.Command.
func (a *arguments) reconstitute() []string {
	// Compute the result size.
	count := len(a.files) + len(a.preCommandArguments)
	if a.command != nil {
		count += 1 + len(a.postCommandArguments)
	}

	// If there are no arguments, then we can return an empty slice.
	if count == 0 {
		return nil
	}

	// Preallocate the result slice.
	result := make([]string, 0, count)

	// Populate the argument slice.
	for _, file := range a.files {
		result = append(result, fmt.Sprintf("--file=%s", file))
	}
	result = append(result, a.preCommandArguments...)
	if a.command != nil {
		result = append(result, *a.command)
		result = append(result, a.postCommandArguments...)
	}

	// Done.
	return result
}

// stringptr takes a string value and returns a pointer to that value. It is a
// utility function for extracting string pointers from loop variables. It is
// guaranteed to return a non-nil result.
func stringptr(value string) *string {
	return &value
}

// parseArguments performs a cursory parsing of Docker Compose arguments and
// flags. It creates a structure that provides access to various elements of the
// arguments. This structure also allows for reconstitution of the flags for the
// purposes of invoking Docker Compose. This function avoids any disruption to
// the behavior of Docker Compose's parsing, so we only filter out file
// specifications and we keep behavioral parity with the command line parsing
// library (docopt) used by Docker Compose.
func parseArguments(rawArguments []string) (*arguments, error) {
	// Set up state tracking.
	var files, preCommandArguments, postCommandArguments []string
	var command, projectDirectory, environmentFile *string
	var nextIsFile, nextIsProjectDirectory, nextIsEnvironmentFile bool

	// Process arguments.
	for _, argument := range rawArguments {
		if nextIsFile {
			files = append(files, argument)
			nextIsFile = false
		} else if nextIsProjectDirectory {
			projectDirectory = stringptr(argument)
			nextIsProjectDirectory = false
		} else if nextIsEnvironmentFile {
			environmentFile = stringptr(argument)
			nextIsEnvironmentFile = false
		} else if command != nil {
			postCommandArguments = append(postCommandArguments, argument)
		} else if argument == "--file" {
			nextIsFile = true
		} else if strings.HasPrefix(argument, "--file=") {
			files = append(files, argument[7:])
		} else if argument == "--project-directory" {
			nextIsProjectDirectory = true
		} else if strings.HasPrefix(argument, "--project-directory=") {
			projectDirectory = stringptr(argument[20:])
		} else if argument == "--env-file" {
			nextIsEnvironmentFile = true
		} else if strings.HasPrefix(argument, "--env-file=") {
			environmentFile = stringptr(argument[11:])
		} else if shorthand := shorthandFileFlagMatcher.FindString(argument); shorthand != "" {
			if len(shorthand) == len(argument) {
				nextIsFile = true
			} else {
				files = append(files, argument[len(shorthand):])
			}
			if shorthand != "-f" {
				preCommandArguments = append(preCommandArguments, argument[:len(shorthand)-1])
			}
		} else if strings.HasPrefix(argument, "-") && argument != "-" && argument != "--" {
			preCommandArguments = append(preCommandArguments, argument)
		} else {
			command = stringptr(argument)
		}
	}

	// Check for unexpected termination.
	if nextIsFile {
		return nil, errors.New("missing file specification")
	} else if nextIsProjectDirectory {
		return nil, errors.New("missing project directory specification")
	} else if nextIsEnvironmentFile {
		return nil, errors.New("missing environment file specification")
	}

	// Success.
	return &arguments{
		files:                files,
		preCommandArguments:  preCommandArguments,
		command:              command,
		postCommandArguments: postCommandArguments,
		projectDirectory:     projectDirectory,
		environmentFile:      environmentFile,
	}, nil
}
<commit_msg>Added parsing of -p/--project-name for compose command.<commit_after>package compose

import (
	"errors"
	"fmt"
	"regexp"
	"strings"
)

var (
	// shorthandFileFlagMatcher matches shorthand flag specifications containing
	// file specifications at their end or as the next argument. This requires
	// supporting other shorthand flags which don't take an argument and which
	// may precede the file flag specification. Such flags are specified in the
	// bracket expression of the regular expression and may need to be updated
	// if Docker Compose's command line interface evolves.
	shorthandFileFlagMatcher = regexp.MustCompile(`^-[hv]*f`)

	// shorthandProjectNameFlagMatcher matches shorthand flag specifications
	// containing project name specifications at their end or as the next
	// argument. The design is the same as that of shorthandFileFlagMatcher.
	shorthandProjectNameFlagMatcher = regexp.MustCompile(`^-[hv]*p`)
)

// arguments stores the components of the (partially) parsed Docker Compose
// command line arguments and flags.
type arguments struct {
	// files are any configuration files explicitly specified using the -f or
	// --file flags, in order of specification.
	files []string
	// preCommandArguments are those arguments and flags preceeding the command
	// specification (if any). It will have file specifications filtered out. If
	// a shorthand file specification is combined with other nullary shorthand
	// flags, then the file specification will be removed from that clustering
	// of flags but the remainder of the cluster will be included. No other
	// flags are filtered out, including those whose value may be extracted.
	preCommandArguments []string
	// command is the command specification (if any).
	command *string
	// postCommandArguments are those arguments and flags following the command
	// specification (if any). If a command was not specified, then this slice
	// will be empty.
	postCommandArguments []string
	// projectName is the project name specified by the -p or --project-name
	// flags, if any.
	projectName *string
	// projectDirectory is the project directory specified by the
	// --project-directory flag, if any.
	projectDirectory *string
	// environmentFile is the environment variables file specified by the
	// --env-file flag, if any.
	environmentFile *string
}

// reconstitute recombines arguments into a single slice suitable for using with
// os/exec.Command.
func (a *arguments) reconstitute() []string {
	// Compute the result size.
	count := len(a.files) + len(a.preCommandArguments)
	if a.command != nil {
		count += 1 + len(a.postCommandArguments)
	}

	// If there are no arguments, then we can return an empty slice.
	if count == 0 {
		return nil
	}

	// Preallocate the result slice.
	result := make([]string, 0, count)

	// Populate the argument slice.
	for _, file := range a.files {
		result = append(result, fmt.Sprintf("--file=%s", file))
	}
	result = append(result, a.preCommandArguments...)
	if a.command != nil {
		result = append(result, *a.command)
		result = append(result, a.postCommandArguments...)
	}

	// Done.
	return result
}

// stringptr takes a string value and returns a pointer to that value. It is a
// utility function for extracting string pointers from loop variables. It is
// guaranteed to return a non-nil result.
func stringptr(value string) *string {
	return &value
}

// parseArguments performs a cursory parsing of Docker Compose arguments and
// flags. It creates a structure that provides access to various elements of the
// arguments. This structure also allows for reconstitution of the flags for the
// purposes of invoking Docker Compose. This function avoids any disruption to
// the behavior of Docker Compose's parsing, so we only filter out file
// specifications and we keep behavioral parity with the command line parsing
// library (docopt) used by Docker Compose.
func parseArguments(rawArguments []string) (*arguments, error) {
	// Set up state tracking.
	var files, preCommandArguments, postCommandArguments []string
	var command, projectName, projectDirectory, environmentFile *string
	var nextIsFile, nextIsProjectName, nextIsProjectDirectory, nextIsEnvironmentFile bool

	// Process arguments.
	for _, argument := range rawArguments {
		if nextIsFile {
			files = append(files, argument)
			nextIsFile = false
		} else if nextIsProjectName {
			projectName = stringptr(argument)
			nextIsProjectName = false
		} else if nextIsProjectDirectory {
			projectDirectory = stringptr(argument)
			nextIsProjectDirectory = false
		} else if nextIsEnvironmentFile {
			environmentFile = stringptr(argument)
			nextIsEnvironmentFile = false
		} else if command != nil {
			postCommandArguments = append(postCommandArguments, argument)
		} else if argument == "--file" {
			nextIsFile = true
		} else if strings.HasPrefix(argument, "--file=") {
			files = append(files, argument[7:])
		} else if argument == "--project-name" {
			nextIsProjectName = true
		} else if strings.HasPrefix(argument, "--project-name=") {
			files = append(files, argument[15:])
		} else if argument == "--project-directory" {
			nextIsProjectDirectory = true
		} else if strings.HasPrefix(argument, "--project-directory=") {
			projectDirectory = stringptr(argument[20:])
		} else if argument == "--env-file" {
			nextIsEnvironmentFile = true
		} else if strings.HasPrefix(argument, "--env-file=") {
			environmentFile = stringptr(argument[11:])
		} else if shorthand := shorthandFileFlagMatcher.FindString(argument); shorthand != "" {
			if len(shorthand) == len(argument) {
				nextIsFile = true
			} else {
				files = append(files, argument[len(shorthand):])
			}
			if shorthand != "-f" {
				preCommandArguments = append(preCommandArguments, argument[:len(shorthand)-1])
			}
		} else if shorthand = shorthandProjectNameFlagMatcher.FindString(argument); shorthand != "" {
			if len(shorthand) == len(argument) {
				nextIsProjectName = true
			} else {
				files = append(files, argument[len(shorthand):])
			}
			if shorthand != "-p" {
				preCommandArguments = append(preCommandArguments, argument[:len(shorthand)-1])
			}
		} else if strings.HasPrefix(argument, "-") && argument != "-" && argument != "--" {
			preCommandArguments = append(preCommandArguments, argument)
		} else {
			command = stringptr(argument)
		}
	}

	// Check for unexpected termination.
	if nextIsFile {
		return nil, errors.New("missing file specification")
	} else if nextIsProjectName {
		return nil, errors.New("missing project name specification")
	} else if nextIsProjectDirectory {
		return nil, errors.New("missing project directory specification")
	} else if nextIsEnvironmentFile {
		return nil, errors.New("missing environment file specification")
	}

	// Success.
	return &arguments{
		files:                files,
		preCommandArguments:  preCommandArguments,
		command:              command,
		postCommandArguments: postCommandArguments,
		projectName:          projectName,
		projectDirectory:     projectDirectory,
		environmentFile:      environmentFile,
	}, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Rook Authors. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

	http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package client

import (
	"fmt"
	"path"
	"time"

	"github.com/rook/rook/pkg/clusterd"
)

// When running the e2e tests, all ceph commands need to be run in the toolbox.
// Everywhere else, the ceph tools are assumed to be in the container where we can shell out.
var RunAllCephCommandsInToolbox = false

const (
	AdminUsername         = "client.admin"
	CephTool              = "ceph"
	RBDTool               = "rbd"
	Kubectl               = "kubectl"
	CrushTool             = "crushtool"
	cmdExecuteTimeout     = 1 * time.Minute
	cephConnectionTimeout = "15" // in seconds
)

func FinalizeCephCommandArgs(command string, args []string, configDir, clusterName string) (string, []string) {
	// the rbd client tool does not support the '--connect-timeout' option
	// so we only use it for the 'ceph' command
	// Also, there is no point of adding that option to 'crushtool' since that CLI does not connect to anything
	// 'crushtool' is a utility that lets you create, compile, decompile and test CRUSH map files.

	// we could use a slice and iterate over it but since we have only 3 elements
	// I don't think this is worth a loop
	if command != "rbd" && command != "crushtool" && command != "radosgw-admin" {
		args = append(args, "--connect-timeout="+cephConnectionTimeout)
	}

	// If the command should be run inside the toolbox pod, include the kubectl args to call the toolbox
	if RunAllCephCommandsInToolbox {
		toolArgs := []string{"-it", "exec", "rook-ceph-tools", "-n", clusterName, "--", command}
		return Kubectl, append(toolArgs, args...)
	}

	// No need to append the args if it's the default ceph cluster
	if clusterName == "ceph" && configDir == "/etc" {
		return command, args
	}

	// Append the args to find the config and keyring
	confFile := fmt.Sprintf("%s.config", clusterName)
	keyringFile := fmt.Sprintf("%s.keyring", AdminUsername)
	configArgs := []string{
		fmt.Sprintf("--cluster=%s", clusterName),
		fmt.Sprintf("--conf=%s", path.Join(configDir, clusterName, confFile)),
		fmt.Sprintf("--keyring=%s", path.Join(configDir, clusterName, keyringFile)),
	}
	return command, append(args, configArgs...)
}

func ExecuteCephCommandDebugLog(context *clusterd.Context, clusterName string, args []string) ([]byte, error) {
	return executeCephCommandWithOutputFile(context, clusterName, true, args)
}

func ExecuteCephCommand(context *clusterd.Context, clusterName string, args []string) ([]byte, error) {
	return executeCephCommandWithOutputFile(context, clusterName, false, args)
}

func ExecuteCephCommandPlain(context *clusterd.Context, clusterName string, args []string) ([]byte, error) {
	command, args := FinalizeCephCommandArgs(CephTool, args, context.ConfigDir, clusterName)
	args = append(args, "--format", "plain")
	return executeCommandWithOutputFile(context, false, command, args)
}

func ExecuteCephCommandPlainNoOutputFile(context *clusterd.Context, clusterName string, args []string) ([]byte, error) {
	command, args := FinalizeCephCommandArgs(CephTool, args, context.ConfigDir, clusterName)
	args = append(args, "--format", "plain")
	return executeCommand(context, command, args)
}

func executeCephCommandWithOutputFile(context *clusterd.Context, clusterName string, debug bool, args []string) ([]byte, error) {
	command, args := FinalizeCephCommandArgs(CephTool, args, context.ConfigDir, clusterName)
	args = append(args, "--format", "json")
	return executeCommandWithOutputFile(context, debug, command, args)
}

func ExecuteRBDCommand(context *clusterd.Context, clusterName string, args []string) ([]byte, error) {
	command, args := FinalizeCephCommandArgs(RBDTool, args, context.ConfigDir, clusterName)
	args = append(args, "--format", "json")
	return executeCommand(context, command, args)
}

func ExecuteRBDCommandNoFormat(context *clusterd.Context, clusterName string, args []string) ([]byte, error) {
	command, args := FinalizeCephCommandArgs(RBDTool, args, context.ConfigDir, clusterName)
	return executeCommand(context, command, args)
}

func ExecuteRBDCommandWithTimeout(context *clusterd.Context, clusterName string, args []string) (string, error) {
	output, err := context.Executor.ExecuteCommandWithTimeout(false, cmdExecuteTimeout, "", RBDTool, args...)
	return output, err
}

func executeCommand(context *clusterd.Context, command string, args []string) ([]byte, error) {
	output, err := context.Executor.ExecuteCommandWithOutput(false, "", command, args...)
	return []byte(output), err
}

func executeCommandWithOutputFile(context *clusterd.Context, debug bool, command string, args []string) ([]byte, error) {
	if command == Kubectl {
		// Kubectl commands targeting the toolbox container generate a temp file in the wrong place, so we will instead capture the output from stdout for the tests
		return executeCommand(context, command, args)
	}
	output, err := context.Executor.ExecuteCommandWithOutputFile(debug, "", command, "--out-file", args...)
	return []byte(output), err
}
<commit_msg>client: ceph - add missing comments<commit_after>/*
Copyright 2018 The Rook Authors. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

	http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package client

import (
	"fmt"
	"path"
	"time"

	"github.com/rook/rook/pkg/clusterd"
)

// RunAllCephCommandsInToolbox - when running the e2e tests, all ceph commands need to be run in the toolbox.
// Everywhere else, the ceph tools are assumed to be in the container where we can shell out.
var RunAllCephCommandsInToolbox = false

const (
	// AdminUsername is the name of the admin user
	AdminUsername = "client.admin"
	// CephTool is the name of the CLI tool for 'ceph'
	CephTool = "ceph"
	// RBDTool is the name of the CLI tool for 'rbd'
	RBDTool = "rbd"
	// Kubectl is the name of the CLI tool for 'kubectl'
	Kubectl = "kubectl"
	// CrushTool is the name of the CLI tool for 'crushtool'
	CrushTool             = "crushtool"
	cmdExecuteTimeout     = 1 * time.Minute
	cephConnectionTimeout = "15" // in seconds
)

// FinalizeCephCommandArgs builds the command line to be called
func FinalizeCephCommandArgs(command string, args []string, configDir, clusterName string) (string, []string) {
	// the rbd client tool does not support the '--connect-timeout' option
	// so we only use it for the 'ceph' command
	// Also, there is no point of adding that option to 'crushtool' since that CLI does not connect to anything
	// 'crushtool' is a utility that lets you create, compile, decompile and test CRUSH map files.

	// we could use a slice and iterate over it but since we have only 3 elements
	// I don't think this is worth a loop
	if command != "rbd" && command != "crushtool" && command != "radosgw-admin" {
		args = append(args, "--connect-timeout="+cephConnectionTimeout)
	}

	// If the command should be run inside the toolbox pod, include the kubectl args to call the toolbox
	if RunAllCephCommandsInToolbox {
		toolArgs := []string{"-it", "exec", "rook-ceph-tools", "-n", clusterName, "--", command}
		return Kubectl, append(toolArgs, args...)
	}

	// No need to append the args if it's the default ceph cluster
	if clusterName == "ceph" && configDir == "/etc" {
		return command, args
	}

	// Append the args to find the config and keyring
	confFile := fmt.Sprintf("%s.config", clusterName)
	keyringFile := fmt.Sprintf("%s.keyring", AdminUsername)
	configArgs := []string{
		fmt.Sprintf("--cluster=%s", clusterName),
		fmt.Sprintf("--conf=%s", path.Join(configDir, clusterName, confFile)),
		fmt.Sprintf("--keyring=%s", path.Join(configDir, clusterName, keyringFile)),
	}
	return command, append(args, configArgs...)
}

// ExecuteCephCommandDebugLog executes the 'ceph' command with 'debug' logs instead of 'info' logs
func ExecuteCephCommandDebugLog(context *clusterd.Context, clusterName string, args []string) ([]byte, error) {
	return executeCephCommandWithOutputFile(context, clusterName, true, args)
}

// ExecuteCephCommand executes the 'ceph' command
func ExecuteCephCommand(context *clusterd.Context, clusterName string, args []string) ([]byte, error) {
	return executeCephCommandWithOutputFile(context, clusterName, false, args)
}

// ExecuteCephCommandPlain executes the 'ceph' command and returns stdout in PLAIN format instead of JSON
func ExecuteCephCommandPlain(context *clusterd.Context, clusterName string, args []string) ([]byte, error) {
	command, args := FinalizeCephCommandArgs(CephTool, args, context.ConfigDir, clusterName)
	args = append(args, "--format", "plain")
	return executeCommandWithOutputFile(context, false, command, args)
}

// ExecuteCephCommandPlainNoOutputFile executes the 'ceph' command and returns stdout in PLAIN format instead of JSON
// with no output file, suppresses '--out-file' option
func ExecuteCephCommandPlainNoOutputFile(context *clusterd.Context, clusterName string, args []string) ([]byte, error) {
	command, args := FinalizeCephCommandArgs(CephTool, args, context.ConfigDir, clusterName)
	args = append(args, "--format", "plain")
	return executeCommand(context, command, args)
}

func executeCephCommandWithOutputFile(context *clusterd.Context, clusterName string, debug bool, args []string) ([]byte, error) {
	command, args := FinalizeCephCommandArgs(CephTool, args, context.ConfigDir, clusterName)
	args = append(args, "--format", "json")
	return executeCommandWithOutputFile(context, debug, command, args)
}

// ExecuteRBDCommand executes the 'rbd' command
func ExecuteRBDCommand(context *clusterd.Context, clusterName string, args []string) ([]byte, error) {
	command, args := FinalizeCephCommandArgs(RBDTool, args, context.ConfigDir, clusterName)
	args = append(args, "--format", "json")
	return executeCommand(context, command, args)
}

// ExecuteRBDCommandNoFormat executes the 'rbd' command and returns stdout in PLAIN format
func ExecuteRBDCommandNoFormat(context *clusterd.Context, clusterName string, args []string) ([]byte, error) {
	command, args := FinalizeCephCommandArgs(RBDTool, args, context.ConfigDir, clusterName)
	return executeCommand(context, command, args)
}

// ExecuteRBDCommandWithTimeout executes the 'rbd' command with a timeout of 1 minute
func ExecuteRBDCommandWithTimeout(context *clusterd.Context, clusterName string, args []string) (string, error) {
	output, err := context.Executor.ExecuteCommandWithTimeout(false, cmdExecuteTimeout, "", RBDTool, args...)
	return output, err
}

func executeCommand(context *clusterd.Context, command string, args []string) ([]byte, error) {
	output, err := context.Executor.ExecuteCommandWithOutput(false, "", command, args...)
	return []byte(output), err
}

func executeCommandWithOutputFile(context *clusterd.Context, debug bool, command string, args []string) ([]byte, error) {
	if command == Kubectl {
		// Kubectl commands targeting the toolbox container generate a temp file in the wrong place, so we will instead capture the output from stdout for the tests
		return executeCommand(context, command, args)
	}
	output, err := context.Executor.ExecuteCommandWithOutputFile(debug, "", command, "--out-file", args...)
	return []byte(output), err
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Kubernetes Authors All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package mustload loads minikube clusters, exiting with user-friendly messages
package mustload

import (
	"fmt"
	"net"
	"os"

	"github.com/docker/machine/libmachine"
	"github.com/docker/machine/libmachine/host"
	"github.com/docker/machine/libmachine/state"
	"k8s.io/klog/v2"
	"k8s.io/minikube/pkg/minikube/bootstrapper/bsutil/kverify"
	"k8s.io/minikube/pkg/minikube/command"
	"k8s.io/minikube/pkg/minikube/config"
	"k8s.io/minikube/pkg/minikube/constants"
	"k8s.io/minikube/pkg/minikube/driver"
	"k8s.io/minikube/pkg/minikube/exit"
	"k8s.io/minikube/pkg/minikube/machine"
	"k8s.io/minikube/pkg/minikube/out"
	"k8s.io/minikube/pkg/minikube/reason"
	"k8s.io/minikube/pkg/minikube/style"
)

// ClusterController holds all the needed information for a minikube cluster
type ClusterController struct {
	Config *config.ClusterConfig
	API    libmachine.API
	CP     ControlPlane
}

// ControlPlane holds all the needed information for the k8s control plane
type ControlPlane struct {
	// Host is the libmachine host object
	Host *host.Host
	// Node is our internal control object
	Node *config.Node
	// Runner provides command execution
	Runner command.Runner
	// Hostname is the host-accesible target for the apiserver
	Hostname string
	// Port is the host-accessible port for the apiserver
	Port int
	// IP is the host-accessible IP for the control plane
	IP net.IP
}

// Partial is a cmd-friendly way to load a cluster which may or may not be running
func Partial(name string, miniHome ...string) (libmachine.API, *config.ClusterConfig) {
	klog.Infof("Loading cluster: %s", name)
	api, err := machine.NewAPIClient(miniHome...)
	if err != nil {
		exit.Error(reason.NewAPIClient, "libmachine failed", err)
	}

	cc, err := config.Load(name, miniHome...)
	if err != nil {
		if config.IsNotExist(err) {
			out.T(style.Shrug, `There is no local cluster named "{{.cluster}}"`, out.V{"cluster": name})
			exitTip("start", name, reason.ExGuestNotFound)
		}
		exit.Error(reason.HostConfigLoad, "Error getting cluster config", err)
	}

	return api, cc
}

// Running is a cmd-friendly way to load a running cluster
func Running(name string) ClusterController {
	api, cc := Partial(name)

	cp, err := config.PrimaryControlPlane(cc)
	if err != nil {
		exit.Error(reason.GuestCpConfig, "Unable to find control plane", err)
	}

	machineName := driver.MachineName(*cc, cp)
	hs, err := machine.Status(api, machineName)
	if err != nil {
		exit.Error(reason.GuestStatus, "Unable to get machine status", err)
	}

	if hs == state.None.String() {
		out.T(style.Shrug, `The control plane node "{{.name}}" does not exist.`, out.V{"name": cp.Name})
		exitTip("start", name, reason.ExGuestNotFound)
	}

	if hs == state.Stopped.String() {
		out.T(style.Shrug, `The control plane node must be running for this command`)
		exitTip("start", name, reason.ExGuestUnavailable)
	}

	if hs != state.Running.String() {
		out.T(style.Shrug, `The control plane node is not running (state={{.state}})`, out.V{"name": cp.Name, "state": hs})
		exitTip("start", name, reason.ExSvcUnavailable)
	}

	host, err := machine.LoadHost(api, name)
	if err != nil {
		exit.Error(reason.GuestLoadHost, "Unable to load host", err)
	}

	cr, err := machine.CommandRunner(host)
	if err != nil {
		exit.Error(reason.InternalCommandRunner, "Unable to get command runner", err)
	}

	hostname, ip, port, err := driver.ControlPlaneEndpoint(cc, &cp, host.DriverName)
	if err != nil {
		exit.Error(reason.DrvCPEndpoint, "Unable to get forwarded endpoint", err)
	}

	return ClusterController{
		API:    api,
		Config: cc,
		CP: ControlPlane{
			Runner:   cr,
			Host:     host,
			Node:     &cp,
			Hostname: hostname,
			IP:       ip,
			Port:     port,
		},
	}
}

// Healthy is a cmd-friendly way to load a healthy cluster
func Healthy(name string) ClusterController {
	co := Running(name)

	as, err := kverify.APIServerStatus(co.CP.Runner, co.CP.Hostname, co.CP.Port)
	if err != nil {
		out.FailureT(`Unable to get control plane status: {{.error}}`, out.V{"error": err})
		exitTip("delete", name, reason.ExSvcError)
	}

	if as == state.Paused {
		out.T(style.Shrug, `The control plane for "{{.name}}" is paused!`, out.V{"name": name})
		exitTip("unpause", name, reason.ExSvcConfig)
	}

	if as != state.Running {
		out.T(style.Shrug, `This control plane is not running! (state={{.state}})`, out.V{"state": as.String()})
		out.WarningT(`This is unusual - you may want to investigate using "{{.command}}"`, out.V{"command": ExampleCmd(name, "logs")})
		exitTip("start", name, reason.ExSvcUnavailable)
	}
	return co
}

// ExampleCmd Return a minikube command containing the current profile name
func ExampleCmd(cname string, action string) string {
	if cname != constants.DefaultClusterName {
		return fmt.Sprintf("minikube %s -p %s", action, cname)
	}
	return fmt.Sprintf("minikube %s", action)
}

// exitTip returns an action tip and exits
func exitTip(action string, profile string, code int) {
	command := ExampleCmd(profile, action)
	out.T(style.Workaround, `To fix this, run: "{{.command}}"`, out.V{"command": command})
	os.Exit(code)
}
<commit_msg>Updating log messages<commit_after>/*
Copyright 2020 The Kubernetes Authors All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package mustload loads minikube clusters, exiting with user-friendly messages
package mustload

import (
	"fmt"
	"net"
	"os"

	"github.com/docker/machine/libmachine"
	"github.com/docker/machine/libmachine/host"
	"github.com/docker/machine/libmachine/state"
	"k8s.io/klog/v2"
	"k8s.io/minikube/pkg/minikube/bootstrapper/bsutil/kverify"
	"k8s.io/minikube/pkg/minikube/command"
	"k8s.io/minikube/pkg/minikube/config"
	"k8s.io/minikube/pkg/minikube/constants"
	"k8s.io/minikube/pkg/minikube/driver"
	"k8s.io/minikube/pkg/minikube/exit"
	"k8s.io/minikube/pkg/minikube/machine"
	"k8s.io/minikube/pkg/minikube/out"
	"k8s.io/minikube/pkg/minikube/reason"
	"k8s.io/minikube/pkg/minikube/style"
)

// ClusterController holds all the needed information for a minikube cluster
type ClusterController struct {
	Config *config.ClusterConfig
	API    libmachine.API
	CP     ControlPlane
}

// ControlPlane holds all the needed information for the k8s control plane
type ControlPlane struct {
	// Host is the libmachine host object
	Host *host.Host
	// Node is our internal control object
	Node *config.Node
	// Runner provides command execution
	Runner command.Runner
	// Hostname is the host-accesible target for the apiserver
	Hostname string
	// Port is the host-accessible port for the apiserver
	Port int
	// IP is the host-accessible IP for the control plane
	IP net.IP
}

// Partial is a cmd-friendly way to load a cluster which may or may not be running
func Partial(name string, miniHome ...string) (libmachine.API, *config.ClusterConfig) {
	klog.Infof("Loading cluster: %s", name)
	api, err := machine.NewAPIClient(miniHome...)
	if err != nil {
		exit.Error(reason.NewAPIClient, "libmachine failed", err)
	}

	cc, err := config.Load(name, miniHome...)
	if err != nil {
		if config.IsNotExist(err) {
			out.T(style.Shrug, `Profile "{{.cluster}}" not found. Run "minikube profile list" to view all profiles.`, out.V{"cluster": name})
			exitTip("start", name, reason.ExGuestNotFound)
		}
		exit.Error(reason.HostConfigLoad, "Error getting cluster config", err)
	}

	return api, cc
}

// Running is a cmd-friendly way to load a running cluster
func Running(name string) ClusterController {
	api, cc := Partial(name)

	cp, err := config.PrimaryControlPlane(cc)
	if err != nil {
		exit.Error(reason.GuestCpConfig, "Unable to find control plane", err)
	}

	machineName := driver.MachineName(*cc, cp)
	hs, err := machine.Status(api, machineName)
	if err != nil {
		exit.Error(reason.GuestStatus, "Unable to get machine status", err)
	}

	if hs == state.None.String() {
		out.T(style.Shrug, `The control plane node "{{.name}}" does not exist.`, out.V{"name": cp.Name})
		exitTip("start", name, reason.ExGuestNotFound)
	}

	if hs == state.Stopped.String() {
		out.T(style.Shrug, `The control plane node must be running for this command`)
		exitTip("start", name, reason.ExGuestUnavailable)
	}

	if hs != state.Running.String() {
		out.T(style.Shrug, `The control plane node is not running (state={{.state}})`, out.V{"name": cp.Name, "state": hs})
		exitTip("start", name, reason.ExSvcUnavailable)
	}

	host, err := machine.LoadHost(api, name)
	if err != nil {
		exit.Error(reason.GuestLoadHost, "Unable to load host", err)
	}

	cr, err := machine.CommandRunner(host)
	if err != nil {
		exit.Error(reason.InternalCommandRunner, "Unable to get command runner", err)
	}

	hostname, ip, port, err := driver.ControlPlaneEndpoint(cc, &cp, host.DriverName)
	if err != nil {
		exit.Error(reason.DrvCPEndpoint, "Unable to get forwarded endpoint", err)
	}

	return ClusterController{
		API:    api,
		Config: cc,
		CP: ControlPlane{
			Runner:   cr,
			Host:     host,
			Node:     &cp,
			Hostname: hostname,
			IP:       ip,
			Port:     port,
		},
	}
}

// Healthy is a cmd-friendly way to load a healthy cluster
func Healthy(name string) ClusterController {
	co := Running(name)

	as, err := kverify.APIServerStatus(co.CP.Runner, co.CP.Hostname, co.CP.Port)
	if err != nil {
		out.FailureT(`Unable to get control plane status: {{.error}}`, out.V{"error": err})
		exitTip("delete", name, reason.ExSvcError)
	}

	if as == state.Paused {
		out.T(style.Shrug, `The control plane for "{{.name}}" is paused!`, out.V{"name": name})
		exitTip("unpause", name, reason.ExSvcConfig)
	}

	if as != state.Running {
		out.T(style.Shrug, `This control plane is not running! (state={{.state}})`, out.V{"state": as.String()})
		out.WarningT(`This is unusual - you may want to investigate using "{{.command}}"`, out.V{"command": ExampleCmd(name, "logs")})
		exitTip("start", name, reason.ExSvcUnavailable)
	}
	return co
}

// ExampleCmd Return a minikube command containing the current profile name
func ExampleCmd(cname string, action string) string {
	if cname != constants.DefaultClusterName {
		return fmt.Sprintf("minikube %s -p %s", action, cname)
	}
	return fmt.Sprintf("minikube %s", action)
}

// exitTip returns an action tip and exits
func exitTip(action string, profile string, code int) {
	command := ExampleCmd(profile, action)
	out.T(style.Workaround, `To start a cluster, run: "{{.command}}"`, out.V{"command": command})
	os.Exit(code)
}
<|endoftext|>