<commit_before>// affinity assigns issues based on team mentions and those team captains.
// The idea is to separate the work of triaging of issues and pull requests
// out to a larger pool of people to make it less of a burden to be involved.
package affinity

import (
	"fmt"
	"strings"

	"github.com/google/go-github/github"
	"github.com/parkr/auto-reply/ctx"
)

var explanation = `We are utilizing a new workflow in our issues and pull requests. Affinity teams have been setup to allow community members to hear about pull requests that may be interesting to them. When a new issue or pull request comes in, we are asking that the author mention the appropriate affinity team. I then assign a random "team captain" or two to the issue who is in charge of triaging it until it is closed or passing it off to another captain. In order to move forward with this new workflow, we need to know: which of the following teams best fits your issue or contribution?`

func assignTeamCaptains(context *ctx.Context, handler Handler, body string, assigneeCount int) error {
	if context.Issue.IsEmpty() {
		context.IncrStat("affinity.error.no_ref")
		return context.NewError("assignTeamCaptains: issue reference was not set; bailing")
	}

	team, err := findAffinityTeam(body, handler.teams)
	if err != nil {
		context.IncrStat("affinity.error.no_team")
		return askForAffinityTeam(context, handler.teams)
	}

	context.Log("all affinity team captains: %q", team.Captains)
	victims := team.RandomCaptainLogins(assigneeCount)
	context.Log("selected affinity team captains for %s: %q", context.Issue, victims)
	_, _, err = context.GitHub.Issues.AddAssignees(
		context.Issue.Owner,
		context.Issue.Repo,
		context.Issue.Num,
		victims,
	)
	if err != nil {
		context.IncrStat("affinity.error.github_api")
		return context.NewError("assignTeamCaptains: problem assigning: %v", err)
	}

	context.IncrStat("affinity.success")
	context.Log("assignTeamCaptains: assigned %q to %s", victims, context.Issue)
	return nil
}

func findAffinityTeam(body string, allTeams []Team) (Team, error) {
	for _, team := range allTeams {
		if strings.Contains(body, team.Mention) {
			return team, nil
		}
	}
	return Team{}, fmt.Errorf("findAffinityTeam: no matching team")
}

func askForAffinityTeam(context *ctx.Context, allTeams []Team) error {
	return nil
	// This is annoying.
	_, _, err := context.GitHub.Issues.CreateComment(
		context.Issue.Owner,
		context.Issue.Repo,
		context.Issue.Num,
		&github.IssueComment{Body: github.String(buildAffinityTeamMessage(context, allTeams))},
	)
	if err != nil {
		return context.NewError("askForAffinityTeam: could not leave comment: %v", err)
	}
	return nil
}

func buildAffinityTeamMessage(context *ctx.Context, allTeams []Team) string {
	var prefix string
	if context.Issue.Author != "" {
		prefix = fmt.Sprintf("Hey, @%s!", context.Issue.Author)
	} else {
		prefix = "Hello!"
	}

	teams := []string{}
	for _, team := range allTeams {
		teams = append(teams, fmt.Sprintf(
			"- `%s` – %s",
			team.Mention, team.Description,
		))
	}

	return fmt.Sprintf(
		"%s %s\n\n%s\n\nMention one of these teams in a comment below and we'll get this sorted. Thanks!",
		prefix, explanation, strings.Join(teams, "\n"),
	)
}
<commit_msg>affinity: assignTeamCaptains should just assign; no asking<commit_after>// affinity assigns issues based on team mentions and those team captains.
// The idea is to separate the work of triaging of issues and pull requests
// out to a larger pool of people to make it less of a burden to be involved.
package affinity

import (
	"fmt"
	"strings"

	"github.com/google/go-github/github"
	"github.com/parkr/auto-reply/ctx"
)

var explanation = `We are utilizing a new workflow in our issues and pull requests. Affinity teams have been setup to allow community members to hear about pull requests that may be interesting to them. When a new issue or pull request comes in, we are asking that the author mention the appropriate affinity team. I then assign a random "team captain" or two to the issue who is in charge of triaging it until it is closed or passing it off to another captain. In order to move forward with this new workflow, we need to know: which of the following teams best fits your issue or contribution?`

func assignTeamCaptains(context *ctx.Context, handler Handler, body string, assigneeCount int) error {
	if context.Issue.IsEmpty() {
		context.IncrStat("affinity.error.no_ref")
		return context.NewError("assignTeamCaptains: issue reference was not set; bailing")
	}

	team, err := findAffinityTeam(body, handler.teams)
	if err != nil {
		context.IncrStat("affinity.error.no_team")
		//return askForAffinityTeam(context, handler.teams)
		return context.NewError("%s: no team in the message body; unable to assign", context.Issue)
	}

	context.Log("all affinity team captains: %q", team.Captains)
	victims := team.RandomCaptainLogins(assigneeCount)
	context.Log("selected affinity team captains for %s: %q", context.Issue, victims)
	_, _, err = context.GitHub.Issues.AddAssignees(
		context.Issue.Owner,
		context.Issue.Repo,
		context.Issue.Num,
		victims,
	)
	if err != nil {
		context.IncrStat("affinity.error.github_api")
		return context.NewError("assignTeamCaptains: problem assigning: %v", err)
	}

	context.IncrStat("affinity.success")
	context.Log("assignTeamCaptains: assigned %q to %s", victims, context.Issue)
	return nil
}

func findAffinityTeam(body string, allTeams []Team) (Team, error) {
	for _, team := range allTeams {
		if strings.Contains(body, team.Mention) {
			return team, nil
		}
	}
	return Team{}, fmt.Errorf("findAffinityTeam: no matching team")
}

func askForAffinityTeam(context *ctx.Context, allTeams []Team) error {
	_, _, err := context.GitHub.Issues.CreateComment(
		context.Issue.Owner,
		context.Issue.Repo,
		context.Issue.Num,
		&github.IssueComment{Body: github.String(buildAffinityTeamMessage(context, allTeams))},
	)
	if err != nil {
		return context.NewError("askForAffinityTeam: could not leave comment: %v", err)
	}
	return nil
}

func buildAffinityTeamMessage(context *ctx.Context, allTeams []Team) string {
	var prefix string
	if context.Issue.Author != "" {
		prefix = fmt.Sprintf("Hey, @%s!", context.Issue.Author)
	} else {
		prefix = "Hello!"
	}

	teams := []string{}
	for _, team := range allTeams {
		teams = append(teams, fmt.Sprintf(
			"- `%s` – %s",
			team.Mention, team.Description,
		))
	}

	return fmt.Sprintf(
		"%s %s\n\n%s\n\nMention one of these teams in a comment below and we'll get this sorted. Thanks!",
		prefix, explanation, strings.Join(teams, "\n"),
	)
}
<|endoftext|>


<commit_before>package engine

import (
	"errors"
	"io"

	"github.com/proullon/ramsql/engine/log"
	"github.com/proullon/ramsql/engine/parser"
	"github.com/proullon/ramsql/engine/protocol"
)

type executor func(*Engine, *parser.Decl, protocol.EngineConn) error

// Engine is the root struct of RamSQL server
type Engine struct {
	endpoint     protocol.EngineEndpoint
	relations    map[string]*Relation
	opsExecutors map[int]executor

	// Any value send to this channel (through Engine.stop)
	// Will stop the listening loop
	stop chan bool
}

// New initialize a new RamSQL server
func New(endpoint protocol.EngineEndpoint) (e *Engine, err error) {

	e = &Engine{
		endpoint: endpoint,
	}

	e.stop = make(chan bool)

	e.opsExecutors = map[int]executor{
		parser.CreateToken:   createExecutor,
		parser.TableToken:    createTableExecutor,
		parser.SelectToken:   selectExecutor,
		parser.InsertToken:   insertIntoTableExecutor,
		parser.DeleteToken:   deleteExecutor,
		parser.UpdateToken:   updateExecutor,
		parser.IfToken:       ifExecutor,
		parser.NotToken:      notExecutor,
		parser.ExistsToken:   existsExecutor,
		parser.TruncateToken: truncateExecutor,
	}

	e.relations = make(map[string]*Relation)

	err = e.start()
	if err != nil {
		return nil, err
	}

	return
}

func (e *Engine) start() (err error) {
	go e.listen()
	return nil
}

// Stop shutdown the RamSQL server
func (e *Engine) Stop() {
	if e.stop == nil {
		// already stopped
		return
	}
	go func() {
		e.stop <- true
		close(e.stop)
		e.stop = nil
	}()
}

func (e *Engine) relation(name string) *Relation {
	// Lock ?
	r := e.relations[name]
	// Unlock ?

	return r
}

func (e *Engine) listen() {
	newConnectionChannel := make(chan protocol.EngineConn)

	go func() {
		for {
			conn, err := e.endpoint.Accept()
			if err != nil {
				e.Stop()
				return
			}

			newConnectionChannel <- conn
		}
	}()

	for {
		select {
		case conn := <-newConnectionChannel:
			go e.handleConnection(conn)
			break

		case <-e.stop:
			e.endpoint.Close()
			return
		}
	}

}

func (e *Engine) handleConnection(conn protocol.EngineConn) {

	for {
		stmt, err := conn.ReadStatement()
		if err == io.EOF {
			// Todo: close engine if there is no conn left
			return
		}
		if err != nil {
			log.Warning("Enginge.handleConnection: cannot read : %s", err)
			return
		}

		instructions, err := parser.ParseInstruction(stmt)
		if err != nil {
			conn.WriteError(err)
			continue
		}

		err = e.executeQueries(instructions, conn)
		if err != nil {
			conn.WriteError(err)
			continue
		}
	}
}

func (e *Engine) executeQueries(instructions []parser.Instruction, conn protocol.EngineConn) error {
	for _, i := range instructions {
		err := e.executeQuery(i, conn)
		if err != nil {
			return err
		}
	}

	return nil
}

func (e *Engine) executeQuery(i parser.Instruction, conn protocol.EngineConn) error {

	if e.opsExecutors[i.Decls[0].Token] != nil {
		return e.opsExecutors[i.Decls[0].Token](e, i.Decls[0], conn)
	}

	return errors.New("Not Implemented")
}

func createExecutor(e *Engine, createDecl *parser.Decl, conn protocol.EngineConn) error {

	if len(createDecl.Decl) == 0 {
		return errors.New("Parsing failed, no declaration after CREATE")
	}

	if e.opsExecutors[createDecl.Decl[0].Token] != nil {
		return e.opsExecutors[createDecl.Decl[0].Token](e, createDecl.Decl[0], conn)
	}

	return errors.New("Parsing failed, unkown token " + createDecl.Decl[0].Lexeme)
}
<commit_msg>fix (engine): ensure engine recover from panic<commit_after>package engine

import (
	"errors"
	"fmt"
	"io"

	"github.com/proullon/ramsql/engine/log"
	"github.com/proullon/ramsql/engine/parser"
	"github.com/proullon/ramsql/engine/protocol"
)

type executor func(*Engine, *parser.Decl, protocol.EngineConn) error

// Engine is the root struct of RamSQL server
type Engine struct {
	endpoint     protocol.EngineEndpoint
	relations    map[string]*Relation
	opsExecutors map[int]executor

	// Any value send to this channel (through Engine.stop)
	// Will stop the listening loop
	stop chan bool
}

// New initialize a new RamSQL server
func New(endpoint protocol.EngineEndpoint) (e *Engine, err error) {

	e = &Engine{
		endpoint: endpoint,
	}

	e.stop = make(chan bool)

	e.opsExecutors = map[int]executor{
		parser.CreateToken:   createExecutor,
		parser.TableToken:    createTableExecutor,
		parser.SelectToken:   selectExecutor,
		parser.InsertToken:   insertIntoTableExecutor,
		parser.DeleteToken:   deleteExecutor,
		parser.UpdateToken:   updateExecutor,
		parser.IfToken:       ifExecutor,
		parser.NotToken:      notExecutor,
		parser.ExistsToken:   existsExecutor,
		parser.TruncateToken: truncateExecutor,
	}

	e.relations = make(map[string]*Relation)

	err = e.start()
	if err != nil {
		return nil, err
	}

	return
}

func (e *Engine) start() (err error) {
	go e.listen()
	return nil
}

// Stop shutdown the RamSQL server
func (e *Engine) Stop() {
	if e.stop == nil {
		// already stopped
		return
	}
	go func() {
		e.stop <- true
		close(e.stop)
		e.stop = nil
	}()
}

func (e *Engine) relation(name string) *Relation {
	// Lock ?
	r := e.relations[name]
	// Unlock ?

	return r
}

func (e *Engine) listen() {
	newConnectionChannel := make(chan protocol.EngineConn)

	go func() {
		for {
			conn, err := e.endpoint.Accept()
			if err != nil {
				e.Stop()
				return
			}

			newConnectionChannel <- conn
		}
	}()

	for {
		select {
		case conn := <-newConnectionChannel:
			go e.handleConnection(conn)
			break

		case <-e.stop:
			e.endpoint.Close()
			return
		}
	}

}

func (e *Engine) handleConnection(conn protocol.EngineConn) {

	for {
		stmt, err := conn.ReadStatement()
		if err == io.EOF {
			// Todo: close engine if there is no conn left
			return
		}
		if err != nil {
			log.Warning("Enginge.handleConnection: cannot read : %s", err)
			return
		}

		instructions, err := parser.ParseInstruction(stmt)
		if err != nil {
			conn.WriteError(err)
			continue
		}

		err = e.executeQueries(instructions, conn)
		if err != nil {
			conn.WriteError(err)
			continue
		}
	}
}

func (e *Engine) executeQueries(instructions []parser.Instruction, conn protocol.EngineConn) (err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("fatal error: %s", r)
			return
		}
	}()

	for _, i := range instructions {
		err = e.executeQuery(i, conn)
		if err != nil {
			return err
		}
	}

	return nil
}

func (e *Engine) executeQuery(i parser.Instruction, conn protocol.EngineConn) error {

	if e.opsExecutors[i.Decls[0].Token] != nil {
		return e.opsExecutors[i.Decls[0].Token](e, i.Decls[0], conn)
	}

	return errors.New("Not Implemented")
}

func createExecutor(e *Engine, createDecl *parser.Decl, conn protocol.EngineConn) error {

	if len(createDecl.Decl) == 0 {
		return errors.New("Parsing failed, no declaration after CREATE")
	}

	if e.opsExecutors[createDecl.Decl[0].Token] != nil {
		return e.opsExecutors[createDecl.Decl[0].Token](e, createDecl.Decl[0], conn)
	}

	return errors.New("Parsing failed, unkown token " + createDecl.Decl[0].Lexeme)
}
<|endoftext|>


<commit_before><commit_msg>tokenize: Fix gometalinter errors<commit_after><|endoftext|>


<commit_before>package auth

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"strings"

	"github.com/cozy/cozy-stack/client/request"
	"github.com/cozy/cozy-stack/pkg/config"
)

type (
	// Client describes the data of an OAuth client
	Client struct {
		ClientID          string   `json:"client_id,omitempty"`
		ClientSecret      string   `json:"client_secret"`
		SecretExpiresAt   int      `json:"client_secret_expires_at"`
		RegistrationToken string   `json:"registration_access_token"`
		RedirectURIs      []string `json:"redirect_uris"`
		ClientName        string   `json:"client_name"`
		ClientKind        string   `json:"client_kind,omitempty"`
		ClientURI         string   `json:"client_uri,omitempty"`
		LogoURI           string   `json:"logo_uri,omitempty"`
		PolicyURI         string   `json:"policy_uri,omitempty"`
		SoftwareID        string   `json:"software_id"`
		SoftwareVersion   string   `json:"software_version,omitempty"`
	}

	// AccessToken describes the content of an access token
	AccessToken struct {
		TokenType    string `json:"token_type"`
		AccessToken  string `json:"access_token"`
		RefreshToken string `json:"refresh_token"`
		Scope        string `json:"scope"`
	}

	// UserAcceptFunc is a function that can be defined by the user of this
	// library to describe how to ask the user for authorizing the client to
	// access to its data.
	//
	// The method should return the url on which the user has been redirected
	// which should contain a registering code and state, or an error .
	UserAcceptFunc func(accessURL string) (*url.URL, error)

	// Request represents an OAuth request with client parameters (*Client) and
	// list of scopes that the application wants to access.
	Request struct {
		ClientParams *Client
		Scopes       []string
		Domain       string
		Scheme       string
		HTTPClient   *http.Client
		UserAgent    string
		UserAccept   UserAcceptFunc
		Storage      Storage

		token  *AccessToken
		client *Client
	}

	// Error represents a client registration error returned by the OAuth server
	Error struct {
		Value       string `json:"error"`
		Description string `json:"error_description,omitempty"`
	}
)

func (e *Error) Error() string {
	return fmt.Sprintf("Authentication error: %s (%s)", e.Description, e.Value)
}

// AuthHeader implements the Tokener interface for the client
func (c *Client) AuthHeader() string {
	return "Bearer " + c.RegistrationToken
}

// AuthHeader implements the Tokener interface for the access token
func (t *AccessToken) AuthHeader() string {
	return "Bearer " + t.AccessToken
}

// AuthHeader implements the Tokener interface for the request
func (r *Request) AuthHeader() string {
	return r.token.AuthHeader()
}

// defaultClient defaults some values of the given client
func defaultClient(c *Client) *Client {
	if c == nil {
		c = &Client{}
	}
	if c.SoftwareID == "" {
		c.SoftwareID = "github.com/cozy/cozy-stack"
	}
	if c.SoftwareVersion == "" {
		c.SoftwareVersion = config.Version
	}
	if c.ClientName == "" {
		c.ClientName = "Cozy Go client"
	}
	if c.ClientKind == "" {
		c.ClientKind = "unknown"
	}
	return c
}

// Authenticate will start the authentication flow.
//
// If the storage has a client and token stored, it is reused and no
// authentication flow is started. Otherwise, a new client is registered and
// the authentication process is started.
func (r *Request) Authenticate() error {
	client, token, err := r.Storage.Load(r.Domain)
	if err != nil && !os.IsNotExist(err) {
		return err
	}
	if client != nil && token != nil {
		r.client, r.token = client, token
		return nil
	}
	if client == nil {
		client, err = r.RegisterClient(defaultClient(r.ClientParams))
		if err != nil {
			return err
		}
	}
	b := make([]byte, 32)
	if _, err = io.ReadFull(rand.Reader, b); err != nil {
		return err
	}
	state := base64.StdEncoding.EncodeToString(b)
	if err = r.Storage.Save(r.Domain, client, nil); err != nil {
		return err
	}
	codeURL, err := r.AuthCodeURL(client, state)
	if err != nil {
		return err
	}
	receivedURL, err := r.UserAccept(codeURL)
	if err != nil {
		return err
	}
	query := receivedURL.Query()
	if state != query.Get("state") {
		return errors.New("Non matching states")
	}
	token, err = r.GetAccessToken(client, query.Get("access_code"))
	if err != nil {
		return err
	}
	if err = r.Storage.Save(r.Domain, client, token); err != nil {
		return err
	}
	r.client, r.token = client, token
	return nil
}

// AuthCodeURL returns the url on which the user is asked to authorize the
// application.
func (r *Request) AuthCodeURL(c *Client, state string) (string, error) {
	query := url.Values{
		"client_id":     {c.ClientID},
		"redirect_uri":  {c.RedirectURIs[0]},
		"state":         {state},
		"response_type": {"code"},
		"scope":         {strings.Join(r.Scopes, " ")},
	}
	u := url.URL{
		Scheme:   "https",
		Host:     r.Domain,
		Path:     "/auth/authorize",
		RawQuery: query.Encode(),
	}
	return u.String(), nil
}

// req performs an authentication HTTP request
func (r *Request) req(opts *request.Options) (*http.Response, error) {
	opts.Domain = r.Domain
	opts.Scheme = r.Scheme
	opts.Client = r.HTTPClient
	opts.ParseError = parseError
	return request.Req(opts)
}

// RegisterClient performs the registration of the specified client.
func (r *Request) RegisterClient(c *Client) (*Client, error) {
	body, err := request.WriteJSON(c)
	if err != nil {
		return nil, err
	}
	res, err := r.req(&request.Options{
		Method: "POST",
		Path:   "/auth/register",
		Headers: request.Headers{
			"Content-Type": "application/json",
			"Accept":       "application/json",
		},
		Body: body,
	})
	if err != nil {
		return nil, err
	}
	return readClient(res.Body)
}

// GetAccessToken fetch the access token using the specified authorization
// code.
func (r *Request) GetAccessToken(c *Client, code string) (*AccessToken, error) {
	q := url.Values{
		"grant_type":    {"authorization_code"},
		"code":          {code},
		"client_id":     {c.ClientID},
		"client_secret": {c.ClientSecret},
	}
	return r.retrieveToken(c, nil, q)
}

// RefreshToken performs a token refresh using the specified client and current
// access token.
func (r *Request) RefreshToken(c *Client, t *AccessToken) (*AccessToken, error) {
	q := url.Values{
		"grant_type":    {"refresh_token"},
		"code":          {t.RefreshToken},
		"client_id":     {c.ClientID},
		"client_secret": {c.ClientSecret},
	}
	return r.retrieveToken(c, t, q)
}

func (r *Request) retrieveToken(c *Client, t *AccessToken, q url.Values) (*AccessToken, error) {
	opts := &request.Options{
		Method: "POST",
		Path:   "/auth/access_token",
		Body:   strings.NewReader(q.Encode()),
		Headers: request.Headers{
			"Content-Type": "application/x-www-form-urlencoded",
			"Accept":       "application/json",
		},
	}
	if t != nil {
		opts.Authorizer = t
	}

	res, err := r.req(opts)

	if err != nil {
		return nil, err
	}
	token := &AccessToken{}
	if err := request.ReadJSON(res.Body, token); err != nil {
		return nil, err
	}
	return token, nil
}

func parseError(res *http.Response, b []byte) error {
	var err Error
	if err := json.Unmarshal(b, &err); err != nil {
		return &request.Error{
			Status: http.StatusText(res.StatusCode),
			Title:  http.StatusText(res.StatusCode),
			Detail: string(b),
		}
	}
	// TODO: handle multi-error
	return &err
}

func readClient(r io.ReadCloser) (*Client, error) {
	client := &Client{}
	if err := request.ReadJSON(r, client); err != nil {
		return nil, err
	}
	return defaultClient(client), nil
}
<commit_msg>Fix refresh token request<commit_after>package auth

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"strings"

	"github.com/cozy/cozy-stack/client/request"
	"github.com/cozy/cozy-stack/pkg/config"
)

type (
	// Client describes the data of an OAuth client
	Client struct {
		ClientID          string   `json:"client_id,omitempty"`
		ClientSecret      string   `json:"client_secret"`
		SecretExpiresAt   int      `json:"client_secret_expires_at"`
		RegistrationToken string   `json:"registration_access_token"`
		RedirectURIs      []string `json:"redirect_uris"`
		ClientName        string   `json:"client_name"`
		ClientKind        string   `json:"client_kind,omitempty"`
		ClientURI         string   `json:"client_uri,omitempty"`
		LogoURI           string   `json:"logo_uri,omitempty"`
		PolicyURI         string   `json:"policy_uri,omitempty"`
		SoftwareID        string   `json:"software_id"`
		SoftwareVersion   string   `json:"software_version,omitempty"`
	}

	// AccessToken describes the content of an access token
	AccessToken struct {
		TokenType    string `json:"token_type"`
		AccessToken  string `json:"access_token"`
		RefreshToken string `json:"refresh_token"`
		Scope        string `json:"scope"`
	}

	// UserAcceptFunc is a function that can be defined by the user of this
	// library to describe how to ask the user for authorizing the client to
	// access to its data.
	//
	// The method should return the url on which the user has been redirected
	// which should contain a registering code and state, or an error .
	UserAcceptFunc func(accessURL string) (*url.URL, error)

	// Request represents an OAuth request with client parameters (*Client) and
	// list of scopes that the application wants to access.
	Request struct {
		ClientParams *Client
		Scopes       []string
		Domain       string
		Scheme       string
		HTTPClient   *http.Client
		UserAgent    string
		UserAccept   UserAcceptFunc
		Storage      Storage

		token  *AccessToken
		client *Client
	}

	// Error represents a client registration error returned by the OAuth server
	Error struct {
		Value       string `json:"error"`
		Description string `json:"error_description,omitempty"`
	}
)

func (e *Error) Error() string {
	return fmt.Sprintf("Authentication error: %s (%s)", e.Description, e.Value)
}

// AuthHeader implements the Tokener interface for the client
func (c *Client) AuthHeader() string {
	return "Bearer " + c.RegistrationToken
}

// AuthHeader implements the Tokener interface for the access token
func (t *AccessToken) AuthHeader() string {
	return "Bearer " + t.AccessToken
}

// AuthHeader implements the Tokener interface for the request
func (r *Request) AuthHeader() string {
	return r.token.AuthHeader()
}

// defaultClient defaults some values of the given client
func defaultClient(c *Client) *Client {
	if c == nil {
		c = &Client{}
	}
	if c.SoftwareID == "" {
		c.SoftwareID = "github.com/cozy/cozy-stack"
	}
	if c.SoftwareVersion == "" {
		c.SoftwareVersion = config.Version
	}
	if c.ClientName == "" {
		c.ClientName = "Cozy Go client"
	}
	if c.ClientKind == "" {
		c.ClientKind = "unknown"
	}
	return c
}

// Authenticate will start the authentication flow.
//
// If the storage has a client and token stored, it is reused and no
// authentication flow is started. Otherwise, a new client is registered and
// the authentication process is started.
func (r *Request) Authenticate() error {
	client, token, err := r.Storage.Load(r.Domain)
	if err != nil && !os.IsNotExist(err) {
		return err
	}
	if client != nil && token != nil {
		r.client, r.token = client, token
		return nil
	}
	if client == nil {
		client, err = r.RegisterClient(defaultClient(r.ClientParams))
		if err != nil {
			return err
		}
	}
	b := make([]byte, 32)
	if _, err = io.ReadFull(rand.Reader, b); err != nil {
		return err
	}
	state := base64.StdEncoding.EncodeToString(b)
	if err = r.Storage.Save(r.Domain, client, nil); err != nil {
		return err
	}
	codeURL, err := r.AuthCodeURL(client, state)
	if err != nil {
		return err
	}
	receivedURL, err := r.UserAccept(codeURL)
	if err != nil {
		return err
	}
	query := receivedURL.Query()
	if state != query.Get("state") {
		return errors.New("Non matching states")
	}
	token, err = r.GetAccessToken(client, query.Get("access_code"))
	if err != nil {
		return err
	}
	if err = r.Storage.Save(r.Domain, client, token); err != nil {
		return err
	}
	r.client, r.token = client, token
	return nil
}

// AuthCodeURL returns the url on which the user is asked to authorize the
// application.
func (r *Request) AuthCodeURL(c *Client, state string) (string, error) {
	query := url.Values{
		"client_id":     {c.ClientID},
		"redirect_uri":  {c.RedirectURIs[0]},
		"state":         {state},
		"response_type": {"code"},
		"scope":         {strings.Join(r.Scopes, " ")},
	}
	u := url.URL{
		Scheme:   "https",
		Host:     r.Domain,
		Path:     "/auth/authorize",
		RawQuery: query.Encode(),
	}
	return u.String(), nil
}

// req performs an authentication HTTP request
func (r *Request) req(opts *request.Options) (*http.Response, error) {
	opts.Domain = r.Domain
	opts.Scheme = r.Scheme
	opts.Client = r.HTTPClient
	opts.ParseError = parseError
	return request.Req(opts)
}

// RegisterClient performs the registration of the specified client.
func (r *Request) RegisterClient(c *Client) (*Client, error) {
	body, err := request.WriteJSON(c)
	if err != nil {
		return nil, err
	}
	res, err := r.req(&request.Options{
		Method: "POST",
		Path:   "/auth/register",
		Headers: request.Headers{
			"Content-Type": "application/json",
			"Accept":       "application/json",
		},
		Body: body,
	})
	if err != nil {
		return nil, err
	}
	return readClient(res.Body)
}

// GetAccessToken fetch the access token using the specified authorization
// code.
func (r *Request) GetAccessToken(c *Client, code string) (*AccessToken, error) {
	q := url.Values{
		"grant_type":    {"authorization_code"},
		"code":          {code},
		"client_id":     {c.ClientID},
		"client_secret": {c.ClientSecret},
	}
	return r.retrieveToken(c, nil, q)
}

// RefreshToken performs a token refresh using the specified client and current
// access token.
func (r *Request) RefreshToken(c *Client, t *AccessToken) (*AccessToken, error) {
	q := url.Values{
		"grant_type":    {"refresh_token"},
		"refresh_token": {t.RefreshToken},
		"client_id":     {c.ClientID},
		"client_secret": {c.ClientSecret},
	}
	return r.retrieveToken(c, t, q)
}

func (r *Request) retrieveToken(c *Client, t *AccessToken, q url.Values) (*AccessToken, error) {
	opts := &request.Options{
		Method: "POST",
		Path:   "/auth/access_token",
		Body:   strings.NewReader(q.Encode()),
		Headers: request.Headers{
			"Content-Type": "application/x-www-form-urlencoded",
			"Accept":       "application/json",
		},
	}
	if t != nil {
		opts.Authorizer = t
	}

	res, err := r.req(opts)

	if err != nil {
		return nil, err
	}
	token := &AccessToken{}
	if err := request.ReadJSON(res.Body, token); err != nil {
		return nil, err
	}
	return token, nil
}

func parseError(res *http.Response, b []byte) error {
	var err Error
	if err := json.Unmarshal(b, &err); err != nil {
		return &request.Error{
			Status: http.StatusText(res.StatusCode),
			Title:  http.StatusText(res.StatusCode),
			Detail: string(b),
		}
	}
	// TODO: handle multi-error
	return &err
}

func readClient(r io.ReadCloser) (*Client, error) {
	client := &Client{}
	if err := request.ReadJSON(r, client); err != nil {
		return nil, err
	}
	return defaultClient(client), nil
}
<|endoftext|>


<commit_before>package topgun_test

import (
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net/http"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"time"

	"golang.org/x/oauth2"

	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/lager/lagertest"

	gclient "code.cloudfoundry.org/garden/client"
	gconn "code.cloudfoundry.org/garden/client/connection"
	sq "github.com/Masterminds/squirrel"
	"github.com/concourse/atc"
	"github.com/concourse/go-concourse/concourse"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"

	"testing"
)

var (
	deploymentName, flyTarget string
	dbIP                      string
	atcIP, atcExternalURL     string
	atcIP2, atcExternalURL2   string

	concourseReleaseVersion, gardenRuncReleaseVersion string
	stemcellVersion                                   string

	pipelineName string

	tmpHome string
	flyBin  string

	logger *lagertest.TestLogger

	boshLogs *gexec.Session
)

var psql = sq.StatementBuilder.PlaceholderFormat(sq.Dollar)

func TestTOPGUN(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "TOPGUN Suite")
}

var _ = SynchronizedBeforeSuite(func() []byte {
	flyBinPath, err := gexec.Build("github.com/concourse/fly")
	Expect(err).ToNot(HaveOccurred())

	return []byte(flyBinPath)
}, func(data []byte) {
	flyBin = string(data)
})

var _ = SynchronizedAfterSuite(func() {
}, func() {
	gexec.CleanupBuildArtifacts()
})

var _ = BeforeEach(func() {
	SetDefaultEventuallyTimeout(5 * time.Minute)
	SetDefaultEventuallyPollingInterval(time.Second)
	SetDefaultConsistentlyDuration(time.Minute)
	SetDefaultConsistentlyPollingInterval(time.Second)

	logger = lagertest.NewTestLogger("test")

	concourseReleaseVersion = os.Getenv("CONCOURSE_RELEASE_VERSION")
	if concourseReleaseVersion == "" {
		concourseReleaseVersion = "latest"
	}

	gardenRuncReleaseVersion = os.Getenv("GARDEN_RUNC_RELEASE_VERSION")
	if gardenRuncReleaseVersion == "" {
		gardenRuncReleaseVersion = "latest"
	}

	stemcellVersion = os.Getenv("STEMCELL_VERSION")
	if stemcellVersion == "" {
		stemcellVersion = "latest"
	}

	s1 := rand.NewSource(time.Now().UnixNano())
	r1 := rand.New(s1)

	randomizer := r1.Intn(100000)

	deploymentName = fmt.Sprintf("concourse-topgun-%d-%d", GinkgoParallelNode(), randomizer)
	flyTarget = deploymentName

	bosh("delete-deployment")

	atcIP = fmt.Sprintf("10.234.%d.2", GinkgoParallelNode())
	atcIP2 = fmt.Sprintf("10.234.%d.3", GinkgoParallelNode())
	dbIP = fmt.Sprintf("10.234.%d.4", GinkgoParallelNode())

	atcExternalURL = fmt.Sprintf("http://%s:8080", atcIP)
	atcExternalURL2 = fmt.Sprintf("http://%s:8080", atcIP2)
})

var _ = AfterEach(func() {
	boshLogs.Signal(os.Interrupt)
	<-boshLogs.Exited
	boshLogs = nil

	deleteAllContainers()

	bosh("delete-deployment")
})

func Deploy(manifest string, operations ...string) {
	opFlags := []string{}
	for _, op := range operations {
		opFlags = append(opFlags, fmt.Sprintf("-o=%s", op))
	}

	bosh(
		append([]string{
			"deploy", manifest,
			"-v", "deployment-name=" + deploymentName,
			"-v", "atc-ip=" + atcIP,
			"-v", "atc-ip-2=" + atcIP2,
			"-v", "db-ip=" + dbIP,
			"-v", "atc-external-url=" + atcExternalURL,
			"-v", "atc-external-url-2=" + atcExternalURL2,
			"-v", "concourse-release-version=" + concourseReleaseVersion,
			"-v", "garden-runc-release-version=" + gardenRuncReleaseVersion,

			// 3363.10 becomes 3363.1 as it's floating point; quotes prevent that
			"-v", "stemcell-version='" + stemcellVersion + "'",
		}, opFlags...)...,
	)

	fly("login", "-c", atcExternalURL)

	boshLogs = spawnBosh("logs", "-f")
}

func bosh(argv ...string) {
	wait(spawnBosh(argv...))
}

func spawnBosh(argv ...string) *gexec.Session {
	return spawn("bosh", append([]string{"-n", "-d", deploymentName}, argv...)...)
}

func fly(argv ...string) {
	wait(spawnFly(argv...))
}

func concourseClient() concourse.Client {
	token, err := getATCToken(atcExternalURL)
	Expect(err).NotTo(HaveOccurred())
	httpClient := oauthClient(token)
	return concourse.NewClient(atcExternalURL, httpClient)
}

func deleteAllContainers() {
	client := concourseClient()
	workers, err := client.ListWorkers()
	Expect(err).NotTo(HaveOccurred())

	containers, err := client.ListContainers(map[string]string{})
	Expect(err).NotTo(HaveOccurred())

	for _, worker := range workers {
		connection := gconn.New("tcp", worker.GardenAddr)
		gardenClient := gclient.New(connection)
		for _, container := range containers {
			if container.WorkerName == worker.Name {
				err = gardenClient.Destroy(container.ID)
				if err != nil {
					logger.Error("failed-to-delete-container", err, lager.Data{"handle": container.ID})
				}
			}
		}
	}
}

func spawnFly(argv ...string) *gexec.Session {
	return spawn(flyBin, append([]string{"-t", flyTarget}, argv...)...)
}

func spawnFlyInteractive(stdin io.Reader, argv ...string) *gexec.Session {
	return spawnInteractive(stdin, flyBin, append([]string{"-t", flyTarget}, argv...)...)
}

func run(argc string, argv ...string) {
	wait(spawn(argc, argv...))
}

func spawn(argc string, argv ...string) *gexec.Session {
	By("running: " + argc + " " + strings.Join(argv, " "))
	cmd := exec.Command(argc, argv...)
	session, err := gexec.Start(cmd, GinkgoWriter, GinkgoWriter)
	Expect(err).ToNot(HaveOccurred())
	return session
}

func spawnInteractive(stdin io.Reader, argc string, argv ...string) *gexec.Session {
	By("interactively running: " + argc + " " + strings.Join(argv, " "))
	cmd := exec.Command(argc, argv...)
	cmd.Stdin = stdin
	session, err := gexec.Start(cmd, GinkgoWriter, GinkgoWriter)
	Expect(err).ToNot(HaveOccurred())
	return session
}

func wait(session *gexec.Session) {
	<-session.Exited
	Expect(session.ExitCode()).To(Equal(0))
}

func getATCToken(atcURL string) (*atc.AuthToken, error) {
	response, err := http.Get(atcURL + "/api/v1/teams/main/auth/token")
	if err != nil {
		return nil, err
	}

	var token *atc.AuthToken
	body, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}

	err = json.Unmarshal(body, &token)
	if err != nil {
		return nil, err
	}

	return token, nil
}

func oauthClient(atcToken *atc.AuthToken) *http.Client {
	return &http.Client{
		Transport: &oauth2.Transport{
			Source: oauth2.StaticTokenSource(&oauth2.Token{
				TokenType:   atcToken.Type,
				AccessToken: atcToken.Value,
			}),
			Base: &http.Transport{
				TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
			},
		},
	}
}

func waitForLandingOrLandedWorker() string {
	return waitForWorkerInState("landing", "landed")
}

func waitForRunningWorker() string {
	return waitForWorkerInState("running")
}

func waitForStalledWorker() string {
	return waitForWorkerInState("stalled")
}

func waitForWorkerInState(desiredStates ...string) string {
	var workerName string
	Eventually(func() string {

		rows := listWorkers()
		for _, row := range rows {
			if row == "" {
				continue
			}

			worker := strings.Fields(row)

			name := worker[0]
			state := worker[len(worker)-1]

			anyMatched := false
			for _, desiredState := range desiredStates {
				if state == desiredState {
					anyMatched = true
				}
			}

			if !anyMatched {
				continue
			}

			if workerName != "" {
				Fail("multiple workers in states: " + strings.Join(desiredStates, ", "))
			}

			workerName = name
		}

		return workerName
	}).ShouldNot(BeEmpty())

	return workerName
}

func flyTable(argv ...string) []map[string]string {
	session := spawnFly(append([]string{"--print-table-headers"}, argv...)...)
	<-session.Exited
	Expect(session.ExitCode()).To(Equal(0))

	result := []map[string]string{}
	var headers []string

	rows := strings.Split(string(session.Out.Contents()), "\n")
	for i, row := range rows {
		if i == 0 {
			headers = splitFlyColumns(row)
			continue
		}
		if row == "" {
			continue
		}

		result = append(result, map[string]string{})
		columns := splitFlyColumns(row)

		Expect(columns).To(HaveLen(len(headers)))

		for j, header := range headers {
			if header == "" || columns[j] == "" {
				continue
			}

			result[i-1][header] = columns[j]
		}
	}

	return result
}

func splitFlyColumns(row string) []string {
	return regexp.MustCompile(`\s{2,}`).Split(strings.TrimSpace(row), -1)
}

func listWorkers() []string {
	session := spawnFly("workers")
	<-session.Exited

	return strings.Split(string(session.Out.Contents()), "\n")
}

func waitForWorkersToBeRunning() {
	Eventually(func() bool {

		rows := listWorkers()
		anyNotRunning := false
		for _, row := range rows {
			if row == "" {
				continue
			}

			worker := strings.Fields(row)

			state := worker[len(worker)-1]

			if state != "running" {
				anyNotRunning = true
			}
		}

		return anyNotRunning
	}).Should(BeFalse())
}

func workersWithContainers() []string {
	client := concourseClient()
	containers, err := client.ListContainers(map[string]string{})
	Expect(err).NotTo(HaveOccurred())

	usedWorkers := map[string]struct{}{}

	for _, container := range containers {
		usedWorkers[container.WorkerName] = struct{}{}
	}

	var workerNames []string
	for worker, _ := range usedWorkers {
		workerNames = append(workerNames, worker)
	}

	return workerNames
}
<commit_msg>added support for an arbitary number of topgun networks<commit_after>package topgun_test

import (
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"time"

	"golang.org/x/oauth2"

	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/lager/lagertest"

	gclient "code.cloudfoundry.org/garden/client"
	gconn "code.cloudfoundry.org/garden/client/connection"
	sq "github.com/Masterminds/squirrel"
	"github.com/concourse/atc"
	"github.com/concourse/go-concourse/concourse"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"

	"strconv"
	"testing"
)

var (
	deploymentName, flyTarget string
	dbIP                      string
	atcIP, atcExternalURL     string
	atcIP2, atcExternalURL2   string

	concourseReleaseVersion, gardenRuncReleaseVersion string
	stemcellVersion                                   string

	pipelineName string

	tmpHome string
	flyBin  string

	logger *lagertest.TestLogger

	boshLogs *gexec.Session
)

var psql = sq.StatementBuilder.PlaceholderFormat(sq.Dollar)

func TestTOPGUN(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "TOPGUN Suite")
}

var _ = SynchronizedBeforeSuite(func() []byte {
	flyBinPath, err := gexec.Build("github.com/concourse/fly")
	Expect(err).ToNot(HaveOccurred())

	return []byte(flyBinPath)
}, func(data []byte) {
	flyBin = string(data)
})

var _ = SynchronizedAfterSuite(func() {
}, func() {
	gexec.CleanupBuildArtifacts()
})

var _ = BeforeEach(func() {
	SetDefaultEventuallyTimeout(5 * time.Minute)
	SetDefaultEventuallyPollingInterval(time.Second)
	SetDefaultConsistentlyDuration(time.Minute)
	SetDefaultConsistentlyPollingInterval(time.Second)

	logger = lagertest.NewTestLogger("test")

	n, found := os.LookupEnv("TOPGUN_NETWORK_OFFSET")
	var networkOffset int
	var err error

	if found {
		networkOffset, err = strconv.Atoi(n)
	}
	Expect(err).NotTo(HaveOccurred())

	concourseReleaseVersion = os.Getenv("CONCOURSE_RELEASE_VERSION")
	if concourseReleaseVersion == "" {
		concourseReleaseVersion = "latest"
	}

	gardenRuncReleaseVersion = os.Getenv("GARDEN_RUNC_RELEASE_VERSION")
	if gardenRuncReleaseVersion == "" {
		gardenRuncReleaseVersion = "latest"
	}

	stemcellVersion = os.Getenv("STEMCELL_VERSION")
	if stemcellVersion == "" {
		stemcellVersion = "latest"
	}

	deploymentNumber := GinkgoParallelNode() + (networkOffset * 4)

	deploymentName = fmt.Sprintf("concourse-topgun-%d", deploymentNumber)
	flyTarget = deploymentName

	bosh("delete-deployment")

	atcIP = fmt.Sprintf("10.234.%d.2", deploymentNumber)
	atcIP2 = fmt.Sprintf("10.234.%d.3", deploymentNumber)
	dbIP = fmt.Sprintf("10.234.%d.4", deploymentNumber)

	atcExternalURL = fmt.Sprintf("http://%s:8080", atcIP)
	atcExternalURL2 = fmt.Sprintf("http://%s:8080", atcIP2)
})

var _ = AfterEach(func() {
	boshLogs.Signal(os.Interrupt)
	<-boshLogs.Exited
	boshLogs = nil

	deleteAllContainers()

	bosh("delete-deployment")
})

func Deploy(manifest string, operations ...string) {
	opFlags := []string{}
	for _, op := range operations {
		opFlags = append(opFlags, fmt.Sprintf("-o=%s", op))
	}

	bosh(
		append([]string{
			"deploy", manifest,
			"-v", "deployment-name=" + deploymentName,
			"-v", "atc-ip=" + atcIP,
			"-v", "atc-ip-2=" + atcIP2,
			"-v", "db-ip=" + dbIP,
			"-v", "atc-external-url=" + atcExternalURL,
			"-v", "atc-external-url-2=" + atcExternalURL2,
			"-v", "concourse-release-version=" + concourseReleaseVersion,
			"-v", "garden-runc-release-version=" + gardenRuncReleaseVersion,

			// 3363.10 becomes 3363.1 as it's floating point; quotes prevent that
			"-v", "stemcell-version='" + stemcellVersion + "'",
		}, opFlags...)...,
	)

	fly("login", "-c", atcExternalURL)

	boshLogs = spawnBosh("logs", "-f")
}

func bosh(argv ...string) {
	wait(spawnBosh(argv...))
}

func spawnBosh(argv ...string) *gexec.Session {
	return spawn("bosh", append([]string{"-n", "-d", deploymentName}, argv...)...)
}

func fly(argv ...string) {
	wait(spawnFly(argv...))
}

func concourseClient() concourse.Client {
	token, err := getATCToken(atcExternalURL)
	Expect(err).NotTo(HaveOccurred())
	httpClient := oauthClient(token)
	return concourse.NewClient(atcExternalURL, httpClient)
}

func deleteAllContainers() {
	client := concourseClient()
	workers, err := client.ListWorkers()
	Expect(err).NotTo(HaveOccurred())

	containers, err := client.ListContainers(map[string]string{})
	Expect(err).NotTo(HaveOccurred())

	for _, worker := range workers {
		connection := gconn.New("tcp", worker.GardenAddr)
		gardenClient := gclient.New(connection)
		for _, container := range containers {
			if container.WorkerName == worker.Name {
				err = gardenClient.Destroy(container.ID)
				if err != nil {
					logger.Error("failed-to-delete-container", err, lager.Data{"handle": container.ID})
				}
			}
		}
	}
}

func spawnFly(argv ...string) *gexec.Session {
	return spawn(flyBin, append([]string{"-t", flyTarget}, argv...)...)
}

func spawnFlyInteractive(stdin io.Reader, argv ...string) *gexec.Session {
	return spawnInteractive(stdin, flyBin, append([]string{"-t", flyTarget}, argv...)...)
}

func run(argc string, argv ...string) {
	wait(spawn(argc, argv...))
}

func spawn(argc string, argv ...string) *gexec.Session {
	By("running: " + argc + " " + strings.Join(argv, " "))
	cmd := exec.Command(argc, argv...)
	session, err := gexec.Start(cmd, GinkgoWriter, GinkgoWriter)
	Expect(err).ToNot(HaveOccurred())
	return session
}

func spawnInteractive(stdin io.Reader, argc string, argv ...string) *gexec.Session {
	By("interactively running: " + argc + " " + strings.Join(argv, " "))
	cmd := exec.Command(argc, argv...)
	cmd.Stdin = stdin
	session, err := gexec.Start(cmd, GinkgoWriter, GinkgoWriter)
	Expect(err).ToNot(HaveOccurred())
	return session
}

func wait(session *gexec.Session) {
	<-session.Exited
	Expect(session.ExitCode()).To(Equal(0))
}

func getATCToken(atcURL string) (*atc.AuthToken, error) {
	response, err := http.Get(atcURL + "/api/v1/teams/main/auth/token")
	if err != nil {
		return nil, err
	}

	var token *atc.AuthToken
	body, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}

	err = json.Unmarshal(body, &token)
	if err != nil {
		return nil, err
	}

	return token, nil
}

func oauthClient(atcToken *atc.AuthToken) *http.Client {
	return &http.Client{
		Transport: &oauth2.Transport{
			Source: oauth2.StaticTokenSource(&oauth2.Token{
				TokenType:   atcToken.Type,
				AccessToken: atcToken.Value,
			}),
			Base: &http.Transport{
				TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
			},
		},
	}
}

func waitForLandingOrLandedWorker() string {
	return waitForWorkerInState("landing", "landed")
}

func waitForRunningWorker() string {
	return waitForWorkerInState("running")
}

func waitForStalledWorker() string {
	return waitForWorkerInState("stalled")
}

func waitForWorkerInState(desiredStates ...string) string {
	var workerName string
	Eventually(func() string {

		rows := listWorkers()
		for _, row := range rows {
			if row == "" {
				continue
			}

			worker := strings.Fields(row)

			name := worker[0]
			state := worker[len(worker)-1]

			anyMatched := false
			for _, desiredState := range desiredStates {
				if state == desiredState {
					anyMatched = true
				}
			}

			if !anyMatched {
				continue
			}

			if workerName != "" {
				Fail("multiple workers in states: " + strings.Join(desiredStates, ", "))
			}

			workerName = name
		}

		return workerName
	}).ShouldNot(BeEmpty())

	return workerName
}

func flyTable(argv ...string) []map[string]string {
	session := spawnFly(append([]string{"--print-table-headers"}, argv...)...)
	<-session.Exited
	Expect(session.ExitCode()).To(Equal(0))

	result := []map[string]string{}
	var headers []string

	rows := strings.Split(string(session.Out.Contents()), "\n")
	for i, row := range rows {
		if i == 0 {
			headers = splitFlyColumns(row)
			continue
		}
		if row == "" {
			continue
		}

		result = append(result, map[string]string{})
		columns := splitFlyColumns(row)

		Expect(columns).To(HaveLen(len(headers)))

		for j, header := range headers {
			if header == "" || columns[j] == "" {
				continue
			}

			result[i-1][header] = columns[j]
		}
	}

	return result
}

func splitFlyColumns(row string) []string {
	return regexp.MustCompile(`\s{2,}`).Split(strings.TrimSpace(row), -1)
}

func listWorkers() []string {
	session := spawnFly("workers")
	<-session.Exited

	return strings.Split(string(session.Out.Contents()), "\n")
}

func waitForWorkersToBeRunning() {
	Eventually(func() bool {

		rows := listWorkers()
		anyNotRunning := false
		for _, row := range rows {
			if row == "" {
				continue
			}

			worker := strings.Fields(row)

			state := worker[len(worker)-1]

			if state != "running" {
				anyNotRunning = true
			}
		}

		return anyNotRunning
	}).Should(BeFalse())
}

func workersWithContainers() []string {
	client := concourseClient()
	containers, err := client.ListContainers(map[string]string{})
	Expect(err).NotTo(HaveOccurred())

	usedWorkers := map[string]struct{}{}

	for _, container := range containers {
		usedWorkers[container.WorkerName] = struct{}{}
	}

	var workerNames []string
	for worker, _ := range usedWorkers {
		workerNames = append(workerNames, worker)
	}

	return workerNames
}
<|endoftext|>


<commit_before>package model

import (
	"io/ioutil"
	"log"
	"os"
	"time"

	"github.com/tinzenite/shared"
)

/*
localRemove initiates a deletion locally, creating all necessary files and
removing the file from the model.
*/
func (m *Model) localRemove(path *shared.RelativePath) error {
	// get stin for notify
	stin, exists := m.StaticInfos[path.SubPath()]
	if !exists {
		m.log("LocalRemove: stin is missing!")
		return shared.ErrIllegalFileState
	}
	// sanity check
	if m.isRemoved(stin.Identification) {
		// shouldn't happen but let's be sure; warn at least
		m.warn("LocalRemove: file removal already begun!")
	}
	// direct remove
	err := m.directRemove(path)
	if err != nil {
		m.log("LocalRemove: failed to directly remove file!")
		return err
	}
	// write peers
	err = m.writeRemovalDir(stin.Identification)
	if err != nil {
		m.log("failed to update removal dir for", stin.Identification)
		return err
	}
	// update removal dir here so that creations etc are sent before notify below!
	err = m.updateLocal(m.Root + "/" + shared.TINZENITEDIR + "/" + shared.REMOVEDIR + "/" + stin.Identification)
	if err != nil {
		m.warn("partial update on local remove failed!")
		// but continue on because the changes will be synchronized later then anyway
	}
	// send notify
	notifyObj := &shared.ObjectInfo{
		Identification: stin.Identification,
		Name:           path.LastElement(),
		Path:           path.SubPath(),
		Content:        stin.Content,
		Version:        stin.Version,
		Directory:      stin.Directory}
	m.notify(shared.OpRemove, notifyObj)
	return nil
}

/*
remoteRemove handles a remote call of remove.

TODO this is buggy, fix it.
*/
func (m *Model) remoteRemove(path *shared.RelativePath, remoteObject *shared.ObjectInfo) error {
	// sanity check
	if remoteObject == nil {
		return shared.ErrIllegalParameters
	}
	// get state information
	localFileExists := m.IsTracked(path.FullPath())
	removalExists := m.isRemoved(remoteObject.Identification)
	// if still exists locally remove it
	if localFileExists {
		// remove file (removedir should already exist, so nothing else to do)
		err := m.directRemove(path)
		if err != nil {
			m.log("couldn't remove file", path.FullPath())
			return err
		}
	}
	// warn if remove dir didn't exist yet
	if !removalExists {
		// we'll create it anyway if it doesn't exist, so all ok but warn
		m.warn("remote file removed but removedir didn't yet exist!")
	}
	// since remote removal --> write peer to done
	err := m.writeRemovalDir(remoteObject.Identification)
	if err != nil {
		m.log("updating removal dir failed!")
		return err
	}
	// send notify (reuse remoteObject)
	m.notify(shared.OpRemove, remoteObject)
	return nil
}

/*
checkRemove checks whether a remove can be finally applied and purged from the
model dependent on the peers in done and check.
*/
func (m *Model) checkRemove() error {
	removeDir := m.Root + "/" + shared.TINZENITEDIR + "/" + shared.REMOVEDIR
	allRemovals, err := ioutil.ReadDir(removeDir)
	if err != nil {
		m.log("reading all removals failed")
		return err
	}
	// check for each removal
	for _, stat := range allRemovals {
		// update removal stats (including writing own peer into DONE for all of them)
		err = m.writeRemovalDir(stat.Name())
		if err != nil {
			log.Println("DEBUG: updating removal dir failed on checkRemove!", err)
			return err
		}
		// check if we can complete the removal
		err := m.completeTrackedRemoval(stat.Name())
		if err != nil {
			// notify of error but don't stop, rest can still be checked
			m.log("completeTrackedRemoval:", err.Error())
		}
		// warn of possible orphans
		if time.Since(stat.ModTime()) > removalTimeout {
			m.warn("Removal may be orphaned! ", stat.Name())
			/*TODO this may be called even if it has just been removed... do better logic!
			Also: is there something we can do in this case?*/
		}
		// warn of possibly unapplied removals:
		subPath, err := m.GetSubPath(stat.Name())
		// if err just skip the check (can happen if the file has been removed, so ok)
		if err == nil && m.IsTracked(m.Root+"/"+subPath) {
			m.warn("Removal may be unapplied!")
		}
	}
	// also remove old local remove notifies:
	localDir := m.Root + "/" + shared.TINZENITEDIR + "/" + shared.LOCALDIR + "/" + shared.REMOVESTOREDIR
	allLocals, err := ioutil.ReadDir(localDir)
	if err != nil {
		m.log("reading of local remove notifies failed!")
		return err
	}
	for _, stat := range allLocals {
		if time.Since(stat.ModTime()) > removalLocal {
			log.Println("DEBUG: local remove of notify can be done!")
			// TODO actually remove them here...
		}
	}
	return nil
}

/*
completeTrackedRemoval checks and if ok, removes the tracked removal dir, replacing
it with a local notify of the removal. This allows the tracked removal to be
purged. After a time out the local notify is also removed.
*/
func (m *Model) completeTrackedRemoval(identification string) error {
	removeDir := m.Root + "/" + shared.TINZENITEDIR + "/" + shared.REMOVEDIR
	// working directory
	objRemovePath := removeDir + "/" + identification
	// read all peers to check for
	allCheck, err := ioutil.ReadDir(objRemovePath + "/" + shared.REMOVECHECKDIR)
	if err != nil {
		m.log("Failed reading check peer list!")
		return err
	}
	// test whether we can remove it
	complete := true
	for _, peerStat := range allCheck {
		if !shared.FileExists(objRemovePath + "/" + shared.REMOVEDONEDIR + "/" + peerStat.Name()) {
			complete = false
			break
		}
	}
	if complete {
		// make local note of removal instead of tracked one so that we can remove it
		err := m.makeLocalRemove(identification)
		if err != nil {
			m.log("failed to write local remove note, will not complete removal!")
			return err
		}
		// HARD delete the entire dir: all peers should do the same (soft delete would make removal recursive)
		err = m.directRemove(shared.CreatePathRoot(m.Root).Apply(objRemovePath))
		if err != nil {
			m.log("Failed to direct remove!")
			return err
		}
		// note that other peers may not HARD delete it yet, but the isLocalRemoved check ensures that the dir isn't reintroduced
	}
	return nil
}

/*
writeRemovalDir is an internal function that writes all known peers to check
and the own peer to done, if not already existing.
*/
func (m *Model) writeRemovalDir(objIdentification string) error {
	removeDirectory := m.Root + "/" + shared.TINZENITEDIR + "/" + shared.REMOVEDIR + "/" + objIdentification
	// make directories if don't exist
	if !shared.FileExists(removeDirectory) {
		err := shared.MakeDirectories(removeDirectory, shared.REMOVECHECKDIR, shared.REMOVEDONEDIR)
		if err != nil {
			m.log("making removedir error")
			return err
		}
	}
	// write peer list to check which must all be notified of removal
	peers, err := m.readPeers()
	if err != nil {
		m.log("Failed to read peers")
		return err
	}
	for _, peer := range peers {
		path := removeDirectory + "/" + shared.REMOVECHECKDIR + "/" + peer
		// if already written don't rewrite
		if shared.FileExists(path) {
			continue
		}
		err = ioutil.WriteFile(path, []byte(""), shared.FILEPERMISSIONMODE)
		if err != nil {
			m.log("Couldn't write peer file", peer, "to", shared.REMOVECHECKDIR, "!")
			return err
		}
	}
	// write own peer into DONE
	path := removeDirectory + "/" + shared.REMOVEDONEDIR + "/" + m.SelfID
	// if already written don't rewrite
	if !shared.FileExists(path) {
		// write own peer file also to done dir as removal already applied locally
		err = ioutil.WriteFile(path, []byte(""), shared.FILEPERMISSIONMODE)
		if err != nil {
			m.log("Couldn't write own peer file to", shared.REMOVEDONEDIR, "!", err.Error())
			return err
		}
	}
	// update model accordingly and return
	return m.updateLocal(removeDirectory)
}

/*
directRemove removes the given path from the model immediately without notifying
the update. NOTE: Do not use unless this is the behaviour you want! This method
is specifically a part of the normal applyRemove method, do not call outside
of it!
*/
func (m *Model) directRemove(path *shared.RelativePath) error {
	objList, err := m.partialPopulateMap(path.FullPath())
	if err != nil {
		m.log("partialPopulateMap failed in directRemove")
		return err
	}
	// iterate over each path
	for obj := range objList {
		relPath := path.Apply(obj)
		// if it still exists --> remove
		if shared.FileExists(relPath.FullPath()) {
			err := os.RemoveAll(relPath.FullPath())
			if err != nil {
				m.log("directRemove failed to remove the file itself!")
				return err
			}
		}
		// remove from model
		delete(m.TrackedPaths, relPath.SubPath())
		delete(m.StaticInfos, relPath.SubPath())
	}
	return nil
}

/*
isRemoved checks whether a file is due for deletion or whether it has already
been locally removed completely.
*/
func (m *Model) isRemoved(identification string) bool {
	path := m.Root + "/" + shared.TINZENITEDIR + "/" + shared.REMOVEDIR + "/" + identification
	return shared.FileExists(path) || m.isLocalRemoved(identification)
}

/*
makeLocalRemove is used to locally remember which removals have been applied
already, meaning the shared tracking of a file deletion has been removed.
*/
func (m *Model) makeLocalRemove(identification string) error {
	path := m.Root + "/" + shared.TINZENITEDIR + "/" + shared.LOCALDIR + "/" + shared.REMOVESTOREDIR + "/" + identification
	return ioutil.WriteFile(path, []byte(""), shared.FILEPERMISSIONMODE)
}

/*
isLocalRemoved notes whether a deletion may be being reintroduced even though it
was completely accepted.
*/
func (m *Model) isLocalRemoved(identification string) bool {
	path := m.Root + "/" + shared.TINZENITEDIR + "/" + shared.LOCALDIR + "/" + shared.REMOVESTOREDIR + "/" + identification
	return shared.FileExists(path)
}
<commit_msg>implemented removal of notify objects after timeout<commit_after>package model

import (
	"io/ioutil"
	"log"
	"os"
	"time"

	"github.com/tinzenite/shared"
)

/*
localRemove initiates a deletion locally, creating all necessary files and
removing the file from the model.
*/
func (m *Model) localRemove(path *shared.RelativePath) error {
	// get stin for notify
	stin, exists := m.StaticInfos[path.SubPath()]
	if !exists {
		m.log("LocalRemove: stin is missing!")
		return shared.ErrIllegalFileState
	}
	// sanity check
	if m.isRemoved(stin.Identification) {
		// shouldn't happen but let's be sure; warn at least
		m.warn("LocalRemove: file removal already begun!")
	}
	// direct remove
	err := m.directRemove(path)
	if err != nil {
		m.log("LocalRemove: failed to directly remove file!")
		return err
	}
	// write peers
	err = m.writeRemovalDir(stin.Identification)
	if err != nil {
		m.log("failed to update removal dir for", stin.Identification)
		return err
	}
	// update removal dir here so that creations etc are sent before notify below!
	err = m.updateLocal(m.Root + "/" + shared.TINZENITEDIR + "/" + shared.REMOVEDIR + "/" + stin.Identification)
	if err != nil {
		m.warn("partial update on local remove failed!")
		// but continue on because the changes will be synchronized later then anyway
	}
	// send notify
	notifyObj := &shared.ObjectInfo{
		Identification: stin.Identification,
		Name:           path.LastElement(),
		Path:           path.SubPath(),
		Content:        stin.Content,
		Version:        stin.Version,
		Directory:      stin.Directory}
	m.notify(shared.OpRemove, notifyObj)
	return nil
}

/*
remoteRemove handles a remote call of remove.

TODO this is buggy, fix it.
*/
func (m *Model) remoteRemove(path *shared.RelativePath, remoteObject *shared.ObjectInfo) error {
	// sanity check
	if remoteObject == nil {
		return shared.ErrIllegalParameters
	}
	// get state information
	localFileExists := m.IsTracked(path.FullPath())
	removalExists := m.isRemoved(remoteObject.Identification)
	// if still exists locally remove it
	if localFileExists {
		// remove file (removedir should already exist, so nothing else to do)
		err := m.directRemove(path)
		if err != nil {
			m.log("couldn't remove file", path.FullPath())
			return err
		}
	}
	// warn if remove dir didn't exist yet
	if !removalExists {
		// we'll create it anyway if it doesn't exist, so all ok but warn
		m.warn("remote file removed but removedir didn't yet exist!")
	}
	// since remote removal --> write peer to done
	err := m.writeRemovalDir(remoteObject.Identification)
	if err != nil {
		m.log("updating removal dir failed!")
		return err
	}
	// send notify (reuse remoteObject)
	m.notify(shared.OpRemove, remoteObject)
	return nil
}

/*
checkRemove checks whether a remove can be finally applied and purged from the
model dependent on the peers in done and check.
*/
func (m *Model) checkRemove() error {
	removeDir := m.Root + "/" + shared.TINZENITEDIR + "/" + shared.REMOVEDIR
	allRemovals, err := ioutil.ReadDir(removeDir)
	if err != nil {
		m.log("reading all removals failed")
		return err
	}
	// check for each removal
	for _, stat := range allRemovals {
		// update removal stats (including writing own peer into DONE for all of them)
		err = m.writeRemovalDir(stat.Name())
		if err != nil {
			log.Println("DEBUG: updating removal dir failed on checkRemove!", err)
			return err
		}
		// check if we can complete the removal
		err := m.completeTrackedRemoval(stat.Name())
		if err != nil {
			// notify of error but don't stop, rest can still be checked
			m.log("completeTrackedRemoval:", err.Error())
		}
		// warn of possible orphans
		if time.Since(stat.ModTime()) > removalTimeout {
			m.warn("Removal may be orphaned! ", stat.Name())
			/*TODO this may be called even if it has just been removed... do better logic!
			Also: is there something we can do in this case?*/
		}
		// warn of possibly unapplied removals:
		subPath, err := m.GetSubPath(stat.Name())
		// if err just skip the check (can happen if the file has been removed, so ok)
		if err == nil && m.IsTracked(m.Root+"/"+subPath) {
			m.warn("Removal may be unapplied!")
		}
	}
	// also remove old local remove notifies:
	localDir := m.Root + "/" + shared.TINZENITEDIR + "/" + shared.LOCALDIR + "/" + shared.REMOVESTOREDIR
	allLocals, err := ioutil.ReadDir(localDir)
	if err != nil {
		m.log("reading of local remove notifies failed!")
		return err
	}
	for _, stat := range allLocals {
		if time.Since(stat.ModTime()) > removalLocal {
			// remove notify
			err := os.Remove(localDir + "/" + stat.Name())
			if err != nil {
				m.warn("Failed to remove notify object:", err.Error())
			}
		}
	}
	return nil
}

/*
completeTrackedRemoval checks and if ok, removes the tracked removal dir, replacing
it with a local notify of the removal. This allows the tracked removal to be
purged. After a time out the local notify is also removed.
*/
func (m *Model) completeTrackedRemoval(identification string) error {
	removeDir := m.Root + "/" + shared.TINZENITEDIR + "/" + shared.REMOVEDIR
	// working directory
	objRemovePath := removeDir + "/" + identification
	// read all peers to check for
	allCheck, err := ioutil.ReadDir(objRemovePath + "/" + shared.REMOVECHECKDIR)
	if err != nil {
		m.log("Failed reading check peer list!")
		return err
	}
	// test whether we can remove it
	complete := true
	for _, peerStat := range allCheck {
		if !shared.FileExists(objRemovePath + "/" + shared.REMOVEDONEDIR + "/" + peerStat.Name()) {
			complete = false
			break
		}
	}
	if complete {
		// make local note of removal instead of tracked one so that we can remove it
		err := m.makeLocalRemove(identification)
		if err != nil {
			m.log("failed to write local remove note, will not complete removal!")
			return err
		}
		// HARD delete the entire dir: all peers should do the same (soft delete would make removal recursive)
		err = m.directRemove(shared.CreatePathRoot(m.Root).Apply(objRemovePath))
		if err != nil {
			m.log("Failed to direct remove!")
			return err
		}
		// note that other peers may not HARD delete it yet, but the isLocalRemoved check ensures that the dir isn't reintroduced
	}
	return nil
}

/*
writeRemovalDir is an internal function that writes all known peers to check
and the own peer to done, if not already existing.
*/
func (m *Model) writeRemovalDir(objIdentification string) error {
	removeDirectory := m.Root + "/" + shared.TINZENITEDIR + "/" + shared.REMOVEDIR + "/" + objIdentification
	// make directories if don't exist
	if !shared.FileExists(removeDirectory) {
		err := shared.MakeDirectories(removeDirectory, shared.REMOVECHECKDIR, shared.REMOVEDONEDIR)
		if err != nil {
			m.log("making removedir error")
			return err
		}
	}
	// write peer list to check which must all be notified of removal
	peers, err := m.readPeers()
	if err != nil {
		m.log("Failed to read peers")
		return err
	}
	for _, peer := range peers {
		path := removeDirectory + "/" + shared.REMOVECHECKDIR + "/" + peer
		// if already written don't rewrite
		if shared.FileExists(path) {
			continue
		}
		err = ioutil.WriteFile(path, []byte(""), shared.FILEPERMISSIONMODE)
		if err != nil {
			m.log("Couldn't write peer file", peer, "to", shared.REMOVECHECKDIR, "!")
			return err
		}
	}
	// write own peer into DONE
	path := removeDirectory + "/" + shared.REMOVEDONEDIR + "/" + m.SelfID
	// if already written don't rewrite
	if !shared.FileExists(path) {
		// write own peer file also to done dir as removal already applied locally
		err = ioutil.WriteFile(path, []byte(""), shared.FILEPERMISSIONMODE)
		if err != nil {
			m.log("Couldn't write own peer file to", shared.REMOVEDONEDIR, "!", err.Error())
			return err
		}
	}
	// update model accordingly and return
	return m.updateLocal(removeDirectory)
}

/*
directRemove removes the given path from the model immediately without notifying
the update. NOTE: Do not use unless this is the behaviour you want! This method
is specifically a part of the normal applyRemove method, do not call outside
of it!
*/
func (m *Model) directRemove(path *shared.RelativePath) error {
	objList, err := m.partialPopulateMap(path.FullPath())
	if err != nil {
		m.log("partialPopulateMap failed in directRemove")
		return err
	}
	// iterate over each path
	for obj := range objList {
		relPath := path.Apply(obj)
		// if it still exists --> remove
		if shared.FileExists(relPath.FullPath()) {
			err := os.RemoveAll(relPath.FullPath())
			if err != nil {
				m.log("directRemove failed to remove the file itself!")
				return err
			}
		}
		// remove from model
		delete(m.TrackedPaths, relPath.SubPath())
		delete(m.StaticInfos, relPath.SubPath())
	}
	return nil
}

/*
isRemoved checks whether a file is due for deletion or whether it has already
been locally removed completely.
*/
func (m *Model) isRemoved(identification string) bool {
	path := m.Root + "/" + shared.TINZENITEDIR + "/" + shared.REMOVEDIR + "/" + identification
	return shared.FileExists(path) || m.isLocalRemoved(identification)
}

/*
makeLocalRemove is used to locally remember which removals have been applied
already, meaning the shared tracking of a file deletion has been removed.
*/
func (m *Model) makeLocalRemove(identification string) error {
	path := m.Root + "/" + shared.TINZENITEDIR + "/" + shared.LOCALDIR + "/" + shared.REMOVESTOREDIR + "/" + identification
	return ioutil.WriteFile(path, []byte(""), shared.FILEPERMISSIONMODE)
}

/*
isLocalRemoved notes whether a deletion may be being reintroduced even though it
was completely accepted.
*/
func (m *Model) isLocalRemoved(identification string) bool {
	path := m.Root + "/" + shared.TINZENITEDIR + "/" + shared.LOCALDIR + "/" + shared.REMOVESTOREDIR + "/" + identification
	return shared.FileExists(path)
}
<|endoftext|>


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

package main

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/disk"
	"github.com/jacobsa/comeback/fs"
	"github.com/jacobsa/comeback/repr"
	"io"
	"log"
	"os"
	"path"
)

const (
	gTarget = "/tmp/restore_target"
)

var blobStore blob.Store

type score struct {
	hash []byte
}

func (s *score) Sha1Hash() []byte {
	return s.hash
}

func fromHexHash(h string) (blob.Score, error) {
	b, err := hex.DecodeString(h)
	if err != nil {
		return nil, fmt.Errorf("Invalid hex string: %s", h)
	}

	return &score{b}, nil
}

// Restore the file whose contents are described by the referenced blobs to the
// supplied target, whose parent must already exist.
func restoreFile(target string, scores []blob.Score) error {
	// Open the file.
	f, err := os.Create(target)
	defer f.Close()

	if err != nil {
		return fmt.Errorf("Create: %v", err)
	}

	// Process each blob.
	for _, score := range scores {
		// Load the blob.
		blob, err := blobStore.Load(score)
		if err != nil {
			return fmt.Errorf("Loading blob: %v", err)
		}

		// Write out its contents.
		_, err = io.Copy(f, bytes.NewReader(blob))
		if err != nil {
			return fmt.Errorf("Copy: %v", err)
		}
	}

	return nil
}

// Restore the directory whose contents are described by the referenced blob to
// the supplied target, which must already exist.
func restoreDir(target string, score blob.Score) error {
	// Load the appropriate blob.
	blob, err := blobStore.Load(score)
	if err != nil {
		return fmt.Errorf("Loading blob: %v", err)
	}

	// Parse its contents.
	entries, err := repr.Unmarshal(blob)
	if err != nil {
		return fmt.Errorf("Parsing blob: %v", err)
	}

	// Deal with each entry.
	for _, entry := range entries {
		entryPath := path.Join(target, entry.Name)

		// Switch on type.
		switch entry.Type {
		case fs.TypeFile:
			err = restoreFile(entryPath, entry.Scores)
			if err != nil {
				return err
			}

		case fs.TypeDirectory:
			if len(entry.Scores) != 1 {
				return fmt.Errorf("Wrong number of scores: %v", entry)
			}

			if err = os.Mkdir(entryPath, 0700); err != nil {
				return err
			}

			if err = restoreDir(entryPath, entry.Scores[0]); err != nil {
				return err
			}

		case fs.TypeSymlink:
			err = os.Symlink(entry.Target, entryPath)
			if err != nil {
				return err
			}

		default:
			return fmt.Errorf("Don't know how to deal with entry: %v", entry)
		}
	}

	return nil
}

func main() {
	var err error

	// Create the blob store.
	blobStore, err = disk.NewBlobStore("/tmp/blobs")
	if err != nil {
		log.Fatalf("Creating store: %v", err)
	}

	// Parse the score.
	score, err := fromHexHash("feb26f438b5706fe93281dd6d2c453eca9416c42")
	if err != nil {
		log.Fatalf("Parsing score: %v", err)
	}

	// Make sure the target doesn't exist.
	err = os.RemoveAll(gTarget)
	if err != nil {
		log.Fatalf("RemoveAll: %v", err)
	}

	// Create the target.
	err = os.Mkdir("/tmp/restore_target", 0700)
	if err != nil {
		log.Fatalf("Mkdir: %v", err)
	}

	// Attempt a restore.
	err = restoreDir(gTarget, score)
	if err != nil {
		log.Fatalf("Restoring: %v", err)
	}
}
<commit_msg>Added the ability to restore named pipes.<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/disk"
	"github.com/jacobsa/comeback/fs"
	"github.com/jacobsa/comeback/repr"
	"io"
	"log"
	"os"
	"path"
	"syscall"
)

const (
	gTarget = "/tmp/restore_target"
)

var blobStore blob.Store

type score struct {
	hash []byte
}

func (s *score) Sha1Hash() []byte {
	return s.hash
}

func fromHexHash(h string) (blob.Score, error) {
	b, err := hex.DecodeString(h)
	if err != nil {
		return nil, fmt.Errorf("Invalid hex string: %s", h)
	}

	return &score{b}, nil
}

// Restore the file whose contents are described by the referenced blobs to the
// supplied target, whose parent must already exist.
func restoreFile(target string, scores []blob.Score) error {
	// Open the file.
	f, err := os.Create(target)
	defer f.Close()

	if err != nil {
		return fmt.Errorf("Create: %v", err)
	}

	// Process each blob.
	for _, score := range scores {
		// Load the blob.
		blob, err := blobStore.Load(score)
		if err != nil {
			return fmt.Errorf("Loading blob: %v", err)
		}

		// Write out its contents.
		_, err = io.Copy(f, bytes.NewReader(blob))
		if err != nil {
			return fmt.Errorf("Copy: %v", err)
		}
	}

	return nil
}

// Restore the directory whose contents are described by the referenced blob to
// the supplied target, which must already exist.
func restoreDir(target string, score blob.Score) error {
	// Load the appropriate blob.
	blob, err := blobStore.Load(score)
	if err != nil {
		return fmt.Errorf("Loading blob: %v", err)
	}

	// Parse its contents.
	entries, err := repr.Unmarshal(blob)
	if err != nil {
		return fmt.Errorf("Parsing blob: %v", err)
	}

	// Deal with each entry.
	for _, entry := range entries {
		entryPath := path.Join(target, entry.Name)

		// Switch on type.
		switch entry.Type {
		case fs.TypeFile:
			err = restoreFile(entryPath, entry.Scores)
			if err != nil {
				return err
			}

		case fs.TypeDirectory:
			if len(entry.Scores) != 1 {
				return fmt.Errorf("Wrong number of scores: %v", entry)
			}

			if err = os.Mkdir(entryPath, 0700); err != nil {
				return err
			}

			if err = restoreDir(entryPath, entry.Scores[0]); err != nil {
				return err
			}

		case fs.TypeSymlink:
			err = os.Symlink(entry.Target, entryPath)
			if err != nil {
				return err
			}

		case fs.TypeNamedPipe:
			err = makeNamedPipe(entryPath, syscallPermissions(entry.Permissions))
			if err != nil {
				return err
			}

		default:
			return fmt.Errorf("Don't know how to deal with entry: %v", entry)
		}
	}

	return nil
}

func syscallPermissions(permissions os.FileMode) (o uint32) {
	// Include r/w/x permission bits.
	o = uint32(permissions & os.ModePerm)

	// Also include setuid/setgid/sticky bits.
	if permissions & os.ModeSetuid != 0 {
		o |= syscall.S_ISUID
	}

	if permissions & os.ModeSetgid != 0 {
		o |= syscall.S_ISGID
	}

	if permissions & os.ModeSticky != 0 {
		o |= syscall.S_ISVTX
	}

	return
}

// Create a named pipe at the supplied path.
func makeNamedPipe(path string, permissions uint32) error {
	return syscall.Mkfifo(path, permissions)
}

func main() {
	var err error

	// Create the blob store.
	blobStore, err = disk.NewBlobStore("/tmp/blobs")
	if err != nil {
		log.Fatalf("Creating store: %v", err)
	}

	// Parse the score.
	score, err := fromHexHash("feb26f438b5706fe93281dd6d2c453eca9416c42")
	if err != nil {
		log.Fatalf("Parsing score: %v", err)
	}

	// Make sure the target doesn't exist.
	err = os.RemoveAll(gTarget)
	if err != nil {
		log.Fatalf("RemoveAll: %v", err)
	}

	// Create the target.
	err = os.Mkdir("/tmp/restore_target", 0700)
	if err != nil {
		log.Fatalf("Mkdir: %v", err)
	}

	// Attempt a restore.
	err = restoreDir(gTarget, score)
	if err != nil {
		log.Fatalf("Restoring: %v", err)
	}
}
<|endoftext|>


<commit_before>package testing

import (
	"fmt"
	"log"

	"github.com/awslabs/aws-sdk-go/gen/swf"
	"github.com/sclasen/swfsm/fsm"
	. "github.com/sclasen/swfsm/sugar"
)

const (
	StubWorkflow      = "stub"
	ShortStubWorkflow = "stub"
	StubVersion       = "1"
)

var (
	StubTaskList      = &swf.TaskList{Name: S(fmt.Sprintf("%s->%s", StubWorkflow, StubVersion))}
	ShortStubTaskList = &swf.TaskList{Name: S(fmt.Sprintf("%s->%s", ShortStubWorkflow, StubVersion))}
)

type DecisionOutcome struct {
	DecisionTask *swf.DecisionTask
	State        string
	Decisions    []swf.Decision
}

func StubFSM(domain string, client fsm.SWFOps) *fsm.FSM {
	f := &fsm.FSM{
		SWF:        client,
		DataType:   make(map[string]interface{}),
		Domain:     domain,
		Name:       StubWorkflow,
		Serializer: fsm.JSONStateSerializer{},
		TaskList:   *StubTaskList.Name,
	}

	f.AddInitialState(&fsm.FSMState{Name: "Initial", Decider: StubState()})
	return f
}

func StubState() fsm.Decider {
	return func(ctx *fsm.FSMContext, h swf.HistoryEvent, data interface{}) fsm.Outcome {
		log.Printf("at=stub-event event=%+v", PrettyHistoryEvent(h))
		return ctx.Stay(data, ctx.EmptyDecisions())
	}
}

func ShortStubFSM(domain string, client fsm.SWFOps) *fsm.FSM {
	f := &fsm.FSM{
		SWF:        client,
		DataType:   make(map[string]interface{}),
		Domain:     domain,
		Name:       ShortStubWorkflow,
		Serializer: fsm.JSONStateSerializer{},
		TaskList:   *StubTaskList.Name,
	}

	f.AddInitialState(&fsm.FSMState{Name: "Initial", Decider: ShortStubState()})
	return f
}

func ShortStubState() fsm.Decider {
	return func(ctx *fsm.FSMContext, h swf.HistoryEvent, data interface{}) fsm.Outcome {
		log.Printf("at=short-stub-event event=%+v", PrettyHistoryEvent(h))
		return ctx.CompleteWorkflow(data)
	}
}

//intercept any attempts to start a workflow and launch the stub workflow instead.
func TestInterceptor(testID string, stubbedWorkflows, stubbedShortWorkflows []string) *fsm.FuncInterceptor {
	stubbed := make(map[string]struct{})
	stubbedShort := make(map[string]struct{})
	v := struct{}{}
	for _, s := range stubbedWorkflows {
		stubbed[s] = v
	}
	for _, s := range stubbedShortWorkflows {
		stubbedShort[s] = v
	}
	return &fsm.FuncInterceptor{
		AfterDecisionFn: func(decision *swf.DecisionTask, ctx *fsm.FSMContext, outcome *fsm.Outcome) {
			for _, d := range outcome.Decisions {
				switch *d.DecisionType {
				case swf.DecisionTypeStartChildWorkflowExecution:
					if _, ok := stubbed[*d.StartChildWorkflowExecutionDecisionAttributes.WorkflowType.Name]; ok {
						d.StartChildWorkflowExecutionDecisionAttributes.WorkflowType.Name = S(StubWorkflow)
						d.StartChildWorkflowExecutionDecisionAttributes.WorkflowType.Version = S(StubVersion)
						d.StartChildWorkflowExecutionDecisionAttributes.ExecutionStartToCloseTimeout = S("360")
						d.StartChildWorkflowExecutionDecisionAttributes.TaskList = StubTaskList
					}
					if _, ok := stubbedShort[*d.StartChildWorkflowExecutionDecisionAttributes.WorkflowType.Name]; ok {
						d.StartChildWorkflowExecutionDecisionAttributes.WorkflowType.Name = S(ShortStubWorkflow)
						d.StartChildWorkflowExecutionDecisionAttributes.WorkflowType.Version = S(StubVersion)
						d.StartChildWorkflowExecutionDecisionAttributes.ExecutionStartToCloseTimeout = S("360")
						d.StartChildWorkflowExecutionDecisionAttributes.TaskList = ShortStubTaskList
					}
				case swf.DecisionTypeScheduleActivityTask:
					d.ScheduleActivityTaskDecisionAttributes.TaskList = &swf.TaskList{Name: S(*d.ScheduleActivityTaskDecisionAttributes.TaskList.Name + testID)}
				}
			}
		},
	}
}

func TestReplicator(decisionOutcomes chan DecisionOutcome) fsm.ReplicationHandler {
	return func(ctx *fsm.FSMContext, task *swf.DecisionTask, outcome *swf.RespondDecisionTaskCompletedInput, state *fsm.SerializedState) error {
		decisionOutcomes <- DecisionOutcome{State: state.StateName, DecisionTask: task, Decisions: outcome.Decisions}
		return nil
	}
}

func TestSWF(client fsm.ClientSWFOps, stubbedWorkflow ...string) fsm.ClientSWFOps {
	stubbed := make(map[string]struct{})
	v := struct{}{}
	for _, s := range stubbedWorkflow {
		stubbed[s] = v
	}
	return &StubSWFClient{
		ClientSWFOps:     client,
		stubbedWorkflows: stubbed,
	}
}

//intercept any attempts to start a workflow and launch the stub workflow instead.
type StubSWFClient struct {
	fsm.ClientSWFOps
	stubbedWorkflows map[string]struct{}
}

func (s *StubSWFClient) StartWorkflowExecution(req *swf.StartWorkflowExecutionInput) (resp *swf.Run, err error) {
	if _, ok := s.stubbedWorkflows[*req.WorkflowType.Name]; ok {
		req.WorkflowType.Name = S(StubWorkflow)
		req.WorkflowType.Version = S(StubVersion)
		req.ExecutionStartToCloseTimeout = S("360")
		req.TaskList = StubTaskList
	}
	return s.ClientSWFOps.StartWorkflowExecution(req)
}
<commit_msg>intercept continueAsNew and set the task list<commit_after>package testing

import (
	"fmt"
	"log"

	"github.com/awslabs/aws-sdk-go/gen/swf"
	"github.com/sclasen/swfsm/fsm"
	. "github.com/sclasen/swfsm/sugar"
)

const (
	StubWorkflow      = "stub"
	ShortStubWorkflow = "stub"
	StubVersion       = "1"
)

var (
	StubTaskList      = &swf.TaskList{Name: S(fmt.Sprintf("%s->%s", StubWorkflow, StubVersion))}
	ShortStubTaskList = &swf.TaskList{Name: S(fmt.Sprintf("%s->%s", ShortStubWorkflow, StubVersion))}
)

type DecisionOutcome struct {
	DecisionTask *swf.DecisionTask
	State        string
	Decisions    []swf.Decision
}

func StubFSM(domain string, client fsm.SWFOps) *fsm.FSM {
	f := &fsm.FSM{
		SWF:        client,
		DataType:   make(map[string]interface{}),
		Domain:     domain,
		Name:       StubWorkflow,
		Serializer: fsm.JSONStateSerializer{},
		TaskList:   *StubTaskList.Name,
	}

	f.AddInitialState(&fsm.FSMState{Name: "Initial", Decider: StubState()})
	return f
}

func StubState() fsm.Decider {
	return func(ctx *fsm.FSMContext, h swf.HistoryEvent, data interface{}) fsm.Outcome {
		log.Printf("at=stub-event event=%+v", PrettyHistoryEvent(h))
		return ctx.Stay(data, ctx.EmptyDecisions())
	}
}

func ShortStubFSM(domain string, client fsm.SWFOps) *fsm.FSM {
	f := &fsm.FSM{
		SWF:        client,
		DataType:   make(map[string]interface{}),
		Domain:     domain,
		Name:       ShortStubWorkflow,
		Serializer: fsm.JSONStateSerializer{},
		TaskList:   *StubTaskList.Name,
	}

	f.AddInitialState(&fsm.FSMState{Name: "Initial", Decider: ShortStubState()})
	return f
}

func ShortStubState() fsm.Decider {
	return func(ctx *fsm.FSMContext, h swf.HistoryEvent, data interface{}) fsm.Outcome {
		log.Printf("at=short-stub-event event=%+v", PrettyHistoryEvent(h))
		return ctx.CompleteWorkflow(data)
	}
}

//intercept any attempts to start a workflow and launch the stub workflow instead.
func TestInterceptor(testID string, stubbedWorkflows, stubbedShortWorkflows []string) *fsm.FuncInterceptor {
	stubbed := make(map[string]struct{})
	stubbedShort := make(map[string]struct{})
	v := struct{}{}
	for _, s := range stubbedWorkflows {
		stubbed[s] = v
	}
	for _, s := range stubbedShortWorkflows {
		stubbedShort[s] = v
	}
	return &fsm.FuncInterceptor{
		AfterDecisionFn: func(decision *swf.DecisionTask, ctx *fsm.FSMContext, outcome *fsm.Outcome) {
			for _, d := range outcome.Decisions {
				switch *d.DecisionType {
				case swf.DecisionTypeStartChildWorkflowExecution:
					if _, ok := stubbed[*d.StartChildWorkflowExecutionDecisionAttributes.WorkflowType.Name]; ok {
						d.StartChildWorkflowExecutionDecisionAttributes.WorkflowType.Name = S(StubWorkflow)
						d.StartChildWorkflowExecutionDecisionAttributes.WorkflowType.Version = S(StubVersion)
						d.StartChildWorkflowExecutionDecisionAttributes.ExecutionStartToCloseTimeout = S("360")
						d.StartChildWorkflowExecutionDecisionAttributes.TaskList = StubTaskList
					}
					if _, ok := stubbedShort[*d.StartChildWorkflowExecutionDecisionAttributes.WorkflowType.Name]; ok {
						d.StartChildWorkflowExecutionDecisionAttributes.WorkflowType.Name = S(ShortStubWorkflow)
						d.StartChildWorkflowExecutionDecisionAttributes.WorkflowType.Version = S(StubVersion)
						d.StartChildWorkflowExecutionDecisionAttributes.ExecutionStartToCloseTimeout = S("360")
						d.StartChildWorkflowExecutionDecisionAttributes.TaskList = ShortStubTaskList
					}
				case swf.DecisionTypeScheduleActivityTask:
					d.ScheduleActivityTaskDecisionAttributes.TaskList = &swf.TaskList{Name: S(*d.ScheduleActivityTaskDecisionAttributes.TaskList.Name + testID)}
				case swf.DecisionTypeContinueAsNewWorkflowExecution:
					d.ContinueAsNewWorkflowExecutionDecisionAttributes.TaskList = &swf.TaskList{Name: S(*d.ContinueAsNewWorkflowExecutionDecisionAttributes.TaskList.Name + testID)}
				}
			}
		},
	}
}

func TestReplicator(decisionOutcomes chan DecisionOutcome) fsm.ReplicationHandler {
	return func(ctx *fsm.FSMContext, task *swf.DecisionTask, outcome *swf.RespondDecisionTaskCompletedInput, state *fsm.SerializedState) error {
		decisionOutcomes <- DecisionOutcome{State: state.StateName, DecisionTask: task, Decisions: outcome.Decisions}
		return nil
	}
}

func TestSWF(client fsm.ClientSWFOps, stubbedWorkflow ...string) fsm.ClientSWFOps {
	stubbed := make(map[string]struct{})
	v := struct{}{}
	for _, s := range stubbedWorkflow {
		stubbed[s] = v
	}
	return &StubSWFClient{
		ClientSWFOps:     client,
		stubbedWorkflows: stubbed,
	}
}

//intercept any attempts to start a workflow and launch the stub workflow instead.
type StubSWFClient struct {
	fsm.ClientSWFOps
	stubbedWorkflows map[string]struct{}
}

func (s *StubSWFClient) StartWorkflowExecution(req *swf.StartWorkflowExecutionInput) (resp *swf.Run, err error) {
	if _, ok := s.stubbedWorkflows[*req.WorkflowType.Name]; ok {
		req.WorkflowType.Name = S(StubWorkflow)
		req.WorkflowType.Version = S(StubVersion)
		req.ExecutionStartToCloseTimeout = S("360")
		req.TaskList = StubTaskList
	}
	return s.ClientSWFOps.StartWorkflowExecution(req)
}
<|endoftext|>


<commit_before>package main

import (
	"net"
	"net/http"
	"runtime"
	"sync"
	"time"

	log "github.com/funkygao/log4go"
	"github.com/julienschmidt/httprouter"
)

type webServer struct {
	name       string
	maxClients int
	gw         *Gateway

	httpListener net.Listener
	httpServer   *http.Server

	httpsListener net.Listener
	httpsServer   *http.Server

	router *httprouter.Router

	waitExitFunc waitExitFunc
}

func newWebServer(name string, httpAddr, httpsAddr string, maxClients int,
	gw *Gateway) *webServer {
	this := &webServer{
		name:       name,
		gw:         gw,
		maxClients: maxClients,
		router:     httprouter.New(),
	}

	if httpAddr != "" {
		this.httpServer = &http.Server{
			Addr:           httpAddr,
			Handler:        this.router,
			ReadTimeout:    httpReadTimeout,
			WriteTimeout:   httpWriteTimeout,
			MaxHeaderBytes: httpHeaderMaxBytes,
		}
	}

	if httpsAddr != "" {
		this.httpsServer = &http.Server{
			Addr:           httpsAddr,
			Handler:        this.router,
			ReadTimeout:    httpReadTimeout,
			WriteTimeout:   httpWriteTimeout,
			MaxHeaderBytes: httpHeaderMaxBytes,
		}
	}

	return this
}

func (this *webServer) Start() {
	if this.httpsServer != nil {
		this.startServer(true)
	}

	if this.httpServer != nil {
		this.startServer(false)
	}

}

func (this *webServer) Router() *httprouter.Router {
	return this.router
}

func (this *webServer) startServer(https bool) {
	var err error
	waitListenerUp := make(chan struct{})
	go func() {
		if options.cpuAffinity {
			runtime.LockOSThread()
		}

		var (
			retryDelay         time.Duration
			theListener        net.Listener
			waitListenerUpOnce sync.Once
			addr               string
		)
		for {
			if https {
				this.httpsListener, err = net.Listen("tcp", this.httpsServer.Addr)
				if err != nil {
					if retryDelay == 0 {
						retryDelay = 50 * time.Millisecond
					} else {
						retryDelay = 2 * retryDelay
					}
					if maxDelay := time.Second; retryDelay > maxDelay {
						retryDelay = maxDelay
					}
					log.Error("%s listener %v, retry in %v", this.name, err, retryDelay)
					time.Sleep(retryDelay)
					continue
				}

				this.httpsListener, err = setupHttpsListener(this.httpsListener,
					this.gw.certFile, this.gw.keyFile)
				if err != nil {
					panic(err)
				}

				theListener = this.httpsListener
				addr = this.httpsServer.Addr
			} else {
				this.httpListener, err = net.Listen("tcp", this.httpServer.Addr)
				theListener = this.httpListener
				addr = this.httpServer.Addr
			}

			if err != nil {
				if retryDelay == 0 {
					retryDelay = 50 * time.Millisecond
				} else {
					retryDelay = 2 * retryDelay
				}
				if maxDelay := time.Second; retryDelay > maxDelay {
					retryDelay = maxDelay
				}
				log.Error("%s listener %v, retry in %v", this.name, err, retryDelay)
				time.Sleep(retryDelay)
				continue
			}

			theListener = LimitListener(this.gw, theListener, this.maxClients)
			waitListenerUpOnce.Do(func() {
				close(waitListenerUp)
			})

			// on non-temporary err, net/http will close the listener
			if https {
				err = this.httpsServer.Serve(theListener)
			} else {
				err = this.httpServer.Serve(theListener)
			}

			select {
			case <-this.gw.shutdownCh:
				log.Trace("%s server stopped on %s", this.name, addr)
				return

			default:
				log.Error("%s server: %v", this.name, err)
			}
		}
	}()

	// FIXME if net.Listen fails, kateway will not be able to stop
	// e,g. start a kateway, then start another, the 2nd will not be able to stop

	// wait for the listener up
	<-waitListenerUp

	this.gw.wg.Add(1)
	if this.waitExitFunc == nil {
		this.waitExitFunc = this.waitExit
	}
	if https {
		go this.waitExitFunc(this.httpsServer, this.httpsListener, this.gw.shutdownCh)
	} else {
		go this.waitExitFunc(this.httpServer, this.httpListener, this.gw.shutdownCh)
	}

	if https {
		log.Info("%s https server ready on %s", this.name, this.httpsServer.Addr)
	} else {
		log.Info("%s http server ready on %s", this.name, this.httpServer.Addr)
	}
}

func (this *webServer) waitExit(server *http.Server, listener net.Listener, exit <-chan struct{}) {
	<-exit

	// HTTP response will have "Connection: close"
	server.SetKeepAlivesEnabled(false)

	// avoid new connections
	if err := listener.Close(); err != nil {
		log.Error(err.Error())
	}

	this.gw.wg.Done()
}
<commit_msg>web server waitExitFunc default val set on start<commit_after>package main

import (
	"net"
	"net/http"
	"runtime"
	"sync"
	"sync/atomic"
	"time"

	log "github.com/funkygao/log4go"
	"github.com/julienschmidt/httprouter"
)

type webServer struct {
	name       string
	maxClients int
	gw         *Gateway

	httpListener net.Listener
	httpServer   *http.Server

	httpsListener net.Listener
	httpsServer   *http.Server

	router *httprouter.Router

	waitExitFunc waitExitFunc
}

func newWebServer(name string, httpAddr, httpsAddr string, maxClients int,
	gw *Gateway) *webServer {
	this := &webServer{
		name:       name,
		gw:         gw,
		maxClients: maxClients,
		router:     httprouter.New(),
	}

	if httpAddr != "" {
		this.httpServer = &http.Server{
			Addr:           httpAddr,
			Handler:        this.router,
			ReadTimeout:    httpReadTimeout,
			WriteTimeout:   httpWriteTimeout,
			MaxHeaderBytes: httpHeaderMaxBytes,
		}
	}

	if httpsAddr != "" {
		this.httpsServer = &http.Server{
			Addr:           httpsAddr,
			Handler:        this.router,
			ReadTimeout:    httpReadTimeout,
			WriteTimeout:   httpWriteTimeout,
			MaxHeaderBytes: httpHeaderMaxBytes,
		}
	}

	return this
}

func (this *webServer) Start() {
	if this.waitExitFunc == nil {
		this.waitExitFunc = this.waitExit
	}

	if this.httpsServer != nil {
		this.startServer(true)
	}

	if this.httpServer != nil {
		this.startServer(false)
	}

}

func (this *webServer) Router() *httprouter.Router {
	return this.router
}

func (this *webServer) startServer(https bool) {
	var err error
	waitListenerUp := make(chan struct{})
	go func() {
		if options.cpuAffinity {
			runtime.LockOSThread()
		}

		var (
			retryDelay         time.Duration
			theListener        net.Listener
			waitListenerUpOnce sync.Once
			addr               string
		)
		for {
			if https {
				this.httpsListener, err = net.Listen("tcp", this.httpsServer.Addr)
				if err != nil {
					if retryDelay == 0 {
						retryDelay = 50 * time.Millisecond
					} else {
						retryDelay = 2 * retryDelay
					}
					if maxDelay := time.Second; retryDelay > maxDelay {
						retryDelay = maxDelay
					}
					log.Error("%s listener %v, retry in %v", this.name, err, retryDelay)
					time.Sleep(retryDelay)
					continue
				}

				this.httpsListener, err = setupHttpsListener(this.httpsListener,
					this.gw.certFile, this.gw.keyFile)
				if err != nil {
					panic(err)
				}

				theListener = this.httpsListener
				addr = this.httpsServer.Addr
			} else {
				this.httpListener, err = net.Listen("tcp", this.httpServer.Addr)
				theListener = this.httpListener
				addr = this.httpServer.Addr
			}

			if err != nil {
				if retryDelay == 0 {
					retryDelay = 50 * time.Millisecond
				} else {
					retryDelay = 2 * retryDelay
				}
				if maxDelay := time.Second; retryDelay > maxDelay {
					retryDelay = maxDelay
				}
				log.Error("%s listener %v, retry in %v", this.name, err, retryDelay)
				time.Sleep(retryDelay)
				continue
			}

			theListener = LimitListener(this.gw, theListener, this.maxClients)
			waitListenerUpOnce.Do(func() {
				close(waitListenerUp)
			})

			// on non-temporary err, net/http will close the listener
			if https {
				err = this.httpsServer.Serve(theListener)
			} else {
				err = this.httpServer.Serve(theListener)
			}

			select {
			case <-this.gw.shutdownCh:
				log.Trace("%s server stopped on %s", this.name, addr)
				return

			default:
				log.Error("%s server: %v", this.name, err)
			}
		}
	}()

	// FIXME if net.Listen fails, kateway will not be able to stop
	// e,g. start a kateway, then start another, the 2nd will not be able to stop

	// wait for the listener up
	<-waitListenerUp

	this.gw.wg.Add(1)
	if https {
		go this.waitExitFunc(this.httpsServer, this.httpsListener, this.gw.shutdownCh)
	} else {
		go this.waitExitFunc(this.httpServer, this.httpListener, this.gw.shutdownCh)
	}

	if https {
		log.Info("%s https server ready on %s", this.name, this.httpsServer.Addr)
	} else {
		log.Info("%s http server ready on %s", this.name, this.httpServer.Addr)
	}
}

func (this *webServer) waitExit(server *http.Server, listener net.Listener, exit <-chan struct{}) {
	<-exit

	// HTTP response will have "Connection: close"
	server.SetKeepAlivesEnabled(false)

	// avoid new connections
	if err := listener.Close(); err != nil {
		log.Error(err.Error())
	}

	this.gw.wg.Done()
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/siddontang/go/num"
	"github.com/siddontang/ledisdb/config"
	"github.com/siddontang/ledisdb/ledis"
	"os"
	"runtime"
	"sync"
	"sync/atomic"
	"time"
)

var KB = config.KB
var MB = config.MB
var GB = config.GB

var name = flag.String("db_name", "goleveldb", "db name")
var number = flag.Int("n", 10000, "request number")
var clients = flag.Int("c", 50, "number of clients")
var round = flag.Int("r", 1, "benchmark round number")
var valueSize = flag.Int("vsize", 100, "kv value size")
var wg sync.WaitGroup

var ldb *ledis.Ledis
var db *ledis.DB

var loop int = 0

func bench(cmd string, f func()) {
	wg.Add(*clients)

	t1 := time.Now()
	for i := 0; i < *clients; i++ {
		go func() {
			for j := 0; j < loop; j++ {
				f()
			}
			wg.Done()
		}()
	}

	wg.Wait()

	t2 := time.Now()

	d := t2.Sub(t1)
	fmt.Printf("%s %s: %0.3f micros/op, %0.2fmb/s %0.2fop/s\n",
		cmd,
		d.String(),
		float64(d.Nanoseconds()/1e3)/float64(*number),
		float64((*valueSize+16)*(*number))/(1024.0*1024.0*(d.Seconds())),
		float64(*number)/d.Seconds())
}

var kvSetBase int64 = 0
var kvGetBase int64 = 0

func benchSet() {
	f := func() {
		value := make([]byte, *valueSize)
		n := atomic.AddInt64(&kvSetBase, 1)

		db.Set(num.Int64ToBytes(n), value)
	}

	bench("set", f)
}

func benchGet() {
	f := func() {
		n := atomic.AddInt64(&kvGetBase, 1)
		v, err := db.Get(num.Int64ToBytes(n))
		if err != nil {
			println(err.Error())
		} else if len(v) != *valueSize {
			println(len(v), *valueSize)
		}
	}

	bench("get", f)
}

func setRocksDB(cfg *config.RocksDBConfig) {
	cfg.BlockSize = 64 * KB
	cfg.WriteBufferSize = 64 * MB
	cfg.MaxWriteBufferNum = 2
	cfg.MaxBytesForLevelBase = 512 * MB
	cfg.TargetFileSizeBase = 64 * MB
	cfg.BackgroundThreads = 4
	cfg.HighPriorityBackgroundThreads = 1
	cfg.MaxBackgroundCompactions = 3
	cfg.MaxBackgroundFlushes = 1
	cfg.CacheSize = 512 * MB
	cfg.EnableStatistics = true
	cfg.StatsDumpPeriodSec = 5
	cfg.Level0FileNumCompactionTrigger = 8
	cfg.MaxBytesForLevelMultiplier = 8
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	flag.Parse()

	cfg := config.NewConfigDefault()
	cfg.DBPath = "./var/db_test"
	cfg.DBName = *name
	os.RemoveAll(cfg.DBPath)

	cfg.LevelDB.BlockSize = 32 * KB
	cfg.LevelDB.CacheSize = 512 * MB
	cfg.LevelDB.WriteBufferSize = 64 * MB
	cfg.LevelDB.MaxOpenFiles = 1000

	setRocksDB(&cfg.RocksDB)

	var err error
	ldb, err = ledis.Open(cfg)
	if err != nil {
		panic(err)
		return
	}

	db, _ = ldb.Select(0)

	if *number <= 0 {
		panic("invalid number")
		return
	}

	if *clients <= 0 || *number < *clients {
		panic("invalid client number")
		return
	}

	loop = *number / *clients

	if *round <= 0 {
		*round = 1
	}

	for i := 0; i < *round; i++ {
		benchSet()
		benchGet()

		println("")
	}
}
<commit_msg>bugfix ledis-dbbench<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/siddontang/go/num"
	"github.com/siddontang/ledisdb/config"
	"github.com/siddontang/ledisdb/ledis"
	"os"
	"runtime"
	"sync"
	"sync/atomic"
	"time"
)

var KB = config.KB
var MB = config.MB
var GB = config.GB

var name = flag.String("db_name", "goleveldb", "db name")
var number = flag.Int("n", 10000, "request number")
var clients = flag.Int("c", 50, "number of clients")
var round = flag.Int("r", 1, "benchmark round number")
var valueSize = flag.Int("vsize", 100, "kv value size")
var wg sync.WaitGroup

var ldb *ledis.Ledis
var db *ledis.DB

var loop int = 0

func bench(cmd string, f func()) {
	wg.Add(*clients)

	t1 := time.Now()
	for i := 0; i < *clients; i++ {
		go func() {
			for j := 0; j < loop; j++ {
				f()
			}
			wg.Done()
		}()
	}

	wg.Wait()

	t2 := time.Now()

	d := t2.Sub(t1)
	fmt.Printf("%s %s: %0.3f micros/op, %0.2fmb/s %0.2fop/s\n",
		cmd,
		d.String(),
		float64(d.Nanoseconds()/1e3)/float64(*number),
		float64((*valueSize+16)*(*number))/(1024.0*1024.0*(d.Seconds())),
		float64(*number)/d.Seconds())
}

var kvSetBase int64 = 0
var kvGetBase int64 = 0

func benchSet() {
	f := func() {
		value := make([]byte, *valueSize)
		n := atomic.AddInt64(&kvSetBase, 1)

		db.Set(num.Int64ToBytes(n), value)
	}

	bench("set", f)
}

func benchGet() {
	f := func() {
		n := atomic.AddInt64(&kvGetBase, 1)
		v, err := db.Get(num.Int64ToBytes(n))
		if err != nil {
			println(err.Error())
		} else if len(v) != *valueSize {
			println(len(v), *valueSize)
		}
	}

	bench("get", f)
}

func setRocksDB(cfg *config.RocksDBConfig) {
	cfg.BlockSize = 64 * KB
	cfg.WriteBufferSize = 64 * MB
	cfg.MaxWriteBufferNum = 2
	cfg.MaxBytesForLevelBase = 512 * MB
	cfg.TargetFileSizeBase = 64 * MB
	cfg.BackgroundThreads = 4
	cfg.HighPriorityBackgroundThreads = 1
	cfg.MaxBackgroundCompactions = 3
	cfg.MaxBackgroundFlushes = 1
	cfg.CacheSize = 512 * MB
	cfg.EnableStatistics = true
	cfg.StatsDumpPeriodSec = 5
	cfg.Level0FileNumCompactionTrigger = 8
	cfg.MaxBytesForLevelMultiplier = 8
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	flag.Parse()

	cfg := config.NewConfigDefault()
	cfg.DataDir = "./var/ledis_dbbench"
	cfg.DBName = *name
	os.RemoveAll(cfg.DBPath)
	defer os.RemoveAll(cfg.DBPath)

	os.MkdirAll(cfg.DBPath, 0755)

	cfg.LevelDB.BlockSize = 32 * KB
	cfg.LevelDB.CacheSize = 512 * MB
	cfg.LevelDB.WriteBufferSize = 64 * MB
	cfg.LevelDB.MaxOpenFiles = 1000

	setRocksDB(&cfg.RocksDB)

	var err error
	ldb, err = ledis.Open(cfg)
	if err != nil {
		println(err.Error())
		return
	}

	db, _ = ldb.Select(0)

	if *number <= 0 {
		panic("invalid number")
		return
	}

	if *clients <= 0 || *number < *clients {
		panic("invalid client number")
		return
	}

	loop = *number / *clients

	if *round <= 0 {
		*round = 1
	}

	for i := 0; i < *round; i++ {
		benchSet()
		benchGet()

		println("")
	}
}
<|endoftext|>