<commit_before>package githubprovider

import (
	"errors"
	"fmt"
	"strings"

	"github.com/google/go-github/github"
	"github.com/hashicorp/terraform/helper/schema"
)

var (
	admin  string = "admin"
	member string = "member"
)

// required field are here for adding a user to the organization
func resourceGithubAddUser() *schema.Resource {
	return &schema.Resource{
		Create: resourceGithubAddUserCreate,
		Read:   resourceGithubAddUserRead,
		Update: resourceGithubAddUserCreate,
		Delete: resourceGithubAddUserDelete,

		Schema: map[string]*schema.Schema{
			"username": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},

			// role is the required for the membership
			// its value is member as default.
			"role": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Default:  "member",
			},

			// repos is the repos that the organization has
			"repos": &schema.Schema{
				Type:     schema.TypeList,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Required: true,
			},

			// organization is the name of the organization
			"organization": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},

			"teams": &schema.Schema{
				Type:     schema.TypeList,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Required: true,
			},

			// title is the title of the SSH Key
			"title": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},

			// SSHKey is the public key of the user
			"SSHKey": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
		},
	}
}

// GetTeamIDs gets the teams id of the organization
func GetTeamIDs(client *github.Client, org string, teamNames []string) ([]int, error) {
	currentPage := 0

	var teamIDs []int

	for {
		options := &github.ListOptions{
			PerPage: 100,
			Page:    currentPage,
		}

		teams, resp, err := client.Organizations.ListTeams(org, options)
		if err != nil {
			return nil, err
		}

		if len(teams) == 0 {
			break
		}
		// Iterate over all teams and add current user to related team(s)
		for i, team := range teams {
			for _, teamName := range teamNames {
				if *team.Name == teamName {
					teamIDs = append(teamIDs, *teams[i].ID)
				}
			}
		}

		if currentPage == resp.LastPage {
			break
		}

		currentPage = resp.NextPage
	}

	return teamIDs, nil
}

// resourceGithubAddUserCreate adds the user to the organization & the teams
func resourceGithubAddUserCreate(d *schema.ResourceData, meta interface{}) error {
	clientOrg := meta.(*Clients).OrgClient
	client := meta.(*Clients).UserClient

	org := d.Get("organization").(string)
	user := d.Get("username").(string)
	teamNames := interfaceToStringSlice(d.Get("teams"))
	role := d.Get("role").(string)

	if err := checkScopePermissions(client, user); err != nil {
		return err
	}

	if len(teamNames) == 0 {
		return errors.New("team name is not defined")
	}

	member, resp, err := client.Organizations.GetOrgMembership("", org)
	// user might be a member of organization key
	// otherwise it will response 404, and we need to ignore that error
	if err != nil && resp.StatusCode != 404 {
		return err
	}

	// override member role here if user is admin of organization
	if member.Role == &admin {
		role = admin
	}

	teamIDs, err := GetTeamIDs(clientOrg, org, teamNames)

	optAddOrgMembership := &github.OrganizationAddTeamMembershipOptions{
		Role: role,
	}

	if len(teamNames) != len(teamIDs) {
		return errors.New("team name is not found")
	}

	for _, teamID := range teamIDs {
		_, _, err := clientOrg.Organizations.AddTeamMembership(teamID, user, optAddOrgMembership)
		if err != nil {
			return err
		}
	}

	active := "active"

	membership := &github.Membership{
		// state should be active to add the user into organization
		State: &active,

		// Role is the required for the membership
		Role: &role,
	}

	// EditOrgMembership edits the membership for user in specified organization.
	// if user is authenticated, we dont need to set 1.parameter as user
	_, _, err = client.Organizations.EditOrgMembership("", org, membership)
	if err != nil {
		return err
	}

	for _, repo := range interfaceToStringSlice(d.Get("repos")) {
		// Creates a fork for the authenticated user.
		_, _, err = client.Repositories.CreateFork(org, repo, nil)
		if err != nil {
			return err
		}
	}

	title := d.Get("title").(string)
	keySSH := d.Get("SSHKey").(string)

	key := &github.Key{
		Title: &title,
		Key:   &keySSH,
	}

	// CreateKey creates a public key. Requires that you are authenticated via Basic Auth,
	// or OAuth with at least `write:public_key` scope.
	//
	// If SSH key is already set up, when u try to add same SSHKEY then
	//you are gonna get 422: Validation error.
	_, _, err = client.Users.CreateKey(key)
	if err != nil && !isErr422ValidationFailed(err) {
		return err
	}

	d.SetId(user)

	return nil
}

func resourceGithubAddUserRead(d *schema.ResourceData, meta interface{}) error {
	org := d.Get("organization").(string)
	user := d.Get("username").(string)
	role := d.Get("role").(string)
	teamNames := interfaceToStringSlice(d.Get("teams"))
	repos := interfaceToStringSlice(d.Get("repos"))
	fmt.Println("org: %v, user: %v,role: %v, teamnames: %v, repos: %v",
		org,
		user,
		role,
		teamNames,
		repos,
	)

	return nil
}

// resourceGithubAddUserCreate removes the user from the organization & the teams
func resourceGithubAddUserDelete(d *schema.ResourceData, meta interface{}) error {
	// We'r not gonna use removemember for now.
	// And then we can simply return here
	return nil

	client := meta.(*Clients).OrgClient

	user := d.Get("username").(string)
	org := d.Get("organization").(string)

	// Removing a user from this list will remove them from all teams and
	// they will no longer have any access to the organization’s repositories.
	_, err := client.Organizations.RemoveMember(org, user)
	return err
}

func getKeyID(client *github.Client, user, title string) (int, error) {
	keys, _, err := client.Users.ListKeys(user, nil)
	if err != nil {
		return 0, err
	}

	for _, key := range keys {
		if *key.Title == title {
			return *key.ID, nil
		}
	}

	return 0, err
}

// interfaceToStringSlice converts the interface to slice of string
func interfaceToStringSlice(s interface{}) []string {
	slice, ok := s.([]interface{})
	if !ok {
		return nil
	}

	sslice := make([]string, len(slice))
	for i := range slice {
		sslice[i] = slice[i].(string)
	}

	return sslice
}

func checkScopePermissions(client *github.Client, username string) error {
	arr, err := getScopes(client, username)
	if err != nil {
		return err
	}

	// we created 2-dimensional array for scopes.
	scopeArray := [][]string{
		// if user enables one of this scopes, then its OK to go..
		{"write:public_key", "admin:public_key"},
		{"user"},
		{"repo", "public_repo"},
		{"admin:org"},
	}

	for _, scopeElement := range scopeArray {
		if !(isInArray(arr, scopeElement)) {
			scopeErr := fmt.Errorf("Could not find required scope :", scopeElement)
			return scopeErr
		}
	}

	return nil
}

func getScopes(client *github.Client, username string) ([]string, error) {
	var scopes []string
	_, resp, err := client.Users.Get(username)
	if err != nil {
		return scopes, err
	}

	list := resp.Header.Get("X-Oauth-Scopes")
	scopes = strings.Split(list, ", ")

	return scopes, nil
}

func isInArray(arr, item []string) bool {
	for _, a := range arr {
		for _, i := range item {
			if a == i {
				return true
			}
		}
	}
	return false
}

// isErr422ValidationFailed return true if error contains the string:
// '422 Validation Failed'. This error is special cased so we can ignore it on
// when it occurs during rebuilding of stack template.
func isErr422ValidationFailed(err error) bool {
	return err != nil && strings.Contains(err.Error(), "422 Validation Failed")
}
<commit_msg>github-provider: add clear comment for error checking<commit_after>package githubprovider

import (
	"errors"
	"fmt"
	"strings"

	"github.com/google/go-github/github"
	"github.com/hashicorp/terraform/helper/schema"
)

var (
	admin  string = "admin"
	member string = "member"
)

// required field are here for adding a user to the organization
func resourceGithubAddUser() *schema.Resource {
	return &schema.Resource{
		Create: resourceGithubAddUserCreate,
		Read:   resourceGithubAddUserRead,
		Update: resourceGithubAddUserCreate,
		Delete: resourceGithubAddUserDelete,

		Schema: map[string]*schema.Schema{
			"username": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},

			// role is the required for the membership
			// its value is member as default.
			"role": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Default:  "member",
			},

			// repos is the repos that the organization has
			"repos": &schema.Schema{
				Type:     schema.TypeList,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Required: true,
			},

			// organization is the name of the organization
			"organization": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},

			"teams": &schema.Schema{
				Type:     schema.TypeList,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Required: true,
			},

			// title is the title of the SSH Key
			"title": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},

			// SSHKey is the public key of the user
			"SSHKey": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
		},
	}
}

// GetTeamIDs gets the teams id of the organization
func GetTeamIDs(client *github.Client, org string, teamNames []string) ([]int, error) {
	currentPage := 0

	var teamIDs []int

	for {
		options := &github.ListOptions{
			PerPage: 100,
			Page:    currentPage,
		}

		teams, resp, err := client.Organizations.ListTeams(org, options)
		if err != nil {
			return nil, err
		}

		if len(teams) == 0 {
			break
		}
		// Iterate over all teams and add current user to related team(s)
		for i, team := range teams {
			for _, teamName := range teamNames {
				if *team.Name == teamName {
					teamIDs = append(teamIDs, *teams[i].ID)
				}
			}
		}

		if currentPage == resp.LastPage {
			break
		}

		currentPage = resp.NextPage
	}

	return teamIDs, nil
}

// resourceGithubAddUserCreate adds the user to the organization & the teams
func resourceGithubAddUserCreate(d *schema.ResourceData, meta interface{}) error {
	clientOrg := meta.(*Clients).OrgClient
	client := meta.(*Clients).UserClient

	org := d.Get("organization").(string)
	user := d.Get("username").(string)
	teamNames := interfaceToStringSlice(d.Get("teams"))
	role := d.Get("role").(string)

	if err := checkScopePermissions(client, user); err != nil {
		return err
	}

	if len(teamNames) == 0 {
		return errors.New("team name is not defined")
	}

	member, resp, err := client.Organizations.GetOrgMembership("", org)
	// user might be a member of organization key or not
	// with that error checking;
	// if user is a member of organization, we can change user's role as admin or member
	// if user is not a member of organization, it will give 404 error, then we need to ignore that
	// error for process (adding user into the organization etc..)
	if err != nil && resp.StatusCode != 404 {
		return err
	}

	// override member role here if user is admin of organization
	if member.Role == &admin {
		role = admin
	}

	teamIDs, err := GetTeamIDs(clientOrg, org, teamNames)

	optAddOrgMembership := &github.OrganizationAddTeamMembershipOptions{
		Role: role,
	}

	if len(teamNames) != len(teamIDs) {
		return errors.New("team name is not found")
	}

	for _, teamID := range teamIDs {
		_, _, err := clientOrg.Organizations.AddTeamMembership(teamID, user, optAddOrgMembership)
		if err != nil {
			return err
		}
	}

	active := "active"

	membership := &github.Membership{
		// state should be active to add the user into organization
		State: &active,

		// Role is the required for the membership
		Role: &role,
	}

	// EditOrgMembership edits the membership for user in specified organization.
	// if user is authenticated, we dont need to set 1.parameter as user
	_, _, err = client.Organizations.EditOrgMembership("", org, membership)
	if err != nil {
		return err
	}

	for _, repo := range interfaceToStringSlice(d.Get("repos")) {
		// Creates a fork for the authenticated user.
		_, _, err = client.Repositories.CreateFork(org, repo, nil)
		if err != nil {
			return err
		}
	}

	title := d.Get("title").(string)
	keySSH := d.Get("SSHKey").(string)

	key := &github.Key{
		Title: &title,
		Key:   &keySSH,
	}

	// CreateKey creates a public key. Requires that you are authenticated via Basic Auth,
	// or OAuth with at least `write:public_key` scope.
	//
	// If SSH key is already set up, when u try to add same SSHKEY then
	//you are gonna get 422: Validation error.
	_, _, err = client.Users.CreateKey(key)
	if err != nil && !isErr422ValidationFailed(err) {
		return err
	}

	d.SetId(user)

	return nil
}

func resourceGithubAddUserRead(d *schema.ResourceData, meta interface{}) error {
	org := d.Get("organization").(string)
	user := d.Get("username").(string)
	role := d.Get("role").(string)
	teamNames := interfaceToStringSlice(d.Get("teams"))
	repos := interfaceToStringSlice(d.Get("repos"))
	fmt.Println("org: %v, user: %v,role: %v, teamnames: %v, repos: %v",
		org,
		user,
		role,
		teamNames,
		repos,
	)

	return nil
}

// resourceGithubAddUserCreate removes the user from the organization & the teams
func resourceGithubAddUserDelete(d *schema.ResourceData, meta interface{}) error {
	// We'r not gonna use removemember for now.
	// And then we can simply return here
	return nil

	client := meta.(*Clients).OrgClient

	user := d.Get("username").(string)
	org := d.Get("organization").(string)

	// Removing a user from this list will remove them from all teams and
	// they will no longer have any access to the organization’s repositories.
	_, err := client.Organizations.RemoveMember(org, user)
	return err
}

func getKeyID(client *github.Client, user, title string) (int, error) {
	keys, _, err := client.Users.ListKeys(user, nil)
	if err != nil {
		return 0, err
	}

	for _, key := range keys {
		if *key.Title == title {
			return *key.ID, nil
		}
	}

	return 0, err
}

// interfaceToStringSlice converts the interface to slice of string
func interfaceToStringSlice(s interface{}) []string {
	slice, ok := s.([]interface{})
	if !ok {
		return nil
	}

	sslice := make([]string, len(slice))
	for i := range slice {
		sslice[i] = slice[i].(string)
	}

	return sslice
}

func checkScopePermissions(client *github.Client, username string) error {
	arr, err := getScopes(client, username)
	if err != nil {
		return err
	}

	// we created 2-dimensional array for scopes.
	scopeArray := [][]string{
		// if user enables one of this scopes, then its OK to go..
		{"write:public_key", "admin:public_key"},
		{"user"},
		{"repo", "public_repo"},
		{"admin:org"},
	}

	for _, scopeElement := range scopeArray {
		if !(isInArray(arr, scopeElement)) {
			scopeErr := fmt.Errorf("Could not find required scope :", scopeElement)
			return scopeErr
		}
	}

	return nil
}

func getScopes(client *github.Client, username string) ([]string, error) {
	var scopes []string
	_, resp, err := client.Users.Get(username)
	if err != nil {
		return scopes, err
	}

	list := resp.Header.Get("X-Oauth-Scopes")
	scopes = strings.Split(list, ", ")

	return scopes, nil
}

func isInArray(arr, item []string) bool {
	for _, a := range arr {
		for _, i := range item {
			if a == i {
				return true
			}
		}
	}
	return false
}

// isErr422ValidationFailed return true if error contains the string:
// '422 Validation Failed'. This error is special cased so we can ignore it on
// when it occurs during rebuilding of stack template.
func isErr422ValidationFailed(err error) bool {
	return err != nil && strings.Contains(err.Error(), "422 Validation Failed")
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"flag"
	"fmt"
	"garden/tools/pgperf/model"
	"io/ioutil"
	"log"
	"math/rand"
	"os"
	"regexp"
	"strings"
	"sync"
	"time"

	_ "github.com/lib/pq"
	yaml "gopkg.in/yaml.v2"
)

var pms = []int{1, 200, 2000}

func main() {
	ps := usage()
	if len(os.Args) <= 1 {
		ps = &params{
			numConcurrency: &pms[0],
			numRows:        &pms[1],
			maxOperation:   &pms[2],
		}
		log.Println(message())
	}
	path, err := os.Getwd()
	if err != nil {
		log.Fatal(err)
	}
	cfg := new(PgDatabase)
	cfg.parse(path + "/config.yaml")
	//fmt.Println("length of channel buffer", runtime.NumGoroutine())
	db, err := sql.Open(cfg.Driver, cfg.getDSN())
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()
	infoHw()
	//fmt.Println(time.Now(), "db number is", runtime.NumGoroutine())
	start := time.Now()

	model.Drop(db, "VcState")
	model.Create(db, "VcState")
	model.Insert(db, "VcState", *ps.numRows)
	end := time.Now()
	dur := end.Sub(start)
	fmt.Printf("CREATE And INSERT %v items total duration %v\n", *ps.numRows, dur)

	//fmt.Println("Concurrency number", *ps.numConcurrency)
	concurrencyKey := make(chan int, *ps.numConcurrency)
	var wg sync.WaitGroup
	if ps.isOnly {
		for n := 0; n < *ps.maxOperation; n++ {
			wg.Add(1)
			concurrencyKey <- n
			go func(i int) {
				defer func() {
					wg.Done()
					<-concurrencyKey
				}()
				vcID := getRandNumber(i+*ps.maxOperation, *ps.numRows)
				model.OnlyUpdate(db, "VcState", vcID)
			}(n)
		}
		wg.Wait()
		close(concurrencyKey)
		tail := time.Now()
		dur = tail.Sub(end)
	} else {
		for n := 0; n < *ps.maxOperation; n++ {
			wg.Add(1)
			concurrencyKey <- n
			go func(i int) {
				defer func() {
					wg.Done()
					<-concurrencyKey
				}()
				vcID := getRandNumber(i+*ps.maxOperation, *ps.numRows)
				model.CmpUpdate(db, "VcState", vcID)
			}(n)
		}
		wg.Wait()
		close(concurrencyKey)
		tail := time.Now()
		dur = tail.Sub(end)
	}
	avg := float64(dur.Nanoseconds()) / float64(*ps.maxOperation) / 1000
	fmt.Println("opertion avg", avg, "macroseconds")
	nps := float64(*ps.maxOperation) * (1.0 / dur.Seconds())
	fmt.Printf("UPDATE %v operations duration %v,  %v n/s  \n", *ps.maxOperation, time.Now().Sub(end), nps)
}

//PgDatabase Database configure information
type PgDatabase struct {
	Driver   string `yaml:"driver"`
	Host     string `yaml:"host"`
	Username string `yaml:"username"`
	Password string `yaml:"password"`
	DBName   string `yaml:"dbname"`
	Switch   string `yaml:"switch"`
	Port     string `yaml:"-"`
	Sslmode  string `yaml:"sslmode"`
}

// Parse config.yaml to data struct
func (cfg *PgDatabase) parse(file string) {
	body, err := ioutil.ReadFile(file)
	if err != nil {
		log.Fatal(err)
	}
	err = yaml.Unmarshal(body, cfg)
	if err != nil {
		log.Fatal(err)
	}
}

func (cfg *PgDatabase) getDSN() string {
	var dsn string
	if strings.ToLower(cfg.Switch) == "on" {
		dsn = fmt.Sprintf("host=%s user=%s password=%s dbname=%s sslmode=%s", cfg.Host, cfg.Username, cfg.Password, cfg.DBName, cfg.Sslmode)
	}
	return dsn
}

type params struct {
	numConcurrency *int
	numRows        *int
	maxOperation   *int
	isOnly         bool
}

var numConcurrency = flag.Int("g", pms[0], "UPDATE goroutine jobs")
var numRows = flag.Int("r", pms[1], "Table ROWs number")
var maxOperation = flag.Int("op", pms[2], "Max UPDATE operation numbers")
var isOnly = flag.String("oy", "false", "Operation ONLY update or not")

func usage() *params {
	flag.Parse()
	ps := new(params)
	ps.numConcurrency = numConcurrency
	ps.numRows = numRows
	ps.maxOperation = maxOperation
	*isOnly = strings.ToLower(*isOnly)
	oyRe := regexp.MustCompile("^y|yes|t|true$")
	ps.isOnly = oyRe.MatchString(*isOnly)
	return ps
}

func message() string {
	return fmt.Sprintf(`Now %s is running with default parameters as follows
 1.Parameter -g |--g : UPDATE goroutine jobs default
 2.Parameter -r |--r : Table ROWs number default 200
 3.Parameter -op|--op: Max UPDATE operation numbers default 2000
 4.Parameter -oy|--oy: Operation ONLY update or not,
                       accept [Yy]|[Yy][Ee][Ss]|[Tt]|[Tt][Rr][Uu][Ee] as true, default false

More manaul with %s -h|--help
`, os.Args[0], os.Args[0])

}

//infoHw Show hardware information
func infoHw() {
	model.InfoCPU()
	model.InfoMem()
}

func getRandNumber(scope, numRows int) int {
	//rand.Seed(time.Now().UnixNano())
	s := rand.Intn(scope)
	return s%numRows + 1
}
<commit_msg>Golang pprof support sample for performance profile<commit_after>package main

import (
	"database/sql"
	"flag"
	"fmt"
	"garden/tools/pgperf/model"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	_ "net/http/pprof"
	"os"
	"regexp"
	"strings"
	"sync"
	"time"

	_ "github.com/lib/pq"
	yaml "gopkg.in/yaml.v2"
)

var pms = []int{1, 200, 2000}

func main() {
	go func() {
		log.Println(http.ListenAndServe("localhost:6060", nil))
	}()
	ps := usage()
	if len(os.Args) <= 1 {
		ps = &params{
			numConcurrency: &pms[0],
			numRows:        &pms[1],
			maxOperation:   &pms[2],
		}
		log.Println(message())
	}
	path, err := os.Getwd()
	if err != nil {
		log.Fatal(err)
	}
	cfg := new(PgDatabase)
	cfg.parse(path + "/config.yaml")
	//fmt.Println("length of channel buffer", runtime.NumGoroutine())
	db, err := sql.Open(cfg.Driver, cfg.getDSN())
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()
	infoHw()
	//fmt.Println(time.Now(), "db number is", runtime.NumGoroutine())
	start := time.Now()

	model.Drop(db, "VcState")
	model.Create(db, "VcState")
	model.Insert(db, "VcState", *ps.numRows)
	end := time.Now()
	dur := end.Sub(start)
	fmt.Printf("CREATE And INSERT %v items total duration %v\n", *ps.numRows, dur)

	//fmt.Println("Concurrency number", *ps.numConcurrency)
	concurrencyKey := make(chan int, *ps.numConcurrency)
	var wg sync.WaitGroup
	if ps.isOnly {
		for n := 0; n < *ps.maxOperation; n++ {
			wg.Add(1)
			concurrencyKey <- n
			go func(i int) {
				defer func() {
					wg.Done()
					<-concurrencyKey
				}()
				vcID := getRandNumber(i+*ps.maxOperation, *ps.numRows)
				model.OnlyUpdate(db, "VcState", vcID)
			}(n)
		}
		wg.Wait()
		close(concurrencyKey)
		tail := time.Now()
		dur = tail.Sub(end)
	} else {
		for n := 0; n < *ps.maxOperation; n++ {
			wg.Add(1)
			concurrencyKey <- n
			go func(i int) {
				defer func() {
					wg.Done()
					<-concurrencyKey
				}()
				vcID := getRandNumber(i+*ps.maxOperation, *ps.numRows)
				model.CmpUpdate(db, "VcState", vcID)
			}(n)
		}
		wg.Wait()
		close(concurrencyKey)
		tail := time.Now()
		dur = tail.Sub(end)
	}
	avg := float64(dur.Nanoseconds()) / float64(*ps.maxOperation) / 1000
	fmt.Println("opertion avg", avg, "macroseconds")
	nps := float64(*ps.maxOperation) * (1.0 / dur.Seconds())
	fmt.Printf("UPDATE %v operations duration %v,  %v n/s  \n", *ps.maxOperation, time.Now().Sub(end), nps)
}

//PgDatabase Database configure information
type PgDatabase struct {
	Driver   string `yaml:"driver"`
	Host     string `yaml:"host"`
	Username string `yaml:"username"`
	Password string `yaml:"password"`
	DBName   string `yaml:"dbname"`
	Switch   string `yaml:"switch"`
	Port     string `yaml:"-"`
	Sslmode  string `yaml:"sslmode"`
}

// Parse config.yaml to data struct
func (cfg *PgDatabase) parse(file string) {
	body, err := ioutil.ReadFile(file)
	if err != nil {
		log.Fatal(err)
	}
	err = yaml.Unmarshal(body, cfg)
	if err != nil {
		log.Fatal(err)
	}
}

func (cfg *PgDatabase) getDSN() string {
	var dsn string
	if strings.ToLower(cfg.Switch) == "on" {
		dsn = fmt.Sprintf("host=%s user=%s password=%s dbname=%s sslmode=%s", cfg.Host, cfg.Username, cfg.Password, cfg.DBName, cfg.Sslmode)
	}
	return dsn
}

type params struct {
	numConcurrency *int
	numRows        *int
	maxOperation   *int
	isOnly         bool
}

var numConcurrency = flag.Int("g", pms[0], "UPDATE goroutine jobs")
var numRows = flag.Int("r", pms[1], "Table ROWs number")
var maxOperation = flag.Int("op", pms[2], "Max UPDATE operation numbers")
var isOnly = flag.String("oy", "false", "Operation ONLY update or not")

func usage() *params {
	flag.Parse()
	ps := new(params)
	ps.numConcurrency = numConcurrency
	ps.numRows = numRows
	ps.maxOperation = maxOperation
	*isOnly = strings.ToLower(*isOnly)
	oyRe := regexp.MustCompile("^y|yes|t|true$")
	ps.isOnly = oyRe.MatchString(*isOnly)
	return ps
}

func message() string {
	return fmt.Sprintf(`Now %s is running with default parameters as follows
 1.Parameter -g |--g : UPDATE goroutine jobs default
 2.Parameter -r |--r : Table ROWs number default 200
 3.Parameter -op|--op: Max UPDATE operation numbers default 2000
 4.Parameter -oy|--oy: Operation ONLY update or not,
                       accept [Yy]|[Yy][Ee][Ss]|[Tt]|[Tt][Rr][Uu][Ee] as true, default false

More manaul with %s -h|--help
`, os.Args[0], os.Args[0])

}

//infoHw Show hardware information
func infoHw() {
	model.InfoCPU()
	model.InfoMem()
}

func getRandNumber(scope, numRows int) int {
	//rand.Seed(time.Now().UnixNano())
	s := rand.Intn(scope)
	return s%numRows + 1
}
<|endoftext|>


<commit_before>// Package red implements server speaking redis serialization protocol.
package red

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"runtime/debug"
	"strings"
	"sync"
	"time"

	"github.com/artyom/resp"
)

// HandlerFunc is a type of command processing function. It should return value
// that's passed to resp.Encode function, so it's expected such function would
// return values of resp package types.
type HandlerFunc func(req Request) (interface{}, error)

// Request holds information about single redis command.
type Request struct {
	Name string   // lowercase command itself (first word)
	Args []string // command arguments
}

// NewServer returns initialized server.
func NewServer() *Server {
	return &Server{log: noopLogger{}}
}

// WithLogger configures server to use provided Logger.
func (s *Server) WithLogger(l Logger) {
	if l != nil {
		s.log = l
	}
}

// WithUnsafeTx disables command processing serialization that guarantees
// redis-like transaction safety.
func (s *Server) WithUnsafeTx() { s.unsafe = true }

// Handle registers handler for command with given name (case-insensitive)
func (s *Server) Handle(name string, h HandlerFunc) {
	if name == "" {
		panic("Handle called with empty name")
	}
	if h == nil {
		panic("Handle called with nil HandlerFunc")
	}
	if s.handlers == nil {
		s.handlers = make(map[string]HandlerFunc)
	}
	s.handlers[strings.ToLower(name)] = h
}

// Server implements server speaking RESP (REdis Serialization Protocol). Server
// automatically handles MULTI, EXEC, DISCARD commands for transactions, QUIT for
// client-initiated disconnect, other commands are expected to be implemented
// separately and registered with Handle method.
type Server struct {
	log      Logger
	handlers map[string]HandlerFunc
	mu       sync.Mutex // used to serialize transactions
	unsafe   bool       // whether to skip transaction serialization
}

func (s *Server) txLock() {
	if s.unsafe {
		return
	}
	s.mu.Lock()
}
func (s *Server) txUnlock() {
	if s.unsafe {
		return
	}
	s.mu.Unlock()
}

// HandleConn processes single client connection, automatically handling
// following commands MULTI/EXEC/DISCARD (transactions), QUIT (client disconnect). It
// calls user-provided handlers for registered commands.
func (s *Server) HandleConn(conn io.ReadWriteCloser) error {
	defer conn.Close()
	rd := bufio.NewReader(conn)
	var tx []Request
	var inTx bool  // if we're inside transaction
	var errTx bool // true if transaction seen error and should be discarded
	var err error
	for {
		if err != nil {
			return err
		}
		req, err := resp.DecodeRequest(rd)
		switch err {
		case nil:
		case resp.ErrInvalidRequest:
			err = resp.Encode(conn, resp.Error("ERR unknown command"))
			continue
		default:
			return err
		}
		cmd := strings.ToLower(req[0])
		switch cmd {
		case "quit":
			return resp.Encode(conn, resp.OK)
		case "discard":
			if len(req) != 1 {
				if inTx {
					errTx = true
				}
				err = resp.Encode(conn, errWrongArgs(cmd))
				continue
			}
			if !inTx {
				err = resp.Encode(conn, resp.Error("ERR DISCARD without MULTI"))
				continue
			}
			inTx, errTx, tx = false, false, tx[:0]
			err = resp.Encode(conn, resp.OK)
			continue
		case "multi":
			if len(req) != 1 {
				if inTx {
					errTx = true
				}
				err = resp.Encode(conn, errWrongArgs(cmd))
				continue
			}
			if inTx {
				errTx = true
				err = resp.Encode(conn, resp.Error("ERR MULTI calls can not be nested"))
				continue
			}
			inTx, errTx = true, false
			err = resp.Encode(conn, resp.OK)
			continue
		case "exec":
			if len(req) != 1 {
				if inTx {
					errTx = true
				}
				err = resp.Encode(conn, errWrongArgs(cmd))
				continue
			}
			if !inTx {
				err = resp.Encode(conn, resp.Error("ERR EXEC without MULTI"))
				continue
			}
			if errTx {
				inTx, errTx = false, false
				tx = tx[:0]
				err = resp.Encode(conn, resp.Error("EXECABORT Transaction discarded because of previous errors."))
				continue
			}
		default:
			h, ok := s.handlers[cmd]
			if !ok {
				if inTx {
					errTx = true
				}
				err = resp.Encode(conn, errNoCmd(cmd))
				continue
			}
			if inTx {
				if !errTx {
					tx = append(tx, Request{Name: cmd, Args: req[1:]})
				}
				err = resp.Encode(conn, resp.SimpleString("QUEUED"))
				continue
			}
			s.txLock()
			err = resp.Encode(conn, runHandler(h, Request{Name: cmd, Args: req[1:]}))
			s.txUnlock()
			continue
		}

		txReplies := make(resp.Array, 0, len(tx))
		s.txLock()
		for _, r := range tx {
			h, ok := s.handlers[r.Name]
			if !ok {
				txReplies = append(txReplies, errNoCmd(r.Name))
				continue
			}
			txReplies = append(txReplies, runHandler(h, r))
		}
		s.txUnlock()
		inTx, errTx = false, false
		tx = tx[:0]
		err = resp.Encode(conn, txReplies)
	}
}

// ListenAndServe listens on TCP network address addr and then calls Serve to
// handle requests on incoming connections.
func (s *Server) ListenAndServe(addr string) error {
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	return s.Serve(tcpKeepAliveListener{ln.(*net.TCPListener)})
}

// Serve accepts incoming connections on the Listener l, creating a new service
// goroutine for each.
func (s *Server) Serve(l net.Listener) error {
	defer l.Close()
	for {
		conn, err := l.Accept()
		if err != nil {
			return err
		}
		go func(c net.Conn) {
			switch err := s.HandleConn(c); err {
			case nil, io.EOF:
			default:
				if s.log != nil {
					s.log.Println(err)
				}
			}
		}(conn)
	}
}

// Logger is a set of methods used to log information. *log.Logger implements
// this interface.
type Logger interface {
	Print(v ...interface{})
	Printf(format string, v ...interface{})
	Println(v ...interface{})
}

type noopLogger struct{}

func (noopLogger) Print(v ...interface{})                 {}
func (noopLogger) Printf(format string, v ...interface{}) {}
func (noopLogger) Println(v ...interface{})               {}

func errNoCmd(name string) resp.Error { return resp.Error("ERR unknown command '" + name + "'") }

// ErrWrongArgs are expected to be returned by HandlerFunc implementations when
// number of arguments are wrong. This error is automatically passed to client
// with command name annotated.
var ErrWrongArgs = errors.New("wrong number of arguments")

func errWrongArgs(name string) resp.Error {
	return resp.Error("ERR wrong number of arguments for '" + name + "' command")
}

// runHandler calls HandlerFunc h with given Request r and returns single value
// that is expected to be passed to resp.Encode.
//
// If h returns a non-nil error or panics, then runHandler returns resp.Error
// type, otherwise it's an interface{} returned by h.
func runHandler(h HandlerFunc, r Request) interface{} {
	var panicked bool
	v, err := func() (interface{}, error) {
		defer func() {
			if x := recover(); x != nil {
				panicked = true
				fmt.Fprintln(os.Stderr, "panic recovered:", x)
				debug.PrintStack()
			}
		}()
		return h(r)
	}()
	if panicked {
		return resp.Error("ERR internal server error (panic)")
	}
	if err != nil {
		if err == ErrWrongArgs {
			return errWrongArgs(r.Name)
		}
		text := err.Error()
		if strings.ContainsAny(text, "\r\n") {
			text = strings.Replace(text, "\r", " ", -1)
			text = strings.Replace(text, "\n", " ", -1)
		}
		return resp.Error("ERR " + text)
	}
	return v
}

// tcpKeepAliveListener sets TCP keep-alive timeouts on accepted
// connections. It's used by ListenAndServe and ListenAndServeTLS so
// dead TCP connections (e.g. closing laptop mid-download) eventually
// go away.
type tcpKeepAliveListener struct {
	*net.TCPListener
}

func (ln tcpKeepAliveListener) Accept() (c net.Conn, err error) {
	tc, err := ln.AcceptTCP()
	if err != nil {
		return
	}
	tc.SetKeepAlive(true)
	tc.SetKeepAlivePeriod(3 * time.Minute)
	return tc, nil
}
<commit_msg>Remove unused noopLogger<commit_after>// Package red implements server speaking redis serialization protocol.
package red

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"runtime/debug"
	"strings"
	"sync"
	"time"

	"github.com/artyom/resp"
)

// HandlerFunc is a type of command processing function. It should return value
// that's passed to resp.Encode function, so it's expected such function would
// return values of resp package types.
type HandlerFunc func(req Request) (interface{}, error)

// Request holds information about single redis command.
type Request struct {
	Name string   // lowercase command itself (first word)
	Args []string // command arguments
}

// NewServer returns initialized server.
func NewServer() *Server { return new(Server) }

// WithLogger configures server to use provided Logger.
func (s *Server) WithLogger(l Logger) {
	if l != nil {
		s.log = l
	}
}

// WithUnsafeTx disables command processing serialization that guarantees
// redis-like transaction safety.
func (s *Server) WithUnsafeTx() { s.unsafe = true }

// Handle registers handler for command with given name (case-insensitive)
func (s *Server) Handle(name string, h HandlerFunc) {
	if name == "" {
		panic("Handle called with empty name")
	}
	if h == nil {
		panic("Handle called with nil HandlerFunc")
	}
	if s.handlers == nil {
		s.handlers = make(map[string]HandlerFunc)
	}
	s.handlers[strings.ToLower(name)] = h
}

// Server implements server speaking RESP (REdis Serialization Protocol). Server
// automatically handles MULTI, EXEC, DISCARD commands for transactions, QUIT for
// client-initiated disconnect, other commands are expected to be implemented
// separately and registered with Handle method.
type Server struct {
	log      Logger
	handlers map[string]HandlerFunc
	mu       sync.Mutex // used to serialize transactions
	unsafe   bool       // whether to skip transaction serialization
}

func (s *Server) txLock() {
	if s.unsafe {
		return
	}
	s.mu.Lock()
}
func (s *Server) txUnlock() {
	if s.unsafe {
		return
	}
	s.mu.Unlock()
}

// HandleConn processes single client connection, automatically handling
// following commands MULTI/EXEC/DISCARD (transactions), QUIT (client disconnect). It
// calls user-provided handlers for registered commands.
func (s *Server) HandleConn(conn io.ReadWriteCloser) error {
	defer conn.Close()
	rd := bufio.NewReader(conn)
	var tx []Request
	var inTx bool  // if we're inside transaction
	var errTx bool // true if transaction seen error and should be discarded
	var err error
	for {
		if err != nil {
			return err
		}
		req, err := resp.DecodeRequest(rd)
		switch err {
		case nil:
		case resp.ErrInvalidRequest:
			err = resp.Encode(conn, resp.Error("ERR unknown command"))
			continue
		default:
			return err
		}
		cmd := strings.ToLower(req[0])
		switch cmd {
		case "quit":
			return resp.Encode(conn, resp.OK)
		case "discard":
			if len(req) != 1 {
				if inTx {
					errTx = true
				}
				err = resp.Encode(conn, errWrongArgs(cmd))
				continue
			}
			if !inTx {
				err = resp.Encode(conn, resp.Error("ERR DISCARD without MULTI"))
				continue
			}
			inTx, errTx, tx = false, false, tx[:0]
			err = resp.Encode(conn, resp.OK)
			continue
		case "multi":
			if len(req) != 1 {
				if inTx {
					errTx = true
				}
				err = resp.Encode(conn, errWrongArgs(cmd))
				continue
			}
			if inTx {
				errTx = true
				err = resp.Encode(conn, resp.Error("ERR MULTI calls can not be nested"))
				continue
			}
			inTx, errTx = true, false
			err = resp.Encode(conn, resp.OK)
			continue
		case "exec":
			if len(req) != 1 {
				if inTx {
					errTx = true
				}
				err = resp.Encode(conn, errWrongArgs(cmd))
				continue
			}
			if !inTx {
				err = resp.Encode(conn, resp.Error("ERR EXEC without MULTI"))
				continue
			}
			if errTx {
				inTx, errTx = false, false
				tx = tx[:0]
				err = resp.Encode(conn, resp.Error("EXECABORT Transaction discarded because of previous errors."))
				continue
			}
		default:
			h, ok := s.handlers[cmd]
			if !ok {
				if inTx {
					errTx = true
				}
				err = resp.Encode(conn, errNoCmd(cmd))
				continue
			}
			if inTx {
				if !errTx {
					tx = append(tx, Request{Name: cmd, Args: req[1:]})
				}
				err = resp.Encode(conn, resp.SimpleString("QUEUED"))
				continue
			}
			s.txLock()
			err = resp.Encode(conn, runHandler(h, Request{Name: cmd, Args: req[1:]}))
			s.txUnlock()
			continue
		}

		txReplies := make(resp.Array, 0, len(tx))
		s.txLock()
		for _, r := range tx {
			h, ok := s.handlers[r.Name]
			if !ok {
				txReplies = append(txReplies, errNoCmd(r.Name))
				continue
			}
			txReplies = append(txReplies, runHandler(h, r))
		}
		s.txUnlock()
		inTx, errTx = false, false
		tx = tx[:0]
		err = resp.Encode(conn, txReplies)
	}
}

// ListenAndServe listens on TCP network address addr and then calls Serve to
// handle requests on incoming connections.
func (s *Server) ListenAndServe(addr string) error {
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	return s.Serve(tcpKeepAliveListener{ln.(*net.TCPListener)})
}

// Serve accepts incoming connections on the Listener l, creating a new service
// goroutine for each.
func (s *Server) Serve(l net.Listener) error {
	defer l.Close()
	for {
		conn, err := l.Accept()
		if err != nil {
			return err
		}
		go func(c net.Conn) {
			switch err := s.HandleConn(c); err {
			case nil, io.EOF:
			default:
				if s.log != nil {
					s.log.Println(err)
				}
			}
		}(conn)
	}
}

// Logger is a set of methods used to log information. *log.Logger implements
// this interface.
type Logger interface {
	Print(v ...interface{})
	Printf(format string, v ...interface{})
	Println(v ...interface{})
}

func errNoCmd(name string) resp.Error { return resp.Error("ERR unknown command '" + name + "'") }

// ErrWrongArgs are expected to be returned by HandlerFunc implementations when
// number of arguments are wrong. This error is automatically passed to client
// with command name annotated.
var ErrWrongArgs = errors.New("wrong number of arguments")

func errWrongArgs(name string) resp.Error {
	return resp.Error("ERR wrong number of arguments for '" + name + "' command")
}

// runHandler calls HandlerFunc h with given Request r and returns single value
// that is expected to be passed to resp.Encode.
//
// If h returns a non-nil error or panics, then runHandler returns resp.Error
// type, otherwise it's an interface{} returned by h.
func runHandler(h HandlerFunc, r Request) interface{} {
	var panicked bool
	v, err := func() (interface{}, error) {
		defer func() {
			if x := recover(); x != nil {
				panicked = true
				fmt.Fprintln(os.Stderr, "panic recovered:", x)
				debug.PrintStack()
			}
		}()
		return h(r)
	}()
	if panicked {
		return resp.Error("ERR internal server error (panic)")
	}
	if err != nil {
		if err == ErrWrongArgs {
			return errWrongArgs(r.Name)
		}
		text := err.Error()
		if strings.ContainsAny(text, "\r\n") {
			text = strings.Replace(text, "\r", " ", -1)
			text = strings.Replace(text, "\n", " ", -1)
		}
		return resp.Error("ERR " + text)
	}
	return v
}

// tcpKeepAliveListener sets TCP keep-alive timeouts on accepted
// connections. It's used by ListenAndServe and ListenAndServeTLS so
// dead TCP connections (e.g. closing laptop mid-download) eventually
// go away.
type tcpKeepAliveListener struct {
	*net.TCPListener
}

func (ln tcpKeepAliveListener) Accept() (c net.Conn, err error) {
	tc, err := ln.AcceptTCP()
	if err != nil {
		return
	}
	tc.SetKeepAlive(true)
	tc.SetKeepAlivePeriod(3 * time.Minute)
	return tc, nil
}
<|endoftext|>


<commit_before>package anidb

import (
	"bytes"
	"github.com/Kovensky/go-anidb/titles"
	"io"
	"net/http"
	"os"
	"time"
)

var titlesDB = &titles.TitlesDatabase{}

// Loads the database from anime-titles.dat.gz in the cache dir.
func RefreshTitles() error {
	if lock, err := Cache.Lock("anime-titles.dat.gz"); err != nil {
		return err
	} else {
		defer lock.Unlock()
	}

	fh, err := Cache.Open("anime-titles.dat.gz")
	if err != nil {
		return err
	}
	defer fh.Close()

	titlesDB.LoadDB(fh)
	return nil
}

// Returns true if the titles database is up-to-date (newer than 24 hours).
func TitlesUpToDate() (ok bool) {
	return time.Now().Sub(titlesDB.UpdateTime) < 24*time.Hour
}

// Downloads a new anime-titles database if the database is outdated.
//
// Saves the database as anime-titles.dat.gz in the cache dir.
func (adb *AniDB) UpdateTitles() error {
	// needs the AniDB for the Logger

	// too new, no need to update
	if TitlesUpToDate() {
		return nil
	}

	switch lock, err := Cache.Lock("anime-titles.dat.gz"); {
	case os.IsNotExist(err):
		// we're creating it now
	case err == nil:
		defer lock.Unlock()
	default:
		return err
	}

	c := &http.Client{Transport: &http.Transport{DisableCompression: true}}

	adb.Logger.Printf("HTTP>>> %s", titles.DataDumpURL)

	resp, err := c.Get(titles.DataDumpURL)
	if err != nil {
		adb.Logger.Printf("HTTP<<< %s", resp.Status)
		return err
	}
	defer resp.Body.Close()

	buf := bytes.Buffer{}
	adb.Logger.Printf("HTTP--- %s", resp.Status)

	_, err = io.Copy(&buf, resp.Body)
	if err != nil {
		adb.Logger.Printf("HTTP--- %v", err)
		return err
	}

	fh, err := Cache.Create("anime-titles.dat.gz")
	if err != nil {
		return err
	}

	_, err = io.Copy(fh, &buf)
	if err != nil {
		return err
	}

	defer func() {
		adb.Logger.Printf("HTTP<<< Titles version %s", titlesDB.UpdateTime)
	}()
	return RefreshTitles()
}
<commit_msg>anidb: Add AnimeCount() call<commit_after>package anidb

import (
	"bytes"
	"github.com/Kovensky/go-anidb/titles"
	"io"
	"net/http"
	"os"
	"time"
)

var titlesDB = &titles.TitlesDatabase{}

// Loads the database from anime-titles.dat.gz in the cache dir.
func RefreshTitles() error {
	if lock, err := Cache.Lock("anime-titles.dat.gz"); err != nil {
		return err
	} else {
		defer lock.Unlock()
	}

	fh, err := Cache.Open("anime-titles.dat.gz")
	if err != nil {
		return err
	}
	defer fh.Close()

	titlesDB.LoadDB(fh)
	return nil
}

// Returns true if the titles database is up-to-date (newer than 24 hours).
func TitlesUpToDate() (ok bool) {
	return time.Now().Sub(titlesDB.UpdateTime) < 24*time.Hour
}

// Returns the number of anime in the titles database
func AnimeCount() int {
	return len(titlesDB.AnimeMap)
}

// Downloads a new anime-titles database if the database is outdated.
//
// Saves the database as anime-titles.dat.gz in the cache dir.
func (adb *AniDB) UpdateTitles() error {
	// needs the AniDB for the Logger

	// too new, no need to update
	if TitlesUpToDate() {
		return nil
	}

	switch lock, err := Cache.Lock("anime-titles.dat.gz"); {
	case os.IsNotExist(err):
		// we're creating it now
	case err == nil:
		defer lock.Unlock()
	default:
		return err
	}

	c := &http.Client{Transport: &http.Transport{DisableCompression: true}}

	adb.Logger.Printf("HTTP>>> %s", titles.DataDumpURL)

	resp, err := c.Get(titles.DataDumpURL)
	if err != nil {
		adb.Logger.Printf("HTTP<<< %s", resp.Status)
		return err
	}
	defer resp.Body.Close()

	buf := bytes.Buffer{}
	adb.Logger.Printf("HTTP--- %s", resp.Status)

	_, err = io.Copy(&buf, resp.Body)
	if err != nil {
		adb.Logger.Printf("HTTP--- %v", err)
		return err
	}

	fh, err := Cache.Create("anime-titles.dat.gz")
	if err != nil {
		return err
	}

	_, err = io.Copy(fh, &buf)
	if err != nil {
		return err
	}

	defer func() {
		adb.Logger.Printf("HTTP<<< Titles version %s", titlesDB.UpdateTime)
	}()
	return RefreshTitles()
}
<|endoftext|>


<commit_before>package tm

import (
	"fmt"
	"testing"
)

func TestMatchTokens1(t *testing.T) {
	delT, insT := LineToTokens("{abc{def}ghi}"), LineToTokens("{def}")
	matA, matB := MatchTokens(delT, insT)
	fmt.Println(delT, matA)
	fmt.Println(insT, matB)

	if matA[2] != 0 {
		t.Errorf("matA[2] should be 0")
	}
	if matB[0] != 2 {
		t.Errorf("matB[0] should be 2")
	}

	if matA[4] != 2 {
		t.Errorf("matA[4] should be 2")
	}
	if matB[2] != 4 {
		t.Errorf("matB[2] should be 4")
	}
}

func TestMatchTokens2(t *testing.T) {
	delT, insT := LineToTokens("(a.b())"), LineToTokens("a.b()")
	matA, matB := MatchTokens(delT, insT)
	fmt.Println(delT, matA)
	fmt.Println(insT, matB)

	if matA[4] != 3 {
		t.Errorf("matA[4] should be 3")
	}
	if matB[3] != 4 {
		t.Errorf("matB[3] should be 4")
	}

	if matA[5] != 4 {
		t.Errorf("matA[5] should be 4")
	}
	if matB[4] != 5 {
		t.Errorf("matB[4] should be 5")
	}

	delT, insT = LineToTokens("(a.b(), c)"), LineToTokens("(a.b, c)")
	matA, matB = MatchTokens(delT, insT)
	fmt.Println(delT, matA)
	fmt.Println(insT, matB)

	if matA[0] != 0 {
		t.Errorf("matA[0] should be 0")
	}

	if matA[9] != 7 {
		t.Errorf("matA[9] should be 7 but got %d, %v, %v", matA[9], delT, insT)
	}

	delT, insT = LineToTokens("[a.b[]]"), LineToTokens("[a.b]")
	matA, matB = MatchTokens(delT, insT)
	fmt.Println(delT, matA)
	fmt.Println(insT, matB)

	if matA[0] != 0 {
		t.Errorf("matA[0] should be 0")
	}
	if matB[0] != 0 {
		t.Errorf("matB[0] should be 0")
	}

	if matA[6] != 4 {
		t.Errorf("matA[5] should be 4 but got %d", matA[6])
	}
	if matB[4] != 6 {
		t.Errorf("matB[4] should be 6 but got %d", matB[4])
	}

	delT, insT = LineToTokens("(), (abc)"), LineToTokens("(abc)")
	matA, matB = MatchTokens(delT, insT)
	fmt.Println(delT, matA)
	fmt.Println(insT, matB)

	if matA[4] != 0 {
		t.Errorf("matA[4] should be 0 but got %d", matA[4])
	}
	if matB[0] != 4 {
		t.Errorf("matB[0] should be 0 but got %d", matB[0])
	}

	if matA[6] != 2 {
		t.Errorf("matA[6] should be 2 but got %d", matA[6])
	}
	if matB[2] != 6 {
		t.Errorf("matB[2] should be 6 but got %d", matB[2])
	}

	delT, insT = LineToTokens(`"", "abc"`), LineToTokens(`"abc"`)
	matA, matB = MatchTokens(delT, insT)
	fmt.Println(delT, matA)
	fmt.Println(insT, matB)

	if matA[4] != 0 {
		t.Errorf("matA[4] should be 0 but got %d", matA[4])
	}
	if matB[0] != 4 {
		t.Errorf("matB[0] should be 0 but got %d", matB[0])
	}

	if matA[6] != 2 {
		t.Errorf("matA[6] should be 2 but got %d", matA[6])
	}
	if matB[2] != 6 {
		t.Errorf("matB[2] should be 6 but got %d", matB[2])
	}

	delT, insT = LineToTokens(`cmd:=exec.Command("go", "abc")`), LineToTokens(`cmd:=villa.Path("go").Command("abc")`)
	matA, matB = MatchTokens(delT, insT)
	fmt.Println(delT, matA)
	fmt.Println(insT, matB)

	if matA[7] != -1 {
		t.Errorf("matA[7] should be -1 but got %d", matA[7])
	}

	if matA[12] != 14 {
		t.Errorf("matA[12] should be 14 but got %d", matA[12])
	}
}

func TestCalcDiffOfSourceLine(t *testing.T) {
	diff := CalcDiffOfSourceLine("return &monitor{", "m := &monitor{", 300)
	fmt.Println("diff", diff)

	diff = CalcDiffOfSourceLine("if (delO[iA] == 0) == (insO[iB] == 0) {", "if delL[iA] {", 1000)
	fmt.Println("diff", diff)
	diff = CalcDiffOfSourceLine("if (delO[iA] == 0) == (insO[iB] == 0) {", "c += diffAt(delT, iA - 1, insT, iB - 1)", 1000)
	fmt.Println("diff", diff)
}
<commit_msg>DEL some useless checking in testcase<commit_after>package tm

import (
	"fmt"
	"testing"
)

func TestMatchTokens1(t *testing.T) {
	delT, insT := LineToTokens("{abc{def}ghi}"), LineToTokens("{def}")
	matA, matB := MatchTokens(delT, insT)
	fmt.Println(delT, matA)
	fmt.Println(insT, matB)

	if matA[2] != 0 {
		t.Errorf("matA[2] should be 0")
	}
	if matB[0] != 2 {
		t.Errorf("matB[0] should be 2")
	}

	if matA[4] != 2 {
		t.Errorf("matA[4] should be 2")
	}
	if matB[2] != 4 {
		t.Errorf("matB[2] should be 4")
	}
}

func TestMatchTokens2(t *testing.T) {
	delT, insT := LineToTokens("(a.b())"), LineToTokens("a.b()")
	matA, matB := MatchTokens(delT, insT)
	fmt.Println(delT, matA)
	fmt.Println(insT, matB)

	if matA[4] != 3 {
		t.Errorf("matA[4] should be 3")
	}

	if matA[5] != 4 {
		t.Errorf("matA[5] should be 4")
	}

	delT, insT = LineToTokens("(a.b(), c)"), LineToTokens("(a.b, c)")
	matA, matB = MatchTokens(delT, insT)
	fmt.Println(delT, matA)
	fmt.Println(insT, matB)

	if matA[0] != 0 {
		t.Errorf("matA[0] should be 0")
	}

	if matA[9] != 7 {
		t.Errorf("matA[9] should be 7 but got %d, %v, %v", matA[9], delT, insT)
	}

	delT, insT = LineToTokens("[a.b[]]"), LineToTokens("[a.b]")
	matA, matB = MatchTokens(delT, insT)
	fmt.Println(delT, matA)
	fmt.Println(insT, matB)

	if matA[0] != 0 {
		t.Errorf("matA[0] should be 0")
	}

	if matA[6] != 4 {
		t.Errorf("matA[5] should be 4 but got %d", matA[6])
	}

	delT, insT = LineToTokens("(), (abc)"), LineToTokens("(abc)")
	matA, matB = MatchTokens(delT, insT)
	fmt.Println(delT, matA)
	fmt.Println(insT, matB)

	if matA[4] != 0 {
		t.Errorf("matA[4] should be 0 but got %d", matA[4])
	}

	if matA[6] != 2 {
		t.Errorf("matA[6] should be 2 but got %d", matA[6])
	}

	delT, insT = LineToTokens(`"", "abc"`), LineToTokens(`"abc"`)
	matA, matB = MatchTokens(delT, insT)
	fmt.Println(delT, matA)
	fmt.Println(insT, matB)

	if matA[4] != 0 {
		t.Errorf("matA[4] should be 0 but got %d", matA[4])
	}

	if matA[6] != 2 {
		t.Errorf("matA[6] should be 2 but got %d", matA[6])
	}

	delT, insT = LineToTokens(`cmd:=exec.Command("go", "abc")`), LineToTokens(`cmd:=villa.Path("go").Command("abc")`)
	matA, matB = MatchTokens(delT, insT)
	fmt.Println(delT, matA)
	fmt.Println(insT, matB)

	if matA[7] != -1 {
		t.Errorf("matA[7] should be -1 but got %d", matA[7])
	}

	if matA[12] != 14 {
		t.Errorf("matA[12] should be 14 but got %d", matA[12])
	}
}

func TestCalcDiffOfSourceLine(t *testing.T) {
	diff := CalcDiffOfSourceLine("return &monitor{", "m := &monitor{", 300)
	fmt.Println("diff", diff)

	diff = CalcDiffOfSourceLine("if (delO[iA] == 0) == (insO[iB] == 0) {", "if delL[iA] {", 1000)
	fmt.Println("diff", diff)
	diff = CalcDiffOfSourceLine("if (delO[iA] == 0) == (insO[iB] == 0) {", "c += diffAt(delT, iA - 1, insT, iB - 1)", 1000)
	fmt.Println("diff", diff)
}
<|endoftext|>


<commit_before>// Copyright 2014 <chaishushan{AT}gmail.com>. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package webp

import (
	"image"
	"image/color"
	"reflect"
)

var (
	_ color.Color = (*RGBColor)(nil)
	_ image.Image = (*RGBImage)(nil)
)

var RGBModel color.Model = color.ModelFunc(rgbModel)

type RGBColor struct {
	R, G, B uint8
}

func (c RGBColor) RGBA() (r, g, b, a uint32) {
	r = uint32(c.R)
	r |= r << 8
	g = uint32(c.G)
	g |= g << 8
	b = uint32(c.B)
	b |= b << 8
	a = 0xffff
	return
}

func rgbModel(c color.Color) color.Color {
	if _, ok := c.(RGBColor); ok {
		return c
	}
	r, g, b, _ := c.RGBA()
	return RGBColor{R: uint8(r >> 8), G: uint8(g >> 8), B: uint8(b >> 8)}
}

type RGBImage struct {
	Pix    []uint8
	Stride int
	Rect   image.Rectangle
}

func (p *RGBImage) ColorModel() color.Model { return color.RGBAModel }

func (p *RGBImage) Bounds() image.Rectangle { return p.Rect }

func (p *RGBImage) At(x, y int) color.Color {
	return p.RGBAt(x, y)
}

func (p *RGBImage) RGBAt(x, y int) RGBColor {
	if !(image.Point{x, y}.In(p.Rect)) {
		return RGBColor{}
	}
	i := p.PixOffset(x, y)
	return RGBColor{
		R: p.Pix[i+0],
		G: p.Pix[i+1],
		B: p.Pix[i+2],
	}
}

// PixOffset returns the index of the first element of Pix that corresponds to
// the pixel at (x, y).
func (p *RGBImage) PixOffset(x, y int) int {
	return (y-p.Rect.Min.Y)*p.Stride + (x-p.Rect.Min.X)*3
}

func (p *RGBImage) Set(x, y int, c color.Color) {
	if !(image.Point{x, y}.In(p.Rect)) {
		return
	}
	i := p.PixOffset(x, y)
	c1 := RGBModel.Convert(c).(RGBColor)
	p.Pix[i+0] = c1.R
	p.Pix[i+1] = c1.G
	p.Pix[i+2] = c1.B
	return
}

func (p *RGBImage) SetRGB(x, y int, c RGBColor) {
	if !(image.Point{x, y}.In(p.Rect)) {
		return
	}
	i := p.PixOffset(x, y)
	p.Pix[i+0] = c.R
	p.Pix[i+1] = c.G
	p.Pix[i+2] = c.B
	return
}

// SubImage returns an image representing the portion of the image p visible
// through r. The returned value shares pixels with the original image.
func (p *RGBImage) SubImage(r image.Rectangle) image.Image {
	r = r.Intersect(p.Rect)
	// If r1 and r2 are Rectangles, r1.Intersect(r2) is not guaranteed to be inside
	// either r1 or r2 if the intersection is empty. Without explicitly checking for
	// this, the Pix[i:] expression below can panic.
	if r.Empty() {
		return &RGBImage{}
	}
	i := p.PixOffset(r.Min.X, r.Min.Y)
	return &RGBImage{
		Pix:    p.Pix[i:],
		Stride: p.Stride,
		Rect:   r,
	}
}

// Opaque scans the entire image and reports whether it is fully opaque.
func (p *RGBImage) Opaque() bool {
	return true
}

// MemPImage return a MemP typed image.
//
// MemP Image Spec (Native Endian), see https://github.com/chai2010/image.
func (p *RGBImage) MemPImage() interface{} {
	type mempImage struct {
		MemPMagic string // MemP
		Rect      image.Rectangle
		Channels  int
		DataType  reflect.Kind
		Pix       []byte

		// Stride is the Pix stride (in bytes, must align with SizeofKind(p.DataType))
		// between vertically adjacent pixels.
		Stride int
	}

	return &mempImage{
		MemPMagic: "MemP",
		Rect:      p.Rect,
		Channels:  3,
		DataType:  reflect.Uint8,
		Pix:       p.Pix,
		Stride:    p.Stride,
	}
}

// MemP Image Spec (Native Endian), see https://github.com/chai2010/image.

// NewRGBImage returns a new RGBImage with the given bounds.
func NewRGBImage(r image.Rectangle) *RGBImage {
	w, h := r.Dx(), r.Dy()
	pix := make([]uint8, 3*w*h)
	return &RGBImage{
		Pix:    pix,
		Stride: 3 * w,
		Rect:   r,
	}
}

func NewRGBImageFrom(m image.Image) *RGBImage {
	if m, ok := m.(*RGBImage); ok {
		return m
	}

	// convert to RGBImage
	b := m.Bounds()
	rgb := NewRGBImage(b)
	for y := b.Min.Y; y < b.Max.Y; y++ {
		for x := b.Min.X; x < b.Max.X; x++ {
			pr, pg, pb, _ := m.At(x, y).RGBA()
			rgb.SetRGB(x, y, RGBColor{
				R: uint8(pr >> 8),
				G: uint8(pg >> 8),
				B: uint8(pb >> 8),
			})
		}
	}
	return rgb
}
<commit_msg>remove RGBImage.MemPImage method<commit_after>// Copyright 2014 <chaishushan{AT}gmail.com>. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package webp

import (
	"image"
	"image/color"
)

var (
	_ color.Color = (*RGBColor)(nil)
	_ image.Image = (*RGBImage)(nil)
)

var RGBModel color.Model = color.ModelFunc(rgbModel)

type RGBColor struct {
	R, G, B uint8
}

func (c RGBColor) RGBA() (r, g, b, a uint32) {
	r = uint32(c.R)
	r |= r << 8
	g = uint32(c.G)
	g |= g << 8
	b = uint32(c.B)
	b |= b << 8
	a = 0xffff
	return
}

func rgbModel(c color.Color) color.Color {
	if _, ok := c.(RGBColor); ok {
		return c
	}
	r, g, b, _ := c.RGBA()
	return RGBColor{R: uint8(r >> 8), G: uint8(g >> 8), B: uint8(b >> 8)}
}

type RGBImage struct {
	Pix    []uint8
	Stride int
	Rect   image.Rectangle
}

func (p *RGBImage) ColorModel() color.Model { return color.RGBAModel }

func (p *RGBImage) Bounds() image.Rectangle { return p.Rect }

func (p *RGBImage) At(x, y int) color.Color {
	return p.RGBAt(x, y)
}

func (p *RGBImage) RGBAt(x, y int) RGBColor {
	if !(image.Point{x, y}.In(p.Rect)) {
		return RGBColor{}
	}
	i := p.PixOffset(x, y)
	return RGBColor{
		R: p.Pix[i+0],
		G: p.Pix[i+1],
		B: p.Pix[i+2],
	}
}

// PixOffset returns the index of the first element of Pix that corresponds to
// the pixel at (x, y).
func (p *RGBImage) PixOffset(x, y int) int {
	return (y-p.Rect.Min.Y)*p.Stride + (x-p.Rect.Min.X)*3
}

func (p *RGBImage) Set(x, y int, c color.Color) {
	if !(image.Point{x, y}.In(p.Rect)) {
		return
	}
	i := p.PixOffset(x, y)
	c1 := RGBModel.Convert(c).(RGBColor)
	p.Pix[i+0] = c1.R
	p.Pix[i+1] = c1.G
	p.Pix[i+2] = c1.B
	return
}

func (p *RGBImage) SetRGB(x, y int, c RGBColor) {
	if !(image.Point{x, y}.In(p.Rect)) {
		return
	}
	i := p.PixOffset(x, y)
	p.Pix[i+0] = c.R
	p.Pix[i+1] = c.G
	p.Pix[i+2] = c.B
	return
}

// SubImage returns an image representing the portion of the image p visible
// through r. The returned value shares pixels with the original image.
func (p *RGBImage) SubImage(r image.Rectangle) image.Image {
	r = r.Intersect(p.Rect)
	// If r1 and r2 are Rectangles, r1.Intersect(r2) is not guaranteed to be inside
	// either r1 or r2 if the intersection is empty. Without explicitly checking for
	// this, the Pix[i:] expression below can panic.
	if r.Empty() {
		return &RGBImage{}
	}
	i := p.PixOffset(r.Min.X, r.Min.Y)
	return &RGBImage{
		Pix:    p.Pix[i:],
		Stride: p.Stride,
		Rect:   r,
	}
}

// Opaque scans the entire image and reports whether it is fully opaque.
func (p *RGBImage) Opaque() bool {
	return true
}

// NewRGBImage returns a new RGBImage with the given bounds.
func NewRGBImage(r image.Rectangle) *RGBImage {
	w, h := r.Dx(), r.Dy()
	pix := make([]uint8, 3*w*h)
	return &RGBImage{
		Pix:    pix,
		Stride: 3 * w,
		Rect:   r,
	}
}

func NewRGBImageFrom(m image.Image) *RGBImage {
	if m, ok := m.(*RGBImage); ok {
		return m
	}

	// convert to RGBImage
	b := m.Bounds()
	rgb := NewRGBImage(b)
	for y := b.Min.Y; y < b.Max.Y; y++ {
		for x := b.Min.X; x < b.Max.X; x++ {
			pr, pg, pb, _ := m.At(x, y).RGBA()
			rgb.SetRGB(x, y, RGBColor{
				R: uint8(pr >> 8),
				G: uint8(pg >> 8),
				B: uint8(pb >> 8),
			})
		}
	}
	return rgb
}
<|endoftext|>


<commit_before>package rin

import (
	"log"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/crowdmob/goamz/aws"
	"github.com/crowdmob/goamz/sqs"
)

var SQS *sqs.SQS
var config *Config
var Debug bool
var Runnable bool

var TrapSignals = []os.Signal{
	syscall.SIGHUP,
	syscall.SIGINT,
	syscall.SIGTERM,
	syscall.SIGQUIT,
}

type NoMessageError struct {
	s string
}

func (e *NoMessageError) Error() string {
	return e.s
}

func Run(configFile string, batchMode bool) error {
	Runnable = true
	var err error
	log.Println("[info] Loading config:", configFile)
	config, err = LoadConfig(configFile)
	if err != nil {
		return err
	}
	for _, target := range config.Targets {
		log.Println("[info] Define target", target)
	}

	auth := aws.Auth{
		AccessKey: config.Credentials.AWS_ACCESS_KEY_ID,
		SecretKey: config.Credentials.AWS_SECRET_ACCESS_KEY,
	}
	region := aws.GetRegion(config.Credentials.AWS_REGION)
	SQS = sqs.New(auth, region)

	shutdownCh := make(chan interface{})
	exitCh := make(chan int)
	signalCh := make(chan os.Signal, 1)
	signal.Notify(signalCh, TrapSignals...)

	// run worker
	if batchMode {
		go func() {
			err := sqsBatch(shutdownCh)
			if err != nil {
				log.Println("[error]", err)
				exitCh <- 1
			}
			exitCh <- 0
		}()
	} else {
		go sqsWorker(shutdownCh)
	}

	// wait for signal
	var exitCode int
	select {
	case s := <-signalCh:
		switch sig := s.(type) {
		case syscall.Signal:
			log.Printf("[info] Got signal: %s(%d)", sig, sig)
		default:
		}
		log.Println("[info] Shutting down worker...")
		close(shutdownCh) // notify shutdown to worker
	case exitCode = <-exitCh:
	}

	log.Println("[info] Shutdown.")
	os.Exit(exitCode)
	return nil
}

func waitForRetry() {
	log.Println("[warn] Retry after 10 sec.")
	time.Sleep(10 * time.Second)
}

func runnable(ch chan interface{}) bool {
	if !Runnable {
		return false
	}
	select {
	case <-ch:
		// ch closed == shutdown
		Runnable = false
		return false
	default:
	}
	return true
}

func sqsBatch(ch chan interface{}) error {
	log.Printf("[info] Starting up SQS Batch")
	defer log.Println("[info] Shutdown SQS Batch")

	log.Println("[info] Connect to SQS:", config.QueueName)
	queue, err := SQS.GetQueue(config.QueueName)
	if err != nil {
		return err
	}
	for runnable(ch) {
		err := handleMessage(queue)
		if err != nil {
			if _, ok := err.(*NoMessageError); ok {
				break
			} else {
				return err
			}
		}
	}
	return nil
}

func sqsWorker(ch chan interface{}) {
	log.Printf("[info] Starting up SQS Worker")
	defer log.Println("[info] Shutdown SQS Worker")

	for runnable(ch) {
		log.Println("[info] Connect to SQS:", config.QueueName)
		queue, err := SQS.GetQueue(config.QueueName)
		if err != nil {
			log.Println("[error] Can't get queue:", err)
			waitForRetry()
			continue
		}
		quit, err := handleQueue(queue, ch)
		if err != nil {
			log.Println("[error] Processing failed:", err)
			waitForRetry()
			continue
		}
		if quit {
			break
		}
	}
}

func handleQueue(queue *sqs.Queue, ch chan interface{}) (bool, error) {
	for runnable(ch) {
		err := handleMessage(queue)
		if err != nil {
			if _, ok := err.(*NoMessageError); ok {
				continue
			} else {
				return false, err
			}
		}
	}
	return true, nil
}

func handleMessage(queue *sqs.Queue) error {
	var completed = false
	res, err := queue.ReceiveMessage(1)
	if err != nil {
		return err
	}
	if len(res.Messages) == 0 {
		return &NoMessageError{"No messages"}
	}
	msg := res.Messages[0]
	log.Printf("[info] [%s] Starting process message.", msg.MessageId)
	if Debug {
		log.Printf("[degug] [%s] handle: %s", msg.MessageId, msg.ReceiptHandle)
		log.Printf("[debug] [%s] body: %s", msg.MessageId, msg.Body)
	}
	defer func() {
		if !completed {
			log.Printf("[info] [%s] Aborted message.", msg.MessageId)
		}
	}()

	event, err := ParseEvent([]byte(msg.Body))
	if err != nil {
		log.Printf("[error] [%s] Can't parse event from Body.", msg.MessageId, err)
		return err
	}
	log.Printf("[info] [%s] Importing event: %s", msg.MessageId, event)
	n, err := Import(event)
	if err != nil {
		log.Printf("[error] [%s] Import failed. %s", msg.MessageId, err)
		return err
	}
	if n == 0 {
		log.Printf("[warn] [%s] All events were not matched for any targets. Ignored.", msg.MessageId)
	} else {
		log.Printf("[info] [%s] %d import action completed.", msg.MessageId, n)
	}
	_, err = queue.DeleteMessage(&msg)
	if err != nil {
		log.Printf("[error] [%s] Can't delete message. %s", msg.MessageId, err)
	}
	completed = true
	log.Printf("[info] [%s] Completed message.", msg.MessageId)
	return nil
}
<commit_msg>retry delete message.<commit_after>package rin

import (
	"log"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/crowdmob/goamz/aws"
	"github.com/crowdmob/goamz/sqs"
)

var SQS *sqs.SQS
var config *Config
var Debug bool
var Runnable bool
var MaxDeleteRetry = 8

var TrapSignals = []os.Signal{
	syscall.SIGHUP,
	syscall.SIGINT,
	syscall.SIGTERM,
	syscall.SIGQUIT,
}

type NoMessageError struct {
	s string
}

func (e *NoMessageError) Error() string {
	return e.s
}

func Run(configFile string, batchMode bool) error {
	Runnable = true
	var err error
	log.Println("[info] Loading config:", configFile)
	config, err = LoadConfig(configFile)
	if err != nil {
		return err
	}
	for _, target := range config.Targets {
		log.Println("[info] Define target", target)
	}

	auth := aws.Auth{
		AccessKey: config.Credentials.AWS_ACCESS_KEY_ID,
		SecretKey: config.Credentials.AWS_SECRET_ACCESS_KEY,
	}
	region := aws.GetRegion(config.Credentials.AWS_REGION)
	SQS = sqs.New(auth, region)

	shutdownCh := make(chan interface{})
	exitCh := make(chan int)
	signalCh := make(chan os.Signal, 1)
	signal.Notify(signalCh, TrapSignals...)

	// run worker
	if batchMode {
		go func() {
			err := sqsBatch(shutdownCh)
			if err != nil {
				log.Println("[error]", err)
				exitCh <- 1
			}
			exitCh <- 0
		}()
	} else {
		go sqsWorker(shutdownCh)
	}

	// wait for signal
	var exitCode int
	select {
	case s := <-signalCh:
		switch sig := s.(type) {
		case syscall.Signal:
			log.Printf("[info] Got signal: %s(%d)", sig, sig)
		default:
		}
		log.Println("[info] Shutting down worker...")
		close(shutdownCh) // notify shutdown to worker
	case exitCode = <-exitCh:
	}

	log.Println("[info] Shutdown.")
	os.Exit(exitCode)
	return nil
}

func waitForRetry() {
	log.Println("[warn] Retry after 10 sec.")
	time.Sleep(10 * time.Second)
}

func runnable(ch chan interface{}) bool {
	if !Runnable {
		return false
	}
	select {
	case <-ch:
		// ch closed == shutdown
		Runnable = false
		return false
	default:
	}
	return true
}

func sqsBatch(ch chan interface{}) error {
	log.Printf("[info] Starting up SQS Batch")
	defer log.Println("[info] Shutdown SQS Batch")

	log.Println("[info] Connect to SQS:", config.QueueName)
	queue, err := SQS.GetQueue(config.QueueName)
	if err != nil {
		return err
	}
	for runnable(ch) {
		err := handleMessage(queue)
		if err != nil {
			if _, ok := err.(*NoMessageError); ok {
				break
			} else {
				return err
			}
		}
	}
	return nil
}

func sqsWorker(ch chan interface{}) {
	log.Printf("[info] Starting up SQS Worker")
	defer log.Println("[info] Shutdown SQS Worker")

	for runnable(ch) {
		log.Println("[info] Connect to SQS:", config.QueueName)
		queue, err := SQS.GetQueue(config.QueueName)
		if err != nil {
			log.Println("[error] Can't get queue:", err)
			waitForRetry()
			continue
		}
		quit, err := handleQueue(queue, ch)
		if err != nil {
			log.Println("[error] Processing failed:", err)
			waitForRetry()
			continue
		}
		if quit {
			break
		}
	}
}

func handleQueue(queue *sqs.Queue, ch chan interface{}) (bool, error) {
	for runnable(ch) {
		err := handleMessage(queue)
		if err != nil {
			if _, ok := err.(*NoMessageError); ok {
				continue
			} else {
				return false, err
			}
		}
	}
	return true, nil
}

func handleMessage(queue *sqs.Queue) error {
	var completed = false
	res, err := queue.ReceiveMessage(1)
	if err != nil {
		return err
	}
	if len(res.Messages) == 0 {
		return &NoMessageError{"No messages"}
	}
	msg := res.Messages[0]
	log.Printf("[info] [%s] Starting process message.", msg.MessageId)
	if Debug {
		log.Printf("[degug] [%s] handle: %s", msg.MessageId, msg.ReceiptHandle)
		log.Printf("[debug] [%s] body: %s", msg.MessageId, msg.Body)
	}
	defer func() {
		if !completed {
			log.Printf("[info] [%s] Aborted message.", msg.MessageId)
		}
	}()

	event, err := ParseEvent([]byte(msg.Body))
	if err != nil {
		log.Printf("[error] [%s] Can't parse event from Body.", msg.MessageId, err)
		return err
	}
	log.Printf("[info] [%s] Importing event: %s", msg.MessageId, event)
	n, err := Import(event)
	if err != nil {
		log.Printf("[error] [%s] Import failed. %s", msg.MessageId, err)
		return err
	}
	if n == 0 {
		log.Printf("[warn] [%s] All events were not matched for any targets. Ignored.", msg.MessageId)
	} else {
		log.Printf("[info] [%s] %d import action completed.", msg.MessageId, n)
	}
	_, err = queue.DeleteMessage(&msg)
	if err != nil {
		log.Printf("[error] [%s] Can't delete message. %s", msg.MessageId, err)
		// retry
		for i := 1; i <= MaxDeleteRetry; i++ {
			log.Printf("[info] [%s] Retry to delete after %d sec.", msg.MessageId, i*i)
			time.Sleep(time.Duration(i*i) * time.Second)
			_, err := queue.DeleteMessage(&msg)
			if err == nil {
				log.Printf("[info] [%s] Message was deleted successfuly.", msg.MessageId)
				break
			}
			log.Printf("[error] [%s] Can't delete message. %s", msg.MessageId, err)
		}
		log.Printf("[error] [%s] Max retry count reached. Giving up.", msg.MessageId)
	}

	completed = true
	log.Printf("[info] [%s] Completed message.", msg.MessageId)
	return nil
}
<|endoftext|>


<commit_before>package raft

import (
  "net/http"
  "fmt"
)

// Used to store any field to interact with the RPC
// functions
type RPC struct {
  Server *http.Server
}

// Start the RPC server
func (rpc *RPC) Start(ok chan struct{}) error {
  // TODO: test the server
  go func() {
    ok <- struct{}{}
  }()

  err := rpc.Server.ListenAndServe()

  return err
}

// Create RPC connection
func NewRPC(config *StateMachineConfiguration) (*RPC, error) {
  mux := http.NewServeMux()

  server := &http.Server{
    Addr: ":8081",
    Handler: mux,
  }

  mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
    fmt.Fprintf(w, "Hello World!")
  })

  rpc := &RPC{
    Server: server,
  }

  return rpc, nil
}
<commit_msg>Added RPC host and port configuration<commit_after>package raft

import (
  "net/http"
  "fmt"
)

// Used to store any field to interact with the RPC
// functions
type RPC struct {
  Server *http.Server
}

// Start the RPC server
func (rpc *RPC) Start(ok chan struct{}) error {
  // TODO: test the server
  go func() {
    ok <- struct{}{}
  }()

  err := rpc.Server.ListenAndServe()

  return err
}

// Create RPC connection
func NewRPC(config *StateMachineConfiguration) (*RPC, error) {
  mux := http.NewServeMux()

  addr := fmt.Sprintf("%s:%d", config.RPCHost, config.RPCPort)

  server := &http.Server{
    Addr: addr,
    Handler: mux,
  }

  mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
    fmt.Fprintf(w, "Hello World!")
  })

  rpc := &RPC{
    Server: server,
  }

  return rpc, nil
}
<|endoftext|>


<commit_before>package bnblog

import (
	"appengine"
	"appengine/datastore"
	"encoding/base64"
	"fmt"
	"github.com/gorilla/feeds"
	"net/http"
	"strings"
	"time"
)

func GetRSS(rw http.ResponseWriter, req *http.Request) {
	now := time.Now()
	feed := &feeds.Feed{
		Title:       "benjojo blog",
		Link:        &feeds.Link{Href: "https://blog.benjojo.co.uk"},
		Description: "Programming, Networking and some things I found hard to fix at some point",
		Author:      &feeds.Author{"Ben Cartwright-Cox", "ben@benjojo.co.uk"},
		Created:     now,
	}

	c := appengine.NewContext(req)
	q := datastore.NewQuery("Post").Order("-Date").Limit(100)
	posts := make([]Post, 0, 100)

	if _, err := q.GetAll(c, &posts); err != nil {
		http.Error(rw, err.Error(), http.StatusInternalServerError)
		return
	}

	feed.Items = []*feeds.Item{}

	for _, v := range posts {
		if !strings.HasPrefix(v.Slug, "DRAFT-") {
			// newpost := PostFormatted{
			// 	Author:  v.Author,
			// 	Content: v.Content,
			// 	Date:    v.Date.Format("2006-01-02 15:04:05"),
			// 	Slug:    v.Slug,
			// 	Title:   v.Title,
			// }
			// FormattedPosts = append(FormattedPosts, newpost)
			postd, _ := base64.StdEncoding.DecodeString(v.Content)
			wot := &feeds.Item{
				Title:       v.Title,
				Link:        &feeds.Link{Href: fmt.Sprintf("https://blog.benjojo.co.uk/post/%s", v.Slug)},
				Description: string(postd[:256]),
				Author:      &feeds.Author{"Ben Cox", "ben@benjojo.co.uk"},
				Created:     v.Date,
			}
			feed.Items = append(feed.Items, wot)
		}
	}

	rss, err := feed.ToRss()
	if err != nil {
		http.Error(rw, fmt.Sprintf("argh %s", err), http.StatusInternalServerError)
		return
	}
	rw.Header().Add("Content-Type", "application/rss+xml")
	rw.Write([]byte(rss))
}
<commit_msg>Fix RSS feeds from not having GUID + Bad email address<commit_after>package bnblog

import (
	"crypto/md5"
	"encoding/base64"
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/gorilla/feeds"

	"appengine"
	"appengine/datastore"
)

func GetRSS(rw http.ResponseWriter, req *http.Request) {
	now := time.Now()
	feed := &feeds.Feed{
		Title:       "benjojo blog",
		Link:        &feeds.Link{Href: "https://blog.benjojo.co.uk"},
		Description: "Programming, Networking and some things I found hard to fix at some point",
		Author:      &feeds.Author{"Ben Cartwright-Cox", "ben@benjojo.co.uk"},
		Created:     now,
	}

	c := appengine.NewContext(req)
	q := datastore.NewQuery("Post").Order("-Date").Limit(100)
	posts := make([]Post, 0, 100)

	if _, err := q.GetAll(c, &posts); err != nil {
		http.Error(rw, err.Error(), http.StatusInternalServerError)
		return
	}

	feed.Items = []*feeds.Item{}

	for _, v := range posts {
		if !strings.HasPrefix(v.Slug, "DRAFT-") {
			// newpost := PostFormatted{
			// 	Author:  v.Author,
			// 	Content: v.Content,
			// 	Date:    v.Date.Format("2006-01-02 15:04:05"),
			// 	Slug:    v.Slug,
			// 	Title:   v.Title,
			// }
			// FormattedPosts = append(FormattedPosts, newpost)
			postd, _ := base64.StdEncoding.DecodeString(v.Content)
			wot := &feeds.Item{
				Title:       v.Title,
				Link:        &feeds.Link{Href: fmt.Sprintf("https://blog.benjojo.co.uk/post/%s", v.Slug)},
				Description: string(postd[:256]),
				Author:      &feeds.Author{"Ben Cox <ben@benjojo.co.uk>", "ben@benjojo.co.uk"},
				Created:     v.Date,
				ID:          generateBadUUID(v.Title),
			}
			feed.Items = append(feed.Items, wot)
		}
	}

	rss, err := feed.ToRss()
	if err != nil {
		http.Error(rw, fmt.Sprintf("argh %s", err), http.StatusInternalServerError)
		return
	}
	rw.Header().Add("Content-Type", "application/rss+xml")
	rw.Write([]byte(rss))
}

// This is a hack as you might have guessed. This blogging system was
// never designed with UUIDs in mind, so I'm sort of just generating one
// out of the title, MD5 is fine since I don't think I am going attack
// myself with colliding titles.
func generateBadUUID(title string) string {
	hashbytes := md5.Sum([]byte(title))
	return fmt.Sprintf("%1x%1x%1x%1x-%1x%1x-40%1x-%1x%1x-%1x%1x%1x%1x%1x%1x",
		hashbytes[0], hashbytes[1], hashbytes[2], hashbytes[3], hashbytes[4],
		hashbytes[5], hashbytes[6], hashbytes[7], hashbytes[8], hashbytes[9],
		hashbytes[10], hashbytes[11], hashbytes[12], hashbytes[13],
		hashbytes[14])

}
<|endoftext|>


<commit_before>package slackapi

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"reflect"

	"golang.org/x/net/websocket"
)

// RTM is the real time messaging.
type RTM struct {
	conn      *websocket.Conn
	connURL   string
	Events    chan Event
	rawEvents chan json.RawMessage
}

// RTMResponse defines the JSON-encoded output for RTM connection.
type RTMResponse struct {
	Response
	// Self                    null `json:"self"`
	// Team                    null `json:"team"`
	// LatestEventTS           null `json:"latest_event_ts"`
	// Channels                null `json:"channels"`
	// Groups                  null `json:"groups"`
	// IMS                     null `json:"ims"`
	// CacheTS                 null `json:"cache_ts"`
	// ReadOnlyChannels        null `json:"read_only_channels"`
	// CanManageSharedChannels null `json:"can_manage_shared_channels"`
	// Subteams                null `json:"subteams"`
	// DND                     null `json:"dnd"`
	// Users                   null `json:"users"`
	// CacheVersion            null `json:"cache_version"`
	// CacheTSVersion          null `json:"cache_ts_version"`
	// Bots                    null `json:"bots"`
	URL string `json:"url"`
}

// ACKMessage is used for messages received in reply to other messages
type ACKMessage struct {
	Response
	ReplyTo   int    `json:"reply_to"`
	Text      string `json:"text"`
	Timestamp string `json:"ts"`
}

// EventTypes represents the data type for each websocket event.
var EventTypes = map[string]interface{}{}

// NewRTM connects to the real time messaging websocket.
func (s *SlackAPI) NewRTM() (*RTM, error) {
	var response RTMResponse
	s.GetRequest(&response, "rtm.start", "token")

	ws, err := websocket.Dial(response.URL, "", "https://api.slack.com")

	if err != nil {
		return &RTM{}, err
	}

	return &RTM{
		conn:      ws,
		connURL:   response.URL,
		Events:    make(chan Event, 50),
		rawEvents: make(chan json.RawMessage),
	}, nil
}

// ManageEvents controls the websocket events.
func (rtm *RTM) ManageEvents() {
	keepalive := make(chan bool)

	go rtm.handleIncomingEvents(keepalive)

	rtm.handleEvents(keepalive)
}

func (rtm *RTM) handleEvents(keepRunning chan bool) {
	for {
		select {
		case rawEvent := <-rtm.rawEvents:
			rtm.handleRawEvent(rawEvent)
		}
	}
}

func (rtm *RTM) handleIncomingEvents(keepalive <-chan bool) {
	for {
		select {
		case <-keepalive:
			return
		default:
			rtm.receiveIncomingEvent()
		}
	}
}

// receiveIncomingEvent inserts the websocket events into a queue.
func (rtm *RTM) receiveIncomingEvent() {
	event := json.RawMessage{}

	if err := websocket.JSON.Receive(rtm.conn, &event); err != nil {
		if err == io.EOF {
			log.Fatal("Connection dropped")
			return
		}

		rtm.Events <- Event{Type: "error", Data: &ErrorEvent{err.Error()}}
		return
	}

	if len(event) == 0 {
		err := fmt.Errorf("invalid event; %s", event)
		rtm.Events <- Event{Type: "error", Data: &ErrorEvent{err.Error()}}
		return
	}

	rtm.rawEvents <- event
}

func (rtm *RTM) handleRawEvent(rawEvent json.RawMessage) {
	event := &Event{}

	if err := json.Unmarshal(rawEvent, event); err != nil {
		rtm.Events <- Event{Type: "error", Data: &ErrorEvent{err.Error()}}
		return
	}

	rtm.handleEvent(event.Type, rawEvent)
}

func (rtm *RTM) handleEvent(_type string, event json.RawMessage) {
	v, exists := EventTypes[_type]

	if !exists {
		err := fmt.Errorf("unsupported event %q: %s", _type, string(event))
		rtm.Events <- Event{Type: "error", Data: &ErrorEvent{err.Error()}}
		return
	}

	t := reflect.TypeOf(v)
	recvEvent := reflect.New(t).Interface()

	if err := json.Unmarshal(event, recvEvent); err != nil {
		err = fmt.Errorf("unmarshall event %q: %s", _type, string(event))
		rtm.Events <- Event{Type: "error", Data: &ErrorEvent{err.Error()}}
		return
	}
	rtm.Events <- Event{_type, recvEvent}
}
<commit_msg>Add HELLO and ACK support for the RTM websocket entity<commit_after>package slackapi

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"reflect"

	"golang.org/x/net/websocket"
)

// RTM is the real time messaging.
type RTM struct {
	conn      *websocket.Conn
	connURL   string
	Events    chan Event
	rawEvents chan json.RawMessage
}

// RTMResponse defines the JSON-encoded output for RTM connection.
type RTMResponse struct {
	Response
	// Self                    null `json:"self"`
	// Team                    null `json:"team"`
	// LatestEventTS           null `json:"latest_event_ts"`
	// Channels                null `json:"channels"`
	// Groups                  null `json:"groups"`
	// IMS                     null `json:"ims"`
	// CacheTS                 null `json:"cache_ts"`
	// ReadOnlyChannels        null `json:"read_only_channels"`
	// CanManageSharedChannels null `json:"can_manage_shared_channels"`
	// Subteams                null `json:"subteams"`
	// DND                     null `json:"dnd"`
	// Users                   null `json:"users"`
	// CacheVersion            null `json:"cache_version"`
	// CacheTSVersion          null `json:"cache_ts_version"`
	// Bots                    null `json:"bots"`
	URL string `json:"url"`
}

// ACKMessage is used for messages received in reply to other messages
type ACKMessage struct {
	Response
	ReplyTo   int    `json:"reply_to"`
	Text      string `json:"text"`
	Timestamp string `json:"ts"`
}

// EventTypes represents the data type for each websocket event.
var EventTypes = map[string]interface{}{}

// NewRTM connects to the real time messaging websocket.
func (s *SlackAPI) NewRTM() (*RTM, error) {
	var response RTMResponse
	s.GetRequest(&response, "rtm.start", "token")

	ws, err := websocket.Dial(response.URL, "", "https://api.slack.com")

	if err != nil {
		return &RTM{}, err
	}

	return &RTM{
		conn:      ws,
		connURL:   response.URL,
		Events:    make(chan Event, 50),
		rawEvents: make(chan json.RawMessage),
	}, nil
}

// ManageEvents controls the websocket events.
func (rtm *RTM) ManageEvents() {
	keepalive := make(chan bool)

	go rtm.handleIncomingEvents(keepalive)

	rtm.handleEvents(keepalive)
}

func (rtm *RTM) handleEvents(keepRunning chan bool) {
	for {
		select {
		case rawEvent := <-rtm.rawEvents:
			rtm.handleRawEvent(rawEvent)
		}
	}
}

func (rtm *RTM) handleIncomingEvents(keepalive <-chan bool) {
	for {
		select {
		case <-keepalive:
			return
		default:
			rtm.receiveIncomingEvent()
		}
	}
}

// receiveIncomingEvent inserts the websocket events into a queue.
func (rtm *RTM) receiveIncomingEvent() {
	event := json.RawMessage{}

	if err := websocket.JSON.Receive(rtm.conn, &event); err != nil {
		if err == io.EOF {
			log.Fatal("Connection dropped")
			return
		}

		rtm.Events <- Event{Type: "error", Data: &ErrorEvent{err.Error()}}
		return
	}

	if len(event) == 0 {
		err := fmt.Errorf("invalid event; %s", event)
		rtm.Events <- Event{Type: "error", Data: &ErrorEvent{err.Error()}}
		return
	}

	rtm.rawEvents <- event
}

func (rtm *RTM) handleRawEvent(rawEvent json.RawMessage) {
	event := &Event{}

	if err := json.Unmarshal(rawEvent, event); err != nil {
		rtm.Events <- Event{Type: "error", Data: &ErrorEvent{err.Error()}}
		return
	}

	switch event.Type {
	case "":
		rtm.handleACK(rawEvent)
	case "hello":
		rtm.Events <- Event{Type: "hello", Data: &HelloEvent{}}
	default:
		rtm.handleEvent(event.Type, rawEvent)
	}
}

func (rtm *RTM) handleEvent(_type string, event json.RawMessage) {
	v, exists := EventTypes[_type]

	if !exists {
		err := fmt.Errorf("unsupported event %q: %s", _type, string(event))
		rtm.Events <- Event{Type: "error", Data: &ErrorEvent{err.Error()}}
		return
	}

	t := reflect.TypeOf(v)
	recvEvent := reflect.New(t).Interface()

	if err := json.Unmarshal(event, recvEvent); err != nil {
		err = fmt.Errorf("unmarshall event %q: %s", _type, string(event))
		rtm.Events <- Event{Type: "error", Data: &ErrorEvent{err.Error()}}
		return
	}
	rtm.Events <- Event{_type, recvEvent}
}

func (rtm *RTM) handleACK(event json.RawMessage) {
	ack := &ACKMessage{}

	if err := json.Unmarshal(event, ack); err != nil {
		err = fmt.Errorf("ack unmarshal; %s: %s", err.Error(), string(event))
		rtm.Events <- Event{Type: "error", Data: &ErrorEvent{err.Error()}}
		return
	}

	if !ack.Ok {
		rtm.Events <- Event{Type: "error", Data: &ErrorEvent{ack.Error}}
		return
	}

	rtm.Events <- Event{Type: "ack", Data: ack}
}
<|endoftext|>


<commit_before>package run

import (
	"bufio"
	"io"
	"os"
	"os/exec"
	"syscall"
)

// Run runs the command and returns a channel of output lines, errors and result of cmd.Wait
func Run(cmd *exec.Cmd) (
	lines chan string, errors chan error, resultCh chan error) {
	lines = make(chan string)
	errors = make(chan error, 1)
	resultCh = make(chan error)

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		errors <- err
		return
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		errors <- err
		return
	}

	err = cmd.Start()
	if err != nil {
		errors <- err
		return
	}

	go tailReader(bufio.NewReader(stdout), lines, errors)
	go tailReader(bufio.NewReader(stderr), lines, errors)

	go func() {
		resultCh <- cmd.Wait()
	}()

	return
}

func tailReader(r *bufio.Reader, ch chan string, errCh chan error) {
	for {
		line, _, err := r.ReadLine()
		if err != nil {
			// log.Printf("** %T | %#v | %v\n", err, err, err)

			// Apparently EOF and EBADF are the two errors we must
			// safely ignore when reading from a Stdout or Stderr pipe
			// of a exec.Command - because they seem to be returned by
			// io.Reader when the command exits. Other errors are
			// considered abnormal, so we return them to the caller
			// via errCh.
			if err == io.EOF {
				break
			}
			if e, ok := err.(*os.PathError); ok && e.Err == syscall.EBADF {
				break
			}
			// unknown error
			errCh <- err
			break
		}
		ch <- string(line)
	}
}
<commit_msg>call Wait only after the stdout/stderr pipes are closed<commit_after>package run

import (
	"log"
	"bufio"
	"io"
	"os/exec"
	"sync"
)

// Run runs the command and returns a channel of output lines, errors and result of cmd.Wait
func Run(cmd *exec.Cmd) (
	lines chan string, errors chan error, resultCh chan error) {
	lines = make(chan string)
	errors = make(chan error, 1)
	resultCh = make(chan error)

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		errors <- err
		return
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		errors <- err
		return
	}

	err = cmd.Start()
	if err != nil {
		errors <- err
		return
	}

	var wg sync.WaitGroup
	go tailReader(bufio.NewReader(stdout), lines, errors, wg)
	go tailReader(bufio.NewReader(stderr), lines, errors, wg)

	go func() {
		wg.Wait()
		resultCh <- cmd.Wait()
	}()

	return
}

func tailReader(
	r *bufio.Reader, ch chan string, errCh chan error,
	wg sync.WaitGroup) {
	defer wg.Done()
	wg.Add(1)

	for {
		line, _, err := r.ReadLine()
		if err != nil {
			log.Printf("** %T | %#v | %v\n", err, err, err)

			if err != io.EOF {
				errCh <- err
			}
			break
		}
		ch <- string(line)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"os/signal"
	"sort"
	"strings"
	"sync"
	"syscall"

	"github.com/mitchellh/iochan"
)

// Run runs swan migrations.
func Run(filename, dir string) error {
	// Open the last migration file
	f, err := os.OpenFile(filename, os.O_CREATE|os.O_RDWR, 0644)
	if err != nil {
		fmt.Println("unable to read or create file", filename)
		return err
	}
	defer f.Close()

	// Read the last migration stored in it.
	var lastB []byte
	lastB, err = ioutil.ReadAll(f)
	if err != nil {

	}
	last := string(lastB)
	last = strings.Trim(last, "\n\r")

	// skipFirstMigration is used to indicate skipping the first migration. This
	// is used because if the last migration is x, we want to run everything after x,
	// not including x. However, if there was no last migration we want to run all
	// migrations and not skip the first one.
	skipFirstMigration := last != ""

	// Read migrations from directory
	files, err := ioutil.ReadDir(dir)
	if err != nil {

	}

	// Extract file names from file infos.
	migrations := make([]string, len(files))
	for i, file := range files {
		migrations[i] = file.Name()
	}

	// Sort migrations so they can be run in order
	sort.Strings(migrations)

	// Execute later migrations
	i := sort.SearchStrings(migrations, last)
	if i > len(migrations) {
		fmt.Println("No such migration", last)
		return errors.New("no such migration")
	}

	if skipFirstMigration {
		i++
	}

	// Are there no migrations to run?
	if i == len(migrations) {
		fmt.Println("up to date")
		return nil
	}

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

	// Execute each migration
	var quit bool
	for ; i < len(migrations); i++ {
		quit = false
		outCh := make(chan string, 1)
		doneCh := make(chan struct{}, 1)
		go func() {
			for line := range outCh {
				fmt.Println(line)
			}
			close(doneCh)
		}()

		migration := "migrations/" + migrations[i]
		outCh <- "==> executing " + migration
		cmd := exec.Command(migration)

		if err := execute(cmd, outCh, sigCh); err != nil {
			outCh <- "==> " + err.Error()
			quit = true
		}

		// Signal that there are no more lines to print
		close(outCh)
		// Wait for lines to finish printing
		<-doneCh
		if quit {
			break
		}
		fmt.Println("==> OK\t" + migration)
	}

	// if i is 0, we haven't run a migration before and we failed to run the first
	// migration.
	if i == 0 {
		return errors.New("no migration executed")
	}

	// Empty the file
	if err := f.Truncate(0); err != nil {
		fmt.Println("unable to truncate file", err)
		return err
	}

	// Seek to the beginning of the file
	if _, err := f.Seek(0, 0); err != nil {
		fmt.Println("unable to seek file", err)
		return err
	}

	// Write the migration to the last migration file
	lastMigration := migrations[i-1]
	if _, err := f.WriteString(lastMigration); err != nil {
		fmt.Println("unable to write last migration", lastMigration, "to", filename)
		return err
	}
	return nil
}

func execute(cmd *exec.Cmd, outCh chan<- string, sigCh <-chan os.Signal) error {
	outPipe, err := cmd.StdoutPipe()
	if err != nil {
		return fmt.Errorf("unable to read stdout from command: %s", err)
	}
	errPipe, err := cmd.StderrPipe()
	if err != nil {
		return fmt.Errorf("unable to read stderr from command: %s", err)
	}

	// Create the channels we'll use for data
	exitCh := make(chan int, 1)
	doneCh := make(chan interface{}, 1)
	stdoutCh := iochan.DelimReader(outPipe, '\n')
	stderrCh := iochan.DelimReader(errPipe, '\n')
	if err := cmd.Start(); err != nil {
		return fmt.Errorf("unable to start command: %s", err)
	}

	go func() {
		select {
		case <-doneCh:
			return
		case sig := <-sigCh:
			cmd.Process.Signal(sig)
		}
	}()

	// Start the goroutine to watch for the exit
	go func() {
		exitStatus := 0

		err := cmd.Wait()
		doneCh <- struct{}{}

		if exitErr, ok := err.(*exec.ExitError); ok {
			exitStatus = 1

			// There is no process-independent way to get the REAL
			// exit status so we just try to go deeper.
			if status, ok := exitErr.Sys().(syscall.WaitStatus); ok {
				exitStatus = status.ExitStatus()
			}
		}

		exitCh <- exitStatus
	}()

	var wg sync.WaitGroup
	wg.Add(2)

	streamFunc := func(ch <-chan string) {
		defer wg.Done()
		for data := range ch {
			if data != "" {
				outCh <- data
			}
		}
	}

	go streamFunc(stdoutCh)
	go streamFunc(stderrCh)

	exitStatus := <-exitCh
	wg.Wait()

	if exitStatus != 0 {
		return fmt.Errorf("non-zero exit code: %d", exitStatus)
	}
	return nil
}
<commit_msg>Use provided directory<commit_after>package main

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"syscall"

	"github.com/mitchellh/iochan"
)

// Run runs swan migrations.
func Run(filename, dir string) error {
	// Open the last migration file
	f, err := os.OpenFile(filename, os.O_CREATE|os.O_RDWR, 0644)
	if err != nil {
		fmt.Println("unable to read or create file", filename)
		return err
	}
	defer f.Close()

	// Read the last migration stored in it.
	var lastB []byte
	lastB, err = ioutil.ReadAll(f)
	if err != nil {

	}
	last := string(lastB)
	last = strings.Trim(last, "\n\r")

	// skipFirstMigration is used to indicate skipping the first migration. This
	// is used because if the last migration is x, we want to run everything after x,
	// not including x. However, if there was no last migration we want to run all
	// migrations and not skip the first one.
	skipFirstMigration := last != ""

	// Read migrations from directory
	files, err := ioutil.ReadDir(dir)
	if err != nil {

	}

	// Extract file names from file infos.
	migrations := make([]string, len(files))
	for i, file := range files {
		migrations[i] = file.Name()
	}

	// Sort migrations so they can be run in order
	sort.Strings(migrations)

	// Execute later migrations
	i := sort.SearchStrings(migrations, last)
	if i > len(migrations) {
		fmt.Println("No such migration", last)
		return errors.New("no such migration")
	}

	if skipFirstMigration {
		i++
	}

	// Are there no migrations to run?
	if i == len(migrations) {
		fmt.Println("up to date")
		return nil
	}

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

	// Execute each migration
	var quit bool
	for ; i < len(migrations); i++ {
		quit = false
		outCh := make(chan string, 1)
		doneCh := make(chan struct{}, 1)
		go func() {
			for line := range outCh {
				fmt.Println(line)
			}
			close(doneCh)
		}()

		migration := filepath.Join(dir, migrations[i])
		outCh <- "==> executing " + migration
		cmd := exec.Command(migration)

		if err := execute(cmd, outCh, sigCh); err != nil {
			outCh <- "==> " + err.Error()
			quit = true
		}

		// Signal that there are no more lines to print
		close(outCh)
		// Wait for lines to finish printing
		<-doneCh
		if quit {
			break
		}
		fmt.Println("==> OK\t" + migration)
	}

	// if i is 0, we haven't run a migration before and we failed to run the first
	// migration.
	if i == 0 {
		return errors.New("no migration executed")
	}

	// Empty the file
	if err := f.Truncate(0); err != nil {
		fmt.Println("unable to truncate file", err)
		return err
	}

	// Seek to the beginning of the file
	if _, err := f.Seek(0, 0); err != nil {
		fmt.Println("unable to seek file", err)
		return err
	}

	// Write the migration to the last migration file
	lastMigration := migrations[i-1]
	if _, err := f.WriteString(lastMigration); err != nil {
		fmt.Println("unable to write last migration", lastMigration, "to", filename)
		return err
	}
	return nil
}

func execute(cmd *exec.Cmd, outCh chan<- string, sigCh <-chan os.Signal) error {
	outPipe, err := cmd.StdoutPipe()
	if err != nil {
		return fmt.Errorf("unable to read stdout from command: %s", err)
	}
	errPipe, err := cmd.StderrPipe()
	if err != nil {
		return fmt.Errorf("unable to read stderr from command: %s", err)
	}

	// Create the channels we'll use for data
	exitCh := make(chan int, 1)
	doneCh := make(chan interface{}, 1)
	stdoutCh := iochan.DelimReader(outPipe, '\n')
	stderrCh := iochan.DelimReader(errPipe, '\n')
	if err := cmd.Start(); err != nil {
		return fmt.Errorf("unable to start command: %s", err)
	}

	go func() {
		select {
		case <-doneCh:
			return
		case sig := <-sigCh:
			cmd.Process.Signal(sig)
		}
	}()

	// Start the goroutine to watch for the exit
	go func() {
		exitStatus := 0

		err := cmd.Wait()
		doneCh <- struct{}{}

		if exitErr, ok := err.(*exec.ExitError); ok {
			exitStatus = 1

			// There is no process-independent way to get the REAL
			// exit status so we just try to go deeper.
			if status, ok := exitErr.Sys().(syscall.WaitStatus); ok {
				exitStatus = status.ExitStatus()
			}
		}

		exitCh <- exitStatus
	}()

	var wg sync.WaitGroup
	wg.Add(2)

	streamFunc := func(ch <-chan string) {
		defer wg.Done()
		for data := range ch {
			if data != "" {
				outCh <- data
			}
		}
	}

	go streamFunc(stdoutCh)
	go streamFunc(stderrCh)

	exitStatus := <-exitCh
	wg.Wait()

	if exitStatus != 0 {
		return fmt.Errorf("non-zero exit code: %d", exitStatus)
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"crypto/tls"
	"io"
	"log"
	"net/url"
	"os"
	"os/exec"
	"os/signal"
	"strings"
)

var (
	detachedRun bool
)

var cmdRun = &Command{
	Run:   runRun,
	Usage: "run [-a app] command [arguments]",
	Short: "run a process",
	Long:  `Run a process on Heroku`,
}

func init() {
	cmdRun.Flag.StringVar(&flagApp, "a", "", "app")
	cmdRun.Flag.BoolVar(&detachedRun, "d", false, "detached")
}

func stty(args ...string) *exec.Cmd {
	c := exec.Command("stty", args...)
	c.Stdin = os.Stdin
	return c
}

func runRun(cmd *Command, args []string) {
	data := make(url.Values)
	if !detachedRun {
		data.Add("attach", "true")
	}
	data.Add("command", strings.Join(args, " "))

	resp := struct {
		Url *string `json:"rendezvous_url,omitempty"`
	}{}

	r := APIReq("POST", "/apps/"+mustApp()+"/ps")
	r.SetBodyForm(data)
	r.Do(&resp)

	if detachedRun {
		return
	}

	u, err := url.Parse(*resp.Url)
	if err != nil {
		log.Fatal(err)
	}

	cn, err := tls.Dial("tcp", u.Host, nil)
	if err != nil {
		log.Fatal(err)
	}
	defer cn.Close()

	br := bufio.NewReader(cn)

	_, err = io.WriteString(cn, u.Path[1:]+"\r\n")
	if err != nil {
		log.Fatal(err)
	}

	for {
		_, pre, err := br.ReadLine()
		if err != nil {
			log.Fatal(err)
		}
		if !pre {
			break
		}
	}

	if isTerminal(os.Stdin) && isTerminal(os.Stdout) {
		err = stty("-icanon", "-echo").Run()
		if err != nil {
			log.Fatal(err)
		}
		defer stty("icanon", "echo").Run()

		sig := make(chan os.Signal)
		signal.Notify(sig)
		go func() {
			<-sig
			// This will get called after the first os.Exit()
			stty("icanon", "echo").Run()
			os.Exit(1)
		}()
	}

	errc := make(chan error)
	cp := func(a io.Writer, b io.Reader) {
		_, err := io.Copy(a, b)
		errc <- err
	}

	go cp(os.Stdout, br)
	go cp(cn, os.Stdin)
	if err = <-errc; err != nil {
		log.Fatal(err)
	}
}
<commit_msg>handle SIGINT/QUIT properly - fixes #4<commit_after>package main

import (
	"bufio"
	"crypto/tls"
	"io"
	"log"
	"net/url"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"syscall"
)

var (
	detachedRun bool
)

var cmdRun = &Command{
	Run:   runRun,
	Usage: "run [-a app] command [arguments]",
	Short: "run a process",
	Long:  `Run a process on Heroku`,
}

func init() {
	cmdRun.Flag.StringVar(&flagApp, "a", "", "app")
	cmdRun.Flag.BoolVar(&detachedRun, "d", false, "detached")
}

func stty(args ...string) *exec.Cmd {
	c := exec.Command("stty", args...)
	c.Stdin = os.Stdin
	return c
}

func runRun(cmd *Command, args []string) {
	data := make(url.Values)
	if !detachedRun {
		data.Add("attach", "true")
	}
	data.Add("command", strings.Join(args, " "))

	resp := struct {
		Url *string `json:"rendezvous_url,omitempty"`
	}{}

	r := APIReq("POST", "/apps/"+mustApp()+"/ps")
	r.SetBodyForm(data)
	r.Do(&resp)

	if detachedRun {
		return
	}

	u, err := url.Parse(*resp.Url)
	if err != nil {
		log.Fatal(err)
	}

	cn, err := tls.Dial("tcp", u.Host, nil)
	if err != nil {
		log.Fatal(err)
	}
	defer cn.Close()

	br := bufio.NewReader(cn)

	_, err = io.WriteString(cn, u.Path[1:]+"\r\n")
	if err != nil {
		log.Fatal(err)
	}

	for {
		_, pre, err := br.ReadLine()
		if err != nil {
			log.Fatal(err)
		}
		if !pre {
			break
		}
	}

	if isTerminal(os.Stdin) && isTerminal(os.Stdout) {
		err = stty("-icanon", "-echo").Run()
		if err != nil {
			log.Fatal(err)
		}
		defer stty("icanon", "echo").Run()

		sig := make(chan os.Signal)
		signal.Notify(sig, os.Signal(syscall.SIGQUIT), os.Interrupt)
		go func() {
			defer stty("icanon", "echo").Run()
			for sg := range sig {
				switch sg {
				case os.Interrupt:
					cn.Write([]byte{3})
				case os.Signal(syscall.SIGQUIT):
					cn.Write([]byte{28})
				default:
					panic("not reached")
				}
			}
		}()
	}

	errc := make(chan error)
	cp := func(a io.Writer, b io.Reader) {
		_, err := io.Copy(a, b)
		errc <- err
	}

	go cp(os.Stdout, br)
	go cp(cn, os.Stdin)
	if err = <-errc; err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>