<commit_before>package dbf

import (
	"bufio"
	"encoding/json"
	"fmt"
	"math"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"sync"
)

const (
	PATH_SEPARATOR = string(os.PathSeparator)

	_PATH_CONF_DIR  = 1
	_PATH_CONF_FILE = 2
	_PATH_VENDOR    = 3
	_PATH_TASK      = 4
	_PATH_CRACK     = 5

	_BENCH_TYPE_CPU     = 0
	_BENCH_TYPE_GPU_AMD = 1
	_BENCH_TYPE_GPU_NV  = 2

	_VENDOR_TYPE_GENERATOR = "generator"
	_VENDOR_TYPE_CRACKER   = "cracker"
)

var (
	PathData          string
	confDbf           StructConfDbf
	activePlatStr     string
	regexpBenchValue  *regexp.Regexp
	regexpBenchFloat  *regexp.Regexp
	regexpBenchSuffix *regexp.Regexp
	wgBench           sync.WaitGroup
)

type StructActivePlatform struct {
	Id        string `json:"id"`
	Benchmark uint64 `json:"benchmark"`
}

func InitConfig() {
	initConfigPlatform()

	check()
}

func getPath(path int) string {
	switch path {
	case _PATH_CONF_DIR:
		return PathData + "config" + PATH_SEPARATOR
	case _PATH_CONF_FILE:
		return PathData + "config" + PATH_SEPARATOR + "dbf.json"
	case _PATH_CRACK:
		return PathData + "crack" + PATH_SEPARATOR
	case _PATH_TASK:
		return PathData + "task" + PATH_SEPARATOR
	case _PATH_VENDOR:
		return PathData + "vendor" + PATH_SEPARATOR
	default:
		Log.Printf("Undefined path id '%d' in getPath()!\n", path)
		return ""
	}
}

func check() {
	err := checkDir(getPath(_PATH_CONF_DIR))
	if err != nil {
		Log.Printf("%s\n", err)
		panic(1)
	}

	// Check config file
	if _, err := os.Stat(getPath(_PATH_CONF_FILE)); err != nil {
		if os.IsNotExist(err) { // Does not exist, so create it

			fmt.Println("Creating initial config file...")

			// Create initial config file
			err = createConfDbf()
			if err != nil {
				Log.Printf("%s\n", err) // Can't create
				panic(1)
			}
		} else {
			Log.Printf("%s\n", err) // Error in accessing
			panic(1)
		}
	}

	// Sync config file
	fmt.Println("Synchronizing config file...")

	err = checkDir(getPath(_PATH_VENDOR))
	if err != nil {
		Log.Printf("%s\n", err)
		panic(1)
	}

	err = checkDir(getPath(_PATH_TASK))
	if err != nil {
		Log.Printf("%s\n", err)
		panic(1)
	}

	err = checkDir(getPath(_PATH_CRACK))
	if err != nil {
		Log.Printf("%s\n", err)
		panic(1)
	}

	setConfDbf()

	initServer()

	// Init regexp
	regexpBenchValue, err = regexp.Compile("\\d*\\.?\\d*[A-Za-z]?")
	if err != nil {
		Log.Printf("%s\n", err)
		panic(1)
	}
	regexpBenchFloat, err = regexp.Compile("\\d*\\.?\\d*")
	if err != nil {
		Log.Printf("%s\n", err)
		panic(1)
	}
	regexpBenchSuffix, err = regexp.Compile("[A-Za-z]$")
	if err != nil {
		Log.Printf("%s\n", err)
		panic(1)
	}

	// Check default vendor files and update benchmarks
	for i, platform := range confDbf.Platform {
		if platform.Active != 0 { // Is active
			if strings.HasPrefix(platform.Id, "cpu") { // CPU
				wgBench.Add(1)
				go func(i int, platformId string) {
					confDbf.Platform[i].Benchmark = getBench(_BENCH_TYPE_CPU, &platformId)
					wgBench.Done()
				}(i, platform.Id)
			} else if strings.HasSuffix(platform.Id, "_amd") { // GPU AMD
				wgBench.Add(1)
				go func(i int, platformId string) {
					confDbf.Platform[i].Benchmark = getBench(_BENCH_TYPE_GPU_AMD, &platformId)
					wgBench.Done()
				}(i, platform.Id)
			} else if strings.HasSuffix(platform.Id, "_nv") { // GPU Nvidia
				wgBench.Add(1)
				go func(i int, platformId string) {
					confDbf.Platform[i].Benchmark = getBench(_BENCH_TYPE_GPU_NV, &platformId)
					wgBench.Done()
				}(i, platform.Id)
			}
		}
	}

	wgBench.Wait() // Wait for all benchmarks to finish

	var activePlat []StructActivePlatform
	for i, platform := range confDbf.Platform {
		if platform.Active != 0 { // Is active
			if confDbf.Platform[i].Benchmark > 0 {
				activePlat = append(activePlat, StructActivePlatform{
					Id:        platform.Id,
					Benchmark: confDbf.Platform[i].Benchmark,
				})
			}
		}
	}

	activePlatByte, err := json.Marshal(activePlat)
	if err != nil {
		Log.Printf("%s\n", err)
	}
	activePlatStr = string(activePlatByte)

	fmt.Println("Updating config file...")

	// Update config file
	saveConfDbf()
}

// checkDir checks if dir exists and is accessible, otherwise tries to create it
func checkDir(path string) error {
	if _, err := os.Stat(path); err != nil {
		if os.IsNotExist(err) { // Does not exist, so create it
			if err = os.MkdirAll(path, 0775); err != nil {
				Log.Printf("%s\n", err) // Error in creating
				return err
			}
		} else {
			Log.Printf("%s\n", err) // Error in accessing
			return err
		}
	}

	return nil
}

func getBench(benchType int, platformId *string) uint64 {
	// Check vendor file
	var vendorBench, benchLinePrefix string
	switch benchType {
	case _BENCH_TYPE_CPU:
		vendorBench = "hashcat"
		benchLinePrefix = "Speed/sec:"
	case _BENCH_TYPE_GPU_AMD:
		vendorBench = "oclHashcat"
		benchLinePrefix = "Speed.GPU.#"
	case _BENCH_TYPE_GPU_NV:
		vendorBench = "cudaHashcat"
		benchLinePrefix = "Speed.GPU.#"
	default:
		return 0
	}

	vendorBenchPath := getPath(_PATH_VENDOR) + _VENDOR_TYPE_CRACKER + PATH_SEPARATOR + vendorBench + PATH_SEPARATOR + *platformId + PATH_SEPARATOR + _VENDOR_TYPE_CRACKER + extExecutable

	if checkVendor(_VENDOR_TYPE_CRACKER, &vendorBench, platformId, &vendorBenchPath) == false {
		return 0
	}

	// Preform benchmark
	fmt.Printf("Preforming benchmark (%s)...\n", *platformId)

	cmd := exec.Command(vendorBenchPath, "-b", "-m 0")

	cmdOut, err := cmd.StdoutPipe()
	if err != nil {
		Log.Printf("%s\n", err)
		return 0
	}

	err = cmd.Start()
	if err != nil {
		Log.Printf("%s\n", err)
		return 0
	}

	var bench, totalBench uint64
	numOfBench := 0

	cmdScanner := bufio.NewScanner(cmdOut)
	cmdScanner.Split(bufio.ScanLines)
	for cmdScanner.Scan() {
		// Get benchmark from line
		if strings.HasPrefix(cmdScanner.Text(), benchLinePrefix) {
			bench = getBenchValue(regexpBenchValue.FindString(strings.Replace(strings.SplitN(cmdScanner.Text(), ":", 2)[1], " ", "", -1)))
			if bench > 0 {
				totalBench += bench
				numOfBench++
			}
		}
	}
	err = cmdScanner.Err()
	if err != nil {
		Log.Printf("%s\n", err)
	}

	err = cmd.Wait()
	if err != nil {
		Log.Printf("%s\n", err)
	}

	if numOfBench > 0 {
		return uint64(math.Floor(float64(totalBench/uint64(numOfBench)) + 0.5))
	}

	return 0
}

func getBenchValue(benchValue string) uint64 {
	benchFloat, err := strconv.ParseFloat(regexpBenchFloat.FindString(benchValue), 64)
	if err != nil {
		Log.Printf("%s\n", err)
		return 0
	}

	benchSuffix := strings.ToUpper(regexpBenchSuffix.FindString(benchValue))

	// Should return bench in Mega

	if len(benchSuffix) == 0 {
		benchFloat /= 1048576 // Byte to Mega
	}

	switch benchSuffix {
	case "K":
		benchFloat /= 1024 // Kilo to Mega
	case "G":
		benchFloat *= 1024 // Giga to Mega
	case "T":
		benchFloat *= 1048576 // Tera to Mega
	case "P":
		benchFloat *= 1073741824 // Peta to Mega
	case "E":
		benchFloat *= 1099511627776 // Exa to Mega
	case "Z":
		benchFloat *= 1125899906842624 // Zetta to Mega
	case "Y":
		benchFloat *= 1152921504606846976 // Yotta to Mega
	}

	return uint64(math.Floor(benchFloat + 0.5))
}

func checkVendor(vendorType string, vendorName *string, platformId *string, vendorPath *string) bool {
	if _, err := os.Stat(*vendorPath); err != nil {
		if os.IsNotExist(err) { // Does not exist, so get it

			fmt.Printf("Downloading %s: %s (%s)...\n", vendorType, *vendorName, *platformId)

			return getVendor(&vendorType, vendorName, platformId, vendorPath)
		} else {
			Log.Printf("%s\n", err) // Error in accessing
			return false
		}
	}

	return true
}
<commit_msg>Warn that download may take long<commit_after>package dbf

import (
	"bufio"
	"encoding/json"
	"fmt"
	"math"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"sync"
)

const (
	PATH_SEPARATOR = string(os.PathSeparator)

	_PATH_CONF_DIR  = 1
	_PATH_CONF_FILE = 2
	_PATH_VENDOR    = 3
	_PATH_TASK      = 4
	_PATH_CRACK     = 5

	_BENCH_TYPE_CPU     = 0
	_BENCH_TYPE_GPU_AMD = 1
	_BENCH_TYPE_GPU_NV  = 2

	_VENDOR_TYPE_GENERATOR = "generator"
	_VENDOR_TYPE_CRACKER   = "cracker"
)

var (
	PathData          string
	confDbf           StructConfDbf
	activePlatStr     string
	regexpBenchValue  *regexp.Regexp
	regexpBenchFloat  *regexp.Regexp
	regexpBenchSuffix *regexp.Regexp
	wgBench           sync.WaitGroup
)

type StructActivePlatform struct {
	Id        string `json:"id"`
	Benchmark uint64 `json:"benchmark"`
}

func InitConfig() {
	initConfigPlatform()

	check()
}

func getPath(path int) string {
	switch path {
	case _PATH_CONF_DIR:
		return PathData + "config" + PATH_SEPARATOR
	case _PATH_CONF_FILE:
		return PathData + "config" + PATH_SEPARATOR + "dbf.json"
	case _PATH_CRACK:
		return PathData + "crack" + PATH_SEPARATOR
	case _PATH_TASK:
		return PathData + "task" + PATH_SEPARATOR
	case _PATH_VENDOR:
		return PathData + "vendor" + PATH_SEPARATOR
	default:
		Log.Printf("Undefined path id '%d' in getPath()!\n", path)
		return ""
	}
}

func check() {
	err := checkDir(getPath(_PATH_CONF_DIR))
	if err != nil {
		Log.Printf("%s\n", err)
		panic(1)
	}

	// Check config file
	if _, err := os.Stat(getPath(_PATH_CONF_FILE)); err != nil {
		if os.IsNotExist(err) { // Does not exist, so create it

			fmt.Println("Creating initial config file...")

			// Create initial config file
			err = createConfDbf()
			if err != nil {
				Log.Printf("%s\n", err) // Can't create
				panic(1)
			}
		} else {
			Log.Printf("%s\n", err) // Error in accessing
			panic(1)
		}
	}

	// Sync config file
	fmt.Println("Synchronizing config file...")

	err = checkDir(getPath(_PATH_VENDOR))
	if err != nil {
		Log.Printf("%s\n", err)
		panic(1)
	}

	err = checkDir(getPath(_PATH_TASK))
	if err != nil {
		Log.Printf("%s\n", err)
		panic(1)
	}

	err = checkDir(getPath(_PATH_CRACK))
	if err != nil {
		Log.Printf("%s\n", err)
		panic(1)
	}

	setConfDbf()

	initServer()

	// Init regexp
	regexpBenchValue, err = regexp.Compile("\\d*\\.?\\d*[A-Za-z]?")
	if err != nil {
		Log.Printf("%s\n", err)
		panic(1)
	}
	regexpBenchFloat, err = regexp.Compile("\\d*\\.?\\d*")
	if err != nil {
		Log.Printf("%s\n", err)
		panic(1)
	}
	regexpBenchSuffix, err = regexp.Compile("[A-Za-z]$")
	if err != nil {
		Log.Printf("%s\n", err)
		panic(1)
	}

	// Check default vendor files and update benchmarks
	for i, platform := range confDbf.Platform {
		if platform.Active != 0 { // Is active
			if strings.HasPrefix(platform.Id, "cpu") { // CPU
				wgBench.Add(1)
				go func(i int, platformId string) {
					confDbf.Platform[i].Benchmark = getBench(_BENCH_TYPE_CPU, &platformId)
					wgBench.Done()
				}(i, platform.Id)
			} else if strings.HasSuffix(platform.Id, "_amd") { // GPU AMD
				wgBench.Add(1)
				go func(i int, platformId string) {
					confDbf.Platform[i].Benchmark = getBench(_BENCH_TYPE_GPU_AMD, &platformId)
					wgBench.Done()
				}(i, platform.Id)
			} else if strings.HasSuffix(platform.Id, "_nv") { // GPU Nvidia
				wgBench.Add(1)
				go func(i int, platformId string) {
					confDbf.Platform[i].Benchmark = getBench(_BENCH_TYPE_GPU_NV, &platformId)
					wgBench.Done()
				}(i, platform.Id)
			}
		}
	}

	wgBench.Wait() // Wait for all benchmarks to finish

	var activePlat []StructActivePlatform
	for i, platform := range confDbf.Platform {
		if platform.Active != 0 { // Is active
			if confDbf.Platform[i].Benchmark > 0 {
				activePlat = append(activePlat, StructActivePlatform{
					Id:        platform.Id,
					Benchmark: confDbf.Platform[i].Benchmark,
				})
			}
		}
	}

	activePlatByte, err := json.Marshal(activePlat)
	if err != nil {
		Log.Printf("%s\n", err)
	}
	activePlatStr = string(activePlatByte)

	fmt.Println("Updating config file...")

	// Update config file
	saveConfDbf()
}

// checkDir checks if dir exists and is accessible, otherwise tries to create it
func checkDir(path string) error {
	if _, err := os.Stat(path); err != nil {
		if os.IsNotExist(err) { // Does not exist, so create it
			if err = os.MkdirAll(path, 0775); err != nil {
				Log.Printf("%s\n", err) // Error in creating
				return err
			}
		} else {
			Log.Printf("%s\n", err) // Error in accessing
			return err
		}
	}

	return nil
}

func getBench(benchType int, platformId *string) uint64 {
	// Check vendor file
	var vendorBench, benchLinePrefix string
	switch benchType {
	case _BENCH_TYPE_CPU:
		vendorBench = "hashcat"
		benchLinePrefix = "Speed/sec:"
	case _BENCH_TYPE_GPU_AMD:
		vendorBench = "oclHashcat"
		benchLinePrefix = "Speed.GPU.#"
	case _BENCH_TYPE_GPU_NV:
		vendorBench = "cudaHashcat"
		benchLinePrefix = "Speed.GPU.#"
	default:
		return 0
	}

	vendorBenchPath := getPath(_PATH_VENDOR) + _VENDOR_TYPE_CRACKER + PATH_SEPARATOR + vendorBench + PATH_SEPARATOR + *platformId + PATH_SEPARATOR + _VENDOR_TYPE_CRACKER + extExecutable

	if checkVendor(_VENDOR_TYPE_CRACKER, &vendorBench, platformId, &vendorBenchPath) == false {
		return 0
	}

	// Preform benchmark
	fmt.Printf("Preforming benchmark (%s)...\n", *platformId)

	cmd := exec.Command(vendorBenchPath, "-b", "-m 0")

	cmdOut, err := cmd.StdoutPipe()
	if err != nil {
		Log.Printf("%s\n", err)
		return 0
	}

	err = cmd.Start()
	if err != nil {
		Log.Printf("%s\n", err)
		return 0
	}

	var bench, totalBench uint64
	numOfBench := 0

	cmdScanner := bufio.NewScanner(cmdOut)
	cmdScanner.Split(bufio.ScanLines)
	for cmdScanner.Scan() {
		// Get benchmark from line
		if strings.HasPrefix(cmdScanner.Text(), benchLinePrefix) {
			bench = getBenchValue(regexpBenchValue.FindString(strings.Replace(strings.SplitN(cmdScanner.Text(), ":", 2)[1], " ", "", -1)))
			if bench > 0 {
				totalBench += bench
				numOfBench++
			}
		}
	}
	err = cmdScanner.Err()
	if err != nil {
		Log.Printf("%s\n", err)
	}

	err = cmd.Wait()
	if err != nil {
		Log.Printf("%s\n", err)
	}

	if numOfBench > 0 {
		return uint64(math.Floor(float64(totalBench/uint64(numOfBench)) + 0.5))
	}

	return 0
}

func getBenchValue(benchValue string) uint64 {
	benchFloat, err := strconv.ParseFloat(regexpBenchFloat.FindString(benchValue), 64)
	if err != nil {
		Log.Printf("%s\n", err)
		return 0
	}

	benchSuffix := strings.ToUpper(regexpBenchSuffix.FindString(benchValue))

	// Should return bench in Mega

	if len(benchSuffix) == 0 {
		benchFloat /= 1048576 // Byte to Mega
	}

	switch benchSuffix {
	case "K":
		benchFloat /= 1024 // Kilo to Mega
	case "G":
		benchFloat *= 1024 // Giga to Mega
	case "T":
		benchFloat *= 1048576 // Tera to Mega
	case "P":
		benchFloat *= 1073741824 // Peta to Mega
	case "E":
		benchFloat *= 1099511627776 // Exa to Mega
	case "Z":
		benchFloat *= 1125899906842624 // Zetta to Mega
	case "Y":
		benchFloat *= 1152921504606846976 // Yotta to Mega
	}

	return uint64(math.Floor(benchFloat + 0.5))
}

func checkVendor(vendorType string, vendorName *string, platformId *string, vendorPath *string) bool {
	if _, err := os.Stat(*vendorPath); err != nil {
		if os.IsNotExist(err) { // Does not exist, so get it

			fmt.Printf("Downloading %s: %s (%s)\nThis may take a few minutes, please wait...\n", vendorType, *vendorName, *platformId)

			return getVendor(&vendorType, vendorName, platformId, vendorPath)
		} else {
			Log.Printf("%s\n", err) // Error in accessing
			return false
		}
	}

	return true
}
<|endoftext|>


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

package ddl

import (
	"time"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/tidb/context"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/meta"
	"github.com/pingcap/tidb/model"
	"github.com/pingcap/tidb/terror"
)

func (d *ddl) startJob(ctx context.Context, job *model.Job) error {
	// for every DDL, we must commit current transaction.
	if err := ctx.FinishTxn(false); err != nil {
		return errors.Trace(err)
	}

	// Create a new job and queue it.
	err := kv.RunInNewTxn(d.store, true, func(txn kv.Transaction) error {
		t := meta.NewMeta(txn)
		var err error
		job.ID, err = t.GenGlobalID()
		if err != nil {
			return errors.Trace(err)
		}

		err = t.EnQueueDDLJob(job)
		return errors.Trace(err)
	})

	if err != nil {
		return errors.Trace(err)
	}

	// notice worker that we push a new job and wait the job done.
	asyncNotify(d.jobCh)

	log.Warnf("start DDL job %v", job)

	jobID := job.ID

	var historyJob *model.Job

	// for a job from start to end, the state of it will be none -> delete only -> write only -> reorganization -> public
	// for every state change, we will wait as lease 2 * lease time, so here the ticker check is 10 * lease.
	ticker := time.NewTicker(chooseLeaseTime(10*d.lease, 10*time.Second))
	defer ticker.Stop()
	for {
		select {
		case <-d.jobDoneCh:
		case <-ticker.C:
		}

		historyJob, err = d.getHistoryJob(jobID)
		if err != nil {
			log.Errorf("get history job err %v, check again", err)
			continue
		} else if historyJob == nil {
			log.Warnf("job %d is not in history, maybe not run", jobID)
			continue
		}

		// if a job is a history table, the state must be JobDone or JobCancel.
		if historyJob.State == model.JobDone {
			return nil
		}

		return errors.Errorf(historyJob.Error)
	}
}

func (d *ddl) getHistoryJob(id int64) (*model.Job, error) {
	var job *model.Job

	err := kv.RunInNewTxn(d.store, false, func(txn kv.Transaction) error {
		t := meta.NewMeta(txn)
		var err1 error
		job, err1 = t.GetHistoryDDLJob(id)
		return errors.Trace(err1)
	})

	return job, errors.Trace(err)
}

func asyncNotify(ch chan struct{}) {
	select {
	case ch <- struct{}{}:
	default:
	}
}

func (d *ddl) checkOwner(t *meta.Meta) (*model.Owner, error) {
	owner, err := t.GetDDLOwner()
	if err != nil {
		return nil, errors.Trace(err)
	}

	if owner == nil {
		owner = &model.Owner{}
		// try to set onwer
		owner.OwnerID = d.uuid
	}

	now := time.Now().UnixNano()
	// we must wait 2 * lease time to guarantee other servers update the schema,
	// the owner will update its owner status every 2 * lease time, so here we use
	// 4 * lease to check its timeout.
	maxTimeout := int64(4 * d.lease)
	if owner.OwnerID == d.uuid || now-owner.LastUpdateTS > maxTimeout {
		owner.OwnerID = d.uuid
		owner.LastUpdateTS = now
		// update status.
		if err = t.SetDDLOwner(owner); err != nil {
			return nil, errors.Trace(err)
		}
		log.Debugf("become owner %s", owner.OwnerID)
	}

	if owner.OwnerID != d.uuid {
		log.Debugf("not owner, owner is %s", owner.OwnerID)
		return nil, errors.Trace(ErrNotOwner)
	}

	return owner, nil
}

func (d *ddl) getFirstJob(t *meta.Meta) (*model.Job, error) {
	job, err := t.GetDDLJob(0)
	return job, errors.Trace(err)
}

// every time we enter another state except final state, we must call this function.
func (d *ddl) updateJob(t *meta.Meta, job *model.Job) error {
	err := t.UpdateDDLJob(0, job)
	return errors.Trace(err)
}

func (d *ddl) finishJob(t *meta.Meta, job *model.Job) error {
	log.Warnf("finish DDL job %v", job)
	// done, notice and run next job.
	_, err := t.DeQueueDDLJob()
	if err != nil {
		return errors.Trace(err)
	}

	err = t.AddHistoryDDLJob(job)
	return errors.Trace(err)
}

// ErrNotOwner means we are not owner and can't handle DDL jobs.
var ErrNotOwner = errors.New("DDL: not owner")

func (d *ddl) handleJobQueue() error {
	for {
		if d.isClosed() {
			return nil
		}

		var job *model.Job
		err := kv.RunInNewTxn(d.store, false, func(txn kv.Transaction) error {
			t := meta.NewMeta(txn)
			owner, err := d.checkOwner(t)
			if terror.ErrorEqual(err, ErrNotOwner) {
				// we are not owner, return and retry checking later.
				return nil
			} else if err != nil {
				return errors.Trace(err)
			}

			// become the owner
			// get the first job and run
			job, err = d.getFirstJob(t)
			if job == nil || err != nil {
				return errors.Trace(err)
			}

			log.Warnf("run DDL job %v", job)

			d.hook.OnJobRunBefore(job)

			// if run job meets error, we will save this error in job Error
			// and retry later if the job is not cancelled.
			d.runJob(t, job)

			d.hook.OnJobRunAfter(job)

			if job.State == model.JobDone || job.State == model.JobCancelled {
				err = d.finishJob(t, job)
				if err == nil {
					asyncNotify(d.jobDoneCh)
				}
			} else {
				err = d.updateJob(t, job)
			}

			if err != nil {
				return errors.Trace(err)
			}

			// running job may cost some time, so here we must update owner status to
			// prevent other become the owner.
			owner.LastUpdateTS = time.Now().UnixNano()
			if err = t.SetDDLOwner(owner); err != nil {
				return errors.Trace(err)
			}
			return errors.Trace(err)
		})

		if err != nil {
			return errors.Trace(err)
		} else if job == nil {
			// no job now, return and retry get later.
			return nil
		}

		// here means the job enters another state (delete only, write only, public, etc...) or is cancelled.
		// if the job is done or still running, we will wait 2 * lease time to guarantee other servers to update
		// the newest schema.
		if job.State == model.JobRunning || job.State == model.JobDone {
			d.waitSchemaChanged()
		}
	}
}

func chooseLeaseTime(n1 time.Duration, n2 time.Duration) time.Duration {
	if n1 > 0 {
		return n1
	}

	return n2
}

// onWorker is for async online schema change, it will try to become the owner first,
// then wait or pull the job queue to handle a schema change job.
func (d *ddl) onWorker() {
	defer d.wait.Done()

	// we use 4 * lease time to check owner's timeout, so here, we will update owner's status
	// every 2 * lease time, if lease is 0, we will use default 10s.
	checkTime := chooseLeaseTime(2*d.lease, 10*time.Second)

	ticker := time.NewTicker(checkTime)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			log.Debugf("wait %s to check DDL status again", checkTime)
		case <-d.jobCh:
		case <-d.quitCh:
			return
		}

		err := d.handleJobQueue()
		if err != nil {
			log.Errorf("handle job err %v", errors.ErrorStack(err))
		}
	}
}

func (d *ddl) runJob(t *meta.Meta, job *model.Job) {
	if job.State == model.JobDone || job.State == model.JobCancelled {
		return
	}

	job.State = model.JobRunning

	var err error
	switch job.Type {
	case model.ActionCreateSchema:
		err = d.onCreateSchema(t, job)
	case model.ActionDropSchema:
		err = d.onDropSchema(t, job)
	case model.ActionCreateTable:
		err = d.onCreateTable(t, job)
	case model.ActionDropTable:
		err = d.onDropTable(t, job)
	case model.ActionAddColumn:
		err = d.onAddColumn(t, job)
	case model.ActionDropColumn:
		err = d.onDropColumn(t, job)
	case model.ActionAddIndex:
		err = d.onCreateIndex(t, job)
	case model.ActionDropIndex:
		err = d.onDropIndex(t, job)
	case model.ActionAddConstraint:
		log.Fatal("Doesn't support change constraint online")
	case model.ActionDropConstraint:
		log.Fatal("Doesn't support change constraint online")
	default:
		// invalid job, cancel it.
		job.State = model.JobCancelled
		err = errors.Errorf("invalid job %v", job)
	}

	// saves error in job, so that others can know error happens.
	if err != nil {
		// if job is not cancelled, we should log this error.
		if job.State != model.JobCancelled {
			log.Errorf("run job err %v", errors.ErrorStack(err))
		}

		job.Error = err.Error()
		job.ErrorCount++
	}
}

// for every lease seconds, we will re-update the whole schema, so we will wait 2 * lease time
// to guarantee that all servers have already updated schema.
func (d *ddl) waitSchemaChanged() {
	if d.lease > 0 {
		select {
		case <-time.After(d.lease * 2):
		case <-d.quitCh:
		}
	}
}
<commit_msg>ddl: adjust job callback place<commit_after>// Copyright 2015 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package ddl

import (
	"time"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/tidb/context"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/meta"
	"github.com/pingcap/tidb/model"
	"github.com/pingcap/tidb/terror"
)

func (d *ddl) startJob(ctx context.Context, job *model.Job) error {
	// for every DDL, we must commit current transaction.
	if err := ctx.FinishTxn(false); err != nil {
		return errors.Trace(err)
	}

	// Create a new job and queue it.
	err := kv.RunInNewTxn(d.store, true, func(txn kv.Transaction) error {
		t := meta.NewMeta(txn)
		var err error
		job.ID, err = t.GenGlobalID()
		if err != nil {
			return errors.Trace(err)
		}

		err = t.EnQueueDDLJob(job)
		return errors.Trace(err)
	})

	if err != nil {
		return errors.Trace(err)
	}

	// notice worker that we push a new job and wait the job done.
	asyncNotify(d.jobCh)

	log.Warnf("start DDL job %v", job)

	jobID := job.ID

	var historyJob *model.Job

	// for a job from start to end, the state of it will be none -> delete only -> write only -> reorganization -> public
	// for every state change, we will wait as lease 2 * lease time, so here the ticker check is 10 * lease.
	ticker := time.NewTicker(chooseLeaseTime(10*d.lease, 10*time.Second))
	defer ticker.Stop()
	for {
		select {
		case <-d.jobDoneCh:
		case <-ticker.C:
		}

		historyJob, err = d.getHistoryJob(jobID)
		if err != nil {
			log.Errorf("get history job err %v, check again", err)
			continue
		} else if historyJob == nil {
			log.Warnf("job %d is not in history, maybe not run", jobID)
			continue
		}

		// if a job is a history table, the state must be JobDone or JobCancel.
		if historyJob.State == model.JobDone {
			return nil
		}

		return errors.Errorf(historyJob.Error)
	}
}

func (d *ddl) getHistoryJob(id int64) (*model.Job, error) {
	var job *model.Job

	err := kv.RunInNewTxn(d.store, false, func(txn kv.Transaction) error {
		t := meta.NewMeta(txn)
		var err1 error
		job, err1 = t.GetHistoryDDLJob(id)
		return errors.Trace(err1)
	})

	return job, errors.Trace(err)
}

func asyncNotify(ch chan struct{}) {
	select {
	case ch <- struct{}{}:
	default:
	}
}

func (d *ddl) checkOwner(t *meta.Meta) (*model.Owner, error) {
	owner, err := t.GetDDLOwner()
	if err != nil {
		return nil, errors.Trace(err)
	}

	if owner == nil {
		owner = &model.Owner{}
		// try to set onwer
		owner.OwnerID = d.uuid
	}

	now := time.Now().UnixNano()
	// we must wait 2 * lease time to guarantee other servers update the schema,
	// the owner will update its owner status every 2 * lease time, so here we use
	// 4 * lease to check its timeout.
	maxTimeout := int64(4 * d.lease)
	if owner.OwnerID == d.uuid || now-owner.LastUpdateTS > maxTimeout {
		owner.OwnerID = d.uuid
		owner.LastUpdateTS = now
		// update status.
		if err = t.SetDDLOwner(owner); err != nil {
			return nil, errors.Trace(err)
		}
		log.Debugf("become owner %s", owner.OwnerID)
	}

	if owner.OwnerID != d.uuid {
		log.Debugf("not owner, owner is %s", owner.OwnerID)
		return nil, errors.Trace(ErrNotOwner)
	}

	return owner, nil
}

func (d *ddl) getFirstJob(t *meta.Meta) (*model.Job, error) {
	job, err := t.GetDDLJob(0)
	return job, errors.Trace(err)
}

// every time we enter another state except final state, we must call this function.
func (d *ddl) updateJob(t *meta.Meta, job *model.Job) error {
	err := t.UpdateDDLJob(0, job)
	return errors.Trace(err)
}

func (d *ddl) finishJob(t *meta.Meta, job *model.Job) error {
	log.Warnf("finish DDL job %v", job)
	// done, notice and run next job.
	_, err := t.DeQueueDDLJob()
	if err != nil {
		return errors.Trace(err)
	}

	err = t.AddHistoryDDLJob(job)
	return errors.Trace(err)
}

// ErrNotOwner means we are not owner and can't handle DDL jobs.
var ErrNotOwner = errors.New("DDL: not owner")

func (d *ddl) handleJobQueue() error {
	for {
		if d.isClosed() {
			return nil
		}

		var job *model.Job
		err := kv.RunInNewTxn(d.store, false, func(txn kv.Transaction) error {
			t := meta.NewMeta(txn)
			owner, err := d.checkOwner(t)
			if terror.ErrorEqual(err, ErrNotOwner) {
				// we are not owner, return and retry checking later.
				return nil
			} else if err != nil {
				return errors.Trace(err)
			}

			// become the owner
			// get the first job and run
			job, err = d.getFirstJob(t)
			if job == nil || err != nil {
				return errors.Trace(err)
			}

			log.Warnf("run DDL job %v", job)

			d.hook.OnJobRunBefore(job)

			// if run job meets error, we will save this error in job Error
			// and retry later if the job is not cancelled.
			d.runJob(t, job)

			if job.State == model.JobDone || job.State == model.JobCancelled {
				err = d.finishJob(t, job)
				if err == nil {
					asyncNotify(d.jobDoneCh)
				}
			} else {
				err = d.updateJob(t, job)
			}

			if err != nil {
				return errors.Trace(err)
			}

			// running job may cost some time, so here we must update owner status to
			// prevent other become the owner.
			owner.LastUpdateTS = time.Now().UnixNano()
			if err = t.SetDDLOwner(owner); err != nil {
				return errors.Trace(err)
			}
			return errors.Trace(err)
		})

		if err != nil {
			return errors.Trace(err)
		} else if job == nil {
			// no job now, return and retry get later.
			return nil
		}

		d.hook.OnJobRunAfter(job)

		// here means the job enters another state (delete only, write only, public, etc...) or is cancelled.
		// if the job is done or still running, we will wait 2 * lease time to guarantee other servers to update
		// the newest schema.
		if job.State == model.JobRunning || job.State == model.JobDone {
			d.waitSchemaChanged()
		}
	}
}

func chooseLeaseTime(n1 time.Duration, n2 time.Duration) time.Duration {
	if n1 > 0 {
		return n1
	}

	return n2
}

// onWorker is for async online schema change, it will try to become the owner first,
// then wait or pull the job queue to handle a schema change job.
func (d *ddl) onWorker() {
	defer d.wait.Done()

	// we use 4 * lease time to check owner's timeout, so here, we will update owner's status
	// every 2 * lease time, if lease is 0, we will use default 10s.
	checkTime := chooseLeaseTime(2*d.lease, 10*time.Second)

	ticker := time.NewTicker(checkTime)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			log.Debugf("wait %s to check DDL status again", checkTime)
		case <-d.jobCh:
		case <-d.quitCh:
			return
		}

		err := d.handleJobQueue()
		if err != nil {
			log.Errorf("handle job err %v", errors.ErrorStack(err))
		}
	}
}

func (d *ddl) runJob(t *meta.Meta, job *model.Job) {
	if job.State == model.JobDone || job.State == model.JobCancelled {
		return
	}

	job.State = model.JobRunning

	var err error
	switch job.Type {
	case model.ActionCreateSchema:
		err = d.onCreateSchema(t, job)
	case model.ActionDropSchema:
		err = d.onDropSchema(t, job)
	case model.ActionCreateTable:
		err = d.onCreateTable(t, job)
	case model.ActionDropTable:
		err = d.onDropTable(t, job)
	case model.ActionAddColumn:
		err = d.onAddColumn(t, job)
	case model.ActionDropColumn:
		err = d.onDropColumn(t, job)
	case model.ActionAddIndex:
		err = d.onCreateIndex(t, job)
	case model.ActionDropIndex:
		err = d.onDropIndex(t, job)
	case model.ActionAddConstraint:
		log.Fatal("Doesn't support change constraint online")
	case model.ActionDropConstraint:
		log.Fatal("Doesn't support change constraint online")
	default:
		// invalid job, cancel it.
		job.State = model.JobCancelled
		err = errors.Errorf("invalid job %v", job)
	}

	// saves error in job, so that others can know error happens.
	if err != nil {
		// if job is not cancelled, we should log this error.
		if job.State != model.JobCancelled {
			log.Errorf("run job err %v", errors.ErrorStack(err))
		}

		job.Error = err.Error()
		job.ErrorCount++
	}
}

// for every lease seconds, we will re-update the whole schema, so we will wait 2 * lease time
// to guarantee that all servers have already updated schema.
func (d *ddl) waitSchemaChanged() {
	if d.lease > 0 {
		select {
		case <-time.After(d.lease * 2):
		case <-d.quitCh:
		}
	}
}
<|endoftext|>


<commit_before>package deb

import (
	"github.com/smira/aptly/aptly"
	"github.com/smira/aptly/utils"
	"os"
	"path/filepath"
	"sort"
	"strings"
)

// CollectPackageFiles walks filesystem collecting all candidates for package files
func CollectPackageFiles(locations []string, reporter aptly.ResultReporter) (packageFiles, failedFiles []string) {
	for _, location := range locations {
		info, err2 := os.Stat(location)
		if err2 != nil {
			reporter.Warning("Unable to process %s: %s", location, err2)
			failedFiles = append(failedFiles, location)
			continue
		}
		if info.IsDir() {
			err2 = filepath.Walk(location, func(path string, info os.FileInfo, err3 error) error {
				if err3 != nil {
					return err3
				}
				if info.IsDir() {
					return nil
				}

				if strings.HasSuffix(info.Name(), ".deb") || strings.HasSuffix(info.Name(), ".udeb") ||
					strings.HasSuffix(info.Name(), ".dsc") {
					packageFiles = append(packageFiles, path)
				}

				return nil
			})
		} else {
			if strings.HasSuffix(info.Name(), ".deb") || strings.HasSuffix(info.Name(), ".udeb") ||
				strings.HasSuffix(info.Name(), ".dsc") {
				packageFiles = append(packageFiles, location)
			} else {
				reporter.Warning("Unknown file extension: %s", location)
				failedFiles = append(failedFiles, location)
				continue
			}
		}
	}

	sort.Strings(packageFiles)

	return
}

// ImportPackageFiles imports files into local repository
func ImportPackageFiles(list *PackageList, packageFiles []string, forceReplace bool, verifier utils.Verifier,
	pool aptly.PackagePool, collection *PackageCollection, reporter aptly.ResultReporter, restriction PackageQuery) (processedFiles []string, failedFiles []string, err error) {
	if forceReplace {
		list.PrepareIndex()
	}

	for _, file := range packageFiles {
		var (
			stanza Stanza
			p      *Package
		)

		candidateProcessedFiles := []string{}
		isSourcePackage := strings.HasSuffix(file, ".dsc")
		isUdebPackage := strings.HasSuffix(file, ".udeb")

		if isSourcePackage {
			stanza, err = GetControlFileFromDsc(file, verifier)

			if err == nil {
				stanza["Package"] = stanza["Source"]
				delete(stanza, "Source")

				p, err = NewSourcePackageFromControlFile(stanza)
			}
		} else {
			stanza, err = GetControlFileFromDeb(file)
			if isUdebPackage {
				p = NewUdebPackageFromControlFile(stanza)
			} else {
				p = NewPackageFromControlFile(stanza)
			}
		}
		if err != nil {
			reporter.Warning("Unable to read file %s: %s", file, err)
			failedFiles = append(failedFiles, file)
			continue
		}

		if p.Name == "" {
			reporter.Warning("Empty package name on %s", file)
			failedFiles = append(failedFiles, file)
			continue
		}

		if p.Version == "" {
			reporter.Warning("Empty version on %s", file)
			failedFiles = append(failedFiles, file)
			continue
		}

		if p.Architecture == "" {
			reporter.Warning("Empty architecture on %s", file)
			failedFiles = append(failedFiles, file)
			continue
		}

		var checksums utils.ChecksumInfo
		checksums, err = utils.ChecksumsForFile(file)
		if err != nil {
			return nil, nil, err
		}

		if isSourcePackage {
			p.UpdateFiles(append(p.Files(), PackageFile{Filename: filepath.Base(file), Checksums: checksums}))
		} else {
			p.UpdateFiles([]PackageFile{{Filename: filepath.Base(file), Checksums: checksums}})
		}

		err = pool.Import(file, checksums.MD5)
		if err != nil {
			reporter.Warning("Unable to import file %s into pool: %s", file, err)
			failedFiles = append(failedFiles, file)
			continue
		}

		candidateProcessedFiles = append(candidateProcessedFiles, file)

		// go over all files, except for the last one (.dsc/.deb itself)
		for _, f := range p.Files() {
			if filepath.Base(f.Filename) == filepath.Base(file) {
				continue
			}
			sourceFile := filepath.Join(filepath.Dir(file), filepath.Base(f.Filename))
			err = pool.Import(sourceFile, f.Checksums.MD5)
			if err != nil {
				reporter.Warning("Unable to import file %s into pool: %s", sourceFile, err)
				failedFiles = append(failedFiles, file)
				break
			}

			candidateProcessedFiles = append(candidateProcessedFiles, sourceFile)
		}
		if err != nil {
			// some files haven't been imported
			continue
		}

		if restriction != nil && !restriction.Matches(p) {
			reporter.Warning("%s has been ignored as it doesn't match restriction", p)
			failedFiles = append(failedFiles, file)
			continue
		}

		err = collection.Update(p)
		if err != nil {
			reporter.Warning("Unable to save package %s: %s", p, err)
			failedFiles = append(failedFiles, file)
			continue
		}

		if forceReplace {
			conflictingPackages := list.Search(Dependency{Pkg: p.Name, Version: p.Version, Relation: VersionEqual, Architecture: p.Architecture}, true)
			for _, cp := range conflictingPackages {
				reporter.Removed("%s removed due to conflict with package being added", cp)
				list.Remove(cp)
			}
		}

		err = list.Add(p)
		if err != nil {
			reporter.Warning("Unable to add package to repo %s: %s", p, err)
			failedFiles = append(failedFiles, file)
			continue
		}

		reporter.Added("%s added", p)
		processedFiles = append(processedFiles, candidateProcessedFiles...)
	}

	err = nil
	return
}
<commit_msg>Update import.go<commit_after>package deb

import (
	"github.com/smira/aptly/aptly"
	"github.com/smira/aptly/utils"
	"os"
	"path/filepath"
	"sort"
	"strings"
)

// CollectPackageFiles walks filesystem collecting all candidates for package files
func CollectPackageFiles(locations []string, reporter aptly.ResultReporter) (packageFiles, failedFiles []string) {
	for _, location := range locations {
		info, err2 := os.Stat(location)
		if err2 != nil {
			reporter.Warning("Unable to process %s: %s", location, err2)
			failedFiles = append(failedFiles, location)
			continue
		}
		if info.IsDir() {
			err2 = filepath.Walk(location, func(path string, info os.FileInfo, err3 error) error {
				if err3 != nil {
					return err3
				}
				if info.IsDir() {
					return nil
				}

				if strings.HasSuffix(info.Name(), ".deb") || strings.HasSuffix(info.Name(), ".udeb") ||
					strings.HasSuffix(info.Name(), ".dsc") || strings.HasSuffix(info.Name(), ".ddeb") {
					packageFiles = append(packageFiles, path)
				}

				return nil
			})
		} else {
			if strings.HasSuffix(info.Name(), ".deb") || strings.HasSuffix(info.Name(), ".udeb") ||
				strings.HasSuffix(info.Name(), ".dsc") || strings.HasSuffix(info.Name(), ".ddeb") {
				packageFiles = append(packageFiles, location)
			} else {
				reporter.Warning("Unknown file extension: %s", location)
				failedFiles = append(failedFiles, location)
				continue
			}
		}
	}

	sort.Strings(packageFiles)

	return
}

// ImportPackageFiles imports files into local repository
func ImportPackageFiles(list *PackageList, packageFiles []string, forceReplace bool, verifier utils.Verifier,
	pool aptly.PackagePool, collection *PackageCollection, reporter aptly.ResultReporter, restriction PackageQuery) (processedFiles []string, failedFiles []string, err error) {
	if forceReplace {
		list.PrepareIndex()
	}

	for _, file := range packageFiles {
		var (
			stanza Stanza
			p      *Package
		)

		candidateProcessedFiles := []string{}
		isSourcePackage := strings.HasSuffix(file, ".dsc")
		isUdebPackage := strings.HasSuffix(file, ".udeb")

		if isSourcePackage {
			stanza, err = GetControlFileFromDsc(file, verifier)

			if err == nil {
				stanza["Package"] = stanza["Source"]
				delete(stanza, "Source")

				p, err = NewSourcePackageFromControlFile(stanza)
			}
		} else {
			stanza, err = GetControlFileFromDeb(file)
			if isUdebPackage {
				p = NewUdebPackageFromControlFile(stanza)
			} else {
				p = NewPackageFromControlFile(stanza)
			}
		}
		if err != nil {
			reporter.Warning("Unable to read file %s: %s", file, err)
			failedFiles = append(failedFiles, file)
			continue
		}

		if p.Name == "" {
			reporter.Warning("Empty package name on %s", file)
			failedFiles = append(failedFiles, file)
			continue
		}

		if p.Version == "" {
			reporter.Warning("Empty version on %s", file)
			failedFiles = append(failedFiles, file)
			continue
		}

		if p.Architecture == "" {
			reporter.Warning("Empty architecture on %s", file)
			failedFiles = append(failedFiles, file)
			continue
		}

		var checksums utils.ChecksumInfo
		checksums, err = utils.ChecksumsForFile(file)
		if err != nil {
			return nil, nil, err
		}

		if isSourcePackage {
			p.UpdateFiles(append(p.Files(), PackageFile{Filename: filepath.Base(file), Checksums: checksums}))
		} else {
			p.UpdateFiles([]PackageFile{{Filename: filepath.Base(file), Checksums: checksums}})
		}

		err = pool.Import(file, checksums.MD5)
		if err != nil {
			reporter.Warning("Unable to import file %s into pool: %s", file, err)
			failedFiles = append(failedFiles, file)
			continue
		}

		candidateProcessedFiles = append(candidateProcessedFiles, file)

		// go over all files, except for the last one (.dsc/.deb itself)
		for _, f := range p.Files() {
			if filepath.Base(f.Filename) == filepath.Base(file) {
				continue
			}
			sourceFile := filepath.Join(filepath.Dir(file), filepath.Base(f.Filename))
			err = pool.Import(sourceFile, f.Checksums.MD5)
			if err != nil {
				reporter.Warning("Unable to import file %s into pool: %s", sourceFile, err)
				failedFiles = append(failedFiles, file)
				break
			}

			candidateProcessedFiles = append(candidateProcessedFiles, sourceFile)
		}
		if err != nil {
			// some files haven't been imported
			continue
		}

		if restriction != nil && !restriction.Matches(p) {
			reporter.Warning("%s has been ignored as it doesn't match restriction", p)
			failedFiles = append(failedFiles, file)
			continue
		}

		err = collection.Update(p)
		if err != nil {
			reporter.Warning("Unable to save package %s: %s", p, err)
			failedFiles = append(failedFiles, file)
			continue
		}

		if forceReplace {
			conflictingPackages := list.Search(Dependency{Pkg: p.Name, Version: p.Version, Relation: VersionEqual, Architecture: p.Architecture}, true)
			for _, cp := range conflictingPackages {
				reporter.Removed("%s removed due to conflict with package being added", cp)
				list.Remove(cp)
			}
		}

		err = list.Add(p)
		if err != nil {
			reporter.Warning("Unable to add package to repo %s: %s", p, err)
			failedFiles = append(failedFiles, file)
			continue
		}

		reporter.Added("%s added", p)
		processedFiles = append(processedFiles, candidateProcessedFiles...)
	}

	err = nil
	return
}
<|endoftext|>


<commit_before>/*
CompileDaemon is a very simple compile daemon for Go.

CompileDaemon watches your .go files in a directory and invokes `go build`
if a file changes.

Examples

In its simplest form, the defaults will do. With the current working directory set
to the source directory you can simply…

    $ CompileDaemon

… and it will recompile your code whenever you save a source file.

If you want it to also run your program each time it builds you might add…

    $ CompileDaemon -command="./MyProgram -my-options"

… and it will also keep a copy of your program running. Killing the old one and
starting a new one each time you build. For advanced usage you can also supply
the changed file to the command by doing…

	$ CompileDaemon -command="./MyProgram -my-options %[1]s"

…but note that this will not be set on the first start.

You may find that you need to exclude some directories and files from
monitoring, such as a .git repository or emacs temporary files…

    $ CompileDaemon -exclude-dir=.git -exclude=".#*"

If you want to monitor files other than .go and .c files you might…

    $ CompileDaemon -include=Makefile -include="*.less" -include="*.tmpl"

Options

There are command line options.

	FILE SELECTION
	-directory=XXX    – which directory to monitor for changes
	-recursive=XXX    – look into subdirectories
	-exclude-dir=XXX  – exclude directories matching glob pattern XXX
	-exlude=XXX       – exclude files whose basename matches glob pattern XXX
	-include=XXX      – include files whose basename matches glob pattern XXX
	-pattern=XXX      – include files whose path matches regexp XXX

	MISC
	-color            - enable colorized output
	-log-prefix       - Enable/disable stdout/stderr labelling for the child process
	-graceful-kill    - On supported platforms, send the child process a SIGTERM to
	                    allow it to exit gracefully if possible.
	ACTIONS
	-build=CCC        – Execute CCC to rebuild when a file changes
	-command=CCC      – Run command CCC after a successful build, stops previous command first

*/
package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"regexp"
	"strings"
	"syscall"
	"time"

	"github.com/fatih/color"
	"gopkg.in/fsnotify.v1"
)

// Milliseconds to wait for the next job to begin after a file change
const WorkDelay = 900

// Default pattern to match files which trigger a build
const FilePattern = `(.+\.go|.+\.c)$`

type globList []string

func (g *globList) String() string {
	return fmt.Sprint(*g)
}
func (g *globList) Set(value string) error {
	*g = append(*g, value)
	return nil
}
func (g *globList) Matches(value string) bool {
	for _, v := range *g {
		if match, err := filepath.Match(v, value); err != nil {
			log.Fatalf("Bad pattern \"%s\": %s", v, err.Error())
		} else if match {
			return true
		}
	}
	return false
}

var (
	flag_directory    = flag.String("directory", ".", "Directory to watch for changes")
	flag_pattern      = flag.String("pattern", FilePattern, "Pattern of watched files")
	flag_command      = flag.String("command", "", "Command to run and restart after build")
	flag_command_stop = flag.Bool("command-stop", false, "Stop command before building")
	flag_recursive    = flag.Bool("recursive", true, "Watch all dirs. recursively")
	flag_build        = flag.String("build", "go build", "Command to rebuild after changes")
	flag_build_dir    = flag.String("build-dir", "", "Directory to run build command in.  Defaults to directory")
	flag_color        = flag.Bool("color", false, "Colorize output for CompileDaemon status messages")
	flag_logprefix    = flag.Bool("log-prefix", true, "Print log timestamps and subprocess stderr/stdout output")
	flag_gracefulkill = flag.Bool("graceful-kill", false, "Gracefully attempt to kill the child process by sending a SIGTERM first")

	// initialized in main() due to custom type.
	flag_excludedDirs  globList
	flag_excludedFiles globList
	flag_includedFiles globList
)

func okColor(format string, args ...interface{}) string {
	if *flag_color {
		return color.GreenString(format, args...)
	} else {
		return fmt.Sprintf(format, args...)
	}
}

func failColor(format string, args ...interface{}) string {
	if *flag_color {
		return color.RedString(format, args...)
	} else {
		return fmt.Sprintf(format, args...)
	}
}

// Run `go build` and print the output if something's gone wrong.
func build() bool {
	log.Println(okColor("Running build command!"))

	args := strings.Split(*flag_build, " ")
	if len(args) == 0 {
		// If the user has specified and empty then we are done.
		return true
	}

	cmd := exec.Command(args[0], args[1:]...)

	if *flag_build_dir != "" {
		cmd.Dir = *flag_build_dir
	} else {
		cmd.Dir = *flag_directory
	}

	output, err := cmd.CombinedOutput()

	if err == nil {
		log.Println(okColor("Build ok."))
	} else {
		log.Println(failColor("Error while building:\n"), failColor(string(output)))
	}

	return err == nil
}

func matchesPattern(pattern *regexp.Regexp, file string) bool {
	return pattern.MatchString(file)
}

// Accept build jobs and start building when there are no jobs rushing in.
// The inrush protection is WorkDelay milliseconds long, in this period
// every incoming job will reset the timer.
func builder(jobs <-chan string, buildStarted chan<- string, buildDone chan<- bool) {
	createThreshold := func() <-chan time.Time {
		return time.After(time.Duration(WorkDelay * time.Millisecond))
	}

	threshold := createThreshold()
	eventPath := ""

	for {
		select {
		case eventPath = <-jobs:
			threshold = createThreshold()
		case <-threshold:
			buildStarted <- eventPath
			buildDone <- build()
		}
	}
}

func logger(pipeChan <-chan io.ReadCloser) {
	dumper := func(pipe io.ReadCloser, prefix string) {
		reader := bufio.NewReader(pipe)

	readloop:
		for {
			line, err := reader.ReadString('\n')

			if err != nil {
				break readloop
			}

			if *flag_logprefix {
				log.Print(prefix, " ", line)
			} else {
				log.Print(line)
			}
		}
	}

	for {
		pipe := <-pipeChan
		go dumper(pipe, "stdout:")

		pipe = <-pipeChan
		go dumper(pipe, "stderr:")
	}
}

// Start the supplied command and return stdout and stderr pipes for logging.
func startCommand(command string) (cmd *exec.Cmd, stdout io.ReadCloser, stderr io.ReadCloser, err error) {
	args := strings.Split(command, " ")
	cmd = exec.Command(args[0], args[1:]...)

	if stdout, err = cmd.StdoutPipe(); err != nil {
		err = fmt.Errorf("can't get stdout pipe for command: %s", err)
		return
	}

	if stderr, err = cmd.StderrPipe(); err != nil {
		err = fmt.Errorf("can't get stderr pipe for command: %s", err)
		return
	}

	if err = cmd.Start(); err != nil {
		err = fmt.Errorf("can't start command: %s", err)
		return
	}

	return
}

// Run the command in the given string and restart it after
// a message was received on the buildDone channel.
func runner(commandTemplate string, buildStarted <-chan string, buildSuccess <-chan bool) {
	var currentProcess *os.Process
	pipeChan := make(chan io.ReadCloser)

	go logger(pipeChan)

	go func() {
		sigChan := make(chan os.Signal, 1)
		signal.Notify(sigChan, fatalSignals...)
		<-sigChan
		log.Println(okColor("Received signal, terminating cleanly."))
		if currentProcess != nil {
			killProcess(currentProcess)
		}
		os.Exit(0)
	}()

	for {
		eventPath := <-buildStarted

		// append %0.s to use format specifier even if not supplied by user
		// to suppress warning in returned string.
		command := fmt.Sprintf("%0.s" + commandTemplate, eventPath)

		if !*flag_command_stop {
			if !<-buildSuccess {
				continue
			}
		}

		if currentProcess != nil {
			killProcess(currentProcess)
		}
		if *flag_command_stop {
			log.Println(okColor("Command stopped. Waiting for build to complete."))
			<-buildSuccess
		}

		log.Println(okColor("Restarting the given command."))
		cmd, stdoutPipe, stderrPipe, err := startCommand(command)

		if err != nil {
			log.Fatal(failColor("Could not start command:", err))
		}

		pipeChan <- stdoutPipe
		pipeChan <- stderrPipe

		currentProcess = cmd.Process
	}
}

func killProcess(process *os.Process) {
	if *flag_gracefulkill {
		killProcessGracefully(process)
	} else {
		killProcessHard(process)
	}
}

func killProcessHard(process *os.Process) {
	log.Println(okColor("Hard stopping the current process.."))

	if err := process.Kill(); err != nil {
		log.Fatal(failColor("Could not kill child process. Aborting due to danger of infinite forks."))
	}

	if _, err := process.Wait(); err != nil {
		log.Fatal(failColor("Could not wait for child process. Aborting due to danger of infinite forks."))
	}
}

func killProcessGracefully(process *os.Process) {
	done := make(chan error, 1)
	go func() {
		log.Println(okColor("Gracefully stopping the current process.."))
		if err := terminateGracefully(process); err != nil {
			done <- err
			return
		}
		_, err := process.Wait()
		done <- err
	}()

	select {
	case <-time.After(3 * time.Second):
		log.Println(failColor("Could not gracefully stop the current process, proceeding to hard stop."))
		killProcessHard(process)
		<-done
	case err := <-done:
		if err != nil {
			log.Fatal(failColor("Could not kill child process. Aborting due to danger of infinite forks."))
		}
	}
}

func flusher(buildStarted <-chan string, buildSuccess <-chan bool) {
	for {
		<-buildStarted
		<-buildSuccess
	}
}

func main() {
	flag.Var(&flag_excludedDirs, "exclude-dir", " Don't watch directories matching this name")
	flag.Var(&flag_excludedFiles, "exclude", " Don't watch files matching this name")
	flag.Var(&flag_includedFiles, "include", " Watch files matching this name")

	flag.Parse()

	if !*flag_logprefix {
		log.SetFlags(0)
	}

	if *flag_directory == "" {
		fmt.Fprintf(os.Stderr, "-directory=... is required.\n")
		os.Exit(1)
	}

	if *flag_gracefulkill && !gracefulTerminationPossible() {
		log.Fatal("Graceful termination is not supported on your platform.")
	}

	watcher, err := fsnotify.NewWatcher()

	if err != nil {
		log.Fatal(err)
	}

	defer watcher.Close()

	if *flag_recursive == true {
		err = filepath.Walk(*flag_directory, func(path string, info os.FileInfo, err error) error {
			if err == nil && info.IsDir() {
				if flag_excludedDirs.Matches(info.Name()) {
					return filepath.SkipDir
				} else {
					return watcher.Add(path)
				}
			}
			return err
		})

		if err != nil {
			log.Fatal("filepath.Walk():", err)
		}

		if err := watcher.Add(*flag_directory); err != nil {
			log.Fatal("watcher.Add():", err)
		}

	} else {
		if err := watcher.Add(*flag_directory); err != nil {
			log.Fatal("watcher.Add():", err)
		}
	}

	pattern := regexp.MustCompile(*flag_pattern)
	jobs := make(chan string)
	buildSuccess := make(chan bool)
	buildStarted := make(chan string)

	go builder(jobs, buildStarted, buildSuccess)

	if *flag_command != "" {
		go runner(*flag_command, buildStarted, buildSuccess)
	} else {
		go flusher(buildStarted, buildSuccess)
	}

	for {
		select {
		case ev := <-watcher.Events:
			if ev.Op&fsnotify.Remove == fsnotify.Remove || ev.Op&fsnotify.Write == fsnotify.Write || ev.Op&fsnotify.Create == fsnotify.Create {
				base := filepath.Base(ev.Name)

				// Assume it is a directory and track it.
				if *flag_recursive == true && !flag_excludedDirs.Matches(ev.Name) {
					watcher.Add(ev.Name)
				}

				if flag_includedFiles.Matches(base) || matchesPattern(pattern, ev.Name) {
					if !flag_excludedFiles.Matches(base) {
						jobs <- ev.Name
					}
				}
			}

		case err := <-watcher.Errors:
			if v, ok := err.(*os.SyscallError); ok {
				if v.Err == syscall.EINTR {
					continue
				}
				log.Fatal("watcher.Error: SyscallError:", v)
			}
			log.Fatal("watcher.Error:", err)
		}
	}
}
<commit_msg>:pencil2: fix invalid format<commit_after>/*
CompileDaemon is a very simple compile daemon for Go.

CompileDaemon watches your .go files in a directory and invokes `go build`
if a file changes.

Examples

In its simplest form, the defaults will do. With the current working directory set
to the source directory you can simply…

    $ CompileDaemon

… and it will recompile your code whenever you save a source file.

If you want it to also run your program each time it builds you might add…

    $ CompileDaemon -command="./MyProgram -my-options"

… and it will also keep a copy of your program running. Killing the old one and
starting a new one each time you build. For advanced usage you can also supply
the changed file to the command by doing…

	$ CompileDaemon -command="./MyProgram -my-options %[1]s"

…but note that this will not be set on the first start.

You may find that you need to exclude some directories and files from
monitoring, such as a .git repository or emacs temporary files…

    $ CompileDaemon -exclude-dir=.git -exclude=".#*"

If you want to monitor files other than .go and .c files you might…

    $ CompileDaemon -include=Makefile -include="*.less" -include="*.tmpl"

Options

There are command line options.

	FILE SELECTION
	-directory=XXX    – which directory to monitor for changes
	-recursive=XXX    – look into subdirectories
	-exclude-dir=XXX  – exclude directories matching glob pattern XXX
	-exlude=XXX       – exclude files whose basename matches glob pattern XXX
	-include=XXX      – include files whose basename matches glob pattern XXX
	-pattern=XXX      – include files whose path matches regexp XXX

	MISC
	-color            - enable colorized output
	-log-prefix       - Enable/disable stdout/stderr labelling for the child process
	-graceful-kill    - On supported platforms, send the child process a SIGTERM to
	                    allow it to exit gracefully if possible.
	ACTIONS
	-build=CCC        – Execute CCC to rebuild when a file changes
	-command=CCC      – Run command CCC after a successful build, stops previous command first

*/
package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"regexp"
	"strings"
	"syscall"
	"time"

	"github.com/fatih/color"
	"gopkg.in/fsnotify.v1"
)

// Milliseconds to wait for the next job to begin after a file change
const WorkDelay = 900

// Default pattern to match files which trigger a build
const FilePattern = `(.+\.go|.+\.c)$`

type globList []string

func (g *globList) String() string {
	return fmt.Sprint(*g)
}
func (g *globList) Set(value string) error {
	*g = append(*g, value)
	return nil
}
func (g *globList) Matches(value string) bool {
	for _, v := range *g {
		if match, err := filepath.Match(v, value); err != nil {
			log.Fatalf("Bad pattern \"%s\": %s", v, err.Error())
		} else if match {
			return true
		}
	}
	return false
}

var (
	flag_directory    = flag.String("directory", ".", "Directory to watch for changes")
	flag_pattern      = flag.String("pattern", FilePattern, "Pattern of watched files")
	flag_command      = flag.String("command", "", "Command to run and restart after build")
	flag_command_stop = flag.Bool("command-stop", false, "Stop command before building")
	flag_recursive    = flag.Bool("recursive", true, "Watch all dirs. recursively")
	flag_build        = flag.String("build", "go build", "Command to rebuild after changes")
	flag_build_dir    = flag.String("build-dir", "", "Directory to run build command in.  Defaults to directory")
	flag_color        = flag.Bool("color", false, "Colorize output for CompileDaemon status messages")
	flag_logprefix    = flag.Bool("log-prefix", true, "Print log timestamps and subprocess stderr/stdout output")
	flag_gracefulkill = flag.Bool("graceful-kill", false, "Gracefully attempt to kill the child process by sending a SIGTERM first")

	// initialized in main() due to custom type.
	flag_excludedDirs  globList
	flag_excludedFiles globList
	flag_includedFiles globList
)

func okColor(format string, args ...interface{}) string {
	if *flag_color {
		return color.GreenString(format, args...)
	} else {
		return fmt.Sprintf(format, args...)
	}
}

func failColor(format string, args ...interface{}) string {
	if *flag_color {
		return color.RedString(format, args...)
	} else {
		return fmt.Sprintf(format, args...)
	}
}

// Run `go build` and print the output if something's gone wrong.
func build() bool {
	log.Println(okColor("Running build command!"))

	args := strings.Split(*flag_build, " ")
	if len(args) == 0 {
		// If the user has specified and empty then we are done.
		return true
	}

	cmd := exec.Command(args[0], args[1:]...)

	if *flag_build_dir != "" {
		cmd.Dir = *flag_build_dir
	} else {
		cmd.Dir = *flag_directory
	}

	output, err := cmd.CombinedOutput()

	if err == nil {
		log.Println(okColor("Build ok."))
	} else {
		log.Println(failColor("Error while building:\n"), failColor(string(output)))
	}

	return err == nil
}

func matchesPattern(pattern *regexp.Regexp, file string) bool {
	return pattern.MatchString(file)
}

// Accept build jobs and start building when there are no jobs rushing in.
// The inrush protection is WorkDelay milliseconds long, in this period
// every incoming job will reset the timer.
func builder(jobs <-chan string, buildStarted chan<- string, buildDone chan<- bool) {
	createThreshold := func() <-chan time.Time {
		return time.After(time.Duration(WorkDelay * time.Millisecond))
	}

	threshold := createThreshold()
	eventPath := ""

	for {
		select {
		case eventPath = <-jobs:
			threshold = createThreshold()
		case <-threshold:
			buildStarted <- eventPath
			buildDone <- build()
		}
	}
}

func logger(pipeChan <-chan io.ReadCloser) {
	dumper := func(pipe io.ReadCloser, prefix string) {
		reader := bufio.NewReader(pipe)

	readloop:
		for {
			line, err := reader.ReadString('\n')

			if err != nil {
				break readloop
			}

			if *flag_logprefix {
				log.Print(prefix, " ", line)
			} else {
				log.Print(line)
			}
		}
	}

	for {
		pipe := <-pipeChan
		go dumper(pipe, "stdout:")

		pipe = <-pipeChan
		go dumper(pipe, "stderr:")
	}
}

// Start the supplied command and return stdout and stderr pipes for logging.
func startCommand(command string) (cmd *exec.Cmd, stdout io.ReadCloser, stderr io.ReadCloser, err error) {
	args := strings.Split(command, " ")
	cmd = exec.Command(args[0], args[1:]...)

	if stdout, err = cmd.StdoutPipe(); err != nil {
		err = fmt.Errorf("can't get stdout pipe for command: %s", err)
		return
	}

	if stderr, err = cmd.StderrPipe(); err != nil {
		err = fmt.Errorf("can't get stderr pipe for command: %s", err)
		return
	}

	if err = cmd.Start(); err != nil {
		err = fmt.Errorf("can't start command: %s", err)
		return
	}

	return
}

// Run the command in the given string and restart it after
// a message was received on the buildDone channel.
func runner(commandTemplate string, buildStarted <-chan string, buildSuccess <-chan bool) {
	var currentProcess *os.Process
	pipeChan := make(chan io.ReadCloser)

	go logger(pipeChan)

	go func() {
		sigChan := make(chan os.Signal, 1)
		signal.Notify(sigChan, fatalSignals...)
		<-sigChan
		log.Println(okColor("Received signal, terminating cleanly."))
		if currentProcess != nil {
			killProcess(currentProcess)
		}
		os.Exit(0)
	}()

	for {
		eventPath := <-buildStarted

		// append %0.s to use format specifier even if not supplied by user
		// to suppress warning in returned string.
		command := fmt.Sprintf("%0.s" + commandTemplate, eventPath)

		if !*flag_command_stop {
			if !<-buildSuccess {
				continue
			}
		}

		if currentProcess != nil {
			killProcess(currentProcess)
		}
		if *flag_command_stop {
			log.Println(okColor("Command stopped. Waiting for build to complete."))
			<-buildSuccess
		}

		log.Println(okColor("Restarting the given command."))
		cmd, stdoutPipe, stderrPipe, err := startCommand(command)

		if err != nil {
			log.Fatal(failColor("Could not start command: %s", err))
		}

		pipeChan <- stdoutPipe
		pipeChan <- stderrPipe

		currentProcess = cmd.Process
	}
}

func killProcess(process *os.Process) {
	if *flag_gracefulkill {
		killProcessGracefully(process)
	} else {
		killProcessHard(process)
	}
}

func killProcessHard(process *os.Process) {
	log.Println(okColor("Hard stopping the current process.."))

	if err := process.Kill(); err != nil {
		log.Fatal(failColor("Could not kill child process. Aborting due to danger of infinite forks."))
	}

	if _, err := process.Wait(); err != nil {
		log.Fatal(failColor("Could not wait for child process. Aborting due to danger of infinite forks."))
	}
}

func killProcessGracefully(process *os.Process) {
	done := make(chan error, 1)
	go func() {
		log.Println(okColor("Gracefully stopping the current process.."))
		if err := terminateGracefully(process); err != nil {
			done <- err
			return
		}
		_, err := process.Wait()
		done <- err
	}()

	select {
	case <-time.After(3 * time.Second):
		log.Println(failColor("Could not gracefully stop the current process, proceeding to hard stop."))
		killProcessHard(process)
		<-done
	case err := <-done:
		if err != nil {
			log.Fatal(failColor("Could not kill child process. Aborting due to danger of infinite forks."))
		}
	}
}

func flusher(buildStarted <-chan string, buildSuccess <-chan bool) {
	for {
		<-buildStarted
		<-buildSuccess
	}
}

func main() {
	flag.Var(&flag_excludedDirs, "exclude-dir", " Don't watch directories matching this name")
	flag.Var(&flag_excludedFiles, "exclude", " Don't watch files matching this name")
	flag.Var(&flag_includedFiles, "include", " Watch files matching this name")

	flag.Parse()

	if !*flag_logprefix {
		log.SetFlags(0)
	}

	if *flag_directory == "" {
		fmt.Fprintf(os.Stderr, "-directory=... is required.\n")
		os.Exit(1)
	}

	if *flag_gracefulkill && !gracefulTerminationPossible() {
		log.Fatal("Graceful termination is not supported on your platform.")
	}

	watcher, err := fsnotify.NewWatcher()

	if err != nil {
		log.Fatal(err)
	}

	defer watcher.Close()

	if *flag_recursive == true {
		err = filepath.Walk(*flag_directory, func(path string, info os.FileInfo, err error) error {
			if err == nil && info.IsDir() {
				if flag_excludedDirs.Matches(info.Name()) {
					return filepath.SkipDir
				} else {
					return watcher.Add(path)
				}
			}
			return err
		})

		if err != nil {
			log.Fatal("filepath.Walk():", err)
		}

		if err := watcher.Add(*flag_directory); err != nil {
			log.Fatal("watcher.Add():", err)
		}

	} else {
		if err := watcher.Add(*flag_directory); err != nil {
			log.Fatal("watcher.Add():", err)
		}
	}

	pattern := regexp.MustCompile(*flag_pattern)
	jobs := make(chan string)
	buildSuccess := make(chan bool)
	buildStarted := make(chan string)

	go builder(jobs, buildStarted, buildSuccess)

	if *flag_command != "" {
		go runner(*flag_command, buildStarted, buildSuccess)
	} else {
		go flusher(buildStarted, buildSuccess)
	}

	for {
		select {
		case ev := <-watcher.Events:
			if ev.Op&fsnotify.Remove == fsnotify.Remove || ev.Op&fsnotify.Write == fsnotify.Write || ev.Op&fsnotify.Create == fsnotify.Create {
				base := filepath.Base(ev.Name)

				// Assume it is a directory and track it.
				if *flag_recursive == true && !flag_excludedDirs.Matches(ev.Name) {
					watcher.Add(ev.Name)
				}

				if flag_includedFiles.Matches(base) || matchesPattern(pattern, ev.Name) {
					if !flag_excludedFiles.Matches(base) {
						jobs <- ev.Name
					}
				}
			}

		case err := <-watcher.Errors:
			if v, ok := err.(*os.SyscallError); ok {
				if v.Err == syscall.EINTR {
					continue
				}
				log.Fatal("watcher.Error: SyscallError:", v)
			}
			log.Fatal("watcher.Error:", err)
		}
	}
}
<|endoftext|>


<commit_before>/*
CompileDaemon is a very simple compile daemon for Go.

CompileDaemon watches your .go files in a directory and invokes `go build`
if a file changes.

Examples

In its simplest form, the defaults will do. With the current working directory set
to the source directory you can simply…

    $ CompileDaemon

… and it will recompile your code whenever you save a source file.

If you want it to also run your program each time it builds you might add…

    $ CompileDaemon -command="./MyProgram -my-options"

… and it will also keep a copy of your program running. Killing the old one and
starting a new one each time you build. For advanced usage you can also supply
the changed file to the command by doing…

	$ CompileDaemon -command="./MyProgram -my-options %[1]s"

…but note that this will not be set on the first start.

You may find that you need to exclude some directories and files from
monitoring, such as a .git repository or emacs temporary files…

    $ CompileDaemon -exclude-dir=.git -exclude=".#*"

If you want to monitor files other than .go and .c files you might…

    $ CompileDaemon -include=Makefile -include="*.less" -include="*.tmpl"

Options

There are command line options.

	FILE SELECTION
	-directory=XXX    – Which directory to monitor for changes
	-recursive=XXX    – Look into subdirectories
	-exclude-dir=XXX  – Exclude directories matching glob pattern XXX
	-exlude=XXX       – Exclude files whose basename matches glob pattern XXX
	-include=XXX      – Include files whose basename matches glob pattern XXX
	-pattern=XXX      – Include files whose path matches regexp XXX

	FILE WATCH
	-polling          - Use polling instead of FS notifications to detect changes. Default is false
	-polling-interval - Milliseconds of interval between polling file changes when polling option is selected

	MISC
	-color            - Enable colorized output
	-log-prefix       - Enable/disable stdout/stderr labelling for the child process
	-graceful-kill    - On supported platforms, send the child process a SIGTERM to
	                    allow it to exit gracefully if possible.
	-graceful-timeout - Duration (in seconds) to wait for graceful kill to complete
	-verbose          - Print information about watched directories.

	ACTIONS
	-build=CCC        – Execute CCC to rebuild when a file changes
	-command=CCC      – Run command CCC after a successful build, stops previous command first

*/
package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/fatih/color"
)

// Milliseconds to wait for the next job to begin after a file change
const WorkDelay = 900

// Default pattern to match files which trigger a build
const FilePattern = `(.+\.go|.+\.c)$`

type globList []string

func (g *globList) String() string {
	return fmt.Sprint(*g)
}
func (g *globList) Set(value string) error {
	*g = append(*g, filepath.Clean(value))
	return nil
}
func (g *globList) Matches(value string) bool {
	for _, v := range *g {
		if match, err := filepath.Match(v, value); err != nil {
			log.Fatalf("Bad pattern \"%s\": %s", v, err.Error())
		} else if match {
			return true
		}
	}
	return false
}

var (
	flagPattern         = flag.String("pattern", FilePattern, "Pattern of watched files")
	flagCommand         = flag.String("command", "", "Command to run and restart after build")
	flagCommandStop     = flag.Bool("command-stop", false, "Stop command before building")
	flagRecursive       = flag.Bool("recursive", true, "Watch all dirs. recursively")
	flagBuild           = flag.String("build", "go build", "Command to rebuild after changes")
	flagBuildDir        = flag.String("build-dir", "", "Directory to run build command in.  Defaults to directory")
	flagRunDir          = flag.String("run-dir", "", "Directory to run command in.  Defaults to directory")
	flagColor           = flag.Bool("color", false, "Colorize output for CompileDaemon status messages")
	flagLogPrefix       = flag.Bool("log-prefix", true, "Print log timestamps and subprocess stderr/stdout output")
	flagGracefulKill    = flag.Bool("graceful-kill", false, "Gracefully attempt to kill the child process by sending a SIGTERM first")
	flagGracefulTimeout = flag.Uint("graceful-timeout", 3, "Duration (in seconds) to wait for graceful kill to complete")
	flagVerbose         = flag.Bool("verbose", false, "Be verbose about which directories are watched.")
	flagPolling         = flag.Bool("polling", false, "Use polling method to watch file change instead of fsnotify")
	flagPollingInterval = flag.Int("polling-interval", 100, "Milliseconds of interval between polling file changes when polling option is selected")

	// initialized in main() due to custom type.
	flagDirectories   globList
	flagExcludedDirs  globList
	flagExcludedFiles globList
	flagIncludedFiles globList
)

func okColor(format string, args ...interface{}) string {
	if *flagColor {
		return color.GreenString(format, args...)
	} else {
		return fmt.Sprintf(format, args...)
	}
}

func failColor(format string, args ...interface{}) string {
	if *flagColor {
		return color.RedString(format, args...)
	} else {
		return fmt.Sprintf(format, args...)
	}
}

// Run `go build` and print the output if something's gone wrong.
func build() bool {
	log.Println(okColor("Running build command!"))

	args := strings.Split(*flagBuild, " ")
	if len(args) == 0 {
		// If the user has specified and empty then we are done.
		return true
	}

	cmd := exec.Command(args[0], args[1:]...)

	if *flagBuildDir != "" {
		cmd.Dir = *flagBuildDir
	} else if len(flagDirectories) > 0 {
		cmd.Dir = flagDirectories[0]
	}

	output, err := cmd.CombinedOutput()

	if err == nil {
		log.Println(okColor("Build ok."))
	} else {
		log.Println(failColor("Error while building:\n"), failColor(string(output)))
	}

	return err == nil
}

func matchesPattern(pattern *regexp.Regexp, file string) bool {
	return pattern.MatchString(file)
}

// Accept build jobs and start building when there are no jobs rushing in.
// The inrush protection is WorkDelay milliseconds long, in this period
// every incoming job will reset the timer.
func builder(jobs <-chan string, buildStarted chan<- string, buildDone chan<- bool) {
	createThreshold := func() <-chan time.Time {
		return time.After(time.Duration(WorkDelay * time.Millisecond))
	}

	threshold := createThreshold()
	eventPath := ""

	for {
		select {
		case eventPath = <-jobs:
			threshold = createThreshold()
		case <-threshold:
			buildStarted <- eventPath
			buildDone <- build()
		}
	}
}

func logger(pipeChan <-chan io.ReadCloser) {
	dumper := func(pipe io.ReadCloser, prefix string) {
		reader := bufio.NewReader(pipe)

	readloop:
		for {
			line, err := reader.ReadString('\n')

			if err != nil {
				break readloop
			}

			if *flagLogPrefix {
				log.Print(prefix, " ", line)
			} else {
				log.Print(line)
			}
		}
	}

	for {
		pipe := <-pipeChan
		go dumper(pipe, "stdout:")

		pipe = <-pipeChan
		go dumper(pipe, "stderr:")
	}
}

// Start the supplied command and return stdout and stderr pipes for logging.
func startCommand(command string) (cmd *exec.Cmd, stdout io.ReadCloser, stderr io.ReadCloser, err error) {
	args := strings.Split(command, " ")
	cmd = exec.Command(args[0], args[1:]...)

	if *flagRunDir != "" {
		cmd.Dir = *flagRunDir
	}

	if stdout, err = cmd.StdoutPipe(); err != nil {
		err = fmt.Errorf("can't get stdout pipe for command: %s", err)
		return
	}

	if stderr, err = cmd.StderrPipe(); err != nil {
		err = fmt.Errorf("can't get stderr pipe for command: %s", err)
		return
	}

	if err = cmd.Start(); err != nil {
		err = fmt.Errorf("can't start command: %s", err)
		return
	}

	return
}

// Run the command in the given string and restart it after
// a message was received on the buildDone channel.
func runner(commandTemplate string, buildStarted <-chan string, buildSuccess <-chan bool) {
	var currentProcess *os.Process
	pipeChan := make(chan io.ReadCloser)

	go logger(pipeChan)

	// Launch concurrent process watching for signals from outside that
	// indicate termination to kill the running process alongside
	// CompileDaemon to prevent orphan processes.
	go func() {
		processSignalChannel := make(chan os.Signal, 1)
		signal.Notify(processSignalChannel, fatalSignals...)
		<-processSignalChannel

		log.Println(okColor("Received signal, terminating cleanly."))
		if currentProcess != nil {
			killProcess(currentProcess)
		}
		os.Exit(0)
	}()

	for {
		eventPath := <-buildStarted

		// prepend %0.s (which adds nothing) to prevent warning about missing
		// format specifier if the user did not supply one.
		command := fmt.Sprintf("%0.s"+commandTemplate, eventPath)

		if !*flagCommandStop {
			if !<-buildSuccess {
				continue
			}
		}

		if currentProcess != nil {
			killProcess(currentProcess)
		}

		if *flagCommandStop {
			log.Println(okColor("Command stopped. Waiting for build to complete."))
			if !<-buildSuccess {
				continue
			}
		}

		log.Println(okColor("Restarting the given command."))
		cmd, stdoutPipe, stderrPipe, err := startCommand(command)

		if err != nil {
			log.Fatal(failColor("Could not start command: %s", err))
		}

		pipeChan <- stdoutPipe
		pipeChan <- stderrPipe

		currentProcess = cmd.Process
	}
}

func killProcess(process *os.Process) {
	if *flagGracefulKill {
		killProcessGracefully(process)
	} else {
		killProcessHard(process)
	}
}

func killProcessHard(process *os.Process) {
	log.Println(okColor("Hard stopping the current process.."))

	if err := process.Kill(); err != nil {
		log.Println(failColor("Warning: could not kill child process.  It may have already exited."))
	}

	if _, err := process.Wait(); err != nil {
		log.Fatal(failColor("Could not wait for child process. Aborting due to danger of infinite forks."))
	}
}

func killProcessGracefully(process *os.Process) {
	done := make(chan error, 1)
	go func() {
		log.Println(okColor("Gracefully stopping the current process.."))
		if err := terminateGracefully(process); err != nil {
			done <- err
			return
		}
		_, err := process.Wait()
		done <- err
	}()

	select {
	case <-time.After(time.Duration(*flagGracefulTimeout) * time.Second):
		log.Println(failColor("Could not gracefully stop the current process, proceeding to hard stop."))
		killProcessHard(process)
		<-done
	case err := <-done:
		if err != nil {
			log.Fatal(failColor("Could not kill child process. Aborting due to danger of infinite forks."))
		}
	}
}

func flusher(buildStarted <-chan string, buildSuccess <-chan bool) {
	for {
		<-buildStarted
		<-buildSuccess
	}
}

func main() {
	flag.Var(&flagDirectories, "directory", "Directory to watch for changes, can be set more than once")
	flag.Var(&flagExcludedDirs, "exclude-dir", " Don't watch directories matching this name, can be set more than once")
	flag.Var(&flagExcludedFiles, "exclude", " Don't watch files matching this name, can be set more than once")
	flag.Var(&flagIncludedFiles, "include", " Watch files matching this name, can be set more than once")

	flag.Parse()

	if !*flagLogPrefix {
		log.SetFlags(0)
	}

	if len(flagDirectories) == 0 {
		flagDirectories = globList([]string{"."})
	}

	if *flagGracefulKill && !gracefulTerminationPossible() {
		log.Fatal("Graceful termination is not supported on your platform.")
	}

	pattern := regexp.MustCompile(*flagPattern)

	cfg := &WatcherConfig{
		flagVerbose:         *flagVerbose,
		flagRecursive:       *flagRecursive,
		flagPolling:         *flagPolling,
		flagPollingInterval: *flagPollingInterval,
		flagDirectories:     flagDirectories,
		flagExcludedDirs:    flagExcludedDirs,
		flagExcludedFiles:   flagExcludedFiles,
		flagIncludedFiles:   flagIncludedFiles,
		pattern:             pattern,
	}
	watcher, err := NewWatcher(cfg)

	if err != nil {
		log.Fatal(err)
	}

	defer watcher.Close()

	err = watcher.AddFiles()
	if err != nil {
		log.Fatal("watcher.Addfiles():", err)
	}

	jobs := make(chan string)
	buildSuccess := make(chan bool)
	buildStarted := make(chan string)

	go builder(jobs, buildStarted, buildSuccess)

	if *flagCommand != "" {
		go runner(*flagCommand, buildStarted, buildSuccess)
	} else {
		go flusher(buildStarted, buildSuccess)
	}

	watcher.Watch(jobs) // start watching files
}
<commit_msg>allow multiple commands to run on build (#64)<commit_after>/*
CompileDaemon is a very simple compile daemon for Go.

CompileDaemon watches your .go files in a directory and invokes `go build`
if a file changes.

Examples

In its simplest form, the defaults will do. With the current working directory set
to the source directory you can simply…

    $ CompileDaemon

… and it will recompile your code whenever you save a source file.

If you want it to also run your program each time it builds you might add…

    $ CompileDaemon -command="./MyProgram -my-options"

… and it will also keep a copy of your program running. Killing the old one and
starting a new one each time you build. For advanced usage you can also supply
the changed file to the command by doing…

	$ CompileDaemon -command="./MyProgram -my-options %[1]s"

…but note that this will not be set on the first start.

You may find that you need to exclude some directories and files from
monitoring, such as a .git repository or emacs temporary files…

    $ CompileDaemon -exclude-dir=.git -exclude=".#*"

If you want to monitor files other than .go and .c files you might…

    $ CompileDaemon -include=Makefile -include="*.less" -include="*.tmpl"

Options

There are command line options.

	FILE SELECTION
	-directory=XXX    – Which directory to monitor for changes
	-recursive=XXX    – Look into subdirectories
	-exclude-dir=XXX  – Exclude directories matching glob pattern XXX
	-exlude=XXX       – Exclude files whose basename matches glob pattern XXX
	-include=XXX      – Include files whose basename matches glob pattern XXX
	-pattern=XXX      – Include files whose path matches regexp XXX

	FILE WATCH
	-polling          - Use polling instead of FS notifications to detect changes. Default is false
	-polling-interval - Milliseconds of interval between polling file changes when polling option is selected

	MISC
	-color            - Enable colorized output
	-log-prefix       - Enable/disable stdout/stderr labelling for the child process
	-graceful-kill    - On supported platforms, send the child process a SIGTERM to
	                    allow it to exit gracefully if possible.
	-graceful-timeout - Duration (in seconds) to wait for graceful kill to complete
	-verbose          - Print information about watched directories.

	ACTIONS
	-build=CCC        – Execute CCC to rebuild when a file changes
	-command=CCC      – Run command CCC after a successful build, stops previous command first

*/
package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/fatih/color"
)

// Milliseconds to wait for the next job to begin after a file change
const WorkDelay = 900

// Default pattern to match files which trigger a build
const FilePattern = `(.+\.go|.+\.c)$`

type globList []string

func (g *globList) String() string {
	return fmt.Sprint(*g)
}
func (g *globList) Set(value string) error {
	*g = append(*g, filepath.Clean(value))
	return nil
}
func (g *globList) Matches(value string) bool {
	for _, v := range *g {
		if match, err := filepath.Match(v, value); err != nil {
			log.Fatalf("Bad pattern \"%s\": %s", v, err.Error())
		} else if match {
			return true
		}
	}
	return false
}

type buildCommandList struct {
	commands []string
}

func (f *buildCommandList) String() string {
	return ""
}

func (f *buildCommandList) Set(s string) error {
	f.commands = append(f.commands, s)
	return nil
}

var (
	flagPattern         = flag.String("pattern", FilePattern, "Pattern of watched files")
	flagCommand         = flag.String("command", "", "Command to run and restart after build")
	flagCommandStop     = flag.Bool("command-stop", false, "Stop command before building")
	flagRecursive       = flag.Bool("recursive", true, "Watch all dirs. recursively")
	flagBuildDir        = flag.String("build-dir", "", "Directory to run build command in.  Defaults to directory")
	flagRunDir          = flag.String("run-dir", "", "Directory to run command in.  Defaults to directory")
	flagColor           = flag.Bool("color", false, "Colorize output for CompileDaemon status messages")
	flagLogPrefix       = flag.Bool("log-prefix", true, "Print log timestamps and subprocess stderr/stdout output")
	flagGracefulKill    = flag.Bool("graceful-kill", false, "Gracefully attempt to kill the child process by sending a SIGTERM first")
	flagGracefulTimeout = flag.Uint("graceful-timeout", 3, "Duration (in seconds) to wait for graceful kill to complete")
	flagVerbose         = flag.Bool("verbose", false, "Be verbose about which directories are watched.")
	flagPolling         = flag.Bool("polling", false, "Use polling method to watch file change instead of fsnotify")
	flagPollingInterval = flag.Int("polling-interval", 100, "Milliseconds of interval between polling file changes when polling option is selected")

	// initialized in main() due to custom type.
	flagDirectories      globList
	flagExcludedDirs     globList
	flagExcludedFiles    globList
	flagIncludedFiles    globList
	flagBuildCommandList buildCommandList
)

func okColor(format string, args ...interface{}) string {
	if *flagColor {
		return color.GreenString(format, args...)
	} else {
		return fmt.Sprintf(format, args...)
	}
}

func failColor(format string, args ...interface{}) string {
	if *flagColor {
		return color.RedString(format, args...)
	} else {
		return fmt.Sprintf(format, args...)
	}
}

// Run `go build` and print the output if something's gone wrong.
func build() bool {
	log.Println(okColor("Running build commands!"))

	for _, c := range flagBuildCommandList.commands {
		err := runBuildCommand(c)
		if err != nil {
			log.Println(failColor("Command failed: "), failColor(c))
			return false
		}
	}
	log.Println(okColor("Build ok."))

	return true
}

func runBuildCommand(c string) error {
	c = strings.TrimSpace(c)
	args := strings.Split(c, " ")
	if len(args) == 0 {
		return nil
	}

	cmd := exec.Command(args[0], args[1:]...)

	if *flagBuildDir != "" {
		cmd.Dir = *flagBuildDir
	} else if len(flagDirectories) > 0 {
		cmd.Dir = flagDirectories[0]
	}

	output, err := cmd.CombinedOutput()

	if err != nil {
		log.Println(failColor("Error while building:\n"), failColor(string(output)))
		return err
	}
	return nil
}

func matchesPattern(pattern *regexp.Regexp, file string) bool {
	return pattern.MatchString(file)
}

// Accept build jobs and start building when there are no jobs rushing in.
// The inrush protection is WorkDelay milliseconds long, in this period
// every incoming job will reset the timer.
func builder(jobs <-chan string, buildStarted chan<- string, buildDone chan<- bool) {
	createThreshold := func() <-chan time.Time {
		return time.After(time.Duration(WorkDelay * time.Millisecond))
	}

	threshold := createThreshold()
	eventPath := ""

	for {
		select {
		case eventPath = <-jobs:
			threshold = createThreshold()
		case <-threshold:
			buildStarted <- eventPath
			buildDone <- build()
		}
	}
}

func logger(pipeChan <-chan io.ReadCloser) {
	dumper := func(pipe io.ReadCloser, prefix string) {
		reader := bufio.NewReader(pipe)

	readloop:
		for {
			line, err := reader.ReadString('\n')

			if err != nil {
				break readloop
			}

			if *flagLogPrefix {
				log.Print(prefix, " ", line)
			} else {
				log.Print(line)
			}
		}
	}

	for {
		pipe := <-pipeChan
		go dumper(pipe, "stdout:")

		pipe = <-pipeChan
		go dumper(pipe, "stderr:")
	}
}

// Start the supplied command and return stdout and stderr pipes for logging.
func startCommand(command string) (cmd *exec.Cmd, stdout io.ReadCloser, stderr io.ReadCloser, err error) {
	args := strings.Split(command, " ")
	cmd = exec.Command(args[0], args[1:]...)

	if *flagRunDir != "" {
		cmd.Dir = *flagRunDir
	}

	if stdout, err = cmd.StdoutPipe(); err != nil {
		err = fmt.Errorf("can't get stdout pipe for command: %s", err)
		return
	}

	if stderr, err = cmd.StderrPipe(); err != nil {
		err = fmt.Errorf("can't get stderr pipe for command: %s", err)
		return
	}

	if err = cmd.Start(); err != nil {
		err = fmt.Errorf("can't start command: %s", err)
		return
	}

	return
}

// Run the command in the given string and restart it after
// a message was received on the buildDone channel.
func runner(commandTemplate string, buildStarted <-chan string, buildSuccess <-chan bool) {
	var currentProcess *os.Process
	pipeChan := make(chan io.ReadCloser)

	go logger(pipeChan)

	// Launch concurrent process watching for signals from outside that
	// indicate termination to kill the running process alongside
	// CompileDaemon to prevent orphan processes.
	go func() {
		processSignalChannel := make(chan os.Signal, 1)
		signal.Notify(processSignalChannel, fatalSignals...)
		<-processSignalChannel

		log.Println(okColor("Received signal, terminating cleanly."))
		if currentProcess != nil {
			killProcess(currentProcess)
		}
		os.Exit(0)
	}()

	for {
		eventPath := <-buildStarted

		// prepend %0.s (which adds nothing) to prevent warning about missing
		// format specifier if the user did not supply one.
		command := fmt.Sprintf("%0.s"+commandTemplate, eventPath)

		if !*flagCommandStop {
			if !<-buildSuccess {
				continue
			}
		}

		if currentProcess != nil {
			killProcess(currentProcess)
		}

		if *flagCommandStop {
			log.Println(okColor("Command stopped. Waiting for build to complete."))
			if !<-buildSuccess {
				continue
			}
		}

		log.Println(okColor("Restarting the given command."))
		cmd, stdoutPipe, stderrPipe, err := startCommand(command)

		if err != nil {
			log.Fatal(failColor("Could not start command: %s", err))
		}

		pipeChan <- stdoutPipe
		pipeChan <- stderrPipe

		currentProcess = cmd.Process
	}
}

func killProcess(process *os.Process) {
	if *flagGracefulKill {
		killProcessGracefully(process)
	} else {
		killProcessHard(process)
	}
}

func killProcessHard(process *os.Process) {
	log.Println(okColor("Hard stopping the current process.."))

	if err := process.Kill(); err != nil {
		log.Println(failColor("Warning: could not kill child process.  It may have already exited."))
	}

	if _, err := process.Wait(); err != nil {
		log.Fatal(failColor("Could not wait for child process. Aborting due to danger of infinite forks."))
	}
}

func killProcessGracefully(process *os.Process) {
	done := make(chan error, 1)
	go func() {
		log.Println(okColor("Gracefully stopping the current process.."))
		if err := terminateGracefully(process); err != nil {
			done <- err
			return
		}
		_, err := process.Wait()
		done <- err
	}()

	select {
	case <-time.After(time.Duration(*flagGracefulTimeout) * time.Second):
		log.Println(failColor("Could not gracefully stop the current process, proceeding to hard stop."))
		killProcessHard(process)
		<-done
	case err := <-done:
		if err != nil {
			log.Fatal(failColor("Could not kill child process. Aborting due to danger of infinite forks."))
		}
	}
}

func flusher(buildStarted <-chan string, buildSuccess <-chan bool) {
	for {
		<-buildStarted
		<-buildSuccess
	}
}

func main() {
	flag.Var(&flagDirectories, "directory", "Directory to watch for changes, can be set more than once")
	flag.Var(&flagExcludedDirs, "exclude-dir", " Don't watch directories matching this name, can be set more than once")
	flag.Var(&flagExcludedFiles, "exclude", " Don't watch files matching this name, can be set more than once")
	flag.Var(&flagIncludedFiles, "include", " Watch files matching this name, can be set more than once")
	flag.Var(&flagBuildCommandList, "build", "Command to rebuild after changes, can be set more than once. Defaults to 'go build'.")

	flag.Parse()

	if !*flagLogPrefix {
		log.SetFlags(0)
	}

	if len(flagDirectories) == 0 {
		flagDirectories = globList([]string{"."})
	}

	if *flagGracefulKill && !gracefulTerminationPossible() {
		log.Fatal("Graceful termination is not supported on your platform.")
	}

	if len(flagBuildCommandList.commands) == 0 {
		_ = flagBuildCommandList.Set("go build")
	}

	pattern := regexp.MustCompile(*flagPattern)

	cfg := &WatcherConfig{
		flagVerbose:         *flagVerbose,
		flagRecursive:       *flagRecursive,
		flagPolling:         *flagPolling,
		flagPollingInterval: *flagPollingInterval,
		flagDirectories:     flagDirectories,
		flagExcludedDirs:    flagExcludedDirs,
		flagExcludedFiles:   flagExcludedFiles,
		flagIncludedFiles:   flagIncludedFiles,
		pattern:             pattern,
	}
	watcher, err := NewWatcher(cfg)

	if err != nil {
		log.Fatal(err)
	}

	defer watcher.Close()

	err = watcher.AddFiles()
	if err != nil {
		log.Fatal("watcher.Addfiles():", err)
	}

	jobs := make(chan string)
	buildSuccess := make(chan bool)
	buildStarted := make(chan string)

	go builder(jobs, buildStarted, buildSuccess)

	if *flagCommand != "" {
		go runner(*flagCommand, buildStarted, buildSuccess)
	} else {
		go flusher(buildStarted, buildSuccess)
	}

	watcher.Watch(jobs) // start watching files
}
<|endoftext|>