<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package libkb

import (
	"sync"

	keybase1 "github.com/keybase/client/go/protocol/keybase1"
)

type SecretRetriever interface {
	RetrieveSecret() ([]byte, error)
}

type SecretStorer interface {
	StoreSecret(secret []byte) error
}

type SecretStore interface {
	SecretRetriever
	SecretStorer
}

type SecretStoreAll interface {
	RetrieveSecret(username NormalizedUsername) ([]byte, error)
	StoreSecret(username NormalizedUsername, secret []byte) error
	ClearSecret(username NormalizedUsername) error
	GetUsersWithStoredSecrets() ([]string, error)
	GetApprovalPrompt() string
	GetTerminalPrompt() string
}

type SecretStoreContext interface {
	GetAllUserNames() (NormalizedUsername, []NormalizedUsername, error)
	GetStoredSecretServiceName() string
	GetStoredSecretAccessGroup() string
}

type SecretStoreImp struct {
	username NormalizedUsername
	store    *SecretStoreLocked
}

func (s *SecretStoreImp) RetrieveSecret() ([]byte, error) {
	return s.store.RetrieveSecret(s.username)
}

func (s *SecretStoreImp) StoreSecret(secret []byte) error {
	return s.store.StoreSecret(s.username, secret)
}

func NewSecretStore(g *GlobalContext, username NormalizedUsername) SecretStore {
	if g.SecretStoreAll != nil {
		return &SecretStoreImp{
			username: username,
			store:    g.SecretStoreAll,
		}
	}
	return nil
}

func GetConfiguredAccounts(c SecretStoreContext, s SecretStoreAll) ([]keybase1.ConfiguredAccount, error) {
	currentUsername, otherUsernames, err := c.GetAllUserNames()
	if err != nil {
		return nil, err
	}

	allUsernames := append(otherUsernames, currentUsername)

	accounts := make(map[NormalizedUsername]keybase1.ConfiguredAccount)

	for _, username := range allUsernames {
		accounts[username] = keybase1.ConfiguredAccount{
			Username: username.String(),
		}
	}
	var storedSecretUsernames []string
	if s != nil {
		storedSecretUsernames, err = s.GetUsersWithStoredSecrets()
	}
	if err != nil {
		return nil, err
	}

	for _, username := range storedSecretUsernames {
		nu := NewNormalizedUsername(username)
		account, ok := accounts[nu]
		if ok {
			account.HasStoredSecret = true
			accounts[nu] = account
		}
	}

	configuredAccounts := make([]keybase1.ConfiguredAccount, 0, len(accounts))
	for _, account := range accounts {
		configuredAccounts = append(configuredAccounts, account)
	}

	return configuredAccounts, nil
}

func ClearStoredSecret(g *GlobalContext, username NormalizedUsername) error {
	if g.SecretStoreAll == nil {
		return nil
	}
	return g.SecretStoreAll.ClearSecret(username)
}

// SecretStoreLocked protects a SecretStoreAll with a mutex.
type SecretStoreLocked struct {
	SecretStoreAll
	sync.Mutex
}

func NewSecretStoreLocked(g *GlobalContext) *SecretStoreLocked {
	return &SecretStoreLocked{
		SecretStoreAll: NewSecretStoreAll(g),
	}
}

func (s *SecretStoreLocked) RetrieveSecret(username NormalizedUsername) ([]byte, error) {
	s.Lock()
	defer s.Unlock()
	return s.SecretStoreAll.RetrieveSecret(username)
}

func (s *SecretStoreLocked) StoreSecret(username NormalizedUsername, secret []byte) error {
	s.Lock()
	defer s.Unlock()
	return s.SecretStoreAll.StoreSecret(username, secret)
}

func (s *SecretStoreLocked) ClearSecret(username NormalizedUsername) error {
	s.Lock()
	defer s.Unlock()
	return s.SecretStoreAll.ClearSecret(username)
}

func (s *SecretStoreLocked) GetUsersWithStoredSecrets() ([]string, error) {
	s.Lock()
	defer s.Unlock()
	if s.SecretStoreAll == nil {
		return []string{}, nil
	}
	return s.SecretStoreAll.GetUsersWithStoredSecrets()
}
<commit_msg>Check for nil SecretStoreAll<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package libkb

import (
	"sync"

	keybase1 "github.com/keybase/client/go/protocol/keybase1"
)

type SecretRetriever interface {
	RetrieveSecret() ([]byte, error)
}

type SecretStorer interface {
	StoreSecret(secret []byte) error
}

type SecretStore interface {
	SecretRetriever
	SecretStorer
}

type SecretStoreAll interface {
	RetrieveSecret(username NormalizedUsername) ([]byte, error)
	StoreSecret(username NormalizedUsername, secret []byte) error
	ClearSecret(username NormalizedUsername) error
	GetUsersWithStoredSecrets() ([]string, error)
	GetApprovalPrompt() string
	GetTerminalPrompt() string
}

type SecretStoreContext interface {
	GetAllUserNames() (NormalizedUsername, []NormalizedUsername, error)
	GetStoredSecretServiceName() string
	GetStoredSecretAccessGroup() string
}

type SecretStoreImp struct {
	username NormalizedUsername
	store    *SecretStoreLocked
}

func (s *SecretStoreImp) RetrieveSecret() ([]byte, error) {
	return s.store.RetrieveSecret(s.username)
}

func (s *SecretStoreImp) StoreSecret(secret []byte) error {
	return s.store.StoreSecret(s.username, secret)
}

func NewSecretStore(g *GlobalContext, username NormalizedUsername) SecretStore {
	if g.SecretStoreAll != nil {
		return &SecretStoreImp{
			username: username,
			store:    g.SecretStoreAll,
		}
	}
	return nil
}

func GetConfiguredAccounts(c SecretStoreContext, s SecretStoreAll) ([]keybase1.ConfiguredAccount, error) {
	currentUsername, otherUsernames, err := c.GetAllUserNames()
	if err != nil {
		return nil, err
	}

	allUsernames := append(otherUsernames, currentUsername)

	accounts := make(map[NormalizedUsername]keybase1.ConfiguredAccount)

	for _, username := range allUsernames {
		accounts[username] = keybase1.ConfiguredAccount{
			Username: username.String(),
		}
	}
	var storedSecretUsernames []string
	if s != nil {
		storedSecretUsernames, err = s.GetUsersWithStoredSecrets()
	}
	if err != nil {
		return nil, err
	}

	for _, username := range storedSecretUsernames {
		nu := NewNormalizedUsername(username)
		account, ok := accounts[nu]
		if ok {
			account.HasStoredSecret = true
			accounts[nu] = account
		}
	}

	configuredAccounts := make([]keybase1.ConfiguredAccount, 0, len(accounts))
	for _, account := range accounts {
		configuredAccounts = append(configuredAccounts, account)
	}

	return configuredAccounts, nil
}

func ClearStoredSecret(g *GlobalContext, username NormalizedUsername) error {
	if g.SecretStoreAll == nil {
		return nil
	}
	return g.SecretStoreAll.ClearSecret(username)
}

// SecretStoreLocked protects a SecretStoreAll with a mutex.
type SecretStoreLocked struct {
	SecretStoreAll
	sync.Mutex
}

func NewSecretStoreLocked(g *GlobalContext) *SecretStoreLocked {
	return &SecretStoreLocked{
		SecretStoreAll: NewSecretStoreAll(g),
	}
}

func (s *SecretStoreLocked) RetrieveSecret(username NormalizedUsername) ([]byte, error) {
	s.Lock()
	defer s.Unlock()
	if s.SecretStoreAll == nil {
		return nil, nil
	}
	return s.SecretStoreAll.RetrieveSecret(username)
}

func (s *SecretStoreLocked) StoreSecret(username NormalizedUsername, secret []byte) error {
	s.Lock()
	defer s.Unlock()
	if s.SecretStoreAll == nil {
		return nil
	}
	return s.SecretStoreAll.StoreSecret(username, secret)
}

func (s *SecretStoreLocked) ClearSecret(username NormalizedUsername) error {
	s.Lock()
	defer s.Unlock()
	if s.SecretStoreAll == nil {
		return nil
	}
	return s.SecretStoreAll.ClearSecret(username)
}

func (s *SecretStoreLocked) GetUsersWithStoredSecrets() ([]string, error) {
	s.Lock()
	defer s.Unlock()
	if s.SecretStoreAll == nil {
		return []string{}, nil
	}
	return s.SecretStoreAll.GetUsersWithStoredSecrets()
}
<|endoftext|>


<commit_before>package mysql

import (
	"fmt"
	"reflect"
	"strings"

	"github.com/freneticmonkey/migrate/migrate/metadata"
	"github.com/freneticmonkey/migrate/migrate/table"
	"github.com/freneticmonkey/migrate/migrate/util"
)

// SQLOperation Stores the data associated with an alter operation for each migration step
type SQLOperation struct {
	Statement string
	Op        int
	Name      string
	Metadata  metadata.Metadata
}

// SQLOperations Slice helper type
type SQLOperations []SQLOperation

// Add Add an SQLOperation to the slice
func (s *SQLOperations) Add(op SQLOperation) {
	*s = append(*s, op)
}

// Merge Merge SQLOperations slices
func (s *SQLOperations) Merge(ops SQLOperations) {
	for _, slice := range ops {
		if len(slice.Statement) > 0 {
			s.Add(slice)
		}
	}
}

// StatementBuilder Helper for building SQL ALTER TABLE statements.
// Assists with string concatenation and SQL specific formatting requirements.
type StatementBuilder struct {
	Components []string
}

// Reset Empties the statements components
func (sb *StatementBuilder) Reset() {
	sb.Components = []string{}
}

// Add Add a new statement component
func (sb *StatementBuilder) Add(component string) {
	if component != "" {
		sb.Components = append(sb.Components, component)
	}
}

// AddFormat Add a new formatted statement component
func (sb *StatementBuilder) AddFormat(format string, info ...interface{}) {
	canAdd := true
	for _, i := range info {
		if reflect.TypeOf(i).Kind() == reflect.String {
			if i == "" {
				canAdd = false
				break
			}
		}
	}
	if canAdd {
		sb.Components = append(sb.Components, fmt.Sprintf(format, info...))
	}
}

// AddQuote Add a new component with `quotes`
func (sb *StatementBuilder) AddQuote(component string) {
	if component != "" {
		sb.Components = append(sb.Components, fmt.Sprintf("`%s`", component))
	}
}

// AddType Add a component which correctly formats type size it it's provided
func (sb *StatementBuilder) AddType(typename string, size []int) {
	// Support for decimal places makes the size a little complicated
	if typename != "" {
		colType := ""

		switch len(size) {
		case 2:
			colType = fmt.Sprintf("%s(%d,%d)", typename, size[0], size[1])
		case 1:
			colType = fmt.Sprintf("%s(%d)", typename, size[0])
		default:
			colType = typename
		}
		sb.Add(colType)
	}

}

// Format Produce a formatted String
func (sb StatementBuilder) Format() string {
	return strings.Join(sb.Components, " ")
}

// generateCreateTable Generate a MySQL CREATE TABLE statement from a
// Table struct
func generateCreateTable(tbl table.Table) (operation SQLOperation) {

	operation.Op = table.Add

	tableTemplate := "CREATE TABLE `%s` (%s%s) ENGINE=%s%s DEFAULT CHARSET=%s;"

	columns := []string{}
	isAutoInc := false

	for _, col := range tbl.Columns {
		columns = append(columns, col.ToSQL())
		if col.AutoInc {
			isAutoInc = true
		}
	}

	indexes := []string{}

	if tbl.PrimaryIndex.IsValid() {
		indexes = append(indexes, tbl.PrimaryIndex.ToSQL())
	}

	for _, ind := range tbl.SecondaryIndexes {
		if ind.IsValid() {
			indexes = append(indexes, ind.ToSQL())
		}
	}

	strIndexes := ""
	if len(indexes) > 0 {
		strIndexes = ", " + strings.Join(indexes, ",")
	}

	autoInc := ""
	if isAutoInc {
		autoInc = fmt.Sprintf(" AUTO_INCREMENT=%d", tbl.AutoInc)
	}

	operation.Statement = fmt.Sprintf(tableTemplate, tbl.Name, strings.Join(columns, ","), strIndexes, tbl.Engine, autoInc, tbl.CharSet)
	operation.Metadata = tbl.Metadata
	return operation
}

// generateAlterColumn Generate a MySQL ALTER COLUMN statement from a
// Table struct
func generateAlterColumn(diff table.Diff) (ops SQLOperations) {
	var operation SQLOperation
	var builder StatementBuilder

	operation.Op = diff.Op
	operation.Metadata = diff.Metadata
	operation.Name = diff.Metadata.Name

	builder.Add("ALTER TABLE")

	switch diff.Op {

	case table.Add:
		builder.AddQuote(diff.Table)
		builder.Add("COLUMN")
		builder.AddQuote(diff.Property)

		column, ok := diff.Value.(table.Column)
		if ok {
			builder.AddType(column.Type, column.Size)

			if !column.Nullable {
				builder.Add("NOT NULL")
			}
		}

	case table.Del:
		builder.AddQuote(diff.Table)
		builder.Add("DROP COLUMN")
		builder.AddQuote(diff.Property)

	case table.Mod:
		// Process modification by type

		diffPair := diff.Value.(table.DiffPair)
		toColumn := diffPair.To.(table.Column)
		fromColumn := diffPair.From.(table.Column)

		builder.AddQuote(diff.Table)

		// Name needs special handling because it requires a different number of components
		// Assuming that we are modifying the column definition by default
		if diff.Property == "Name" {
			builder.Add("CHANGE COLUMN")
			builder.AddFormat("`%s` `%s`", fromColumn.Name, toColumn.Name)
			operation.Name = toColumn.Name

		} else {
			builder.Add("MODIFY COLUMN")
			builder.AddQuote(fromColumn.Name)
		}

		// Support for decimal places makes the size a little complicated
		builder.AddType(toColumn.Type, toColumn.Size)

		// if Nullable is T or F
		if !toColumn.Nullable {
			builder.Add("NOT NULL")
		}

		// if AutoInc is T or F
		if toColumn.AutoInc {
			builder.Add("AUTO_INCREMENT")
		}

		// if a Default value is defined
		if len(toColumn.Default) > 0 {
			if toColumn.Default == NULL {
				builder.Add("DEFAULT NULL")
			} else {
				builder.AddFormat("DEFAULT '%s'", toColumn.Default)
			}
		}
	}

	operation.Statement = builder.Format()

	ops.Add(operation)
	return ops
}

// generateAlterIndex Generate a MySQL ALTER INDEX statement from a
// Table struct
func generateAlterIndex(diff table.Diff) (ops SQLOperations) {

	var builder StatementBuilder

	// Obtain Index Object
	diffPair := diff.Value.(table.DiffPair)
	toIndex, ok := diffPair.To.(table.Index)

	if ok {
		indexName := ""
		columns := fmt.Sprintf("(`%s`)", strings.Join(toIndex.Columns, "`, `"))

		if diff.Field == "PrimaryIndex" {
			indexName = "PRIMARY KEY"

		} else if diff.Field == "SecondaryIndexes" {
			indexName = fmt.Sprintf("%s", toIndex.Name)
		}

		// Drop
		builder.Add("DROP INDEX")
		builder.AddQuote(indexName)
		builder.Add("ON")
		builder.AddQuote(diff.Table)

		removeOp := SQLOperation{
			Statement: builder.Format(),
			Op:        table.Del,
			Metadata:  diff.Metadata,
		}

		builder.Reset()
		builder.Add("CREATE INDEX")
		builder.AddQuote(indexName)
		builder.Add("ON")
		builder.AddQuote(diff.Table)
		builder.Add(columns)

		addOp := SQLOperation{
			Statement: builder.Format(),
			Op:        table.Add,
			Metadata:  diff.Metadata,
		}

		switch diff.Op {

		case table.Add:
			ops.Add(addOp)

		case table.Del:
			ops.Add(removeOp)

		case table.Mod:
			// Process modification by type
			if diff.Property == "Name" {

				fromIndex, ok := diffPair.From.(table.Index)
				if ok {
					builder.Reset()
					builder.Add("ALTER TABLE")
					builder.AddQuote(diff.Table)
					builder.Add("RENAME")
					builder.AddFormat("%s %s", fromIndex.Name, toIndex.Name)

					ops.Add(SQLOperation{
						Statement: builder.Format(),
						Op:        table.Mod,
						Metadata:  diff.Metadata,
					})
				} else {
					util.LogError("Gen SQL: ALTER INDEX: MOD: Could not obtain from index")
				}
			} else {
				// if anything other than a rename, we need to drop the index and re-add
				ops.Add(removeOp)
				ops.Add(addOp)
			}
		}
	} else {
		util.LogError("Obtaining Index FAILED")
	}
	return ops
}

// generateAlterTable Generate a MySQL CREATE TABLE, DROP TABLE or ALTER TABLE statement from a
// Table struct
func generateAlterTable(diff table.Diff) (ops SQLOperations) {

	if diff.Field == "*" && diff.Property == "*" {
		switch diff.Op {
		case table.Add:
			// generate the create table
			tbl, ok := diff.Value.(table.Table)
			if ok {
				ops.Add(generateCreateTable(tbl))
			} else {
				util.LogError("ISSUES obtaining table object: " + diff.Table)
			}
		case table.Del:
			// generate the drop table
			ops.Add(SQLOperation{
				Statement: fmt.Sprintf("DROP TABLE `%s`;", diff.Table),
				Op:        table.Del,
				Name:      diff.Table,
				Metadata:  diff.Metadata,
			})
		}

	} else {
		switch diff.Property {

		case "Name":
			newTableName, ok := diff.Value.(string)
			if ok {
				ops.Add(SQLOperation{
					Statement: fmt.Sprintf("ALTER TABLE `%s` RENAME TO `%s`;", diff.Table, newTableName),
					Op:        table.Mod,
					Name:      newTableName,
					Metadata:  diff.Metadata,
				})
			} else {
				util.LogError("ISSUES obtaining table name for rename: " + diff.Table)
			}

		case "AutoInc":
			ops.Add(SQLOperation{
				Statement: fmt.Sprintf("ALTER TABLE `%s` AUTO_INCREMENT=%d;", diff.Table, diff.Value),
				Op:        table.Mod,
				Name:      diff.Table,
				Metadata:  diff.Metadata,
			})

		case "Engine":
			ops.Add(SQLOperation{
				Statement: fmt.Sprintf("ALTER TABLE `%s` ENGINE=%s;", diff.Table, diff.Value),
				Op:        table.Mod,
				Name:      diff.Table,
				Metadata:  diff.Metadata,
			})

		case "CharSet":
			ops.Add(SQLOperation{
				Statement: fmt.Sprintf("ALTER TABLE `%s` DEFAULT CHARACTER SET `%s`;", diff.Table, diff.Value),
				Op:        table.Mod,
				Name:      diff.Table,
				Metadata:  diff.Metadata,
			})
		}
	}

	return ops

}

// GenerateAlters Generate MySQL ALTER TABLE statements from the Differences
// between Table structs
func GenerateAlters(differences table.Differences) (operations SQLOperations) {

	for _, diff := range differences.Slice {
		var alter SQLOperations

		// Check if the Diff is for a table
		if diff.Field == "Columns" {
			// It's a column change.
			alter = generateAlterColumn(diff)

		} else if diff.Field == "PrimaryIndex" || diff.Field == "SecondaryIndexes" {
			// It's an index change.
			alter = generateAlterIndex(diff)

		} else {
			alter = generateAlterTable(diff)
		}
		operations.Merge(alter)
	}

	util.LogInfo("Generated MySQL")
	for _, ops := range operations {
		table.FormatOperation(ops.Statement, ops.Op)
	}

	return operations
}
<commit_msg>Limited Table AUTO_INCREMENT property to only be set if the AUTO_INCREMENT value is > 0<commit_after>package mysql

import (
	"fmt"
	"reflect"
	"strings"

	"github.com/freneticmonkey/migrate/migrate/metadata"
	"github.com/freneticmonkey/migrate/migrate/table"
	"github.com/freneticmonkey/migrate/migrate/util"
)

// SQLOperation Stores the data associated with an alter operation for each migration step
type SQLOperation struct {
	Statement string
	Op        int
	Name      string
	Metadata  metadata.Metadata
}

// SQLOperations Slice helper type
type SQLOperations []SQLOperation

// Add Add an SQLOperation to the slice
func (s *SQLOperations) Add(op SQLOperation) {
	*s = append(*s, op)
}

// Merge Merge SQLOperations slices
func (s *SQLOperations) Merge(ops SQLOperations) {
	for _, slice := range ops {
		if len(slice.Statement) > 0 {
			s.Add(slice)
		}
	}
}

// StatementBuilder Helper for building SQL ALTER TABLE statements.
// Assists with string concatenation and SQL specific formatting requirements.
type StatementBuilder struct {
	Components []string
}

// Reset Empties the statements components
func (sb *StatementBuilder) Reset() {
	sb.Components = []string{}
}

// Add Add a new statement component
func (sb *StatementBuilder) Add(component string) {
	if component != "" {
		sb.Components = append(sb.Components, component)
	}
}

// AddFormat Add a new formatted statement component
func (sb *StatementBuilder) AddFormat(format string, info ...interface{}) {
	canAdd := true
	for _, i := range info {
		if reflect.TypeOf(i).Kind() == reflect.String {
			if i == "" {
				canAdd = false
				break
			}
		}
	}
	if canAdd {
		sb.Components = append(sb.Components, fmt.Sprintf(format, info...))
	}
}

// AddQuote Add a new component with `quotes`
func (sb *StatementBuilder) AddQuote(component string) {
	if component != "" {
		sb.Components = append(sb.Components, fmt.Sprintf("`%s`", component))
	}
}

// AddType Add a component which correctly formats type size it it's provided
func (sb *StatementBuilder) AddType(typename string, size []int) {
	// Support for decimal places makes the size a little complicated
	if typename != "" {
		colType := ""

		switch len(size) {
		case 2:
			colType = fmt.Sprintf("%s(%d,%d)", typename, size[0], size[1])
		case 1:
			colType = fmt.Sprintf("%s(%d)", typename, size[0])
		default:
			colType = typename
		}
		sb.Add(colType)
	}

}

// Format Produce a formatted String
func (sb StatementBuilder) Format() string {
	return strings.Join(sb.Components, " ")
}

// generateCreateTable Generate a MySQL CREATE TABLE statement from a
// Table struct
func generateCreateTable(tbl table.Table) (operation SQLOperation) {

	operation.Op = table.Add

	tableTemplate := "CREATE TABLE `%s` (%s%s) ENGINE=%s%s DEFAULT CHARSET=%s;"

	columns := []string{}
	isAutoInc := false

	for _, col := range tbl.Columns {
		columns = append(columns, col.ToSQL())
		if col.AutoInc {
			isAutoInc = true
		}
	}

	indexes := []string{}

	if tbl.PrimaryIndex.IsValid() {
		indexes = append(indexes, tbl.PrimaryIndex.ToSQL())
	}

	for _, ind := range tbl.SecondaryIndexes {
		if ind.IsValid() {
			indexes = append(indexes, ind.ToSQL())
		}
	}

	strIndexes := ""
	if len(indexes) > 0 {
		strIndexes = ", " + strings.Join(indexes, ",")
	}

	autoInc := ""
	// If AUTO_INCREMENT is being used and it has a non-zero value
	if isAutoInc && tbl.AutoInc > 0 {
		autoInc = fmt.Sprintf(" AUTO_INCREMENT=%d", tbl.AutoInc)
	}

	operation.Statement = fmt.Sprintf(tableTemplate, tbl.Name, strings.Join(columns, ","), strIndexes, tbl.Engine, autoInc, tbl.CharSet)
	operation.Metadata = tbl.Metadata
	return operation
}

// generateAlterColumn Generate a MySQL ALTER COLUMN statement from a
// Table struct
func generateAlterColumn(diff table.Diff) (ops SQLOperations) {
	var operation SQLOperation
	var builder StatementBuilder

	operation.Op = diff.Op
	operation.Metadata = diff.Metadata
	operation.Name = diff.Metadata.Name

	builder.Add("ALTER TABLE")

	switch diff.Op {

	case table.Add:
		builder.AddQuote(diff.Table)
		builder.Add("COLUMN")
		builder.AddQuote(diff.Property)

		column, ok := diff.Value.(table.Column)
		if ok {
			builder.AddType(column.Type, column.Size)

			if !column.Nullable {
				builder.Add("NOT NULL")
			}
		}

	case table.Del:
		builder.AddQuote(diff.Table)
		builder.Add("DROP COLUMN")
		builder.AddQuote(diff.Property)

	case table.Mod:
		// Process modification by type

		diffPair := diff.Value.(table.DiffPair)
		toColumn := diffPair.To.(table.Column)
		fromColumn := diffPair.From.(table.Column)

		builder.AddQuote(diff.Table)

		// Name needs special handling because it requires a different number of components
		// Assuming that we are modifying the column definition by default
		if diff.Property == "Name" {
			builder.Add("CHANGE COLUMN")
			builder.AddFormat("`%s` `%s`", fromColumn.Name, toColumn.Name)
			operation.Name = toColumn.Name

		} else {
			builder.Add("MODIFY COLUMN")
			builder.AddQuote(fromColumn.Name)
		}

		// Support for decimal places makes the size a little complicated
		builder.AddType(toColumn.Type, toColumn.Size)

		// if Nullable is T or F
		if !toColumn.Nullable {
			builder.Add("NOT NULL")
		}

		// if AutoInc is T or F
		if toColumn.AutoInc {
			builder.Add("AUTO_INCREMENT")
		}

		// if a Default value is defined
		if len(toColumn.Default) > 0 {
			if toColumn.Default == NULL {
				builder.Add("DEFAULT NULL")
			} else {
				builder.AddFormat("DEFAULT '%s'", toColumn.Default)
			}
		}
	}

	operation.Statement = builder.Format()

	ops.Add(operation)
	return ops
}

// generateAlterIndex Generate a MySQL ALTER INDEX statement from a
// Table struct
func generateAlterIndex(diff table.Diff) (ops SQLOperations) {

	var builder StatementBuilder

	// Obtain Index Object
	diffPair := diff.Value.(table.DiffPair)
	toIndex, ok := diffPair.To.(table.Index)

	if ok {
		indexName := ""
		columns := fmt.Sprintf("(`%s`)", strings.Join(toIndex.Columns, "`, `"))

		if diff.Field == "PrimaryIndex" {
			indexName = "PRIMARY KEY"

		} else if diff.Field == "SecondaryIndexes" {
			indexName = fmt.Sprintf("%s", toIndex.Name)
		}

		// Drop
		builder.Add("DROP INDEX")
		builder.AddQuote(indexName)
		builder.Add("ON")
		builder.AddQuote(diff.Table)

		removeOp := SQLOperation{
			Statement: builder.Format(),
			Op:        table.Del,
			Metadata:  diff.Metadata,
		}

		builder.Reset()
		builder.Add("CREATE INDEX")
		builder.AddQuote(indexName)
		builder.Add("ON")
		builder.AddQuote(diff.Table)
		builder.Add(columns)

		addOp := SQLOperation{
			Statement: builder.Format(),
			Op:        table.Add,
			Metadata:  diff.Metadata,
		}

		switch diff.Op {

		case table.Add:
			ops.Add(addOp)

		case table.Del:
			ops.Add(removeOp)

		case table.Mod:
			// Process modification by type
			if diff.Property == "Name" {

				fromIndex, ok := diffPair.From.(table.Index)
				if ok {
					builder.Reset()
					builder.Add("ALTER TABLE")
					builder.AddQuote(diff.Table)
					builder.Add("RENAME")
					builder.AddFormat("%s %s", fromIndex.Name, toIndex.Name)

					ops.Add(SQLOperation{
						Statement: builder.Format(),
						Op:        table.Mod,
						Metadata:  diff.Metadata,
					})
				} else {
					util.LogError("Gen SQL: ALTER INDEX: MOD: Could not obtain from index")
				}
			} else {
				// if anything other than a rename, we need to drop the index and re-add
				ops.Add(removeOp)
				ops.Add(addOp)
			}
		}
	} else {
		util.LogError("Obtaining Index FAILED")
	}
	return ops
}

// generateAlterTable Generate a MySQL CREATE TABLE, DROP TABLE or ALTER TABLE statement from a
// Table struct
func generateAlterTable(diff table.Diff) (ops SQLOperations) {

	if diff.Field == "*" && diff.Property == "*" {
		switch diff.Op {
		case table.Add:
			// generate the create table
			tbl, ok := diff.Value.(table.Table)
			if ok {
				ops.Add(generateCreateTable(tbl))
			} else {
				util.LogError("ISSUES obtaining table object: " + diff.Table)
			}
		case table.Del:
			// generate the drop table
			ops.Add(SQLOperation{
				Statement: fmt.Sprintf("DROP TABLE `%s`;", diff.Table),
				Op:        table.Del,
				Name:      diff.Table,
				Metadata:  diff.Metadata,
			})
		}

	} else {
		switch diff.Property {

		case "Name":
			newTableName, ok := diff.Value.(string)
			if ok {
				ops.Add(SQLOperation{
					Statement: fmt.Sprintf("ALTER TABLE `%s` RENAME TO `%s`;", diff.Table, newTableName),
					Op:        table.Mod,
					Name:      newTableName,
					Metadata:  diff.Metadata,
				})
			} else {
				util.LogError("ISSUES obtaining table name for rename: " + diff.Table)
			}

		case "AutoInc":
			ops.Add(SQLOperation{
				Statement: fmt.Sprintf("ALTER TABLE `%s` AUTO_INCREMENT=%d;", diff.Table, diff.Value),
				Op:        table.Mod,
				Name:      diff.Table,
				Metadata:  diff.Metadata,
			})

		case "Engine":
			ops.Add(SQLOperation{
				Statement: fmt.Sprintf("ALTER TABLE `%s` ENGINE=%s;", diff.Table, diff.Value),
				Op:        table.Mod,
				Name:      diff.Table,
				Metadata:  diff.Metadata,
			})

		case "CharSet":
			ops.Add(SQLOperation{
				Statement: fmt.Sprintf("ALTER TABLE `%s` DEFAULT CHARACTER SET `%s`;", diff.Table, diff.Value),
				Op:        table.Mod,
				Name:      diff.Table,
				Metadata:  diff.Metadata,
			})
		}
	}

	return ops

}

// GenerateAlters Generate MySQL ALTER TABLE statements from the Differences
// between Table structs
func GenerateAlters(differences table.Differences) (operations SQLOperations) {

	for _, diff := range differences.Slice {
		var alter SQLOperations

		// Check if the Diff is for a table
		if diff.Field == "Columns" {
			// It's a column change.
			alter = generateAlterColumn(diff)

		} else if diff.Field == "PrimaryIndex" || diff.Field == "SecondaryIndexes" {
			// It's an index change.
			alter = generateAlterIndex(diff)

		} else {
			alter = generateAlterTable(diff)
		}
		operations.Merge(alter)
	}

	util.LogInfo("Generated MySQL")
	for _, ops := range operations {
		table.FormatOperation(ops.Statement, ops.Op)
	}

	return operations
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess Authors.

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

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

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

// Package servenv contains functionality that is common for all
// Vitess server programs.  It defines and initializes command line
// flags that control the runtime environment.
//
// After a server program has called flag.Parse, it needs to call
// env.Init to make env use the command line variables to initialize
// the environment. It also needs to call env.Close before exiting.
//
// Note: If you need to plug in any custom initialization/cleanup for
// a vitess distribution, register them using OnInit and onClose. A
// clean way of achieving that is adding to this package a file with
// an init() function that registers the hooks.
package servenv

import (
	"flag"
	"net/url"
	"os"
	"runtime"
	"strings"
	"sync"
	"syscall"
	"time"

	// register the HTTP handlers for profiling
	_ "net/http/pprof"

	"vitess.io/vitess/go/event"
	"vitess.io/vitess/go/netutil"
	"vitess.io/vitess/go/stats"
	"vitess.io/vitess/go/vt/log"

	// register the proper init and shutdown hooks for logging
	_ "vitess.io/vitess/go/vt/logutil"
)

var (
	// Port is part of the flags used when calling RegisterDefaultFlags.
	Port *int

	// Flags to alter the behavior of the library.
	lameduckPeriod       = flag.Duration("lameduck-period", 50*time.Millisecond, "keep running at least this long after SIGTERM before stopping")
	onTermTimeout        = flag.Duration("onterm_timeout", 10*time.Second, "wait no more than this for OnTermSync handlers before stopping")
	memProfileRate       = flag.Int("mem-profile-rate", 512*1024, "profile every n bytes allocated")
	mutexProfileFraction = flag.Int("mutex-profile-fraction", 0, "profile every n mutex contention events (see runtime.SetMutexProfileFraction)")

	// mutex used to protect the Init function
	mu sync.Mutex

	onInitHooks     event.Hooks
	onTermHooks     event.Hooks
	onTermSyncHooks event.Hooks
	onRunHooks      event.Hooks
	inited          bool

	// ListeningURL is filled in when calling Run, contains the server URL.
	ListeningURL url.URL
)

// Init is the first phase of the server startup.
func Init() {
	mu.Lock()
	defer mu.Unlock()

	// Add version tag to every info log
	log.Infof(AppVersion.String())
	if inited {
		log.Fatal("servenv.Init called second time")
	}
	inited = true

	// Once you run as root, you pretty much destroy the chances of a
	// non-privileged user starting the program correctly.
	if uid := os.Getuid(); uid == 0 {
		log.Exitf("servenv.Init: running this as root makes no sense")
	}

	runtime.MemProfileRate = *memProfileRate

	if *mutexProfileFraction != 0 {
		log.Infof("setting mutex profile fraction to %v", *mutexProfileFraction)
		runtime.SetMutexProfileFraction(*mutexProfileFraction)
	}

	// We used to set this limit directly, but you pretty much have to
	// use a root account to allow increasing a limit reliably. Dropping
	// privileges is also tricky. The best strategy is to make a shell
	// script set up the limits as root and switch users before starting
	// the server.
	fdLimit := &syscall.Rlimit{}
	if err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, fdLimit); err != nil {
		log.Errorf("max-open-fds failed: %v", err)
	}
	fdl := stats.NewGauge("MaxFds", "File descriptor limit")
	fdl.Set(int64(fdLimit.Cur))

	onInitHooks.Fire()
}

func populateListeningURL(port int32) {
	host, err := netutil.FullyQualifiedHostname()
	if err != nil {
		host, err = os.Hostname()
		if err != nil {
			log.Exitf("os.Hostname() failed: %v", err)
		}
	}
	ListeningURL = url.URL{
		Scheme: "http",
		Host:   netutil.JoinHostPort(host, port),
		Path:   "/",
	}
}

// OnInit registers f to be run at the beginning of the app
// lifecycle. It should be called in an init() function.
func OnInit(f func()) {
	onInitHooks.Add(f)
}

// OnTerm registers a function to be run when the process receives a SIGTERM.
// This allows the program to change its behavior during the lameduck period.
//
// All hooks are run in parallel, and there is no guarantee that the process
// will wait for them to finish before dying when the lameduck period expires.
//
// See also: OnTermSync
func OnTerm(f func()) {
	onTermHooks.Add(f)
}

// OnTermSync registers a function to be run when the process receives SIGTERM.
// This allows the program to change its behavior during the lameduck period.
//
// All hooks are run in parallel, and the process will do its best to wait
// (up to -onterm_timeout) for all of them to finish before dying.
//
// See also: OnTerm
func OnTermSync(f func()) {
	onTermSyncHooks.Add(f)
}

// fireOnTermSyncHooks returns true iff all the hooks finish before the timeout.
func fireOnTermSyncHooks(timeout time.Duration) bool {
	defer log.Flush()
	log.Infof("Firing synchronous OnTermSync hooks and waiting up to %v for them", timeout)

	timer := time.NewTimer(timeout)
	defer timer.Stop()

	done := make(chan struct{})
	go func() {
		onTermSyncHooks.Fire()
		close(done)
	}()

	select {
	case <-done:
		log.Infof("OnTermSync hooks finished")
		return true
	case <-timer.C:
		log.Infof("OnTermSync hooks timed out")
		return false
	}
}

// OnRun registers f to be run right at the beginning of Run. All
// hooks are run in parallel.
func OnRun(f func()) {
	onRunHooks.Add(f)
}

// FireRunHooks fires the hooks registered by OnHook.
// Use this in a non-server to run the hooks registered
// by servenv.OnRun().
func FireRunHooks() {
	onRunHooks.Fire()
}

// RegisterDefaultFlags registers the default flags for
// listening to a given port for standard connections.
// If calling this, then call RunDefault()
func RegisterDefaultFlags() {
	Port = flag.Int("port", 0, "port for the server")
}

// RunDefault calls Run() with the parameters from the flags.
func RunDefault() {
	Run(*Port)
}

// ParseFlags initializes flags and handles the common case when no positional
// arguments are expected.
func ParseFlags(cmd string) {
	flag.Parse()

	if *Version {
		AppVersion.Print()
		os.Exit(0)
	}

	args := flag.Args()
	if len(args) > 0 {
		flag.Usage()
		log.Exitf("%s doesn't take any positional arguments, got '%s'", cmd, strings.Join(args, " "))
	}
}

// ParseFlagsWithArgs initializes flags and returns the positional arguments
func ParseFlagsWithArgs(cmd string) []string {
	flag.Parse()

	if *Version {
		AppVersion.Print()
		os.Exit(0)
	}

	args := flag.Args()
	if len(args) == 0 {
		log.Exitf("%s expected at least one positional argument", cmd)
	}

	return args
}
<commit_msg>Add flag to catch and ignore SIGPIPE<commit_after>/*
Copyright 2019 The Vitess Authors.

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

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

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

// Package servenv contains functionality that is common for all
// Vitess server programs.  It defines and initializes command line
// flags that control the runtime environment.
//
// After a server program has called flag.Parse, it needs to call
// env.Init to make env use the command line variables to initialize
// the environment. It also needs to call env.Close before exiting.
//
// Note: If you need to plug in any custom initialization/cleanup for
// a vitess distribution, register them using OnInit and onClose. A
// clean way of achieving that is adding to this package a file with
// an init() function that registers the hooks.
package servenv

import (
	"flag"
	"net/url"
	"os"
	"os/signal"
	"runtime"
	"strings"
	"sync"
	"syscall"
	"time"

	// register the HTTP handlers for profiling
	_ "net/http/pprof"

	"vitess.io/vitess/go/event"
	"vitess.io/vitess/go/netutil"
	"vitess.io/vitess/go/stats"
	"vitess.io/vitess/go/vt/log"

	// register the proper init and shutdown hooks for logging
	_ "vitess.io/vitess/go/vt/logutil"
)

var (
	// Port is part of the flags used when calling RegisterDefaultFlags.
	Port *int

	// Flags to alter the behavior of the library.
	lameduckPeriod       = flag.Duration("lameduck-period", 50*time.Millisecond, "keep running at least this long after SIGTERM before stopping")
	onTermTimeout        = flag.Duration("onterm_timeout", 10*time.Second, "wait no more than this for OnTermSync handlers before stopping")
	memProfileRate       = flag.Int("mem-profile-rate", 512*1024, "profile every n bytes allocated")
	mutexProfileFraction = flag.Int("mutex-profile-fraction", 0, "profile every n mutex contention events (see runtime.SetMutexProfileFraction)")
	catchSigpipe         = flag.Bool("catch-sigpipe", false, "catch and ignore SIGPIPE on stdout and stderr if specified")

	// mutex used to protect the Init function
	mu sync.Mutex

	onInitHooks     event.Hooks
	onTermHooks     event.Hooks
	onTermSyncHooks event.Hooks
	onRunHooks      event.Hooks
	inited          bool

	// ListeningURL is filled in when calling Run, contains the server URL.
	ListeningURL url.URL
)

// Init is the first phase of the server startup.
func Init() {
	mu.Lock()
	defer mu.Unlock()

	// Ignore SIGPIPE if specified
	// The Go runtime catches SIGPIPE for us on all fds except stdout/stderr
	// See https://golang.org/pkg/os/signal/#hdr-SIGPIPE
	if *catchSigpipe {
		sigChan := make(chan os.Signal, 1)
		signal.Notify(sigChan, syscall.SIGPIPE)
		go func() {
			<-sigChan
			log.Warning("Caught SIGPIPE (ignoring all future SIGPIPEs)")
			signal.Ignore(syscall.SIGPIPE)
		}()
	}

	// Add version tag to every info log
	log.Infof(AppVersion.String())
	if inited {
		log.Fatal("servenv.Init called second time")
	}
	inited = true

	// Once you run as root, you pretty much destroy the chances of a
	// non-privileged user starting the program correctly.
	if uid := os.Getuid(); uid == 0 {
		log.Exitf("servenv.Init: running this as root makes no sense")
	}

	runtime.MemProfileRate = *memProfileRate

	if *mutexProfileFraction != 0 {
		log.Infof("setting mutex profile fraction to %v", *mutexProfileFraction)
		runtime.SetMutexProfileFraction(*mutexProfileFraction)
	}

	// We used to set this limit directly, but you pretty much have to
	// use a root account to allow increasing a limit reliably. Dropping
	// privileges is also tricky. The best strategy is to make a shell
	// script set up the limits as root and switch users before starting
	// the server.
	fdLimit := &syscall.Rlimit{}
	if err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, fdLimit); err != nil {
		log.Errorf("max-open-fds failed: %v", err)
	}
	fdl := stats.NewGauge("MaxFds", "File descriptor limit")
	fdl.Set(int64(fdLimit.Cur))

	onInitHooks.Fire()
}

func populateListeningURL(port int32) {
	host, err := netutil.FullyQualifiedHostname()
	if err != nil {
		host, err = os.Hostname()
		if err != nil {
			log.Exitf("os.Hostname() failed: %v", err)
		}
	}
	ListeningURL = url.URL{
		Scheme: "http",
		Host:   netutil.JoinHostPort(host, port),
		Path:   "/",
	}
}

// OnInit registers f to be run at the beginning of the app
// lifecycle. It should be called in an init() function.
func OnInit(f func()) {
	onInitHooks.Add(f)
}

// OnTerm registers a function to be run when the process receives a SIGTERM.
// This allows the program to change its behavior during the lameduck period.
//
// All hooks are run in parallel, and there is no guarantee that the process
// will wait for them to finish before dying when the lameduck period expires.
//
// See also: OnTermSync
func OnTerm(f func()) {
	onTermHooks.Add(f)
}

// OnTermSync registers a function to be run when the process receives SIGTERM.
// This allows the program to change its behavior during the lameduck period.
//
// All hooks are run in parallel, and the process will do its best to wait
// (up to -onterm_timeout) for all of them to finish before dying.
//
// See also: OnTerm
func OnTermSync(f func()) {
	onTermSyncHooks.Add(f)
}

// fireOnTermSyncHooks returns true iff all the hooks finish before the timeout.
func fireOnTermSyncHooks(timeout time.Duration) bool {
	defer log.Flush()
	log.Infof("Firing synchronous OnTermSync hooks and waiting up to %v for them", timeout)

	timer := time.NewTimer(timeout)
	defer timer.Stop()

	done := make(chan struct{})
	go func() {
		onTermSyncHooks.Fire()
		close(done)
	}()

	select {
	case <-done:
		log.Infof("OnTermSync hooks finished")
		return true
	case <-timer.C:
		log.Infof("OnTermSync hooks timed out")
		return false
	}
}

// OnRun registers f to be run right at the beginning of Run. All
// hooks are run in parallel.
func OnRun(f func()) {
	onRunHooks.Add(f)
}

// FireRunHooks fires the hooks registered by OnHook.
// Use this in a non-server to run the hooks registered
// by servenv.OnRun().
func FireRunHooks() {
	onRunHooks.Fire()
}

// RegisterDefaultFlags registers the default flags for
// listening to a given port for standard connections.
// If calling this, then call RunDefault()
func RegisterDefaultFlags() {
	Port = flag.Int("port", 0, "port for the server")
}

// RunDefault calls Run() with the parameters from the flags.
func RunDefault() {
	Run(*Port)
}

// ParseFlags initializes flags and handles the common case when no positional
// arguments are expected.
func ParseFlags(cmd string) {
	flag.Parse()

	if *Version {
		AppVersion.Print()
		os.Exit(0)
	}

	args := flag.Args()
	if len(args) > 0 {
		flag.Usage()
		log.Exitf("%s doesn't take any positional arguments, got '%s'", cmd, strings.Join(args, " "))
	}
}

// ParseFlagsWithArgs initializes flags and returns the positional arguments
func ParseFlagsWithArgs(cmd string) []string {
	flag.Parse()

	if *Version {
		AppVersion.Print()
		os.Exit(0)
	}

	args := flag.Args()
	if len(args) == 0 {
		log.Exitf("%s expected at least one positional argument", cmd)
	}

	return args
}
<|endoftext|>


<commit_before>// Copyright 2016, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package worker

import (
	"fmt"
	"time"

	"github.com/youtube/vitess/go/vt/discovery"
	"github.com/youtube/vitess/go/vt/throttler"
	"github.com/youtube/vitess/go/vt/topo/topoproto"
	"github.com/youtube/vitess/go/vt/wrangler"
	"golang.org/x/net/context"

	topodatapb "github.com/youtube/vitess/go/vt/proto/topodata"
)

// executor takes care of the write-side of the copy.
// There is one executor for each destination shard and writer thread.
// To-be-written data will be passed in through a channel.
// The main purpose of this struct is to aggregate the objects which won't
// change during the execution and remove them from method signatures.
type executor struct {
	wr        *wrangler.Wrangler
	tsc       *discovery.TabletStatsCache
	throttler *throttler.Throttler
	keyspace  string
	shard     string
	threadID  int
	// statsKey is the cached metric key which we need when we increment the stats
	// variable when we get throttled.
	statsKey []string
}

func newExecutor(wr *wrangler.Wrangler, tsc *discovery.TabletStatsCache, throttler *throttler.Throttler, keyspace, shard string, threadID int) *executor {
	return &executor{
		wr:        wr,
		tsc:       tsc,
		throttler: throttler,
		keyspace:  keyspace,
		shard:     shard,
		threadID:  threadID,
		statsKey:  []string{keyspace, shard, fmt.Sprint(threadID)},
	}
}

// fetchLoop loops over the provided insertChannel and sends the commands to the
// current master.
func (e *executor) fetchLoop(ctx context.Context, insertChannel chan string) error {
	for {
		select {
		case cmd, ok := <-insertChannel:
			if !ok {
				// no more to read, we're done
				return nil
			}
			if err := e.fetchWithRetries(ctx, cmd); err != nil {
				return fmt.Errorf("ExecuteFetch failed: %v", err)
			}
		case <-ctx.Done():
			// Doesn't really matter if this select gets starved, because the other case
			// will also return an error due to executeFetch's context being closed. This case
			// does prevent us from blocking indefinitely on insertChannel when the worker is canceled.
			return nil
		}
	}
}

// fetchWithRetries will attempt to run ExecuteFetch for a single command, with
// a reasonably small timeout.
// If will keep retrying the ExecuteFetch (for a finite but longer duration) if
// it fails due to a timeout or a retriable application error.
//
// executeFetchWithRetries will always get the current MASTER tablet from the
// TabletStatsCache instance. If no MASTER is available, it will keep retrying.
func (e *executor) fetchWithRetries(ctx context.Context, command string) error {
	retryDuration := *retryDuration
	// We should keep retrying up until the retryCtx runs out.
	retryCtx, retryCancel := context.WithTimeout(ctx, retryDuration)
	defer retryCancel()
	// Is this current attempt a retry of a previous attempt?
	isRetry := false
	for {
		var master *discovery.TabletStats
		var err error

		// Get the current master from the TabletStatsCache.
		masters := e.tsc.GetHealthyTabletStats(e.keyspace, e.shard, topodatapb.TabletType_MASTER)
		if len(masters) == 0 {
			e.wr.Logger().Warningf("ExecuteFetch failed for keyspace/shard %v/%v because no MASTER is available; will retry until there is MASTER again", e.keyspace, e.shard)
			statsRetryCount.Add(1)
			statsRetryCounters.Add(retryCategoryNoMasterAvailable, 1)
			goto retry
		}
		master = &masters[0]

		// Block if we are throttled.
		if e.throttler != nil {
			for {
				backoff := e.throttler.Throttle(e.threadID)
				if backoff == throttler.NotThrottled {
					break
				}
				statsThrottledCounters.Add(e.statsKey, 1)
				time.Sleep(backoff)
			}
		}

		// Run the command (in a block since goto above does not allow to introduce
		// new variables until the label is reached.)
		{
			tryCtx, cancel := context.WithTimeout(retryCtx, 2*time.Minute)
			_, err = e.wr.TabletManagerClient().ExecuteFetchAsApp(tryCtx, master.Tablet, true, []byte(command), 0)
			cancel()

			if err == nil {
				// success!
				return nil
			}

			succeeded, finalErr := e.checkError(tryCtx, err, isRetry, master)
			if succeeded {
				// We can ignore the error and don't have to retry.
				return nil
			}
			if finalErr != nil {
				// Non-retryable error.
				return finalErr
			}
		}

	retry:
		masterAlias := "no-master-was-available"
		if master != nil {
			masterAlias = topoproto.TabletAliasString(master.Tablet.Alias)
		}
		tabletString := fmt.Sprintf("%v (%v/%v)", masterAlias, e.keyspace, e.shard)

		select {
		case <-retryCtx.Done():
			if retryCtx.Err() == context.DeadlineExceeded {
				return fmt.Errorf("failed to connect to destination tablet %v after retrying for %v", tabletString, retryDuration)
			}
			return fmt.Errorf("interrupted while trying to run %v on tablet %v", command, tabletString)
		case <-time.After(*executeFetchRetryTime):
			// Retry 30s after the failure using the current master seen by the HealthCheck.
		}
		isRetry = true
	}
}

// checkError returns true if the error can be ignored and the command
// succeeded, false if the error is retryable and a non-nil error if the
// command must not be retried.
func (e *executor) checkError(ctx context.Context, err error, isRetry bool, master *discovery.TabletStats) (bool, error) {
	tabletString := fmt.Sprintf("%v (%v/%v)", topoproto.TabletAliasString(master.Tablet.Alias), e.keyspace, e.shard)

	// first see if it was a context timeout.
	select {
	case <-ctx.Done():
		if ctx.Err() == context.DeadlineExceeded {
			e.wr.Logger().Warningf("ExecuteFetch failed on %v; will retry because it was a timeout error on the context", tabletString)
			statsRetryCount.Add(1)
			statsRetryCounters.Add(retryCategoryTimeoutError, 1)
			return false, nil
		}
	default:
	}

	// If the ExecuteFetch call failed because of an application error, we will try to figure out why.
	// We need to extract the MySQL error number, and will attempt to retry if we think the error is recoverable.
	match := errExtract.FindStringSubmatch(err.Error())
	var errNo string
	if len(match) == 2 {
		errNo = match[1]
	}
	switch {
	case errNo == "1290":
		e.wr.Logger().Warningf("ExecuteFetch failed on %v; will reresolve and retry because it's due to a MySQL read-only error: %v", tabletString, err)
		statsRetryCount.Add(1)
		statsRetryCounters.Add(retryCategoryReadOnly, 1)
	case errNo == "2002" || errNo == "2006":
		e.wr.Logger().Warningf("ExecuteFetch failed on %v; will reresolve and retry because it's due to a MySQL connection error: %v", tabletString, err)
		statsRetryCount.Add(1)
		statsRetryCounters.Add(retryCategoryConnectionError, 1)
	case errNo == "1062":
		if !isRetry {
			return false, fmt.Errorf("ExecuteFetch failed on %v on the first attempt; not retrying as this is not a recoverable error: %v", tabletString, err)
		}
		e.wr.Logger().Infof("ExecuteFetch failed on %v with a duplicate entry error; marking this as a success, because of the likelihood that this query has already succeeded before being retried: %v", tabletString, err)
		return true, nil
	default:
		// Unknown error.
		return false, err
	}
	return false, nil
}
<commit_msg>worker: Retry writes if the MySQL error code is 2013 ("Lost connection").<commit_after>// Copyright 2016, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package worker

import (
	"fmt"
	"time"

	"github.com/youtube/vitess/go/vt/discovery"
	"github.com/youtube/vitess/go/vt/throttler"
	"github.com/youtube/vitess/go/vt/topo/topoproto"
	"github.com/youtube/vitess/go/vt/wrangler"
	"golang.org/x/net/context"

	topodatapb "github.com/youtube/vitess/go/vt/proto/topodata"
)

// executor takes care of the write-side of the copy.
// There is one executor for each destination shard and writer thread.
// To-be-written data will be passed in through a channel.
// The main purpose of this struct is to aggregate the objects which won't
// change during the execution and remove them from method signatures.
type executor struct {
	wr        *wrangler.Wrangler
	tsc       *discovery.TabletStatsCache
	throttler *throttler.Throttler
	keyspace  string
	shard     string
	threadID  int
	// statsKey is the cached metric key which we need when we increment the stats
	// variable when we get throttled.
	statsKey []string
}

func newExecutor(wr *wrangler.Wrangler, tsc *discovery.TabletStatsCache, throttler *throttler.Throttler, keyspace, shard string, threadID int) *executor {
	return &executor{
		wr:        wr,
		tsc:       tsc,
		throttler: throttler,
		keyspace:  keyspace,
		shard:     shard,
		threadID:  threadID,
		statsKey:  []string{keyspace, shard, fmt.Sprint(threadID)},
	}
}

// fetchLoop loops over the provided insertChannel and sends the commands to the
// current master.
func (e *executor) fetchLoop(ctx context.Context, insertChannel chan string) error {
	for {
		select {
		case cmd, ok := <-insertChannel:
			if !ok {
				// no more to read, we're done
				return nil
			}
			if err := e.fetchWithRetries(ctx, cmd); err != nil {
				return fmt.Errorf("ExecuteFetch failed: %v", err)
			}
		case <-ctx.Done():
			// Doesn't really matter if this select gets starved, because the other case
			// will also return an error due to executeFetch's context being closed. This case
			// does prevent us from blocking indefinitely on insertChannel when the worker is canceled.
			return nil
		}
	}
}

// fetchWithRetries will attempt to run ExecuteFetch for a single command, with
// a reasonably small timeout.
// If will keep retrying the ExecuteFetch (for a finite but longer duration) if
// it fails due to a timeout or a retriable application error.
//
// executeFetchWithRetries will always get the current MASTER tablet from the
// TabletStatsCache instance. If no MASTER is available, it will keep retrying.
func (e *executor) fetchWithRetries(ctx context.Context, command string) error {
	retryDuration := *retryDuration
	// We should keep retrying up until the retryCtx runs out.
	retryCtx, retryCancel := context.WithTimeout(ctx, retryDuration)
	defer retryCancel()
	// Is this current attempt a retry of a previous attempt?
	isRetry := false
	for {
		var master *discovery.TabletStats
		var err error

		// Get the current master from the TabletStatsCache.
		masters := e.tsc.GetHealthyTabletStats(e.keyspace, e.shard, topodatapb.TabletType_MASTER)
		if len(masters) == 0 {
			e.wr.Logger().Warningf("ExecuteFetch failed for keyspace/shard %v/%v because no MASTER is available; will retry until there is MASTER again", e.keyspace, e.shard)
			statsRetryCount.Add(1)
			statsRetryCounters.Add(retryCategoryNoMasterAvailable, 1)
			goto retry
		}
		master = &masters[0]

		// Block if we are throttled.
		if e.throttler != nil {
			for {
				backoff := e.throttler.Throttle(e.threadID)
				if backoff == throttler.NotThrottled {
					break
				}
				statsThrottledCounters.Add(e.statsKey, 1)
				time.Sleep(backoff)
			}
		}

		// Run the command (in a block since goto above does not allow to introduce
		// new variables until the label is reached.)
		{
			tryCtx, cancel := context.WithTimeout(retryCtx, 2*time.Minute)
			_, err = e.wr.TabletManagerClient().ExecuteFetchAsApp(tryCtx, master.Tablet, true, []byte(command), 0)
			cancel()

			if err == nil {
				// success!
				return nil
			}

			succeeded, finalErr := e.checkError(tryCtx, err, isRetry, master)
			if succeeded {
				// We can ignore the error and don't have to retry.
				return nil
			}
			if finalErr != nil {
				// Non-retryable error.
				return finalErr
			}
		}

	retry:
		masterAlias := "no-master-was-available"
		if master != nil {
			masterAlias = topoproto.TabletAliasString(master.Tablet.Alias)
		}
		tabletString := fmt.Sprintf("%v (%v/%v)", masterAlias, e.keyspace, e.shard)

		select {
		case <-retryCtx.Done():
			if retryCtx.Err() == context.DeadlineExceeded {
				return fmt.Errorf("failed to connect to destination tablet %v after retrying for %v", tabletString, retryDuration)
			}
			return fmt.Errorf("interrupted while trying to run %v on tablet %v", command, tabletString)
		case <-time.After(*executeFetchRetryTime):
			// Retry 30s after the failure using the current master seen by the HealthCheck.
		}
		isRetry = true
	}
}

// checkError returns true if the error can be ignored and the command
// succeeded, false if the error is retryable and a non-nil error if the
// command must not be retried.
func (e *executor) checkError(ctx context.Context, err error, isRetry bool, master *discovery.TabletStats) (bool, error) {
	tabletString := fmt.Sprintf("%v (%v/%v)", topoproto.TabletAliasString(master.Tablet.Alias), e.keyspace, e.shard)

	// first see if it was a context timeout.
	select {
	case <-ctx.Done():
		if ctx.Err() == context.DeadlineExceeded {
			e.wr.Logger().Warningf("ExecuteFetch failed on %v; will retry because it was a timeout error on the context", tabletString)
			statsRetryCount.Add(1)
			statsRetryCounters.Add(retryCategoryTimeoutError, 1)
			return false, nil
		}
	default:
	}

	// If the ExecuteFetch call failed because of an application error, we will try to figure out why.
	// We need to extract the MySQL error number, and will attempt to retry if we think the error is recoverable.
	match := errExtract.FindStringSubmatch(err.Error())
	var errNo string
	if len(match) == 2 {
		errNo = match[1]
	}
	switch {
	case errNo == "1290":
		e.wr.Logger().Warningf("ExecuteFetch failed on %v; will reresolve and retry because it's due to a MySQL read-only error: %v", tabletString, err)
		statsRetryCount.Add(1)
		statsRetryCounters.Add(retryCategoryReadOnly, 1)
	case errNo == "2002" || errNo == "2006" || errNo == "2013":
		// Note:
		// "2006" happens if the connection is already dead. Retrying a query in
		// this case is safe.
		// "2013" happens if the connection dies in the middle of a query. This is
		// also safe to retry because either the query went through on the server or
		// it was aborted. If we retry the query and get a duplicate entry error, we
		// assume that the previous execution was successful and ignore the error.
		// See below for the handling of duplicate entry error "1062".
		e.wr.Logger().Warningf("ExecuteFetch failed on %v; will reresolve and retry because it's due to a MySQL connection error: %v", tabletString, err)
		statsRetryCount.Add(1)
		statsRetryCounters.Add(retryCategoryConnectionError, 1)
	case errNo == "1062":
		if !isRetry {
			return false, fmt.Errorf("ExecuteFetch failed on %v on the first attempt; not retrying as this is not a recoverable error: %v", tabletString, err)
		}
		e.wr.Logger().Infof("ExecuteFetch failed on %v with a duplicate entry error; marking this as a success, because of the likelihood that this query has already succeeded before being retried: %v", tabletString, err)
		return true, nil
	default:
		// Unknown error.
		return false, err
	}
	return false, nil
}
<|endoftext|>


<commit_before>package gocd

import (
	"context"
	"fmt"
	"net/url"
	"strings"
)

// PipelineTemplatesService describes the HAL _link resource for the api response object for a pipeline configuration objects.
type PipelineTemplatesService service

// PipelineTemplatesLinks describes a single pipeline template config object HAL links
//go:generate gocd-response-links-generator -type=PipelineTemplatesLinks,PipelineTemplateLinks
type PipelineTemplatesLinks struct {
	Self *url.URL `json:"self"`
	Doc  *url.URL `json:"doc"`
	Find *url.URL `json:"find"`
}

// PipelineTemplateLinks describes multiple pipeline template config object HAL links
type PipelineTemplateLinks struct {
	Self *url.URL `json:"self"`
	Doc  *url.URL `json:"doc"`
	Find *url.URL `json:"find"`
}

// PipelineTemplateRequest describes a PipelineTemplate
type PipelineTemplateRequest struct {
	Name   string   `json:"name"`
	Stages []*Stage `json:"stages"`
}

// PipelineTemplateResponse describes an api response for a single pipeline templates
type PipelineTemplateResponse struct {
	Name     string `json:"name"`
	Embedded *struct {
		Pipelines []*struct {
			Name string `json:"name"`
		}
	} `json:"_embedded,omitempty"`
}

// PipelineTemplatesResponse describes an api response for multiple pipeline templates
type PipelineTemplatesResponse struct {
	Links    PipelineTemplatesLinks `json:"_links,omitempty"`
	Embedded *struct {
		Templates []*PipelineTemplate `json:"templates"`
	} `json:"_embedded,omitempty"`
}

type embeddedPipelineTemplate struct {
	Pipelines []*Pipeline `json:"pipelines,omitempty"`
}

// PipelineTemplate describes a response from the API for a pipeline template object.
type PipelineTemplate struct {
	Links    *PipelineTemplateLinks    `json:"_links,omitempty"`
	Name     string                    `json:"name"`
	Embedded *embeddedPipelineTemplate `json:"_embedded,omitempty"`
	Version  string                    `json:"template_version"`
	Stages   []*Stage                  `json:"stages,omitempty"`
}

// RemoveLinks gets the PipelineTemplate ready to be submitted to the GoCD API.
func (pt *PipelineTemplate) RemoveLinks() {
	pt.Links = nil
}

// Pipelines returns a list of Pipelines attached to this PipelineTemplate object.
func (pt *PipelineTemplate) Pipelines() []*Pipeline {
	return pt.Embedded.Pipelines
}

// Exists ensures whether or not a PipelineTemplate is present in the API.
func (pts *PipelineTemplatesService) Exists(ctx context.Context, name string) (bool, *APIResponse, error) {
	return pts.client.genericHeadAction(ctx, fmt.Sprintf("admin/templates/%s", name), apiV3)
}

// Get a single PipelineTemplate object in the GoCD API.
func (pts *PipelineTemplatesService) Get(ctx context.Context, name string) (*PipelineTemplate, *APIResponse, error) {
	pt := PipelineTemplate{}
	_, resp, err := pts.client.getAction(ctx, &APIClientRequest{
		Path:         "admin/templates/" + name,
		APIVersion:   apiV3,
		ResponseBody: &pt,
	})
	return &pt, resp, err
}

// List all PipelineTemplate objects in the GoCD API.
func (pts *PipelineTemplatesService) List(ctx context.Context) ([]*PipelineTemplate, *APIResponse, error) {
	ptr := PipelineTemplatesResponse{}

	_, resp, err := pts.client.getAction(ctx, &APIClientRequest{
		Path:         "admin/templates",
		APIVersion:   apiV3,
		ResponseBody: &ptr,
	})

	return ptr.Embedded.Templates, resp, err
}

// Create a new PipelineTemplate object in the GoCD API.
func (pts *PipelineTemplatesService) Create(ctx context.Context, name string, st []*Stage) (*PipelineTemplate, *APIResponse, error) {

	pt := PipelineTemplateRequest{
		Name:   name,
		Stages: st,
	}
	ptr := PipelineTemplate{}

	_, resp, err := pts.client.putAction(ctx, &APIClientRequest{
		Path:         "admin/templates",
		APIVersion:   apiV3,
		RequestBody:  pt,
		ResponseBody: &ptr,
	})

	etag := resp.HTTP.Header.Get("Etag")
	ptr.Version = strings.Replace(etag, "\"", "", -1)

	return &ptr, resp, err

}

// Update an PipelineTemplate object in the GoCD API.
func (pts *PipelineTemplatesService) Update(ctx context.Context, name string, version string, st []*Stage) (*PipelineTemplate, *APIResponse, error) {
	pt := &PipelineTemplateRequest{
		Name:   name,
		Stages: st,
	}
	ptr := &PipelineTemplate{}

	_, resp, err := pts.client.putAction(ctx, &APIClientRequest{
		Path:         "admin/templates/" + name,
		APIVersion:   apiV3,
		RequestBody:  pt,
		ResponseBody: &ptr,
	})

	return ptr, resp, err

}

// Delete a PipelineTemplate from the GoCD API.
func (pts *PipelineTemplatesService) Delete(ctx context.Context, name string) (string, *APIResponse, error) {
	return pts.client.deleteAction(ctx, fmt.Sprintf("admin/templates/%s", name), apiV3)
}
<commit_msg>GoCD API does not respond to head requests<commit_after>package gocd

import (
	"context"
	"fmt"
	"net/url"
	"strings"
)

// PipelineTemplatesService describes the HAL _link resource for the api response object for a pipeline configuration objects.
type PipelineTemplatesService service

// PipelineTemplatesLinks describes a single pipeline template config object HAL links
//go:generate gocd-response-links-generator -type=PipelineTemplatesLinks,PipelineTemplateLinks
type PipelineTemplatesLinks struct {
	Self *url.URL `json:"self"`
	Doc  *url.URL `json:"doc"`
	Find *url.URL `json:"find"`
}

// PipelineTemplateLinks describes multiple pipeline template config object HAL links
type PipelineTemplateLinks struct {
	Self *url.URL `json:"self"`
	Doc  *url.URL `json:"doc"`
	Find *url.URL `json:"find"`
}

// PipelineTemplateRequest describes a PipelineTemplate
type PipelineTemplateRequest struct {
	Name   string   `json:"name"`
	Stages []*Stage `json:"stages"`
}

// PipelineTemplateResponse describes an api response for a single pipeline templates
type PipelineTemplateResponse struct {
	Name     string `json:"name"`
	Embedded *struct {
		Pipelines []*struct {
			Name string `json:"name"`
		}
	} `json:"_embedded,omitempty"`
}

// PipelineTemplatesResponse describes an api response for multiple pipeline templates
type PipelineTemplatesResponse struct {
	Links    PipelineTemplatesLinks `json:"_links,omitempty"`
	Embedded *struct {
		Templates []*PipelineTemplate `json:"templates"`
	} `json:"_embedded,omitempty"`
}

type embeddedPipelineTemplate struct {
	Pipelines []*Pipeline `json:"pipelines,omitempty"`
}

// PipelineTemplate describes a response from the API for a pipeline template object.
type PipelineTemplate struct {
	Links    *PipelineTemplateLinks    `json:"_links,omitempty"`
	Name     string                    `json:"name"`
	Embedded *embeddedPipelineTemplate `json:"_embedded,omitempty"`
	Version  string                    `json:"template_version"`
	Stages   []*Stage                  `json:"stages,omitempty"`
}

// RemoveLinks gets the PipelineTemplate ready to be submitted to the GoCD API.
func (pt *PipelineTemplate) RemoveLinks() {
	pt.Links = nil
}

// Pipelines returns a list of Pipelines attached to this PipelineTemplate object.
func (pt *PipelineTemplate) Pipelines() []*Pipeline {
	return pt.Embedded.Pipelines
}

// Get a single PipelineTemplate object in the GoCD API.
func (pts *PipelineTemplatesService) Get(ctx context.Context, name string) (*PipelineTemplate, *APIResponse, error) {
	pt := PipelineTemplate{}
	_, resp, err := pts.client.getAction(ctx, &APIClientRequest{
		Path:         "admin/templates/" + name,
		APIVersion:   apiV3,
		ResponseBody: &pt,
	})
	return &pt, resp, err
}

// List all PipelineTemplate objects in the GoCD API.
func (pts *PipelineTemplatesService) List(ctx context.Context) ([]*PipelineTemplate, *APIResponse, error) {
	ptr := PipelineTemplatesResponse{}

	_, resp, err := pts.client.getAction(ctx, &APIClientRequest{
		Path:         "admin/templates",
		APIVersion:   apiV3,
		ResponseBody: &ptr,
	})

	return ptr.Embedded.Templates, resp, err
}

// Create a new PipelineTemplate object in the GoCD API.
func (pts *PipelineTemplatesService) Create(ctx context.Context, name string, st []*Stage) (*PipelineTemplate, *APIResponse, error) {

	pt := PipelineTemplateRequest{
		Name:   name,
		Stages: st,
	}
	ptr := PipelineTemplate{}

	_, resp, err := pts.client.putAction(ctx, &APIClientRequest{
		Path:         "admin/templates",
		APIVersion:   apiV3,
		RequestBody:  pt,
		ResponseBody: &ptr,
	})

	etag := resp.HTTP.Header.Get("Etag")
	ptr.Version = strings.Replace(etag, "\"", "", -1)

	return &ptr, resp, err

}

// Update an PipelineTemplate object in the GoCD API.
func (pts *PipelineTemplatesService) Update(ctx context.Context, name string, version string, st []*Stage) (*PipelineTemplate, *APIResponse, error) {
	pt := &PipelineTemplateRequest{
		Name:   name,
		Stages: st,
	}
	ptr := &PipelineTemplate{}

	_, resp, err := pts.client.putAction(ctx, &APIClientRequest{
		Path:         "admin/templates/" + name,
		APIVersion:   apiV3,
		RequestBody:  pt,
		ResponseBody: &ptr,
	})

	return ptr, resp, err

}

// Delete a PipelineTemplate from the GoCD API.
func (pts *PipelineTemplatesService) Delete(ctx context.Context, name string) (string, *APIResponse, error) {
	return pts.client.deleteAction(ctx, fmt.Sprintf("admin/templates/%s", name), apiV3)
}
<|endoftext|>


<commit_before>package progress

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/olekukonko/ts"
)

// ProgressMeter provides a progress bar type output for the TransferQueue. It
// is given an estimated file count and size up front and tracks the number of
// files and bytes transferred as well as the number of files and bytes that
// get skipped because the transfer is unnecessary.
type ProgressMeter struct {
	finishedFiles     int64 // int64s must come first for struct alignment
	skippedFiles      int64
	transferringFiles int64
	estimatedBytes    int64
	currentBytes      int64
	skippedBytes      int64
	started           int32
	estimatedFiles    int32
	startTime         time.Time
	finished          chan interface{}
	logger            *progressLogger
	fileIndex         map[string]int64 // Maps a file name to its transfer number
	fileIndexMutex    *sync.Mutex
	dryRun            bool
}

type env interface {
	Get(key string) (val string, ok bool)
}

type meterOption func(*ProgressMeter)

// DryRun is an option for NewMeter() that determines whether updates should be
// sent to stdout.
func DryRun(dryRun bool) meterOption {
	return func(m *ProgressMeter) {
		m.dryRun = dryRun
	}
}

// WithLogFile is an option for NewMeter() that sends updates to a text file.
func WithLogFile(name string) meterOption {
	printErr := func(err string) {
		fmt.Fprintf(os.Stderr, "Error creating progress logger: %s\n", err)
	}

	return func(m *ProgressMeter) {
		if len(name) == 0 {
			return
		}

		if !filepath.IsAbs(name) {
			printErr("GIT_LFS_PROGRESS must be an absolute path")
			return
		}

		cbDir := filepath.Dir(name)
		if err := os.MkdirAll(cbDir, 0755); err != nil {
			printErr(err.Error())
			return
		}

		file, err := os.OpenFile(name, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
		if err != nil {
			printErr(err.Error())
			return
		}

		m.logger.writeData = true
		m.logger.log = file
	}
}

// WithOSEnv is an option for NewMeter() that sends updates to the text file
// path specified in the OS Env.
func WithOSEnv(os env) meterOption {
	name, _ := os.Get("GIT_LFS_PROGRESS")
	return WithLogFile(name)
}

// NewMeter creates a new ProgressMeter.
func NewMeter(options ...meterOption) *ProgressMeter {
	m := &ProgressMeter{
		logger:         &progressLogger{},
		startTime:      time.Now(),
		fileIndex:      make(map[string]int64),
		fileIndexMutex: &sync.Mutex{},
		finished:       make(chan interface{}),
	}

	for _, opt := range options {
		opt(m)
	}

	return m
}

// Start begins sending status updates to the optional log file, and stdout.
func (p *ProgressMeter) Start() {
	if atomic.CompareAndSwapInt32(&p.started, 0, 1) {
		go p.writer()
	}
}

// Pause stops sending status updates temporarily, until Start() is called again.
func (p *ProgressMeter) Pause() {
	if atomic.CompareAndSwapInt32(&p.started, 1, 0) {
		p.finished <- true
	}
}

// Add tells the progress meter that a single file of the given size will
// possibly be transferred. If a file doesn't need to be transferred for some
// reason, be sure to call Skip(int64) with the same size.
func (p *ProgressMeter) Add(size int64) {
	atomic.AddInt32(&p.estimatedFiles, 1)
	atomic.AddInt64(&p.estimatedBytes, size)
}

// Skip tells the progress meter that a file of size `size` is being skipped
// because the transfer is unnecessary.
func (p *ProgressMeter) Skip(size int64) {
	atomic.AddInt64(&p.skippedFiles, 1)
	atomic.AddInt64(&p.skippedBytes, size)
	// Reduce bytes and files so progress easier to parse
	atomic.AddInt32(&p.estimatedFiles, -1)
	atomic.AddInt64(&p.estimatedBytes, -size)
}

// StartTransfer tells the progress meter that a transferring file is being
// added to the TransferQueue.
func (p *ProgressMeter) StartTransfer(name string) {
	idx := atomic.AddInt64(&p.transferringFiles, 1)
	p.fileIndexMutex.Lock()
	p.fileIndex[name] = idx
	p.fileIndexMutex.Unlock()
}

// TransferBytes increments the number of bytes transferred
func (p *ProgressMeter) TransferBytes(direction, name string, read, total int64, current int) {
	atomic.AddInt64(&p.currentBytes, int64(current))
	p.logBytes(direction, name, read, total)
}

// FinishTransfer increments the finished transfer count
func (p *ProgressMeter) FinishTransfer(name string) {
	atomic.AddInt64(&p.finishedFiles, 1)
	p.fileIndexMutex.Lock()
	delete(p.fileIndex, name)
	p.fileIndexMutex.Unlock()
}

// Finish shuts down the ProgressMeter
func (p *ProgressMeter) Finish() {
	close(p.finished)
	p.update()
	p.logger.Close()
	if !p.dryRun && p.estimatedBytes > 0 {
		fmt.Fprintf(os.Stdout, "\n")
	}
}

func (p *ProgressMeter) logBytes(direction, name string, read, total int64) {
	p.fileIndexMutex.Lock()
	idx := p.fileIndex[name]
	p.fileIndexMutex.Unlock()
	line := fmt.Sprintf("%s %d/%d %d/%d %s\n", direction, idx, p.estimatedFiles, read, total, name)
	if err := p.logger.Write([]byte(line)); err != nil {
		p.logger.Shutdown()
	}
}

func (p *ProgressMeter) writer() {
	p.update()
	for {
		select {
		case <-p.finished:
			return
		case <-time.After(time.Millisecond * 200):
			p.update()
		}
	}
}

func (p *ProgressMeter) update() {
	if p.dryRun || (p.estimatedFiles == 0 && p.skippedFiles == 0) {
		return
	}

	// (%d of %d files, %d skipped) %f B / %f B, %f B skipped
	// skipped counts only show when > 0

	out := fmt.Sprintf("\rGit LFS: (%d of %d files", p.finishedFiles, p.estimatedFiles)
	if p.skippedFiles > 0 {
		out += fmt.Sprintf(", %d skipped", p.skippedFiles)
	}
	out += fmt.Sprintf(") %s / %s", formatBytes(p.currentBytes), formatBytes(p.estimatedBytes))
	if p.skippedBytes > 0 {
		out += fmt.Sprintf(", %s skipped", formatBytes(p.skippedBytes))
	}

	fmt.Fprintf(os.Stdout, pad(out))
}

func formatBytes(i int64) string {
	switch {
	case i > 1099511627776:
		return fmt.Sprintf("%#0.2f TB", float64(i)/1099511627776)
	case i > 1073741824:
		return fmt.Sprintf("%#0.2f GB", float64(i)/1073741824)
	case i > 1048576:
		return fmt.Sprintf("%#0.2f MB", float64(i)/1048576)
	case i > 1024:
		return fmt.Sprintf("%#0.2f KB", float64(i)/1024)
	}

	return fmt.Sprintf("%d B", i)
}

const defaultWidth = 80

// pad pads the given message to occupy the entire maximum width of the terminal
// LFS is attached to. In doing so, this safeguards subsequent prints of shorter
// messages from leaving stray characters from the previous message on the
// screen by writing over them with whitespace padding.
func pad(msg string) string {
	width := defaultWidth
	size, err := ts.GetSize()
	if err == nil {
		// If `ts.GetSize()` was successful, set the width to the number
		// of columns present in the terminal LFS is attached to.
		// Otherwise, fall-back to `defaultWidth`.
		width = size.Col()
	}

	// Pad the string with whitespace so that printing at the start of the
	// line removes all traces from the last print.removes all traces from
	// the last print.
	padding := strings.Repeat(" ", maxInt(0, width-len(msg)))

	return msg + padding
}

// maxInt returns the greater of two `int`s, "a", or "b". This function
// originally comes from `github.com/git-lfs/git-lfs/tools#MaxInt`, but would
// introduce an import cycle if depended on directly.
func maxInt(a, b int) int {
	if a > b {
		return a
	}

	return b
}
<commit_msg>progress: make ProgressMeter implement log.Task<commit_after>package progress

import (
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"sync/atomic"
	"time"

	"github.com/git-lfs/git-lfs/git/githistory/log"
	"github.com/git-lfs/git-lfs/tools/humanize"
)

// ProgressMeter provides a progress bar type output for the TransferQueue. It
// is given an estimated file count and size up front and tracks the number of
// files and bytes transferred as well as the number of files and bytes that
// get skipped because the transfer is unnecessary.
type ProgressMeter struct {
	finishedFiles     int64 // int64s must come first for struct alignment
	skippedFiles      int64
	transferringFiles int64
	estimatedBytes    int64
	currentBytes      int64
	skippedBytes      int64
	estimatedFiles    int32
	paused            uint32
	logger            *progressLogger
	fileIndex         map[string]int64 // Maps a file name to its transfer number
	fileIndexMutex    *sync.Mutex
	dryRun            bool
	updates           chan *log.Update
}

type env interface {
	Get(key string) (val string, ok bool)
}

type meterOption func(*ProgressMeter)

// DryRun is an option for NewMeter() that determines whether updates should be
// sent to stdout.
func DryRun(dryRun bool) meterOption {
	return func(m *ProgressMeter) {
		m.dryRun = dryRun
	}
}

// WithLogFile is an option for NewMeter() that sends updates to a text file.
func WithLogFile(name string) meterOption {
	printErr := func(err string) {
		fmt.Fprintf(os.Stderr, "Error creating progress logger: %s\n", err)
	}

	return func(m *ProgressMeter) {
		if len(name) == 0 {
			return
		}

		if !filepath.IsAbs(name) {
			printErr("GIT_LFS_PROGRESS must be an absolute path")
			return
		}

		cbDir := filepath.Dir(name)
		if err := os.MkdirAll(cbDir, 0755); err != nil {
			printErr(err.Error())
			return
		}

		file, err := os.OpenFile(name, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
		if err != nil {
			printErr(err.Error())
			return
		}

		m.logger.writeData = true
		m.logger.log = file
	}
}

// WithOSEnv is an option for NewMeter() that sends updates to the text file
// path specified in the OS Env.
func WithOSEnv(os env) meterOption {
	name, _ := os.Get("GIT_LFS_PROGRESS")
	return WithLogFile(name)
}

// NewMeter creates a new ProgressMeter.
func NewMeter(options ...meterOption) *ProgressMeter {
	m := &ProgressMeter{
		logger:         &progressLogger{},
		fileIndex:      make(map[string]int64),
		fileIndexMutex: &sync.Mutex{},
		updates:        make(chan *log.Update),
	}

	for _, opt := range options {
		opt(m)
	}

	return m
}

// Start begins sending status updates to the optional log file, and stdout.
func (p *ProgressMeter) Start() {
	atomic.StoreUint32(&p.paused, 0)
}

// Pause stops sending status updates temporarily, until Start() is called again.
func (p *ProgressMeter) Pause() {
	atomic.StoreUint32(&p.paused, 1)
}

// Add tells the progress meter that a single file of the given size will
// possibly be transferred. If a file doesn't need to be transferred for some
// reason, be sure to call Skip(int64) with the same size.
func (p *ProgressMeter) Add(size int64) {
	defer p.update()
	atomic.AddInt32(&p.estimatedFiles, 1)
	atomic.AddInt64(&p.estimatedBytes, size)
}

// Skip tells the progress meter that a file of size `size` is being skipped
// because the transfer is unnecessary.
func (p *ProgressMeter) Skip(size int64) {
	defer p.update()
	atomic.AddInt64(&p.skippedFiles, 1)
	atomic.AddInt64(&p.skippedBytes, size)
	// Reduce bytes and files so progress easier to parse
	atomic.AddInt32(&p.estimatedFiles, -1)
	atomic.AddInt64(&p.estimatedBytes, -size)
}

// StartTransfer tells the progress meter that a transferring file is being
// added to the TransferQueue.
func (p *ProgressMeter) StartTransfer(name string) {
	defer p.update()
	idx := atomic.AddInt64(&p.transferringFiles, 1)
	p.fileIndexMutex.Lock()
	p.fileIndex[name] = idx
	p.fileIndexMutex.Unlock()
}

// TransferBytes increments the number of bytes transferred
func (p *ProgressMeter) TransferBytes(direction, name string, read, total int64, current int) {
	defer p.update()
	atomic.AddInt64(&p.currentBytes, int64(current))
	p.logBytes(direction, name, read, total)
}

// FinishTransfer increments the finished transfer count
func (p *ProgressMeter) FinishTransfer(name string) {
	defer p.update()
	atomic.AddInt64(&p.finishedFiles, 1)
	p.fileIndexMutex.Lock()
	delete(p.fileIndex, name)
	p.fileIndexMutex.Unlock()
}

// Finish shuts down the ProgressMeter
func (p *ProgressMeter) Finish() {
	p.update()
	close(p.updates)
}

func (p *ProgressMeter) Updates() <-chan *log.Update {
	return p.updates
}

func (p *ProgressMeter) Throttled() bool {
	return true
}

func (p *ProgressMeter) update() {
	if p.skipUpdate() {
		return
	}

	p.updates <- &log.Update{
		S:  p.str(),
		At: time.Now(),
	}
}

func (p *ProgressMeter) skipUpdate() bool {
	return p.dryRun ||
		(p.estimatedFiles == 0 && p.skippedFiles == 0) ||
		atomic.LoadUint32(&p.paused) == 1
}

func (p *ProgressMeter) str() string {
	// (%d of %d files, %d skipped) %f B / %f B, %f B skipped
	// skipped counts only show when > 0

	out := fmt.Sprintf("\rGit LFS: (%d of %d files",
		p.finishedFiles,
		p.estimatedFiles)
	if p.skippedFiles > 0 {
		out += fmt.Sprintf(", %d skipped", p.skippedFiles)
	}
	out += fmt.Sprintf(") %s / %s",
		humanize.FormatBytes(uint64(p.currentBytes)),
		humanize.FormatBytes(uint64(p.estimatedBytes)))
	if p.skippedBytes > 0 {
		out += fmt.Sprintf(", %s skipped",
			humanize.FormatBytes(uint64(p.skippedBytes)))
	}

	return out
}

func (p *ProgressMeter) logBytes(direction, name string, read, total int64) {
	p.fileIndexMutex.Lock()
	idx := p.fileIndex[name]
	p.fileIndexMutex.Unlock()
	line := fmt.Sprintf("%s %d/%d %d/%d %s\n", direction, idx, p.estimatedFiles, read, total, name)
	if err := p.logger.Write([]byte(line)); err != nil {
		p.logger.Shutdown()
	}
}
<|endoftext|>


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

package s2

import (
	"math"
	"math/rand"
	"reflect"
	"testing"
)

func TestCovererRandomCells(t *testing.T) {
	rc := &RegionCoverer{MinLevel: 0, MaxLevel: 30, LevelMod: 1, MaxCells: 1}

	// Test random cell ids at all levels.
	for i := 0; i < 10000; i++ {
		id := CellID(randomUint64())
		for !id.IsValid() {
			id = CellID(randomUint64())
		}
		covering := rc.Covering(Region(CellFromCellID(id)))
		if len(covering) != 1 {
			t.Errorf("Iteration %d, cell ID token %s, got covering size = %d, want covering size = 1", i, id.ToToken(), len(covering))
		}
		if (covering)[0] != id {
			t.Errorf("Iteration %d, cell ID token %s, got covering = %v, want covering = %v", i, id.ToToken(), covering, id)
		}
	}
}

// checkCovering reports whether covering is a valid cover for the region.
func checkCovering(t *testing.T, rc *RegionCoverer, r Region, covering CellUnion, interior bool) {
	// Keep track of how many cells have the same rc.MinLevel ancestor.
	minLevelCells := map[CellID]int{}
	var tempCover CellUnion
	for _, ci := range covering {
		level := ci.Level()
		if level < rc.MinLevel {
			t.Errorf("CellID(%s).Level() = %d, want >= %d", ci.ToToken(), level, rc.MinLevel)
		}
		if level > rc.MaxLevel {
			t.Errorf("CellID(%s).Level() = %d, want <= %d", ci.ToToken(), level, rc.MaxLevel)
		}
		if rem := (level - rc.MinLevel) % rc.LevelMod; rem != 0 {
			t.Errorf("(CellID(%s).Level() - MinLevel) mod LevelMod = %d, want = %d", ci.ToToken(), rem, 0)
		}
		tempCover = append(tempCover, ci)
		minLevelCells[ci.Parent(rc.MinLevel)]++
	}
	if len(covering) > rc.MaxCells {
		// If the covering has more than the requested number of cells, then check
		// that the cell count cannot be reduced by using the parent of some cell.
		for ci, count := range minLevelCells {
			if count > 1 {
				t.Errorf("Min level CellID %s, count = %d, want = %d", ci.ToToken(), count, 1)
			}
		}
	}
	if interior {
		for _, ci := range covering {
			if !r.ContainsCell(CellFromCellID(ci)) {
				t.Errorf("Region(%v).ContainsCell(%v) = %t, want = %t", r, CellFromCellID(ci), false, true)
			}
		}
	} else {
		tempCover.Normalize()
		checkCoveringTight(t, r, tempCover, true, 0)
	}
}

// checkCoveringTight checks that "cover" completely covers the given region.
// If "checkTight" is true, also checks that it does not contain any cells that
// do not intersect the given region. ("id" is only used internally.)
func checkCoveringTight(t *testing.T, r Region, cover CellUnion, checkTight bool, id CellID) {
	if !id.IsValid() {
		for f := 0; f < 6; f++ {
			checkCoveringTight(t, r, cover, checkTight, CellIDFromFace(f))
		}
		return
	}

	if !r.IntersectsCell(CellFromCellID(id)) {
		// If region does not intersect id, then neither should the covering.
		if got := cover.IntersectsCellID(id); checkTight && got {
			t.Errorf("CellUnion(%v).IntersectsCellID(%s) = %t; want = %t", cover, id.ToToken(), got, false)
		}
	} else if !cover.ContainsCellID(id) {
		// The region may intersect id, but we can't assert that the covering
		// intersects id because we may discover that the region does not actually
		// intersect upon further subdivision.  (IntersectsCell is not exact.)
		if got := r.ContainsCell(CellFromCellID(id)); got {
			t.Errorf("Region(%v).ContainsCell(%v) = %t; want = %t", r, CellFromCellID(id), got, false)
		}
		if got := id.IsLeaf(); got {
			t.Errorf("CellID(%s).IsLeaf() = %t; want = %t", id.ToToken(), got, false)
		}

		for child := id.ChildBegin(); child != id.ChildEnd(); child = child.Next() {
			checkCoveringTight(t, r, cover, checkTight, child)
		}
	}
}

func TestCovererRandomCaps(t *testing.T) {
	rc := &RegionCoverer{}
	for i := 0; i < 1000; i++ {
		rc.MinLevel = int(rand.Int31n(maxLevel + 1))
		rc.MaxLevel = int(rand.Int31n(maxLevel + 1))
		for rc.MinLevel > rc.MaxLevel {
			rc.MinLevel = int(rand.Int31n(maxLevel + 1))
			rc.MaxLevel = int(rand.Int31n(maxLevel + 1))
		}
		rc.LevelMod = int(1 + rand.Int31n(3))
		rc.MaxCells = skewedInt(10)

		maxArea := math.Min(4*math.Pi, float64(3*rc.MaxCells+1)*AvgAreaMetric.Value(rc.MinLevel))
		r := Region(randomCap(0.1*AvgAreaMetric.Value(maxLevel), maxArea))

		covering := rc.Covering(r)
		checkCovering(t, rc, r, covering, false)
		interior := rc.InteriorCovering(r)
		checkCovering(t, rc, r, interior, true)

		// Check that Covering is deterministic.
		covering2 := rc.Covering(r)
		if !reflect.DeepEqual(covering, covering2) {
			t.Errorf("Iteration %d, got covering = %v, want covering = %v", i, covering2, covering)
		}

		// Also check Denormalize. The denormalized covering
		// may still be different and smaller than "covering" because
		// s2.RegionCoverer does not guarantee that it will not output all four
		// children of the same parent.
		covering.Denormalize(rc.MinLevel, rc.LevelMod)
		checkCovering(t, rc, r, covering, false)
	}
}

func TestRegionCovererInteriorCovering(t *testing.T) {
	// We construct the region the following way. Start with Cell of level l.
	// Remove from it one of its grandchildren (level l+2). If we then set
	//   minLevel < l + 1
	//   maxLevel > l + 2
	//   maxCells = 3
	// the best interior covering should contain 3 children of the initial cell,
	// that were not effected by removal of a grandchild.
	const level = 12
	smallCell := cellIDFromPoint(randomPoint()).Parent(level + 2)
	largeCell := smallCell.Parent(level)

	smallCellUnion := CellUnion([]CellID{smallCell})
	largeCellUnion := CellUnion([]CellID{largeCell})
	diff := CellUnionFromDifference(largeCellUnion, smallCellUnion)

	coverer := &RegionCoverer{
		MaxCells: 3,
		MaxLevel: level + 3,
		MinLevel: level,
	}

	interior := coverer.InteriorCovering(&diff)
	if len(interior) != 3 {
		t.Fatalf("len(coverer.InteriorCovering(%v)) = %v, want 3", diff, len(interior))
	}
	for i := 0; i < 3; i++ {
		if got, want := interior[i].Level(), level+1; got != want {
			t.Errorf("interior[%d].Level() = %v, want %v", i, got, want)
		}
	}
}

func TestRegionCovererSimpleRegionCovering(t *testing.T) {
	const maxLevel = maxLevel
	for i := 0; i < 100; i++ {
		level := randomUniformInt(maxLevel + 1)
		maxArea := math.Min(4*math.Pi, 1000.0*AvgAreaMetric.Value(level))
		c := randomCap(0.1*AvgAreaMetric.Value(maxLevel), maxArea)
		covering := SimpleRegionCovering(c, c.Center(), level)
		rc := &RegionCoverer{MaxLevel: level, MinLevel: level, MaxCells: math.MaxInt32, LevelMod: 1}
		checkCovering(t, rc, c, covering, false)
	}
}

// TODO(roberts): Differences from C++
//  func TestRegionCovererAccuracy(t *testing.T) {
//  func TestRegionCovererFastCoveringHugeFixedLevelCovering(t *testing.T) {
//  func TestRegionCovererIsCanonicalInvalidS2CellId(t *testing.T) {
//  func TestRegionCovererIsCanonicalUnsorted(t *testing.T) {
//  func TestRegionCovererIsCanonicalOverlapping(t *testing.T) {
//  func TestRegionCovererIsCanonicalMinLevel(t *testing.T) {
//  func TestRegionCovererIsCanonicalMaxLevel(t *testing.T) {
//  func TestRegionCovererIsCanonicalLevelMod(t *testing.T) {
//  func TestRegionCovererIsCanonicalMaxCells(t *testing.T) {
//  func TestRegionCovererIsCanonicalNormalized(t *testing.T) {
//  func TestRegionCovererCanonicalizeCoveringUnsortedDuplicateCells(t *testing.T) {
//  func TestRegionCovererCanonicalizeCoveringMaxLevelExceeded(t *testing.T) {
//  func TestRegionCovererCanonicalizeCoveringWrongLevelMod(t *testing.T) {
//  func TestRegionCovererCanonicalizeCoveringReplacedByParent(t *testing.T) {
//  func TestRegionCovererCanonicalizeCoveringDenormalizedCellUnion(t *testing.T) {
//  func TestRegionCovererCanonicalizeCoveringMaxCellsMergesSmallest(t *testing.T) {
//  func TestRegionCovererCanonicalizeCoveringMaxCellsMergesRepeatedly(t *testing.T) {
<commit_msg>s2: Add Benchmarks for RegionCoverer.<commit_after>// Copyright 2015 Google Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package s2

import (
	"fmt"
	"math"
	"math/rand"
	"reflect"
	"testing"
)

func TestCovererRandomCells(t *testing.T) {
	rc := &RegionCoverer{MinLevel: 0, MaxLevel: 30, LevelMod: 1, MaxCells: 1}

	// Test random cell ids at all levels.
	for i := 0; i < 10000; i++ {
		id := CellID(randomUint64())
		for !id.IsValid() {
			id = CellID(randomUint64())
		}
		covering := rc.Covering(Region(CellFromCellID(id)))
		if len(covering) != 1 {
			t.Errorf("Iteration %d, cell ID token %s, got covering size = %d, want covering size = 1", i, id.ToToken(), len(covering))
		}
		if (covering)[0] != id {
			t.Errorf("Iteration %d, cell ID token %s, got covering = %v, want covering = %v", i, id.ToToken(), covering, id)
		}
	}
}

// checkCovering reports whether covering is a valid cover for the region.
func checkCovering(t *testing.T, rc *RegionCoverer, r Region, covering CellUnion, interior bool) {
	// Keep track of how many cells have the same rc.MinLevel ancestor.
	minLevelCells := map[CellID]int{}
	var tempCover CellUnion
	for _, ci := range covering {
		level := ci.Level()
		if level < rc.MinLevel {
			t.Errorf("CellID(%s).Level() = %d, want >= %d", ci.ToToken(), level, rc.MinLevel)
		}
		if level > rc.MaxLevel {
			t.Errorf("CellID(%s).Level() = %d, want <= %d", ci.ToToken(), level, rc.MaxLevel)
		}
		if rem := (level - rc.MinLevel) % rc.LevelMod; rem != 0 {
			t.Errorf("(CellID(%s).Level() - MinLevel) mod LevelMod = %d, want = %d", ci.ToToken(), rem, 0)
		}
		tempCover = append(tempCover, ci)
		minLevelCells[ci.Parent(rc.MinLevel)]++
	}
	if len(covering) > rc.MaxCells {
		// If the covering has more than the requested number of cells, then check
		// that the cell count cannot be reduced by using the parent of some cell.
		for ci, count := range minLevelCells {
			if count > 1 {
				t.Errorf("Min level CellID %s, count = %d, want = %d", ci.ToToken(), count, 1)
			}
		}
	}
	if interior {
		for _, ci := range covering {
			if !r.ContainsCell(CellFromCellID(ci)) {
				t.Errorf("Region(%v).ContainsCell(%v) = %t, want = %t", r, CellFromCellID(ci), false, true)
			}
		}
	} else {
		tempCover.Normalize()
		checkCoveringTight(t, r, tempCover, true, 0)
	}
}

// checkCoveringTight checks that "cover" completely covers the given region.
// If "checkTight" is true, also checks that it does not contain any cells that
// do not intersect the given region. ("id" is only used internally.)
func checkCoveringTight(t *testing.T, r Region, cover CellUnion, checkTight bool, id CellID) {
	if !id.IsValid() {
		for f := 0; f < 6; f++ {
			checkCoveringTight(t, r, cover, checkTight, CellIDFromFace(f))
		}
		return
	}

	if !r.IntersectsCell(CellFromCellID(id)) {
		// If region does not intersect id, then neither should the covering.
		if got := cover.IntersectsCellID(id); checkTight && got {
			t.Errorf("CellUnion(%v).IntersectsCellID(%s) = %t; want = %t", cover, id.ToToken(), got, false)
		}
	} else if !cover.ContainsCellID(id) {
		// The region may intersect id, but we can't assert that the covering
		// intersects id because we may discover that the region does not actually
		// intersect upon further subdivision.  (IntersectsCell is not exact.)
		if got := r.ContainsCell(CellFromCellID(id)); got {
			t.Errorf("Region(%v).ContainsCell(%v) = %t; want = %t", r, CellFromCellID(id), got, false)
		}
		if got := id.IsLeaf(); got {
			t.Errorf("CellID(%s).IsLeaf() = %t; want = %t", id.ToToken(), got, false)
		}

		for child := id.ChildBegin(); child != id.ChildEnd(); child = child.Next() {
			checkCoveringTight(t, r, cover, checkTight, child)
		}
	}
}

func TestCovererRandomCaps(t *testing.T) {
	rc := &RegionCoverer{}
	for i := 0; i < 1000; i++ {
		rc.MinLevel = int(rand.Int31n(maxLevel + 1))
		rc.MaxLevel = int(rand.Int31n(maxLevel + 1))
		for rc.MinLevel > rc.MaxLevel {
			rc.MinLevel = int(rand.Int31n(maxLevel + 1))
			rc.MaxLevel = int(rand.Int31n(maxLevel + 1))
		}
		rc.LevelMod = int(1 + rand.Int31n(3))
		rc.MaxCells = skewedInt(10)

		maxArea := math.Min(4*math.Pi, float64(3*rc.MaxCells+1)*AvgAreaMetric.Value(rc.MinLevel))
		r := Region(randomCap(0.1*AvgAreaMetric.Value(maxLevel), maxArea))

		covering := rc.Covering(r)
		checkCovering(t, rc, r, covering, false)
		interior := rc.InteriorCovering(r)
		checkCovering(t, rc, r, interior, true)

		// Check that Covering is deterministic.
		covering2 := rc.Covering(r)
		if !reflect.DeepEqual(covering, covering2) {
			t.Errorf("Iteration %d, got covering = %v, want covering = %v", i, covering2, covering)
		}

		// Also check Denormalize. The denormalized covering
		// may still be different and smaller than "covering" because
		// s2.RegionCoverer does not guarantee that it will not output all four
		// children of the same parent.
		covering.Denormalize(rc.MinLevel, rc.LevelMod)
		checkCovering(t, rc, r, covering, false)
	}
}

func TestRegionCovererInteriorCovering(t *testing.T) {
	// We construct the region the following way. Start with Cell of level l.
	// Remove from it one of its grandchildren (level l+2). If we then set
	//   minLevel < l + 1
	//   maxLevel > l + 2
	//   maxCells = 3
	// the best interior covering should contain 3 children of the initial cell,
	// that were not effected by removal of a grandchild.
	const level = 12
	smallCell := cellIDFromPoint(randomPoint()).Parent(level + 2)
	largeCell := smallCell.Parent(level)

	smallCellUnion := CellUnion([]CellID{smallCell})
	largeCellUnion := CellUnion([]CellID{largeCell})
	diff := CellUnionFromDifference(largeCellUnion, smallCellUnion)

	coverer := &RegionCoverer{
		MaxCells: 3,
		MaxLevel: level + 3,
		MinLevel: level,
	}

	interior := coverer.InteriorCovering(&diff)
	if len(interior) != 3 {
		t.Fatalf("len(coverer.InteriorCovering(%v)) = %v, want 3", diff, len(interior))
	}
	for i := 0; i < 3; i++ {
		if got, want := interior[i].Level(), level+1; got != want {
			t.Errorf("interior[%d].Level() = %v, want %v", i, got, want)
		}
	}
}

func TestRegionCovererSimpleRegionCovering(t *testing.T) {
	const maxLevel = maxLevel
	for i := 0; i < 100; i++ {
		level := randomUniformInt(maxLevel + 1)
		maxArea := math.Min(4*math.Pi, 1000.0*AvgAreaMetric.Value(level))
		c := randomCap(0.1*AvgAreaMetric.Value(maxLevel), maxArea)
		covering := SimpleRegionCovering(c, c.Center(), level)
		rc := &RegionCoverer{MaxLevel: level, MinLevel: level, MaxCells: math.MaxInt32, LevelMod: 1}
		checkCovering(t, rc, c, covering, false)
	}
}

const numCoveringBMRegions = 1000

func BenchmarkRegionCovererCoveringCap(b *testing.B) {
	benchmarkRegionCovererCovering(b, func(n int) string {
		return fmt.Sprintf("Cap%d", n)
	},
		func(n int) []Region {
			regions := make([]Region, numCoveringBMRegions)
			for i := 0; i < numCoveringBMRegions; i++ {
				regions[i] = randomCap(0.1*AvgAreaMetric.Value(maxLevel), 4*math.Pi)
			}
			return regions
		})
}

func BenchmarkRegionCovererCoveringCell(b *testing.B) {
	benchmarkRegionCovererCovering(b, func(n int) string {
		return fmt.Sprintf("Cell%d", n)
	},
		func(n int) []Region {
			regions := make([]Region, numCoveringBMRegions)
			for i := 0; i < numCoveringBMRegions; i++ {
				regions[i] = CellFromCellID(randomCellIDForLevel(maxLevel - randomUniformInt(n)))
			}
			return regions
		})
}

func BenchmarkRegionCovererCoveringLoop(b *testing.B) {
	benchmarkRegionCovererCovering(b, func(n int) string {
		return fmt.Sprintf("Loop-%d-edges", int(math.Pow(2.0, float64(n))))
	},
		func(n int) []Region {
			size := int(math.Pow(2.0, float64(n)))
			regions := make([]Region, numCoveringBMRegions)
			for i := 0; i < numCoveringBMRegions; i++ {
				regions[i] = RegularLoop(randomPoint(), kmToAngle(10.0), size)
			}
			return regions
		})
}

func BenchmarkRegionCovererCoveringCellUnion(b *testing.B) {
	benchmarkRegionCovererCovering(b, func(n int) string {
		return fmt.Sprintf("CellUnion-%d-cells", int(math.Pow(2.0, float64(n))))
	},
		func(n int) []Region {
			size := int(math.Pow(2.0, float64(n)))
			regions := make([]Region, numCoveringBMRegions)
			for i := 0; i < numCoveringBMRegions; i++ {
				cu := randomCellUnion(size)
				regions[i] = &cu
			}
			return regions
		})
}

// TODO(roberts): Add more benchmarking that changes the values in the coverer (min/max level, # cells).

// benchmark Covering using the supplied func to generate a slice of random Regions of
// the given type to choose from for the benchmark.
//
// e.g. Loops with [4..~2^n] edges, CellUnions of 2^n random Cells, random Cells and Caps
func benchmarkRegionCovererCovering(b *testing.B, genLabel func(n int) string, genRegions func(n int) []Region) {
	rc := &RegionCoverer{MinLevel: 0, MaxLevel: 30, LevelMod: 1, MaxCells: 8}

	// Range over a variety of region complexities.
	for n := 2; n <= 16; n++ {
		b.Run(genLabel(n),
			func(b *testing.B) {
				b.StopTimer()
				regions := genRegions(n)
				l := len(regions)
				b.StartTimer()
				for i := 0; i < b.N; i++ {
					rc.Covering(regions[i%l])
				}
			})
	}
}

// TODO(roberts): Differences from C++
//  func TestRegionCovererAccuracy(t *testing.T) {
//  func TestRegionCovererFastCoveringHugeFixedLevelCovering(t *testing.T) {
//  func TestRegionCovererIsCanonicalInvalidS2CellId(t *testing.T) {
//  func TestRegionCovererIsCanonicalUnsorted(t *testing.T) {
//  func TestRegionCovererIsCanonicalOverlapping(t *testing.T) {
//  func TestRegionCovererIsCanonicalMinLevel(t *testing.T) {
//  func TestRegionCovererIsCanonicalMaxLevel(t *testing.T) {
//  func TestRegionCovererIsCanonicalLevelMod(t *testing.T) {
//  func TestRegionCovererIsCanonicalMaxCells(t *testing.T) {
//  func TestRegionCovererIsCanonicalNormalized(t *testing.T) {
//  func TestRegionCovererCanonicalizeCoveringUnsortedDuplicateCells(t *testing.T) {
//  func TestRegionCovererCanonicalizeCoveringMaxLevelExceeded(t *testing.T) {
//  func TestRegionCovererCanonicalizeCoveringWrongLevelMod(t *testing.T) {
//  func TestRegionCovererCanonicalizeCoveringReplacedByParent(t *testing.T) {
//  func TestRegionCovererCanonicalizeCoveringDenormalizedCellUnion(t *testing.T) {
//  func TestRegionCovererCanonicalizeCoveringMaxCellsMergesSmallest(t *testing.T) {
//  func TestRegionCovererCanonicalizeCoveringMaxCellsMergesRepeatedly(t *testing.T) {
<|endoftext|>