<commit_before>// Copyright 2018 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

// Package build contains helper functions for building kernels/images.
package build

import (
	"bytes"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/google/syzkaller/pkg/osutil"
	"github.com/google/syzkaller/pkg/report"
	"github.com/google/syzkaller/pkg/vcs"
	"github.com/google/syzkaller/sys/targets"
)

// Params is input arguments for the Image function.
type Params struct {
	TargetOS     string
	TargetArch   string
	VMType       string
	KernelDir    string
	OutputDir    string
	Compiler     string
	Ccache       string
	UserspaceDir string
	CmdlineFile  string
	SysctlFile   string
	Config       []byte
}

// Information that is returned from the Image function.
type ImageDetails struct {
	Signature  string
	CompilerID string
}

// Image creates a disk image for the specified OS/ARCH/VM.
// Kernel is taken from KernelDir, userspace system is taken from UserspaceDir.
// If CmdlineFile is not empty, contents of the file are appended to the kernel command line.
// If SysctlFile is not empty, contents of the file are appended to the image /etc/sysctl.conf.
// Output is stored in OutputDir and includes (everything except for image is optional):
//  - image: the image
//  - key: ssh key for the image
//  - kernel: kernel for injected boot
//  - initrd: initrd for injected boot
//  - kernel.config: actual kernel config used during build
//  - obj/: directory with kernel object files (this should match KernelObject
//    specified in sys/targets, e.g. vmlinux for linux)
// The returned structure contains a kernel ID that will be the same for kernels
// with the same runtime behavior, and different for kernels with different runtime
// behavior. Binary equal builds, or builds that differ only in e.g. debug info,
// have the same ID. The ID may be empty if OS implementation does not have
// a way to calculate such IDs.
// Also that structure provides a compiler ID field that contains the name and
// the version of the compiler/toolchain that was used to build the kernel.
// The CompilerID field is not guaranteed to be non-empty.
func Image(params Params) (details ImageDetails, err error) {
	builder, err := getBuilder(params.TargetOS, params.TargetArch, params.VMType)
	if err != nil {
		return
	}
	if err = osutil.MkdirAll(filepath.Join(params.OutputDir, "obj")); err != nil {
		return
	}
	if len(params.Config) != 0 {
		// Write kernel config early, so that it's captured on build failures.
		if err = osutil.WriteFile(filepath.Join(params.OutputDir, "kernel.config"), params.Config); err != nil {
			err = fmt.Errorf("failed to write config file: %v", err)
			return
		}
	}
	details, err = builder.build(params)
	if details.CompilerID == "" {
		// Fill in the compiler info even if the build failed.
		var idErr error
		details.CompilerID, idErr = compilerIdentity(params.Compiler)
		if err == nil {
			err = idErr
		} // Try to preserve the build error otherwise.
	}
	if err != nil {
		err = extractRootCause(err, params.TargetOS, params.KernelDir)
		return
	}
	if key := filepath.Join(params.OutputDir, "key"); osutil.IsExist(key) {
		if err := os.Chmod(key, 0600); err != nil {
			return details, fmt.Errorf("failed to chmod 0600 %v: %v", key, err)
		}
	}
	return
}

func Clean(targetOS, targetArch, vmType, kernelDir string) error {
	builder, err := getBuilder(targetOS, targetArch, vmType)
	if err != nil {
		return err
	}
	return builder.clean(kernelDir, targetArch)
}

type KernelError struct {
	Report     []byte
	Output     []byte
	Recipients vcs.Recipients
	guiltyFile string
}

func (err *KernelError) Error() string {
	return string(err.Report)
}

type builder interface {
	build(params Params) (ImageDetails, error)
	clean(kernelDir, targetArch string) error
}

func getBuilder(targetOS, targetArch, vmType string) (builder, error) {
	var supported = []struct {
		OS    string
		archs []string
		vms   []string
		b     builder
	}{
		{targets.Linux, []string{targets.AMD64}, []string{"gvisor"}, gvisor{}},
		{targets.Linux, []string{targets.AMD64}, []string{"gce", "qemu"}, linux{}},
		{targets.Linux, []string{targets.ARM, targets.ARM64, targets.I386, targets.MIPS64LE,
			targets.PPC64LE, targets.S390x, targets.RiscV64}, []string{"qemu"}, linux{}},
		{targets.Fuchsia, []string{targets.AMD64, targets.ARM64}, []string{"qemu"}, fuchsia{}},
		{targets.Akaros, []string{targets.AMD64}, []string{"qemu"}, akaros{}},
		{targets.OpenBSD, []string{targets.AMD64}, []string{"gce", "vmm"}, openbsd{}},
		{targets.NetBSD, []string{targets.AMD64}, []string{"gce", "qemu"}, netbsd{}},
		{targets.FreeBSD, []string{targets.AMD64}, []string{"gce", "qemu"}, freebsd{}},
		{targets.Darwin, []string{targets.AMD64}, []string{"qemu"}, darwin{}},
		{targets.TestOS, []string{targets.TestArch64}, []string{"qemu"}, test{}},
	}
	for _, s := range supported {
		if targetOS == s.OS {
			for _, arch := range s.archs {
				if targetArch == arch {
					for _, vm := range s.vms {
						if vmType == vm {
							return s.b, nil
						}
					}
				}
			}
		}
	}
	return nil, fmt.Errorf("unsupported image type %v/%v/%v", targetOS, targetArch, vmType)
}

func compilerIdentity(compiler string) (string, error) {
	if compiler == "" {
		return "", nil
	}

	bazel := strings.HasSuffix(compiler, "bazel")

	arg, timeout := "--version", time.Minute
	if bazel {
		// Bazel episodically fails with 1 min timeout.
		arg, timeout = "", 10*time.Minute
	}
	output, err := osutil.RunCmd(timeout, "", compiler, arg)
	if err != nil {
		return "", err
	}
	for _, line := range strings.Split(string(output), "\n") {
		if bazel {
			// Strip extracting and log lines...
			if strings.Contains(line, "Extracting Bazel") {
				continue
			}
			if strings.HasPrefix(line, "INFO: ") {
				continue
			}
			if strings.HasPrefix(line, "WARNING: ") {
				continue
			}
		}

		return strings.TrimSpace(line), nil
	}
	return "", fmt.Errorf("no output from compiler --version")
}

func extractRootCause(err error, OS, kernelSrc string) error {
	if err == nil {
		return nil
	}
	verr, ok := err.(*osutil.VerboseError)
	if !ok {
		return err
	}
	reason, file := extractCauseInner(verr.Output, kernelSrc)
	if len(reason) == 0 {
		return err
	}
	kernelErr := &KernelError{
		Report:     reason,
		Output:     verr.Output,
		guiltyFile: file,
	}
	if file != "" && OS == targets.Linux {
		maintainers, err := report.GetLinuxMaintainers(kernelSrc, file)
		if err != nil {
			kernelErr.Output = append(kernelErr.Output, err.Error()...)
		}
		kernelErr.Recipients = maintainers
	}
	return kernelErr
}

func extractCauseInner(s []byte, kernelSrc string) ([]byte, string) {
	lines := extractCauseRaw(s)
	const maxLines = 20
	if len(lines) > maxLines {
		lines = lines[:maxLines]
	}
	var stripPrefix []byte
	if kernelSrc != "" {
		stripPrefix = []byte(kernelSrc)
		if stripPrefix[len(stripPrefix)-1] != filepath.Separator {
			stripPrefix = append(stripPrefix, filepath.Separator)
		}
	}
	file := ""
	for i := range lines {
		if stripPrefix != nil {
			lines[i] = bytes.Replace(lines[i], stripPrefix, nil, -1)
		}
		if file == "" {
			for _, fileRe := range fileRes {
				match := fileRe.FindSubmatch(lines[i])
				if match != nil {
					file = string(match[1])
					if file[0] != '/' {
						break
					}
					// We already removed kernel source prefix,
					// if we still have an absolute path, it's probably pointing
					// to compiler/system libraries (not going to work).
					file = ""
				}
			}
		}
	}
	file = strings.TrimPrefix(file, "./")
	if strings.HasSuffix(file, ".o") {
		// Linker may point to object files instead.
		file = strings.TrimSuffix(file, ".o") + ".c"
	}
	res := bytes.Join(lines, []byte{'\n'})
	// gcc uses these weird quotes around identifiers, which may be
	// mis-rendered by systems that don't understand utf-8.
	res = bytes.Replace(res, []byte("‘"), []byte{'\''}, -1)
	res = bytes.Replace(res, []byte("’"), []byte{'\''}, -1)
	return res, file
}

func extractCauseRaw(s []byte) [][]byte {
	weak := true
	var cause [][]byte
	dedup := make(map[string]bool)
	for _, line := range bytes.Split(s, []byte{'\n'}) {
		for _, pattern := range buildFailureCauses {
			if !pattern.pattern.Match(line) {
				continue
			}
			if weak && !pattern.weak {
				cause = nil
				dedup = make(map[string]bool)
			}
			if dedup[string(line)] {
				continue
			}
			dedup[string(line)] = true
			if cause == nil {
				weak = pattern.weak
			}
			cause = append(cause, line)
			break
		}
	}
	return cause
}

type buildFailureCause struct {
	pattern *regexp.Regexp
	weak    bool
}

var buildFailureCauses = [...]buildFailureCause{
	{pattern: regexp.MustCompile(`: error: `)},
	{pattern: regexp.MustCompile(`ERROR: `)},
	{pattern: regexp.MustCompile(`: fatal error: `)},
	{pattern: regexp.MustCompile(`: undefined reference to`)},
	{pattern: regexp.MustCompile(`: multiple definition of`)},
	{pattern: regexp.MustCompile(`: Permission denied`)},
	{pattern: regexp.MustCompile(`: not found`)},
	{pattern: regexp.MustCompile(`^([a-zA-Z0-9_\-/.]+):[0-9]+:([0-9]+:)?.*(error|invalid|fatal|wrong)`)},
	{pattern: regexp.MustCompile(`FAILED unresolved symbol`)},
	{pattern: regexp.MustCompile(`No rule to make target`)},
	{weak: true, pattern: regexp.MustCompile(`: final link failed: `)},
	{weak: true, pattern: regexp.MustCompile(`collect2: error: `)},
	{weak: true, pattern: regexp.MustCompile(`FAILED: Build did NOT complete`)},
}

var fileRes = []*regexp.Regexp{
	regexp.MustCompile(`^([a-zA-Z0-9_\-/.]+):[0-9]+:([0-9]+:)? `),
	regexp.MustCompile(`^(?:ld: )?(([a-zA-Z0-9_\-/.]+?)\.o):`),
	regexp.MustCompile(`; (([a-zA-Z0-9_\-/.]+?)\.o):`),
}
<commit_msg>pkg/build: simplify the code<commit_after>// Copyright 2018 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

// Package build contains helper functions for building kernels/images.
package build

import (
	"bytes"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/google/syzkaller/pkg/osutil"
	"github.com/google/syzkaller/pkg/report"
	"github.com/google/syzkaller/pkg/vcs"
	"github.com/google/syzkaller/sys/targets"
)

// Params is input arguments for the Image function.
type Params struct {
	TargetOS     string
	TargetArch   string
	VMType       string
	KernelDir    string
	OutputDir    string
	Compiler     string
	Ccache       string
	UserspaceDir string
	CmdlineFile  string
	SysctlFile   string
	Config       []byte
}

// Information that is returned from the Image function.
type ImageDetails struct {
	Signature  string
	CompilerID string
}

// Image creates a disk image for the specified OS/ARCH/VM.
// Kernel is taken from KernelDir, userspace system is taken from UserspaceDir.
// If CmdlineFile is not empty, contents of the file are appended to the kernel command line.
// If SysctlFile is not empty, contents of the file are appended to the image /etc/sysctl.conf.
// Output is stored in OutputDir and includes (everything except for image is optional):
//  - image: the image
//  - key: ssh key for the image
//  - kernel: kernel for injected boot
//  - initrd: initrd for injected boot
//  - kernel.config: actual kernel config used during build
//  - obj/: directory with kernel object files (this should match KernelObject
//    specified in sys/targets, e.g. vmlinux for linux)
// The returned structure contains a kernel ID that will be the same for kernels
// with the same runtime behavior, and different for kernels with different runtime
// behavior. Binary equal builds, or builds that differ only in e.g. debug info,
// have the same ID. The ID may be empty if OS implementation does not have
// a way to calculate such IDs.
// Also that structure provides a compiler ID field that contains the name and
// the version of the compiler/toolchain that was used to build the kernel.
// The CompilerID field is not guaranteed to be non-empty.
func Image(params Params) (details ImageDetails, err error) {
	builder, err := getBuilder(params.TargetOS, params.TargetArch, params.VMType)
	if err != nil {
		return
	}
	if err = osutil.MkdirAll(filepath.Join(params.OutputDir, "obj")); err != nil {
		return
	}
	if len(params.Config) != 0 {
		// Write kernel config early, so that it's captured on build failures.
		if err = osutil.WriteFile(filepath.Join(params.OutputDir, "kernel.config"), params.Config); err != nil {
			err = fmt.Errorf("failed to write config file: %v", err)
			return
		}
	}
	details, err = builder.build(params)
	if details.CompilerID == "" {
		// Fill in the compiler info even if the build failed.
		var idErr error
		details.CompilerID, idErr = compilerIdentity(params.Compiler)
		if err == nil {
			err = idErr
		} // Try to preserve the build error otherwise.
	}
	if err != nil {
		err = extractRootCause(err, params.TargetOS, params.KernelDir)
		return
	}
	if key := filepath.Join(params.OutputDir, "key"); osutil.IsExist(key) {
		if err := os.Chmod(key, 0600); err != nil {
			return details, fmt.Errorf("failed to chmod 0600 %v: %v", key, err)
		}
	}
	return
}

func Clean(targetOS, targetArch, vmType, kernelDir string) error {
	builder, err := getBuilder(targetOS, targetArch, vmType)
	if err != nil {
		return err
	}
	return builder.clean(kernelDir, targetArch)
}

type KernelError struct {
	Report     []byte
	Output     []byte
	Recipients vcs.Recipients
	guiltyFile string
}

func (err *KernelError) Error() string {
	return string(err.Report)
}

type builder interface {
	build(params Params) (ImageDetails, error)
	clean(kernelDir, targetArch string) error
}

func getBuilder(targetOS, targetArch, vmType string) (builder, error) {
	if targetOS == targets.Linux && vmType == "gvisor" {
		return gvisor{}, nil
	}
	builders := map[string]builder{
		targets.Linux:   linux{},
		targets.Fuchsia: fuchsia{},
		targets.Akaros:  akaros{},
		targets.OpenBSD: openbsd{},
		targets.NetBSD:  netbsd{},
		targets.FreeBSD: freebsd{},
		targets.Darwin:  darwin{},
		targets.TestOS:  test{},
	}
	if builder, ok := builders[targetOS]; ok {
		return builder, nil
	}
	return nil, fmt.Errorf("unsupported image type %v/%v/%v", targetOS, targetArch, vmType)
}

func compilerIdentity(compiler string) (string, error) {
	if compiler == "" {
		return "", nil
	}

	bazel := strings.HasSuffix(compiler, "bazel")

	arg, timeout := "--version", time.Minute
	if bazel {
		// Bazel episodically fails with 1 min timeout.
		arg, timeout = "", 10*time.Minute
	}
	output, err := osutil.RunCmd(timeout, "", compiler, arg)
	if err != nil {
		return "", err
	}
	for _, line := range strings.Split(string(output), "\n") {
		if bazel {
			// Strip extracting and log lines...
			if strings.Contains(line, "Extracting Bazel") {
				continue
			}
			if strings.HasPrefix(line, "INFO: ") {
				continue
			}
			if strings.HasPrefix(line, "WARNING: ") {
				continue
			}
		}

		return strings.TrimSpace(line), nil
	}
	return "", fmt.Errorf("no output from compiler --version")
}

func extractRootCause(err error, OS, kernelSrc string) error {
	if err == nil {
		return nil
	}
	verr, ok := err.(*osutil.VerboseError)
	if !ok {
		return err
	}
	reason, file := extractCauseInner(verr.Output, kernelSrc)
	if len(reason) == 0 {
		return err
	}
	kernelErr := &KernelError{
		Report:     reason,
		Output:     verr.Output,
		guiltyFile: file,
	}
	if file != "" && OS == targets.Linux {
		maintainers, err := report.GetLinuxMaintainers(kernelSrc, file)
		if err != nil {
			kernelErr.Output = append(kernelErr.Output, err.Error()...)
		}
		kernelErr.Recipients = maintainers
	}
	return kernelErr
}

func extractCauseInner(s []byte, kernelSrc string) ([]byte, string) {
	lines := extractCauseRaw(s)
	const maxLines = 20
	if len(lines) > maxLines {
		lines = lines[:maxLines]
	}
	var stripPrefix []byte
	if kernelSrc != "" {
		stripPrefix = []byte(kernelSrc)
		if stripPrefix[len(stripPrefix)-1] != filepath.Separator {
			stripPrefix = append(stripPrefix, filepath.Separator)
		}
	}
	file := ""
	for i := range lines {
		if stripPrefix != nil {
			lines[i] = bytes.Replace(lines[i], stripPrefix, nil, -1)
		}
		if file == "" {
			for _, fileRe := range fileRes {
				match := fileRe.FindSubmatch(lines[i])
				if match != nil {
					file = string(match[1])
					if file[0] != '/' {
						break
					}
					// We already removed kernel source prefix,
					// if we still have an absolute path, it's probably pointing
					// to compiler/system libraries (not going to work).
					file = ""
				}
			}
		}
	}
	file = strings.TrimPrefix(file, "./")
	if strings.HasSuffix(file, ".o") {
		// Linker may point to object files instead.
		file = strings.TrimSuffix(file, ".o") + ".c"
	}
	res := bytes.Join(lines, []byte{'\n'})
	// gcc uses these weird quotes around identifiers, which may be
	// mis-rendered by systems that don't understand utf-8.
	res = bytes.Replace(res, []byte("‘"), []byte{'\''}, -1)
	res = bytes.Replace(res, []byte("’"), []byte{'\''}, -1)
	return res, file
}

func extractCauseRaw(s []byte) [][]byte {
	weak := true
	var cause [][]byte
	dedup := make(map[string]bool)
	for _, line := range bytes.Split(s, []byte{'\n'}) {
		for _, pattern := range buildFailureCauses {
			if !pattern.pattern.Match(line) {
				continue
			}
			if weak && !pattern.weak {
				cause = nil
				dedup = make(map[string]bool)
			}
			if dedup[string(line)] {
				continue
			}
			dedup[string(line)] = true
			if cause == nil {
				weak = pattern.weak
			}
			cause = append(cause, line)
			break
		}
	}
	return cause
}

type buildFailureCause struct {
	pattern *regexp.Regexp
	weak    bool
}

var buildFailureCauses = [...]buildFailureCause{
	{pattern: regexp.MustCompile(`: error: `)},
	{pattern: regexp.MustCompile(`ERROR: `)},
	{pattern: regexp.MustCompile(`: fatal error: `)},
	{pattern: regexp.MustCompile(`: undefined reference to`)},
	{pattern: regexp.MustCompile(`: multiple definition of`)},
	{pattern: regexp.MustCompile(`: Permission denied`)},
	{pattern: regexp.MustCompile(`: not found`)},
	{pattern: regexp.MustCompile(`^([a-zA-Z0-9_\-/.]+):[0-9]+:([0-9]+:)?.*(error|invalid|fatal|wrong)`)},
	{pattern: regexp.MustCompile(`FAILED unresolved symbol`)},
	{pattern: regexp.MustCompile(`No rule to make target`)},
	{weak: true, pattern: regexp.MustCompile(`: final link failed: `)},
	{weak: true, pattern: regexp.MustCompile(`collect2: error: `)},
	{weak: true, pattern: regexp.MustCompile(`FAILED: Build did NOT complete`)},
}

var fileRes = []*regexp.Regexp{
	regexp.MustCompile(`^([a-zA-Z0-9_\-/.]+):[0-9]+:([0-9]+:)? `),
	regexp.MustCompile(`^(?:ld: )?(([a-zA-Z0-9_\-/.]+?)\.o):`),
	regexp.MustCompile(`; (([a-zA-Z0-9_\-/.]+?)\.o):`),
}
<|endoftext|>


<commit_before>package jobs

import (
	"fmt"
	"math/rand"
	"sync/atomic"
	"time"

	log "github.com/Sirupsen/logrus"
)

var (
	defaultConcurrency  = 1
	defaultMaxExecCount = 3
	defaultMaxExecTime  = 60 * time.Second
	defaultRetryDelay   = 60 * time.Millisecond
	defaultTimeout      = 10 * time.Second

	maxMaxExecCount = 5
	maxMaxExecTime  = 5 * time.Minute
	maxTimeout      = 1 * time.Minute
)

type (
	// WorkerFunc represent the work function that a worker should implement.
	WorkerFunc func(msg *Message, timeout <-chan time.Time) error

	// Worker is a unit of work that will consume from a queue and execute the do
	// method for each jobs it pulls.
	Worker struct {
		Domain string
		Type   string
		Conf   *WorkerConfig

		jobs    Queue
		started int32
	}
)

// Start is used to start the worker consumption of messages from its queue.
func (w *Worker) Start(q Queue) {
	if !atomic.CompareAndSwapInt32(&w.started, 0, 1) {
		return
	}
	w.jobs = q
	for i := 0; i < int(w.Conf.Concurrency); i++ {
		name := fmt.Sprintf("%s/%s/%d", w.Domain, w.Type, i)
		go w.work(name)
	}
}

func (w *Worker) work(workerID string) {
	// TODO: err handling and persistence
	for {
		job, err := w.jobs.Consume()
		if err != nil {
			if err != ErrQueueClosed {
				log.Errorf("[job] %s: error while consuming queue (%s)",
					workerID, err.Error())
			}
			return
		}
		infos := job.Infos()
		if err = job.AckConsumed(); err != nil {
			log.Errorf("[job] %s: error acking consume job %s (%s)",
				workerID, infos.ID, err.Error())
			continue
		}
		t := &task{
			infos: infos,
			conf:  w.defaultedConf(infos.Options),
		}
		if err = t.run(); err != nil {
			log.Errorf("[job] %s: error while performing job %s (%s)",
				workerID, infos.ID, err.Error())
			err = job.Nack(err)
		} else {
			err = job.Ack()
		}
		if err != nil {
			log.Errorf("[job] %s: error while acking job done %s (%s)",
				workerID, infos.ID, err.Error())
		}
	}
}

func (w *Worker) defaultedConf(opts *JobOptions) *WorkerConfig {
	c := &(*w.Conf)
	if opts != nil {
		if opts.MaxExecCount != 0 {
			c.MaxExecCount = opts.MaxExecCount
		}
		if opts.MaxExecTime > 0 {
			c.MaxExecTime = opts.MaxExecTime
		}
		if opts.Timeout > 0 {
			c.Timeout = opts.Timeout
		}
	}
	if c.Concurrency == 0 {
		c.Concurrency = uint(defaultConcurrency)
	}
	if c.MaxExecCount == 0 {
		c.MaxExecCount = uint(defaultMaxExecCount)
	} else if c.MaxExecCount > uint(maxMaxExecCount) {
		c.MaxExecCount = uint(maxMaxExecCount)
	}
	if c.MaxExecTime == 0 {
		c.MaxExecTime = defaultMaxExecTime
	} else if c.MaxExecTime > maxMaxExecTime {
		c.MaxExecTime = maxMaxExecTime
	}
	if c.RetryDelay == 0 {
		c.RetryDelay = defaultRetryDelay
	}
	if c.Timeout == 0 {
		c.Timeout = defaultTimeout
	} else if c.Timeout > maxTimeout {
		c.Timeout = maxTimeout
	}
	return c
}

// Stop will stop the worker's consumption of its queue. It will also close the
// associated queue.
func (w *Worker) Stop() {
	if !atomic.CompareAndSwapInt32(&w.started, 1, 0) {
		return
	}
	w.jobs.Close()
}

type task struct {
	infos *JobInfos
	conf  *WorkerConfig

	startTime time.Time
	execCount uint
}

func (t *task) run() (err error) {
	t.startTime = time.Now()
	t.execCount = 0
	for {
		retry, delay, timeout := t.nextDelay()
		if !retry {
			return err
		}
		if err != nil {
			log.Warnf("[job] %s: %s (retry in %s)", t.infos.ID, err.Error(), delay)
		}
		if delay > 0 {
			time.Sleep(delay)
		}
		log.Debugf("[job] %s: run %d (timeout %s)", t.infos.ID, t.execCount, timeout)
		err = t.conf.WorkerFunc(t.infos.Message, time.After(timeout))
		if err == nil {
			break
		}
		t.execCount++
	}
	return nil
}

func (t *task) nextDelay() (bool, time.Duration, time.Duration) {
	c := t.conf
	execTime := time.Since(t.startTime)

	if t.execCount >= c.MaxExecCount || execTime > c.MaxExecTime {
		return false, 0, 0
	}

	// the worker timeout should take into account the maximum execution time
	// allowed to the task
	timeout := c.Timeout
	if execTime+timeout > c.MaxExecTime {
		timeout = execTime - c.MaxExecTime
	}

	var nextDelay time.Duration
	if t.execCount == 0 {
		// on first execution, execute immediatly
		nextDelay = 0
	} else {
		nextDelay = c.RetryDelay << (t.execCount - 1)

		// fuzzDelay number between delay * (1 +/- 0.1)
		fuzzDelay := int(0.1 * float64(nextDelay))
		nextDelay = nextDelay + time.Duration((rand.Intn(2*fuzzDelay) - fuzzDelay))
	}

	if execTime+nextDelay > c.MaxExecTime {
		return false, 0, 0
	}

	return true, nextDelay, timeout
}
<commit_msg>Allow worker config to be superior to default global config<commit_after>package jobs

import (
	"fmt"
	"math/rand"
	"sync/atomic"
	"time"

	log "github.com/Sirupsen/logrus"
)

var (
	defaultConcurrency  = 1
	defaultMaxExecCount = 3
	defaultMaxExecTime  = 60 * time.Second
	defaultRetryDelay   = 60 * time.Millisecond
	defaultTimeout      = 10 * time.Second

	maxMaxExecCount = 5
	maxMaxExecTime  = 5 * time.Minute
	maxTimeout      = 1 * time.Minute
)

type (
	// WorkerFunc represent the work function that a worker should implement.
	WorkerFunc func(msg *Message, timeout <-chan time.Time) error

	// Worker is a unit of work that will consume from a queue and execute the do
	// method for each jobs it pulls.
	Worker struct {
		Domain string
		Type   string
		Conf   *WorkerConfig

		jobs    Queue
		started int32
	}
)

// Start is used to start the worker consumption of messages from its queue.
func (w *Worker) Start(q Queue) {
	if !atomic.CompareAndSwapInt32(&w.started, 0, 1) {
		return
	}
	w.jobs = q
	for i := 0; i < int(w.Conf.Concurrency); i++ {
		name := fmt.Sprintf("%s/%s/%d", w.Domain, w.Type, i)
		go w.work(name)
	}
}

func (w *Worker) work(workerID string) {
	// TODO: err handling and persistence
	for {
		job, err := w.jobs.Consume()
		if err != nil {
			if err != ErrQueueClosed {
				log.Errorf("[job] %s: error while consuming queue (%s)",
					workerID, err.Error())
			}
			return
		}
		infos := job.Infos()
		if err = job.AckConsumed(); err != nil {
			log.Errorf("[job] %s: error acking consume job %s (%s)",
				workerID, infos.ID, err.Error())
			continue
		}
		t := &task{
			infos: infos,
			conf:  w.defaultedConf(infos.Options),
		}
		if err = t.run(); err != nil {
			log.Errorf("[job] %s: error while performing job %s (%s)",
				workerID, infos.ID, err.Error())
			err = job.Nack(err)
		} else {
			err = job.Ack()
		}
		if err != nil {
			log.Errorf("[job] %s: error while acking job done %s (%s)",
				workerID, infos.ID, err.Error())
		}
	}
}

func (w *Worker) defaultedConf(opts *JobOptions) *WorkerConfig {
	c := &(*w.Conf)
	if c.Concurrency == 0 {
		c.Concurrency = uint(defaultConcurrency)
	}
	if c.MaxExecCount == 0 {
		c.MaxExecCount = uint(defaultMaxExecCount)
	}
	if c.MaxExecTime == 0 {
		c.MaxExecTime = defaultMaxExecTime
	}
	if c.RetryDelay == 0 {
		c.RetryDelay = defaultRetryDelay
	}
	if c.Timeout == 0 {
		c.Timeout = defaultTimeout
	}
	if opts == nil {
		return c
	}
	if opts.MaxExecCount != 0 && opts.MaxExecCount < c.MaxExecCount {
		c.MaxExecCount = opts.MaxExecCount
	}
	if opts.MaxExecTime > 0 && opts.MaxExecTime < c.MaxExecTime {
		c.MaxExecTime = opts.MaxExecTime
	}
	if opts.Timeout > 0 && opts.Timeout < c.Timeout {
		c.Timeout = opts.Timeout
	}
	return c
}

// Stop will stop the worker's consumption of its queue. It will also close the
// associated queue.
func (w *Worker) Stop() {
	if !atomic.CompareAndSwapInt32(&w.started, 1, 0) {
		return
	}
	w.jobs.Close()
}

type task struct {
	infos *JobInfos
	conf  *WorkerConfig

	startTime time.Time
	execCount uint
}

func (t *task) run() (err error) {
	t.startTime = time.Now()
	t.execCount = 0
	for {
		retry, delay, timeout := t.nextDelay()
		if !retry {
			return err
		}
		if err != nil {
			log.Warnf("[job] %s: %s (retry in %s)", t.infos.ID, err.Error(), delay)
		}
		if delay > 0 {
			time.Sleep(delay)
		}
		log.Debugf("[job] %s: run %d (timeout %s)", t.infos.ID, t.execCount, timeout)
		err = t.conf.WorkerFunc(t.infos.Message, time.After(timeout))
		if err == nil {
			break
		}
		t.execCount++
	}
	return nil
}

func (t *task) nextDelay() (bool, time.Duration, time.Duration) {
	c := t.conf
	execTime := time.Since(t.startTime)

	if t.execCount >= c.MaxExecCount || execTime > c.MaxExecTime {
		return false, 0, 0
	}

	// the worker timeout should take into account the maximum execution time
	// allowed to the task
	timeout := c.Timeout
	if execTime+timeout > c.MaxExecTime {
		timeout = execTime - c.MaxExecTime
	}

	var nextDelay time.Duration
	if t.execCount == 0 {
		// on first execution, execute immediatly
		nextDelay = 0
	} else {
		nextDelay = c.RetryDelay << (t.execCount - 1)

		// fuzzDelay number between delay * (1 +/- 0.1)
		fuzzDelay := int(0.1 * float64(nextDelay))
		nextDelay = nextDelay + time.Duration((rand.Intn(2*fuzzDelay) - fuzzDelay))
	}

	if execTime+nextDelay > c.MaxExecTime {
		return false, 0, 0
	}

	return true, nextDelay, timeout
}
<|endoftext|>


<commit_before>package proto

import (
	"bufio"
	"strconv"
	"strings"
	"os"
)

type Request struct {
	Parts [][]byte
	Err   os.Error
}

var (
	ProtocolError = os.NewError("multi bulk protocol error")
)

func scanNumber(data *bufio.Reader, after byte) (n uint64, err os.Error) {

	for {
		var c byte

		c, err = data.ReadByte()
		if err != nil {
			return
		}

		switch c {
		default:
			err = ProtocolError
			return
		case '\r', '\n':
			continue
		case after:
			var sn string
			sn, err = data.ReadString('\n')
			if err != nil {
				return
			}

			return strconv.Btoui64(strings.TrimSpace(sn), 10)
		}
	}

	panic("This should never be reached!")
}

func Scan(data *bufio.Reader, ch chan *Request) {

	for {
		count, err := scanNumber(data, '*')
		if err != nil {
			ch <- &Request{Err: err}
			switch err {
			case os.EOF:
				return
			case ProtocolError:
				data.ReadString('\n')
				continue
			}
		}

		if count == 0 {
			continue
		}

		parts := make([][]byte, count)

		for count > 0 {
			size, err := scanNumber(data, '$')
			if err != nil {
				ch <- &Request{Err: err}
				if err == os.EOF {
					return
				}
			}

			// Read the data
			bytes := make([]byte, size)
			_, err = data.Read(bytes)
			if err != nil {
				ch <- &Request{Err: err}
				if err == os.EOF {
					return
				}
			}

			parts[len(parts)-int(count)] = bytes

			count--
		}

		ch <- &Request{Parts: parts}
	}

}
<commit_msg>prevent buffer overflow<commit_after>package proto

import (
	"bufio"
	"strconv"
	"strings"
	"os"
)

type Request struct {
	Parts [][]byte
	Err   os.Error
}

var (
	ProtocolError = os.NewError("multi bulk protocol error")
)

func scanNumber(data *bufio.Reader, after byte) (n uint64, err os.Error) {

	for {
		var c byte

		c, err = data.ReadByte()
		if err != nil {
			return
		}

		switch c {
		default:
			err = ProtocolError
			return
		case '\r', '\n':
			continue
		case after:
			var sn string
			sn, err = data.ReadString('\n')
			if err != nil {
				return
			}

			return strconv.Btoui64(strings.TrimSpace(sn), 10)
		}
	}

	panic("This should never be reached!")
}

func skipBytes(buf *bufio.Reader, delim byte) os.Error {
	for {
		c, err := buf.ReadByte()
		switch {
			case err != nil: return err
			case c == delim: return nil
		}
	}

	panic("can't happen")
}

func Scan(data *bufio.Reader, ch chan *Request) {

	for {
		count, err := scanNumber(data, '*')
		if err != nil {
			ch <- &Request{Err: err}
			switch err {
			case os.EOF:
				return
			case ProtocolError:
				skipBytes(data, '\n')
				continue
			}
		}

		if count == 0 {
			continue
		}

		parts := make([][]byte, count)

		for count > 0 {
			size, err := scanNumber(data, '$')
			if err != nil {
				ch <- &Request{Err: err}
				if err == os.EOF {
					return
				}
			}

			// Read the data
			bytes := make([]byte, size)
			_, err = data.Read(bytes)
			if err != nil {
				ch <- &Request{Err: err}
				if err == os.EOF {
					return
				}
			}

			parts[len(parts)-int(count)] = bytes

			count--
		}

		ch <- &Request{Parts: parts}
	}

}
<|endoftext|>


<commit_before>package store

import (
	"bytes"
	"gob"
	"doozer/util"
	"os"
	"regexp"
	"strings"
)

const (
	Clobber = ""
	Missing = "0"
	Dir     = "dir"
)

// TODO revisit this when package regexp is more complete (e.g. do Unicode)
const (
	charPat = `([a-zA-Z0-9.]|-)`
	partPat = "/" + charPat + "+"
	pathPat = "^/$|^(" + partPat + ")+$"
)

var pathRe = regexp.MustCompile(pathPat)

var (
	ErrBadMutation = os.NewError("bad mutation")
	ErrBadSnapshot = os.NewError("bad snapshot")
	ErrTooLate     = os.NewError("too late")
	ErrCasMismatch = os.NewError("cas mismatch")
)

type BadPathError struct {
	Path string
}

func (e *BadPathError) String() string {
	return "bad path: " + e.Path
}

// Applies mutations sent on Ops in sequence according to field Seqn. Any
// errors that occur will be written to ErrorPath. Duplicate operations at a
// given position are sliently ignored.
type Store struct {
	Ops     chan<- Op
	Seqns   <-chan uint64
	watchCh chan watch
	watches []watch
	todo    map[uint64]Op
	state   *state
	log     map[uint64]Event
	cleanCh chan uint64
}

// Represents an operation to apply to the store at position Seqn.
//
// If Mut is a snapshot, notifications will not be sent.
//
// If Mut is Nop, no change will be made, but a dummy event will still be sent.
type Op struct {
	Seqn uint64
	Mut  string
}

type state struct {
	ver  uint64
	root node
}

type watch struct {
	in, out chan Event
	re      *regexp.Regexp
}

// Creates a new, empty data store. Mutations will be applied in order,
// starting at number 1 (number 0 can be thought of as the creation of the
// store).
func New() *Store {
	ops   := make(chan Op)
	seqns := make(chan uint64)

	st := &Store{
		Ops:     ops,
		Seqns:   seqns,
		watchCh: make(chan watch),
		todo:    make(map[uint64]Op),
		watches: []watch{},
		state:   &state{0, emptyDir},
		log:     make(map[uint64]Event),
		cleanCh: make(chan uint64),
	}

	go st.process(ops, seqns)
	return st
}

func split(path string) []string {
	if path == "/" {
		return []string{}
	}
	return strings.Split(path[1:], "/", -1)
}

func join(parts []string) string {
	return "/" + strings.Join(parts, "/")
}

func checkPath(k string) os.Error {
	if !pathRe.MatchString(k) {
		return &BadPathError{k}
	}
	return nil
}

// Returns a mutation that can be applied to a `Store`. The mutation will set
// the contents of the file at `path` to `body` iff the CAS token of that file
// matches `cas` at the time of application.
//
// If `path` is not valid, returns a `BadPathError`.
func EncodeSet(path, body string, cas string) (mutation string, err os.Error) {
	if err = checkPath(path); err != nil {
		return
	}
	return cas + ":" + path + "=" + body, nil
}

// Returns a mutation that can be applied to a `Store`. The mutation will cause
// the file at `path` to be deleted iff the CAS token of that file matches
// `cas` at the time of application.
//
// If `path` is not valid, returns a `BadPathError`.
func EncodeDel(path string, cas string) (mutation string, err os.Error) {
	if err := checkPath(path); err != nil {
		return
	}
	return cas + ":" + path, nil
}

// MustEncodeSet is like EncodeSet but panics if the mutation cannot be
// encoded. It simplifies safe initialization of global variables holding
// mutations.
func MustEncodeSet(path, body, cas string) (mutation string) {
	m, err := EncodeSet(path, body, cas)
	if err != nil {
		panic(err)
	}
	return m
}

// MustEncodeDel is like EncodeDel but panics if the mutation cannot be
// encoded. It simplifies safe initialization of global variables holding
// mutations.
func MustEncodeDel(path, cas string) (mutation string) {
	m, err := EncodeDel(path, cas)
	if err != nil {
		panic(err)
	}
	return m
}

func decode(mutation string) (path, v, cas string, keep bool, err os.Error) {
	cm := strings.Split(mutation, ":", 2)

	if len(cm) != 2 {
		err = ErrBadMutation
		return
	}

	kv := strings.Split(cm[1], "=", 2)

	if err = checkPath(kv[0]); err != nil {
		return
	}

	switch len(kv) {
	case 1:
		return kv[0], "", cm[0], false, nil
	case 2:
		return kv[0], kv[1], cm[0], true, nil
	}
	panic("unreachable")
}

func (st *Store) notify(e Event) {
	nwatches := make([]watch, len(st.watches))

	i := 0
	for _, w := range st.watches {
		if closed(w.out) {
			continue
		}

		nwatches[i] = w
		i++

		if w.re.MatchString(e.Path) {
			w.out <- e
		}
	}

	st.watches = nwatches[0:i]
}

func (st *Store) closeWatches() {
	for _, w := range st.watches {
		close(w.out)
	}
}

func (st *Store) process(ops <-chan Op, seqns chan<-uint64) {
	logger := util.NewLogger("store")
	defer st.closeWatches()

	var head uint64

	for {
		ver, values := st.state.ver, st.state.root

		// Take any incoming requests and queue them up.
		select {
		case a := <-ops:
			if closed(ops) {
				return
			}

			if a.Seqn > ver {
				st.todo[a.Seqn] = a
			}
		case w := <-st.watchCh:
			st.watches = append(st.watches, w)
		case seqn := <-st.cleanCh:
			for ; head <= seqn; head++ {
				st.log[head] = Event{}, false
			}
		case seqns <- ver:
			// nothing to do here
		}

		// If we have any mutations that can be applied, do them.
		for t, ok := st.todo[ver+1]; ok; t, ok = st.todo[ver+1] {
			var ev Event
			values, ev = values.apply(t.Seqn, t.Mut)
			logger.Printf("apply %s %v %v %v %v %v", ev.Desc(), ev.Seqn, ev.Path, ev.Body, ev.Cas, ev.Err)
			st.state = &state{ev.Seqn, values}
			st.log[t.Seqn] = ev
			st.notify(ev)
			for ver < ev.Seqn {
				ver++
				st.todo[ver] = Op{}, false
			}
		}
	}
}

// Gets the value stored at `path`, if any.
//
// If no value is stored at `path`, `cas` will be `Missing` and `value` will be
// nil.
//
// if `path` is a directory, `cas` will be `Dir` and `value` will be a list of
// entries.
//
// Otherwise, `cas` is the CAS token and `value[0]` is the body.
func (st *Store) Get(path string) (value []string, cas string) {
	// WARNING: Be sure to read the pointer value of st.state only once. If you
	// need multiple accesses, copy the pointer first.
	return st.state.root.Get(path)
}

// Encodes the entire storage state, including the current sequence number, as
// a mutation. This mutation can be applied to an empty store to reproduce the
// state of `s`.
//
// Returns the sequence number of the snapshot and the mutation itself.
//
// A snapshot must be applied at sequence number 1. Once a snapshot has been
// applied, the store's sequence number will be set to `seqn`.
//
// Note that applying a snapshot does not send notifications.
func (st *Store) Snapshot() (seqn uint64, mutation string) {
	w := new(bytes.Buffer)

	// WARNING: Be sure to read the pointer value of st.state only once. If you
	// need multiple accesses, copy the pointer first.
	ss := st.state

	err := gob.NewEncoder(w).Encode(ss.ver)
	if err != nil {
		panic(err)
	}

	err = gob.NewEncoder(w).Encode(ss.root)
	if err != nil {
		panic(err)
	}

	return ss.ver, w.String()
}

// Subscribes `ch` to receive notifications when mutations are applied to paths
// in the store. One event will be sent for each mutation iff the event's path
// matches `pattern`, a Unix-style glob pattern.
//
// Glob notation:
//  - "?" matches a single char in a single path component
//  - "*" matches zero or more chars in a single path component
//  - "**" matches zero or more chars in zero or more components
//  - any other sequence matches itself
//
// Notifications will not be sent for changes made as the result of applying a
// snapshot.
func (st *Store) WatchOn(pattern string, ch chan Event) {
	re, _ := compileGlob(pattern)
	st.watchCh <- watch{out: ch, re: re}
}

func (st *Store) Watch(pattern string) <-chan Event {
	ch := make(chan Event)
	st.WatchOn(pattern, ch)
	return ch
}

// Returns a read-only chan that will receive a single event representing the
// change made at position `seqn`.
//
// If `seqn` was applied before the call to `Wait`, a dummy event will be
// sent with its `Err` set to `ErrTooLate`.
func (st *Store) Wait(seqn uint64) <-chan Event {
	ch, all := make(chan Event, 1), st.Watch("**")

	// Reading shared state. This must happen after the call to st.Watch.
	if st.state.ver >= seqn {
		close(all)
		if ev, ok := st.log[seqn]; ok {
			ch <- ev
		} else {
			ch <- Event{Seqn: seqn, Err: ErrTooLate}
		}
	}

	go func() {
		for e := range all {
			if e.Seqn == seqn {
				close(all)
				ch <- e
			}
		}
	}()
	return ch
}

// Ensures that the application of mutation at `seqn` happens before the call
// to `Sync` returns.
//
// See http://golang.org/doc/go_mem.html for the meaning of "happens before" in
// Go.
func (st *Store) Sync(seqn uint64) {
	<-st.Wait(seqn)
}

// Returns an immutable copy of `st` in which `path` exists as a regular file
// (not a dir). Waits for `path` to be set, if necessary.
func (st *Store) SyncPath(path string) Getter {
	evs := st.Watch(path)
	defer func() {
		close(evs)
		<-evs
	}()

	g := st.state.root // TODO make this use a public method
	_, cas := g.Get(path)
	if cas != Dir && cas != Missing {
		return g
	}

	for ev := range evs {
		if ev.IsSet() {
			return ev
		}
	}

	panic("unreachable")
}

// Lists out the contents of `path` as dummy events on `ch`. Also subscribes
// `ch` to receive future events for changes to `path+"/*"`.
//
// The subscription is made before listing the directory entries. This
// guarantees no entry will be missed, but one or more of the dummy events may
// duplicate a true event.
func (st *Store) GetDirAndWatch(path string, ch chan Event) {
	st.WatchOn(path+"/*", ch)
	go func() {
		for _, ent := range GetDir(st, path) {
			p := path + "/" + ent
			v, cas := st.Get(p)
			if cas != Missing && cas != Dir {
				ch <- Event{0, p, v[0], cas, "", nil, nil}
			}
		}
	}()
}

func (st *Store) Clean(seqn uint64) {
	st.cleanCh <- seqn
}
<commit_msg>store: remove unused field<commit_after>package store

import (
	"bytes"
	"gob"
	"doozer/util"
	"os"
	"regexp"
	"strings"
)

const (
	Clobber = ""
	Missing = "0"
	Dir     = "dir"
)

// TODO revisit this when package regexp is more complete (e.g. do Unicode)
const (
	charPat = `([a-zA-Z0-9.]|-)`
	partPat = "/" + charPat + "+"
	pathPat = "^/$|^(" + partPat + ")+$"
)

var pathRe = regexp.MustCompile(pathPat)

var (
	ErrBadMutation = os.NewError("bad mutation")
	ErrBadSnapshot = os.NewError("bad snapshot")
	ErrTooLate     = os.NewError("too late")
	ErrCasMismatch = os.NewError("cas mismatch")
)

type BadPathError struct {
	Path string
}

func (e *BadPathError) String() string {
	return "bad path: " + e.Path
}

// Applies mutations sent on Ops in sequence according to field Seqn. Any
// errors that occur will be written to ErrorPath. Duplicate operations at a
// given position are sliently ignored.
type Store struct {
	Ops     chan<- Op
	Seqns   <-chan uint64
	watchCh chan watch
	watches []watch
	todo    map[uint64]Op
	state   *state
	log     map[uint64]Event
	cleanCh chan uint64
}

// Represents an operation to apply to the store at position Seqn.
//
// If Mut is a snapshot, notifications will not be sent.
//
// If Mut is Nop, no change will be made, but a dummy event will still be sent.
type Op struct {
	Seqn uint64
	Mut  string
}

type state struct {
	ver  uint64
	root node
}

type watch struct {
	out chan Event
	re  *regexp.Regexp
}

// Creates a new, empty data store. Mutations will be applied in order,
// starting at number 1 (number 0 can be thought of as the creation of the
// store).
func New() *Store {
	ops   := make(chan Op)
	seqns := make(chan uint64)

	st := &Store{
		Ops:     ops,
		Seqns:   seqns,
		watchCh: make(chan watch),
		todo:    make(map[uint64]Op),
		watches: []watch{},
		state:   &state{0, emptyDir},
		log:     make(map[uint64]Event),
		cleanCh: make(chan uint64),
	}

	go st.process(ops, seqns)
	return st
}

func split(path string) []string {
	if path == "/" {
		return []string{}
	}
	return strings.Split(path[1:], "/", -1)
}

func join(parts []string) string {
	return "/" + strings.Join(parts, "/")
}

func checkPath(k string) os.Error {
	if !pathRe.MatchString(k) {
		return &BadPathError{k}
	}
	return nil
}

// Returns a mutation that can be applied to a `Store`. The mutation will set
// the contents of the file at `path` to `body` iff the CAS token of that file
// matches `cas` at the time of application.
//
// If `path` is not valid, returns a `BadPathError`.
func EncodeSet(path, body string, cas string) (mutation string, err os.Error) {
	if err = checkPath(path); err != nil {
		return
	}
	return cas + ":" + path + "=" + body, nil
}

// Returns a mutation that can be applied to a `Store`. The mutation will cause
// the file at `path` to be deleted iff the CAS token of that file matches
// `cas` at the time of application.
//
// If `path` is not valid, returns a `BadPathError`.
func EncodeDel(path string, cas string) (mutation string, err os.Error) {
	if err := checkPath(path); err != nil {
		return
	}
	return cas + ":" + path, nil
}

// MustEncodeSet is like EncodeSet but panics if the mutation cannot be
// encoded. It simplifies safe initialization of global variables holding
// mutations.
func MustEncodeSet(path, body, cas string) (mutation string) {
	m, err := EncodeSet(path, body, cas)
	if err != nil {
		panic(err)
	}
	return m
}

// MustEncodeDel is like EncodeDel but panics if the mutation cannot be
// encoded. It simplifies safe initialization of global variables holding
// mutations.
func MustEncodeDel(path, cas string) (mutation string) {
	m, err := EncodeDel(path, cas)
	if err != nil {
		panic(err)
	}
	return m
}

func decode(mutation string) (path, v, cas string, keep bool, err os.Error) {
	cm := strings.Split(mutation, ":", 2)

	if len(cm) != 2 {
		err = ErrBadMutation
		return
	}

	kv := strings.Split(cm[1], "=", 2)

	if err = checkPath(kv[0]); err != nil {
		return
	}

	switch len(kv) {
	case 1:
		return kv[0], "", cm[0], false, nil
	case 2:
		return kv[0], kv[1], cm[0], true, nil
	}
	panic("unreachable")
}

func (st *Store) notify(e Event) {
	nwatches := make([]watch, len(st.watches))

	i := 0
	for _, w := range st.watches {
		if closed(w.out) {
			continue
		}

		nwatches[i] = w
		i++

		if w.re.MatchString(e.Path) {
			w.out <- e
		}
	}

	st.watches = nwatches[0:i]
}

func (st *Store) closeWatches() {
	for _, w := range st.watches {
		close(w.out)
	}
}

func (st *Store) process(ops <-chan Op, seqns chan<-uint64) {
	logger := util.NewLogger("store")
	defer st.closeWatches()

	var head uint64

	for {
		ver, values := st.state.ver, st.state.root

		// Take any incoming requests and queue them up.
		select {
		case a := <-ops:
			if closed(ops) {
				return
			}

			if a.Seqn > ver {
				st.todo[a.Seqn] = a
			}
		case w := <-st.watchCh:
			st.watches = append(st.watches, w)
		case seqn := <-st.cleanCh:
			for ; head <= seqn; head++ {
				st.log[head] = Event{}, false
			}
		case seqns <- ver:
			// nothing to do here
		}

		// If we have any mutations that can be applied, do them.
		for t, ok := st.todo[ver+1]; ok; t, ok = st.todo[ver+1] {
			var ev Event
			values, ev = values.apply(t.Seqn, t.Mut)
			logger.Printf("apply %s %v %v %v %v %v", ev.Desc(), ev.Seqn, ev.Path, ev.Body, ev.Cas, ev.Err)
			st.state = &state{ev.Seqn, values}
			st.log[t.Seqn] = ev
			st.notify(ev)
			for ver < ev.Seqn {
				ver++
				st.todo[ver] = Op{}, false
			}
		}
	}
}

// Gets the value stored at `path`, if any.
//
// If no value is stored at `path`, `cas` will be `Missing` and `value` will be
// nil.
//
// if `path` is a directory, `cas` will be `Dir` and `value` will be a list of
// entries.
//
// Otherwise, `cas` is the CAS token and `value[0]` is the body.
func (st *Store) Get(path string) (value []string, cas string) {
	// WARNING: Be sure to read the pointer value of st.state only once. If you
	// need multiple accesses, copy the pointer first.
	return st.state.root.Get(path)
}

// Encodes the entire storage state, including the current sequence number, as
// a mutation. This mutation can be applied to an empty store to reproduce the
// state of `s`.
//
// Returns the sequence number of the snapshot and the mutation itself.
//
// A snapshot must be applied at sequence number 1. Once a snapshot has been
// applied, the store's sequence number will be set to `seqn`.
//
// Note that applying a snapshot does not send notifications.
func (st *Store) Snapshot() (seqn uint64, mutation string) {
	w := new(bytes.Buffer)

	// WARNING: Be sure to read the pointer value of st.state only once. If you
	// need multiple accesses, copy the pointer first.
	ss := st.state

	err := gob.NewEncoder(w).Encode(ss.ver)
	if err != nil {
		panic(err)
	}

	err = gob.NewEncoder(w).Encode(ss.root)
	if err != nil {
		panic(err)
	}

	return ss.ver, w.String()
}

// Subscribes `ch` to receive notifications when mutations are applied to paths
// in the store. One event will be sent for each mutation iff the event's path
// matches `pattern`, a Unix-style glob pattern.
//
// Glob notation:
//  - "?" matches a single char in a single path component
//  - "*" matches zero or more chars in a single path component
//  - "**" matches zero or more chars in zero or more components
//  - any other sequence matches itself
//
// Notifications will not be sent for changes made as the result of applying a
// snapshot.
func (st *Store) WatchOn(pattern string, ch chan Event) {
	re, _ := compileGlob(pattern)
	st.watchCh <- watch{out: ch, re: re}
}

func (st *Store) Watch(pattern string) <-chan Event {
	ch := make(chan Event)
	st.WatchOn(pattern, ch)
	return ch
}

// Returns a read-only chan that will receive a single event representing the
// change made at position `seqn`.
//
// If `seqn` was applied before the call to `Wait`, a dummy event will be
// sent with its `Err` set to `ErrTooLate`.
func (st *Store) Wait(seqn uint64) <-chan Event {
	ch, all := make(chan Event, 1), st.Watch("**")

	// Reading shared state. This must happen after the call to st.Watch.
	if st.state.ver >= seqn {
		close(all)
		if ev, ok := st.log[seqn]; ok {
			ch <- ev
		} else {
			ch <- Event{Seqn: seqn, Err: ErrTooLate}
		}
	}

	go func() {
		for e := range all {
			if e.Seqn == seqn {
				close(all)
				ch <- e
			}
		}
	}()
	return ch
}

// Ensures that the application of mutation at `seqn` happens before the call
// to `Sync` returns.
//
// See http://golang.org/doc/go_mem.html for the meaning of "happens before" in
// Go.
func (st *Store) Sync(seqn uint64) {
	<-st.Wait(seqn)
}

// Returns an immutable copy of `st` in which `path` exists as a regular file
// (not a dir). Waits for `path` to be set, if necessary.
func (st *Store) SyncPath(path string) Getter {
	evs := st.Watch(path)
	defer func() {
		close(evs)
		<-evs
	}()

	g := st.state.root // TODO make this use a public method
	_, cas := g.Get(path)
	if cas != Dir && cas != Missing {
		return g
	}

	for ev := range evs {
		if ev.IsSet() {
			return ev
		}
	}

	panic("unreachable")
}

// Lists out the contents of `path` as dummy events on `ch`. Also subscribes
// `ch` to receive future events for changes to `path+"/*"`.
//
// The subscription is made before listing the directory entries. This
// guarantees no entry will be missed, but one or more of the dummy events may
// duplicate a true event.
func (st *Store) GetDirAndWatch(path string, ch chan Event) {
	st.WatchOn(path+"/*", ch)
	go func() {
		for _, ent := range GetDir(st, path) {
			p := path + "/" + ent
			v, cas := st.Get(p)
			if cas != Missing && cas != Dir {
				ch <- Event{0, p, v[0], cas, "", nil, nil}
			}
		}
	}()
}

func (st *Store) Clean(seqn uint64) {
	st.cleanCh <- seqn
}
<|endoftext|>


<commit_before>package store

import (
	"net/url"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/garyburd/redigo/redis"
)

const (
	// InMemory storage
	InMemory = "memory"
	// Redis storage
	Redis = "redis"
	// None Nullable storage
	None = "none"
)

// Store is the common interface for datastores.
type Store interface {
	Exists(key string) (bool, error)
	Get(key string) (string, error)
	Remove(key string) error
	Set(key string, value string, expire int64) error
}

// Subscriber holds the basic methods to subscribe to a topic
type Subscriber interface {
	Subscribe(topic string) *Subscription
}

// Publisher holds the basic methods to publish a message
type Publisher interface {
	Publish(topic string, data []byte) error
}

// Options are options for store.
type Options struct {
	// Prefix is the prefix to use for the key.
	Prefix string

	// MaxRetry is the maximum number of retry under race conditions.
	MaxRetry int

	// CleanUpInterval is the interval for cleanup.
	CleanUpInterval time.Duration
}

// Message represents the message that comes
// form an update
type Message []byte

// Subscription holds a message channel
type Subscription struct {
	Message chan Message
}

// NewSubscription creates a new instance of Subscription
func NewSubscription() *Subscription {
	return &Subscription{
		Message: make(chan Message),
	}
}

// Build creates a new storage based on the provided DSN
func Build(dsn string) (Store, error) {
	url, err := url.Parse(dsn)
	if nil != err {
		return nil, err
	}
	log.WithField("type", url.Scheme).Debug("Initializing storage")

	switch url.Scheme {
	case InMemory:
		return NewInMemoryStore(), nil
	case Redis:
		// Create a Redis pool.
		pool := &redis.Pool{
			MaxIdle:     3,
			IdleTimeout: 240 * time.Second,
			Dial:        func() (redis.Conn, error) { return redis.DialURL(dsn) },
		}

		log.Debugf("Trying to connect to redis pool: %s", dsn)
		return NewRedisStore(pool)
	}

	return nil, ErrUnknownStorage
}
<commit_msg>Fixed typo<commit_after>package store

import (
	"net/url"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/garyburd/redigo/redis"
)

const (
	// InMemory storage
	InMemory = "memory"
	// Redis storage
	Redis = "redis"
	// None Nullable storage
	None = "none"
)

// Store is the common interface for datastores.
type Store interface {
	Exists(key string) (bool, error)
	Get(key string) (string, error)
	Remove(key string) error
	Set(key string, value string, expire int64) error
}

// Subscriber holds the basic methods to subscribe to a topic
type Subscriber interface {
	Subscribe(topic string) *Subscription
}

// Publisher holds the basic methods to publish a message
type Publisher interface {
	Publish(topic string, data []byte) error
}

// Options are options for store.
type Options struct {
	// Prefix is the prefix to use for the key.
	Prefix string

	// MaxRetry is the maximum number of retry under race conditions.
	MaxRetry int

	// CleanUpInterval is the interval for cleanup.
	CleanUpInterval time.Duration
}

// Message represents the message that comes
// from an update
type Message []byte

// Subscription holds a message channel
type Subscription struct {
	Message chan Message
}

// NewSubscription creates a new instance of Subscription
func NewSubscription() *Subscription {
	return &Subscription{
		Message: make(chan Message),
	}
}

// Build creates a new storage based on the provided DSN
func Build(dsn string) (Store, error) {
	url, err := url.Parse(dsn)
	if nil != err {
		return nil, err
	}
	log.WithField("type", url.Scheme).Debug("Initializing storage")

	switch url.Scheme {
	case InMemory:
		return NewInMemoryStore(), nil
	case Redis:
		// Create a Redis pool.
		pool := &redis.Pool{
			MaxIdle:     3,
			IdleTimeout: 240 * time.Second,
			Dial:        func() (redis.Conn, error) { return redis.DialURL(dsn) },
		}

		log.Debugf("Trying to connect to redis pool: %s", dsn)
		return NewRedisStore(pool)
	}

	return nil, ErrUnknownStorage
}
<|endoftext|>


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

package prjcfg

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"sort"
	"strings"
	"time"

	"google.golang.org/protobuf/proto"

	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/gae/service/datastore"

	cfgpb "go.chromium.org/luci/cv/api/config/v2"
)

const projectConfigKind string = "ProjectConfig"

// schemaVersion is the current DS schema version.
//
// Bump it to force-update Project configs and their Config Groups after the
// next deployment.
const schemaVersion = 1

// ProjectConfig is the root entity that keeps track of the latest version
// info of the CV config for a LUCI Project. It only contains high-level
// metadata about the config. The actual content of config is stored in the
// `ConfigGroup` entities which can be looked up by constructing IDs using
// `ConfigGroupNames` field.
type ProjectConfig struct {
	_kind string `gae:"$kind,ProjectConfig"`
	// Project is the name of this LUCI Project.
	Project string `gae:"$id"`
	// SchemaVersion is the version of the schema.
	//
	// It is used to force-update old entities to newest format.
	// See schemaVersion const.
	SchemaVersion int `gae:",noindex"`
	// Enabled indicates whether CV is enabled for this LUCI Project.
	//
	// Project is disabled if it is de-registered in LUCI Config or it no longer
	// has CV config file.
	Enabled bool
	// UpdateTime is the timestamp when this ProjectConfig was last updated.
	UpdateTime time.Time `gae:",noindex"`
	// EVersion is the latest version number of this ProjectConfig.
	//
	// It increments by 1 every time a new config change is imported to CV for
	// this LUCI Project.
	EVersion int64 `gae:",noindex"`
	// Hash is a string computed from the content of latest imported CV Config
	// using `computeHash()`.
	Hash string `gae:",noindex"`
	// ExternalHash is the hash string of this CV config in the external source
	// of truth (currently, LUCI Config). Used to quickly decided whether the
	// Config has been updated without fetching the full content.
	ExternalHash string `gae:",noindex"`
	// ConfigGroupNames are the names of all ConfigGroups in the current version
	// of CV Config.
	ConfigGroupNames []string `gae:",noindex"`
}

// computeHash computes the hash string of given CV Config and prefixed with
// hash algorithm string. (e.g. sha256:deadbeefdeadbeef)
//
// The hash string is an hex-encoded string of the first 8 bytes (i.e. 16
// char in length) of sha256(deterministically binary serialized Config proto).
// Note that, deterministic marshalling does NOT guarantee the same output
// for the equal proto message  across different language or event builds.
// Therefore, in worst case scenario, when a newer version of proto lib is
// deployed, CV may re-ingest functionally equivalent config.
// See: https://godoc.org/google.golang.org/protobuf/proto#MarshalOptions
func computeHash(cfg *cfgpb.Config) string {
	b, err := proto.MarshalOptions{Deterministic: true}.Marshal(cfg)
	if err != nil {
		panic(fmt.Sprintf("failed to marshal config: %s", err))
	}
	sha := sha256.New()
	sha.Write(b)
	return fmt.Sprintf("sha256:%s", hex.EncodeToString(sha.Sum(nil)[:8]))
}

// GetAllProjectIDs returns the names of all projects available in datastore.
func GetAllProjectIDs(ctx context.Context, enabledOnly bool) ([]string, error) {
	var projects []*ProjectConfig
	query := datastore.NewQuery(projectConfigKind).Project("Enabled")
	if err := datastore.GetAll(ctx, query, &projects); err != nil {
		return nil, errors.Annotate(err, "failed to query all projects").Tag(transient.Tag).Err()
	}
	ret := make([]string, 0, len(projects))
	for _, p := range projects {
		if enabledOnly && !p.Enabled {
			continue
		}
		ret = append(ret, p.Project)
	}
	sort.Strings(ret)
	return ret, nil
}

// ConfigHashInfo stores high-level info about a ProjectConfig `Hash`.
//
// It is primarily used for cleanup purpose to decide which `Hash` and
// its corresponding `ConfigGroup`s can be safely deleted.
type ConfigHashInfo struct {
	_kind string `gae:"$kind,ProjectConfigHashInfo"`
	// Hash is the `Hash` of a `ProjectConfig` that CV has imported.
	Hash    string         `gae:"$id"`
	Project *datastore.Key `gae:"$parent"`
	// SchemaVersion is the version of the schema.
	//
	// It is used to force-update old entities to newest format.
	// See schemaVersion const.
	SchemaVersion int `gae:",noindex"`
	// GitRevision is the git revision (commit hash) of the imported config.
	GitRevision string `gae:",noindex"`
	// ProjectEVersion is largest version of ProjectConfig that this `Hash`
	// maps to.
	//
	// It is possible for a ConfigHash maps to multiple EVersions (e.g. a CV
	// Config change is landed then reverted which results in two new EVersions
	// but only one new Hash). Only the largest EVersion matters when cleanup
	// job runs (i.e. CV will keep the last 5 EVersions).
	ProjectEVersion int64 `gae:",noindex"`
	// UpdateTime is the timestamp when this ConfigHashInfo was last updated.
	UpdateTime time.Time `gae:",noindex"`
	// ConfigGroupNames are the names of all ConfigGroups with this `Hash`.
	ConfigGroupNames []string `gae:",noindex"`
}

// ConfigGroupID is the ID for ConfigGroup Entity.
//
// It is in the format of "hash/name" where
//   - `hash` is the `Hash` field in the containing `ProjectConfig`.
//   - `name` is the value of `ConfigGroup.Name`.
type ConfigGroupID string

// Returns Hash of the corresponding project config.
func (c ConfigGroupID) Hash() string {
	s := string(c)
	if i := strings.IndexRune(s, '/'); i >= 0 {
		return s[:i]
	}
	panic(fmt.Errorf("invalid ConfigGroupID %q", c))
}

// Returns name component only.
func (c ConfigGroupID) Name() string {
	s := string(c)
	if i := strings.IndexRune(s, '/'); i >= 0 {
		return s[i+1:]
	}
	panic(fmt.Errorf("invalid ConfigGroupID %q", c))
}

func MakeConfigGroupID(hash, name string) ConfigGroupID {
	if name == "" {
		panic(fmt.Errorf("name must be given"))
	}
	return ConfigGroupID(fmt.Sprintf("%s/%s", hash, name))
}

// ConfigGroup is an entity that represents a ConfigGroup defined in CV config.
type ConfigGroup struct {
	_kind   string         `gae:"$kind,ProjectConfigGroup"`
	Project *datastore.Key `gae:"$parent"`
	ID      ConfigGroupID  `gae:"$id"`
	// SchemaVersion is the version of the schema.
	//
	// It is used to force-update old entities to newest format.
	// See schemaVersion const.
	SchemaVersion int `gae:",noindex"`
	// DrainingStartTime represents `draining_start_time` field in the CV config.
	//
	// Note that this is a project-level field. Therefore, all ConfigGroups in a
	// single version of config should have the same value.
	DrainingStartTime string `gae:",noindex"`
	// SubmitOptions represents `submit_options` field in the CV config.
	//
	// Note that this is currently a project-level field. Therefore, all
	// ConfigGroups in a single version of Config should have the same value.
	SubmitOptions *cfgpb.SubmitOptions
	// Content represents a `pb.ConfigGroup` proto message defined in the CV
	// config
	Content *cfgpb.ConfigGroup
	// CQStatusHost is the URL of the CQ status app. Optional.
	//
	// Deprecated.
	// TODO(crbug/1233963): remove this field.
	CQStatusHost string `gae:",noindex"`
}

// ProjectString returns LUCI Project as a string.
func (c *ConfigGroup) ProjectString() string {
	return c.Project.StringID()
}

// putConfigGroups puts the ConfigGroups in the given CV config to datastore.
//
// It checks for existence of each ConfigGroup first to avoid unnecessary puts.
// It is also idempotent so it is safe to retry and can be called out of a
// transactional context.
func putConfigGroups(ctx context.Context, cfg *cfgpb.Config, project, hash string) error {
	cgLen := len(cfg.GetConfigGroups())
	if cgLen == 0 {
		return nil
	}

	// Check if there are any existing entities with the current schema version
	// such that we can skip updating them.
	projKey := datastore.MakeKey(ctx, projectConfigKind, project)
	entities := make([]*ConfigGroup, cgLen)
	for i, cg := range cfg.GetConfigGroups() {
		entities[i] = &ConfigGroup{
			ID:      MakeConfigGroupID(hash, cg.GetName()),
			Project: projKey,
		}
	}
	err := datastore.Get(ctx, entities)
	errs, ok := err.(errors.MultiError)
	switch {
	case err != nil && !ok:
		return errors.Annotate(err, "failed to check the existence of ConfigGroups").Tag(transient.Tag).Err()
	case err == nil:
		errs = make(errors.MultiError, cgLen)
	}
	toPut := entities[:0] // re-use the slice
	for i, err := range errs {
		ent := entities[i]
		switch {
		case err == datastore.ErrNoSuchEntity:
			// proceed to put below.
		case err != nil:
			return errors.Annotate(err, "failed to check the existence of one of ConfigGroups").Tag(transient.Tag).Err()
		case ent.SchemaVersion != schemaVersion:
			// Intentionally using != here s.t. rollbacks result in downgrading of the
			// schema. Given that project configs are checked and potentially updated
			// every ~1 minute, this if OK.
		default:
			continue // up to date
		}
		ent.SchemaVersion = schemaVersion
		ent.DrainingStartTime = cfg.GetDrainingStartTime()
		ent.SubmitOptions = cfg.GetSubmitOptions()
		ent.Content = cfg.GetConfigGroups()[i]
		ent.CQStatusHost = cfg.GetCqStatusHost()
		toPut = append(toPut, ent)
	}

	if err := datastore.Put(ctx, toPut); err != nil {
		return errors.Annotate(err, "failed to put ConfigGroups").Tag(transient.Tag).Err()
	}
	return nil
}
<commit_msg>[cv] cosmetic changes<commit_after>// Copyright 2020 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package prjcfg

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"sort"
	"strings"
	"time"

	"google.golang.org/protobuf/proto"

	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/gae/service/datastore"

	cfgpb "go.chromium.org/luci/cv/api/config/v2"
)

const projectConfigKind string = "ProjectConfig"

// schemaVersion is the current DS schema version.
//
// Bump it to force-update Project configs and their Config Groups after the
// next deployment.
const schemaVersion = 1

// ProjectConfig is the root entity that keeps track of the latest version
// info of the CV config for a LUCI Project. It only contains high-level
// metadata about the config. The actual content of config is stored in the
// `ConfigGroup` entities which can be looked up by constructing IDs using
// `ConfigGroupNames` field.
type ProjectConfig struct {
	_kind string `gae:"$kind,ProjectConfig"`
	// Project is the name of this LUCI Project.
	Project string `gae:"$id"`
	// SchemaVersion is the version of the schema.
	//
	// It is used to force-update old entities to newest format.
	// See schemaVersion const.
	SchemaVersion int `gae:",noindex"`
	// Enabled indicates whether CV is enabled for this LUCI Project.
	//
	// Project is disabled if it is de-registered in LUCI Config or it no longer
	// has CV config file.
	Enabled bool
	// UpdateTime is the timestamp when this ProjectConfig was last updated.
	UpdateTime time.Time `gae:",noindex"`
	// EVersion is the latest version number of this ProjectConfig.
	//
	// It increments by 1 every time a new config change is imported to CV for
	// this LUCI Project.
	EVersion int64 `gae:",noindex"`
	// Hash is a string computed from the content of latest imported CV Config
	// using `computeHash()`.
	Hash string `gae:",noindex"`
	// ExternalHash is the hash string of this CV config in the external source
	// of truth (currently, LUCI Config). Used to quickly decided whether the
	// Config has been updated without fetching the full content.
	ExternalHash string `gae:",noindex"`
	// ConfigGroupNames are the names of all ConfigGroups in the current version
	// of CV Config.
	ConfigGroupNames []string `gae:",noindex"`
}

// computeHash computes the hash string of given CV Config and prefixed with
// hash algorithm string. (e.g. sha256:deadbeefdeadbeef)
//
// The hash string is an hex-encoded string of the first 8 bytes (i.e. 16
// char in length) of sha256(deterministically binary serialized Config proto).
// Note that, deterministic marshalling does NOT guarantee the same output
// for the equal proto message  across different language or event builds.
// Therefore, in worst case scenario, when a newer version of proto lib is
// deployed, CV may re-ingest functionally equivalent config.
// See: https://godoc.org/google.golang.org/protobuf/proto#MarshalOptions
func computeHash(cfg *cfgpb.Config) string {
	b, err := proto.MarshalOptions{Deterministic: true}.Marshal(cfg)
	if err != nil {
		panic(fmt.Sprintf("failed to marshal config: %s", err))
	}
	sha := sha256.New()
	sha.Write(b)
	return fmt.Sprintf("sha256:%s", hex.EncodeToString(sha.Sum(nil)[:8]))
}

// GetAllProjectIDs returns the names of all projects available in datastore.
func GetAllProjectIDs(ctx context.Context, enabledOnly bool) ([]string, error) {
	var projects []*ProjectConfig
	query := datastore.NewQuery(projectConfigKind).Project("Enabled")
	if err := datastore.GetAll(ctx, query, &projects); err != nil {
		return nil, errors.Annotate(err, "failed to query all projects").Tag(transient.Tag).Err()
	}
	ret := make([]string, 0, len(projects))
	for _, p := range projects {
		if enabledOnly && !p.Enabled {
			continue
		}
		ret = append(ret, p.Project)
	}
	sort.Strings(ret)
	return ret, nil
}

// ConfigHashInfo stores high-level info about a ProjectConfig `Hash`.
//
// It is primarily used for cleanup purpose to decide which `Hash` and
// its corresponding `ConfigGroup`s can be safely deleted.
type ConfigHashInfo struct {
	_kind string `gae:"$kind,ProjectConfigHashInfo"`
	// Hash is the `Hash` of a `ProjectConfig` that CV has imported.
	Hash    string         `gae:"$id"`
	Project *datastore.Key `gae:"$parent"`
	// SchemaVersion is the version of the schema.
	//
	// It is used to force-update old entities to newest format.
	// See schemaVersion const.
	SchemaVersion int `gae:",noindex"`
	// GitRevision is the git revision (commit hash) of the imported config.
	GitRevision string `gae:",noindex"`
	// ProjectEVersion is largest version of ProjectConfig that this `Hash`
	// maps to.
	//
	// It is possible for a ConfigHash maps to multiple EVersions (e.g. a CV
	// Config change is landed then reverted which results in two new EVersions
	// but only one new Hash). Only the largest EVersion matters when cleanup
	// job runs (i.e. CV will keep the last 5 EVersions).
	ProjectEVersion int64 `gae:",noindex"`
	// UpdateTime is the timestamp when this ConfigHashInfo was last updated.
	UpdateTime time.Time `gae:",noindex"`
	// ConfigGroupNames are the names of all ConfigGroups with this `Hash`.
	ConfigGroupNames []string `gae:",noindex"`
}

// ConfigGroupID is the ID for ConfigGroup Entity.
//
// It is in the format of "hash/name" where
//   - `hash` is the `Hash` field in the containing `ProjectConfig`.
//   - `name` is the value of `ConfigGroup.Name`.
type ConfigGroupID string

// Hash returns Hash of the corresponding project config.
func (c ConfigGroupID) Hash() string {
	s := string(c)
	if i := strings.IndexRune(s, '/'); i >= 0 {
		return s[:i]
	}
	panic(fmt.Errorf("invalid ConfigGroupID %q", c))
}

// Name returns name component only.
func (c ConfigGroupID) Name() string {
	s := string(c)
	if i := strings.IndexRune(s, '/'); i >= 0 {
		return s[i+1:]
	}
	panic(fmt.Errorf("invalid ConfigGroupID %q", c))
}

// MakeConfigGroupID creates ConfigGroupID.
func MakeConfigGroupID(hash, name string) ConfigGroupID {
	if name == "" {
		panic(fmt.Errorf("name must be given"))
	}
	return ConfigGroupID(fmt.Sprintf("%s/%s", hash, name))
}

// ConfigGroup is an entity that represents a ConfigGroup defined in CV config.
type ConfigGroup struct {
	_kind   string         `gae:"$kind,ProjectConfigGroup"`
	Project *datastore.Key `gae:"$parent"`
	ID      ConfigGroupID  `gae:"$id"`
	// SchemaVersion is the version of the schema.
	//
	// It is used to force-update old entities to newest format.
	// See schemaVersion const.
	SchemaVersion int `gae:",noindex"`
	// DrainingStartTime represents `draining_start_time` in the CV config.
	//
	// Note that this is a project-level field. Therefore, all ConfigGroups in a
	// single version of config should have the same value.
	DrainingStartTime string `gae:",noindex"`
	// SubmitOptions represents `submit_options` field in the CV config.
	//
	// Note that this is currently a project-level field. Therefore, all
	// ConfigGroups in a single version of Config should have the same value.
	SubmitOptions *cfgpb.SubmitOptions
	// Content represents a `pb.ConfigGroup` proto message defined in the CV
	// config
	Content *cfgpb.ConfigGroup
	// CQStatusHost is the URL of the CQ status app. Optional.
	//
	// Deprecated.
	// TODO(crbug/1233963): remove this field.
	CQStatusHost string `gae:",noindex"`
}

// ProjectString returns LUCI Project as a string.
func (c *ConfigGroup) ProjectString() string {
	return c.Project.StringID()
}

// putConfigGroups puts the ConfigGroups in the given CV config to datastore.
//
// It checks for existence of each ConfigGroup first to avoid unnecessary puts.
// It is also idempotent so it is safe to retry and can be called out of a
// transactional context.
func putConfigGroups(ctx context.Context, cfg *cfgpb.Config, project, hash string) error {
	cgLen := len(cfg.GetConfigGroups())
	if cgLen == 0 {
		return nil
	}

	// Check if there are any existing entities with the current schema version
	// such that we can skip updating them.
	projKey := datastore.MakeKey(ctx, projectConfigKind, project)
	entities := make([]*ConfigGroup, cgLen)
	for i, cg := range cfg.GetConfigGroups() {
		entities[i] = &ConfigGroup{
			ID:      MakeConfigGroupID(hash, cg.GetName()),
			Project: projKey,
		}
	}
	err := datastore.Get(ctx, entities)
	errs, ok := err.(errors.MultiError)
	switch {
	case err != nil && !ok:
		return errors.Annotate(err, "failed to check the existence of ConfigGroups").Tag(transient.Tag).Err()
	case err == nil:
		errs = make(errors.MultiError, cgLen)
	}
	toPut := entities[:0] // re-use the slice
	for i, err := range errs {
		ent := entities[i]
		switch {
		case err == datastore.ErrNoSuchEntity:
			// proceed to put below.
		case err != nil:
			return errors.Annotate(err, "failed to check the existence of one of ConfigGroups").Tag(transient.Tag).Err()
		case ent.SchemaVersion != schemaVersion:
			// Intentionally using != here s.t. rollbacks result in downgrading
			// of the schema. Given that project configs are checked and
			// potentially updated every ~1 minute, this if OK.
		default:
			continue // up to date
		}
		ent.SchemaVersion = schemaVersion
		ent.DrainingStartTime = cfg.GetDrainingStartTime()
		ent.SubmitOptions = cfg.GetSubmitOptions()
		ent.Content = cfg.GetConfigGroups()[i]
		ent.CQStatusHost = cfg.GetCqStatusHost()
		toPut = append(toPut, ent)
	}

	if err := datastore.Put(ctx, toPut); err != nil {
		return errors.Annotate(err, "failed to put ConfigGroups").Tag(transient.Tag).Err()
	}
	return nil
}
<|endoftext|>


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

package heal

import (
	"fmt"
	"launchpad.net/gocheck"
	"log/syslog"
	"net/http"
	"net/http/httptest"
	"os"
	"sync/atomic"
	"time"
)

type CallerSuite struct {
	instId string
	token  string
}

var _ = gocheck.Suite(&CallerSuite{})

func (s *CallerSuite) SetUpSuite(c *gocheck.C) {
	var err error
	log, err = syslog.New(syslog.LOG_INFO, "tsuru-healer")
	c.Assert(err, gocheck.IsNil)
}

func (s *CallerSuite) TestHealersFromResource(c *gocheck.C) {
	os.Setenv("TSURU_TOKEN", "token123")
	defer os.Setenv("TSURU_TOKEN", "")
	reqs := []*http.Request{}
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		reqs = append(reqs, r)
		w.Write([]byte(`{"bootstrap":"/bootstrap"}`))
	}))
	defer ts.Close()
	expected := map[string]*healer{
		"bootstrap": {url: fmt.Sprintf("%s/bootstrap", ts.URL)},
	}
	healers, err := healersFromResource(ts.URL)
	c.Assert(err, gocheck.IsNil)
	c.Assert(healers, gocheck.DeepEquals, expected)
	c.Assert(reqs, gocheck.HasLen, 1)
	c.Assert(reqs[0].Header.Get("Authorization"), gocheck.Equals, "bearer token123")
}

func (s *CallerSuite) TestTsuruHealer(c *gocheck.C) {
	os.Setenv("TSURU_TOKEN", "token123")
	defer os.Setenv("TSURU_TOKEN", "")
	var reqs []*http.Request
	var called bool
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		called = true
		reqs = append(reqs, r)
	}))
	defer ts.Close()
	h := healer{url: ts.URL}
	err := h.heal()
	c.Assert(err, gocheck.IsNil)
	c.Assert(called, gocheck.Equals, true)
	c.Assert(reqs, gocheck.HasLen, 1)
	c.Assert(reqs[0].Header.Get("Authorization"), gocheck.Equals, "bearer token123")
}

func (s *CallerSuite) TestSetAndGetHealers(c *gocheck.C) {
	h := &healer{url: ""}
	setHealers(map[string]*healer{"test-healer": h})
	healers := getHealers()
	healer, ok := healers["test-healer"]
	c.Assert(healer, gocheck.DeepEquals, h)
	c.Assert(ok, gocheck.Equals, true)
}

func (s *S) TestHealTicker(c *gocheck.C) {
	var called int32
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		atomic.StoreInt32(&called, 1)
	}))
	defer ts.Close()
	h := &healer{url: ts.URL}
	setHealers(map[string]*healer{"ticker-healer": h})
	ch := make(chan time.Time)
	ok := make(chan bool)
	go func() {
		HealTicker(ch)
		ok <- true
	}()
	ch <- time.Now()
	time.Sleep(1 * time.Second)
	close(ch)
	<-ok
	c.Assert(atomic.LoadInt32(&called), gocheck.Equals, int32(1))
}

func (s *S) TestRegisterTicker(c *gocheck.C) {
	var called int32
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		atomic.StoreInt32(&called, 1)
	}))
	defer ts.Close()
	ch := make(chan time.Time)
	ok := make(chan bool)
	go func() {
		RegisterHealerTicker(ch, ts.URL)
		ok <- true
	}()
	ch <- time.Now()
	time.Sleep(1 * time.Second)
	close(ch)
	c.Assert(atomic.LoadInt32(&called), gocheck.Equals, int32(1))
}
<commit_msg>heal/caller_test: initializing tsuru log instead of syslog<commit_after>// Copyright 2013 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package heal

import (
	"fmt"
	"github.com/globocom/tsuru/log"
	"launchpad.net/gocheck"
	"net/http"
	"net/http/httptest"
	"os"
	"sync/atomic"
	"time"
)

type CallerSuite struct {
	instId string
	token  string
}

var _ = gocheck.Suite(&CallerSuite{})

func (s *CallerSuite) SetUpSuite(c *gocheck.C) {
	log.Init()
}

func (s *CallerSuite) TestHealersFromResource(c *gocheck.C) {
	os.Setenv("TSURU_TOKEN", "token123")
	defer os.Setenv("TSURU_TOKEN", "")
	reqs := []*http.Request{}
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		reqs = append(reqs, r)
		w.Write([]byte(`{"bootstrap":"/bootstrap"}`))
	}))
	defer ts.Close()
	expected := map[string]*healer{
		"bootstrap": {url: fmt.Sprintf("%s/bootstrap", ts.URL)},
	}
	healers, err := healersFromResource(ts.URL)
	c.Assert(err, gocheck.IsNil)
	c.Assert(healers, gocheck.DeepEquals, expected)
	c.Assert(reqs, gocheck.HasLen, 1)
	c.Assert(reqs[0].Header.Get("Authorization"), gocheck.Equals, "bearer token123")
}

func (s *CallerSuite) TestTsuruHealer(c *gocheck.C) {
	os.Setenv("TSURU_TOKEN", "token123")
	defer os.Setenv("TSURU_TOKEN", "")
	var reqs []*http.Request
	var called bool
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		called = true
		reqs = append(reqs, r)
	}))
	defer ts.Close()
	h := healer{url: ts.URL}
	err := h.heal()
	c.Assert(err, gocheck.IsNil)
	c.Assert(called, gocheck.Equals, true)
	c.Assert(reqs, gocheck.HasLen, 1)
	c.Assert(reqs[0].Header.Get("Authorization"), gocheck.Equals, "bearer token123")
}

func (s *CallerSuite) TestSetAndGetHealers(c *gocheck.C) {
	h := &healer{url: ""}
	setHealers(map[string]*healer{"test-healer": h})
	healers := getHealers()
	healer, ok := healers["test-healer"]
	c.Assert(healer, gocheck.DeepEquals, h)
	c.Assert(ok, gocheck.Equals, true)
}

func (s *S) TestHealTicker(c *gocheck.C) {
	var called int32
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		atomic.StoreInt32(&called, 1)
	}))
	defer ts.Close()
	h := &healer{url: ts.URL}
	setHealers(map[string]*healer{"ticker-healer": h})
	ch := make(chan time.Time)
	ok := make(chan bool)
	go func() {
		HealTicker(ch)
		ok <- true
	}()
	ch <- time.Now()
	time.Sleep(1 * time.Second)
	close(ch)
	<-ok
	c.Assert(atomic.LoadInt32(&called), gocheck.Equals, int32(1))
}

func (s *S) TestRegisterTicker(c *gocheck.C) {
	var called int32
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		atomic.StoreInt32(&called, 1)
	}))
	defer ts.Close()
	ch := make(chan time.Time)
	ok := make(chan bool)
	go func() {
		RegisterHealerTicker(ch, ts.URL)
		ok <- true
	}()
	ch <- time.Now()
	time.Sleep(1 * time.Second)
	close(ch)
	c.Assert(atomic.LoadInt32(&called), gocheck.Equals, int32(1))
}
<|endoftext|>


<commit_before>package exec

import (
	"io"
	"log"
	"os"
	"os/exec"
	"strings"

	"github.com/hashicorp/otto/ui"
)

// Run runs the given command and streams all the output to the
// given UI. It also connects stdin properly so that input works as
// expected.
func Run(uiVal ui.Ui, cmd *exec.Cmd) error {
	out_r, out_w := io.Pipe()
	cmd.Stdin = os.Stdin
	cmd.Stdout = out_w
	cmd.Stderr = out_w

	// Copy output to the UI until we can't.
	uiDone := make(chan struct{})
	go func() {
		defer close(uiDone)
		var buf [1024]byte
		for {
			n, err := out_r.Read(buf[:])
			if n > 0 {
				uiVal.Raw(string(buf[:n]))
			}

			// We just break on any error. io.EOF is not an error and
			// is our true exit case, but any other error we don't really
			// handle here. It probably means something went wrong
			// somewhere else anyways.
			if err != nil {
				break
			}
		}
	}()

	// Run the command
	log.Printf("[DEBUG] execDir: %s", cmd.Dir)
	log.Printf("[DEBUG] exec: %s %s", cmd.Path, strings.Join(cmd.Args[1:], " "))
	err := cmd.Run()

	// Wait for all the output to finish
	out_w.Close()
	<-uiDone

	// Output one extra newline to separate output from Otto
	uiVal.Message("")

	// Return the output from the command
	return err
}

// OttoSkipCleanupEnvVar, when set, tells Otto to avoid cleaning up its
// temporary workspace files, which can be helpful for debugging.
const OttoSkipCleanupEnvVar = "OTTO_SKIP_CLEANUP"

// ShouldCleanup returns true for normal operation. It returns false if the
// user requested that Otto avoid cleaning up its temporary files for
// debugging purposes.
func ShouldCleanup() bool {
	return os.Getenv(OttoSkipCleanupEnvVar) == ""
}
<commit_msg>helper/exec: easier debugging<commit_after>package exec

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"strings"

	"github.com/hashicorp/otto/ui"
)

// Run runs the given command and streams all the output to the
// given UI. It also connects stdin properly so that input works as
// expected.
func Run(uiVal ui.Ui, cmd *exec.Cmd) error {
	out_r, out_w := io.Pipe()
	cmd.Stdin = os.Stdin
	cmd.Stdout = out_w
	cmd.Stderr = out_w

	// Copy output to the UI until we can't.
	uiDone := make(chan struct{})
	go func() {
		defer close(uiDone)
		var buf [1024]byte
		for {
			n, err := out_r.Read(buf[:])
			if n > 0 {
				uiVal.Raw(string(buf[:n]))
			}

			// We just break on any error. io.EOF is not an error and
			// is our true exit case, but any other error we don't really
			// handle here. It probably means something went wrong
			// somewhere else anyways.
			if err != nil {
				break
			}
		}
	}()

	// Run the command
	log.Printf("[DEBUG] execDir: %s", cmd.Dir)
	log.Printf("[DEBUG] exec: %s %s", cmd.Path, strings.Join(cmd.Args[1:], " "))

	// Build a runnable command that we can log out to make things easier
	// for debugging. This lets debuging devs just copy and paste the command.
	var debugBuf bytes.Buffer
	for _, env := range cmd.Env {
		parts := strings.SplitN(env, "=", 2)
		debugBuf.WriteString(fmt.Sprintf("%s=%q ", parts[0], parts[1]))
	}
	debugBuf.WriteString(cmd.Path + " ")
	for _, arg := range cmd.Args[1:] {
		if strings.Contains(arg, " ") {
			debugBuf.WriteString(fmt.Sprintf("'%s' ", arg))
		} else {
			debugBuf.WriteString(fmt.Sprintf("%s ", arg))
		}
	}
	log.Printf("[DEBUG] exec runnable: %s", debugBuf.String())
	debugBuf.Reset()

	// Run
	err := cmd.Run()

	// Wait for all the output to finish
	out_w.Close()
	<-uiDone

	// Output one extra newline to separate output from Otto
	uiVal.Message("")

	// Return the output from the command
	return err
}

// OttoSkipCleanupEnvVar, when set, tells Otto to avoid cleaning up its
// temporary workspace files, which can be helpful for debugging.
const OttoSkipCleanupEnvVar = "OTTO_SKIP_CLEANUP"

// ShouldCleanup returns true for normal operation. It returns false if the
// user requested that Otto avoid cleaning up its temporary files for
// debugging purposes.
func ShouldCleanup() bool {
	return os.Getenv(OttoSkipCleanupEnvVar) == ""
}
<|endoftext|>


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

package httputil

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

// Create an HTTP request with the supplied information.
//
// Unlike http.NewRequest:
//
//  *  This function doesn't mangle the supplied URL by round tripping it to a
//     string. For example, the Opaque field will continue to differentiate
//     between actual slashes in the path and escaped ones (cf.
//     http://goo.gl/rWX6ps).
//
//  *  This function doesn't magically re-interpret an io.Reader as an
//     io.ReadCloser when possible.
//
//  *  This function provides a convenient choke point to ensure we don't
//     forget to set a user agent header.
//
func NewRequest(
	method string,
	url *url.URL,
	body io.ReadCloser,
	userAgent string) (req *http.Request, err error) {
	// Create the request.
	req = &http.Request{
		Method:     method,
		URL:        url,
		Proto:      "HTTP/1.1",
		ProtoMajor: 1,
		ProtoMinor: 1,
		Header:     make(http.Header),
		Body:       body,
		Host:       url.Host,
	}

	// Set the User-Agent header.
	req.Header.Set("User-Agent", userAgent)

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

package httputil

import (
	"io"
	"net/http"
	"net/url"

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

// Create an HTTP request with the supplied information.
//
// Unlike http.NewRequest:
//
//  *  This function configures the request to be cancelled when the supplied
//     context is.
//
//  *  This function doesn't mangle the supplied URL by round tripping it to a
//     string. For example, the Opaque field will continue to differentiate
//     between actual slashes in the path and escaped ones (cf.
//     http://goo.gl/rWX6ps).
//
//  *  This function doesn't magically re-interpret an io.Reader as an
//     io.ReadCloser when possible.
//
//  *  This function provides a convenient choke point to ensure we don't
//     forget to set a user agent header.
//
func NewRequest(
	ctx context.Context,
	method string,
	url *url.URL,
	body io.ReadCloser,
	userAgent string) (req *http.Request, err error) {
	// Create the request.
	req = &http.Request{
		Method:     method,
		URL:        url,
		Proto:      "HTTP/1.1",
		ProtoMajor: 1,
		ProtoMinor: 1,
		Header:     make(http.Header),
		Body:       body,
		Host:       url.Host,
		Cancel:     ctx.Done(),
	}

	// Set the User-Agent header.
	req.Header.Set("User-Agent", userAgent)

	return
}
<|endoftext|>


<commit_before>package sudoku

import (
	"testing"
	"time"
)

func BenchmarkHumanSolve(b *testing.B) {
	for i := 0; i < b.N; i++ {
		grid := NewGrid()
		defer grid.Done()
		grid.Load(TEST_GRID)
		grid.HumanSolve()
	}
}

func TestHumanSolve(t *testing.T) {
	grid := NewGrid()
	defer grid.Done()
	grid.Load(TEST_GRID)

	steps := grid.HumanSolution()

	if steps == nil {
		t.Log("Human solution returned 0 techniques.")
		t.Fail()
	}

	if grid.Solved() {
		t.Log("Human Solutions mutated the grid.")
		t.Fail()
	}

	steps = grid.HumanSolve()
	//TODO: test to make sure that we use a wealth of different techniques. This will require a cooked random for testing.
	if steps == nil {
		t.Log("Human solve returned 0 techniques")
		t.Fail()
	}
	if !grid.Solved() {
		t.Log("Human solve failed to solve the simple grid.")
		t.Fail()
	}

}

func TestHint(t *testing.T) {
	//TODO: explicitly test hints that end in a guess, too.

	grid := NewGrid()
	defer grid.Done()

	grid.Load(TEST_GRID)

	steps := grid.Hint()

	if steps == nil || len(steps) == 0 {
		t.Error("No steps returned from Hint")
	}

	for count, step := range steps {
		if count == len(steps)-1 {
			//Last one
			if !step.Technique.IsFill() {
				t.Error("Non-fill step as last step in Hint: ", step.Technique.Name())
			}
		} else {
			//Not last one
			if step.Technique.IsFill() {
				t.Error("Fill step as non-last step in Hint: ", count, step.Technique.Name())
			}
		}
	}

}

func TestHumanSolveWithGuess(t *testing.T) {

	grid := NewGrid()
	defer grid.Done()

	if !grid.LoadFromFile(puzzlePath("harddifficulty.sdk")) {
		t.Fatal("harddifficulty.sdk wasn't loaded")
	}

	steps := grid.HumanSolution()

	if steps == nil {
		t.Fatal("Didn't find a solution to a grid that should have needed a guess")
	}

	foundGuess := false
	for i, step := range steps {
		if step.Technique.Name() == "Guess" {
			foundGuess = true
		}
		step.Apply(grid)
		if grid.Invalid() {
			t.Fatal("A solution with a guess in it got us into an invalid grid state. step", i)
		}
	}

	if !foundGuess {
		t.Error("Solution that should have used guess didn't have any guess.")
	}

	if !grid.Solved() {
		t.Error("A solution with a guess said it should solve the puzzle, but it didn't.")
	}

}

func TestStepsDescription(t *testing.T) {

	grid := NewGrid()
	defer grid.Done()

	//It's really brittle that we load techniques in this way... it changes every time we add a new early technique!
	steps := SolveDirections{
		&SolveStep{
			techniquesByName["Only Legal Number"],
			CellSlice{
				grid.Cell(0, 0),
			},
			IntSlice{1},
			nil,
			nil,
			nil,
		},
		&SolveStep{
			techniquesByName["Pointing Pair Col"],
			CellSlice{
				grid.Cell(1, 0),
				grid.Cell(1, 1),
			},
			IntSlice{1, 2},
			CellSlice{
				grid.Cell(1, 3),
				grid.Cell(1, 4),
			},
			nil,
			nil,
		},
		&SolveStep{
			techniquesByName["Only Legal Number"],
			CellSlice{
				grid.Cell(2, 0),
			},
			IntSlice{2},
			nil,
			nil,
			nil,
		},
	}

	descriptions := steps.Description()

	GOLDEN_DESCRIPTIONS := []string{
		"First, we put 1 in cell (0,0) because 1 is the only remaining valid number for that cell.",
		"Next, we remove the possibilities 1 and 2 from cells (1,0) and (1,1) because 1 is only possible in column 0 of block 1, which means it can't be in any other cell in that column not in that block.",
		"Finally, we put 2 in cell (2,0) because 2 is the only remaining valid number for that cell.",
	}

	for i := 0; i < len(GOLDEN_DESCRIPTIONS); i++ {
		if descriptions[i] != GOLDEN_DESCRIPTIONS[i] {
			t.Log("Got wrong human solve description: ", descriptions[i])
			t.Fail()
		}
	}
}

//TODO: this is useful. Should we use this in other tests?
func cellRefsToCells(refs []cellRef, grid *Grid) CellSlice {
	var result CellSlice
	for _, ref := range refs {
		result = append(result, ref.Cell(grid))
	}
	return result
}

func TestTweakChainedStepsWeights(t *testing.T) {

	//TODO: test other, harder cases as well.
	grid := NewGrid()
	lastStep := &SolveStep{
		nil,
		cellRefsToCells([]cellRef{
			{0, 0},
		}, grid),
		nil,
		nil,
		nil,
		nil,
	}
	possibilities := []*SolveStep{
		{
			nil,
			cellRefsToCells([]cellRef{
				{1, 0},
			}, grid),
			nil,
			nil,
			nil,
			nil,
		},
		{
			nil,
			cellRefsToCells([]cellRef{
				{2, 2},
			}, grid),
			nil,
			nil,
			nil,
			nil,
		},
		{
			nil,
			cellRefsToCells([]cellRef{
				{7, 7},
			}, grid),
			nil,
			nil,
			nil,
			nil,
		},
	}
	weights := []float64{
		10.0,
		10.0,
		10.0,
	}

	tweakChainedStepsWeights(lastStep, possibilities, weights)

	lastWeight := 0.0
	for i, weight := range weights {
		if weight <= lastWeight {
			t.Error("Tweak Chained Steps Weights didn't tweak things in the right direction: ", weights, "at", i)
		}
		lastWeight = weight
	}
}

func TestPuzzleDifficulty(t *testing.T) {
	grid := NewGrid()
	defer grid.Done()
	grid.Load(TEST_GRID)

	//We use the cheaper one for testing so it completes faster.
	difficulty := grid.calcluateDifficulty(false)

	if grid.Solved() {
		t.Log("Difficulty shouldn't have changed the underlying grid, but it did.")
		t.Fail()
	}

	if difficulty < 0.0 || difficulty > 1.0 {
		t.Log("The grid's difficulty was outside of allowed bounds.")
		t.Fail()
	}

	puzzleFilenames := []string{"harddifficulty.sdk", "harddifficulty2.sdk"}

	for _, filename := range puzzleFilenames {
		puzzleDifficultyHelper(filename, t)
	}
}

func puzzleDifficultyHelper(filename string, t *testing.T) {
	otherGrid := NewGrid()
	if !otherGrid.LoadFromFile(puzzlePath(filename)) {
		t.Log("Whoops, couldn't load the file to test:", filename)
		t.Fail()
	}

	after := time.After(time.Second * 60)

	done := make(chan bool)

	go func() {
		//We use the cheaper one for testing so it completes faster
		_ = otherGrid.calcluateDifficulty(false)
		done <- true
	}()

	select {
	case <-done:
		//totally fine.
	case <-after:
		//Uh oh.
		t.Log("We never finished solving the hard difficulty puzzle: ", filename)
		t.Fail()
	}
}
<commit_msg>Factored TestHint's meat into hintTestHelper.<commit_after>package sudoku

import (
	"testing"
	"time"
)

func BenchmarkHumanSolve(b *testing.B) {
	for i := 0; i < b.N; i++ {
		grid := NewGrid()
		defer grid.Done()
		grid.Load(TEST_GRID)
		grid.HumanSolve()
	}
}

func TestHumanSolve(t *testing.T) {
	grid := NewGrid()
	defer grid.Done()
	grid.Load(TEST_GRID)

	steps := grid.HumanSolution()

	if steps == nil {
		t.Log("Human solution returned 0 techniques.")
		t.Fail()
	}

	if grid.Solved() {
		t.Log("Human Solutions mutated the grid.")
		t.Fail()
	}

	steps = grid.HumanSolve()
	//TODO: test to make sure that we use a wealth of different techniques. This will require a cooked random for testing.
	if steps == nil {
		t.Log("Human solve returned 0 techniques")
		t.Fail()
	}
	if !grid.Solved() {
		t.Log("Human solve failed to solve the simple grid.")
		t.Fail()
	}

}

func TestHint(t *testing.T) {
	//TODO: explicitly test hints that end in a guess, too.
	hintTestHelper(t)
}

func hintTestHelper(t *testing.T) {
	grid := NewGrid()
	defer grid.Done()

	grid.Load(TEST_GRID)

	steps := grid.Hint()

	if steps == nil || len(steps) == 0 {
		t.Error("No steps returned from Hint")
	}

	for count, step := range steps {
		if count == len(steps)-1 {
			//Last one
			if !step.Technique.IsFill() {
				t.Error("Non-fill step as last step in Hint: ", step.Technique.Name())
			}
		} else {
			//Not last one
			if step.Technique.IsFill() {
				t.Error("Fill step as non-last step in Hint: ", count, step.Technique.Name())
			}
		}
	}
}

func TestHumanSolveWithGuess(t *testing.T) {

	grid := NewGrid()
	defer grid.Done()

	if !grid.LoadFromFile(puzzlePath("harddifficulty.sdk")) {
		t.Fatal("harddifficulty.sdk wasn't loaded")
	}

	steps := grid.HumanSolution()

	if steps == nil {
		t.Fatal("Didn't find a solution to a grid that should have needed a guess")
	}

	foundGuess := false
	for i, step := range steps {
		if step.Technique.Name() == "Guess" {
			foundGuess = true
		}
		step.Apply(grid)
		if grid.Invalid() {
			t.Fatal("A solution with a guess in it got us into an invalid grid state. step", i)
		}
	}

	if !foundGuess {
		t.Error("Solution that should have used guess didn't have any guess.")
	}

	if !grid.Solved() {
		t.Error("A solution with a guess said it should solve the puzzle, but it didn't.")
	}

}

func TestStepsDescription(t *testing.T) {

	grid := NewGrid()
	defer grid.Done()

	//It's really brittle that we load techniques in this way... it changes every time we add a new early technique!
	steps := SolveDirections{
		&SolveStep{
			techniquesByName["Only Legal Number"],
			CellSlice{
				grid.Cell(0, 0),
			},
			IntSlice{1},
			nil,
			nil,
			nil,
		},
		&SolveStep{
			techniquesByName["Pointing Pair Col"],
			CellSlice{
				grid.Cell(1, 0),
				grid.Cell(1, 1),
			},
			IntSlice{1, 2},
			CellSlice{
				grid.Cell(1, 3),
				grid.Cell(1, 4),
			},
			nil,
			nil,
		},
		&SolveStep{
			techniquesByName["Only Legal Number"],
			CellSlice{
				grid.Cell(2, 0),
			},
			IntSlice{2},
			nil,
			nil,
			nil,
		},
	}

	descriptions := steps.Description()

	GOLDEN_DESCRIPTIONS := []string{
		"First, we put 1 in cell (0,0) because 1 is the only remaining valid number for that cell.",
		"Next, we remove the possibilities 1 and 2 from cells (1,0) and (1,1) because 1 is only possible in column 0 of block 1, which means it can't be in any other cell in that column not in that block.",
		"Finally, we put 2 in cell (2,0) because 2 is the only remaining valid number for that cell.",
	}

	for i := 0; i < len(GOLDEN_DESCRIPTIONS); i++ {
		if descriptions[i] != GOLDEN_DESCRIPTIONS[i] {
			t.Log("Got wrong human solve description: ", descriptions[i])
			t.Fail()
		}
	}
}

//TODO: this is useful. Should we use this in other tests?
func cellRefsToCells(refs []cellRef, grid *Grid) CellSlice {
	var result CellSlice
	for _, ref := range refs {
		result = append(result, ref.Cell(grid))
	}
	return result
}

func TestTweakChainedStepsWeights(t *testing.T) {

	//TODO: test other, harder cases as well.
	grid := NewGrid()
	lastStep := &SolveStep{
		nil,
		cellRefsToCells([]cellRef{
			{0, 0},
		}, grid),
		nil,
		nil,
		nil,
		nil,
	}
	possibilities := []*SolveStep{
		{
			nil,
			cellRefsToCells([]cellRef{
				{1, 0},
			}, grid),
			nil,
			nil,
			nil,
			nil,
		},
		{
			nil,
			cellRefsToCells([]cellRef{
				{2, 2},
			}, grid),
			nil,
			nil,
			nil,
			nil,
		},
		{
			nil,
			cellRefsToCells([]cellRef{
				{7, 7},
			}, grid),
			nil,
			nil,
			nil,
			nil,
		},
	}
	weights := []float64{
		10.0,
		10.0,
		10.0,
	}

	tweakChainedStepsWeights(lastStep, possibilities, weights)

	lastWeight := 0.0
	for i, weight := range weights {
		if weight <= lastWeight {
			t.Error("Tweak Chained Steps Weights didn't tweak things in the right direction: ", weights, "at", i)
		}
		lastWeight = weight
	}
}

func TestPuzzleDifficulty(t *testing.T) {
	grid := NewGrid()
	defer grid.Done()
	grid.Load(TEST_GRID)

	//We use the cheaper one for testing so it completes faster.
	difficulty := grid.calcluateDifficulty(false)

	if grid.Solved() {
		t.Log("Difficulty shouldn't have changed the underlying grid, but it did.")
		t.Fail()
	}

	if difficulty < 0.0 || difficulty > 1.0 {
		t.Log("The grid's difficulty was outside of allowed bounds.")
		t.Fail()
	}

	puzzleFilenames := []string{"harddifficulty.sdk", "harddifficulty2.sdk"}

	for _, filename := range puzzleFilenames {
		puzzleDifficultyHelper(filename, t)
	}
}

func puzzleDifficultyHelper(filename string, t *testing.T) {
	otherGrid := NewGrid()
	if !otherGrid.LoadFromFile(puzzlePath(filename)) {
		t.Log("Whoops, couldn't load the file to test:", filename)
		t.Fail()
	}

	after := time.After(time.Second * 60)

	done := make(chan bool)

	go func() {
		//We use the cheaper one for testing so it completes faster
		_ = otherGrid.calcluateDifficulty(false)
		done <- true
	}()

	select {
	case <-done:
		//totally fine.
	case <-after:
		//Uh oh.
		t.Log("We never finished solving the hard difficulty puzzle: ", filename)
		t.Fail()
	}
}
<|endoftext|>