<commit_before>// Copyright 2020 The LUCI 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 ledcmd

import (
	"context"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"sync"
	"time"

	"github.com/bazelbuild/remote-apis-sdks/go/pkg/client"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/command"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/digest"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/filemetadata"
	"github.com/mattn/go-tty"

	"go.chromium.org/luci/auth"
	"go.chromium.org/luci/client/cas"
	"go.chromium.org/luci/client/downloader"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/isolated"
	"go.chromium.org/luci/common/isolatedclient"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/retry"
	apipb "go.chromium.org/luci/swarming/proto/api"

	"go.chromium.org/luci/led/job"
)

// IsolatedTransformer is a function which receives a directory on the local
// disk with the contents of an isolate and is expected to manipulate the
// contents of that directory however it chooses.
//
// EditIsolated takes these functions as a callback in order to manipulate the
// isolated content of a job.Definition.
type IsolatedTransformer func(ctx context.Context, directory string) error

// ProgramIsolatedTransformer returns an IsolatedTransformer which alters the
// contents of the isolated by running a program specified with `args` in the
// directory where the isolated content has been unpacked.
func ProgramIsolatedTransformer(args ...string) IsolatedTransformer {
	return func(ctx context.Context, dir string) error {
		logging.Infof(ctx, "Invoking transform_program: %q", args)
		tProg := exec.CommandContext(ctx, args[0], args[1:]...)
		tProg.Stdout = os.Stderr
		tProg.Stderr = os.Stderr
		tProg.Dir = dir
		return errors.Annotate(tProg.Run(), "running transform_program").Err()
	}
}

// PromptIsolatedTransformer returns an IsolatedTransformer which prompts the
// user to navigate to the directory with the isolated content and manipulate
// it manually. When the user is done they should press "enter" to indicate that
// they're finished.
func PromptIsolatedTransformer() IsolatedTransformer {
	return func(ctx context.Context, dir string) error {
		logging.Infof(ctx, "")
		logging.Infof(ctx, "Edit files as you wish in:")
		logging.Infof(ctx, "\t%s", dir)

		term, err := tty.Open()
		if err != nil {
			return errors.Annotate(err, "opening terminal").Err()
		}
		defer term.Close()

		logging.Infof(ctx, "When finished, press <enter> here to isolate it.")
		_, err = term.ReadString()
		return errors.Annotate(err, "reading <enter>").Err()
	}
}

// EditIsolated allows you to edit the isolated (input_ref or cas_input_root)
// contents of the job.Definition.
//
// This implicitly collapses all isolated sources in the job.Definition into
// a single isolated source.
// The output job.Definition always has cas_user_payload and no user_payload.
func EditIsolated(ctx context.Context, authClient *http.Client, authOpts auth.Options, jd *job.Definition, xform IsolatedTransformer) error {
	logging.Infof(ctx, "editing isolated")

	tdir, err := ioutil.TempDir("", "led-edit-isolated")
	if err != nil {
		return errors.Annotate(err, "failed to create tempdir").Err()
	}
	defer func() {
		if err = os.RemoveAll(tdir); err != nil {
			logging.Errorf(ctx, "failed to cleanup temp dir %q: %s", tdir, err)
		}
	}()

	if err := ConsolidateIsolateSources(ctx, authClient, jd); err != nil {
		return err
	}
	if err := ConsolidateRbeCasSources(ctx, authOpts, jd); err != nil {
		return err
	}

	current, err := jd.Info().CurrentIsolated()
	if err != nil {
		return err
	}
	if current.CASTree != nil && current.CASReference != nil {
		return errors.Reason("job uses isolate and RBE-CAS at the same time - iso: %v\ncas: %v", current.CASTree, current.CASReference).Err()
	}

	err = jd.Edit(func(je job.Editor) {
		je.ClearCurrentIsolated()
	})
	if err != nil {
		return err
	}

	casInstance := current.GetCasInstance()
	if casInstance == "" {
		if casInstance, err = jd.CasInstance(); err != nil {
			return err
		}
	}
	casClient, err := cas.NewClient(ctx, casInstance, authOpts, false)
	if err != nil {
		return err
	}
	defer casClient.Close()

	if err = downloadFromIso(ctx, current.CASTree, authClient, tdir); err != nil {
		return err
	}
	if err = downloadFromCas(ctx, current.CASReference, casClient, tdir); err != nil {
		return err
	}

	if err := xform(ctx, tdir); err != nil {
		return err
	}

	logging.Infof(ctx, "uploading new isolated to RBE-CAS")
	digest, err := uploadToCas(ctx, casClient, tdir)
	if err != nil {
		return errors.Annotate(err, "errors in uploadToCas").Err()
	}
	logging.Infof(ctx, "isolated upload: done")
	jd.CasUserPayload = &apipb.CASReference{
		CasInstance: casInstance,
		Digest: &apipb.Digest{
			Hash:      digest.Hash,
			SizeBytes: digest.Size,
		},
	}
	jd.UserPayload = nil
	return nil
}

func downloadFromCas(ctx context.Context, casRef *apipb.CASReference, casClient *client.Client, tdir string) error {
	if casRef.GetDigest().GetHash() == "" {
		return nil
	}
	d := digest.Digest{
		Hash: casRef.Digest.Hash,
		Size: casRef.Digest.SizeBytes,
	}
	logging.Infof(ctx, "downloading from RBE-CAS...")
	_, err := casClient.DownloadDirectory(ctx, d, tdir, filemetadata.NewNoopCache())
	if err != nil {
		return errors.Annotate(err, "failed to download directory").Err()
	}
	return nil
}

func downloadFromIso(ctx context.Context, iso *apipb.CASTree, authClient *http.Client, tdir string) error {
	if iso.GetDigest() == "" {
		return nil
	}
	rawIsoClient := isolatedclient.NewClient(
		iso.Server,
		isolatedclient.WithAuthClient(authClient),
		isolatedclient.WithNamespace(iso.Namespace),
		isolatedclient.WithRetryFactory(retry.Default))
	var statMu sync.Mutex
	var previousStats *downloader.FileStats

	logging.Infof(ctx, "downloading from isolate...")
	dl := downloader.New(ctx, rawIsoClient, isolated.HexDigest(iso.GetDigest()), tdir, &downloader.Options{
		FileStatsCallback: func(s downloader.FileStats, span time.Duration) {
			logging.Infof(ctx, "%s", s.StatLine(previousStats, span))
			statMu.Lock()
			previousStats = &s
			statMu.Unlock()
		},
	})
	if err := dl.Wait(); err != nil {
		return err
	}
	return nil
}

func uploadToCas(ctx context.Context, client *client.Client, dir string) (*digest.Digest, error) {
	is := command.InputSpec{
		Inputs: []string{"."}, // entire dir
	}
	rootDg, entries, _, err := client.ComputeMerkleTree(dir, &is, filemetadata.NewNoopCache())
	if err != nil {
		return nil, errors.Annotate(err, "failed to compute Merkle Tree").Err()
	}

	_, _, err = client.UploadIfMissing(ctx, entries...)
	if err != nil {
		return nil, errors.Annotate(err, "failed to upload items").Err()
	}
	return &rootDg, nil
}
<commit_msg>led: build fix for deps roll<commit_after>// Copyright 2020 The LUCI 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 ledcmd

import (
	"context"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"sync"
	"time"

	"github.com/bazelbuild/remote-apis-sdks/go/pkg/client"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/command"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/digest"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/filemetadata"
	"github.com/mattn/go-tty"

	"go.chromium.org/luci/auth"
	"go.chromium.org/luci/client/cas"
	"go.chromium.org/luci/client/downloader"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/isolated"
	"go.chromium.org/luci/common/isolatedclient"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/retry"
	apipb "go.chromium.org/luci/swarming/proto/api"

	"go.chromium.org/luci/led/job"
)

// IsolatedTransformer is a function which receives a directory on the local
// disk with the contents of an isolate and is expected to manipulate the
// contents of that directory however it chooses.
//
// EditIsolated takes these functions as a callback in order to manipulate the
// isolated content of a job.Definition.
type IsolatedTransformer func(ctx context.Context, directory string) error

// ProgramIsolatedTransformer returns an IsolatedTransformer which alters the
// contents of the isolated by running a program specified with `args` in the
// directory where the isolated content has been unpacked.
func ProgramIsolatedTransformer(args ...string) IsolatedTransformer {
	return func(ctx context.Context, dir string) error {
		logging.Infof(ctx, "Invoking transform_program: %q", args)
		tProg := exec.CommandContext(ctx, args[0], args[1:]...)
		tProg.Stdout = os.Stderr
		tProg.Stderr = os.Stderr
		tProg.Dir = dir
		return errors.Annotate(tProg.Run(), "running transform_program").Err()
	}
}

// PromptIsolatedTransformer returns an IsolatedTransformer which prompts the
// user to navigate to the directory with the isolated content and manipulate
// it manually. When the user is done they should press "enter" to indicate that
// they're finished.
func PromptIsolatedTransformer() IsolatedTransformer {
	return func(ctx context.Context, dir string) error {
		logging.Infof(ctx, "")
		logging.Infof(ctx, "Edit files as you wish in:")
		logging.Infof(ctx, "\t%s", dir)

		term, err := tty.Open()
		if err != nil {
			return errors.Annotate(err, "opening terminal").Err()
		}
		defer term.Close()

		logging.Infof(ctx, "When finished, press <enter> here to isolate it.")
		_, err = term.ReadString()
		return errors.Annotate(err, "reading <enter>").Err()
	}
}

// EditIsolated allows you to edit the isolated (input_ref or cas_input_root)
// contents of the job.Definition.
//
// This implicitly collapses all isolated sources in the job.Definition into
// a single isolated source.
// The output job.Definition always has cas_user_payload and no user_payload.
func EditIsolated(ctx context.Context, authClient *http.Client, authOpts auth.Options, jd *job.Definition, xform IsolatedTransformer) error {
	logging.Infof(ctx, "editing isolated")

	tdir, err := ioutil.TempDir("", "led-edit-isolated")
	if err != nil {
		return errors.Annotate(err, "failed to create tempdir").Err()
	}
	defer func() {
		if err = os.RemoveAll(tdir); err != nil {
			logging.Errorf(ctx, "failed to cleanup temp dir %q: %s", tdir, err)
		}
	}()

	if err := ConsolidateIsolateSources(ctx, authClient, jd); err != nil {
		return err
	}
	if err := ConsolidateRbeCasSources(ctx, authOpts, jd); err != nil {
		return err
	}

	current, err := jd.Info().CurrentIsolated()
	if err != nil {
		return err
	}
	if current.CASTree != nil && current.CASReference != nil {
		return errors.Reason("job uses isolate and RBE-CAS at the same time - iso: %v\ncas: %v", current.CASTree, current.CASReference).Err()
	}

	err = jd.Edit(func(je job.Editor) {
		je.ClearCurrentIsolated()
	})
	if err != nil {
		return err
	}

	casInstance := current.GetCasInstance()
	if casInstance == "" {
		if casInstance, err = jd.CasInstance(); err != nil {
			return err
		}
	}
	casClient, err := cas.NewClient(ctx, casInstance, authOpts, false)
	if err != nil {
		return err
	}
	defer casClient.Close()

	if err = downloadFromIso(ctx, current.CASTree, authClient, tdir); err != nil {
		return err
	}
	if err = downloadFromCas(ctx, current.CASReference, casClient, tdir); err != nil {
		return err
	}

	if err := xform(ctx, tdir); err != nil {
		return err
	}

	logging.Infof(ctx, "uploading new isolated to RBE-CAS")
	digest, err := uploadToCas(ctx, casClient, tdir)
	if err != nil {
		return errors.Annotate(err, "errors in uploadToCas").Err()
	}
	logging.Infof(ctx, "isolated upload: done")
	jd.CasUserPayload = &apipb.CASReference{
		CasInstance: casInstance,
		Digest: &apipb.Digest{
			Hash:      digest.Hash,
			SizeBytes: digest.Size,
		},
	}
	jd.UserPayload = nil
	return nil
}

func downloadFromCas(ctx context.Context, casRef *apipb.CASReference, casClient *client.Client, tdir string) error {
	if casRef.GetDigest().GetHash() == "" {
		return nil
	}
	d := digest.Digest{
		Hash: casRef.Digest.Hash,
		Size: casRef.Digest.SizeBytes,
	}
	logging.Infof(ctx, "downloading from RBE-CAS...")
	_, _, err := casClient.DownloadDirectory(ctx, d, tdir, filemetadata.NewNoopCache())
	if err != nil {
		return errors.Annotate(err, "failed to download directory").Err()
	}
	return nil
}

func downloadFromIso(ctx context.Context, iso *apipb.CASTree, authClient *http.Client, tdir string) error {
	if iso.GetDigest() == "" {
		return nil
	}
	rawIsoClient := isolatedclient.NewClient(
		iso.Server,
		isolatedclient.WithAuthClient(authClient),
		isolatedclient.WithNamespace(iso.Namespace),
		isolatedclient.WithRetryFactory(retry.Default))
	var statMu sync.Mutex
	var previousStats *downloader.FileStats

	logging.Infof(ctx, "downloading from isolate...")
	dl := downloader.New(ctx, rawIsoClient, isolated.HexDigest(iso.GetDigest()), tdir, &downloader.Options{
		FileStatsCallback: func(s downloader.FileStats, span time.Duration) {
			logging.Infof(ctx, "%s", s.StatLine(previousStats, span))
			statMu.Lock()
			previousStats = &s
			statMu.Unlock()
		},
	})
	if err := dl.Wait(); err != nil {
		return err
	}
	return nil
}

func uploadToCas(ctx context.Context, client *client.Client, dir string) (*digest.Digest, error) {
	is := command.InputSpec{
		Inputs: []string{"."}, // entire dir
	}
	rootDg, entries, _, err := client.ComputeMerkleTree(dir, &is, filemetadata.NewNoopCache())
	if err != nil {
		return nil, errors.Annotate(err, "failed to compute Merkle Tree").Err()
	}

	_, _, err = client.UploadIfMissing(ctx, entries...)
	if err != nil {
		return nil, errors.Annotate(err, "failed to upload items").Err()
	}
	return &rootDg, nil
}
<|endoftext|>


<commit_before>// Copyright 2016 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package plan

import (
	"github.com/juju/errors"
	"github.com/pingcap/tidb/expression"
)

func retrieveColumnsInExpression(expr expression.Expression, schema expression.Schema) (
	expression.Expression, error) {
	switch v := expr.(type) {
	case *expression.ScalarFunction:
		for i, arg := range v.Args {
			newExpr, err := retrieveColumnsInExpression(arg, schema)
			if err != nil {
				return nil, errors.Trace(err)
			}
			v.Args[i] = newExpr
		}
	case *expression.Column:
		if !v.Correlated {
			newColumn := schema.RetrieveColumn(v)
			if newColumn == nil {
				return nil, errors.Errorf("Can't Find column %s.", expr.ToString())
			}
			return newColumn, nil
		}
	}
	return expr, nil
}

func makeUsedList(usedCols []*expression.Column, schema expression.Schema) []bool {
	used := make([]bool, len(schema))
	for _, col := range usedCols {
		idx := schema.GetIndex(col)
		used[idx] = true
	}
	return used
}

// PruneColumnsAndResolveIndices prunes unused columns and resolves index for columns.
// This function returns a column slice representing outer columns and an error.
func pruneColumnsAndResolveIndices(p Plan, parentUsedCols []*expression.Column) ([]*expression.Column, error) {
	switch v := p.(type) {
	case *Projection:
		// Prune
		var cols, outerCols []*expression.Column
		used := makeUsedList(parentUsedCols, p.GetSchema())
		for i := len(used) - 1; i >= 0; i-- {
			if !used[i] {
				v.schema = append(v.schema[:i], v.schema[i+1:]...)
				v.Exprs = append(v.Exprs[:i], v.Exprs[i+1:]...)
			}
		}
		v.schema.InitIndices()
		for _, expr := range v.Exprs {
			cols, outerCols = extractColumn(expr, cols, outerCols)
		}
		outer, err := pruneColumnsAndResolveIndices(p.GetChildByIndex(0), cols)
		if err != nil {
			return nil, errors.Trace(err)
		}
		for i, expr := range v.Exprs {
			v.Exprs[i], err = retrieveColumnsInExpression(expr, p.GetChildByIndex(0).GetSchema())
			if err != nil {
				return nil, errors.Trace(err)
			}
		}
		return append(outer, outerCols...), nil
	case *Selection:
		var outerCols []*expression.Column
		for _, cond := range v.Conditions {
			parentUsedCols, outerCols = extractColumn(cond, parentUsedCols, outerCols)
		}
		outer, err := pruneColumnsAndResolveIndices(p.GetChildByIndex(0), parentUsedCols)
		if err != nil {
			return nil, errors.Trace(err)
		}
		v.SetSchema(p.GetChildByIndex(0).GetSchema())
		for i, cond := range v.Conditions {
			v.Conditions[i], err = retrieveColumnsInExpression(cond, p.GetChildByIndex(0).GetSchema())
			if err != nil {
				return nil, errors.Trace(err)
			}
		}
		return append(outer, outerCols...), nil
	case *Apply:
		return pruneApply(v, parentUsedCols)
	case *Aggregation:
		used := makeUsedList(parentUsedCols, p.GetSchema())
		for i := len(used) - 1; i >= 0; i-- {
			if !used[i] {
				v.schema = append(v.schema[:i], v.schema[i+1:]...)
				v.AggFuncs = append(v.AggFuncs[:i], v.AggFuncs[i+1:]...)
			}
		}
		var cols, outerCols []*expression.Column
		for _, aggrFunc := range v.AggFuncs {
			for _, arg := range aggrFunc.GetArgs() {
				cols, outerCols = extractColumn(arg, cols, outerCols)
			}
		}
		for _, expr := range v.GroupByItems {
			cols, outerCols = extractColumn(expr, cols, outerCols)
		}
		outer, err := pruneColumnsAndResolveIndices(p.GetChildByIndex(0), cols)
		if err != nil {
			return nil, errors.Trace(err)
		}
		for _, aggrFunc := range v.AggFuncs {
			for i, arg := range aggrFunc.GetArgs() {
				var newArg expression.Expression
				newArg, err = retrieveColumnsInExpression(arg, p.GetChildByIndex(0).GetSchema())
				if err != nil {
					return nil, errors.Trace(err)
				}
				aggrFunc.SetArgs(i, newArg)
			}
		}
		v.schema.InitIndices()
		return append(outer, outerCols...), nil
	case *NewSort:
		var outerCols []*expression.Column
		for _, item := range v.ByItems {
			parentUsedCols, outerCols = extractColumn(item.Expr, parentUsedCols, outerCols)
		}
		outer, err := pruneColumnsAndResolveIndices(p.GetChildByIndex(0), parentUsedCols)
		if err != nil {
			return nil, errors.Trace(err)
		}
		v.SetSchema(p.GetChildByIndex(0).GetSchema())
		for _, item := range v.ByItems {
			item.Expr, err = retrieveColumnsInExpression(item.Expr, p.GetChildByIndex(0).GetSchema())
			if err != nil {
				return nil, errors.Trace(err)
			}
		}
		return append(outer, outerCols...), nil
	case *Union:
		var outerCols []*expression.Column
		used := makeUsedList(parentUsedCols, p.GetSchema())
		for i := len(used) - 1; i >= 0; i-- {
			if !used[i] {
				v.schema = append(v.schema[:i], v.schema[i+1:]...)
			}
		}
		v.schema.InitIndices()
		for _, child := range p.GetChildren() {
			schema := child.GetSchema()
			var newSchema []*expression.Column
			for i, use := range used {
				if use {
					newSchema = append(newSchema, schema[i])
				}
			}
			outer, err := pruneColumnsAndResolveIndices(child, newSchema)
			outerCols = append(outerCols, outer...)
			if err != nil {
				return nil, errors.Trace(err)
			}
		}
		return outerCols, nil
	case *NewTableScan:
		used := makeUsedList(parentUsedCols, p.GetSchema())
		for i := len(used) - 1; i >= 0; i-- {
			if !used[i] {
				v.schema = append(v.schema[:i], v.schema[i+1:]...)
				v.Columns = append(v.Columns[:i], v.Columns[i+1:]...)
			}
		}
		v.schema.InitIndices()
		return nil, nil
	case *Limit, *MaxOneRow, *Distinct:
		outer, err := pruneColumnsAndResolveIndices(p.GetChildByIndex(0), parentUsedCols)
		p.SetSchema(p.GetChildByIndex(0).GetSchema())
		return outer, errors.Trace(err)
	case *Trim:
		used := makeUsedList(parentUsedCols, v.schema)
		for i := len(used) - 1; i >= 0; i-- {
			if !used[i] {
				v.schema = append(v.schema[:i], v.schema[i+1:]...)
			}
		}
		outer, err := pruneColumnsAndResolveIndices(p.GetChildByIndex(0), parentUsedCols)
		return outer, errors.Trace(err)
	case *Exists:
		return pruneColumnsAndResolveIndices(p.GetChildByIndex(0), nil)
	case *Join:
		var outerCols []*expression.Column
		for _, eqCond := range v.EqualConditions {
			parentUsedCols, outerCols = extractColumn(eqCond, parentUsedCols, outerCols)
		}
		for _, leftCond := range v.LeftConditions {
			parentUsedCols, outerCols = extractColumn(leftCond, parentUsedCols, outerCols)
		}
		for _, rightCond := range v.RightConditions {
			parentUsedCols, outerCols = extractColumn(rightCond, parentUsedCols, outerCols)
		}
		for _, otherCond := range v.OtherConditions {
			parentUsedCols, outerCols = extractColumn(otherCond, parentUsedCols, outerCols)
		}
		var leftCols, rightCols []*expression.Column
		for _, col := range parentUsedCols {
			if p.GetChildByIndex(0).GetSchema().GetIndex(col) != -1 {
				leftCols = append(leftCols, col)
			} else {
				rightCols = append(rightCols, col)
			}
		}
		outerLeft, err := pruneColumnsAndResolveIndices(p.GetChildByIndex(0), leftCols)
		outerCols = append(outerCols, outerLeft...)
		if err != nil {
			return nil, errors.Trace(err)
		}
		for i, leftCond := range v.LeftConditions {
			v.LeftConditions[i], err = retrieveColumnsInExpression(leftCond, p.GetChildByIndex(0).GetSchema())
			if err != nil {
				return nil, errors.Trace(err)
			}
		}
		outerRight, err := pruneColumnsAndResolveIndices(p.GetChildByIndex(1), rightCols)
		outerCols = append(outerCols, outerRight...)
		if err != nil {
			return nil, errors.Trace(err)
		}
		for i, rightCond := range v.RightConditions {
			v.RightConditions[i], err = retrieveColumnsInExpression(rightCond, p.GetChildByIndex(1).GetSchema())
			if err != nil {
				return nil, errors.Trace(err)
			}
		}
		v.schema = append(v.GetChildByIndex(0).GetSchema().DeepCopy(), v.GetChildByIndex(1).GetSchema().DeepCopy()...)
		v.schema.InitIndices()
		for i, otherCond := range v.OtherConditions {
			v.OtherConditions[i], err = retrieveColumnsInExpression(otherCond, p.GetSchema())
			if err != nil {
				return nil, errors.Trace(err)
			}
		}
		for _, eqCond := range v.EqualConditions {
			eqCond.Args[0], err = retrieveColumnsInExpression(eqCond.Args[0], p.GetChildByIndex(0).GetSchema())
			if err != nil {
				return nil, errors.Trace(err)
			}
			eqCond.Args[1], err = retrieveColumnsInExpression(eqCond.Args[1], p.GetChildByIndex(1).GetSchema())
			if err != nil {
				return nil, errors.Trace(err)
			}
		}
		return outerCols, nil
	default:
		return nil, nil
	}
}

// e.g. For query select b.c ,(select count(*) from a where a.id = b.id) from b. Its plan is Projection->Apply->TableScan.
// The schema of b is (a,b,c,id). When Pruning Apply, the parentUsedCols is (c, extra), outerSchema is (a,b,c,id).
// Then after pruning inner plan, the outer schema in apply becomes (id).
// Now there're two columns in parentUsedCols, c is the column from Apply's child ---- TableScan, but extra isn't.
// So only c in parentUsedCols and id in outerSchema can be passed to TableScan.
func pruneApply(v *Apply, parentUsedCols []*expression.Column) ([]*expression.Column, error) {
	outer, err := pruneColumnsAndResolveIndices(v.InnerPlan, v.InnerPlan.GetSchema())
	if err != nil {
		return nil, errors.Trace(err)
	}
	used := makeUsedList(outer, v.OuterSchema)
	for i := len(used) - 1; i >= 0; i-- {
		if !used[i] {
			v.OuterSchema = append(v.OuterSchema[:i], v.OuterSchema[i+1:]...)
		}
	}
	newUsedCols := v.OuterSchema
	for _, used := range parentUsedCols {
		if v.GetChildByIndex(0).GetSchema().GetIndex(used) != -1 {
			newUsedCols = append(newUsedCols, used)
		}
	}
	if v.Checker != nil {
		condUsedCols, _ := extractColumn(v.Checker.Condition, nil, nil)
		for _, used := range condUsedCols {
			if v.GetChildByIndex(0).GetSchema().GetIndex(used) != -1 {
				newUsedCols = append(newUsedCols, used)
			}
		}
	}
	outer, err = pruneColumnsAndResolveIndices(v.GetChildByIndex(0), newUsedCols)
	for _, col := range v.OuterSchema {
		col.Index = v.GetChildByIndex(0).GetSchema().GetIndex(col)
	}
	if err != nil {
		return nil, errors.Trace(err)
	}
	combinedSchema := append(v.GetChildByIndex(0).GetSchema().DeepCopy(), v.InnerPlan.GetSchema().DeepCopy()...)
	if v.Checker == nil {
		v.schema = combinedSchema
	} else {
		combinedSchema.InitIndices()
		v.Checker.Condition, err = retrieveColumnsInExpression(v.Checker.Condition, combinedSchema)
		if err != nil {
			return nil, errors.Trace(err)
		}
		v.schema = append(v.GetChildByIndex(0).GetSchema().DeepCopy(), v.schema[len(v.schema)-1])
	}
	v.schema.InitIndices()
	return outer, nil
}
<commit_msg>plan: tiny clean up<commit_after>// Copyright 2016 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package plan

import (
	"github.com/juju/errors"
	"github.com/pingcap/tidb/expression"
)

func retrieveColumnsInExpression(expr expression.Expression, schema expression.Schema) (
	expression.Expression, error) {
	switch v := expr.(type) {
	case *expression.ScalarFunction:
		for i, arg := range v.Args {
			newExpr, err := retrieveColumnsInExpression(arg, schema)
			if err != nil {
				return nil, errors.Trace(err)
			}
			v.Args[i] = newExpr
		}
	case *expression.Column:
		if !v.Correlated {
			newColumn := schema.RetrieveColumn(v)
			if newColumn == nil {
				return nil, errors.Errorf("Can't Find column %s.", expr.ToString())
			}
			return newColumn, nil
		}
	}
	return expr, nil
}

func makeUsedList(usedCols []*expression.Column, schema expression.Schema) []bool {
	used := make([]bool, len(schema))
	for _, col := range usedCols {
		idx := schema.GetIndex(col)
		used[idx] = true
	}
	return used
}

// pruneColumnsAndResolveIndices prunes unused columns and resolves index for columns.
// This function returns a column slice representing outer columns and an error.
func pruneColumnsAndResolveIndices(p Plan, parentUsedCols []*expression.Column) ([]*expression.Column, error) {
	switch v := p.(type) {
	case *Projection:
		// Prune
		var cols, outerCols []*expression.Column
		used := makeUsedList(parentUsedCols, p.GetSchema())
		for i := len(used) - 1; i >= 0; i-- {
			if !used[i] {
				v.schema = append(v.schema[:i], v.schema[i+1:]...)
				v.Exprs = append(v.Exprs[:i], v.Exprs[i+1:]...)
			}
		}
		v.schema.InitIndices()
		for _, expr := range v.Exprs {
			cols, outerCols = extractColumn(expr, cols, outerCols)
		}
		outer, err := pruneColumnsAndResolveIndices(p.GetChildByIndex(0), cols)
		if err != nil {
			return nil, errors.Trace(err)
		}
		for i, expr := range v.Exprs {
			v.Exprs[i], err = retrieveColumnsInExpression(expr, p.GetChildByIndex(0).GetSchema())
			if err != nil {
				return nil, errors.Trace(err)
			}
		}
		return append(outer, outerCols...), nil
	case *Selection:
		var outerCols []*expression.Column
		for _, cond := range v.Conditions {
			parentUsedCols, outerCols = extractColumn(cond, parentUsedCols, outerCols)
		}
		outer, err := pruneColumnsAndResolveIndices(p.GetChildByIndex(0), parentUsedCols)
		if err != nil {
			return nil, errors.Trace(err)
		}
		v.SetSchema(p.GetChildByIndex(0).GetSchema())
		for i, cond := range v.Conditions {
			v.Conditions[i], err = retrieveColumnsInExpression(cond, p.GetChildByIndex(0).GetSchema())
			if err != nil {
				return nil, errors.Trace(err)
			}
		}
		return append(outer, outerCols...), nil
	case *Apply:
		return pruneApply(v, parentUsedCols)
	case *Aggregation:
		used := makeUsedList(parentUsedCols, p.GetSchema())
		for i := len(used) - 1; i >= 0; i-- {
			if !used[i] {
				v.schema = append(v.schema[:i], v.schema[i+1:]...)
				v.AggFuncs = append(v.AggFuncs[:i], v.AggFuncs[i+1:]...)
			}
		}
		var cols, outerCols []*expression.Column
		for _, aggrFunc := range v.AggFuncs {
			for _, arg := range aggrFunc.GetArgs() {
				cols, outerCols = extractColumn(arg, cols, outerCols)
			}
		}
		for _, expr := range v.GroupByItems {
			cols, outerCols = extractColumn(expr, cols, outerCols)
		}
		outer, err := pruneColumnsAndResolveIndices(p.GetChildByIndex(0), cols)
		if err != nil {
			return nil, errors.Trace(err)
		}
		for _, aggrFunc := range v.AggFuncs {
			for i, arg := range aggrFunc.GetArgs() {
				var newArg expression.Expression
				newArg, err = retrieveColumnsInExpression(arg, p.GetChildByIndex(0).GetSchema())
				if err != nil {
					return nil, errors.Trace(err)
				}
				aggrFunc.SetArgs(i, newArg)
			}
		}
		v.schema.InitIndices()
		return append(outer, outerCols...), nil
	case *NewSort:
		var outerCols []*expression.Column
		for _, item := range v.ByItems {
			parentUsedCols, outerCols = extractColumn(item.Expr, parentUsedCols, outerCols)
		}
		outer, err := pruneColumnsAndResolveIndices(p.GetChildByIndex(0), parentUsedCols)
		if err != nil {
			return nil, errors.Trace(err)
		}
		v.SetSchema(p.GetChildByIndex(0).GetSchema())
		for _, item := range v.ByItems {
			item.Expr, err = retrieveColumnsInExpression(item.Expr, p.GetChildByIndex(0).GetSchema())
			if err != nil {
				return nil, errors.Trace(err)
			}
		}
		return append(outer, outerCols...), nil
	case *Union:
		var outerCols []*expression.Column
		used := makeUsedList(parentUsedCols, p.GetSchema())
		for i := len(used) - 1; i >= 0; i-- {
			if !used[i] {
				v.schema = append(v.schema[:i], v.schema[i+1:]...)
			}
		}
		v.schema.InitIndices()
		for _, child := range p.GetChildren() {
			schema := child.GetSchema()
			var newSchema []*expression.Column
			for i, use := range used {
				if use {
					newSchema = append(newSchema, schema[i])
				}
			}
			outer, err := pruneColumnsAndResolveIndices(child, newSchema)
			outerCols = append(outerCols, outer...)
			if err != nil {
				return nil, errors.Trace(err)
			}
		}
		return outerCols, nil
	case *NewTableScan:
		used := makeUsedList(parentUsedCols, p.GetSchema())
		for i := len(used) - 1; i >= 0; i-- {
			if !used[i] {
				v.schema = append(v.schema[:i], v.schema[i+1:]...)
				v.Columns = append(v.Columns[:i], v.Columns[i+1:]...)
			}
		}
		v.schema.InitIndices()
		return nil, nil
	case *Limit, *MaxOneRow, *Distinct:
		outer, err := pruneColumnsAndResolveIndices(p.GetChildByIndex(0), parentUsedCols)
		p.SetSchema(p.GetChildByIndex(0).GetSchema())
		return outer, errors.Trace(err)
	case *Trim:
		used := makeUsedList(parentUsedCols, v.schema)
		for i := len(used) - 1; i >= 0; i-- {
			if !used[i] {
				v.schema = append(v.schema[:i], v.schema[i+1:]...)
			}
		}
		outer, err := pruneColumnsAndResolveIndices(p.GetChildByIndex(0), parentUsedCols)
		return outer, errors.Trace(err)
	case *Exists:
		return pruneColumnsAndResolveIndices(p.GetChildByIndex(0), nil)
	case *Join:
		var outerCols []*expression.Column
		for _, eqCond := range v.EqualConditions {
			parentUsedCols, outerCols = extractColumn(eqCond, parentUsedCols, outerCols)
		}
		for _, leftCond := range v.LeftConditions {
			parentUsedCols, outerCols = extractColumn(leftCond, parentUsedCols, outerCols)
		}
		for _, rightCond := range v.RightConditions {
			parentUsedCols, outerCols = extractColumn(rightCond, parentUsedCols, outerCols)
		}
		for _, otherCond := range v.OtherConditions {
			parentUsedCols, outerCols = extractColumn(otherCond, parentUsedCols, outerCols)
		}
		var leftCols, rightCols []*expression.Column
		for _, col := range parentUsedCols {
			if p.GetChildByIndex(0).GetSchema().GetIndex(col) != -1 {
				leftCols = append(leftCols, col)
			} else {
				rightCols = append(rightCols, col)
			}
		}
		outerLeft, err := pruneColumnsAndResolveIndices(p.GetChildByIndex(0), leftCols)
		outerCols = append(outerCols, outerLeft...)
		if err != nil {
			return nil, errors.Trace(err)
		}
		for i, leftCond := range v.LeftConditions {
			v.LeftConditions[i], err = retrieveColumnsInExpression(leftCond, p.GetChildByIndex(0).GetSchema())
			if err != nil {
				return nil, errors.Trace(err)
			}
		}
		outerRight, err := pruneColumnsAndResolveIndices(p.GetChildByIndex(1), rightCols)
		outerCols = append(outerCols, outerRight...)
		if err != nil {
			return nil, errors.Trace(err)
		}
		for i, rightCond := range v.RightConditions {
			v.RightConditions[i], err = retrieveColumnsInExpression(rightCond, p.GetChildByIndex(1).GetSchema())
			if err != nil {
				return nil, errors.Trace(err)
			}
		}
		v.schema = append(v.GetChildByIndex(0).GetSchema().DeepCopy(), v.GetChildByIndex(1).GetSchema().DeepCopy()...)
		v.schema.InitIndices()
		for i, otherCond := range v.OtherConditions {
			v.OtherConditions[i], err = retrieveColumnsInExpression(otherCond, p.GetSchema())
			if err != nil {
				return nil, errors.Trace(err)
			}
		}
		for _, eqCond := range v.EqualConditions {
			eqCond.Args[0], err = retrieveColumnsInExpression(eqCond.Args[0], p.GetChildByIndex(0).GetSchema())
			if err != nil {
				return nil, errors.Trace(err)
			}
			eqCond.Args[1], err = retrieveColumnsInExpression(eqCond.Args[1], p.GetChildByIndex(1).GetSchema())
			if err != nil {
				return nil, errors.Trace(err)
			}
		}
		return outerCols, nil
	default:
		return nil, nil
	}
}

// e.g. For query select b.c ,(select count(*) from a where a.id = b.id) from b. Its plan is Projection->Apply->TableScan.
// The schema of b is (a,b,c,id). When Pruning Apply, the parentUsedCols is (c, extra), outerSchema is (a,b,c,id).
// Then after pruning inner plan, the outer schema in apply becomes (id).
// Now there're two columns in parentUsedCols, c is the column from Apply's child ---- TableScan, but extra isn't.
// So only c in parentUsedCols and id in outerSchema can be passed to TableScan.
func pruneApply(v *Apply, parentUsedCols []*expression.Column) ([]*expression.Column, error) {
	outer, err := pruneColumnsAndResolveIndices(v.InnerPlan, v.InnerPlan.GetSchema())
	if err != nil {
		return nil, errors.Trace(err)
	}
	used := makeUsedList(outer, v.OuterSchema)
	for i := len(used) - 1; i >= 0; i-- {
		if !used[i] {
			v.OuterSchema = append(v.OuterSchema[:i], v.OuterSchema[i+1:]...)
		}
	}
	newUsedCols := v.OuterSchema
	for _, used := range parentUsedCols {
		if v.GetChildByIndex(0).GetSchema().GetIndex(used) != -1 {
			newUsedCols = append(newUsedCols, used)
		}
	}
	if v.Checker != nil {
		condUsedCols, _ := extractColumn(v.Checker.Condition, nil, nil)
		for _, used := range condUsedCols {
			if v.GetChildByIndex(0).GetSchema().GetIndex(used) != -1 {
				newUsedCols = append(newUsedCols, used)
			}
		}
	}
	outer, err = pruneColumnsAndResolveIndices(v.GetChildByIndex(0), newUsedCols)
	for _, col := range v.OuterSchema {
		col.Index = v.GetChildByIndex(0).GetSchema().GetIndex(col)
	}
	if err != nil {
		return nil, errors.Trace(err)
	}
	combinedSchema := append(v.GetChildByIndex(0).GetSchema().DeepCopy(), v.InnerPlan.GetSchema().DeepCopy()...)
	if v.Checker == nil {
		v.schema = combinedSchema
	} else {
		combinedSchema.InitIndices()
		v.Checker.Condition, err = retrieveColumnsInExpression(v.Checker.Condition, combinedSchema)
		if err != nil {
			return nil, errors.Trace(err)
		}
		v.schema = append(v.GetChildByIndex(0).GetSchema().DeepCopy(), v.schema[len(v.schema)-1])
	}
	v.schema.InitIndices()
	return outer, nil
}
<|endoftext|>


<commit_before>package terminal

import (
	"bytes"
	"fmt"
	"io"
	"strconv"
	"strings"
	"syscall"
	"unsafe"

	"github.com/mattn/go-isatty"
)

var (
	cursorFunctions = map[rune]func(c *Cursor) func(int){
		'A': func(c *Cursor) func(int) { return c.Up },
		'B': func(c *Cursor) func(int) { return c.Down },
		'C': func(c *Cursor) func(int) { return c.Forward },
		'D': func(c *Cursor) func(int) { return c.Back },
		'E': func(c *Cursor) func(int) { return c.NextLine },
		'F': func(c *Cursor) func(int) { return c.PreviousLine },
		'G': func(c *Cursor) func(int) { return c.HorizontalAbsolute },
	}
)

const (
	foregroundBlue      = 0x1
	foregroundGreen     = 0x2
	foregroundRed       = 0x4
	foregroundIntensity = 0x8
	foregroundMask      = (foregroundRed | foregroundBlue | foregroundGreen | foregroundIntensity)
	backgroundBlue      = 0x10
	backgroundGreen     = 0x20
	backgroundRed       = 0x40
	backgroundIntensity = 0x80
	backgroundMask      = (backgroundRed | backgroundBlue | backgroundGreen | backgroundIntensity)
)

type Writer struct {
	out     FileWriter
	handle  syscall.Handle
	orgAttr word
}

func NewAnsiStdout(out FileWriter) io.Writer {
	var csbi consoleScreenBufferInfo
	if !isatty.IsTerminal(out.Fd()) {
		return out
	}
	handle := syscall.Handle(out.Fd())
	procGetConsoleScreenBufferInfo.Call(uintptr(handle), uintptr(unsafe.Pointer(&csbi)))
	return &Writer{out: out, handle: handle, orgAttr: csbi.attributes}
}

func NewAnsiStderr(out FileWriter) io.Writer {
	var csbi consoleScreenBufferInfo
	if !isatty.IsTerminal(out.Fd()) {
		return out
	}
	handle := syscall.Handle(out.Fd())
	procGetConsoleScreenBufferInfo.Call(uintptr(handle), uintptr(unsafe.Pointer(&csbi)))
	return &Writer{out: out, handle: handle, orgAttr: csbi.attributes}
}

func (w *Writer) Write(data []byte) (n int, err error) {
	r := bytes.NewReader(data)

	for {
		var ch rune
		var size int
		ch, size, err = r.ReadRune()
		if err != nil {
			return
		}
		n += size

		switch ch {
		case '\x1b':
			size, err = w.handleEscape(r)
			n += size
			if err != nil {
				return
			}
		default:
			_, err = fmt.Fprint(w.out, string(ch))
			if err != nil {
				return
			}
		}
	}
}

func (w *Writer) handleEscape(r *bytes.Reader) (n int, err error) {
	buf := make([]byte, 0, 10)
	buf = append(buf, "\x1b"...)

	// Check '[' continues after \x1b
	ch, size, err := r.ReadRune()
	if err != nil {
		fmt.Fprint(w.out, string(buf))
		return
	}
	n += size
	if ch != '[' {
		fmt.Fprint(w.out, string(buf))
		return
	}

	// Parse escape code
	var code rune
	argBuf := make([]byte, 0, 10)
	for {
		ch, size, err = r.ReadRune()
		if err != nil {
			fmt.Fprint(w.out, string(buf))
			return
		}
		n += size
		if ('a' <= ch && ch <= 'z') || ('A' <= ch && ch <= 'Z') {
			code = ch
			break
		}
		argBuf = append(argBuf, string(ch)...)
	}

	w.applyEscapeCode(buf, string(argBuf), code)
	return
}

func (w *Writer) applyEscapeCode(buf []byte, arg string, code rune) {
	c := &Cursor{Out: w.out}

	switch arg + string(code) {
	case "?25h":
		c.Show()
		return
	case "?25l":
		c.Hide()
		return
	}

	if f, ok := cursorFunctions[code]; ok {
		if n, err := strconv.Atoi(arg); err == nil {
			f(c)(n)
			return
		}
	}

	switch code {
	case 'm':
		w.applySelectGraphicRendition(arg)
	default:
		buf = append(buf, string(code)...)
		fmt.Fprint(w.out, string(buf))
	}
}

// Original implementation: https://github.com/mattn/go-colorable
func (w *Writer) applySelectGraphicRendition(arg string) {
	if arg == "" {
		procSetConsoleTextAttribute.Call(uintptr(w.handle), uintptr(w.orgAttr))
		return
	}

	var csbi consoleScreenBufferInfo
	procGetConsoleScreenBufferInfo.Call(uintptr(w.handle), uintptr(unsafe.Pointer(&csbi)))
	attr := csbi.attributes

	for _, param := range strings.Split(arg, ";") {
		n, err := strconv.Atoi(param)
		if err != nil {
			continue
		}

		switch {
		case n == 0 || n == 100:
			attr = w.orgAttr
		case 1 <= n && n <= 5:
			attr |= foregroundIntensity
		case 30 <= n && n <= 37:
			attr = (attr & backgroundMask)
			if (n-30)&1 != 0 {
				attr |= foregroundRed
			}
			if (n-30)&2 != 0 {
				attr |= foregroundGreen
			}
			if (n-30)&4 != 0 {
				attr |= foregroundBlue
			}
		case 40 <= n && n <= 47:
			attr = (attr & foregroundMask)
			if (n-40)&1 != 0 {
				attr |= backgroundRed
			}
			if (n-40)&2 != 0 {
				attr |= backgroundGreen
			}
			if (n-40)&4 != 0 {
				attr |= backgroundBlue
			}
		case 90 <= n && n <= 97:
			attr = (attr & backgroundMask)
			attr |= foregroundIntensity
			if (n-90)&1 != 0 {
				attr |= foregroundRed
			}
			if (n-90)&2 != 0 {
				attr |= foregroundGreen
			}
			if (n-90)&4 != 0 {
				attr |= foregroundBlue
			}
		case 100 <= n && n <= 107:
			attr = (attr & foregroundMask)
			attr |= backgroundIntensity
			if (n-100)&1 != 0 {
				attr |= backgroundRed
			}
			if (n-100)&2 != 0 {
				attr |= backgroundGreen
			}
			if (n-100)&4 != 0 {
				attr |= backgroundBlue
			}
		}
	}

	procSetConsoleTextAttribute.Call(uintptr(w.handle), uintptr(attr))
}
<commit_msg>Fix Survey output on Windows (#413)<commit_after>package terminal

import (
	"bytes"
	"fmt"
	"io"
	"strconv"
	"strings"
	"syscall"
	"unsafe"

	"github.com/mattn/go-isatty"
)

var (
	cursorFunctions = map[rune]func(c *Cursor) func(int){
		'A': func(c *Cursor) func(int) { return c.Up },
		'B': func(c *Cursor) func(int) { return c.Down },
		'C': func(c *Cursor) func(int) { return c.Forward },
		'D': func(c *Cursor) func(int) { return c.Back },
		'E': func(c *Cursor) func(int) { return c.NextLine },
		'F': func(c *Cursor) func(int) { return c.PreviousLine },
		'G': func(c *Cursor) func(int) { return c.HorizontalAbsolute },
	}
)

const (
	foregroundBlue      = 0x1
	foregroundGreen     = 0x2
	foregroundRed       = 0x4
	foregroundIntensity = 0x8
	foregroundMask      = (foregroundRed | foregroundBlue | foregroundGreen | foregroundIntensity)
	backgroundBlue      = 0x10
	backgroundGreen     = 0x20
	backgroundRed       = 0x40
	backgroundIntensity = 0x80
	backgroundMask      = (backgroundRed | backgroundBlue | backgroundGreen | backgroundIntensity)
)

type Writer struct {
	out     FileWriter
	handle  syscall.Handle
	orgAttr word
}

func NewAnsiStdout(out FileWriter) io.Writer {
	var csbi consoleScreenBufferInfo
	if !isatty.IsTerminal(out.Fd()) {
		return out
	}
	handle := syscall.Handle(out.Fd())
	procGetConsoleScreenBufferInfo.Call(uintptr(handle), uintptr(unsafe.Pointer(&csbi)))
	return &Writer{out: out, handle: handle, orgAttr: csbi.attributes}
}

func NewAnsiStderr(out FileWriter) io.Writer {
	var csbi consoleScreenBufferInfo
	if !isatty.IsTerminal(out.Fd()) {
		return out
	}
	handle := syscall.Handle(out.Fd())
	procGetConsoleScreenBufferInfo.Call(uintptr(handle), uintptr(unsafe.Pointer(&csbi)))
	return &Writer{out: out, handle: handle, orgAttr: csbi.attributes}
}

func (w *Writer) Write(data []byte) (n int, err error) {
	r := bytes.NewReader(data)

	for {
		var ch rune
		var size int
		ch, size, err = r.ReadRune()
		if err != nil {
			if err == io.EOF {
				err = nil
			}
			return
		}
		n += size

		switch ch {
		case '\x1b':
			size, err = w.handleEscape(r)
			n += size
			if err != nil {
				return
			}
		default:
			_, err = fmt.Fprint(w.out, string(ch))
			if err != nil {
				return
			}
		}
	}
}

func (w *Writer) handleEscape(r *bytes.Reader) (n int, err error) {
	buf := make([]byte, 0, 10)
	buf = append(buf, "\x1b"...)

	// Check '[' continues after \x1b
	ch, size, err := r.ReadRune()
	if err != nil {
		fmt.Fprint(w.out, string(buf))
		return
	}
	n += size
	if ch != '[' {
		fmt.Fprint(w.out, string(buf))
		return
	}

	// Parse escape code
	var code rune
	argBuf := make([]byte, 0, 10)
	for {
		ch, size, err = r.ReadRune()
		if err != nil {
			fmt.Fprint(w.out, string(buf))
			return
		}
		n += size
		if ('a' <= ch && ch <= 'z') || ('A' <= ch && ch <= 'Z') {
			code = ch
			break
		}
		argBuf = append(argBuf, string(ch)...)
	}

	w.applyEscapeCode(buf, string(argBuf), code)
	return
}

func (w *Writer) applyEscapeCode(buf []byte, arg string, code rune) {
	c := &Cursor{Out: w.out}

	switch arg + string(code) {
	case "?25h":
		c.Show()
		return
	case "?25l":
		c.Hide()
		return
	}

	if f, ok := cursorFunctions[code]; ok {
		if n, err := strconv.Atoi(arg); err == nil {
			f(c)(n)
			return
		}
	}

	switch code {
	case 'm':
		w.applySelectGraphicRendition(arg)
	default:
		buf = append(buf, string(code)...)
		fmt.Fprint(w.out, string(buf))
	}
}

// Original implementation: https://github.com/mattn/go-colorable
func (w *Writer) applySelectGraphicRendition(arg string) {
	if arg == "" {
		procSetConsoleTextAttribute.Call(uintptr(w.handle), uintptr(w.orgAttr))
		return
	}

	var csbi consoleScreenBufferInfo
	procGetConsoleScreenBufferInfo.Call(uintptr(w.handle), uintptr(unsafe.Pointer(&csbi)))
	attr := csbi.attributes

	for _, param := range strings.Split(arg, ";") {
		n, err := strconv.Atoi(param)
		if err != nil {
			continue
		}

		switch {
		case n == 0 || n == 100:
			attr = w.orgAttr
		case 1 <= n && n <= 5:
			attr |= foregroundIntensity
		case 30 <= n && n <= 37:
			attr = (attr & backgroundMask)
			if (n-30)&1 != 0 {
				attr |= foregroundRed
			}
			if (n-30)&2 != 0 {
				attr |= foregroundGreen
			}
			if (n-30)&4 != 0 {
				attr |= foregroundBlue
			}
		case 40 <= n && n <= 47:
			attr = (attr & foregroundMask)
			if (n-40)&1 != 0 {
				attr |= backgroundRed
			}
			if (n-40)&2 != 0 {
				attr |= backgroundGreen
			}
			if (n-40)&4 != 0 {
				attr |= backgroundBlue
			}
		case 90 <= n && n <= 97:
			attr = (attr & backgroundMask)
			attr |= foregroundIntensity
			if (n-90)&1 != 0 {
				attr |= foregroundRed
			}
			if (n-90)&2 != 0 {
				attr |= foregroundGreen
			}
			if (n-90)&4 != 0 {
				attr |= foregroundBlue
			}
		case 100 <= n && n <= 107:
			attr = (attr & foregroundMask)
			attr |= backgroundIntensity
			if (n-100)&1 != 0 {
				attr |= backgroundRed
			}
			if (n-100)&2 != 0 {
				attr |= backgroundGreen
			}
			if (n-100)&4 != 0 {
				attr |= backgroundBlue
			}
		}
	}

	procSetConsoleTextAttribute.Call(uintptr(w.handle), uintptr(attr))
}
<|endoftext|>


<commit_before>package lib

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"

	"github.com/BytemarkHosting/bytemark-client/util/log"
)

// Request is the lower-level type that makes HTTP calls
type Request interface {
	// GetURL gets the URL this request will be / has been made to
	GetURL() url.URL
	// AllowInsecure allows the Request to be made over http instead of https
	AllowInsecure()
	// JSON-marshals the input and runs the request, unmarshalling the response to out if it's not nil
	MarshalAndRun(in interface{}, out interface{}) (statusCode int, responseBody []byte, err error)
	// Runs the request, reading the request body to reader and unmarshalling the response to responseObject (if not nil)
	Run(body io.Reader, responseObject interface{}) (statusCode int, responseBody []byte, err error)
}

// TokenType returns what kind of token is used for the given endpoint.
// token for endpoints which expect an authorization header like "Token token=blah"
// bearer for endpoints which expect one like "Bearer blah"
func TokenType(ep Endpoint) string {
	switch ep {
	case BillingEndpoint:
		return "token"
	}
	return "bearer"
}

// RequestAlreadyRunError is returned if the Run method was already called for this Request.
type RequestAlreadyRunError struct {
	Request Request
}

// InsecureConnectionError is returned if the endpoint isn't https but AllowInsecure was not called.
type InsecureConnectionError struct {
	Request Request
}

func (e RequestAlreadyRunError) Error() string {
	return "A Request was Run twice"
}

func (e InsecureConnectionError) Error() string {
	return "A Request to an insecure endpoint was attempted when AllowInsecureRequests had not been called."
}

// internalRequest is the workhorse of the bytemark-client/lib - it builds up a request, then Run can be called to get its results.
type internalRequest struct {
	authenticate  bool
	client        Client
	endpoint      Endpoint
	url           *url.URL
	method        string
	body          []byte
	allowInsecure bool
	hasRun        bool
}

// GetURL returns the URL that the Request is for.
func (r *internalRequest) GetURL() url.URL {
	if r.url == nil {
		return url.URL{}
	}
	return *r.url
}

// BuildRequestNoAuth creates a new Request with the intention of not authenticating.
func (c *bytemarkClient) BuildRequestNoAuth(method string, endpoint Endpoint, path string, parts ...string) (r Request, err error) {
	url, err := c.BuildURL(endpoint, path, parts...)
	if err != nil {
		return
	}
	return &internalRequest{
		client:        c,
		endpoint:      endpoint,
		url:           url,
		method:        method,
		allowInsecure: c.allowInsecure,
	}, nil
}

// BuildRequest builds a request that will be authenticated by the endpoint given.
func (c *bytemarkClient) BuildRequest(method string, endpoint Endpoint, path string, parts ...string) (r Request, err error) {
	url, err := c.BuildURL(endpoint, path, parts...)
	if err != nil {
		return
	}
	return &internalRequest{
		authenticate:  true,
		client:        c,
		endpoint:      endpoint,
		url:           url,
		method:        method,
		allowInsecure: c.allowInsecure,
	}, nil
}

// AllowInsecure tells the Request that it's ok if the endpoint isn't communicated with over HTTPS.
func (r *internalRequest) AllowInsecure() {
	r.allowInsecure = true
}

// mkHTTPClient creates an http.Client for this request. If the staging endpoint is used, InsecureSkipVerify is used because I guess we don't have a good cert for that brain.
func (r *internalRequest) mkHTTPClient() (c *http.Client) {
	c = new(http.Client)
	if r.url.Host == "staging.bigv.io" {
		c.Transport = &http.Transport{
			// disable gas lint for this line (gas looks for insecure TLS settings, among other things)
			/* #nosec */
			TLSClientConfig: &tls.Config{
				InsecureSkipVerify: true,
			},
		}
	}
	return c
}

// mkHTTPRequest assembles an http.Request for this Request, adding Authorization headers as needed, setting the Content-Type correctly for whichever endpoint it's talking to.
func (r *internalRequest) mkHTTPRequest(body io.Reader) (req *http.Request, err error) {
	req, err = http.NewRequest(r.method, r.url.String(), body)
	if err != nil {
		return nil, err
	}
	req.Close = true
	req.Header.Add("User-Agent", "bytemark-client-"+Version)

	if r.endpoint == SPPEndpoint {
		req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	} else {
		req.Header.Add("Accept", "application/json")
		req.Header.Add("Content-Type", "application/json")
	}
	if r.authenticate {
		if r.client.GetSessionToken() == "" {
			return nil, NilAuthError{}
		}
		// if we could settle on a single standard
		// rather than two basically-identical ones that'd be cool
		switch TokenType(r.endpoint) {
		case "token":
			req.Header.Add("Authorization", "Token token="+r.client.GetSessionToken())
		case "bearer":
			req.Header.Add("Authorization", "Bearer "+r.client.GetSessionToken())
		}
	}
	return
}

// MarshalAndRun marshals the 'in' object and passes that and 'out' to Run as the body and responseObject, respectively.
func (r *internalRequest) MarshalAndRun(in interface{}, out interface{}) (statusCode int, responseBody []byte, err error) {
	var b bytes.Buffer
	err = json.NewEncoder(&b).Encode(in)
	if err != nil {
		return
	}

	return r.Run(&b, out)
}

// Run performs the request with the given body, and attempts to unmarshal a successful response into responseObject
func (r *internalRequest) Run(body io.Reader, responseObject interface{}) (statusCode int, responseBody []byte, err error) {
	if r.hasRun {
		err = RequestAlreadyRunError{r}
		return
	}
	r.hasRun = true

	if !r.allowInsecure && r.url.Scheme == "http" {
		err = InsecureConnectionError{r}
		return
	}
	var rb []byte
	if body != nil {

		rb, err = ioutil.ReadAll(body)
		if err != nil {
			return 0, nil, err
		}
		log.Debugf(log.LvlHTTPData, "request body: '%s'\r\n", string(rb))
	}

	cli := r.mkHTTPClient()

	// This isnt actually returning errors if not authenticated. (when Auth or preflight has not been run, NilAuthError)
	// TODO: dig deep and find out it does not return this error (NilAuthError)
	req, err := r.mkHTTPRequest(bytes.NewBuffer(rb))
	if err != nil {
		return
	}
	if len(rb) > 0 {
		req.Header.Add("Content-Length", fmt.Sprintf("%d", len(rb)))
	}
	res, err := cli.Do(req)
	if err != nil {
		return
	}

	statusCode = res.StatusCode

	log.Debugf(log.LvlOutline, "%s %s: %d\r\n", r.method, req.URL, res.StatusCode)

	responseBody, err = r.handleResponse(req, rb, res, responseObject)
	return
}

// handleResponse deals with the response coming back from the request - creating an error if required, unmarshalling responseObject if necessary
func (r *internalRequest) handleResponse(req *http.Request, requestBody []byte, res *http.Response, responseObject interface{}) (body []byte, err error) {
	body, err = ioutil.ReadAll(res.Body)
	log.Debugf(log.LvlHTTPData, "response body: '%s'\r\n", body)
	if err != nil {
		return
	}

	baseErr := APIError{
		Method:      r.method,
		URL:         req.URL,
		StatusCode:  res.StatusCode,
		RequestBody: string(requestBody),
	}

	baseErr.ResponseBody = string(body)

	switch res.StatusCode {
	case http.StatusBadRequest:
		// because we need to reference fields specific to BadRequestError later
		err = newBadRequestError(baseErr, body)
	case http.StatusUnauthorized:
		err = UnauthorizedError{baseErr}
	case http.StatusForbidden:
		err = ForbiddenError{baseErr}
	case http.StatusNotFound:
		err = NotFoundError{baseErr}
	case http.StatusInternalServerError:
		err = InternalServerError{baseErr}
	case http.StatusServiceUnavailable:
		err = ServiceUnavailableError{baseErr}
	default:
		if 200 <= res.StatusCode && res.StatusCode <= 299 {
			if responseObject != nil {
				jsonErr := json.Unmarshal(body, responseObject)
				if jsonErr != nil {
					return body, jsonErr
				}
			}
		} else {
			err = UnknownStatusCodeError{baseErr}
		}
	}
	return
}

// BuildURL pieces together a URL from parts, escaping as necessary..
func (c *bytemarkClient) BuildURL(endpoint Endpoint, format string, args ...string) (*url.URL, error) {
	arr := make([]interface{}, len(args))
	for i, str := range args {
		arr[i] = url.QueryEscape(str)
	}
	endpointURL := ""
	switch endpoint {
	case APIEndpoint:
		endpointURL = c.urls.API
	case BrainEndpoint:
		endpointURL = c.urls.Brain
	case BillingEndpoint:
		endpointURL = c.urls.Billing
	case SPPEndpoint:
		endpointURL = c.urls.SPP
	default:
		return nil, UnsupportedEndpointError(endpoint)
	}
	if !strings.HasPrefix(format, "/") {
		return nil, UnsupportedEndpointError(-1)
	}
	return url.Parse(endpointURL + fmt.Sprintf(format, arr...))
}
<commit_msg>changed a file on gitlab, pull that in<commit_after>package lib

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"

	"github.com/BytemarkHosting/bytemark-client/util/log"
)

// Request is the lower-level type that makes HTTP calls
type Request interface {
	// GetURL gets the URL this request will be / has been made to
	GetURL() url.URL
	// AllowInsecure allows the Request to be made over http instead of https
	AllowInsecure()
	// JSON-marshals the input and runs the request, unmarshalling the response to out if it's not nil
	MarshalAndRun(in interface{}, out interface{}) (statusCode int, responseBody []byte, err error)
	// Runs the request, reading the request body to reader and unmarshalling the response to responseObject (if not nil)
	Run(body io.Reader, responseObject interface{}) (statusCode int, responseBody []byte, err error)
}

// TokenType returns what kind of token is used for the given endpoint.
// token for endpoints which expect an authorization header like "Token token=blah"
// bearer for endpoints which expect one like "Bearer blah"
func TokenType(ep Endpoint) string {
	switch ep {
	case BillingEndpoint:
		return "token"
	}
	return "bearer"
}

// RequestAlreadyRunError is returned if the Run method was already called for this Request.
type RequestAlreadyRunError struct {
	Request Request
}

// InsecureConnectionError is returned if the endpoint isn't https but AllowInsecure was not called.
type InsecureConnectionError struct {
	Request Request
}

func (e RequestAlreadyRunError) Error() string {
	return "A Request was Run twice"
}

func (e InsecureConnectionError) Error() string {
	return "A Request to an insecure endpoint was attempted when AllowInsecureRequests had not been called."
}

// internalRequest is the workhorse of the bytemark-client/lib - it builds up a request, then Run can be called to get its results.
type internalRequest struct {
	authenticate  bool
	client        Client
	endpoint      Endpoint
	url           *url.URL
	method        string
	body          []byte
	allowInsecure bool
	hasRun        bool
}

// GetURL returns the URL that the Request is for.
func (r *internalRequest) GetURL() url.URL {
	if r.url == nil {
		return url.URL{}
	}
	return *r.url
}

// BuildRequestNoAuth creates a new Request with the intention of not authenticating.
func (c *bytemarkClient) BuildRequestNoAuth(method string, endpoint Endpoint, path string, parts ...string) (r Request, err error) {
	url, err := c.BuildURL(endpoint, path, parts...)
	if err != nil {
		return
	}
	return &internalRequest{
		client:        c,
		endpoint:      endpoint,
		url:           url,
		method:        method,
		allowInsecure: c.allowInsecure,
	}, nil
}

// BuildRequest builds a request that will be authenticated by the endpoint given.
func (c *bytemarkClient) BuildRequest(method string, endpoint Endpoint, path string, parts ...string) (r Request, err error) {
	url, err := c.BuildURL(endpoint, path, parts...)
	if err != nil {
		return
	}
	return &internalRequest{
		authenticate:  true,
		client:        c,
		endpoint:      endpoint,
		url:           url,
		method:        method,
		allowInsecure: c.allowInsecure,
	}, nil
}

// AllowInsecure tells the Request that it's ok if the endpoint isn't communicated with over HTTPS.
func (r *internalRequest) AllowInsecure() {
	r.allowInsecure = true
}

// mkHTTPClient creates an http.Client for this request. If the staging endpoint is used, InsecureSkipVerify is used because I guess we don't have a good cert for that brain.
func (r *internalRequest) mkHTTPClient() (c *http.Client) {
	c = new(http.Client)
	if r.url.Host == "staging.bigv.io" {
		c.Transport = &http.Transport{
			// disable gas lint for this line (gas looks for insecure TLS settings, among other things)
			/* #nosec */
			TLSClientConfig: &tls.Config{
				InsecureSkipVerify: true,
			},
		}
	}
	return c
}

// mkHTTPRequest assembles an http.Request for this Request, adding Authorization headers as needed, setting the Content-Type correctly for whichever endpoint it's talking to.
func (r *internalRequest) mkHTTPRequest(body io.Reader) (req *http.Request, err error) {
	req, err = http.NewRequest(r.method, r.url.String(), body)
	if err != nil {
		return nil, err
	}
	req.Close = true
	req.Header.Add("User-Agent", "bytemark-client-"+Version)

	if r.endpoint == SPPEndpoint {
		req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	} else {
		req.Header.Add("Accept", "application/json")
		req.Header.Add("Content-Type", "application/json")
	}
	if r.authenticate {
		if r.client.GetSessionToken() == "" {
			return nil, NilAuthError{}
		}
		// if we could settle on a single standard
		// rather than two basically-identical ones that'd be cool
		switch TokenType(r.endpoint) {
		case "token":
			req.Header.Add("Authorization", "Token token="+r.client.GetSessionToken())
		case "bearer":
			req.Header.Add("Authorization", "Bearer "+r.client.GetSessionToken())
		}
	}
	return
}

// MarshalAndRun marshals the 'in' object and passes that and 'out' to Run as the body and responseObject, respectively.
func (r *internalRequest) MarshalAndRun(in interface{}, out interface{}) (statusCode int, responseBody []byte, err error) {
	var b bytes.Buffer
	err = json.NewEncoder(&b).Encode(in)
	if err != nil {
		return
	}

	return r.Run(&b, out)
}

// Run performs the request with the given body, and attempts to unmarshal a successful response into responseObject
func (r *internalRequest) Run(body io.Reader, responseObject interface{}) (statusCode int, responseBody []byte, err error) {
	if r.hasRun {
		err = RequestAlreadyRunError{r}
		return
	}
	r.hasRun = true

	if !r.allowInsecure && r.url.Scheme == "http" {
		err = InsecureConnectionError{r}
		return
	}
	var rb []byte
	if body != nil {

		rb, err = ioutil.ReadAll(body)
		if err != nil {
			return 0, nil, err
		}
		log.Debugf(log.LvlHTTPData, "request body: '%s'\r\n", string(rb))
	}

	cli := r.mkHTTPClient()

	// This isnt actually returning errors if not authenticated. (when Auth or preflight has not been run)
	// TODO: dig deep and find out why it does not return this error (NilAuthError)
	req, err := r.mkHTTPRequest(bytes.NewBuffer(rb))
	if err != nil {
		return
	}
	if len(rb) > 0 {
		req.Header.Add("Content-Length", fmt.Sprintf("%d", len(rb)))
	}
	res, err := cli.Do(req)
	if err != nil {
		return
	}

	statusCode = res.StatusCode

	log.Debugf(log.LvlOutline, "%s %s: %d\r\n", r.method, req.URL, res.StatusCode)

	responseBody, err = r.handleResponse(req, rb, res, responseObject)
	return
}

// handleResponse deals with the response coming back from the request - creating an error if required, unmarshalling responseObject if necessary
func (r *internalRequest) handleResponse(req *http.Request, requestBody []byte, res *http.Response, responseObject interface{}) (body []byte, err error) {
	body, err = ioutil.ReadAll(res.Body)
	log.Debugf(log.LvlHTTPData, "response body: '%s'\r\n", body)
	if err != nil {
		return
	}

	baseErr := APIError{
		Method:      r.method,
		URL:         req.URL,
		StatusCode:  res.StatusCode,
		RequestBody: string(requestBody),
	}

	baseErr.ResponseBody = string(body)

	switch res.StatusCode {
	case http.StatusBadRequest:
		// because we need to reference fields specific to BadRequestError later
		err = newBadRequestError(baseErr, body)
	case http.StatusUnauthorized:
		err = UnauthorizedError{baseErr}
	case http.StatusForbidden:
		err = ForbiddenError{baseErr}
	case http.StatusNotFound:
		err = NotFoundError{baseErr}
	case http.StatusInternalServerError:
		err = InternalServerError{baseErr}
	case http.StatusServiceUnavailable:
		err = ServiceUnavailableError{baseErr}
	default:
		if 200 <= res.StatusCode && res.StatusCode <= 299 {
			if responseObject != nil {
				jsonErr := json.Unmarshal(body, responseObject)
				if jsonErr != nil {
					return body, jsonErr
				}
			}
		} else {
			err = UnknownStatusCodeError{baseErr}
		}
	}
	return
}

// BuildURL pieces together a URL from parts, escaping as necessary..
func (c *bytemarkClient) BuildURL(endpoint Endpoint, format string, args ...string) (*url.URL, error) {
	arr := make([]interface{}, len(args))
	for i, str := range args {
		arr[i] = url.QueryEscape(str)
	}
	endpointURL := ""
	switch endpoint {
	case APIEndpoint:
		endpointURL = c.urls.API
	case BrainEndpoint:
		endpointURL = c.urls.Brain
	case BillingEndpoint:
		endpointURL = c.urls.Billing
	case SPPEndpoint:
		endpointURL = c.urls.SPP
	default:
		return nil, UnsupportedEndpointError(endpoint)
	}
	if !strings.HasPrefix(format, "/") {
		return nil, UnsupportedEndpointError(-1)
	}
	return url.Parse(endpointURL + fmt.Sprintf(format, arr...))
}
<|endoftext|>


<commit_before>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package gob

import (
	"bytes"
	"io"
	"os"
	"runtime"
	"testing"
)

type Bench struct {
	A int
	B float64
	C string
	D []byte
}

func benchmarkEndToEnd(b *testing.B, v interface{}, pipe func() (r io.Reader, w io.Writer, err error)) {
	b.RunParallel(func(pb *testing.PB) {
		r, w, err := pipe()
		if err != nil {
			b.Fatal("can't get pipe:", err)
		}
		enc := NewEncoder(w)
		dec := NewDecoder(r)
		for pb.Next() {
			if err := enc.Encode(v); err != nil {
				b.Fatal("encode error:", err)
			}
			if err := dec.Decode(v); err != nil {
				b.Fatal("decode error:", err)
			}
		}
	})
}

func BenchmarkEndToEndPipe(b *testing.B) {
	v := &Bench{7, 3.2, "now is the time", bytes.Repeat([]byte("for all good men"), 100)}
	benchmarkEndToEnd(b, v, func() (r io.Reader, w io.Writer, err error) {
		r, w, err = os.Pipe()
		return
	})
}

func BenchmarkEndToEndByteBuffer(b *testing.B) {
	v := &Bench{7, 3.2, "now is the time", bytes.Repeat([]byte("for all good men"), 100)}
	benchmarkEndToEnd(b, v, func() (r io.Reader, w io.Writer, err error) {
		var buf bytes.Buffer
		return &buf, &buf, nil
	})
}

func BenchmarkEndToEndSliceByteBuffer(b *testing.B) {
	v := &Bench{7, 3.2, "now is the time", nil}
	Register(v)
	arr := make([]interface{}, 100)
	for i := range arr {
		arr[i] = v
	}
	benchmarkEndToEnd(b, &arr, func() (r io.Reader, w io.Writer, err error) {
		var buf bytes.Buffer
		return &buf, &buf, nil
	})
}

func TestCountEncodeMallocs(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping malloc count in short mode")
	}
	if runtime.GOMAXPROCS(0) > 1 {
		t.Skip("skipping; GOMAXPROCS>1")
	}

	const N = 1000

	var buf bytes.Buffer
	enc := NewEncoder(&buf)
	bench := &Bench{7, 3.2, "now is the time", []byte("for all good men")}

	allocs := testing.AllocsPerRun(N, func() {
		err := enc.Encode(bench)
		if err != nil {
			t.Fatal("encode:", err)
		}
	})
	if allocs != 0 {
		t.Fatalf("mallocs per encode of type Bench: %v; wanted 0\n", allocs)
	}
}

func TestCountDecodeMallocs(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping malloc count in short mode")
	}
	if runtime.GOMAXPROCS(0) > 1 {
		t.Skip("skipping; GOMAXPROCS>1")
	}

	const N = 1000

	var buf bytes.Buffer
	enc := NewEncoder(&buf)
	bench := &Bench{7, 3.2, "now is the time", []byte("for all good men")}

	// Fill the buffer with enough to decode
	testing.AllocsPerRun(N, func() {
		err := enc.Encode(bench)
		if err != nil {
			t.Fatal("encode:", err)
		}
	})

	dec := NewDecoder(&buf)
	allocs := testing.AllocsPerRun(N, func() {
		*bench = Bench{}
		err := dec.Decode(&bench)
		if err != nil {
			t.Fatal("decode:", err)
		}
	})
	if allocs != 4 {
		t.Fatalf("mallocs per decode of type Bench: %v; wanted 4\n", allocs)
	}
}
<commit_msg>encoding/gob: fix data races in benchmarks All goroutines decode into the same value.<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package gob

import (
	"bytes"
	"io"
	"os"
	"runtime"
	"testing"
)

type Bench struct {
	A int
	B float64
	C string
	D []byte
}

func benchmarkEndToEnd(b *testing.B, ctor func() interface{}, pipe func() (r io.Reader, w io.Writer, err error)) {
	b.RunParallel(func(pb *testing.PB) {
		r, w, err := pipe()
		if err != nil {
			b.Fatal("can't get pipe:", err)
		}
		v := ctor()
		enc := NewEncoder(w)
		dec := NewDecoder(r)
		for pb.Next() {
			if err := enc.Encode(v); err != nil {
				b.Fatal("encode error:", err)
			}
			if err := dec.Decode(v); err != nil {
				b.Fatal("decode error:", err)
			}
		}
	})
}

func BenchmarkEndToEndPipe(b *testing.B) {
	benchmarkEndToEnd(b, func() interface{} {
		return &Bench{7, 3.2, "now is the time", bytes.Repeat([]byte("for all good men"), 100)}
	}, func() (r io.Reader, w io.Writer, err error) {
		r, w, err = os.Pipe()
		return
	})
}

func BenchmarkEndToEndByteBuffer(b *testing.B) {
	benchmarkEndToEnd(b, func() interface{} {
		return &Bench{7, 3.2, "now is the time", bytes.Repeat([]byte("for all good men"), 100)}
	}, func() (r io.Reader, w io.Writer, err error) {
		var buf bytes.Buffer
		return &buf, &buf, nil
	})
}

func BenchmarkEndToEndSliceByteBuffer(b *testing.B) {
	benchmarkEndToEnd(b, func() interface{} {
		v := &Bench{7, 3.2, "now is the time", nil}
		Register(v)
		arr := make([]interface{}, 100)
		for i := range arr {
			arr[i] = v
		}
		return &arr
	}, func() (r io.Reader, w io.Writer, err error) {
		var buf bytes.Buffer
		return &buf, &buf, nil
	})
}

func TestCountEncodeMallocs(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping malloc count in short mode")
	}
	if runtime.GOMAXPROCS(0) > 1 {
		t.Skip("skipping; GOMAXPROCS>1")
	}

	const N = 1000

	var buf bytes.Buffer
	enc := NewEncoder(&buf)
	bench := &Bench{7, 3.2, "now is the time", []byte("for all good men")}

	allocs := testing.AllocsPerRun(N, func() {
		err := enc.Encode(bench)
		if err != nil {
			t.Fatal("encode:", err)
		}
	})
	if allocs != 0 {
		t.Fatalf("mallocs per encode of type Bench: %v; wanted 0\n", allocs)
	}
}

func TestCountDecodeMallocs(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping malloc count in short mode")
	}
	if runtime.GOMAXPROCS(0) > 1 {
		t.Skip("skipping; GOMAXPROCS>1")
	}

	const N = 1000

	var buf bytes.Buffer
	enc := NewEncoder(&buf)
	bench := &Bench{7, 3.2, "now is the time", []byte("for all good men")}

	// Fill the buffer with enough to decode
	testing.AllocsPerRun(N, func() {
		err := enc.Encode(bench)
		if err != nil {
			t.Fatal("encode:", err)
		}
	})

	dec := NewDecoder(&buf)
	allocs := testing.AllocsPerRun(N, func() {
		*bench = Bench{}
		err := dec.Decode(&bench)
		if err != nil {
			t.Fatal("decode:", err)
		}
	})
	if allocs != 4 {
		t.Fatalf("mallocs per decode of type Bench: %v; wanted 4\n", allocs)
	}
}
<|endoftext|>


<commit_before>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package types

import (
	"go/ast"
	"go/build"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"
)

var gcPath string // Go compiler path

func init() {
	// determine compiler
	var gc string
	switch runtime.GOARCH {
	case "386":
		gc = "8g"
	case "amd64":
		gc = "6g"
	case "arm":
		gc = "5g"
	default:
		gcPath = "unknown-GOARCH-compiler"
		return
	}
	gcPath = filepath.Join(build.ToolDir, gc)
}

func compile(t *testing.T, dirname, filename string) string {
	cmd := exec.Command(gcPath, filename)
	cmd.Dir = dirname
	out, err := cmd.CombinedOutput()
	if err != nil {
		t.Logf("%s", out)
		t.Fatalf("%s %s failed: %s", gcPath, filename, err)
	}
	archCh, _ := build.ArchChar(runtime.GOARCH)
	// filename should end with ".go"
	return filepath.Join(dirname, filename[:len(filename)-2]+archCh)
}

// Use the same global imports map for all tests. The effect is
// as if all tested packages were imported into a single package.
var imports = make(map[string]*ast.Object)

func testPath(t *testing.T, path string) bool {
	_, err := GcImport(imports, path)
	if err != nil {
		t.Errorf("testPath(%s): %s", path, err)
		return false
	}
	return true
}

const maxTime = 3 * time.Second

func testDir(t *testing.T, dir string, endTime time.Time) (nimports int) {
	dirname := filepath.Join(runtime.GOROOT(), "pkg", runtime.GOOS+"_"+runtime.GOARCH, dir)
	list, err := ioutil.ReadDir(dirname)
	if err != nil {
		t.Errorf("testDir(%s): %s", dirname, err)
	}
	for _, f := range list {
		if time.Now().After(endTime) {
			t.Log("testing time used up")
			return
		}
		switch {
		case !f.IsDir():
			// try extensions
			for _, ext := range pkgExts {
				if strings.HasSuffix(f.Name(), ext) {
					name := f.Name()[0 : len(f.Name())-len(ext)] // remove extension
					if testPath(t, filepath.Join(dir, name)) {
						nimports++
					}
				}
			}
		case f.IsDir():
			nimports += testDir(t, filepath.Join(dir, f.Name()), endTime)
		}
	}
	return
}

func TestGcImport(t *testing.T) {
	// On cross-compile builds, the path will not exist.
	// Need to use GOHOSTOS, which is not available.
	if _, err := os.Stat(gcPath); err != nil {
		t.Logf("skipping test: %v", err)
		return
	}

	if outFn := compile(t, "testdata", "exports.go"); outFn != "" {
		defer os.Remove(outFn)
	}

	nimports := 0
	if testPath(t, "./testdata/exports") {
		nimports++
	}
	nimports += testDir(t, "", time.Now().Add(maxTime)) // installed packages
	t.Logf("tested %d imports", nimports)
}

var importedObjectTests = []struct {
	name string
	kind ast.ObjKind
	typ  string
}{
	{"unsafe.Pointer", ast.Typ, "Pointer"},
	{"math.Pi", ast.Con, "untyped float"},
	{"io.Reader", ast.Typ, "interface{Read(p []byte) (n int, err error)}"},
	{"io.ReadWriter", ast.Typ, "interface{Read(p []byte) (n int, err error); Write(p []byte) (n int, err error)}"},
	{"math.Sin", ast.Fun, "func(x·2 float64) (_ float64)"},
	// TODO(gri) add more tests
}

func TestGcImportedTypes(t *testing.T) {
	// This package does not yet know how to read gccgo export data.
	if runtime.Compiler == "gccgo" {
		return
	}
	for _, test := range importedObjectTests {
		s := strings.Split(test.name, ".")
		if len(s) != 2 {
			t.Fatal("inconsistent test data")
		}
		importPath := s[0]
		objName := s[1]

		pkg, err := GcImport(imports, importPath)
		if err != nil {
			t.Error(err)
			continue
		}

		obj := pkg.Data.(*ast.Scope).Lookup(objName)
		if obj.Kind != test.kind {
			t.Errorf("%s: got kind = %q; want %q", test.name, obj.Kind, test.kind)
		}
		typ := typeString(underlying(obj.Type.(Type)))
		if typ != test.typ {
			t.Errorf("%s: got type = %q; want %q", test.name, typ, test.typ)
		}
	}
}
<commit_msg>go/types: add more logging to gc import test<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package types

import (
	"go/ast"
	"go/build"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"
)

var gcPath string // Go compiler path

func init() {
	// determine compiler
	var gc string
	switch runtime.GOARCH {
	case "386":
		gc = "8g"
	case "amd64":
		gc = "6g"
	case "arm":
		gc = "5g"
	default:
		gcPath = "unknown-GOARCH-compiler"
		return
	}
	gcPath = filepath.Join(build.ToolDir, gc)
}

func compile(t *testing.T, dirname, filename string) string {
	cmd := exec.Command(gcPath, filename)
	cmd.Dir = dirname
	out, err := cmd.CombinedOutput()
	if err != nil {
		t.Logf("%s", out)
		t.Fatalf("%s %s failed: %s", gcPath, filename, err)
	}
	archCh, _ := build.ArchChar(runtime.GOARCH)
	// filename should end with ".go"
	return filepath.Join(dirname, filename[:len(filename)-2]+archCh)
}

// Use the same global imports map for all tests. The effect is
// as if all tested packages were imported into a single package.
var imports = make(map[string]*ast.Object)

func testPath(t *testing.T, path string) bool {
	t0 := time.Now()
	_, err := GcImport(imports, path)
	if err != nil {
		t.Errorf("testPath(%s): %s", path, err)
		return false
	}
	t.Logf("testPath(%s): %v", path, time.Since(t0))
	return true
}

const maxTime = 30 * time.Second

func testDir(t *testing.T, dir string, endTime time.Time) (nimports int) {
	dirname := filepath.Join(runtime.GOROOT(), "pkg", runtime.GOOS+"_"+runtime.GOARCH, dir)
	list, err := ioutil.ReadDir(dirname)
	if err != nil {
		t.Fatalf("testDir(%s): %s", dirname, err)
	}
	for _, f := range list {
		if time.Now().After(endTime) {
			t.Log("testing time used up")
			return
		}
		switch {
		case !f.IsDir():
			// try extensions
			for _, ext := range pkgExts {
				if strings.HasSuffix(f.Name(), ext) {
					name := f.Name()[0 : len(f.Name())-len(ext)] // remove extension
					if testPath(t, filepath.Join(dir, name)) {
						nimports++
					}
				}
			}
		case f.IsDir():
			nimports += testDir(t, filepath.Join(dir, f.Name()), endTime)
		}
	}
	return
}

func TestGcImport(t *testing.T) {
	// On cross-compile builds, the path will not exist.
	// Need to use GOHOSTOS, which is not available.
	if _, err := os.Stat(gcPath); err != nil {
		t.Logf("skipping test: %v", err)
		return
	}

	if outFn := compile(t, "testdata", "exports.go"); outFn != "" {
		defer os.Remove(outFn)
	}

	nimports := 0
	if testPath(t, "./testdata/exports") {
		nimports++
	}
	nimports += testDir(t, "", time.Now().Add(maxTime)) // installed packages
	t.Logf("tested %d imports", nimports)
}

var importedObjectTests = []struct {
	name string
	kind ast.ObjKind
	typ  string
}{
	{"unsafe.Pointer", ast.Typ, "Pointer"},
	{"math.Pi", ast.Con, "untyped float"},
	{"io.Reader", ast.Typ, "interface{Read(p []byte) (n int, err error)}"},
	{"io.ReadWriter", ast.Typ, "interface{Read(p []byte) (n int, err error); Write(p []byte) (n int, err error)}"},
	{"math.Sin", ast.Fun, "func(x·2 float64) (_ float64)"},
	// TODO(gri) add more tests
}

func TestGcImportedTypes(t *testing.T) {
	// This package does not yet know how to read gccgo export data.
	if runtime.Compiler == "gccgo" {
		return
	}
	for _, test := range importedObjectTests {
		s := strings.Split(test.name, ".")
		if len(s) != 2 {
			t.Fatal("inconsistent test data")
		}
		importPath := s[0]
		objName := s[1]

		pkg, err := GcImport(imports, importPath)
		if err != nil {
			t.Error(err)
			continue
		}

		obj := pkg.Data.(*ast.Scope).Lookup(objName)
		if obj.Kind != test.kind {
			t.Errorf("%s: got kind = %q; want %q", test.name, obj.Kind, test.kind)
		}
		typ := typeString(underlying(obj.Type.(Type)))
		if typ != test.typ {
			t.Errorf("%s: got type = %q; want %q", test.name, typ, test.typ)
		}
	}
}
<|endoftext|>


<commit_before><commit_msg>auto-update<commit_after><|endoftext|>


<commit_before><commit_msg>add documentation on rest of the functions<commit_after><|endoftext|>


<commit_before><commit_msg>fix potential null pointer<commit_after><|endoftext|>


<commit_before><commit_msg>Song parsing now strips punctuation from first line when using it as title<commit_after><|endoftext|>


<commit_before>// Copyright (c) 2013 ActiveState Software Inc. All rights reserved.

package tail

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"time"

	"github.com/ActiveState/tail/ratelimiter"
	"github.com/ActiveState/tail/util"
	"github.com/ActiveState/tail/watch"
	"gopkg.in/tomb.v1"
)

var (
	ErrStop = fmt.Errorf("tail should now stop")
)

type Line struct {
	Text string
	Time time.Time
	Err  error // Error from tail
}

// NewLine returns a Line with present time.
func NewLine(text string) *Line {
	return &Line{text, time.Now(), nil}
}

// SeekInfo represents arguments to `os.Seek`
type SeekInfo struct {
	Offset int64
	Whence int // os.SEEK_*
}

// Config is used to specify how a file must be tailed.
type Config struct {
	// File-specifc
	Location    *SeekInfo // Seek to this location before tailing
	ReOpen      bool      // Reopen recreated files (tail -F)
	MustExist   bool      // Fail early if the file does not exist
	Poll        bool      // Poll for file changes instead of using inotify
	RateLimiter *ratelimiter.LeakyBucket

	// Generic IO
	Follow      bool // Continue looking for new lines (tail -f)
	MaxLineSize int  // If non-zero, split longer lines into multiple lines

	// Logger, when nil, is set to tail.DefaultLogger
	// To disable logging: set field to tail.DiscardingLogger
	Logger *log.Logger
}

type Tail struct {
	Filename string
	Lines    chan *Line
	Config

	file    *os.File
	reader  *bufio.Reader
	tracker *watch.InotifyTracker

	watcher watch.FileWatcher
	changes *watch.FileChanges

	tomb.Tomb // provides: Done, Kill, Dying
}

var (
	// DefaultLogger is used when Config.Logger == nil
	DefaultLogger = log.New(os.Stderr, "", log.LstdFlags)
	// DiscardingLogger can be used to disable logging output
	DiscardingLogger = log.New(ioutil.Discard, "", 0)
)

// TailFile begins tailing the file. Output stream is made available
// via the `Tail.Lines` channel. To handle errors during tailing,
// invoke the `Wait` or `Err` method after finishing reading from the
// `Lines` channel.
func TailFile(filename string, config Config) (*Tail, error) {
	if config.ReOpen && !config.Follow {
		util.Fatal("cannot set ReOpen without Follow.")
	}

	t := &Tail{
		Filename: filename,
		Lines:    make(chan *Line),
		Config:   config,
	}

	// when Logger was not specified in config, use default logger
	if t.Logger == nil {
		t.Logger = log.New(os.Stderr, "", log.LstdFlags)
	}

	if t.Poll {
		t.watcher = watch.NewPollingFileWatcher(filename)
	} else {
		t.tracker = watch.NewInotifyTracker()
		w, err := t.tracker.NewWatcher()
		if err != nil {
			return nil, err
		}
		t.watcher = watch.NewInotifyFileWatcher(filename, w)
	}

	if t.MustExist {
		var err error
		t.file, err = OpenFile(t.Filename)
		if err != nil {
			return nil, err
		}
	}

	go t.tailFileSync()

	return t, nil
}

// Return the file's current position, like stdio's ftell().
// But this value is not very accurate.
// it may readed one line in the chan(tail.Lines),
// so it may lost one line.
func (tail *Tail) Tell() (offset int64, err error) {
	if tail.file == nil {
		return
	}
	offset, err = tail.file.Seek(0, os.SEEK_CUR)
	if err == nil {
		offset -= int64(tail.reader.Buffered())
	}
	return
}

// Stop stops the tailing activity.
func (tail *Tail) Stop() error {
	tail.Kill(nil)
	return tail.Wait()
}

func (tail *Tail) close() {
	close(tail.Lines)
	if tail.file != nil {
		tail.file.Close()
	}
}

func (tail *Tail) reopen() error {
	if tail.file != nil {
		tail.file.Close()
	}
	for {
		var err error
		tail.file, err = OpenFile(tail.Filename)
		if err != nil {
			if os.IsNotExist(err) {
				tail.Logger.Printf("Waiting for %s to appear...", tail.Filename)
				if err := tail.watcher.BlockUntilExists(&tail.Tomb); err != nil {
					if err == tomb.ErrDying {
						return err
					}
					return fmt.Errorf("Failed to detect creation of %s: %s", tail.Filename, err)
				}
				continue
			}
			return fmt.Errorf("Unable to open file %s: %s", tail.Filename, err)
		}
		break
	}
	return nil
}

func (tail *Tail) readLine() (string, error) {
	line, err := tail.reader.ReadString('\n')
	if err != nil {
		// Note ReadString "returns the data read before the error" in
		// case of an error, including EOF, so we return it as is. The
		// caller is expected to process it if err is EOF.
		return line, err
	}

	line = strings.TrimRight(line, "\n")

	return line, err
}

func (tail *Tail) tailFileSync() {
	defer tail.Done()
	defer tail.close()

	if !tail.MustExist {
		// deferred first open.
		err := tail.reopen()
		if err != nil {
			if err != tomb.ErrDying {
				tail.Kill(err)
			}
			return
		}
	}

	// Seek to requested location on first open of the file.
	if tail.Location != nil {
		_, err := tail.file.Seek(tail.Location.Offset, tail.Location.Whence)
		tail.Logger.Printf("Seeked %s - %+v\n", tail.Filename, tail.Location)
		if err != nil {
			tail.Killf("Seek error on %s: %s", tail.Filename, err)
			return
		}
	}

	tail.openReader()

	// Read line by line.
	for {
		// grab the position in case we need to back up in the event of a half-line
		offset, err := tail.Tell()
		if err != nil {
			tail.Kill(err)
			return
		}

		line, err := tail.readLine()

		// Process `line` even if err is EOF.
		if err == nil {
			cooloff := !tail.sendLine(line)
			if cooloff {
				// Wait a second before seeking till the end of
				// file when rate limit is reached.
				msg := fmt.Sprintf(
					"Too much log activity; waiting a second " +
						"before resuming tailing")
				tail.Lines <- &Line{msg, time.Now(), fmt.Errorf(msg)}
				select {
				case <-time.After(time.Second):
				case <-tail.Dying():
					return
				}
				err = tail.seekEnd()
				if err != nil {
					tail.Kill(err)
					return
				}
			}
		} else if err == io.EOF {
			if !tail.Follow {
				if line != "" {
					tail.sendLine(line)
				}
				return
			}

			if tail.Follow && line != "" {
				// this has the potential to never return the last line if
				// it's not followed by a newline; seems a fair trade here
				err := tail.seekTo(SeekInfo{Offset: offset, Whence: 0})
				if err != nil {
					tail.Kill(err)
					return
				}
			}

			// When EOF is reached, wait for more data to become
			// available. Wait strategy is based on the `tail.watcher`
			// implementation (inotify or polling).
			err := tail.waitForChanges()
			if err != nil {
				if err != ErrStop {
					tail.Kill(err)
				}
				return
			}
		} else {
			// non-EOF error
			tail.Killf("Error reading %s: %s", tail.Filename, err)
			return
		}

		select {
		case <-tail.Dying():
			return
		default:
		}
	}
}

// waitForChanges waits until the file has been appended, deleted,
// moved or truncated. When moved or deleted - the file will be
// reopened if ReOpen is true. Truncated files are always reopened.
func (tail *Tail) waitForChanges() error {
	if tail.changes == nil {
		st, err := tail.file.Stat()
		if err != nil {
			return err
		}
		tail.changes = tail.watcher.ChangeEvents(&tail.Tomb, st)
	}

	select {
	case <-tail.changes.Modified:
		return nil
	case <-tail.changes.Deleted:
		tail.changes = nil
		if tail.ReOpen {
			// XXX: we must not log from a library.
			tail.Logger.Printf("Re-opening moved/deleted file %s ...", tail.Filename)
			if err := tail.reopen(); err != nil {
				return err
			}
			tail.Logger.Printf("Successfully reopened %s", tail.Filename)
			tail.openReader()
			return nil
		} else {
			tail.Logger.Printf("Stopping tail as file no longer exists: %s", tail.Filename)
			return ErrStop
		}
	case <-tail.changes.Truncated:
		// Always reopen truncated files (Follow is true)
		tail.Logger.Printf("Re-opening truncated file %s ...", tail.Filename)
		if err := tail.reopen(); err != nil {
			return err
		}
		tail.Logger.Printf("Successfully reopened truncated %s", tail.Filename)
		tail.openReader()
		return nil
	case <-tail.Dying():
		return ErrStop
	}
	panic("unreachable")
}

func (tail *Tail) openReader() {
	if tail.MaxLineSize > 0 {
		// add 2 to account for newline characters
		tail.reader = bufio.NewReaderSize(tail.file, tail.MaxLineSize+2)
	} else {
		tail.reader = bufio.NewReader(tail.file)
	}
}

func (tail *Tail) seekEnd() error {
	return tail.seekTo(SeekInfo{Offset: 0, Whence: 2})
}

func (tail *Tail) seekTo(pos SeekInfo) error {
	_, err := tail.file.Seek(pos.Offset, pos.Whence)
	if err != nil {
		return fmt.Errorf("Seek error on %s: %s", tail.Filename, err)
	}
	// Reset the read buffer whenever the file is re-seek'ed
	tail.reader.Reset(tail.file)
	return nil
}

// sendLine sends the line(s) to Lines channel, splitting longer lines
// if necessary. Return false if rate limit is reached.
func (tail *Tail) sendLine(line string) bool {
	now := time.Now()
	lines := []string{line}

	// Split longer lines
	if tail.MaxLineSize > 0 && len(line) > tail.MaxLineSize {
		lines = util.PartitionString(line, tail.MaxLineSize)
	}

	for _, line := range lines {
		tail.Lines <- &Line{line, now, nil}
	}

	if tail.Config.RateLimiter != nil {
		ok := tail.Config.RateLimiter.Pour(uint16(len(lines)))
		if !ok {
			tail.Logger.Printf("Leaky bucket full (%v); entering 1s cooloff period.\n",
				tail.Filename)
			return false
		}
	}

	return true
}

// Cleanup removes inotify watches added by the tail package. This function is
// meant to be invoked from a process's exit handler. Linux kernel may not
// automatically remove inotify watches after the process exits.
func (tail *Tail) Cleanup() {
	if tail.tracker != nil {
		tail.tracker.CloseAll()
	}
}
<commit_msg>Do not seek in named pipe<commit_after>// Copyright (c) 2013 ActiveState Software Inc. All rights reserved.

package tail

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"time"

	"github.com/ActiveState/tail/ratelimiter"
	"github.com/ActiveState/tail/util"
	"github.com/ActiveState/tail/watch"
	"gopkg.in/tomb.v1"
)

var (
	ErrStop = fmt.Errorf("tail should now stop")
)

type Line struct {
	Text string
	Time time.Time
	Err  error // Error from tail
}

// NewLine returns a Line with present time.
func NewLine(text string) *Line {
	return &Line{text, time.Now(), nil}
}

// SeekInfo represents arguments to `os.Seek`
type SeekInfo struct {
	Offset int64
	Whence int // os.SEEK_*
}

// Config is used to specify how a file must be tailed.
type Config struct {
	// File-specifc
	Location    *SeekInfo // Seek to this location before tailing
	ReOpen      bool      // Reopen recreated files (tail -F)
	MustExist   bool      // Fail early if the file does not exist
	Poll        bool      // Poll for file changes instead of using inotify
	Pipe        bool      // Is a named pipe (mkfifo)
	RateLimiter *ratelimiter.LeakyBucket

	// Generic IO
	Follow      bool // Continue looking for new lines (tail -f)
	MaxLineSize int  // If non-zero, split longer lines into multiple lines

	// Logger, when nil, is set to tail.DefaultLogger
	// To disable logging: set field to tail.DiscardingLogger
	Logger *log.Logger
}

type Tail struct {
	Filename string
	Lines    chan *Line
	Config

	file    *os.File
	reader  *bufio.Reader
	tracker *watch.InotifyTracker

	watcher watch.FileWatcher
	changes *watch.FileChanges

	tomb.Tomb // provides: Done, Kill, Dying
}

var (
	// DefaultLogger is used when Config.Logger == nil
	DefaultLogger = log.New(os.Stderr, "", log.LstdFlags)
	// DiscardingLogger can be used to disable logging output
	DiscardingLogger = log.New(ioutil.Discard, "", 0)
)

// TailFile begins tailing the file. Output stream is made available
// via the `Tail.Lines` channel. To handle errors during tailing,
// invoke the `Wait` or `Err` method after finishing reading from the
// `Lines` channel.
func TailFile(filename string, config Config) (*Tail, error) {
	if config.ReOpen && !config.Follow {
		util.Fatal("cannot set ReOpen without Follow.")
	}

	t := &Tail{
		Filename: filename,
		Lines:    make(chan *Line),
		Config:   config,
	}

	// when Logger was not specified in config, use default logger
	if t.Logger == nil {
		t.Logger = log.New(os.Stderr, "", log.LstdFlags)
	}

	if t.Poll {
		t.watcher = watch.NewPollingFileWatcher(filename)
	} else {
		t.tracker = watch.NewInotifyTracker()
		w, err := t.tracker.NewWatcher()
		if err != nil {
			return nil, err
		}
		t.watcher = watch.NewInotifyFileWatcher(filename, w)
	}

	if t.MustExist {
		var err error
		t.file, err = OpenFile(t.Filename)
		if err != nil {
			return nil, err
		}
	}

	go t.tailFileSync()

	return t, nil
}

// Return the file's current position, like stdio's ftell().
// But this value is not very accurate.
// it may readed one line in the chan(tail.Lines),
// so it may lost one line.
func (tail *Tail) Tell() (offset int64, err error) {
	if tail.file == nil {
		return
	}
	offset, err = tail.file.Seek(0, os.SEEK_CUR)
	if err == nil {
		offset -= int64(tail.reader.Buffered())
	}
	return
}

// Stop stops the tailing activity.
func (tail *Tail) Stop() error {
	tail.Kill(nil)
	return tail.Wait()
}

func (tail *Tail) close() {
	close(tail.Lines)
	if tail.file != nil {
		tail.file.Close()
	}
}

func (tail *Tail) reopen() error {
	if tail.file != nil {
		tail.file.Close()
	}
	for {
		var err error
		tail.file, err = OpenFile(tail.Filename)
		if err != nil {
			if os.IsNotExist(err) {
				tail.Logger.Printf("Waiting for %s to appear...", tail.Filename)
				if err := tail.watcher.BlockUntilExists(&tail.Tomb); err != nil {
					if err == tomb.ErrDying {
						return err
					}
					return fmt.Errorf("Failed to detect creation of %s: %s", tail.Filename, err)
				}
				continue
			}
			return fmt.Errorf("Unable to open file %s: %s", tail.Filename, err)
		}
		break
	}
	return nil
}

func (tail *Tail) readLine() (string, error) {
	line, err := tail.reader.ReadString('\n')
	if err != nil {
		// Note ReadString "returns the data read before the error" in
		// case of an error, including EOF, so we return it as is. The
		// caller is expected to process it if err is EOF.
		return line, err
	}

	line = strings.TrimRight(line, "\n")

	return line, err
}

func (tail *Tail) tailFileSync() {
	defer tail.Done()
	defer tail.close()

	if !tail.MustExist {
		// deferred first open.
		err := tail.reopen()
		if err != nil {
			if err != tomb.ErrDying {
				tail.Kill(err)
			}
			return
		}
	}

	// Seek to requested location on first open of the file.
	if tail.Location != nil {
		_, err := tail.file.Seek(tail.Location.Offset, tail.Location.Whence)
		tail.Logger.Printf("Seeked %s - %+v\n", tail.Filename, tail.Location)
		if err != nil {
			tail.Killf("Seek error on %s: %s", tail.Filename, err)
			return
		}
	}

	tail.openReader()

	var offset int64 = 0
	var err error
	// Read line by line.
	for {
		// do not seek in named pipes
		if !tail.Pipe {
			// grab the position in case we need to back up in the event of a half-line
			offset, err = tail.Tell()
			if err != nil {
				tail.Kill(err)
				return
			}
		}

		line, err := tail.readLine()

		// Process `line` even if err is EOF.
		if err == nil {
			cooloff := !tail.sendLine(line)
			if cooloff {
				// Wait a second before seeking till the end of
				// file when rate limit is reached.
				msg := fmt.Sprintf(
					"Too much log activity; waiting a second " +
						"before resuming tailing")
				tail.Lines <- &Line{msg, time.Now(), fmt.Errorf(msg)}
				select {
				case <-time.After(time.Second):
				case <-tail.Dying():
					return
				}
				err = tail.seekEnd()
				if err != nil {
					tail.Kill(err)
					return
				}
			}
		} else if err == io.EOF {
			if !tail.Follow {
				if line != "" {
					tail.sendLine(line)
				}
				return
			}

			if tail.Follow && line != "" {
				// this has the potential to never return the last line if
				// it's not followed by a newline; seems a fair trade here
				err := tail.seekTo(SeekInfo{Offset: offset, Whence: 0})
				if err != nil {
					tail.Kill(err)
					return
				}
			}

			// When EOF is reached, wait for more data to become
			// available. Wait strategy is based on the `tail.watcher`
			// implementation (inotify or polling).
			err := tail.waitForChanges()
			if err != nil {
				if err != ErrStop {
					tail.Kill(err)
				}
				return
			}
		} else {
			// non-EOF error
			tail.Killf("Error reading %s: %s", tail.Filename, err)
			return
		}

		select {
		case <-tail.Dying():
			return
		default:
		}
	}
}

// waitForChanges waits until the file has been appended, deleted,
// moved or truncated. When moved or deleted - the file will be
// reopened if ReOpen is true. Truncated files are always reopened.
func (tail *Tail) waitForChanges() error {
	if tail.changes == nil {
		st, err := tail.file.Stat()
		if err != nil {
			return err
		}
		tail.changes = tail.watcher.ChangeEvents(&tail.Tomb, st)
	}

	select {
	case <-tail.changes.Modified:
		return nil
	case <-tail.changes.Deleted:
		tail.changes = nil
		if tail.ReOpen {
			// XXX: we must not log from a library.
			tail.Logger.Printf("Re-opening moved/deleted file %s ...", tail.Filename)
			if err := tail.reopen(); err != nil {
				return err
			}
			tail.Logger.Printf("Successfully reopened %s", tail.Filename)
			tail.openReader()
			return nil
		} else {
			tail.Logger.Printf("Stopping tail as file no longer exists: %s", tail.Filename)
			return ErrStop
		}
	case <-tail.changes.Truncated:
		// Always reopen truncated files (Follow is true)
		tail.Logger.Printf("Re-opening truncated file %s ...", tail.Filename)
		if err := tail.reopen(); err != nil {
			return err
		}
		tail.Logger.Printf("Successfully reopened truncated %s", tail.Filename)
		tail.openReader()
		return nil
	case <-tail.Dying():
		return ErrStop
	}
	panic("unreachable")
}

func (tail *Tail) openReader() {
	if tail.MaxLineSize > 0 {
		// add 2 to account for newline characters
		tail.reader = bufio.NewReaderSize(tail.file, tail.MaxLineSize+2)
	} else {
		tail.reader = bufio.NewReader(tail.file)
	}
}

func (tail *Tail) seekEnd() error {
	return tail.seekTo(SeekInfo{Offset: 0, Whence: 2})
}

func (tail *Tail) seekTo(pos SeekInfo) error {
	_, err := tail.file.Seek(pos.Offset, pos.Whence)
	if err != nil {
		return fmt.Errorf("Seek error on %s: %s", tail.Filename, err)
	}
	// Reset the read buffer whenever the file is re-seek'ed
	tail.reader.Reset(tail.file)
	return nil
}

// sendLine sends the line(s) to Lines channel, splitting longer lines
// if necessary. Return false if rate limit is reached.
func (tail *Tail) sendLine(line string) bool {
	now := time.Now()
	lines := []string{line}

	// Split longer lines
	if tail.MaxLineSize > 0 && len(line) > tail.MaxLineSize {
		lines = util.PartitionString(line, tail.MaxLineSize)
	}

	for _, line := range lines {
		tail.Lines <- &Line{line, now, nil}
	}

	if tail.Config.RateLimiter != nil {
		ok := tail.Config.RateLimiter.Pour(uint16(len(lines)))
		if !ok {
			tail.Logger.Printf("Leaky bucket full (%v); entering 1s cooloff period.\n",
				tail.Filename)
			return false
		}
	}

	return true
}

// Cleanup removes inotify watches added by the tail package. This function is
// meant to be invoked from a process's exit handler. Linux kernel may not
// automatically remove inotify watches after the process exits.
func (tail *Tail) Cleanup() {
	if tail.tracker != nil {
		tail.tracker.CloseAll()
	}
}
<|endoftext|>


<commit_before><commit_msg>A test binary.<commit_after><|endoftext|>


<commit_before><commit_msg>make the field property of NoPropertyError public<commit_after><|endoftext|>