<commit_before>package main

import (
	"crypto/sha1"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/dickeyxxx/golock"
	"github.com/franela/goreq"
	"github.com/heroku/heroku-cli/gode"
	"github.com/kardianos/osext"
	"github.com/ulikunitz/xz"
)

var updateTopic = &Topic{
	Name:        "update",
	Description: "update heroku-cli",
}

var updateCmd = &Command{
	Topic:            "update",
	Hidden:           true,
	Description:      "updates heroku-cli",
	DisableAnalytics: true,
	Args:             []Arg{{Name: "channel", Optional: true}},
	Flags:            []Flag{{Name: "background", Hidden: true}},
	Run: func(ctx *Context) {
		channel := ctx.Args.(map[string]string)["channel"]
		if channel == "" {
			channel = Channel
		}
		t := "foreground"
		if ctx.Flags["background"] == true {
			t = "background"
		}
		Update(channel, t)
	},
}

var binPath string
var updateLockPath = filepath.Join(AppDir(), "updating.lock")
var autoupdateFile = filepath.Join(AppDir(), "autoupdate")
var tmpPath = filepath.Join(AppDir(), "tmp")

func init() {
	binPath, _ = osext.Executable()
}

// Update updates the CLI and plugins
func Update(channel string, t string) {
	if !IsUpdateNeeded(t) {
		return
	}
	touchAutoupdateFile()
	updateCLI(channel)
	SetupNode()
	SetupBuiltinPlugins()
	updatePlugins()
	SubmitAnalytics()
	truncateErrorLog()
	cleanTmpDir()
}

func updatePlugins() {
	action("heroku-cli: Updating plugins", "", func() {
		plugins := PluginNamesNotSymlinked()
		if len(plugins) == 0 {
			return
		}
		packages, err := gode.OutdatedPackages(plugins...)
		WarnIfError(err)
		if len(packages) > 0 {
			for name, version := range packages {
				lockPlugin(name)
				WarnIfError(gode.InstallPackages(name + "@" + version))
				plugin, err := ParsePlugin(name)
				WarnIfError(err)
				AddPluginsToCache(plugin)
				unlockPlugin(name)
			}
			Errf(" done. Updated %d %s.\n", len(packages), plural("package", len(packages)))
		} else {
			Errln(" no plugins to update.")
		}
	})
}

func updateCLI(channel string) {
	if channel == "?" {
		// do not update dev version
		return
	}
	manifest, err := getUpdateManifest(channel)
	if err != nil {
		Warn("Error updating CLI")
		WarnIfError(err)
		return
	}
	if manifest.Version == Version && manifest.Channel == Channel {
		return
	}
	locked, err := golock.IsLocked(updateLockPath)
	LogIfError(err)
	if locked {
		Warn("Update in progress")
		return
	}
	LogIfError(golock.Lock(updateLockPath))
	unlock := func() {
		golock.Unlock(updateLockPath)
	}
	defer unlock()
	msg := fmt.Sprintf("heroku-cli: Updating to %s", manifest.Version)
	if manifest.Channel != "master" {
		msg = fmt.Sprintf("%s (%s)", msg, manifest.Channel)
	}
	action(msg, "done", func() {
		build := manifest.Builds[runtime.GOOS][runtime.GOARCH]
		// on windows we can't remove an existing file or remove the running binary
		// so we download the file to binName.new
		// move the running binary to binName.old (deleting any existing file first)
		// rename the downloaded file to binName
		tmpBinPathNew := binPath + ".new"
		tmpBinPathOld := binPath + ".old"
		if err := downloadBin(tmpBinPathNew, build.URL); err != nil {
			panic(err)
		}
		if fileSha1(tmpBinPathNew) != build.Sha1 {
			panic("SHA mismatch")
		}
		os.Remove(tmpBinPathOld)
		os.Rename(binPath, tmpBinPathOld)
		if err := os.Rename(tmpBinPathNew, binPath); err != nil {
			panic(err)
		}
		os.Remove(tmpBinPathOld)
		unlock()
		clearAutoupdateFile() // force full update
	})
	reexec() // reexec to finish updating with new code
}

// IsUpdateNeeded checks if an update is available
func IsUpdateNeeded(t string) bool {
	f, err := os.Stat(autoupdateFile)
	if err != nil {
		return true
	}
	if t == "background" {
		return time.Since(f.ModTime()) > 4*time.Hour
	} else if t == "block" {
		return time.Since(f.ModTime()) > 2160*time.Hour // 90 days
	}
	return true
}

func touchAutoupdateFile() {
	out, err := os.OpenFile(autoupdateFile, os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		panic(err)
	}
	out.WriteString(time.Now().String())
	out.Close()
}

// forces a full update on the next run
func clearAutoupdateFile() {
	WarnIfError(os.Remove(autoupdateFile))
}

type manifest struct {
	Channel, Version string
	Builds           map[string]map[string]struct {
		URL, Sha1 string
	}
}

func getUpdateManifest(channel string) (*manifest, error) {
	res, err := goreq.Request{
		Uri:       "https://cli-assets.heroku.com/" + channel + "/manifest.json",
		Timeout:   30 * time.Minute,
		ShowDebug: debugging,
	}.Do()
	if err != nil {
		return nil, err
	}
	var m manifest
	res.Body.FromJsonTo(&m)
	return &m, nil
}

func downloadBin(path, url string) error {
	out, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0755)
	if err != nil {
		return err
	}
	defer out.Close()
	res, err := goreq.Request{
		Uri:       url + ".xz",
		Timeout:   30 * time.Minute,
		ShowDebug: debugging,
	}.Do()
	if err != nil {
		return err
	}
	if res.StatusCode != 200 {
		b, _ := res.Body.ToString()
		return errors.New(b)
	}
	defer res.Body.Close()
	uncompressed, err := xz.NewReader(res.Body)
	if err != nil {
		return err
	}
	_, err = io.Copy(out, uncompressed)
	return err
}

func fileSha1(path string) string {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		panic(err)
	}
	return fmt.Sprintf("%x", sha1.Sum(data))
}

// TriggerBackgroundUpdate will trigger an update to the client in the background
func TriggerBackgroundUpdate() {
	wg.Add(1)
	go func() {
		wg.Done()
		if IsUpdateNeeded("background") {
			exec.Command(binPath, "update", "--background").Start()
		}
	}()
}

// restarts the CLI with the same arguments
func reexec() {
	Debugln("reexecing new CLI...")
	cmd := exec.Command(binPath, os.Args[1:]...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	Exit(getExitCode(cmd.Run()))
}

func truncateErrorLog() {
	Debugln("truncating error log...")
	body, err := ioutil.ReadFile(ErrLogPath)
	if err != nil {
		if !os.IsNotExist(err) {
			WarnIfError(err)
		}
		return
	}
	lines := strings.Split(string(body), "\n")
	lines = lines[maxint(len(lines)-1000, 0) : len(lines)-1]
	err = ioutil.WriteFile(ErrLogPath, []byte(strings.Join(lines, "\n")+"\n"), 0644)
	WarnIfError(err)
}

func maxint(a, b int) int {
	if a > b {
		return a
	}
	return b
}

func cleanTmpDir() {
	Debugln("cleaning up tmp dirs...")
	dirs, err := ioutil.ReadDir(tmpPath)
	if err != nil {
		WarnIfError(err)
		return
	}
	for _, dir := range dirs {
		if time.Since(dir.ModTime()) > 24*time.Hour {
			path := filepath.Join(tmpPath, dir.Name())
			Debugln("deleting " + path)
			WarnIfError(os.RemoveAll(path))
		}
	}
}
<commit_msg>wait before reexecing<commit_after>package main

import (
	"crypto/sha1"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/dickeyxxx/golock"
	"github.com/franela/goreq"
	"github.com/heroku/heroku-cli/gode"
	"github.com/kardianos/osext"
	"github.com/ulikunitz/xz"
)

var updateTopic = &Topic{
	Name:        "update",
	Description: "update heroku-cli",
}

var updateCmd = &Command{
	Topic:            "update",
	Hidden:           true,
	Description:      "updates heroku-cli",
	DisableAnalytics: true,
	Args:             []Arg{{Name: "channel", Optional: true}},
	Flags:            []Flag{{Name: "background", Hidden: true}},
	Run: func(ctx *Context) {
		channel := ctx.Args.(map[string]string)["channel"]
		if channel == "" {
			channel = Channel
		}
		t := "foreground"
		if ctx.Flags["background"] == true {
			t = "background"
		}
		Update(channel, t)
	},
}

var binPath string
var updateLockPath = filepath.Join(AppDir(), "updating.lock")
var autoupdateFile = filepath.Join(AppDir(), "autoupdate")
var tmpPath = filepath.Join(AppDir(), "tmp")

func init() {
	binPath, _ = osext.Executable()
}

// Update updates the CLI and plugins
func Update(channel string, t string) {
	if !IsUpdateNeeded(t) {
		return
	}
	touchAutoupdateFile()
	SubmitAnalytics()
	updateCLI(channel)
	SetupNode()
	SetupBuiltinPlugins()
	updatePlugins()
	truncateErrorLog()
	cleanTmpDir()
}

func updatePlugins() {
	action("heroku-cli: Updating plugins", "", func() {
		plugins := PluginNamesNotSymlinked()
		if len(plugins) == 0 {
			return
		}
		packages, err := gode.OutdatedPackages(plugins...)
		WarnIfError(err)
		if len(packages) > 0 {
			for name, version := range packages {
				lockPlugin(name)
				WarnIfError(gode.InstallPackages(name + "@" + version))
				plugin, err := ParsePlugin(name)
				WarnIfError(err)
				AddPluginsToCache(plugin)
				unlockPlugin(name)
			}
			Errf(" done. Updated %d %s.\n", len(packages), plural("package", len(packages)))
		} else {
			Errln(" no plugins to update.")
		}
	})
}

func updateCLI(channel string) {
	if channel == "?" {
		// do not update dev version
		return
	}
	manifest, err := getUpdateManifest(channel)
	if err != nil {
		Warn("Error updating CLI")
		WarnIfError(err)
		return
	}
	if manifest.Version == Version && manifest.Channel == Channel {
		return
	}
	locked, err := golock.IsLocked(updateLockPath)
	LogIfError(err)
	if locked {
		Warn("Update in progress")
		return
	}
	LogIfError(golock.Lock(updateLockPath))
	unlock := func() {
		golock.Unlock(updateLockPath)
	}
	defer unlock()
	msg := fmt.Sprintf("heroku-cli: Updating to %s", manifest.Version)
	if manifest.Channel != "master" {
		msg = fmt.Sprintf("%s (%s)", msg, manifest.Channel)
	}
	action(msg, "done", func() {
		build := manifest.Builds[runtime.GOOS][runtime.GOARCH]
		// on windows we can't remove an existing file or remove the running binary
		// so we download the file to binName.new
		// move the running binary to binName.old (deleting any existing file first)
		// rename the downloaded file to binName
		tmpBinPathNew := binPath + ".new"
		tmpBinPathOld := binPath + ".old"
		if err := downloadBin(tmpBinPathNew, build.URL); err != nil {
			panic(err)
		}
		if fileSha1(tmpBinPathNew) != build.Sha1 {
			panic("SHA mismatch")
		}
		os.Remove(tmpBinPathOld)
		os.Rename(binPath, tmpBinPathOld)
		if err := os.Rename(tmpBinPathNew, binPath); err != nil {
			panic(err)
		}
		os.Remove(tmpBinPathOld)
		unlock()
		clearAutoupdateFile() // force full update
	})
	wg.Wait()
	reexec() // reexec to finish updating with new code
}

// IsUpdateNeeded checks if an update is available
func IsUpdateNeeded(t string) bool {
	f, err := os.Stat(autoupdateFile)
	if err != nil {
		return true
	}
	if t == "background" {
		return time.Since(f.ModTime()) > 4*time.Hour
	} else if t == "block" {
		return time.Since(f.ModTime()) > 2160*time.Hour // 90 days
	}
	return true
}

func touchAutoupdateFile() {
	out, err := os.OpenFile(autoupdateFile, os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		panic(err)
	}
	out.WriteString(time.Now().String())
	out.Close()
}

// forces a full update on the next run
func clearAutoupdateFile() {
	WarnIfError(os.Remove(autoupdateFile))
}

type manifest struct {
	Channel, Version string
	Builds           map[string]map[string]struct {
		URL, Sha1 string
	}
}

func getUpdateManifest(channel string) (*manifest, error) {
	res, err := goreq.Request{
		Uri:       "https://cli-assets.heroku.com/" + channel + "/manifest.json",
		Timeout:   30 * time.Minute,
		ShowDebug: debugging,
	}.Do()
	if err != nil {
		return nil, err
	}
	var m manifest
	res.Body.FromJsonTo(&m)
	return &m, nil
}

func downloadBin(path, url string) error {
	out, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0755)
	if err != nil {
		return err
	}
	defer out.Close()
	res, err := goreq.Request{
		Uri:       url + ".xz",
		Timeout:   30 * time.Minute,
		ShowDebug: debugging,
	}.Do()
	if err != nil {
		return err
	}
	if res.StatusCode != 200 {
		b, _ := res.Body.ToString()
		return errors.New(b)
	}
	defer res.Body.Close()
	uncompressed, err := xz.NewReader(res.Body)
	if err != nil {
		return err
	}
	_, err = io.Copy(out, uncompressed)
	return err
}

func fileSha1(path string) string {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		panic(err)
	}
	return fmt.Sprintf("%x", sha1.Sum(data))
}

// TriggerBackgroundUpdate will trigger an update to the client in the background
func TriggerBackgroundUpdate() {
	wg.Add(1)
	go func() {
		wg.Done()
		if IsUpdateNeeded("background") {
			exec.Command(binPath, "update", "--background").Start()
		}
	}()
}

// restarts the CLI with the same arguments
func reexec() {
	Debugln("reexecing new CLI...")
	cmd := exec.Command(binPath, os.Args[1:]...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	Exit(getExitCode(cmd.Run()))
}

func truncateErrorLog() {
	Debugln("truncating error log...")
	body, err := ioutil.ReadFile(ErrLogPath)
	if err != nil {
		if !os.IsNotExist(err) {
			WarnIfError(err)
		}
		return
	}
	lines := strings.Split(string(body), "\n")
	lines = lines[maxint(len(lines)-1000, 0) : len(lines)-1]
	err = ioutil.WriteFile(ErrLogPath, []byte(strings.Join(lines, "\n")+"\n"), 0644)
	WarnIfError(err)
}

func maxint(a, b int) int {
	if a > b {
		return a
	}
	return b
}

func cleanTmpDir() {
	Debugln("cleaning up tmp dirs...")
	dirs, err := ioutil.ReadDir(tmpPath)
	if err != nil {
		WarnIfError(err)
		return
	}
	for _, dir := range dirs {
		if time.Since(dir.ModTime()) > 24*time.Hour {
			path := filepath.Join(tmpPath, dir.Name())
			Debugln("deleting " + path)
			WarnIfError(os.RemoveAll(path))
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 The Linux Foundation
//
// 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 schema_test

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net/url"
	"os"
	"strings"
	"testing"

	"github.com/opencontainers/image-spec/schema"
	"github.com/pkg/errors"
	"github.com/russross/blackfriday"
)

var (
	errFormatInvalid = errors.New("format: invalid")
)

func TestValidateManifest(t *testing.T) {
	validate(t, "../manifest.md")
}

func TestValidateSerialization(t *testing.T) {
	validate(t, "../serialization.md")
}

// TODO(sur): include examples from all specification files
func validate(t *testing.T, name string) {
	m, err := os.Open(name)
	if err != nil {
		t.Fatal(err)
	}

	examples, err := extractExamples(m)
	if err != nil {
		t.Fatal(err)
	}

	for _, example := range examples {
		if example.Err == errFormatInvalid { // ignore
			continue
		}

		if example.Err != nil {
			printFields(t, "error", example.Mediatype, example.Title, example.Err)
			t.Error(err)
			continue
		}

		err = schema.Validator(example.Mediatype).Validate(strings.NewReader(example.Body))
		if err == nil {
			printFields(t, "ok", example.Mediatype, example.Title)
			t.Log(example.Body, "---")
			continue
		}

		var errs []error
		if verr, ok := errors.Cause(err).(schema.ValidationError); ok {
			errs = verr.Errs
		} else {
			printFields(t, "error", example.Mediatype, example.Title, err)
			t.Error(err)
			t.Log(example.Body, "---")
			continue
		}

		for _, err := range errs {
			// TOOD(stevvooe): This is nearly useless without file, line no.
			printFields(t, "invalid", example.Mediatype, example.Title)
			t.Error(err)
			fmt.Println(example.Body, "---")
			continue
		}
	}
}

// renderer allows one to incercept fenced blocks in markdown documents.
type renderer struct {
	blackfriday.Renderer
	fn func(text []byte, lang string)
}

func (r *renderer) BlockCode(out *bytes.Buffer, text []byte, lang string) {
	r.fn(text, lang)
	r.Renderer.BlockCode(out, text, lang)
}

type example struct {
	Lang      string // gets raw "lang" field
	Title     string
	Mediatype string
	Body      string
	Err       error

	// TODO(stevvooe): Figure out how to keep track of revision, file, line so
	// that we can trace back verification output.
}

// parseExample treats the field as a syntax,attribute tuple separated by a comma.
// Attributes are encoded as a url values.
//
// An example of this is `json,title=Foo%20Bar&mediatype=application/json. We
// get that the "lang" is json, the title is "Foo Bar" and the mediatype is
// "application/json".
//
// This preserves syntax highlighting and lets us tag examples with further
// metadata.
func parseExample(lang, body string) (e example) {
	e.Lang = lang
	e.Body = body

	parts := strings.SplitN(lang, ",", 2)
	if len(parts) < 2 {
		e.Err = errFormatInvalid
		return
	}

	m, err := url.ParseQuery(parts[1])
	if err != nil {
		e.Err = err
		return
	}

	e.Mediatype = m.Get("mediatype")
	e.Title = m.Get("title")
	return
}

func extractExamples(rd io.Reader) ([]example, error) {
	p, err := ioutil.ReadAll(rd)
	if err != nil {
		return nil, err
	}

	var examples []example
	renderer := &renderer{
		Renderer: blackfriday.HtmlRenderer(0, "test test", ""),
		fn: func(text []byte, lang string) {
			examples = append(examples, parseExample(lang, string(text)))
		},
	}

	// just pass over the markdown and ignore the rendered result. We just want
	// the side-effect of calling back for each code block.
	// TODO(stevvooe): Consider just parsing these with a scanner. It will be
	// faster and we can retain file, line no.
	blackfriday.MarkdownOptions(p, renderer, blackfriday.Options{
		Extensions: blackfriday.EXTENSION_FENCED_CODE,
	})

	return examples, nil
}

// printFields prints each value tab separated.
func printFields(t *testing.T, vs ...interface{}) {
	var ss []string
	for _, f := range vs {
		ss = append(ss, fmt.Sprint(f))
	}
	t.Log(strings.Join(ss, "\t"))
}
<commit_msg>schema/spec_test: ignore examples having no media type<commit_after>// Copyright 2016 The Linux Foundation
//
// 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 schema_test

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net/url"
	"os"
	"strings"
	"testing"

	"github.com/opencontainers/image-spec/schema"
	"github.com/pkg/errors"
	"github.com/russross/blackfriday"
)

var (
	errFormatInvalid = errors.New("format: invalid")
)

func TestValidateManifest(t *testing.T) {
	validate(t, "../manifest.md")
}

func TestValidateSerialization(t *testing.T) {
	validate(t, "../serialization.md")
}

// TODO(sur): include examples from all specification files
func validate(t *testing.T, name string) {
	m, err := os.Open(name)
	if err != nil {
		t.Fatal(err)
	}

	examples, err := extractExamples(m)
	if err != nil {
		t.Fatal(err)
	}

	for _, example := range examples {
		if example.Err == errFormatInvalid && example.Mediatype == "" { // ignore
			continue
		}

		if example.Err != nil {
			printFields(t, "error", example.Mediatype, example.Title, example.Err)
			t.Error(err)
			continue
		}

		err = schema.Validator(example.Mediatype).Validate(strings.NewReader(example.Body))
		if err == nil {
			printFields(t, "ok", example.Mediatype, example.Title)
			t.Log(example.Body, "---")
			continue
		}

		var errs []error
		if verr, ok := errors.Cause(err).(schema.ValidationError); ok {
			errs = verr.Errs
		} else {
			printFields(t, "error", example.Mediatype, example.Title, err)
			t.Error(err)
			t.Log(example.Body, "---")
			continue
		}

		for _, err := range errs {
			// TOOD(stevvooe): This is nearly useless without file, line no.
			printFields(t, "invalid", example.Mediatype, example.Title)
			t.Error(err)
			fmt.Println(example.Body, "---")
			continue
		}
	}
}

// renderer allows one to incercept fenced blocks in markdown documents.
type renderer struct {
	blackfriday.Renderer
	fn func(text []byte, lang string)
}

func (r *renderer) BlockCode(out *bytes.Buffer, text []byte, lang string) {
	r.fn(text, lang)
	r.Renderer.BlockCode(out, text, lang)
}

type example struct {
	Lang      string // gets raw "lang" field
	Title     string
	Mediatype string
	Body      string
	Err       error

	// TODO(stevvooe): Figure out how to keep track of revision, file, line so
	// that we can trace back verification output.
}

// parseExample treats the field as a syntax,attribute tuple separated by a comma.
// Attributes are encoded as a url values.
//
// An example of this is `json,title=Foo%20Bar&mediatype=application/json. We
// get that the "lang" is json, the title is "Foo Bar" and the mediatype is
// "application/json".
//
// This preserves syntax highlighting and lets us tag examples with further
// metadata.
func parseExample(lang, body string) (e example) {
	e.Lang = lang
	e.Body = body

	parts := strings.SplitN(lang, ",", 2)
	if len(parts) < 2 {
		e.Err = errFormatInvalid
		return
	}

	m, err := url.ParseQuery(parts[1])
	if err != nil {
		e.Err = err
		return
	}

	e.Mediatype = m.Get("mediatype")
	e.Title = m.Get("title")
	return
}

func extractExamples(rd io.Reader) ([]example, error) {
	p, err := ioutil.ReadAll(rd)
	if err != nil {
		return nil, err
	}

	var examples []example
	renderer := &renderer{
		Renderer: blackfriday.HtmlRenderer(0, "test test", ""),
		fn: func(text []byte, lang string) {
			examples = append(examples, parseExample(lang, string(text)))
		},
	}

	// just pass over the markdown and ignore the rendered result. We just want
	// the side-effect of calling back for each code block.
	// TODO(stevvooe): Consider just parsing these with a scanner. It will be
	// faster and we can retain file, line no.
	blackfriday.MarkdownOptions(p, renderer, blackfriday.Options{
		Extensions: blackfriday.EXTENSION_FENCED_CODE,
	})

	return examples, nil
}

// printFields prints each value tab separated.
func printFields(t *testing.T, vs ...interface{}) {
	var ss []string
	for _, f := range vs {
		ss = append(ss, fmt.Sprint(f))
	}
	t.Log(strings.Join(ss, "\t"))
}
<|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 tikv

import (
	"container/list"
	"sync"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/kvproto/pkg/kvrpcpb"
	"github.com/pingcap/tidb/store/tikv/tikvrpc"
	goctx "golang.org/x/net/context"
)

const resolvedCacheSize = 512

// LockResolver resolves locks and also caches resolved txn status.
type LockResolver struct {
	store *tikvStore
	mu    struct {
		sync.RWMutex
		// Cache resolved txns (FIFO, txn id -> txnStatus).
		resolved       map[uint64]TxnStatus
		recentResolved *list.List
	}
}

func newLockResolver(store *tikvStore) *LockResolver {
	r := &LockResolver{
		store: store,
	}
	r.mu.resolved = make(map[uint64]TxnStatus)
	r.mu.recentResolved = list.New()
	return r
}

// TxnStatus represents a txn's final status. It should be Commit or Rollback.
type TxnStatus uint64

// IsCommitted returns true if the txn's final status is Commit.
func (s TxnStatus) IsCommitted() bool { return s > 0 }

// CommitTS returns the txn's commitTS. It is valid iff `IsCommitted` is true.
func (s TxnStatus) CommitTS() uint64 { return uint64(s) }

// By default, locks after 3000ms is considered unusual (the client created the
// lock might be dead). Other client may cleanup this kind of lock.
// For locks created recently, we will do backoff and retry.
var defaultLockTTL uint64 = 3000

// TODO: Consider if it's appropriate.
var maxLockTTL uint64 = 120000

// ttl = ttlFactor * sqrt(writeSizeInMiB)
var ttlFactor = 6000

// Lock represents a lock from tikv server.
type Lock struct {
	Key     []byte
	Primary []byte
	TxnID   uint64
	TTL     uint64
}

func newLock(l *kvrpcpb.LockInfo) *Lock {
	ttl := l.GetLockTtl()
	if ttl == 0 {
		ttl = defaultLockTTL
	}
	return &Lock{
		Key:     l.GetKey(),
		Primary: l.GetPrimaryLock(),
		TxnID:   l.GetLockVersion(),
		TTL:     ttl,
	}
}

func (lr *LockResolver) saveResolved(txnID uint64, status TxnStatus) {
	lr.mu.Lock()
	defer lr.mu.Unlock()

	if _, ok := lr.mu.resolved[txnID]; ok {
		return
	}
	lr.mu.resolved[txnID] = status
	lr.mu.recentResolved.PushBack(txnID)
	if len(lr.mu.resolved) > resolvedCacheSize {
		front := lr.mu.recentResolved.Front()
		delete(lr.mu.resolved, front.Value.(uint64))
		lr.mu.recentResolved.Remove(front)
	}
}

func (lr *LockResolver) getResolved(txnID uint64) (TxnStatus, bool) {
	lr.mu.RLock()
	defer lr.mu.RUnlock()

	s, ok := lr.mu.resolved[txnID]
	return s, ok
}

// ResolveLocks tries to resolve Locks. The resolving process is in 3 steps:
// 1) Use the `lockTTL` to pick up all expired locks. Only locks that are too
//    old are considered orphan locks and will be handled later. If all locks
//    are expired then all locks will be resolved so the returned `ok` will be
//    true, otherwise caller should sleep a while before retry.
// 2) For each lock, query the primary key to get txn(which left the lock)'s
//    commit status.
// 3) Send `ResolveLock` cmd to the lock's region to resolve all locks belong to
//    the same transaction.
func (lr *LockResolver) ResolveLocks(bo *Backoffer, locks []*Lock) (ok bool, err error) {
	if len(locks) == 0 {
		return true, nil
	}

	lockResolverCounter.WithLabelValues("resolve").Inc()

	var expiredLocks []*Lock
	for _, l := range locks {
		if lr.store.oracle.IsExpired(l.TxnID, l.TTL) {
			lockResolverCounter.WithLabelValues("expired").Inc()
			expiredLocks = append(expiredLocks, l)
		} else {
			lockResolverCounter.WithLabelValues("not_expired").Inc()
		}
	}
	if len(expiredLocks) == 0 {
		return false, nil
	}

	// TxnID -> []Region, record resolved Regions.
	// TODO: Maybe put it in LockResolver and share by all txns.
	cleanTxns := make(map[uint64]map[RegionVerID]struct{})
	for _, l := range expiredLocks {
		status, err := lr.getTxnStatus(bo, l.TxnID, l.Primary)
		if err != nil {
			return false, errors.Trace(err)
		}

		cleanRegions := cleanTxns[l.TxnID]
		if cleanRegions == nil {
			cleanRegions = make(map[RegionVerID]struct{})
			cleanTxns[l.TxnID] = cleanRegions
		}

		err = lr.resolveLock(bo, l, status, cleanRegions)
		if err != nil {
			return false, errors.Trace(err)
		}
	}
	return len(expiredLocks) == len(locks), nil
}

// GetTxnStatus queries tikv-server for a txn's status (commit/rollback).
// If the primary key is still locked, it will launch a Rollback to abort it.
// To avoid unnecessarily aborting too many txns, it is wiser to wait a few
// seconds before calling it after Prewrite.
func (lr *LockResolver) GetTxnStatus(txnID uint64, primary []byte) (TxnStatus, error) {
	bo := NewBackoffer(cleanupMaxBackoff, goctx.Background())
	status, err := lr.getTxnStatus(bo, txnID, primary)
	return status, errors.Trace(err)
}

func (lr *LockResolver) getTxnStatus(bo *Backoffer, txnID uint64, primary []byte) (TxnStatus, error) {
	if s, ok := lr.getResolved(txnID); ok {
		return s, nil
	}

	lockResolverCounter.WithLabelValues("query_txn_status").Inc()

	var status TxnStatus
	req := &tikvrpc.Request{
		Type: tikvrpc.CmdCleanup,
		Cleanup: &kvrpcpb.CleanupRequest{
			Key:          primary,
			StartVersion: txnID,
		},
	}
	for {
		loc, err := lr.store.regionCache.LocateKey(bo, primary)
		if err != nil {
			return status, errors.Trace(err)
		}
		resp, err := lr.store.SendReq(bo, req, loc.Region, readTimeoutShort)
		if err != nil {
			return status, errors.Trace(err)
		}
		regionErr, err := resp.GetRegionError()
		if err != nil {
			return status, errors.Trace(err)
		}
		if regionErr != nil {
			err = bo.Backoff(boRegionMiss, errors.New(regionErr.String()))
			if err != nil {
				return status, errors.Trace(err)
			}
			continue
		}
		cmdResp := resp.Cleanup
		if cmdResp == nil {
			return status, errors.Trace(errBodyMissing)
		}
		if keyErr := cmdResp.GetError(); keyErr != nil {
			err = errors.Errorf("unexpected cleanup err: %s, tid: %v", keyErr, txnID)
			log.Error(err)
			return status, err
		}
		if cmdResp.CommitVersion != 0 {
			status = TxnStatus(cmdResp.GetCommitVersion())
			lockResolverCounter.WithLabelValues("query_txn_status_committed").Inc()
		} else {
			lockResolverCounter.WithLabelValues("query_txn_status_rolled_back").Inc()
		}
		lr.saveResolved(txnID, status)
		return status, nil
	}
}

func (lr *LockResolver) resolveLock(bo *Backoffer, l *Lock, status TxnStatus, cleanRegions map[RegionVerID]struct{}) error {
	lockResolverCounter.WithLabelValues("query_resolve_locks").Inc()
	for {
		loc, err := lr.store.regionCache.LocateKey(bo, l.Key)
		if err != nil {
			return errors.Trace(err)
		}
		if _, ok := cleanRegions[loc.Region]; ok {
			return nil
		}
		req := &tikvrpc.Request{
			Type: tikvrpc.CmdResolveLock,
			ResolveLock: &kvrpcpb.ResolveLockRequest{
				StartVersion: l.TxnID,
			},
		}
		if status.IsCommitted() {
			req.ResolveLock.CommitVersion = status.CommitTS()
		}
		resp, err := lr.store.SendReq(bo, req, loc.Region, readTimeoutShort)
		if err != nil {
			return errors.Trace(err)
		}
		regionErr, err := resp.GetRegionError()
		if err != nil {
			return errors.Trace(err)
		}
		if regionErr != nil {
			err = bo.Backoff(boRegionMiss, errors.New(regionErr.String()))
			if err != nil {
				return errors.Trace(err)
			}
			continue
		}
		cmdResp := resp.ResolveLock
		if cmdResp == nil {
			return errors.Trace(errBodyMissing)
		}
		if keyErr := cmdResp.GetError(); keyErr != nil {
			err = errors.Errorf("unexpected resolve err: %s, lock: %v", keyErr, l)
			log.Error(err)
			return err
		}
		cleanRegions[loc.Region] = struct{}{}
		return nil
	}
}
<commit_msg>store/tikv: export LockResolver. (#3585)<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 tikv

import (
	"container/list"
	"fmt"
	"sync"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/kvproto/pkg/kvrpcpb"
	"github.com/pingcap/pd/pd-client"
	"github.com/pingcap/tidb/store/tikv/tikvrpc"
	goctx "golang.org/x/net/context"
)

const resolvedCacheSize = 512

// LockResolver resolves locks and also caches resolved txn status.
type LockResolver struct {
	store *tikvStore
	mu    struct {
		sync.RWMutex
		// Cache resolved txns (FIFO, txn id -> txnStatus).
		resolved       map[uint64]TxnStatus
		recentResolved *list.List
	}
}

func newLockResolver(store *tikvStore) *LockResolver {
	r := &LockResolver{
		store: store,
	}
	r.mu.resolved = make(map[uint64]TxnStatus)
	r.mu.recentResolved = list.New()
	return r
}

// NewLockResolver creates a LockResolver.
// It is exported for other services to use. For instance, binlog service needs
// to determine a transaction's commit state.
func NewLockResolver(etcdAddrs []string) (*LockResolver, error) {
	pdCli, err := pd.NewClient(etcdAddrs)
	if err != nil {
		return nil, errors.Trace(err)
	}
	uuid := fmt.Sprintf("tikv-%v", pdCli.GetClusterID(goctx.TODO()))
	s, err := newTikvStore(uuid, &codecPDClient{pdCli}, newRPCClient(), false)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return s.lockResolver, nil
}

// TxnStatus represents a txn's final status. It should be Commit or Rollback.
type TxnStatus uint64

// IsCommitted returns true if the txn's final status is Commit.
func (s TxnStatus) IsCommitted() bool { return s > 0 }

// CommitTS returns the txn's commitTS. It is valid iff `IsCommitted` is true.
func (s TxnStatus) CommitTS() uint64 { return uint64(s) }

// By default, locks after 3000ms is considered unusual (the client created the
// lock might be dead). Other client may cleanup this kind of lock.
// For locks created recently, we will do backoff and retry.
var defaultLockTTL uint64 = 3000

// TODO: Consider if it's appropriate.
var maxLockTTL uint64 = 120000

// ttl = ttlFactor * sqrt(writeSizeInMiB)
var ttlFactor = 6000

// Lock represents a lock from tikv server.
type Lock struct {
	Key     []byte
	Primary []byte
	TxnID   uint64
	TTL     uint64
}

func newLock(l *kvrpcpb.LockInfo) *Lock {
	ttl := l.GetLockTtl()
	if ttl == 0 {
		ttl = defaultLockTTL
	}
	return &Lock{
		Key:     l.GetKey(),
		Primary: l.GetPrimaryLock(),
		TxnID:   l.GetLockVersion(),
		TTL:     ttl,
	}
}

func (lr *LockResolver) saveResolved(txnID uint64, status TxnStatus) {
	lr.mu.Lock()
	defer lr.mu.Unlock()

	if _, ok := lr.mu.resolved[txnID]; ok {
		return
	}
	lr.mu.resolved[txnID] = status
	lr.mu.recentResolved.PushBack(txnID)
	if len(lr.mu.resolved) > resolvedCacheSize {
		front := lr.mu.recentResolved.Front()
		delete(lr.mu.resolved, front.Value.(uint64))
		lr.mu.recentResolved.Remove(front)
	}
}

func (lr *LockResolver) getResolved(txnID uint64) (TxnStatus, bool) {
	lr.mu.RLock()
	defer lr.mu.RUnlock()

	s, ok := lr.mu.resolved[txnID]
	return s, ok
}

// ResolveLocks tries to resolve Locks. The resolving process is in 3 steps:
// 1) Use the `lockTTL` to pick up all expired locks. Only locks that are too
//    old are considered orphan locks and will be handled later. If all locks
//    are expired then all locks will be resolved so the returned `ok` will be
//    true, otherwise caller should sleep a while before retry.
// 2) For each lock, query the primary key to get txn(which left the lock)'s
//    commit status.
// 3) Send `ResolveLock` cmd to the lock's region to resolve all locks belong to
//    the same transaction.
func (lr *LockResolver) ResolveLocks(bo *Backoffer, locks []*Lock) (ok bool, err error) {
	if len(locks) == 0 {
		return true, nil
	}

	lockResolverCounter.WithLabelValues("resolve").Inc()

	var expiredLocks []*Lock
	for _, l := range locks {
		if lr.store.oracle.IsExpired(l.TxnID, l.TTL) {
			lockResolverCounter.WithLabelValues("expired").Inc()
			expiredLocks = append(expiredLocks, l)
		} else {
			lockResolverCounter.WithLabelValues("not_expired").Inc()
		}
	}
	if len(expiredLocks) == 0 {
		return false, nil
	}

	// TxnID -> []Region, record resolved Regions.
	// TODO: Maybe put it in LockResolver and share by all txns.
	cleanTxns := make(map[uint64]map[RegionVerID]struct{})
	for _, l := range expiredLocks {
		status, err := lr.getTxnStatus(bo, l.TxnID, l.Primary)
		if err != nil {
			return false, errors.Trace(err)
		}

		cleanRegions := cleanTxns[l.TxnID]
		if cleanRegions == nil {
			cleanRegions = make(map[RegionVerID]struct{})
			cleanTxns[l.TxnID] = cleanRegions
		}

		err = lr.resolveLock(bo, l, status, cleanRegions)
		if err != nil {
			return false, errors.Trace(err)
		}
	}
	return len(expiredLocks) == len(locks), nil
}

// GetTxnStatus queries tikv-server for a txn's status (commit/rollback).
// If the primary key is still locked, it will launch a Rollback to abort it.
// To avoid unnecessarily aborting too many txns, it is wiser to wait a few
// seconds before calling it after Prewrite.
func (lr *LockResolver) GetTxnStatus(txnID uint64, primary []byte) (TxnStatus, error) {
	bo := NewBackoffer(cleanupMaxBackoff, goctx.Background())
	status, err := lr.getTxnStatus(bo, txnID, primary)
	return status, errors.Trace(err)
}

func (lr *LockResolver) getTxnStatus(bo *Backoffer, txnID uint64, primary []byte) (TxnStatus, error) {
	if s, ok := lr.getResolved(txnID); ok {
		return s, nil
	}

	lockResolverCounter.WithLabelValues("query_txn_status").Inc()

	var status TxnStatus
	req := &tikvrpc.Request{
		Type: tikvrpc.CmdCleanup,
		Cleanup: &kvrpcpb.CleanupRequest{
			Key:          primary,
			StartVersion: txnID,
		},
	}
	for {
		loc, err := lr.store.regionCache.LocateKey(bo, primary)
		if err != nil {
			return status, errors.Trace(err)
		}
		resp, err := lr.store.SendReq(bo, req, loc.Region, readTimeoutShort)
		if err != nil {
			return status, errors.Trace(err)
		}
		regionErr, err := resp.GetRegionError()
		if err != nil {
			return status, errors.Trace(err)
		}
		if regionErr != nil {
			err = bo.Backoff(boRegionMiss, errors.New(regionErr.String()))
			if err != nil {
				return status, errors.Trace(err)
			}
			continue
		}
		cmdResp := resp.Cleanup
		if cmdResp == nil {
			return status, errors.Trace(errBodyMissing)
		}
		if keyErr := cmdResp.GetError(); keyErr != nil {
			err = errors.Errorf("unexpected cleanup err: %s, tid: %v", keyErr, txnID)
			log.Error(err)
			return status, err
		}
		if cmdResp.CommitVersion != 0 {
			status = TxnStatus(cmdResp.GetCommitVersion())
			lockResolverCounter.WithLabelValues("query_txn_status_committed").Inc()
		} else {
			lockResolverCounter.WithLabelValues("query_txn_status_rolled_back").Inc()
		}
		lr.saveResolved(txnID, status)
		return status, nil
	}
}

func (lr *LockResolver) resolveLock(bo *Backoffer, l *Lock, status TxnStatus, cleanRegions map[RegionVerID]struct{}) error {
	lockResolverCounter.WithLabelValues("query_resolve_locks").Inc()
	for {
		loc, err := lr.store.regionCache.LocateKey(bo, l.Key)
		if err != nil {
			return errors.Trace(err)
		}
		if _, ok := cleanRegions[loc.Region]; ok {
			return nil
		}
		req := &tikvrpc.Request{
			Type: tikvrpc.CmdResolveLock,
			ResolveLock: &kvrpcpb.ResolveLockRequest{
				StartVersion: l.TxnID,
			},
		}
		if status.IsCommitted() {
			req.ResolveLock.CommitVersion = status.CommitTS()
		}
		resp, err := lr.store.SendReq(bo, req, loc.Region, readTimeoutShort)
		if err != nil {
			return errors.Trace(err)
		}
		regionErr, err := resp.GetRegionError()
		if err != nil {
			return errors.Trace(err)
		}
		if regionErr != nil {
			err = bo.Backoff(boRegionMiss, errors.New(regionErr.String()))
			if err != nil {
				return errors.Trace(err)
			}
			continue
		}
		cmdResp := resp.ResolveLock
		if cmdResp == nil {
			return errors.Trace(errBodyMissing)
		}
		if keyErr := cmdResp.GetError(); keyErr != nil {
			err = errors.Errorf("unexpected resolve err: %s, lock: %v", keyErr, l)
			log.Error(err)
			return err
		}
		cleanRegions[loc.Region] = struct{}{}
		return nil
	}
}
<|endoftext|>


<commit_before>package box

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"io"
	"regexp"
	"sort"
	"strings"

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

var AnsiCutter = regexp.MustCompile("\x1B[^a-zA-Z]*[A-Za-z]")

func Print(ctx context.Context, nodes []string, out io.Writer) bool {
	b := New()
	b.Height = 0
	value, _, _ := b.Print(ctx, nodes, 0, out)
	return value
}

func (b *box_t) Print(ctx context.Context,
	nodes []string,
	offset int,
	out io.Writer) (bool, int, int) {

	selected, columns, nlines := b.PrintNoLastLineFeed(ctx, nodes, offset, out)

	// append last linefeed.
	if nlines > 0 {
		fmt.Fprintln(out)
	}
	b.Cache = nil
	return selected, columns, nlines
}

func (b *box_t) PrintNoLastLineFeed(ctx context.Context,
	nodes []string,
	offset int,
	out io.Writer) (bool, int, int) {

	if len(nodes) <= 0 {
		return true, 0, 0
	}

	maxLen := 1
	for _, finfo := range nodes {
		length := runewidth.StringWidth(AnsiCutter.ReplaceAllString(finfo, ""))
		if length > maxLen {
			maxLen = length
		}
	}
	nodePerLine := (b.Width - 1) / (maxLen + 1)
	if nodePerLine <= 0 {
		nodePerLine = 1
	}
	nlines := (len(nodes) + nodePerLine - 1) / nodePerLine

	lines := make([][]byte, nlines)
	row := 0
	for _, finfo := range nodes {
		lines[row] = append(lines[row], finfo...)
		w := runewidth.StringWidth(AnsiCutter.ReplaceAllString(finfo, ""))
		if maxLen < b.Width {
			for i := maxLen + 1; i > w; i-- {
				lines[row] = append(lines[row], ' ')
			}
		}
		row++
		if row >= nlines {
			row = 0
		}
	}
	i_end := len(lines)
	if b.Height > 0 {
		if i_end >= offset+b.Height {
			i_end = offset + b.Height
		}
	}

	if b.Cache == nil {
		b.Cache = make([][]byte, b.Height)
	}
	i := offset
	y := 0
	for {
		if y >= len(b.Cache) {
			b.Cache = append(b.Cache, []byte{})
		}
		// assertion
		if i >= len(lines) {
			panic(fmt.Sprintf("len(lines)==%d i==%d", len(lines), i))
		}
		if bytes.Compare(lines[i], b.Cache[y]) != 0 {
			fmt.Fprint(out, strings.TrimRight(string(lines[i]), " "))
			if len(b.Cache[y]) > 0 {
				fmt.Fprint(out, ERASE_LINE)
			}
			b.Cache[y] = lines[i]
		}
		y++
		if ctx != nil {
			select {
			case <-ctx.Done():
				return false, nodePerLine, nlines
			default:
			}
		}
		i++
		if i >= i_end {
			break
		}
		fmt.Fprintln(out)
	}
	return true, nodePerLine, nlines
}

const (
	CURSOR_OFF = "\x1B[?25l"
	CURSOR_ON  = "\x1B[?25h"
	BOLD_ON    = "\x1B[0;47;30m"
	BOLD_ON2   = "\x1B[0;1;7m"
	BOLD_OFF   = "\x1B[0m"
	UP_N       = "\x1B[%dA"
	ERASE_LINE = "\x1B[0K"
)

func truncate(s string, w int) string {
	return runewidth.Truncate(strings.TrimSpace(s), w, "")
}

const (
	NONE         = 0
	LEFT         = 1
	DOWN         = 2
	UP           = 3
	RIGHT        = 4
	ENTER        = 5
	LEAVE        = 6
	SELECT_DOWN  = 7
	SELECT_UP    = 8
	SELECT_LEFT  = 9
	SELECT_RIGHT = 10
)

type nodeT struct {
	Index int
	Text  string
}

// Choice returns selected string
func Choice(sources []string, out io.Writer) string {
	n := Choose(sources, out)
	if n < 0 {
		return ""
	}
	return sources[n]
}

func ChoiceMulti(sources []string, out io.Writer) []string {
	list := ChooseMulti(sources, out)
	values := make([]string, 0, len(list))
	for _, index := range list {
		values = append(values, sources[index])
	}
	return values
}

// Choice returns the index of selected string
func ChooseMulti(sources []string, out io.Writer) []int {
	cursor := 0
	selected := make(map[int]struct{})

	nodes := make([]*nodeT, 0, len(sources))
	draws := make([]string, 0, len(sources))
	b := New()
	defer b.Close()
	for i, text := range sources {
		val := truncate(text, b.Width-1)
		if val != "" {
			nodes = append(nodes, &nodeT{Index: i, Text: val})
			draws = append(draws, val)
		}
	}
	io.WriteString(out, CURSOR_OFF)
	defer io.WriteString(out, CURSOR_ON)

	if len(nodes) <= 0 {
		nodes = []*nodeT{&nodeT{-1, ""}}
		draws = []string{""}
	}

	offset := 0
	for {
		for index := range selected {
			draws[index] = BOLD_ON + truncate(nodes[index].Text, b.Width-2) + BOLD_OFF
		}
		draws[cursor] = BOLD_ON2 + truncate(nodes[cursor].Text, b.Width-2) + BOLD_OFF
		status, _, h := b.PrintNoLastLineFeed(nil, draws, offset, out)
		if !status {
			return []int{}
		}
		for index, _ := range selected {
			draws[index] = truncate(nodes[index].Text, b.Width-2)
		}
		draws[cursor] = truncate(nodes[cursor].Text, b.Width-2)
		last := cursor

		doSelect := func() {
			if _, ok := selected[cursor]; ok {
				delete(selected, cursor)
			} else {
				selected[cursor] = struct{}{}
			}
		}

		for last == cursor {
			if bw, ok := out.(*bufio.Writer); ok {
				bw.Flush()
			}
			key, err := b.GetKey()
			if err != nil {
				continue
			}
			switch key {
			case "h", K_CTRL_B, K_LEFT:
				cursor = (cursor + len(nodes) - h) % len(nodes)
			case "H", K_CTRL_LEFT:
				cursor = (cursor + len(nodes) - h) % len(nodes)
				doSelect()
			case "L", K_CTRL_RIGHT:
				doSelect()
				fallthrough
			case "l", K_CTRL_F, K_RIGHT:
				cursor = (cursor + h) % len(nodes)
			case " ", "J", K_CTRL_DOWN:
				doSelect()
				fallthrough
			case "j", K_CTRL_N, K_DOWN:
				cursor = (cursor + 1) % len(nodes)
			case "k", K_CTRL_P, K_UP:
				cursor = (cursor + len(nodes) - 1) % len(nodes)
			case "\b", "K", K_CTRL_UP:
				cursor = (cursor + len(nodes) - 1) % len(nodes)
				doSelect()
			case "\r", "\n":
				var result []int
				if len(selected) > 0 {
					result = make([]int, 0, len(selected))
					for index := range selected {
						result = append(result, index)
					}
					sort.Ints(result)
				} else {
					result = []int{cursor}
				}
				return result
			case "\x1B", K_CTRL_G:
				return []int{}
			}

			// x := cursor / h
			y := cursor % h
			if y < offset {
				offset = y
				// offset--
			} else if y >= offset+b.Height {
				offset = y - b.Height + 1
				// offset++
			}
		}
		if h < b.Height {
			if h > 1 {
				fmt.Fprintf(out, UP_N, h-1)
			}
		} else {
			if b.Height > 1 {
				fmt.Fprintf(out, UP_N, b.Height-1)
			}
		}
		fmt.Fprint(out, "\r")
	}
}

func Choose(sources []string, out io.Writer) int {
	selected := ChooseMulti(sources, out)
	if selected == nil || len(selected) <= 0 {
		return -1
	}
	return selected[0]
}
<commit_msg>Support WindowsTerminal<commit_after>package box

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"io"
	"os"
	"regexp"
	"sort"
	"strings"

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

var wtRuneWidth *runewidth.Condition

func init(){
	wtRuneWidth = runewidth.NewCondition()
	if os.Getenv("WT_SESSION") != "" && os.Getenv("WT_PROFILE_ID") != "" {
		wtRuneWidth.EastAsianWidth = false
	}
}

var AnsiCutter = regexp.MustCompile("\x1B[^a-zA-Z]*[A-Za-z]")

func Print(ctx context.Context, nodes []string, out io.Writer) bool {
	b := New()
	b.Height = 0
	value, _, _ := b.Print(ctx, nodes, 0, out)
	return value
}

func (b *box_t) Print(ctx context.Context,
	nodes []string,
	offset int,
	out io.Writer) (bool, int, int) {

	selected, columns, nlines := b.PrintNoLastLineFeed(ctx, nodes, offset, out)

	// append last linefeed.
	if nlines > 0 {
		fmt.Fprintln(out)
	}
	b.Cache = nil
	return selected, columns, nlines
}

func (b *box_t) PrintNoLastLineFeed(ctx context.Context,
	nodes []string,
	offset int,
	out io.Writer) (bool, int, int) {

	if len(nodes) <= 0 {
		return true, 0, 0
	}

	maxLen := 1
	for _, finfo := range nodes {
		length := wtRuneWidth.StringWidth(AnsiCutter.ReplaceAllString(finfo, ""))
		if length > maxLen {
			maxLen = length
		}
	}
	nodePerLine := (b.Width - 1) / (maxLen + 1)
	if nodePerLine <= 0 {
		nodePerLine = 1
	}
	nlines := (len(nodes) + nodePerLine - 1) / nodePerLine

	lines := make([][]byte, nlines)
	row := 0
	for _, finfo := range nodes {
		lines[row] = append(lines[row], finfo...)
		w := wtRuneWidth.StringWidth(AnsiCutter.ReplaceAllString(finfo, ""))
		if maxLen < b.Width {
			for i := maxLen + 1; i > w; i-- {
				lines[row] = append(lines[row], ' ')
			}
		}
		row++
		if row >= nlines {
			row = 0
		}
	}
	i_end := len(lines)
	if b.Height > 0 {
		if i_end >= offset+b.Height {
			i_end = offset + b.Height
		}
	}

	if b.Cache == nil {
		b.Cache = make([][]byte, b.Height)
	}
	i := offset
	y := 0
	for {
		if y >= len(b.Cache) {
			b.Cache = append(b.Cache, []byte{})
		}
		// assertion
		if i >= len(lines) {
			panic(fmt.Sprintf("len(lines)==%d i==%d", len(lines), i))
		}
		if bytes.Compare(lines[i], b.Cache[y]) != 0 {
			fmt.Fprint(out, strings.TrimRight(string(lines[i]), " "))
			if len(b.Cache[y]) > 0 {
				fmt.Fprint(out, ERASE_LINE)
			}
			b.Cache[y] = lines[i]
		}
		y++
		if ctx != nil {
			select {
			case <-ctx.Done():
				return false, nodePerLine, nlines
			default:
			}
		}
		i++
		if i >= i_end {
			break
		}
		fmt.Fprintln(out)
	}
	return true, nodePerLine, nlines
}

const (
	CURSOR_OFF = "\x1B[?25l"
	CURSOR_ON  = "\x1B[?25h"
	BOLD_ON    = "\x1B[0;47;30m"
	BOLD_ON2   = "\x1B[0;1;7m"
	BOLD_OFF   = "\x1B[0m"
	UP_N       = "\x1B[%dA"
	ERASE_LINE = "\x1B[0K"
)

func truncate(s string, w int) string {
	return wtRuneWidth.Truncate(strings.TrimSpace(s), w, "")
}

const (
	NONE         = 0
	LEFT         = 1
	DOWN         = 2
	UP           = 3
	RIGHT        = 4
	ENTER        = 5
	LEAVE        = 6
	SELECT_DOWN  = 7
	SELECT_UP    = 8
	SELECT_LEFT  = 9
	SELECT_RIGHT = 10
)

type nodeT struct {
	Index int
	Text  string
}

// Choice returns selected string
func Choice(sources []string, out io.Writer) string {
	n := Choose(sources, out)
	if n < 0 {
		return ""
	}
	return sources[n]
}

func ChoiceMulti(sources []string, out io.Writer) []string {
	list := ChooseMulti(sources, out)
	values := make([]string, 0, len(list))
	for _, index := range list {
		values = append(values, sources[index])
	}
	return values
}

// Choice returns the index of selected string
func ChooseMulti(sources []string, out io.Writer) []int {
	cursor := 0
	selected := make(map[int]struct{})

	nodes := make([]*nodeT, 0, len(sources))
	draws := make([]string, 0, len(sources))
	b := New()
	defer b.Close()
	for i, text := range sources {
		val := truncate(text, b.Width-1)
		if val != "" {
			nodes = append(nodes, &nodeT{Index: i, Text: val})
			draws = append(draws, val)
		}
	}
	io.WriteString(out, CURSOR_OFF)
	defer io.WriteString(out, CURSOR_ON)

	if len(nodes) <= 0 {
		nodes = []*nodeT{&nodeT{-1, ""}}
		draws = []string{""}
	}

	offset := 0
	for {
		for index := range selected {
			draws[index] = BOLD_ON + truncate(nodes[index].Text, b.Width-2) + BOLD_OFF
		}
		draws[cursor] = BOLD_ON2 + truncate(nodes[cursor].Text, b.Width-2) + BOLD_OFF
		status, _, h := b.PrintNoLastLineFeed(nil, draws, offset, out)
		if !status {
			return []int{}
		}
		for index, _ := range selected {
			draws[index] = truncate(nodes[index].Text, b.Width-2)
		}
		draws[cursor] = truncate(nodes[cursor].Text, b.Width-2)
		last := cursor

		doSelect := func() {
			if _, ok := selected[cursor]; ok {
				delete(selected, cursor)
			} else {
				selected[cursor] = struct{}{}
			}
		}

		for last == cursor {
			if bw, ok := out.(*bufio.Writer); ok {
				bw.Flush()
			}
			key, err := b.GetKey()
			if err != nil {
				continue
			}
			switch key {
			case "h", K_CTRL_B, K_LEFT:
				cursor = (cursor + len(nodes) - h) % len(nodes)
			case "H", K_CTRL_LEFT:
				cursor = (cursor + len(nodes) - h) % len(nodes)
				doSelect()
			case "L", K_CTRL_RIGHT:
				doSelect()
				fallthrough
			case "l", K_CTRL_F, K_RIGHT:
				cursor = (cursor + h) % len(nodes)
			case " ", "J", K_CTRL_DOWN:
				doSelect()
				fallthrough
			case "j", K_CTRL_N, K_DOWN:
				cursor = (cursor + 1) % len(nodes)
			case "k", K_CTRL_P, K_UP:
				cursor = (cursor + len(nodes) - 1) % len(nodes)
			case "\b", "K", K_CTRL_UP:
				cursor = (cursor + len(nodes) - 1) % len(nodes)
				doSelect()
			case "\r", "\n":
				var result []int
				if len(selected) > 0 {
					result = make([]int, 0, len(selected))
					for index := range selected {
						result = append(result, index)
					}
					sort.Ints(result)
				} else {
					result = []int{cursor}
				}
				return result
			case "\x1B", K_CTRL_G:
				return []int{}
			}

			// x := cursor / h
			y := cursor % h
			if y < offset {
				offset = y
				// offset--
			} else if y >= offset+b.Height {
				offset = y - b.Height + 1
				// offset++
			}
		}
		if h < b.Height {
			if h > 1 {
				fmt.Fprintf(out, UP_N, h-1)
			}
		} else {
			if b.Height > 1 {
				fmt.Fprintf(out, UP_N, b.Height-1)
			}
		}
		fmt.Fprint(out, "\r")
	}
}

func Choose(sources []string, out io.Writer) int {
	selected := ChooseMulti(sources, out)
	if selected == nil || len(selected) <= 0 {
		return -1
	}
	return selected[0]
}
<|endoftext|>


<commit_before>package sol

import (
	"encoding/json"
	"reflect"
	"sort"
)

// Values is a map of column names to parameters.
type Values map[string]interface{}

// Diff returns the values in v that differ from the values in other.
// ISO 31-11: v \ other
func (v Values) Diff(other Values) Values {
	diff := Values{}
	for key, value := range v {
		if !ObjectsAreEqual(value, other[key]) {
			diff[key] = value
		}
	}
	return diff
}

// Exclude removes the given keys and returns the remaining Values
func (v Values) Exclude(keys ...string) Values {
	safe := Values{}
ValueLoop:
	for key, value := range v {
		for _, k := range keys {
			if k == key {
				continue ValueLoop
			}
		}
		safe[key] = value
	}
	return safe
}

// Keys returns the keys of the Values map in alphabetical order.
func (v Values) Keys() []string {
	keys := make([]string, len(v))
	var i int
	for key := range v {
		keys[i] = key
		i += 1
	}
	sort.Strings(keys)
	return keys
}

// MarshalJSON converts Values to JSON after converting all byte slices to
// a string type.
// By default, byte slices are JSON unmarshaled as base64.
// This is an issue since the postgres driver will scan string/varchar
// types as byte slices. Since Values{} should rarely be used within
// Go code, we're only modifying the JSON marshaler.
func (v Values) MarshalJSON() ([]byte, error) {
	for key, value := range v {
		if val, ok := value.([]byte); ok {
			v[key] = string(val)
		}
	}

	// Convert to prevent recursive marshaling
	return json.Marshal(map[string]interface{}(v))
}

// Values converts the given object to a Values{} type
func ValuesOf(obj interface{}) Values {
	values := Values{}

	elem := reflect.Indirect(reflect.ValueOf(obj))
	switch elem.Kind() {
	case reflect.Struct:
		fields := SelectFieldsFromElem(elem.Type())
		// TODO how to convert to db column name? Show Values even care?
		for _, field := range fields {
			var fieldElem reflect.Value = elem
			for _, name := range field.names {
				fieldElem = fieldElem.FieldByName(name)
			}
			// TODO Skip empty if omit empty...?
			values[field.column] = fieldElem.Interface()
		}
	case reflect.Map:
		// TODO Convert to Values - generalized map iteration?
	default:
		// TODO Return an error, panic, or silent?
	}
	return values
}
<commit_msg>Values.Merge method for merging multiple Values<commit_after>package sol

import (
	"encoding/json"
	"reflect"
	"sort"
)

// Values is a map of column names to parameters.
type Values map[string]interface{}

// Diff returns the values in v that differ from the values in other.
// ISO 31-11: v \ other
func (v Values) Diff(other Values) Values {
	diff := Values{}
	for key, value := range v {
		if !ObjectsAreEqual(value, other[key]) {
			diff[key] = value
		}
	}
	return diff
}

// Exclude removes the given keys and returns the remaining Values
func (v Values) Exclude(keys ...string) Values {
	safe := Values{}
ValueLoop:
	for key, value := range v {
		for _, k := range keys {
			if k == key {
				continue ValueLoop
			}
		}
		safe[key] = value
	}
	return safe
}

// Keys returns the keys of the Values map in alphabetical order.
func (v Values) Keys() []string {
	keys := make([]string, len(v))
	var i int
	for key := range v {
		keys[i] = key
		i += 1
	}
	sort.Strings(keys)
	return keys
}

// Merge combines the Values{} types, with precedence to the Values given as
// a parameter
func (v Values) Merge(other Values) Values {
	merged := Values{}
	for key, value := range v {
		merged[key] = value
	}
	for key, value := range other {
		merged[key] = value
	}
	return merged
}

// MarshalJSON converts Values to JSON after converting all byte slices to
// a string type.
// By default, byte slices are JSON unmarshaled as base64.
// This is an issue since the postgres driver will scan string/varchar
// types as byte slices. Since Values{} should rarely be used within
// Go code, we're only modifying the JSON marshaler.
func (v Values) MarshalJSON() ([]byte, error) {
	for key, value := range v {
		if val, ok := value.([]byte); ok {
			v[key] = string(val)
		}
	}

	// Convert to prevent recursive marshaling
	return json.Marshal(map[string]interface{}(v))
}

// Values converts the given object to a Values{} type
func ValuesOf(obj interface{}) Values {
	values := Values{}

	elem := reflect.Indirect(reflect.ValueOf(obj))
	switch elem.Kind() {
	case reflect.Struct:
		fields := SelectFieldsFromElem(elem.Type())
		// TODO how to convert to db column name? Show Values even care?
		for _, field := range fields {
			var fieldElem reflect.Value = elem
			for _, name := range field.names {
				fieldElem = fieldElem.FieldByName(name)
			}
			// TODO Skip empty if omit empty...?
			values[field.column] = fieldElem.Interface()
		}
	case reflect.Map:
		// TODO Convert to Values - generalized map iteration?
	default:
		// TODO Return an error, panic, or silent?
	}
	return values
}
<|endoftext|>


<commit_before>package mathgl

import "errors"

type Vec2f [2]float32
type Vec2d [2]float64
type Vec2i [2]int32
type Vec2l [2]int64

type Vec3f [3]float32
type Vec3d [3]float64
type Vec3i [3]int32
type Vec3l [3]int64

type Vecf  []float32
type Vecd []float64
type Veci []int32
type Vecl []int64

func (v Vec2f) Dot2f(u Vec2f) float32 {
	return v[0] * u[0] + v[1] * u[1]
}

func (v Vec2d) Dot2d(u Vec2d) float64 {
	return v[0] * u[0] + v[1] * u[1]
}

func (v Vec2i) Dot2i(u Vec2i) int32 {
	return v[0] * u[0] + v[1] * u[1]
}

func (v Vec2l) Dot2l(u Vec2l) int64 {
	return v[0] * u[0] + v[1] * u[1]
}

func (v Vec3f) Dot3f(u Vec3f) float32 {
	return v[0] * u[0] + v[1] * u[1] + v[2] * u[2]
}

func (v Vec3d) Dot3d(u Vec3d) float64 {
	return v[0] * u[0] + v[1] * u[1] + v[2] * u[2]
}

func (v Vec3i) Dot3i(u Vec3i) int32 {
	return v[0] * u[0] + v[1] * u[1] + v[2] * u[2]
}

func (v Vec3l) Dot3l(u Vec3l) int64 {
	return v[0] * u[0] + v[1] * u[1] + v[2] * u[2]
}

func (v Vecf) Dotf(u Vecf) (f float32, err error) {
	if len(v) != len(u) {
		return 0.0 ,errors.New("Vectors aren't the same length")
	}

	f = 0.0
	for i := range v {
		f = f + u[i]*v[i]
	}

	return f,nil
}

func (v Vecd) Dotd(u Vecd) (f float64, err error) {
	if len(v) != len(u) {
		return 0.0 ,errors.New("Vectors aren't the same length")
	}

	f = 0.0
	for i := range v {
		f = f + u[i]*v[i]
	}

	return f,nil
}

func (v Veci) Doti(u Veci) (f int32, err error) {
	if len(v) != len(u) {
		return 0.0 ,errors.New("Vectors aren't the same length")
	}

	f = 0.0
	for i := range v {
		f = f + u[i]*v[i]
	}

	return f,nil
}

func (v Vecl) Dotl(u Vecl) (f int64, err error) {
	if len(v) != len(u) {
		return 0.0 ,errors.New("Vectors aren't the same length")
	}

	f = 0.0
	for i := range v {
		f = f + u[i]*v[i]
	}

	return f,nil
}

func (v Vec3f) Cross3f(u Vec3f) (r Vec3f) {
	r[0],r[1],r[2] = v[1]*u[2]-v[2]*u[1], v[2]*u[0]-v[0]*u[2], v[0]*u[1]-v[1]*u[0]

	return r
}

func (v Vec3d) Cross3d(u Vec3d) (r Vec3d) {
	r[0],r[1],r[2] = v[1]*u[2]-v[2]*u[1], v[2]*u[0]-v[0]*u[2], v[0]*u[1]-v[1]*u[0]

	return r
}

func (v Vec3i) Cross3i(u Vec3i) (r Vec3i) {
	r[0],r[1],r[2] = v[1]*u[2]-v[2]*u[1], v[2]*u[0]-v[0]*u[2], v[0]*u[1]-v[1]*u[0]

	return r
}

func (v Vec3l) Cross3l(u Vec3l) (r Vec3l) {
	r[0],r[1],r[2] = v[1]*u[2]-v[2]*u[1], v[2]*u[0]-v[0]*u[2], v[0]*u[1]-v[1]*u[0]

	return r
}

// Allow 7-dimensional cross product? (It's also binary).
func (v Vecf) Crossf(u Vecf) (r Vecf, err error) {
	if len(v) != 3 || len(u) != 3 {
		return nil, errors.New("Cross product requires both inputs to be 3-dimensional")
	}

	r = make([]float32,3,3)
	r[0],r[1],r[2] = v[1]*u[2]-v[2]*u[1], v[2]*u[0]-v[0]*u[2], v[0]*u[1]-v[1]*u[0]
	return r, nil
}

func (v Vecd) Crossd(u Vecd) (r Vecd, err error) {
	if len(v) != 3 || len(u) != 3 {
		return nil, errors.New("Cross product requires both inputs to be 3-dimensional")
	}

	r = make([]float64,3,3)
	r[0],r[1],r[2] = v[1]*u[2]-v[2]*u[1], v[2]*u[0]-v[0]*u[2], v[0]*u[1]-v[1]*u[0]
	return r, nil
}

func (v Veci) Crossf(u Veci) (r Veci, err error) {
	if len(v) != 3 || len(u) != 3 {
		return nil, errors.New("Cross product requires both inputs to be 3-dimensional")
	}

	r = make([]int32,3,3)
	r[0],r[1],r[2] = v[1]*u[2]-v[2]*u[1], v[2]*u[0]-v[0]*u[2], v[0]*u[1]-v[1]*u[0]
	return r, nil
}

func (v Vecl) Crossf(u Vecl) (r Vecl, err error) {
	if len(v) != 3 || len(u) != 3 {
		return nil, errors.New("Cross product requires both inputs to be 3-dimensional")
	}

	r = make([]int64,3,3)
	r[0],r[1],r[2] = v[1]*u[2]-v[2]*u[1], v[2]*u[0]-v[0]*u[2], v[0]*u[1]-v[1]*u[0]
	return r, nil
}<commit_msg>Really ugly, redundant code<commit_after><|endoftext|>


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

// A command that reads all blobs necessary for verifying the directory
// structure rooted at a set of backup root scores, ensuring that the entire
// directory structure is intact in GCS.
//
// Optionally, all file content is also read and verified. This is less
// important than verifying directory connectedness if we trust that GCS does
// not corrupt object metadata (where we store expected CRC32C and MD5) and
// does correctly report the object's CRC32C and MD5 sums in listings,
// verifying them periodically.
//
// Output is of the following form:
//
//     <timestamp> <node> [<child node> ...]
//
// where:
//
//  *  Timestamps are formatted according to time.RFC3339.
//
//  *  Node names have one of two forms:
//
//     *   Nodes of the form "d:<hex score>" represent the directory listing
//         contained within the blob of the given score.
//
//     *   Nodes of the form "f:<hex score>" represent a piece of a file,
//         contained within the blob of the given score.
//
// An output line for a directory node means that at the given timestamp we
// certified that a piece of content with the given score was parseable as a
// directory listing that referred to the given scores for its direct children.
//
// An output line for a file node means that at the given timestamp we
// certified that a piece of content with the given score was parseable as a
// piece of a file. File nodes never have children.

package main

import (
	"bytes"
	"fmt"
	"log"
	"runtime"
	"strings"
	"time"

	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/graph"
	"github.com/jacobsa/comeback/util"
	"github.com/jacobsa/comeback/verify"
	"github.com/jacobsa/comeback/wiring"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/syncutil"
	"golang.org/x/net/context"
)

var cmdVerify = &Command{
	Name: "verify",
}

// TODO(jacobsa): Get these automatically from the registry.
var fRoots = cmdVerify.Flags.String(
	"roots",
	"",
	"Comma-separated list of backup root scores to verify.")

var fFast = cmdVerify.Flags.Bool(
	"fast",
	false,
	"When set, don't verify file content.")

func init() {
	cmdVerify.Run = runVerify // Break flag-related dependency loop.
}

////////////////////////////////////////////////////////////////////////
// Visitor types
////////////////////////////////////////////////////////////////////////

// A record certifying something we confirmed about a node at a certain time.
// See the notes at the top of verify.go for details.
type verifyRecord struct {
	t        time.Time
	node     string
	adjacent []string
}

// A visitor that writes the information it gleans from the wrapped visitor to
// a channel.
type snoopingVisitor struct {
	records chan<- verifyRecord
	wrapped graph.Visitor
}

func (v *snoopingVisitor) Visit(
	ctx context.Context,
	node string) (adjacent []string, err error) {
	// Call through.
	adjacent, err = v.wrapped.Visit(ctx, node)
	if err != nil {
		return
	}

	// Write out a record.
	r := verifyRecord{
		t:        time.Now(),
		node:     node,
		adjacent: adjacent,
	}

	select {
	case <-ctx.Done():
		err = ctx.Err()
		return

	case v.records <- r:
	}

	return
}

////////////////////////////////////////////////////////////////////////
// Output
////////////////////////////////////////////////////////////////////////

// Print output based on the visitor results arriving on the supplied channel.
func formatVerifyOutput(r verifyRecord) (s string) {
	var extra string
	if len(r.adjacent) != 0 {
		extra = fmt.Sprintf(" %s", strings.Join(r.adjacent, " "))
	}

	s = fmt.Sprintf(
		"%s %s%s",
		r.t.Format(time.RFC3339),
		r.node,
		extra)

	return
}

// Parse the supplied line (without line break) previously output by the verify
// command.
func parseVerifyRecord(line []byte) (r verifyRecord, err error) {
	// We expect space-separate components.
	components := bytes.Split(line, []byte{' '})
	if len(components) < 2 {
		err = fmt.Errorf(
			"Expected at least two components, got %d.",
			len(components))

		return
	}

	// The first should be the timestmap.
	r.t, err = time.Parse(time.RFC3339, string(components[0]))
	if err != nil {
		err = fmt.Errorf("time.Parse(%q): %v", components[0], err)
		return
	}

	// The next should be the node name.
	r.node = string(components[1])

	// The rest should be adjacent node names.
	for i := 2; i < len(components); i++ {
		r.adjacent = append(r.adjacent, string(components[i]))
	}

	// Make sure all of the node names are legal.
	allNodes := make([]string, 1+len(r.adjacent))
	allNodes[0] = r.node
	copy(allNodes[1:], r.adjacent)

	for _, n := range allNodes {
		_, _, err = verify.ParseNodeName(n)
		if err != nil {
			err = fmt.Errorf("ParseNodeName(%q): %v", n, err)
			return
		}
	}

	return
}

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

// List blob.ListScores, but returns a slice instead of writing into a channel.
func listAllScores(
	ctx context.Context,
	bucket gcs.Bucket,
	namePrefix string) (scores []blob.Score, err error) {
	b := syncutil.NewBundle(context.Background())
	defer func() { err = b.Join() }()

	// List scores into a channel.
	scoreChan := make(chan blob.Score, 100)
	b.Add(func(ctx context.Context) (err error) {
		defer close(scoreChan)
		err = blob.ListScores(ctx, bucket, wiring.BlobObjectNamePrefix, scoreChan)
		if err != nil {
			err = fmt.Errorf("ListScores: %v", err)
			return
		}

		return
	})

	// Accumulate into the slice.
	b.Add(func(ctx context.Context) (err error) {
		for score := range scoreChan {
			scores = append(scores, score)
		}

		return
	})

	return
}

////////////////////////////////////////////////////////////////////////
// Verify
////////////////////////////////////////////////////////////////////////

// Run the verification pipeline. Return a count of the number of scores
// verified and the number skipped due to readFiles being false.
func verifyImpl(
	ctx context.Context,
	readFiles bool,
	rootScores []blob.Score,
	knownScores []blob.Score,
	blobStore blob.Store) (nodesVerified uint64, nodesSkipped uint64, err error) {
	b := syncutil.NewBundle(ctx)

	// Visit every node in the graph, snooping on the graph structure into a
	// channel.
	visitorRecords := make(chan verifyRecord, 100)
	b.Add(func(ctx context.Context) (err error) {
		defer close(visitorRecords)

		visitor := verify.NewVisitor(
			readFiles,
			knownScores,
			blobStore)

		visitor = &snoopingVisitor{
			wrapped: visitor,
			records: visitorRecords,
		}

		// Format root node names.
		var roots []string
		for _, score := range rootScores {
			roots = append(roots, verify.FormatNodeName(true, score))
		}

		// Traverse starting at the specified roots. Use an "experimentally
		// determined" parallelism, which in theory should depend on bandwidth-delay
		// products but in practice comes down to when the OS gets cranky about open
		// files.
		const parallelism = 128

		err = graph.Traverse(
			context.Background(),
			parallelism,
			roots,
			visitor)

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

		return
	})

	// Count and output the nodes visited, filtering out file nodes if we're not
	// actually reading and verifying them.
	b.Add(func(ctx context.Context) (err error) {
		for r := range visitorRecords {
			var dir bool
			dir, _, err = verify.ParseNodeName(r.node)
			if err != nil {
				err = fmt.Errorf("ParseNodeName(%q): %v", r.node, err)
				return
			}

			// Skip files if appropriate.
			if !readFiles && !dir {
				nodesSkipped++
				continue
			}

			// Increment the count and output the information.
			nodesVerified++
			fmt.Println(formatVerifyOutput(r))
		}

		return
	})

	err = b.Join()
	return
}

func runVerify(args []string) {
	// Allow parallelism.
	runtime.GOMAXPROCS(runtime.NumCPU())

	// Die on error.
	var err error
	defer func() {
		if err != nil {
			log.Fatalln(err)
		}
	}()

	// Read flags.
	readFiles := !*fFast

	if *fRoots == "" {
		err = fmt.Errorf("You must set --roots.")
		return
	}

	rootHexScores := strings.Split(*fRoots, ",")
	var rootScores []blob.Score
	for _, hexScore := range rootHexScores {
		var score blob.Score
		score, err = blob.ParseHexScore(hexScore)
		if err != nil {
			err = fmt.Errorf("Invalid root %q: %v", hexScore, err)
			return
		}

		rootScores = append(rootScores, score)
	}

	// Grab dependencies.
	bucket := getBucket()
	crypter := getCrypter()

	// Create a blob store.
	blobStore, err := wiring.MakeBlobStore(
		bucket,
		crypter,
		util.NewStringSet())

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

	// List all scores in the bucket, verifying the object record metadata in the
	// process.
	log.Println("Listing scores...")
	knownScores, err := listAllScores(
		context.Background(),
		bucket,
		wiring.BlobObjectNamePrefix)

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

	log.Printf("Listed %d scores.", len(knownScores))

	// Run the rest of the pipeline.
	nodesVerified, nodesSkipped, err := verifyImpl(
		context.Background(),
		readFiles,
		rootScores,
		knownScores,
		blobStore)

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

	log.Printf(
		"Successfully verified %d nodes (%d skipped due to fast mode).",
		nodesVerified,
		nodesSkipped)

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

// A command that reads all blobs necessary for verifying the directory
// structure rooted at a set of backup root scores, ensuring that the entire
// directory structure is intact in GCS.
//
// Optionally, all file content is also read and verified. This is less
// important than verifying directory connectedness if we trust that GCS does
// not corrupt object metadata (where we store expected CRC32C and MD5) and
// does correctly report the object's CRC32C and MD5 sums in listings,
// verifying them periodically.
//
// Output is of the following form:
//
//     <timestamp> <node> [<child node> ...]
//
// where:
//
//  *  Timestamps are formatted according to time.RFC3339.
//
//  *  Node names have one of two forms:
//
//     *   Nodes of the form "d:<hex score>" represent the directory listing
//         contained within the blob of the given score.
//
//     *   Nodes of the form "f:<hex score>" represent a piece of a file,
//         contained within the blob of the given score.
//
// An output line for a directory node means that at the given timestamp we
// certified that a piece of content with the given score was parseable as a
// directory listing that referred to the given scores for its direct children.
//
// An output line for a file node means that at the given timestamp we
// certified that a piece of content with the given score was parseable as a
// piece of a file. File nodes never have children.

package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"log"
	"runtime"
	"strings"
	"time"

	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/graph"
	"github.com/jacobsa/comeback/util"
	"github.com/jacobsa/comeback/verify"
	"github.com/jacobsa/comeback/wiring"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/syncutil"
	"golang.org/x/net/context"
)

var cmdVerify = &Command{
	Name: "verify",
}

// TODO(jacobsa): Get these automatically from the registry.
var fRoots = cmdVerify.Flags.String(
	"roots",
	"",
	"Comma-separated list of backup root scores to verify.")

var fFast = cmdVerify.Flags.Bool(
	"fast",
	false,
	"When set, don't verify file content.")

func init() {
	cmdVerify.Run = runVerify // Break flag-related dependency loop.
}

////////////////////////////////////////////////////////////////////////
// Visitor types
////////////////////////////////////////////////////////////////////////

// A record certifying something we confirmed about a node at a certain time.
// See the notes at the top of verify.go for details.
type verifyRecord struct {
	t        time.Time
	node     string
	adjacent []string
}

// A visitor that writes the information it gleans from the wrapped visitor to
// a channel.
type snoopingVisitor struct {
	records chan<- verifyRecord
	wrapped graph.Visitor
}

func (v *snoopingVisitor) Visit(
	ctx context.Context,
	node string) (adjacent []string, err error) {
	// Call through.
	adjacent, err = v.wrapped.Visit(ctx, node)
	if err != nil {
		return
	}

	// Write out a record.
	r := verifyRecord{
		t:        time.Now(),
		node:     node,
		adjacent: adjacent,
	}

	select {
	case <-ctx.Done():
		err = ctx.Err()
		return

	case v.records <- r:
	}

	return
}

////////////////////////////////////////////////////////////////////////
// Output
////////////////////////////////////////////////////////////////////////

// Print output based on the visitor results arriving on the supplied channel.
func formatVerifyOutput(r verifyRecord) (s string) {
	var extra string
	if len(r.adjacent) != 0 {
		extra = fmt.Sprintf(" %s", strings.Join(r.adjacent, " "))
	}

	s = fmt.Sprintf(
		"%s %s%s",
		r.t.Format(time.RFC3339),
		r.node,
		extra)

	return
}

// Parse the supplied line (without line break) previously output by the verify
// command.
func parseVerifyRecord(line []byte) (r verifyRecord, err error) {
	// We expect space-separate components.
	components := bytes.Split(line, []byte{' '})
	if len(components) < 2 {
		err = fmt.Errorf(
			"Expected at least two components, got %d.",
			len(components))

		return
	}

	// The first should be the timestmap.
	r.t, err = time.Parse(time.RFC3339, string(components[0]))
	if err != nil {
		err = fmt.Errorf("time.Parse(%q): %v", components[0], err)
		return
	}

	// The next should be the node name.
	r.node = string(components[1])

	// The rest should be adjacent node names.
	for i := 2; i < len(components); i++ {
		r.adjacent = append(r.adjacent, string(components[i]))
	}

	// Make sure all of the node names are legal.
	allNodes := make([]string, 1+len(r.adjacent))
	allNodes[0] = r.node
	copy(allNodes[1:], r.adjacent)

	for _, n := range allNodes {
		_, _, err = verify.ParseNodeName(n)
		if err != nil {
			err = fmt.Errorf("ParseNodeName(%q): %v", n, err)
			return
		}
	}

	return
}

// Parse the supplied output from a run of the verify command, writing records
// into the supplied channel.
func parseVerifyOutput(
	ctx context.Context,
	in io.Reader,
	records chan<- verifyRecord) (err error) {
	reader := bufio.NewReader(in)

	for {
		// Find the next line. EOF with no data means we are done; otherwise ignore
		// EOF in case the file doesn't end with a newline.
		var line []byte
		line, err = reader.ReadBytes('\n')
		if err == io.EOF {
			err = nil
			if len(line) == 0 {
				break
			}
		}

		// Propagate other errors.
		if err != nil {
			err = fmt.Errorf("ReadBytes: %v", err)
			return
		}

		// Trim the delimiter.
		line = line[:len(line)-1]

		// Parse the line.
		var r verifyRecord
		r, err = parseVerifyRecord(line)
		if err != nil {
			err = fmt.Errorf("parseVerifyRecord(%q): %v", line, err)
			return
		}

		// Write out the record.
		select {
		case <-ctx.Done():
			err = ctx.Err()
			return

		case records <- r:
		}
	}

	return
}

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

// List blob.ListScores, but returns a slice instead of writing into a channel.
func listAllScores(
	ctx context.Context,
	bucket gcs.Bucket,
	namePrefix string) (scores []blob.Score, err error) {
	b := syncutil.NewBundle(context.Background())
	defer func() { err = b.Join() }()

	// List scores into a channel.
	scoreChan := make(chan blob.Score, 100)
	b.Add(func(ctx context.Context) (err error) {
		defer close(scoreChan)
		err = blob.ListScores(ctx, bucket, wiring.BlobObjectNamePrefix, scoreChan)
		if err != nil {
			err = fmt.Errorf("ListScores: %v", err)
			return
		}

		return
	})

	// Accumulate into the slice.
	b.Add(func(ctx context.Context) (err error) {
		for score := range scoreChan {
			scores = append(scores, score)
		}

		return
	})

	return
}

////////////////////////////////////////////////////////////////////////
// Verify
////////////////////////////////////////////////////////////////////////

// Run the verification pipeline. Return a count of the number of scores
// verified and the number skipped due to readFiles being false.
func verifyImpl(
	ctx context.Context,
	readFiles bool,
	rootScores []blob.Score,
	knownScores []blob.Score,
	blobStore blob.Store) (nodesVerified uint64, nodesSkipped uint64, err error) {
	b := syncutil.NewBundle(ctx)

	// Visit every node in the graph, snooping on the graph structure into a
	// channel.
	visitorRecords := make(chan verifyRecord, 100)
	b.Add(func(ctx context.Context) (err error) {
		defer close(visitorRecords)

		visitor := verify.NewVisitor(
			readFiles,
			knownScores,
			blobStore)

		visitor = &snoopingVisitor{
			wrapped: visitor,
			records: visitorRecords,
		}

		// Format root node names.
		var roots []string
		for _, score := range rootScores {
			roots = append(roots, verify.FormatNodeName(true, score))
		}

		// Traverse starting at the specified roots. Use an "experimentally
		// determined" parallelism, which in theory should depend on bandwidth-delay
		// products but in practice comes down to when the OS gets cranky about open
		// files.
		const parallelism = 128

		err = graph.Traverse(
			context.Background(),
			parallelism,
			roots,
			visitor)

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

		return
	})

	// Count and output the nodes visited, filtering out file nodes if we're not
	// actually reading and verifying them.
	b.Add(func(ctx context.Context) (err error) {
		for r := range visitorRecords {
			var dir bool
			dir, _, err = verify.ParseNodeName(r.node)
			if err != nil {
				err = fmt.Errorf("ParseNodeName(%q): %v", r.node, err)
				return
			}

			// Skip files if appropriate.
			if !readFiles && !dir {
				nodesSkipped++
				continue
			}

			// Increment the count and output the information.
			nodesVerified++
			fmt.Println(formatVerifyOutput(r))
		}

		return
	})

	err = b.Join()
	return
}

func runVerify(args []string) {
	// Allow parallelism.
	runtime.GOMAXPROCS(runtime.NumCPU())

	// Die on error.
	var err error
	defer func() {
		if err != nil {
			log.Fatalln(err)
		}
	}()

	// Read flags.
	readFiles := !*fFast

	if *fRoots == "" {
		err = fmt.Errorf("You must set --roots.")
		return
	}

	rootHexScores := strings.Split(*fRoots, ",")
	var rootScores []blob.Score
	for _, hexScore := range rootHexScores {
		var score blob.Score
		score, err = blob.ParseHexScore(hexScore)
		if err != nil {
			err = fmt.Errorf("Invalid root %q: %v", hexScore, err)
			return
		}

		rootScores = append(rootScores, score)
	}

	// Grab dependencies.
	bucket := getBucket()
	crypter := getCrypter()

	// Create a blob store.
	blobStore, err := wiring.MakeBlobStore(
		bucket,
		crypter,
		util.NewStringSet())

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

	// List all scores in the bucket, verifying the object record metadata in the
	// process.
	log.Println("Listing scores...")
	knownScores, err := listAllScores(
		context.Background(),
		bucket,
		wiring.BlobObjectNamePrefix)

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

	log.Printf("Listed %d scores.", len(knownScores))

	// Run the rest of the pipeline.
	nodesVerified, nodesSkipped, err := verifyImpl(
		context.Background(),
		readFiles,
		rootScores,
		knownScores,
		blobStore)

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

	log.Printf(
		"Successfully verified %d nodes (%d skipped due to fast mode).",
		nodesVerified,
		nodesSkipped)

	return
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"sync"
	"unicode/utf8"
)

var ErrInvalidText = errors.New("unavailable encoding")

type InternalError struct {
	e error
}

func (ei *InternalError) Error() string {
	return fmt.Sprintf("internal error:%v", ei.e)
}

type Line struct {
	Num uint
	Str string
}

// cap(q) is required always greater than 1
type LineQueue struct{ q chan *Line }

func NewLineQueue(capacity uint) (*LineQueue, error) {
	if capacity == 0 {
		return nil, errors.New("capacity is 0")
	}
	return &LineQueue{make(chan *Line, capacity)}, nil
}
func (lq *LineQueue) Len() int { return len(lq.q) }
func (lq *LineQueue) Cap() int { return cap(lq.q) }
func (lq *LineQueue) Push(l *Line) {
	select {
	case lq.q <- l:
	default:
		<-lq.q
		lq.q <- l
	}
}
func (lq *LineQueue) PopAll() []*Line {
	lines := make([]*Line, 0, cap(lq.q))
	for len(lq.q) != 0 {
		lines = append(lines, <-lq.q)
	}
	return lines
}

// change to struct{ index int, pos int, lines []*Line }?
type Context struct {
	line   *Line
	before []*Line
	after  []*Line
}

func FprintContexts(writer io.Writer, prefix string, cs []*Context) error {
	if cs == nil {
		return nil
	}
	var err error
	f := func(ls []*Line) {
		for _, l := range ls {
			_, err = fmt.Fprintf(writer, "%s%d-%s\n", prefix, l.Num, l.Str)
		}
	}
	for _, c := range cs {
		f(c.before)
		if err != nil {
			return err
		}
		_, err = fmt.Fprintf(writer, "%s%d:%s\n", prefix, c.line.Num, c.line.Str)
		if err != nil {
			return err
		}
		f(c.after)
		if err != nil {
			return err
		}
	}
	return nil
}

type File struct {
	path string
	cs   []*Context
}

func (f *File) Fprint(writer io.Writer) error {
	var err error
	if len(f.cs) == 0 {
		return nil
	}
	_, err = fmt.Fprintln(writer, f.path)
	if err != nil {
		return err
	}
	err = FprintContexts(writer, "", f.cs)
	if err != nil {
		return err
	}
	_, err = fmt.Fprintln(writer)
	if err != nil {
		return err
	}
	return nil
}

func (f *File) FprintVerbose(writer io.Writer) error {
	var err error
	if len(f.cs) == 0 {
		return nil
	}
	err = FprintContexts(writer, f.path, f.cs)
	if err != nil {
		return err
	}
	_, err = fmt.Fprintln(writer)
	if err != nil {
		return err
	}
	return nil
}

func FprintFiles(writer io.Writer, fs ...*File) error {
	var err error
	for i := range fs {
		err = fs[i].Fprint(writer)
		if err != nil {
			return err
		}
	}
	return nil
}

func FprintFilesVerbose(writer io.Writer, fs ...*File) error {
	var err error
	for i := range fs {
		err = fs[i].FprintVerbose(writer)
		if err != nil {
			return err
		}
	}
	return nil
}

type Walker struct {
	fileQueue chan string
	dirQueue  chan []string

	regexp *regexp.Regexp

	nworker       int
	checked       map[string]bool
	mu            sync.Mutex
	wg            sync.WaitGroup
	once          sync.Once
	internalError *InternalError

	log *log.Logger
}

func NewWalker() *Walker {
	nworker := runtime.NumCPU() / 4
	if nworker < 2 {
		nworker = 2
	}
	return &Walker{
		fileQueue:     make(chan string, 128),
		dirQueue:      make(chan []string, nworker),
		nworker:       nworker,
		checked:       make(map[string]bool),
		log:           log.New(ioutil.Discard, Name+":", 0),
		internalError: nil,
	}
}

func (w *Walker) SetLogOutput(writer io.Writer) { w.log.SetOutput(writer) }

func (w *Walker) setInternalError(err error) {
	w.once.Do(func() { w.internalError = &InternalError{e: err} })
}

func (w *Walker) sendQueue(paths ...string) {
	var dirs []string
	for i := range paths {
		abs, err := filepath.Abs(paths[i])
		if err != nil {
			w.setInternalError(err)
			w.log.Printf("[Err]:%v", err)
			continue
		}
		fi, err := os.Stat(abs)
		if err != nil {
			w.setInternalError(err)
			w.log.Printf("[Errr]:%v", err)
			continue
		}
		if fi.IsDir() {
			dirs = append(dirs, abs)
		} else if fi.Mode().IsRegular() {
			w.wg.Add(1)
			w.fileQueue <- abs
		}
	}
	w.wg.Add(1)
	w.dirQueue <- dirs
}

func (w *Walker) Start(pat string, nlines uint, paths ...string) (<-chan *File, func() error, error) {
	var re *regexp.Regexp
	re, err := regexp.Compile(pat)
	if err != nil {
		return nil, nil, err
	}
	w.regexp = re

	resultQueue := make(chan *File, cap(w.fileQueue))
	done := make(chan struct{})
	wait := func() error {
		w.wg.Wait()
		return w.internalError
	}

	for i := 0; i != w.nworker; i++ {
		go w.dirWalker(done)
		go w.fileWalker(done, resultQueue, nlines)
	}
	w.sendQueue(paths...)
	go func() {
		w.wg.Wait()
		close(done)
		close(resultQueue)
	}()
	return resultQueue, wait, nil
}

// for goroutine
// send tasks to {file,dir}Queue
func (w *Walker) dirWalker(done <-chan struct{}) {
	var nextDirs []string
	var dirs []string
	for ; true; w.wg.Done() {
		select {
		case dirs = <-w.dirQueue:
			for i := range dirs {
				fis, err := ioutil.ReadDir(dirs[i])
				if err != nil {
					w.setInternalError(err)
					w.log.Printf("[Err]:%s:%v", dirs[i], err)
					if os.IsNotExist(err) || os.IsPermission(err) {
						continue
					}
					// unexpected error
					panic(err)
				}
				for _, fi := range fis {
					if fi.Mode().IsRegular() {
						w.wg.Add(1)
						w.fileQueue <- filepath.Join(dirs[i], fi.Name())
					} else if fi.IsDir() {
						nextDirs = append(nextDirs, filepath.Join(dirs[i], fi.Name()))
					}
				}
			}
			if nextDirs != nil {
				w.wg.Add(1)
				w.dirQueue <- nextDirs
				nextDirs = nil
			}
		case <-done:
			return
		}
	}
}

// for goroutine
func (w *Walker) fileWalker(done <-chan struct{}, resultQueue chan<- *File, nlines uint) {
	var lq *LineQueue
	if nlines != 0 {
		// not need error check
		lq, _ = NewLineQueue(nlines)
	}
	var file string
	var err error
	var cs []*Context
	for ; true; w.wg.Done() {
		select {
		case file = <-w.fileQueue:
			w.mu.Lock()
			if w.checked[file] {
				w.mu.Unlock()
				continue
			}
			w.checked[file] = true
			w.mu.Unlock()

			cs, err = w.readFile(file, lq)
			if err != nil {
				w.setInternalError(err)
				w.log.Printf("[Err]:%s:%v", file, err)
				if os.IsNotExist(err) || os.IsPermission(err) {
					continue
				}
				if err == bufio.ErrTooLong {
					continue
				}
				if err == ErrInvalidText {
					continue
				}
				// unexpected error
				panic(err)
			}
			w.log.Println(file)
			if len(cs) != 0 {
				resultQueue <- &File{
					path: file,
					cs:   cs,
				}
			}
		case <-done:
			return
		}
	}
}

// TODO? readFile(f *File, file string) error
func (w *Walker) readFile(file string, lq *LineQueue) ([]*Context, error) {
	f, err := os.Open(file)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	var cs []*Context
	var c = new(Context)
	var txt string
	var i uint
	var matched bool

	var csAdd func()
	if lq == nil {
		csAdd = func() {
			if matched {
				cs = append(cs, &Context{
					line:   &Line{i, txt},
					before: []*Line{},
					after:  []*Line{},
				})
			}
		}
	} else {
		defer lq.PopAll()
		csAdd = func() {
			if c.line != nil {
				if matched {
					c.after = lq.PopAll()
					cs = append(cs, c)
					c = &Context{
						before: []*Line{},
						line:   &Line{i, txt},
						after:  []*Line{},
					}
					return
				}
				if lq.Len() == lq.Cap() {
					c.after = lq.PopAll()
					cs = append(cs, c)
					c = new(Context)
				}
			} else if matched {
				c.before = lq.PopAll()
				c.line = &Line{i, txt}
				return
			}
			lq.Push(&Line{i, txt})
		}
	}

	sc := bufio.NewScanner(f)
	for i = uint(1); sc.Scan(); i++ {
		if i == 0 {
			return nil, errors.New("too many lines")
		}
		txt = sc.Text()
		if !utf8.ValidString(txt) {
			return nil, ErrInvalidText
		}
		matched = w.regexp.MatchString(txt)
		csAdd()
	}
	if err = sc.Err(); err != nil {
		return nil, err
	}

	// append last one for nlines != 0
	if c.line != nil {
		c.after = lq.PopAll()
		cs = append(cs, c)
	}
	return cs, nil
}
<commit_msg>fix InternalError<commit_after>package main

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"sync"
	"unicode/utf8"
)

var ErrInvalidText = errors.New("unavailable encoding")

type InternalError struct {
	path string
	e    error
}

func (ei *InternalError) Error() string {
	return fmt.Sprintf("internal error:%v:%s", ei.e, ei.path)
}

type Line struct {
	Num uint
	Str string
}

// cap(q) is required always greater than 1
type LineQueue struct{ q chan *Line }

func NewLineQueue(capacity uint) (*LineQueue, error) {
	if capacity == 0 {
		return nil, errors.New("capacity is 0")
	}
	return &LineQueue{make(chan *Line, capacity)}, nil
}
func (lq *LineQueue) Len() int { return len(lq.q) }
func (lq *LineQueue) Cap() int { return cap(lq.q) }
func (lq *LineQueue) Push(l *Line) {
	select {
	case lq.q <- l:
	default:
		<-lq.q
		lq.q <- l
	}
}
func (lq *LineQueue) PopAll() []*Line {
	lines := make([]*Line, 0, cap(lq.q))
	for len(lq.q) != 0 {
		lines = append(lines, <-lq.q)
	}
	return lines
}

// change to struct{ index int, pos int, lines []*Line }?
type Context struct {
	line   *Line
	before []*Line
	after  []*Line
}

func FprintContexts(writer io.Writer, prefix string, cs []*Context) error {
	if cs == nil {
		return nil
	}
	var err error
	f := func(ls []*Line) {
		for _, l := range ls {
			_, err = fmt.Fprintf(writer, "%s%d-%s\n", prefix, l.Num, l.Str)
		}
	}
	for _, c := range cs {
		f(c.before)
		if err != nil {
			return err
		}
		_, err = fmt.Fprintf(writer, "%s%d:%s\n", prefix, c.line.Num, c.line.Str)
		if err != nil {
			return err
		}
		f(c.after)
		if err != nil {
			return err
		}
	}
	return nil
}

type File struct {
	path string
	cs   []*Context
}

func (f *File) Fprint(writer io.Writer) error {
	var err error
	if len(f.cs) == 0 {
		return nil
	}
	_, err = fmt.Fprintln(writer, f.path)
	if err != nil {
		return err
	}
	err = FprintContexts(writer, "", f.cs)
	if err != nil {
		return err
	}
	_, err = fmt.Fprintln(writer)
	if err != nil {
		return err
	}
	return nil
}

func (f *File) FprintVerbose(writer io.Writer) error {
	var err error
	if len(f.cs) == 0 {
		return nil
	}
	err = FprintContexts(writer, f.path, f.cs)
	if err != nil {
		return err
	}
	_, err = fmt.Fprintln(writer)
	if err != nil {
		return err
	}
	return nil
}

func FprintFiles(writer io.Writer, fs ...*File) error {
	var err error
	for i := range fs {
		err = fs[i].Fprint(writer)
		if err != nil {
			return err
		}
	}
	return nil
}

func FprintFilesVerbose(writer io.Writer, fs ...*File) error {
	var err error
	for i := range fs {
		err = fs[i].FprintVerbose(writer)
		if err != nil {
			return err
		}
	}
	return nil
}

type Walker struct {
	fileQueue chan string
	dirQueue  chan []string

	regexp *regexp.Regexp

	nworker       int
	checked       map[string]bool
	mu            sync.Mutex
	wg            sync.WaitGroup
	once          sync.Once
	internalError *InternalError

	log *log.Logger
}

func NewWalker() *Walker {
	nworker := runtime.NumCPU() / 4
	if nworker < 2 {
		nworker = 2
	}
	return &Walker{
		fileQueue:     make(chan string, 128),
		dirQueue:      make(chan []string, nworker),
		nworker:       nworker,
		checked:       make(map[string]bool),
		log:           log.New(ioutil.Discard, Name+":", 0),
		internalError: nil,
	}
}

func (w *Walker) SetLogOutput(writer io.Writer) { w.log.SetOutput(writer) }

func (w *Walker) setInternalError(err error, path string) {
	w.once.Do(func() { w.internalError = &InternalError{e: err, path: path} })
}

func (w *Walker) sendQueue(paths ...string) {
	var dirs []string
	for i := range paths {
		abs, err := filepath.Abs(paths[i])
		if err != nil {
			w.setInternalError(err, abs)
			w.log.Printf("[Err]:%v", err)
			continue
		}
		fi, err := os.Stat(abs)
		if err != nil {
			w.setInternalError(err, abs)
			w.log.Printf("[Errr]:%v", err)
			continue
		}
		if fi.IsDir() {
			dirs = append(dirs, abs)
		} else if fi.Mode().IsRegular() {
			w.wg.Add(1)
			w.fileQueue <- abs
		}
	}
	w.wg.Add(1)
	w.dirQueue <- dirs
}

func (w *Walker) Start(pat string, nlines uint, paths ...string) (<-chan *File, func() error, error) {
	var re *regexp.Regexp
	re, err := regexp.Compile(pat)
	if err != nil {
		return nil, nil, err
	}
	w.regexp = re

	resultQueue := make(chan *File, cap(w.fileQueue))
	done := make(chan struct{})
	wait := func() error {
		w.wg.Wait()
		return w.internalError
	}

	for i := 0; i != w.nworker; i++ {
		go w.dirWalker(done)
		go w.fileWalker(done, resultQueue, nlines)
	}
	w.sendQueue(paths...)
	go func() {
		w.wg.Wait()
		close(done)
		close(resultQueue)
	}()
	return resultQueue, wait, nil
}

// for goroutine
// send tasks to {file,dir}Queue
func (w *Walker) dirWalker(done <-chan struct{}) {
	var nextDirs []string
	var dirs []string
	for ; true; w.wg.Done() {
		select {
		case <-done:
			return
		case dirs = <-w.dirQueue:
			for i := range dirs {
				fis, err := ioutil.ReadDir(dirs[i])
				if err != nil {
					w.setInternalError(err, dirs[i])
					w.log.Printf("[Err]:%s:%v", dirs[i], err)
					if os.IsNotExist(err) || os.IsPermission(err) {
						continue
					}
					// unexpected error
					panic(err)
				}
				for _, fi := range fis {
					if fi.Mode().IsRegular() {
						w.wg.Add(1)
						w.fileQueue <- filepath.Join(dirs[i], fi.Name())
					} else if fi.IsDir() {
						nextDirs = append(nextDirs, filepath.Join(dirs[i], fi.Name()))
					}
				}
			}
			if nextDirs != nil {
				w.wg.Add(1)
				w.dirQueue <- nextDirs
				nextDirs = nil
			}
		case <-done:
			return
		}
	}
}

// for goroutine
func (w *Walker) fileWalker(done <-chan struct{}, resultQueue chan<- *File, nlines uint) {
	var lq *LineQueue
	if nlines != 0 {
		// not need error check
		lq, _ = NewLineQueue(nlines)
	}
	var file string
	var err error
	var cs []*Context
	for ; true; w.wg.Done() {
		select {
		case <-done:
			return
		case file = <-w.fileQueue:
			w.mu.Lock()
			if w.checked[file] {
				w.mu.Unlock()
				continue
			}
			w.checked[file] = true
			w.mu.Unlock()

			cs, err = w.readFile(file, lq)
			if err != nil {
				w.setInternalError(err, file)
				w.log.Printf("[Err]:%s:%v", file, err)
				if os.IsNotExist(err) || os.IsPermission(err) {
					continue
				}
				if err == bufio.ErrTooLong {
					continue
				}
				if err == ErrInvalidText {
					continue
				}
				// unexpected error
				panic(err)
			}
			w.log.Println(file)
			if len(cs) != 0 {
				resultQueue <- &File{
					path: file,
					cs:   cs,
				}
			}
		}
	}
}

// TODO? readFile(f *File, file string) error
func (w *Walker) readFile(file string, lq *LineQueue) ([]*Context, error) {
	f, err := os.Open(file)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	var cs []*Context
	var c = new(Context)
	var txt string
	var i uint
	var matched bool

	var csAdd func()
	if lq == nil {
		csAdd = func() {
			if matched {
				cs = append(cs, &Context{
					line:   &Line{i, txt},
					before: []*Line{},
					after:  []*Line{},
				})
			}
		}
	} else {
		defer lq.PopAll()
		csAdd = func() {
			if c.line != nil {
				if matched {
					c.after = lq.PopAll()
					cs = append(cs, c)
					c = &Context{
						before: []*Line{},
						line:   &Line{i, txt},
						after:  []*Line{},
					}
					return
				}
				if lq.Len() == lq.Cap() {
					c.after = lq.PopAll()
					cs = append(cs, c)
					c = new(Context)
				}
			} else if matched {
				c.before = lq.PopAll()
				c.line = &Line{i, txt}
				return
			}
			lq.Push(&Line{i, txt})
		}
	}

	sc := bufio.NewScanner(f)
	for i = uint(1); sc.Scan(); i++ {
		if i == 0 {
			return nil, errors.New("too many lines")
		}
		txt = sc.Text()
		if !utf8.ValidString(txt) {
			return nil, ErrInvalidText
		}
		matched = w.regexp.MatchString(txt)
		csAdd()
	}
	if err = sc.Err(); err != nil {
		return nil, err
	}

	// append last one for nlines != 0
	if c.line != nil {
		c.after = lq.PopAll()
		cs = append(cs, c)
	}
	return cs, nil
}
<|endoftext|>