<commit_before>package date

import (
	"database/sql/driver"
	"errors"
	"time"
)

const (
	MillisecondsInSecond = 1000
	MicrosecondsInSecond = 1000 * MillisecondsInSecond
	NanosecondsInSecond  = 1000 * MicrosecondsInSecond
)

// Date is a plain date, without time or timezone info (use time.Time for those!)
type Date struct {
	Year     int
	Month    time.Month
	Day      int
	Location *time.Location
}

// Create a Date from a time.Time object
func From(t time.Time) Date {
	y, m, d := t.Date()
	return Date{y, m, d, t.Location()}
}

func (d Date) PrevDay() Date {
	if d.Day > 1 {
		return Date{d.Year, d.Month, d.Day - 1, d.Location}
	} else if d.Month > 1 {
		return Date{d.Year, d.Month - 1, DaysInMonth(d.Month-1, d.Year), d.Location}
	} else {
		return Date{d.Year - 1, time.December, 31, d.Location}
	}
}

func (d Date) NextDay() Date {
	if d.Day < DaysInMonth(d.Month, d.Year) {
		return Date{d.Year, d.Month, d.Day + 1, d.Location}
	} else if d.Month < time.December {
		return Date{d.Year, d.Month + 1, 1, d.Location}
	} else {
		return Date{d.Year + 1, time.January, 1, d.Location}
	}
}

func (d Date) Before(other Date) bool {
	return (d.Year < other.Year ||
		(d.Year == other.Year &&
			(d.Month < other.Month ||
				(d.Month == other.Month && d.Day < other.Day))))
}

func (d Date) BeginningOfDay() time.Time {
	return time.Date(d.Year, d.Month, d.Day, 0, 0, 0, 0, d.Location)
}

func (d Date) BeginningOfDayIn(timezone *time.Location) time.Time {
	return time.Date(d.Year, d.Month, d.Day, 0, 0, 0, 0, timezone)
}

func (d Date) EndOfDay() time.Time {
	return time.Date(d.Year, d.Month, d.Day, 23, 59, 59, NanosecondsInSecond-1, d.Location)
}

func (d Date) EndOfDayIn(timezone *time.Location) time.Time {
	return time.Date(d.Year, d.Month, d.Day, 23, 59, 59, NanosecondsInSecond-1, timezone)
}

func (d Date) String() string {
	return d.BeginningOfDayIn(time.UTC).Format(`"2006-01-02"`)
}

// Implement sql.Scanner interface
func (d Date) Scan(src interface{}) error {
	t, ok := src.(time.Time)
	if !ok {
		return errors.New("unstable/date: scan value was not a Time")
	}

	d = From(t)
	return nil
}

// Implement sql.driver.Valuer interface
func (d Date) Value() (driver.Value, error) {
	return d.BeginningOfDayIn(time.UTC), nil
}

// Implement json.Marshaler interface
func (d Date) MarshalJSON() ([]byte, error) {
	return []byte(d.String()), nil
}

func IsLeapYear(year int) bool {
	return year%4 == 0 && (year%100 != 0 || year%400 == 0)
}

var DaysInNonLeapMonth = [12]int32{31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31}

func DaysInMonth(m time.Month, year int) int {
	if m == time.February && IsLeapYear(year) {
		return 29
	} else {
		return int(DaysInNonLeapMonth[m-1])
	}
}
<commit_msg>date: fix string function<commit_after>package date

import (
	"database/sql/driver"
	"errors"
	"time"
)

const (
	MillisecondsInSecond = 1000
	MicrosecondsInSecond = 1000 * MillisecondsInSecond
	NanosecondsInSecond  = 1000 * MicrosecondsInSecond
)

// Date is a plain date, without time or timezone info (use time.Time for those!)
type Date struct {
	Year     int
	Month    time.Month
	Day      int
	Location *time.Location
}

// Create a Date from a time.Time object
func From(t time.Time) Date {
	y, m, d := t.Date()
	return Date{y, m, d, t.Location()}
}

func (d Date) PrevDay() Date {
	if d.Day > 1 {
		return Date{d.Year, d.Month, d.Day - 1, d.Location}
	} else if d.Month > 1 {
		return Date{d.Year, d.Month - 1, DaysInMonth(d.Month-1, d.Year), d.Location}
	} else {
		return Date{d.Year - 1, time.December, 31, d.Location}
	}
}

func (d Date) NextDay() Date {
	if d.Day < DaysInMonth(d.Month, d.Year) {
		return Date{d.Year, d.Month, d.Day + 1, d.Location}
	} else if d.Month < time.December {
		return Date{d.Year, d.Month + 1, 1, d.Location}
	} else {
		return Date{d.Year + 1, time.January, 1, d.Location}
	}
}

func (d Date) Before(other Date) bool {
	return (d.Year < other.Year ||
		(d.Year == other.Year &&
			(d.Month < other.Month ||
				(d.Month == other.Month && d.Day < other.Day))))
}

func (d Date) BeginningOfDay() time.Time {
	return time.Date(d.Year, d.Month, d.Day, 0, 0, 0, 0, d.Location)
}

func (d Date) BeginningOfDayIn(timezone *time.Location) time.Time {
	return time.Date(d.Year, d.Month, d.Day, 0, 0, 0, 0, timezone)
}

func (d Date) EndOfDay() time.Time {
	return time.Date(d.Year, d.Month, d.Day, 23, 59, 59, NanosecondsInSecond-1, d.Location)
}

func (d Date) EndOfDayIn(timezone *time.Location) time.Time {
	return time.Date(d.Year, d.Month, d.Day, 23, 59, 59, NanosecondsInSecond-1, timezone)
}

func (d Date) String() string {
	return d.BeginningOfDayIn(time.UTC).Format("2006-01-02")
}

// Implement sql.Scanner interface
func (d Date) Scan(src interface{}) error {
	t, ok := src.(time.Time)
	if !ok {
		return errors.New("unstable/date: scan value was not a Time")
	}

	d = From(t)
	return nil
}

// Implement sql.driver.Valuer interface
func (d Date) Value() (driver.Value, error) {
	return d.BeginningOfDayIn(time.UTC), nil
}

// Implement json.Marshaler interface
func (d Date) MarshalJSON() ([]byte, error) {
	return []byte(d.BeginningOfDayIn(time.UTC).Format(`"2006-01-02"`)), nil
}

func IsLeapYear(year int) bool {
	return year%4 == 0 && (year%100 != 0 || year%400 == 0)
}

var DaysInNonLeapMonth = [12]int32{31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31}

func DaysInMonth(m time.Month, year int) int {
	if m == time.February && IsLeapYear(year) {
		return 29
	} else {
		return int(DaysInNonLeapMonth[m-1])
	}
}
<|endoftext|>


<commit_before>package cron

import (
	"fmt"
	"github.com/open-falcon/alarm/g"
	"github.com/open-falcon/common/model"
	"github.com/open-falcon/common/utils"
)

func BuildCommonSMSContent(event *model.Event) string {
	return fmt.Sprintf(
		"[P%d][%s][%s][][%s %s %s %s %s%s%s][O%d %s]",
		event.Priority(),
		event.Status,
		event.Endpoint,
		event.Note(),
		event.Func(),
		event.Metric(),
		utils.SortedTags(event.PushedTags),
		utils.ReadableFloat(event.LeftValue),
		event.Operator(),
		utils.ReadableFloat(event.RightValue()),
		event.CurrentStep,
		event.FormattedTime(),
	)
}

func BuildCommonMailContent(event *model.Event) string {
	link := g.Link(event)
	return fmt.Sprintf(
		"%s\r\nP%d\r\nEndpoint:%s\r\nMetric:%s\r\nTags:%s\r\n%s: %s%s%s\r\nNote:%s\r\nMax:%d, Current:%d\r\nTimestamp:%s\r\n%s\r\n",
		event.Status,
		event.Priority(),
		event.Endpoint,
		event.Metric(),
		utils.SortedTags(event.PushedTags),
		event.Func(),
		utils.ReadableFloat(event.LeftValue),
		event.Operator(),
		utils.ReadableFloat(event.RightValue()),
		event.Note(),
		event.MaxStep(),
		event.CurrentStep,
		event.FormattedTime(),
		link,
	)
}

func GenerateSmsContent(event *model.Event) string {
	return BuildCommonSMSContent(event)
}

func GenerateMailContent(event *model.Event) string {
	return BuildCommonMailContent(event)
}
<commit_msg>Change mail content of html table.<commit_after>package cron

import (
	"fmt"
	"github.com/open-falcon/alarm/g"
	"github.com/open-falcon/common/model"
	"github.com/open-falcon/common/utils"
)

func BuildCommonSMSContent(event *model.Event) string {
	return fmt.Sprintf(
		"[P%d][%s][%s][][%s %s %s %s %s%s%s][O%d %s]",
		event.Priority(),
		event.Status,
		event.Endpoint,
		event.Note(),
		event.Func(),
		event.Metric(),
		utils.SortedTags(event.PushedTags),
		utils.ReadableFloat(event.LeftValue),
		event.Operator(),
		utils.ReadableFloat(event.RightValue()),
		event.CurrentStep,
		event.FormattedTime(),
	)
}

func BuildCommonMailContent(event *model.Event) string {
	link := g.Link(event)
	return fmt.Sprintf(
		`

    				%s
    				%d

    				Endpoint:
    				%s

    				Metric:
    				%s

    				Tags:
    				%s

    				%s
    				%s%s%s

    				Note:
    				%s

    				Max:
    				%d

    				Current:
    				%d

    				Timesramp:
    				%s

		%s`,
		event.Status,
		event.Priority(),
		event.Endpoint,
		event.Metric(),
		utils.SortedTags(event.PushedTags),
		event.Func(),
		utils.ReadableFloat(event.LeftValue),
		event.Operator(),
		utils.ReadableFloat(event.RightValue()),
		event.Note(),
		event.MaxStep(),
		event.CurrentStep,
		event.FormattedTime(),
		link,
		link,
	)
}

func GenerateSmsContent(event *model.Event) string {
	return BuildCommonSMSContent(event)
}

func GenerateMailContent(event *model.Event) string {
	return BuildCommonMailContent(event)
}
<|endoftext|>


<commit_before>package main

import (
	"testing"
)

func assert(expression bool, message string, t *testing.T) {
	if !expression {
		t.Errorf(message)
		return
	}
	t.Log(message)
}

func assertEqual(expected, actual interface{}, message string, t *testing.T) {
	if expected != actual {
		t.Errorf("Expected %#v, but got %#v - %s", expected, actual, message)
		return
	}
	t.Log(message)
}

func TestRegionTypePercent(t *testing.T) {
	r := StringToRegion("pct:41.6,7.5,40,70")
	assert(r.Type == RTPercent, "r.Type == RTPercent", t)
	assertEqual(41.6, r.X, "r.X", t)
	assertEqual(7.5, r.Y, "r.Y", t)
	assertEqual(40.0, r.W, "r.W", t)
	assertEqual(70.0, r.H, "r.H", t)
}

func TestRegionTypePixels(t *testing.T) {
	r := StringToRegion("10,10,40,70")
	assert(r.Type == RTPixel, "r.Type == RTPixel", t)
	assertEqual(10.0, r.X, "r.X", t)
	assertEqual(10.0, r.Y, "r.Y", t)
	assertEqual(40.0, r.W, "r.W", t)
	assertEqual(70.0, r.H, "r.H", t)
}

func TestRegionTypeFull(t *testing.T) {
	r := StringToRegion("full")
	assert(r.Type == RTFull, "r.Type == RTFull", t)
}
<commit_msg>Make test output more readable<commit_after>package main

import (
	"fmt"
	"testing"
	"regexp"
	"runtime"
)

var re = regexp.MustCompile(`^.*jp2tileserver\.(.*)$`)

type Caller struct {
	Func     *runtime.Func
	Name     string
	Filename string
	Line     int
}

func getCallerName(skip int) *Caller {
	// Increase skip since they surely don't want *this* function
	pc, file, line, _ := runtime.Caller(skip + 1)
	fn := runtime.FuncForPC(pc)
	return &Caller{
		Func:     fn,
		Name:     re.ReplaceAllString(fn.Name(), "$1"),
		Filename: file,
		Line:     line,
	}
}

func success(caller *Caller, message string, t *testing.T) {
	fmt.Printf("\033[32mok\033[0m        %s(): %s\n", caller.Name, message)
}

func failure(caller *Caller, message string, t *testing.T) {
	fmt.Printf("\033[31;1mnot ok\033[0m    %s(): %s\n", caller.Name, message)
	fmt.Printf("          - %s:%d\n", caller.Filename, caller.Line)
	t.Fail()
}

func assert(expression bool, message string, t *testing.T) {
	caller := getCallerName(1)
	if !expression {
		failure(caller, message, t)
		return
	}
	success(caller, message, t)
}

func assertEqual(expected, actual interface{}, message string, t *testing.T) {
	caller := getCallerName(1)
	if expected != actual {
		failure(caller, fmt.Sprintf("Expected %#v, but got %#v - %s", expected, actual, message), t)
		return
	}
	success(caller, message, t)
}

func TestRegionTypePercent(t *testing.T) {
	r := StringToRegion("pct:41.6,7.5,40,70")
	assert(r.Type == RTPercent, "r.Type == RTPercent", t)
	assertEqual(41.6, r.X, "r.X", t)
	assertEqual(7.5, r.Y, "r.Y", t)
	assertEqual(40.0, r.W, "r.W", t)
	assertEqual(70.0, r.H, "r.H", t)
}

func TestRegionTypePixels(t *testing.T) {
	r := StringToRegion("10,10,40,70")
	assert(r.Type == RTPixel, "r.Type == RTPixel", t)
	assertEqual(10.0, r.X, "r.X", t)
	assertEqual(10.0, r.Y, "r.Y", t)
	assertEqual(40.0, r.W, "r.W", t)
	assertEqual(70.0, r.H, "r.H", t)
}

func TestRegionTypeFull(t *testing.T) {
	r := StringToRegion("full")
	assert(r.Type == RTFull, "r.Type == RTFull", t)
}
<|endoftext|>


<commit_before>package moh

import (
	"io/ioutil"
	"net/http"
)

// Replier is the counterpart of Requester.
// It is a HTTP server that responds to HTTP requests with it's handler function.
type Replier struct {
	MessagingServer
}

// NewReplier starts a new HTTP server on addr and returns a pointer to the Replier.
// All request will be replied by the handler function.
func NewReplier(addr string, handler func([]byte) []byte) (*Replier, error) {
	s, err := NewMessagingServer(addr)
	if err != nil {
		return nil, err
	}

	s.Mux.HandleFunc("/", makeHTTPHandler(handler))
	go s.Serve()
	return &Replier{MessagingServer: *s}, nil
}

func makeHTTPHandler(handler func([]byte) []byte) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		defer r.Body.Close()
		body, err := ioutil.ReadAll(r.Body)
		if err != nil {
			http.Error(w, "Cannot read request body", 500)
			return
		}

		reply := handler(body)

		_, err = w.Write(reply)
		if err != nil {
			http.Error(w, "Cannot write reply", 500)
			return
		}
	}
}
<commit_msg>Replier implements http.Handler<commit_after>package moh

import (
	"io/ioutil"
	"net/http"
)

// Replier is the counterpart of Requester.
// It is a HTTP server that responds to HTTP requests with it's handler function.
type Replier struct {
	MessagingServer
	Handler func([]byte) []byte
}

// NewReplier starts a new HTTP server on addr and returns a pointer to the Replier.
// All request will be replied by the handler function.
func NewReplier(addr string, handler func([]byte) []byte) (*Replier, error) {
	s, err := NewMessagingServer(addr)
	if err != nil {
		return nil, err
	}

	r := &Replier{
		MessagingServer: *s,
		Handler:         handler,
	}

	s.Mux.Handle("/", r)
	go s.Serve()
	return r, nil
}

// ServeHTTP implements the http.Handler interface for Replier.
func (rep *Replier) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	defer req.Body.Close()
	body, err := ioutil.ReadAll(req.Body)
	if err != nil {
		http.Error(w, "Cannot read request body", 500)
		return
	}

	reply := rep.Handler(body)

	_, err = w.Write(reply)
	if err != nil {
		http.Error(w, "Cannot write reply", 500)
		return
	}
}
<|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 schemamanager

import (
	"fmt"
	"sync"
	"time"

	"golang.org/x/net/context"

	"vitess.io/vitess/go/sync2"
	"vitess.io/vitess/go/vt/schema"
	"vitess.io/vitess/go/vt/sqlparser"
	"vitess.io/vitess/go/vt/topo"
	"vitess.io/vitess/go/vt/wrangler"

	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
)

// TabletExecutor applies schema changes to all tablets.
type TabletExecutor struct {
	wr                   *wrangler.Wrangler
	tablets              []*topodatapb.Tablet
	isClosed             bool
	allowBigSchemaChange bool
	onlineSchemaChange   bool
	keyspace             string
	waitReplicasTimeout  time.Duration
}

// NewTabletExecutor creates a new TabletExecutor instance
func NewTabletExecutor(wr *wrangler.Wrangler, waitReplicasTimeout time.Duration) *TabletExecutor {
	return &TabletExecutor{
		wr:                   wr,
		isClosed:             true,
		allowBigSchemaChange: false,
		onlineSchemaChange:   false,
		waitReplicasTimeout:  waitReplicasTimeout,
	}
}

// AllowBigSchemaChange changes TabletExecutor such that big schema changes
// will no longer be rejected.
func (exec *TabletExecutor) AllowBigSchemaChange() {
	exec.allowBigSchemaChange = true
}

// DisallowBigSchemaChange enables the check for big schema changes such that
// TabletExecutor will reject these.
func (exec *TabletExecutor) DisallowBigSchemaChange() {
	exec.allowBigSchemaChange = false
}

// SetOnlineSchemaChange sets TabletExecutor such that it initiates online schema change migrations
func (exec *TabletExecutor) SetOnlineSchemaChange() {
	exec.onlineSchemaChange = true
}

// Open opens a connection to the master for every shard.
func (exec *TabletExecutor) Open(ctx context.Context, keyspace string) error {
	if !exec.isClosed {
		return nil
	}
	exec.keyspace = keyspace
	shardNames, err := exec.wr.TopoServer().GetShardNames(ctx, keyspace)
	if err != nil {
		return fmt.Errorf("unable to get shard names for keyspace: %s, error: %v", keyspace, err)
	}
	exec.tablets = make([]*topodatapb.Tablet, len(shardNames))
	for i, shardName := range shardNames {
		shardInfo, err := exec.wr.TopoServer().GetShard(ctx, keyspace, shardName)
		if err != nil {
			return fmt.Errorf("unable to get shard info, keyspace: %s, shard: %s, error: %v", keyspace, shardName, err)
		}
		if !shardInfo.HasMaster() {
			return fmt.Errorf("shard: %s does not have a master", shardName)
		}
		tabletInfo, err := exec.wr.TopoServer().GetTablet(ctx, shardInfo.MasterAlias)
		if err != nil {
			return fmt.Errorf("unable to get master tablet info, keyspace: %s, shard: %s, error: %v", keyspace, shardName, err)
		}
		exec.tablets[i] = tabletInfo.Tablet
	}

	if len(exec.tablets) == 0 {
		return fmt.Errorf("keyspace: %s does not contain any master tablets", keyspace)
	}
	exec.isClosed = false
	return nil
}

// Validate validates a list of sql statements.
func (exec *TabletExecutor) Validate(ctx context.Context, sqls []string) error {
	if exec.isClosed {
		return fmt.Errorf("executor is closed")
	}

	// We ignore DATABASE-level DDLs here because detectBigSchemaChanges doesn't
	// look at them anyway.
	parsedDDLs, _, err := exec.parseDDLs(sqls)
	if err != nil {
		return err
	}

	bigSchemaChange, err := exec.detectBigSchemaChanges(ctx, parsedDDLs)
	if bigSchemaChange && exec.allowBigSchemaChange {
		exec.wr.Logger().Warningf("Processing big schema change. This may cause visible MySQL downtime.")
		return nil
	}
	return err
}

func (exec *TabletExecutor) parseDDLs(sqls []string) ([]*sqlparser.DDL, []*sqlparser.DBDDL, error) {
	parsedDDLs := make([]*sqlparser.DDL, 0)
	parsedDBDDLs := make([]*sqlparser.DBDDL, 0)
	for _, sql := range sqls {
		stat, err := sqlparser.Parse(sql)
		if err != nil {
			return nil, nil, fmt.Errorf("failed to parse sql: %s, got error: %v", sql, err)
		}
		switch ddl := stat.(type) {
		case *sqlparser.DDL:
			parsedDDLs = append(parsedDDLs, ddl)
		case *sqlparser.DBDDL:
			parsedDBDDLs = append(parsedDBDDLs, ddl)
		default:
			if len(exec.tablets) != 1 {
				return nil, nil, fmt.Errorf("non-ddl statements can only be executed for single shard keyspaces: %s", sql)
			}
		}
	}
	return parsedDDLs, parsedDBDDLs, nil
}

// a schema change that satisfies any following condition is considered
// to be a big schema change and will be rejected.
//   1. Alter more than 100,000 rows.
//   2. Change a table with more than 2,000,000 rows (Drops are fine).
func (exec *TabletExecutor) detectBigSchemaChanges(ctx context.Context, parsedDDLs []*sqlparser.DDL) (bool, error) {
	// exec.tablets is guaranteed to have at least one element;
	// Otherwise, Open should fail and executor should fail.
	masterTabletInfo := exec.tablets[0]
	// get database schema, excluding views.
	dbSchema, err := exec.wr.TabletManagerClient().GetSchema(
		ctx, masterTabletInfo, []string{}, []string{}, false)
	if err != nil {
		return false, fmt.Errorf("unable to get database schema, error: %v", err)
	}
	tableWithCount := make(map[string]uint64, len(dbSchema.TableDefinitions))
	for _, tableSchema := range dbSchema.TableDefinitions {
		tableWithCount[tableSchema.Name] = tableSchema.RowCount
	}
	for _, ddl := range parsedDDLs {
		switch ddl.Action {
		case sqlparser.DropStr, sqlparser.CreateStr, sqlparser.TruncateStr, sqlparser.RenameStr:
			continue
		case sqlparser.AlterStr:
			if exec.onlineSchemaChange {
				// Seeing that we intend to run an online-schema-change, we can skip the "big change" check.
				continue
			}
		}
		tableName := ddl.Table.Name.String()
		if rowCount, ok := tableWithCount[tableName]; ok {
			if rowCount > 100000 && ddl.Action == sqlparser.AlterStr {
				return true, fmt.Errorf(
					"big schema change detected. Disable check with -allow_long_unavailability. ddl: %s alters a table with more than 100 thousand rows", sqlparser.String(ddl))
			}
			if rowCount > 2000000 {
				return true, fmt.Errorf(
					"big schema change detected. Disable check with -allow_long_unavailability. ddl: %s changes a table with more than 2 million rows", sqlparser.String(ddl))
			}
		}
	}
	return false, nil
}

func (exec *TabletExecutor) preflightSchemaChanges(ctx context.Context, sqls []string) error {
	_, err := exec.wr.TabletManagerClient().PreflightSchema(ctx, exec.tablets[0], sqls)
	return err
}

// Execute applies schema changes
func (exec *TabletExecutor) Execute(ctx context.Context, sqls []string) *ExecuteResult {
	execResult := ExecuteResult{}
	execResult.Sqls = sqls
	if exec.isClosed {
		execResult.ExecutorErr = "executor is closed"
		return &execResult
	}
	startTime := time.Now()
	defer func() { execResult.TotalTimeSpent = time.Since(startTime) }()

	// Lock the keyspace so our schema change doesn't overlap with other
	// keyspace-wide operations like resharding migrations.
	ctx, unlock, lockErr := exec.wr.TopoServer().LockKeyspace(ctx, exec.keyspace, "ApplySchemaKeyspace")
	if lockErr != nil {
		execResult.ExecutorErr = lockErr.Error()
		return &execResult
	}
	defer func() {
		// This is complicated because execResult.ExecutorErr
		// is not of type error.
		var unlockErr error
		unlock(&unlockErr)
		if execResult.ExecutorErr == "" && unlockErr != nil {
			execResult.ExecutorErr = unlockErr.Error()
		}
	}()

	// Make sure the schema changes introduce a table definition change.
	if err := exec.preflightSchemaChanges(ctx, sqls); err != nil {
		execResult.ExecutorErr = err.Error()
		return &execResult
	}

	for index, sql := range sqls {
		execResult.CurSQLIndex = index

		stat, err := sqlparser.Parse(sql)
		if err != nil {
			execResult.ExecutorErr = err.Error()
			return &execResult
		}
		strategy := schema.DDLStrategyNormal
		switch ddl := stat.(type) {
		case *sqlparser.DDL:
			if ddl.Action == sqlparser.AlterStr && exec.onlineSchemaChange {
				strategy = ddl.Strategy
			}
		}
		exec.wr.Logger().Infof("Received DDL request. strategy = %+v", strategy)
		exec.executeOnAllTablets(ctx, &execResult, sql, strategy)
		if len(execResult.FailedShards) > 0 {
			break
		}
	}
	return &execResult
}

func (exec *TabletExecutor) executeOnAllTablets(
	ctx context.Context, execResult *ExecuteResult, sql string,
	strategy sqlparser.DDLStrategy,
) {
	if strategy != schema.DDLStrategyNormal {
		onlineDDL, err := schema.NewOnlineDDL(exec.keyspace, "", sql, strategy)
		if err != nil {
			execResult.ExecutorErr = err.Error()
			return
		}
		conn, err := exec.wr.TopoServer().ConnForCell(ctx, topo.GlobalCell)
		if err != nil {
			execResult.ExecutorErr = fmt.Sprintf("online DDL ConnForCell error:%s", err.Error())
			return
		}
		err = onlineDDL.WriteTopo(ctx, conn, schema.MigrationRequestsPath())
		if err != nil {
			execResult.ExecutorErr = err.Error()
		}
		fmt.Println(onlineDDL.UUID)
		return
	}

	var wg sync.WaitGroup
	numOfMasterTablets := len(exec.tablets)
	wg.Add(numOfMasterTablets)
	errChan := make(chan ShardWithError, numOfMasterTablets)
	successChan := make(chan ShardResult, numOfMasterTablets)
	for _, tablet := range exec.tablets {
		go func(tablet *topodatapb.Tablet) {
			defer wg.Done()
			exec.executeOneTablet(ctx, tablet, sql, errChan, successChan)
		}(tablet)
	}
	wg.Wait()
	close(errChan)
	close(successChan)
	execResult.FailedShards = make([]ShardWithError, 0, len(errChan))
	execResult.SuccessShards = make([]ShardResult, 0, len(successChan))
	for e := range errChan {
		execResult.FailedShards = append(execResult.FailedShards, e)
	}
	for r := range successChan {
		execResult.SuccessShards = append(execResult.SuccessShards, r)
	}

	if len(execResult.FailedShards) > 0 {
		return
	}

	// If all shards succeeded, wait (up to waitReplicasTimeout) for replicas to
	// execute the schema change via replication. This is best-effort, meaning
	// we still return overall success if the timeout expires.
	concurrency := sync2.NewSemaphore(10, 0)
	reloadCtx, cancel := context.WithTimeout(ctx, exec.waitReplicasTimeout)
	defer cancel()
	for _, result := range execResult.SuccessShards {
		wg.Add(1)
		go func(result ShardResult) {
			defer wg.Done()
			exec.wr.ReloadSchemaShard(reloadCtx, exec.keyspace, result.Shard, result.Position, concurrency, false /* includeMaster */)
		}(result)
	}
	wg.Wait()
}

func (exec *TabletExecutor) executeOneTablet(
	ctx context.Context,
	tablet *topodatapb.Tablet,
	sql string,
	errChan chan ShardWithError,
	successChan chan ShardResult) {
	result, err := exec.wr.TabletManagerClient().ExecuteFetchAsDba(ctx, tablet, false, []byte(sql), 10, false, true)
	if err != nil {
		errChan <- ShardWithError{Shard: tablet.Shard, Err: err.Error()}
		return
	}
	// Get a replication position that's guaranteed to be after the schema change
	// was applied on the master.
	pos, err := exec.wr.TabletManagerClient().MasterPosition(ctx, tablet)
	if err != nil {
		errChan <- ShardWithError{
			Shard: tablet.Shard,
			Err:   fmt.Sprintf("couldn't get replication position after applying schema change on master: %v", err),
		}
		return
	}
	successChan <- ShardResult{
		Shard:    tablet.Shard,
		Result:   result,
		Position: pos,
	}
}

// Close clears tablet executor states
func (exec *TabletExecutor) Close() {
	if !exec.isClosed {
		exec.tablets = nil
		exec.isClosed = true
	}
}

var _ Executor = (*TabletExecutor)(nil)
<commit_msg>extracting table name<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 schemamanager

import (
	"fmt"
	"sync"
	"time"

	"golang.org/x/net/context"

	"vitess.io/vitess/go/sync2"
	"vitess.io/vitess/go/vt/schema"
	"vitess.io/vitess/go/vt/sqlparser"
	"vitess.io/vitess/go/vt/topo"
	"vitess.io/vitess/go/vt/wrangler"

	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
)

// TabletExecutor applies schema changes to all tablets.
type TabletExecutor struct {
	wr                   *wrangler.Wrangler
	tablets              []*topodatapb.Tablet
	isClosed             bool
	allowBigSchemaChange bool
	onlineSchemaChange   bool
	keyspace             string
	waitReplicasTimeout  time.Duration
}

// NewTabletExecutor creates a new TabletExecutor instance
func NewTabletExecutor(wr *wrangler.Wrangler, waitReplicasTimeout time.Duration) *TabletExecutor {
	return &TabletExecutor{
		wr:                   wr,
		isClosed:             true,
		allowBigSchemaChange: false,
		onlineSchemaChange:   false,
		waitReplicasTimeout:  waitReplicasTimeout,
	}
}

// AllowBigSchemaChange changes TabletExecutor such that big schema changes
// will no longer be rejected.
func (exec *TabletExecutor) AllowBigSchemaChange() {
	exec.allowBigSchemaChange = true
}

// DisallowBigSchemaChange enables the check for big schema changes such that
// TabletExecutor will reject these.
func (exec *TabletExecutor) DisallowBigSchemaChange() {
	exec.allowBigSchemaChange = false
}

// SetOnlineSchemaChange sets TabletExecutor such that it initiates online schema change migrations
func (exec *TabletExecutor) SetOnlineSchemaChange() {
	exec.onlineSchemaChange = true
}

// Open opens a connection to the master for every shard.
func (exec *TabletExecutor) Open(ctx context.Context, keyspace string) error {
	if !exec.isClosed {
		return nil
	}
	exec.keyspace = keyspace
	shardNames, err := exec.wr.TopoServer().GetShardNames(ctx, keyspace)
	if err != nil {
		return fmt.Errorf("unable to get shard names for keyspace: %s, error: %v", keyspace, err)
	}
	exec.tablets = make([]*topodatapb.Tablet, len(shardNames))
	for i, shardName := range shardNames {
		shardInfo, err := exec.wr.TopoServer().GetShard(ctx, keyspace, shardName)
		if err != nil {
			return fmt.Errorf("unable to get shard info, keyspace: %s, shard: %s, error: %v", keyspace, shardName, err)
		}
		if !shardInfo.HasMaster() {
			return fmt.Errorf("shard: %s does not have a master", shardName)
		}
		tabletInfo, err := exec.wr.TopoServer().GetTablet(ctx, shardInfo.MasterAlias)
		if err != nil {
			return fmt.Errorf("unable to get master tablet info, keyspace: %s, shard: %s, error: %v", keyspace, shardName, err)
		}
		exec.tablets[i] = tabletInfo.Tablet
	}

	if len(exec.tablets) == 0 {
		return fmt.Errorf("keyspace: %s does not contain any master tablets", keyspace)
	}
	exec.isClosed = false
	return nil
}

// Validate validates a list of sql statements.
func (exec *TabletExecutor) Validate(ctx context.Context, sqls []string) error {
	if exec.isClosed {
		return fmt.Errorf("executor is closed")
	}

	// We ignore DATABASE-level DDLs here because detectBigSchemaChanges doesn't
	// look at them anyway.
	parsedDDLs, _, err := exec.parseDDLs(sqls)
	if err != nil {
		return err
	}

	bigSchemaChange, err := exec.detectBigSchemaChanges(ctx, parsedDDLs)
	if bigSchemaChange && exec.allowBigSchemaChange {
		exec.wr.Logger().Warningf("Processing big schema change. This may cause visible MySQL downtime.")
		return nil
	}
	return err
}

func (exec *TabletExecutor) parseDDLs(sqls []string) ([]*sqlparser.DDL, []*sqlparser.DBDDL, error) {
	parsedDDLs := make([]*sqlparser.DDL, 0)
	parsedDBDDLs := make([]*sqlparser.DBDDL, 0)
	for _, sql := range sqls {
		stat, err := sqlparser.Parse(sql)
		if err != nil {
			return nil, nil, fmt.Errorf("failed to parse sql: %s, got error: %v", sql, err)
		}
		switch ddl := stat.(type) {
		case *sqlparser.DDL:
			parsedDDLs = append(parsedDDLs, ddl)
		case *sqlparser.DBDDL:
			parsedDBDDLs = append(parsedDBDDLs, ddl)
		default:
			if len(exec.tablets) != 1 {
				return nil, nil, fmt.Errorf("non-ddl statements can only be executed for single shard keyspaces: %s", sql)
			}
		}
	}
	return parsedDDLs, parsedDBDDLs, nil
}

// a schema change that satisfies any following condition is considered
// to be a big schema change and will be rejected.
//   1. Alter more than 100,000 rows.
//   2. Change a table with more than 2,000,000 rows (Drops are fine).
func (exec *TabletExecutor) detectBigSchemaChanges(ctx context.Context, parsedDDLs []*sqlparser.DDL) (bool, error) {
	// exec.tablets is guaranteed to have at least one element;
	// Otherwise, Open should fail and executor should fail.
	masterTabletInfo := exec.tablets[0]
	// get database schema, excluding views.
	dbSchema, err := exec.wr.TabletManagerClient().GetSchema(
		ctx, masterTabletInfo, []string{}, []string{}, false)
	if err != nil {
		return false, fmt.Errorf("unable to get database schema, error: %v", err)
	}
	tableWithCount := make(map[string]uint64, len(dbSchema.TableDefinitions))
	for _, tableSchema := range dbSchema.TableDefinitions {
		tableWithCount[tableSchema.Name] = tableSchema.RowCount
	}
	for _, ddl := range parsedDDLs {
		switch ddl.Action {
		case sqlparser.DropStr, sqlparser.CreateStr, sqlparser.TruncateStr, sqlparser.RenameStr:
			continue
		case sqlparser.AlterStr:
			if exec.onlineSchemaChange {
				// Seeing that we intend to run an online-schema-change, we can skip the "big change" check.
				continue
			}
		}
		tableName := ddl.Table.Name.String()
		if rowCount, ok := tableWithCount[tableName]; ok {
			if rowCount > 100000 && ddl.Action == sqlparser.AlterStr {
				return true, fmt.Errorf(
					"big schema change detected. Disable check with -allow_long_unavailability. ddl: %s alters a table with more than 100 thousand rows", sqlparser.String(ddl))
			}
			if rowCount > 2000000 {
				return true, fmt.Errorf(
					"big schema change detected. Disable check with -allow_long_unavailability. ddl: %s changes a table with more than 2 million rows", sqlparser.String(ddl))
			}
		}
	}
	return false, nil
}

func (exec *TabletExecutor) preflightSchemaChanges(ctx context.Context, sqls []string) error {
	_, err := exec.wr.TabletManagerClient().PreflightSchema(ctx, exec.tablets[0], sqls)
	return err
}

// Execute applies schema changes
func (exec *TabletExecutor) Execute(ctx context.Context, sqls []string) *ExecuteResult {
	execResult := ExecuteResult{}
	execResult.Sqls = sqls
	if exec.isClosed {
		execResult.ExecutorErr = "executor is closed"
		return &execResult
	}
	startTime := time.Now()
	defer func() { execResult.TotalTimeSpent = time.Since(startTime) }()

	// Lock the keyspace so our schema change doesn't overlap with other
	// keyspace-wide operations like resharding migrations.
	ctx, unlock, lockErr := exec.wr.TopoServer().LockKeyspace(ctx, exec.keyspace, "ApplySchemaKeyspace")
	if lockErr != nil {
		execResult.ExecutorErr = lockErr.Error()
		return &execResult
	}
	defer func() {
		// This is complicated because execResult.ExecutorErr
		// is not of type error.
		var unlockErr error
		unlock(&unlockErr)
		if execResult.ExecutorErr == "" && unlockErr != nil {
			execResult.ExecutorErr = unlockErr.Error()
		}
	}()

	// Make sure the schema changes introduce a table definition change.
	if err := exec.preflightSchemaChanges(ctx, sqls); err != nil {
		execResult.ExecutorErr = err.Error()
		return &execResult
	}

	for index, sql := range sqls {
		execResult.CurSQLIndex = index

		stat, err := sqlparser.Parse(sql)
		if err != nil {
			execResult.ExecutorErr = err.Error()
			return &execResult
		}
		strategy := schema.DDLStrategyNormal
		tableName := ""
		switch ddl := stat.(type) {
		case *sqlparser.DDL:
			if ddl.Action == sqlparser.AlterStr && exec.onlineSchemaChange {
				strategy = ddl.Strategy
			}
			tableName = ddl.Table.Name.String()
		}
		exec.wr.Logger().Infof("Received DDL request. strategy = %+v", strategy)
		exec.executeOnAllTablets(ctx, &execResult, sql, tableName, strategy)
		if len(execResult.FailedShards) > 0 {
			break
		}
	}
	return &execResult
}

func (exec *TabletExecutor) executeOnAllTablets(
	ctx context.Context, execResult *ExecuteResult, sql string,
	tableName string, strategy sqlparser.DDLStrategy,
) {
	if strategy != schema.DDLStrategyNormal {
		onlineDDL, err := schema.NewOnlineDDL(exec.keyspace, tableName, sql, strategy)
		if err != nil {
			execResult.ExecutorErr = err.Error()
			return
		}
		conn, err := exec.wr.TopoServer().ConnForCell(ctx, topo.GlobalCell)
		if err != nil {
			execResult.ExecutorErr = fmt.Sprintf("online DDL ConnForCell error:%s", err.Error())
			return
		}
		err = onlineDDL.WriteTopo(ctx, conn, schema.MigrationRequestsPath())
		if err != nil {
			execResult.ExecutorErr = err.Error()
		}
		fmt.Println(onlineDDL.UUID)
		return
	}

	var wg sync.WaitGroup
	numOfMasterTablets := len(exec.tablets)
	wg.Add(numOfMasterTablets)
	errChan := make(chan ShardWithError, numOfMasterTablets)
	successChan := make(chan ShardResult, numOfMasterTablets)
	for _, tablet := range exec.tablets {
		go func(tablet *topodatapb.Tablet) {
			defer wg.Done()
			exec.executeOneTablet(ctx, tablet, sql, errChan, successChan)
		}(tablet)
	}
	wg.Wait()
	close(errChan)
	close(successChan)
	execResult.FailedShards = make([]ShardWithError, 0, len(errChan))
	execResult.SuccessShards = make([]ShardResult, 0, len(successChan))
	for e := range errChan {
		execResult.FailedShards = append(execResult.FailedShards, e)
	}
	for r := range successChan {
		execResult.SuccessShards = append(execResult.SuccessShards, r)
	}

	if len(execResult.FailedShards) > 0 {
		return
	}

	// If all shards succeeded, wait (up to waitReplicasTimeout) for replicas to
	// execute the schema change via replication. This is best-effort, meaning
	// we still return overall success if the timeout expires.
	concurrency := sync2.NewSemaphore(10, 0)
	reloadCtx, cancel := context.WithTimeout(ctx, exec.waitReplicasTimeout)
	defer cancel()
	for _, result := range execResult.SuccessShards {
		wg.Add(1)
		go func(result ShardResult) {
			defer wg.Done()
			exec.wr.ReloadSchemaShard(reloadCtx, exec.keyspace, result.Shard, result.Position, concurrency, false /* includeMaster */)
		}(result)
	}
	wg.Wait()
}

func (exec *TabletExecutor) executeOneTablet(
	ctx context.Context,
	tablet *topodatapb.Tablet,
	sql string,
	errChan chan ShardWithError,
	successChan chan ShardResult) {
	result, err := exec.wr.TabletManagerClient().ExecuteFetchAsDba(ctx, tablet, false, []byte(sql), 10, false, true)
	if err != nil {
		errChan <- ShardWithError{Shard: tablet.Shard, Err: err.Error()}
		return
	}
	// Get a replication position that's guaranteed to be after the schema change
	// was applied on the master.
	pos, err := exec.wr.TabletManagerClient().MasterPosition(ctx, tablet)
	if err != nil {
		errChan <- ShardWithError{
			Shard: tablet.Shard,
			Err:   fmt.Sprintf("couldn't get replication position after applying schema change on master: %v", err),
		}
		return
	}
	successChan <- ShardResult{
		Shard:    tablet.Shard,
		Result:   result,
		Position: pos,
	}
}

// Close clears tablet executor states
func (exec *TabletExecutor) Close() {
	if !exec.isClosed {
		exec.tablets = nil
		exec.isClosed = true
	}
}

var _ Executor = (*TabletExecutor)(nil)
<|endoftext|>


<commit_before>/*
Copyright 2022 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 planbuilder

import (
	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
	"vitess.io/vitess/go/vt/sqlparser"
	"vitess.io/vitess/go/vt/vterrors"
	"vitess.io/vitess/go/vt/vtgate/engine"
	"vitess.io/vitess/go/vt/vtgate/evalengine"
	"vitess.io/vitess/go/vt/vtgate/semantics"
)

type projection struct {
	gen4Plan
	source      logicalPlan
	columnNames []string
	columns     []sqlparser.Expr
	primitive   *engine.Projection
	// unorderedColumnIdx is used to find the index at which we should add any column output from projection
	// we don't care for the ordering of. It should also be updated when such a column is added
	unorderedColumnIdx int
}

var _ logicalPlan = (*projection)(nil)

// WireupGen4 implements the logicalPlan interface
func (p *projection) WireupGen4(semTable *semantics.SemTable) error {
	columns := make([]evalengine.Expr, 0, len(p.columns))
	for _, expr := range p.columns {
		convert, err := evalengine.Convert(expr, semTable)
		if err != nil {
			return err
		}
		columns = append(columns, convert)
	}
	p.primitive = &engine.Projection{
		Cols:  p.columnNames,
		Exprs: columns,
		Input: p.source.Primitive(),
	}

	return p.source.WireupGen4(semTable)
}

// Inputs implements the logicalPlan interface
func (p *projection) Inputs() []logicalPlan {
	return []logicalPlan{p.source}
}

// Rewrite implements the logicalPlan interface
func (p *projection) Rewrite(inputs ...logicalPlan) error {
	if len(inputs) != 1 {
		return vterrors.Errorf(vtrpcpb.Code_INTERNAL, "wrong number of inputs")
	}
	p.source = inputs[0]
	return nil
}

// ContainsTables implements the logicalPlan interface
func (p *projection) ContainsTables() semantics.TableSet {
	return p.source.ContainsTables()
}

// OutputColumns implements the logicalPlan interface
func (p *projection) OutputColumns() []sqlparser.SelectExpr {
	columns := p.source.OutputColumns()
	for i, expr := range p.columns {
		columns = append(columns, &sqlparser.AliasedExpr{
			Expr: expr,
			As:   sqlparser.NewColIdent(p.columnNames[i]),
		})
	}
	return columns
}

// Primitive implements the logicalPlan interface
func (p *projection) Primitive() engine.Primitive {
	if p.primitive == nil {
		panic("WireUp not yet run")
	}
	return p.primitive
}

// addColumn is used to add a column output for the projection.
// This is the only function that should be used to add  columns to projection
func (p *projection) addColumn(idx *int, column sqlparser.Expr, columnName string) (int, error) {
	var offset int
	if idx == nil {
		p.unorderedColumnIdx++
		offset = len(p.columns) - p.unorderedColumnIdx
	} else {
		offset = *idx
	}
	if p.columnNames[offset] != "" || p.columns[offset] != nil {
		return -1, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "overwriting columns in projection is not permitted")
	}
	p.columns[offset] = column
	p.columnNames[offset] = columnName
	return offset, nil
}
<commit_msg>fix: cleanup after rebase<commit_after>/*
Copyright 2022 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 planbuilder

import (
	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
	"vitess.io/vitess/go/vt/sqlparser"
	"vitess.io/vitess/go/vt/vterrors"
	"vitess.io/vitess/go/vt/vtgate/engine"
	"vitess.io/vitess/go/vt/vtgate/evalengine"
	"vitess.io/vitess/go/vt/vtgate/semantics"
)

type projection struct {
	gen4Plan
	source      logicalPlan
	columnNames []string
	columns     []sqlparser.Expr
	primitive   *engine.Projection
	// unorderedColumnIdx is used to find the index at which we should add any column output from projection
	// we don't care for the ordering of. It should also be updated when such a column is added
	unorderedColumnIdx int
}

var _ logicalPlan = (*projection)(nil)

// WireupGen4 implements the logicalPlan interface
func (p *projection) WireupGen4(semTable *semantics.SemTable) error {
	columns := make([]evalengine.Expr, 0, len(p.columns))
	for _, expr := range p.columns {
		convert, err := evalengine.Translate(expr, semTable)
		if err != nil {
			return err
		}
		columns = append(columns, convert)
	}
	p.primitive = &engine.Projection{
		Cols:  p.columnNames,
		Exprs: columns,
		Input: p.source.Primitive(),
	}

	return p.source.WireupGen4(semTable)
}

// Inputs implements the logicalPlan interface
func (p *projection) Inputs() []logicalPlan {
	return []logicalPlan{p.source}
}

// Rewrite implements the logicalPlan interface
func (p *projection) Rewrite(inputs ...logicalPlan) error {
	if len(inputs) != 1 {
		return vterrors.Errorf(vtrpcpb.Code_INTERNAL, "wrong number of inputs")
	}
	p.source = inputs[0]
	return nil
}

// ContainsTables implements the logicalPlan interface
func (p *projection) ContainsTables() semantics.TableSet {
	return p.source.ContainsTables()
}

// OutputColumns implements the logicalPlan interface
func (p *projection) OutputColumns() []sqlparser.SelectExpr {
	columns := p.source.OutputColumns()
	for i, expr := range p.columns {
		columns = append(columns, &sqlparser.AliasedExpr{
			Expr: expr,
			As:   sqlparser.NewColIdent(p.columnNames[i]),
		})
	}
	return columns
}

// Primitive implements the logicalPlan interface
func (p *projection) Primitive() engine.Primitive {
	if p.primitive == nil {
		panic("WireUp not yet run")
	}
	return p.primitive
}

// addColumn is used to add a column output for the projection.
// This is the only function that should be used to add  columns to projection
func (p *projection) addColumn(idx *int, column sqlparser.Expr, columnName string) (int, error) {
	var offset int
	if idx == nil {
		p.unorderedColumnIdx++
		offset = len(p.columns) - p.unorderedColumnIdx
	} else {
		offset = *idx
	}
	if p.columnNames[offset] != "" || p.columns[offset] != nil {
		return -1, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "overwriting columns in projection is not permitted")
	}
	p.columns[offset] = column
	p.columnNames[offset] = columnName
	return offset, nil
}
<|endoftext|>


<commit_before>package db

import (
	"container/list"
	"sync"
)

// A FIFO queue for table cache.
type TableQueue struct {
	list     *list.List
	capacity int
}

// Add a new table Id into queue. Return the oldest one in the queue, or
// -1 if there are less than @capacity elements in the queue
func (tq *TableQueue) Add(id int64) (int64, *list.Element) {
	element := &list.Element{Value: id}
	tq.list.PushFront(element)
	if tq.list.Len() < tq.capacity {
		return -1, element
	} else {
		tmp := tq.list.Back()
		ret := tmp.Value.(int64)
		tq.list.Remove(tmp)
		return ret, element
	}
}

// Move the accessed element to the front of the queue
func (tq *TableQueue) Access(element *list.Element) {
	tq.list.MoveToFront(element)
}

type tblInfo struct {
	table   *Table
	element *list.Element
	done    chan bool
}

// Manage cache of tables.
type TableCache struct {
	mutex    sync.Mutex
	tableMap map[int64]tblInfo
	queue    TableQueue
	impl     *DbImpl
}

func MakeTableCache(impl *DbImpl, capacity int) *TableCache {
	return &TableCache{
		tableMap: map[int64]tblInfo{},
		queue:    TableQueue{capacity: capacity},
		impl:     impl,
	}
}

// Add a table into cache.
func (tc *TableCache) Add(table *Table, id int64) {
	tc.mutex.Lock()
	defer tc.mutex.Unlock()

	_, found := tc.tableMap[id]
	if found {
		panic("New table should not be in the cache!")
	}

	oldId, element := tc.queue.Add(id)
	tc.tableMap[id] = tblInfo{table: table, element: element, done: make(chan bool)}
	if oldId >= 0 {
		delete(tc.tableMap, oldId)
	}
}

func (tc *TableCache) Get(id int64) *Table {
	tc.mutex.Lock()

	// Find if the table is already in cache.
	info, found := tc.tableMap[id]
	if found && info.table != nil {
		tc.queue.Access(info.element)
		tc.mutex.Unlock()
		return info.table
	}

	// Someone else is trying to load the table, so let's wait.
	if found && info.table == nil {
		tc.mutex.Unlock()
		<-info.done
		return tc.Get(id)
	}

	// Cache item cannot be found, put a place holder so that
	// other go routine that uses it can wait.
	tc.tableMap[id] = tblInfo{done: make(chan bool)}
	tc.mutex.Unlock()

	finfo, found := tc.impl.GetManifest().GetFileInfo(id)
	if !found {
		panic("Expected id does not found")
	}

	file, status := tc.impl.GetEnv().NewSequentialFile(finfo.Location)
	if !status.Ok() {
		panic("File does not exist!")
	}

	var fsize uint64
	fsize, status = tc.impl.GetEnv().GetFileSize(finfo.Location)
	if !status.Ok() {
		panic("Cannot stat a file!")
	}

	buf := make([]byte, fsize)
	tbl := RecoverTable(file, buf, tc.impl.GetComparator())

	if tbl == nil {
		panic("Fails to recover a table!")
	}

	// Add the table into cache.
	tc.mutex.Lock()
	defer tc.mutex.Unlock()

	info, found = tc.tableMap[id]
	if !found {
		panic("Fails to find previously reserved entry!")
	}

	oldId, element := tc.queue.Add(id)

	info.table = tbl
	info.element = element

	tc.tableMap[id] = info

	// Notify all pending go routines that this table is loaded.
	close(info.done)

	if oldId >= 0 {
		delete(tc.tableMap, oldId)
	}

	return tbl
}

// The real DB type.
type DbImpl struct {
	path      string
	env       Env
	comp      Comparator
	skipList  *Skiplist
	tmpList   *Skiplist
	manifest  *Manifest
	tblCache  *TableCache
	compactor *Compactor
	mutex     sync.RWMutex
}

func (db *DbImpl) Put(opt WriteOptions, key, value []byte) Status {
	db.mutex.Lock()
	db.skipList.Put(key, value)
	db.mutex.Unlock()
	return MakeStatusOk()
}

func (db *DbImpl) Delete(opt WriteOptions, key []byte) Status {
	return db.Put(opt, key, []byte(""))
}

func (db *DbImpl) Write(opt WriteOptions, updates WriteBatch) Status {
	return MakeStatusNotFound(NOT_IMPLEMENTED)
}

func (db *DbImpl) Get(opt ReadOptions, key []byte) ([]byte, Status) {
	iter := db.NewIterator(opt)
	iter.Seek(key)
	defer iter.Close()

	if iter.Valid() && db.comp.Compare(iter.Key(), key) == 0 {
		return iter.Value(), MakeStatusOk()
	} else {
		return nil, MakeStatusNotFound(KEY_NOT_FOUND)
	}
}

func (db *DbImpl) NewIterator(opt ReadOptions) Iterator {
	iterList := make([]Iterator, 0)

	db.mutex.RLock()
	defer db.mutex.RUnlock()

	iterList = append(iterList, db.skipList.NewIterator(&opt))

	if db.tmpList != nil {
		iterList = append(iterList, db.tmpList.NewIterator(&opt))
	}

	sId := db.manifest.GetCurrentSnapshot()
	sinfo := db.manifest.GetSnapshotInfo(sId)

	// Add iterators from L0 tables.
	for idx, infos := range sinfo {
		if idx >= db.compactor.maxL0Levels {
			break
		}

		tid := infos[0].id
		tbl := db.GetTableCache().Get(tid)
		if tbl == nil {
			panic("Expected table is not found")
		}

		iter := tbl.NewIterator()
		iterList = append(iterList, iter)
	}

	// Add iterators from Ln tables.
	for i := db.compactor.maxL0Levels; i < len(sinfo); i++ {
		iter := MakeBinarySearchIterator(db, sinfo[i])
		iterList = append(iterList, iter)
	}

	return MakeHeapIterator(iterList, db.comp)
}

func (db *DbImpl) GetSnapshot() Snapshot {
	panic(NOT_IMPLEMENTED)
	return 0
}

func (db *DbImpl) ReleaseSnapshot(snap Snapshot) {
	panic(NOT_IMPLEMENTED)
}

func (db *DbImpl) GetApproximateSizes(ranges []Range) []uint64 {
	panic(NOT_IMPLEMENTED)
	return nil
}

func (db *DbImpl) CompactRange(start, limit []byte) {
	panic(NOT_IMPLEMENTED)
}

// Freeze current skiplist, push it down to tmpList, create a new skiplist
func (db *DbImpl) RotateSkiplist() (*Skiplist, *Skiplist) {
	db.mutex.Lock()
	defer db.mutex.Unlock()

	if db.tmpList != nil {
		panic("tmpList is not empty during rotation!")
	}

	db.tmpList = db.skipList
	db.skipList = MakeSkiplist()

	return db.skipList, db.tmpList
}

// Get parent dir of this db.
func (db *DbImpl) GetPath() string {
	return db.path
}

func (db *DbImpl) GetEnv() Env {
	return db.env
}

func (db *DbImpl) GetComparator() Comparator {
	return db.comp
}

func (db *DbImpl) GetManifest() *Manifest {
	return db.manifest
}

func (db *DbImpl) GetTableCache() *TableCache {
	return db.tblCache
}
<commit_msg>Add log writer to dbImpl<commit_after>package db

import (
	"bytes"
	"container/list"
	"encoding/gob"
	"path"
	"sync"
)

// A FIFO queue for table cache.
type TableQueue struct {
	list     *list.List
	capacity int
}

// Add a new table Id into queue. Return the oldest one in the queue, or
// -1 if there are less than @capacity elements in the queue
func (tq *TableQueue) Add(id int64) (int64, *list.Element) {
	element := &list.Element{Value: id}
	tq.list.PushFront(element)
	if tq.list.Len() < tq.capacity {
		return -1, element
	} else {
		tmp := tq.list.Back()
		ret := tmp.Value.(int64)
		tq.list.Remove(tmp)
		return ret, element
	}
}

// Move the accessed element to the front of the queue
func (tq *TableQueue) Access(element *list.Element) {
	tq.list.MoveToFront(element)
}

type tblInfo struct {
	table   *Table
	element *list.Element
	done    chan bool
}

// Manage cache of tables.
type TableCache struct {
	mutex    sync.Mutex
	tableMap map[int64]tblInfo
	queue    TableQueue
	impl     *DbImpl
}

func MakeTableCache(impl *DbImpl, capacity int) *TableCache {
	return &TableCache{
		tableMap: map[int64]tblInfo{},
		queue:    TableQueue{capacity: capacity},
		impl:     impl,
	}
}

// Add a table into cache.
func (tc *TableCache) Add(table *Table, id int64) {
	tc.mutex.Lock()
	defer tc.mutex.Unlock()

	_, found := tc.tableMap[id]
	if found {
		panic("New table should not be in the cache!")
	}

	oldId, element := tc.queue.Add(id)
	tc.tableMap[id] = tblInfo{table: table, element: element, done: make(chan bool)}
	if oldId >= 0 {
		delete(tc.tableMap, oldId)
	}
}

func (tc *TableCache) Get(id int64) *Table {
	tc.mutex.Lock()

	// Find if the table is already in cache.
	info, found := tc.tableMap[id]
	if found && info.table != nil {
		tc.queue.Access(info.element)
		tc.mutex.Unlock()
		return info.table
	}

	// Someone else is trying to load the table, so let's wait.
	if found && info.table == nil {
		tc.mutex.Unlock()
		<-info.done
		return tc.Get(id)
	}

	// Cache item cannot be found, put a place holder so that
	// other go routine that uses it can wait.
	tc.tableMap[id] = tblInfo{done: make(chan bool)}
	tc.mutex.Unlock()

	finfo, found := tc.impl.GetManifest().GetFileInfo(id)
	if !found {
		panic("Expected id does not found")
	}

	file, status := tc.impl.GetEnv().NewSequentialFile(finfo.Location)
	if !status.Ok() {
		panic("File does not exist!")
	}

	var fsize uint64
	fsize, status = tc.impl.GetEnv().GetFileSize(finfo.Location)
	if !status.Ok() {
		panic("Cannot stat a file!")
	}

	buf := make([]byte, fsize)
	tbl := RecoverTable(file, buf, tc.impl.GetComparator())

	if tbl == nil {
		panic("Fails to recover a table!")
	}

	// Add the table into cache.
	tc.mutex.Lock()
	defer tc.mutex.Unlock()

	info, found = tc.tableMap[id]
	if !found {
		panic("Fails to find previously reserved entry!")
	}

	oldId, element := tc.queue.Add(id)

	info.table = tbl
	info.element = element

	tc.tableMap[id] = info

	// Notify all pending go routines that this table is loaded.
	close(info.done)

	if oldId >= 0 {
		delete(tc.tableMap, oldId)
	}

	return tbl
}

// The real DB type.
type DbImpl struct {
	path      string
	env       Env
	comp      Comparator
	writer    *LogWriter
	skipList  *Skiplist
	tmpList   *Skiplist
	manifest  *Manifest
	tblCache  *TableCache
	compactor *Compactor
	mutex     sync.RWMutex
}

// Create a brand new Db.
func MakeDefaultDb(path string, env Env, com Comparator) *DbImpl {
	return nil
}

func (db *DbImpl) Put(opt WriteOptions, key, value []byte) Status {
	var buf bytes.Buffer
	enc := gob.NewEncoder(&buf)
	enc.Encode(key)
	enc.Encode(value)

	db.mutex.Lock()
	db.skipList.Put(key, value)
	status := db.writer.AddRecord(buf.Bytes())
	db.mutex.Unlock()

	return status
}

func (db *DbImpl) Delete(opt WriteOptions, key []byte) Status {
	return db.Put(opt, key, []byte(""))
}

func (db *DbImpl) Write(opt WriteOptions, updates WriteBatch) Status {
	return MakeStatusNotFound(NOT_IMPLEMENTED)
}

func (db *DbImpl) Get(opt ReadOptions, key []byte) ([]byte, Status) {
	iter := db.NewIterator(opt)
	iter.Seek(key)
	defer iter.Close()

	if iter.Valid() && db.comp.Compare(iter.Key(), key) == 0 {
		return iter.Value(), MakeStatusOk()
	} else {
		return nil, MakeStatusNotFound(KEY_NOT_FOUND)
	}
}

func (db *DbImpl) NewIterator(opt ReadOptions) Iterator {
	iterList := make([]Iterator, 0)

	db.mutex.RLock()
	defer db.mutex.RUnlock()

	iterList = append(iterList, db.skipList.NewIterator(&opt))

	if db.tmpList != nil {
		iterList = append(iterList, db.tmpList.NewIterator(&opt))
	}

	sId := db.manifest.GetCurrentSnapshot()
	sinfo := db.manifest.GetSnapshotInfo(sId)

	// Add iterators from L0 tables.
	for idx, infos := range sinfo {
		if idx >= db.compactor.maxL0Levels {
			break
		}

		tid := infos[0].id
		tbl := db.GetTableCache().Get(tid)
		if tbl == nil {
			panic("Expected table is not found")
		}

		iter := tbl.NewIterator()
		iterList = append(iterList, iter)
	}

	// Add iterators from Ln tables.
	for i := db.compactor.maxL0Levels; i < len(sinfo); i++ {
		iter := MakeBinarySearchIterator(db, sinfo[i])
		iterList = append(iterList, iter)
	}

	return MakeHeapIterator(iterList, db.comp)
}

func (db *DbImpl) GetSnapshot() Snapshot {
	panic(NOT_IMPLEMENTED)
	return 0
}

func (db *DbImpl) ReleaseSnapshot(snap Snapshot) {
	panic(NOT_IMPLEMENTED)
}

func (db *DbImpl) GetApproximateSizes(ranges []Range) []uint64 {
	panic(NOT_IMPLEMENTED)
	return nil
}

func (db *DbImpl) CompactRange(start, limit []byte) {
	panic(NOT_IMPLEMENTED)
}

// Freeze current skiplist, push it down to tmpList, create a new skiplist
func (db *DbImpl) RotateSkiplist() (*Skiplist, *Skiplist) {
	// Create a new log file.
	newLogFileNumber := db.manifest.CreateFile(false)
	logBaseName := GetLogName(newLogFileNumber)
	newLogPath := path.Join(db.path, logBaseName)
	writer := MakeLogWriter(db.env, newLogPath)
	if writer == nil {
		panic("Fails to create a new log file")
	}

	db.mutex.Lock()
	defer db.mutex.Unlock()

	if db.tmpList != nil {
		panic("tmpList is not empty during rotation!")
	}

	db.writer = writer
	db.tmpList = db.skipList
	db.skipList = MakeSkiplist()

	return db.skipList, db.tmpList
}

// Get parent dir of this db.
func (db *DbImpl) GetPath() string {
	return db.path
}

func (db *DbImpl) GetEnv() Env {
	return db.env
}

func (db *DbImpl) GetComparator() Comparator {
	return db.comp
}

func (db *DbImpl) GetManifest() *Manifest {
	return db.manifest
}

func (db *DbImpl) GetTableCache() *TableCache {
	return db.tblCache
}
<|endoftext|>


<commit_before>package consul

import (
	"context"
	"fmt"
	"log"
	"time"

	"github.com/hashicorp/consul/api"
	"github.com/hashicorp/nomad/nomad/structs"
)

const (
	// defaultPollFreq is the default rate to poll the Consul Checks API
	defaultPollFreq = 900 * time.Millisecond
)

// ChecksAPI is the part of the Consul API the checkWatcher requires.
type ChecksAPI interface {
	// Checks returns a list of all checks.
	Checks() (map[string]*api.AgentCheck, error)
}

// TaskRestarter allows the checkWatcher to restart tasks.
type TaskRestarter interface {
	Restart(source, reason string, failure bool)
}

// checkRestart handles restarting a task if a check is unhealthy.
type checkRestart struct {
	allocID   string
	taskName  string
	checkID   string
	checkName string
	taskKey   string // composite of allocID + taskName for uniqueness

	task           TaskRestarter
	grace          time.Duration
	interval       time.Duration
	timeLimit      time.Duration
	ignoreWarnings bool

	// Mutable fields

	// unhealthyState is the time a check first went unhealthy. Set to the
	// zero value if the check passes before timeLimit.
	unhealthyState time.Time

	// graceUntil is when the check's grace period expires and unhealthy
	// checks should be counted.
	graceUntil time.Time

	logger *log.Logger
}

// apply restart state for check and restart task if necessary. Currrent
// timestamp is passed in so all check updates have the same view of time (and
// to ease testing).
//
// Returns true if a restart was triggered in which case this check should be
// removed (checks are added on task startup).
func (c *checkRestart) apply(now time.Time, status string) bool {
	healthy := func() {
		if !c.unhealthyState.IsZero() {
			c.logger.Printf("[DEBUG] consul.health: alloc %q task %q check %q became healthy; canceling restart",
				c.allocID, c.taskName, c.checkName)
			c.unhealthyState = time.Time{}
		}
	}
	switch status {
	case api.HealthCritical:
	case api.HealthWarning:
		if c.ignoreWarnings {
			// Warnings are ignored, reset state and exit
			healthy()
			return false
		}
	default:
		// All other statuses are ok, reset state and exit
		healthy()
		return false
	}

	if now.Before(c.graceUntil) {
		// In grace period, exit
		return false
	}

	if c.unhealthyState.IsZero() {
		// First failure, set restart deadline
		if c.timeLimit != 0 {
			c.logger.Printf("[DEBUG] consul.health: alloc %q task %q check %q became unhealthy. Restarting in %s if not healthy",
				c.allocID, c.taskName, c.checkName, c.timeLimit)
		}
		c.unhealthyState = now
	}

	// restart timeLimit after start of this check becoming unhealthy
	restartAt := c.unhealthyState.Add(c.timeLimit)

	// Must test >= because if limit=1, restartAt == first failure
	if now.Equal(restartAt) || now.After(restartAt) {
		// hasn't become healthy by deadline, restart!
		c.logger.Printf("[DEBUG] consul.health: restarting alloc %q task %q due to unhealthy check %q", c.allocID, c.taskName, c.checkName)

		// Tell TaskRunner to restart due to failure
		const failure = true
		c.task.Restart("healthcheck", fmt.Sprintf("check %q unhealthy", c.checkName), failure)
		return true
	}

	return false
}

// checkWatchUpdates add or remove checks from the watcher
type checkWatchUpdate struct {
	checkID      string
	remove       bool
	checkRestart *checkRestart
}

// checkWatcher watches Consul checks and restarts tasks when they're
// unhealthy.
type checkWatcher struct {
	consul ChecksAPI

	// pollFreq is how often to poll the checks API and defaults to
	// defaultPollFreq
	pollFreq time.Duration

	// checkUpdateCh is how watches (and removals) are sent to the main
	// watching loop
	checkUpdateCh chan checkWatchUpdate

	// done is closed when Run has exited
	done chan struct{}

	// lastErr is true if the last Consul call failed. It is used to
	// squelch repeated error messages.
	lastErr bool

	logger *log.Logger
}

// newCheckWatcher creates a new checkWatcher but does not call its Run method.
func newCheckWatcher(logger *log.Logger, consul ChecksAPI) *checkWatcher {
	return &checkWatcher{
		consul:        consul,
		pollFreq:      defaultPollFreq,
		checkUpdateCh: make(chan checkWatchUpdate, 8),
		done:          make(chan struct{}),
		logger:        logger,
	}
}

// Run the main Consul checks watching loop to restart tasks when their checks
// fail. Blocks until context is canceled.
func (w *checkWatcher) Run(ctx context.Context) {
	defer close(w.done)

	// map of check IDs to their metadata
	checks := map[string]*checkRestart{}

	// timer for check polling
	checkTimer := time.NewTimer(0)
	defer checkTimer.Stop() // ensure timer is never leaked

	stopTimer := func() {
		checkTimer.Stop()
		select {
		case <-checkTimer.C:
		default:
		}
	}

	// disable by default
	stopTimer()

	// Main watch loop
	for {
		// disable polling if there are no checks
		if len(checks) == 0 {
			stopTimer()
		}

		select {
		case update := <-w.checkUpdateCh:
			if update.remove {
				// Remove a check
				delete(checks, update.checkID)
				continue
			}

			// Add/update a check
			checks[update.checkID] = update.checkRestart
			w.logger.Printf("[DEBUG] consul.health: watching alloc %q task %q check %q",
				update.checkRestart.allocID, update.checkRestart.taskName, update.checkRestart.checkName)

			// if first check was added make sure polling is enabled
			if len(checks) == 1 {
				stopTimer()
				checkTimer.Reset(w.pollFreq)
			}

		case <-ctx.Done():
			return

		case <-checkTimer.C:
			checkTimer.Reset(w.pollFreq)

			// Set "now" as the point in time the following check results represent
			now := time.Now()

			results, err := w.consul.Checks()
			if err != nil {
				if !w.lastErr {
					w.lastErr = true
					w.logger.Printf("[ERR] consul.health: error retrieving health checks: %q", err)
				}
				continue
			}

			w.lastErr = false

			// Keep track of tasks restarted this period so they
			// are only restarted once and all of their checks are
			// removed.
			restartedTasks := map[string]struct{}{}

			// Loop over watched checks and update their status from results
			for cid, check := range checks {
				if _, ok := restartedTasks[check.taskKey]; ok {
					// Check for this task already restarted; remove and skip check
					delete(checks, cid)
					continue
				}

				result, ok := results[cid]
				if !ok {
					// Only warn if outside grace period to avoid races with check registration
					if now.After(check.graceUntil) {
						w.logger.Printf("[WARN] consul.health: watched check %q (%s) not found in Consul", check.checkName, cid)
					}
					continue
				}

				restarted := check.apply(now, result.Status)
				if restarted {
					// Checks are registered+watched on
					// startup, so it's safe to remove them
					// whenever they're restarted
					delete(checks, cid)

					restartedTasks[check.taskKey] = struct{}{}
				}
			}

			// Ensure even passing checks for restartedTasks are removed
			if len(restartedTasks) > 0 {
				for cid, check := range checks {
					if _, ok := restartedTasks[check.taskKey]; ok {
						delete(checks, cid)
					}
				}
			}
		}
	}
}

// Watch a task and restart it if unhealthy.
func (w *checkWatcher) Watch(allocID, taskName, checkID string, check *structs.ServiceCheck, restarter TaskRestarter) {
	if !check.TriggersRestarts() {
		// Not watched, noop
		return
	}

	c := &checkRestart{
		allocID:        allocID,
		taskName:       taskName,
		checkID:        checkID,
		checkName:      check.Name,
		taskKey:        fmt.Sprintf("%s%s", allocID, taskName), // unique task ID
		task:           restarter,
		interval:       check.Interval,
		grace:          check.CheckRestart.Grace,
		graceUntil:     time.Now().Add(check.CheckRestart.Grace),
		timeLimit:      check.Interval * time.Duration(check.CheckRestart.Limit-1),
		ignoreWarnings: check.CheckRestart.IgnoreWarnings,
		logger:         w.logger,
	}

	update := checkWatchUpdate{
		checkID:      checkID,
		checkRestart: c,
	}

	select {
	case w.checkUpdateCh <- update:
		// sent watch
	case <-w.done:
		// exited; nothing to do
	}
}

// Unwatch a task.
func (w *checkWatcher) Unwatch(cid string) {
	c := checkWatchUpdate{
		checkID: cid,
		remove:  true,
	}
	select {
	case w.checkUpdateCh <- c:
		// sent remove watch
	case <-w.done:
		// exited; nothing to do
	}
}
<commit_msg>Fix comments: task -> check<commit_after>package consul

import (
	"context"
	"fmt"
	"log"
	"time"

	"github.com/hashicorp/consul/api"
	"github.com/hashicorp/nomad/nomad/structs"
)

const (
	// defaultPollFreq is the default rate to poll the Consul Checks API
	defaultPollFreq = 900 * time.Millisecond
)

// ChecksAPI is the part of the Consul API the checkWatcher requires.
type ChecksAPI interface {
	// Checks returns a list of all checks.
	Checks() (map[string]*api.AgentCheck, error)
}

// TaskRestarter allows the checkWatcher to restart tasks.
type TaskRestarter interface {
	Restart(source, reason string, failure bool)
}

// checkRestart handles restarting a task if a check is unhealthy.
type checkRestart struct {
	allocID   string
	taskName  string
	checkID   string
	checkName string
	taskKey   string // composite of allocID + taskName for uniqueness

	task           TaskRestarter
	grace          time.Duration
	interval       time.Duration
	timeLimit      time.Duration
	ignoreWarnings bool

	// Mutable fields

	// unhealthyState is the time a check first went unhealthy. Set to the
	// zero value if the check passes before timeLimit.
	unhealthyState time.Time

	// graceUntil is when the check's grace period expires and unhealthy
	// checks should be counted.
	graceUntil time.Time

	logger *log.Logger
}

// apply restart state for check and restart task if necessary. Currrent
// timestamp is passed in so all check updates have the same view of time (and
// to ease testing).
//
// Returns true if a restart was triggered in which case this check should be
// removed (checks are added on task startup).
func (c *checkRestart) apply(now time.Time, status string) bool {
	healthy := func() {
		if !c.unhealthyState.IsZero() {
			c.logger.Printf("[DEBUG] consul.health: alloc %q task %q check %q became healthy; canceling restart",
				c.allocID, c.taskName, c.checkName)
			c.unhealthyState = time.Time{}
		}
	}
	switch status {
	case api.HealthCritical:
	case api.HealthWarning:
		if c.ignoreWarnings {
			// Warnings are ignored, reset state and exit
			healthy()
			return false
		}
	default:
		// All other statuses are ok, reset state and exit
		healthy()
		return false
	}

	if now.Before(c.graceUntil) {
		// In grace period, exit
		return false
	}

	if c.unhealthyState.IsZero() {
		// First failure, set restart deadline
		if c.timeLimit != 0 {
			c.logger.Printf("[DEBUG] consul.health: alloc %q task %q check %q became unhealthy. Restarting in %s if not healthy",
				c.allocID, c.taskName, c.checkName, c.timeLimit)
		}
		c.unhealthyState = now
	}

	// restart timeLimit after start of this check becoming unhealthy
	restartAt := c.unhealthyState.Add(c.timeLimit)

	// Must test >= because if limit=1, restartAt == first failure
	if now.Equal(restartAt) || now.After(restartAt) {
		// hasn't become healthy by deadline, restart!
		c.logger.Printf("[DEBUG] consul.health: restarting alloc %q task %q due to unhealthy check %q", c.allocID, c.taskName, c.checkName)

		// Tell TaskRunner to restart due to failure
		const failure = true
		c.task.Restart("healthcheck", fmt.Sprintf("check %q unhealthy", c.checkName), failure)
		return true
	}

	return false
}

// checkWatchUpdates add or remove checks from the watcher
type checkWatchUpdate struct {
	checkID      string
	remove       bool
	checkRestart *checkRestart
}

// checkWatcher watches Consul checks and restarts tasks when they're
// unhealthy.
type checkWatcher struct {
	consul ChecksAPI

	// pollFreq is how often to poll the checks API and defaults to
	// defaultPollFreq
	pollFreq time.Duration

	// checkUpdateCh is how watches (and removals) are sent to the main
	// watching loop
	checkUpdateCh chan checkWatchUpdate

	// done is closed when Run has exited
	done chan struct{}

	// lastErr is true if the last Consul call failed. It is used to
	// squelch repeated error messages.
	lastErr bool

	logger *log.Logger
}

// newCheckWatcher creates a new checkWatcher but does not call its Run method.
func newCheckWatcher(logger *log.Logger, consul ChecksAPI) *checkWatcher {
	return &checkWatcher{
		consul:        consul,
		pollFreq:      defaultPollFreq,
		checkUpdateCh: make(chan checkWatchUpdate, 8),
		done:          make(chan struct{}),
		logger:        logger,
	}
}

// Run the main Consul checks watching loop to restart tasks when their checks
// fail. Blocks until context is canceled.
func (w *checkWatcher) Run(ctx context.Context) {
	defer close(w.done)

	// map of check IDs to their metadata
	checks := map[string]*checkRestart{}

	// timer for check polling
	checkTimer := time.NewTimer(0)
	defer checkTimer.Stop() // ensure timer is never leaked

	stopTimer := func() {
		checkTimer.Stop()
		select {
		case <-checkTimer.C:
		default:
		}
	}

	// disable by default
	stopTimer()

	// Main watch loop
	for {
		// disable polling if there are no checks
		if len(checks) == 0 {
			stopTimer()
		}

		select {
		case update := <-w.checkUpdateCh:
			if update.remove {
				// Remove a check
				delete(checks, update.checkID)
				continue
			}

			// Add/update a check
			checks[update.checkID] = update.checkRestart
			w.logger.Printf("[DEBUG] consul.health: watching alloc %q task %q check %q",
				update.checkRestart.allocID, update.checkRestart.taskName, update.checkRestart.checkName)

			// if first check was added make sure polling is enabled
			if len(checks) == 1 {
				stopTimer()
				checkTimer.Reset(w.pollFreq)
			}

		case <-ctx.Done():
			return

		case <-checkTimer.C:
			checkTimer.Reset(w.pollFreq)

			// Set "now" as the point in time the following check results represent
			now := time.Now()

			results, err := w.consul.Checks()
			if err != nil {
				if !w.lastErr {
					w.lastErr = true
					w.logger.Printf("[ERR] consul.health: error retrieving health checks: %q", err)
				}
				continue
			}

			w.lastErr = false

			// Keep track of tasks restarted this period so they
			// are only restarted once and all of their checks are
			// removed.
			restartedTasks := map[string]struct{}{}

			// Loop over watched checks and update their status from results
			for cid, check := range checks {
				if _, ok := restartedTasks[check.taskKey]; ok {
					// Check for this task already restarted; remove and skip check
					delete(checks, cid)
					continue
				}

				result, ok := results[cid]
				if !ok {
					// Only warn if outside grace period to avoid races with check registration
					if now.After(check.graceUntil) {
						w.logger.Printf("[WARN] consul.health: watched check %q (%s) not found in Consul", check.checkName, cid)
					}
					continue
				}

				restarted := check.apply(now, result.Status)
				if restarted {
					// Checks are registered+watched on
					// startup, so it's safe to remove them
					// whenever they're restarted
					delete(checks, cid)

					restartedTasks[check.taskKey] = struct{}{}
				}
			}

			// Ensure even passing checks for restartedTasks are removed
			if len(restartedTasks) > 0 {
				for cid, check := range checks {
					if _, ok := restartedTasks[check.taskKey]; ok {
						delete(checks, cid)
					}
				}
			}
		}
	}
}

// Watch a check and restart its task if unhealthy.
func (w *checkWatcher) Watch(allocID, taskName, checkID string, check *structs.ServiceCheck, restarter TaskRestarter) {
	if !check.TriggersRestarts() {
		// Not watched, noop
		return
	}

	c := &checkRestart{
		allocID:        allocID,
		taskName:       taskName,
		checkID:        checkID,
		checkName:      check.Name,
		taskKey:        fmt.Sprintf("%s%s", allocID, taskName), // unique task ID
		task:           restarter,
		interval:       check.Interval,
		grace:          check.CheckRestart.Grace,
		graceUntil:     time.Now().Add(check.CheckRestart.Grace),
		timeLimit:      check.Interval * time.Duration(check.CheckRestart.Limit-1),
		ignoreWarnings: check.CheckRestart.IgnoreWarnings,
		logger:         w.logger,
	}

	update := checkWatchUpdate{
		checkID:      checkID,
		checkRestart: c,
	}

	select {
	case w.checkUpdateCh <- update:
		// sent watch
	case <-w.done:
		// exited; nothing to do
	}
}

// Unwatch a check.
func (w *checkWatcher) Unwatch(cid string) {
	c := checkWatchUpdate{
		checkID: cid,
		remove:  true,
	}
	select {
	case w.checkUpdateCh <- c:
		// sent remove watch
	case <-w.done:
		// exited; nothing to do
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/doozr/guac"
)

func TestDispatchRunsReceiverInBackground(t *testing.T) {
	dispatcher := func(events guac.EventChan, done DoneChan) error {
		return nil
	}
	done := make(DoneChan)
	events := make(guac.EventChan)
	waitGroup := sync.WaitGroup{}

	abort := dispatch(dispatcher, events, done, &waitGroup)
	select {
	case <-abort:
	case <-time.After(2 * time.Second):
		t.Fatal("Expected event within 2 seconds")
	}

	waitGroup.Wait()
}

func TestDispatchShutDownCleanlyWithErrors(t *testing.T) {
	dispatcher := func(events guac.EventChan, done DoneChan) error {
		return fmt.Errorf("Error!")
	}
	done := make(DoneChan)
	events := make(guac.EventChan)
	waitGroup := sync.WaitGroup{}

	abort := dispatch(dispatcher, events, done, &waitGroup)
	select {
	case <-abort:
	case <-time.After(2 * time.Second):
		t.Fatal("Expected event within 2 seconds")
	}

	waitGroup.Wait()
}

func TestDispatcherSendsMessagesToMessageHandler(t *testing.T) {
	done := make(DoneChan)
	events := make(guac.EventChan)
	var received *guac.MessageEvent

	handleMessage := func(event guac.MessageEvent) error {
		if received != nil {
			t.Fatal("Already received a message ", event)
		}
		received = &event
		close(done)
		return nil
	}
	handleUserChange := func(event guac.UserInfo) {
		t.Fatal("Unexpected call to UserChangeHandler")
	}
	dispatcher := createDispatcher(1*time.Millisecond, handleMessage, handleUserChange)

	go func() {
		events <- guac.MessageEvent{
			Text: "test event",
		}
	}()

	err := dispatcher(events, done)
	if err != nil {
		t.Fatal("Unexpected error ", err)
	}

	if received == nil || received.Text != "test event" {
		t.Fatal("Did not receive expected message")
	}
}

func TestDispatcherReturnsErrorIfMessageFails(t *testing.T) {
	done := make(DoneChan)
	events := make(guac.EventChan)

	handleMessage := func(event guac.MessageEvent) error {
		return fmt.Errorf("Error!")
	}
	handleUserChange := func(event guac.UserInfo) {
		t.Fatal("Unexpected call to UserChangeHandler")
	}
	dispatcher := createDispatcher(1*time.Millisecond, handleMessage, handleUserChange)

	go func() {
		events <- guac.MessageEvent{
			Text: "test event",
		}
	}()

	err := dispatcher(events, done)
	if err == nil {
		t.Fatal("Expected error ", err)
	}
}

func TestDispatcherSendsUserChangesToUserChangeHandler(t *testing.T) {
	done := make(DoneChan)
	events := make(guac.EventChan)
	var received *guac.UserInfo

	handleMessage := func(event guac.MessageEvent) error {
		t.Fatal("Unexpected call to MessageHandler")
		return nil
	}
	handleUserChange := func(event guac.UserInfo) {
		if received != nil {
			t.Fatal("Already received a user change ", event)
		}
		received = &event
	}
	dispatcher := createDispatcher(1*time.Millisecond, handleMessage, handleUserChange)

	go func() {
		events <- guac.UserChangeEvent{
			UserInfo: guac.UserInfo{Name: "test event"},
		}
		close(done)
	}()

	err := dispatcher(events, done)
	if err != nil {
		t.Fatal("Unexpected error ", err)
	}

	if received == nil || received.Name != "test event" {
		t.Fatal("Did not receive expected message")
	}
}

func TestDispatcherReturnsErrorIfUserChangeFails(t *testing.T) {
	done := make(DoneChan)
	events := make(guac.EventChan)

	handleMessage := func(event guac.MessageEvent) error {
		t.Fatal("Unexpected call to MessageHandler")
		return nil
	}
	handleUserChange := func(event guac.UserInfo) {
	}
	dispatcher := createDispatcher(1*time.Millisecond, handleMessage, handleUserChange)

	go func() {
		events <- guac.UserChangeEvent{
			UserInfo: guac.UserInfo{Name: "test event"},
		}
	}()

	err := dispatcher(events, done)
	if err == nil {
		t.Fatal("Expected error ", err)
	}
}

func TestDispatcherReturnsErrorOnTimeout(t *testing.T) {
	done := make(DoneChan)
	events := make(guac.EventChan)

	handleMessage := func(event guac.MessageEvent) error {
		t.Fatal("Unexpected call to MessageHandler")
		return nil
	}
	handleUserChange := func(event guac.UserInfo) {
		t.Fatal("Unexpected call to MessageHandler")
	}
	dispatcher := createDispatcher(1*time.Millisecond, handleMessage, handleUserChange)

	err := dispatcher(events, done)
	if err == nil {
		t.Fatal("Expected error ", err)
	}
}

func TestDispatcherReturnsNoErrorIfDone(t *testing.T) {
	done := make(DoneChan)
	events := make(guac.EventChan)

	handleMessage := func(event guac.MessageEvent) error {
		t.Fatal("Unexpected call to MessageHandler")
		return nil
	}
	handleUserChange := func(event guac.UserInfo) {
		t.Fatal("Unexpected call to MessageHandler")
	}
	dispatcher := createDispatcher(1*time.Millisecond, handleMessage, handleUserChange)

	close(done)
	err := dispatcher(events, done)
	if err != nil {
		t.Fatal("Unexpected error ", err)
	}
}

func TestDispatcherSwallowsUnknownEvents(t *testing.T) {
	done := make(DoneChan)
	events := make(guac.EventChan)

	handleMessage := func(event guac.MessageEvent) error {
		t.Fatal("Unexpected call to MessageHandler")
		return nil
	}
	handleUserChange := func(event guac.UserInfo) {
		t.Fatal("Unexpected call to MessageHandler")
	}
	dispatcher := createDispatcher(1*time.Second, handleMessage, handleUserChange)

	// events is blocking so these things must be read in sequence
	go func() {
		events <- "unknown event"
		close(done)
	}()
	err := dispatcher(events, done)
	if err != nil {
		t.Fatal("Unexpected error ", err)
	}

	select {
	case <-events:
		t.Fatal("Expected event to have been swallowed")
	default:
	}
}
<commit_msg>Clean up dispatch tests<commit_after>package main

import (
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/doozr/guac"
)

func testDispatchCleanShutDown(t *testing.T, dispatcher Dispatcher) {
	done := make(DoneChan)
	events := make(guac.EventChan)
	waitGroup := sync.WaitGroup{}

	abort := dispatch(dispatcher, events, done, &waitGroup)
	select {
	case <-abort:
	case <-time.After(2 * time.Second):
		t.Fatal("Expected abort within 2 seconds")
	}

	waitGroup.Wait()
}

func TestDispatchRunsDispatcherInBackground(t *testing.T) {
	called := false
	dispatcher := func(events guac.EventChan, done DoneChan) error {
		called = true
		return nil
	}

	testDispatchCleanShutDown(t, dispatcher)
	if !called {
		t.Fatal("Dispatcher not called")
	}
}

func TestDispatchShutDownCleanlyWhenDispatcherReturnsError(t *testing.T) {
	dispatcher := func(events guac.EventChan, done DoneChan) error {
		return fmt.Errorf("Error!")
	}
	testDispatchCleanShutDown(t, dispatcher)
}

func testMessageDispatch(handleMessage MessageHandler, handleUserChange UserChangeHandler) error {
	done := make(DoneChan)
	events := make(guac.EventChan)

	go func() {
		events <- guac.MessageEvent{
			Text: "test event",
		}
		close(done)
	}()

	dispatcher := createDispatcher(1*time.Millisecond, handleMessage, handleUserChange)
	return dispatcher(events, done)
}

func TestDispatcherSendsMessagesToMessageHandler(t *testing.T) {
	var received *guac.MessageEvent
	handleMessage := func(event guac.MessageEvent) error {
		received = &event
		return nil
	}
	handleUserChange := func(event guac.UserInfo) {
	}

	testMessageDispatch(handleMessage, handleUserChange)
	if received == nil || received.Text != "test event" {
		t.Fatal("Did not receive expected message")
	}
}

func TestDispatcherSendsMessagesOnlyOnce(t *testing.T) {
	calls := 0
	handleMessage := func(event guac.MessageEvent) error {
		calls++
		return nil
	}
	handleUserChange := func(event guac.UserInfo) {
	}

	testMessageDispatch(handleMessage, handleUserChange)
	if calls != 1 {
		t.Fatalf("Expected handler to be called exactly once, was called %d times", calls)
	}
}

func TestDispatcherDoesNotSendMessageToUserChangeHandler(t *testing.T) {
	calls := 0
	handleMessage := func(event guac.MessageEvent) error {
		return nil
	}
	handleUserChange := func(event guac.UserInfo) {
		calls++
	}

	testMessageDispatch(handleMessage, handleUserChange)
	if calls != 0 {
		t.Fatalf("Expected handler not to be called, was called %d times", calls)
	}
}

func TestDispatcherReturnsErrorIfMessageFails(t *testing.T) {
	handleMessage := func(event guac.MessageEvent) error {
		return fmt.Errorf("Error!")
	}
	handleUserChange := func(event guac.UserInfo) {
	}

	err := testMessageDispatch(handleMessage, handleUserChange)
	if err == nil {
		t.Fatal("Expected error ", err)
	}
}

func testUserChangeDispatch(handleMessage MessageHandler, handleUserChange UserChangeHandler) error {
	done := make(DoneChan)
	events := make(guac.EventChan)

	go func() {
		events <- guac.UserChangeEvent{
			UserInfo: guac.UserInfo{Name: "test event"},
		}
		close(done)
	}()

	dispatcher := createDispatcher(1*time.Millisecond, handleMessage, handleUserChange)
	return dispatcher(events, done)
}

func TestDispatcherSendsUserChangesToUserChangeHandler(t *testing.T) {
	var received *guac.UserInfo
	handleMessage := func(event guac.MessageEvent) error {
		return nil
	}
	handleUserChange := func(event guac.UserInfo) {
		received = &event
	}

	testUserChangeDispatch(handleMessage, handleUserChange)
	if received == nil || received.Name != "test event" {
		t.Fatal("Did not receive expected message")
	}
}

func TestDispatcherDoesNotSendUserChangeToMessageHandler(t *testing.T) {
	calls := 0
	handleMessage := func(event guac.MessageEvent) error {
		calls++
		return nil
	}
	handleUserChange := func(event guac.UserInfo) {
	}

	testUserChangeDispatch(handleMessage, handleUserChange)
	if calls != 0 {
		t.Fatalf("Expected handler not to be called, was called %d times", calls)
	}
}

func TestDispatcherSendsUserChangeOnlyOnce(t *testing.T) {
	calls := 0
	handleMessage := func(event guac.MessageEvent) error {
		return nil
	}
	handleUserChange := func(event guac.UserInfo) {
		calls++
	}

	testUserChangeDispatch(handleMessage, handleUserChange)
	if calls != 1 {
		t.Fatalf("Expected handler to be called exactly once, was called %d times", calls)
	}
}

func TestDispatcherReturnsErrorOnTimeout(t *testing.T) {
	done := make(DoneChan)
	events := make(guac.EventChan)

	handleMessage := func(event guac.MessageEvent) error {
		return nil
	}
	handleUserChange := func(event guac.UserInfo) {
	}
	dispatcher := createDispatcher(1*time.Millisecond, handleMessage, handleUserChange)

	err := dispatcher(events, done)
	if err == nil {
		t.Fatal("Expected error ", err)
	}
}

func TestDispatcherReturnsNoErrorIfDone(t *testing.T) {
	done := make(DoneChan)
	events := make(guac.EventChan)

	handleMessage := func(event guac.MessageEvent) error {
		return nil
	}
	handleUserChange := func(event guac.UserInfo) {
	}
	dispatcher := createDispatcher(1*time.Millisecond, handleMessage, handleUserChange)

	close(done)
	err := dispatcher(events, done)
	if err != nil {
		t.Fatal("Unexpected error ", err)
	}
}

func TestDispatcherSwallowsUnknownEvents(t *testing.T) {
	done := make(DoneChan)
	events := make(guac.EventChan)

	handleMessage := func(event guac.MessageEvent) error {
		t.Fatal("Unexpected call to MessageHandler")
		return nil
	}
	handleUserChange := func(event guac.UserInfo) {
		t.Fatal("Unexpected call to MessageHandler")
	}
	dispatcher := createDispatcher(1*time.Second, handleMessage, handleUserChange)

	// events is blocking so these things must be read in sequence
	go func() {
		events <- "unknown event"
		close(done)
	}()
	dispatcher(events, done)

	select {
	case <-events:
		t.Fatal("Expected event to have been swallowed")
	default:
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/spf13/pflag"
	"golang.org/x/sync/errgroup"
)

var opts = struct {
	Verbose   bool
	SourceDir string
	OutputDir string
	Version   string
}{}

func init() {
	pflag.BoolVarP(&opts.Verbose, "verbose", "v", false, "be verbose")
	pflag.StringVarP(&opts.SourceDir, "source", "s", "/restic", "path to the source code `directory`")
	pflag.StringVarP(&opts.OutputDir, "output", "o", "/output", "path to the output `directory`")
	pflag.StringVar(&opts.Version, "version", "", "use `x.y.z` as the version for output files")
	pflag.Parse()
}

func die(f string, args ...interface{}) {
	if !strings.HasSuffix(f, "\n") {
		f += "\n"
	}
	f = "\x1b[31m" + f + "\x1b[0m"
	fmt.Fprintf(os.Stderr, f, args...)
	os.Exit(1)
}

func msg(f string, args ...interface{}) {
	if !strings.HasSuffix(f, "\n") {
		f += "\n"
	}
	f = "\x1b[32m" + f + "\x1b[0m"
	fmt.Printf(f, args...)
}

func verbose(f string, args ...interface{}) {
	if !opts.Verbose {
		return
	}
	if !strings.HasSuffix(f, "\n") {
		f += "\n"
	}
	f = "\x1b[32m" + f + "\x1b[0m"
	fmt.Printf(f, args...)
}

func rm(file string) {
	err := os.Remove(file)

	if os.IsNotExist(err) {
		err = nil
	}

	if err != nil {
		die("error removing %v: %v", file, err)
	}
}

func mkdir(dir string) {
	err := os.MkdirAll(dir, 0755)
	if err != nil {
		die("mkdir %v: %v", dir, err)
	}
}

func abs(dir string) string {
	absDir, err := filepath.Abs(dir)
	if err != nil {
		die("unable to find absolute path for %v: %v", dir, err)
	}
	return absDir
}

func build(sourceDir, outputDir, goos, goarch string) (filename string) {
	filename = fmt.Sprintf("%v_%v_%v", "restic", goos, goarch)

	if opts.Version != "" {
		filename = fmt.Sprintf("%v_%v_%v_%v", "restic", opts.Version, goos, goarch)
	}

	if goos == "windows" {
		filename += ".exe"
	}
	outputFile := filepath.Join(outputDir, filename)

	c := exec.Command("go", "build",
		"-o", outputFile,
		"-ldflags", "-s -w",
		"-tags", "selfupdate",
		"./cmd/restic",
	)
	c.Stdout = os.Stdout
	c.Stderr = os.Stderr

	verbose("run %v %v in %v", "go", c.Args, c.Dir)

	c.Dir = sourceDir
	c.Env = append(os.Environ(),
		"CGO_ENABLED=0",
		"GOOS="+goos,
		"GOARCH="+goarch,
	)

	err := c.Run()
	if err != nil {
		die("error building %v/%v: %v", goos, goarch, err)
	}

	return filename
}

func modTime(file string) time.Time {
	fi, err := os.Lstat(file)
	if err != nil {
		die("unable to get modtime of %v: %v", file, err)
	}

	return fi.ModTime()
}

func touch(file string, t time.Time) {
	err := os.Chtimes(file, t, t)
	if err != nil {
		die("unable to update timestamps for %v: %v", file, err)
	}
}

func chmod(file string, mode os.FileMode) {
	err := os.Chmod(file, mode)
	if err != nil {
		die("unable to chmod %v to %s: %v", file, mode, err)
	}
}

func compress(goos, inputDir, filename string) (outputFile string) {
	var c *exec.Cmd
	switch goos {
	case "windows":
		outputFile = strings.TrimSuffix(filename, ".exe") + ".zip"
		c = exec.Command("zip", "-q", "-X", outputFile, filename)
	default:
		outputFile = filename + ".bz2"
		c = exec.Command("bzip2", filename)
	}

	rm(filepath.Join(inputDir, outputFile))

	c.Stdout = os.Stdout
	c.Stderr = os.Stderr
	c.Dir = inputDir

	verbose("run %v %v in %v", "go", c.Args, c.Dir)

	err := c.Run()
	if err != nil {
		die("error compressing: %v", err)
	}

	rm(filepath.Join(inputDir, filename))

	return outputFile
}

func buildForTarget(sourceDir, outputDir, goos, goarch string) (filename string) {
	mtime := modTime(filepath.Join(sourceDir, "VERSION"))

	filename = build(sourceDir, outputDir, goos, goarch)
	touch(filepath.Join(outputDir, filename), mtime)
	chmod(filepath.Join(outputDir, filename), 0755)
	filename = compress(goos, outputDir, filename)
	return filename
}

func buildTargets(sourceDir, outputDir string, targets map[string][]string) {
	start := time.Now()
	msg("building with %d workers", runtime.NumCPU())

	type Job struct{ GOOS, GOARCH string }

	var wg errgroup.Group
	ch := make(chan Job)

	for i := 0; i < runtime.NumCPU(); i++ {
		wg.Go(func() error {
			for job := range ch {
				start := time.Now()
				verbose("build %v/%v", job.GOOS, job.GOARCH)
				buildForTarget(sourceDir, outputDir, job.GOOS, job.GOARCH)
				msg("built %v/%v in %.3fs", job.GOOS, job.GOARCH, time.Since(start).Seconds())
			}
			return nil
		})
	}

	wg.Go(func() error {
		for goos, archs := range targets {
			for _, goarch := range archs {
				ch <- Job{goos, goarch}
			}
		}
		close(ch)
		return nil
	})

	_ = wg.Wait()
	msg("build finished in %.3fs", time.Since(start).Seconds())
}

// ATTENTION: the list of architectures must be in sync with .github/workflows/tests.yml!
var defaultBuildTargets = map[string][]string{
	"aix":     {"ppc64"},
	"darwin":  {"amd64", "arm64"},
	"freebsd": {"386", "amd64", "arm"},
	"linux":   {"386", "amd64", "arm", "arm64", "ppc64le", "mips", "mipsle", "mips64", "mips64le", "s390x"},
	"netbsd":  {"386", "amd64"},
	"openbsd": {"386", "amd64"},
	"windows": {"386", "amd64"},
	"solaris": {"amd64"},
}

func main() {
	if len(pflag.Args()) != 0 {
		die("USAGE: build-release-binaries [OPTIONS]")
	}

	sourceDir := abs(opts.SourceDir)
	outputDir := abs(opts.OutputDir)
	mkdir(outputDir)

	buildTargets(sourceDir, outputDir, defaultBuildTargets)
}
<commit_msg>helper: cleanups<commit_after>package main

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/spf13/pflag"
	"golang.org/x/sync/errgroup"
)

var opts = struct {
	Verbose   bool
	SourceDir string
	OutputDir string
	Version   string
}{}

func init() {
	pflag.BoolVarP(&opts.Verbose, "verbose", "v", false, "be verbose")
	pflag.StringVarP(&opts.SourceDir, "source", "s", "/restic", "path to the source code `directory`")
	pflag.StringVarP(&opts.OutputDir, "output", "o", "/output", "path to the output `directory`")
	pflag.StringVar(&opts.Version, "version", "", "use `x.y.z` as the version for output files")
	pflag.Parse()
}

func die(f string, args ...interface{}) {
	if !strings.HasSuffix(f, "\n") {
		f += "\n"
	}
	f = "\x1b[31m" + f + "\x1b[0m"
	fmt.Fprintf(os.Stderr, f, args...)
	os.Exit(1)
}

func msg(f string, args ...interface{}) {
	if !strings.HasSuffix(f, "\n") {
		f += "\n"
	}
	f = "\x1b[32m" + f + "\x1b[0m"
	fmt.Printf(f, args...)
}

func verbose(f string, args ...interface{}) {
	if !opts.Verbose {
		return
	}
	if !strings.HasSuffix(f, "\n") {
		f += "\n"
	}
	f = "\x1b[32m" + f + "\x1b[0m"
	fmt.Printf(f, args...)
}

func rm(file string) {
	err := os.Remove(file)

	if os.IsNotExist(err) {
		err = nil
	}

	if err != nil {
		die("error removing %v: %v", file, err)
	}
}

func mkdir(dir string) {
	err := os.MkdirAll(dir, 0755)
	if err != nil {
		die("mkdir %v: %v", dir, err)
	}
}

func abs(dir string) string {
	absDir, err := filepath.Abs(dir)
	if err != nil {
		die("unable to find absolute path for %v: %v", dir, err)
	}
	return absDir
}

func build(sourceDir, outputDir, goos, goarch string) (filename string) {
	filename = fmt.Sprintf("%v_%v_%v", "restic", goos, goarch)

	if opts.Version != "" {
		filename = fmt.Sprintf("%v_%v_%v_%v", "restic", opts.Version, goos, goarch)
	}

	if goos == "windows" {
		filename += ".exe"
	}
	outputFile := filepath.Join(outputDir, filename)

	c := exec.Command("go", "build",
		"-o", outputFile,
		"-ldflags", "-s -w",
		"-tags", "selfupdate",
		"./cmd/restic",
	)
	c.Stdout = os.Stdout
	c.Stderr = os.Stderr
	c.Dir = sourceDir
	c.Env = append(os.Environ(),
		"CGO_ENABLED=0",
		"GOOS="+goos,
		"GOARCH="+goarch,
	)
	verbose("run %v %v in %v", "go", c.Args, c.Dir)

	err := c.Run()
	if err != nil {
		die("error building %v/%v: %v", goos, goarch, err)
	}

	return filename
}

func modTime(file string) time.Time {
	fi, err := os.Lstat(file)
	if err != nil {
		die("unable to get modtime of %v: %v", file, err)
	}

	return fi.ModTime()
}

func touch(file string, t time.Time) {
	err := os.Chtimes(file, t, t)
	if err != nil {
		die("unable to update timestamps for %v: %v", file, err)
	}
}

func chmod(file string, mode os.FileMode) {
	err := os.Chmod(file, mode)
	if err != nil {
		die("unable to chmod %v to %s: %v", file, mode, err)
	}
}

func compress(goos, inputDir, filename string) (outputFile string) {
	var c *exec.Cmd
	switch goos {
	case "windows":
		outputFile = strings.TrimSuffix(filename, ".exe") + ".zip"
		c = exec.Command("zip", "-q", "-X", outputFile, filename)
	default:
		outputFile = filename + ".bz2"
		c = exec.Command("bzip2", filename)
	}

	rm(filepath.Join(inputDir, outputFile))

	c.Stdout = os.Stdout
	c.Stderr = os.Stderr
	c.Dir = inputDir
	verbose("run %v %v in %v", "go", c.Args, c.Dir)

	err := c.Run()
	if err != nil {
		die("error compressing: %v", err)
	}

	rm(filepath.Join(inputDir, filename))

	return outputFile
}

func buildForTarget(sourceDir, outputDir, goos, goarch string) (filename string) {
	mtime := modTime(filepath.Join(sourceDir, "VERSION"))

	filename = build(sourceDir, outputDir, goos, goarch)
	touch(filepath.Join(outputDir, filename), mtime)
	chmod(filepath.Join(outputDir, filename), 0755)
	filename = compress(goos, outputDir, filename)
	return filename
}

func buildTargets(sourceDir, outputDir string, targets map[string][]string) {
	start := time.Now()
	msg("building with %d workers", runtime.NumCPU())

	type Job struct{ GOOS, GOARCH string }

	var wg errgroup.Group
	ch := make(chan Job)

	for i := 0; i < runtime.NumCPU(); i++ {
		wg.Go(func() error {
			for job := range ch {
				start := time.Now()
				verbose("build %v/%v", job.GOOS, job.GOARCH)
				buildForTarget(sourceDir, outputDir, job.GOOS, job.GOARCH)
				msg("built %v/%v in %.3fs", job.GOOS, job.GOARCH, time.Since(start).Seconds())
			}
			return nil
		})
	}

	wg.Go(func() error {
		for goos, archs := range targets {
			for _, goarch := range archs {
				ch <- Job{goos, goarch}
			}
		}
		close(ch)
		return nil
	})

	_ = wg.Wait()
	msg("build finished in %.3fs", time.Since(start).Seconds())
}

// ATTENTION: the list of architectures must be in sync with .github/workflows/tests.yml!
var defaultBuildTargets = map[string][]string{
	"aix":     {"ppc64"},
	"darwin":  {"amd64", "arm64"},
	"freebsd": {"386", "amd64", "arm"},
	"linux":   {"386", "amd64", "arm", "arm64", "ppc64le", "mips", "mipsle", "mips64", "mips64le", "s390x"},
	"netbsd":  {"386", "amd64"},
	"openbsd": {"386", "amd64"},
	"windows": {"386", "amd64"},
	"solaris": {"amd64"},
}

func main() {
	if len(pflag.Args()) != 0 {
		die("USAGE: build-release-binaries [OPTIONS]")
	}

	sourceDir := abs(opts.SourceDir)
	outputDir := abs(opts.OutputDir)
	mkdir(outputDir)

	buildTargets(sourceDir, outputDir, defaultBuildTargets)
}
<|endoftext|>