<commit_before>package format

import (
	"bytes"
	"fmt"
	"margo.sh/mg"
	"os/exec"
)

// FmtFunc is a reducer for generic fmt functions
//
// it takes care of reading the view src and properly reporting any errors to the editor
type FmtFunc struct {
	// Fmt receives a copy of the view src and returns the fmt'ed src.
	//
	// Fmt should ideally fail in the face of any uncertainty
	// e.g. if running a command to do the formatting and it prints anything to stderr;
	// it should return an error because commands do not reliably return an error status.
	Fmt func(mx *mg.Ctx, src []byte) ([]byte, error)

	// Langs is the list of languages in which the reducer should run
	Langs []mg.Lang

	// Actions is a list of actions on which the reducer is allowed to run.
	// The reducer always runs on the ViewFmt action, even if this list is empty.
	Actions []mg.Action
}

// ReducerCond returns true if Langs and Actions matches the Ctx
func (ff FmtFunc) ReducerCond(mx *mg.Ctx) bool {
	return mx.ActionIs(ff.Actions...) && mx.LangIs(ff.Langs...)
}

// Reduce implements the FmtFunc reducer.
func (ff FmtFunc) Reduce(mx *mg.Ctx) *mg.State {
	fn := mx.View.Filename()
	src, err := mx.View.ReadAll()
	if err != nil {
		return mx.AddErrorf("failed to read %s: %s\n", fn, err)
	}
	if len(src) == 0 {
		return mx.State
	}

	src, err = ff.Fmt(mx, src)
	if err != nil {
		return mx.AddErrorf("failed to fmt %s: %s\n", fn, err)
	}
	return mx.SetViewSrc(src)
}

// FmtCmd is wrapper around FmtFunc for generic fmt commands.
//
// The view src is passed to the command's stdin.
// It takes care of handling command failure e.g. output on stderr or no output on stdout.
type FmtCmd struct {
	// Name is the command name or path
	Name string

	// Args is a list of args to pass to the command.
	Args []string

	// Env is a map of additional env vars to pass to the command.
	Env mg.EnvMap

	// Langs is the list of languages in which the reducer should run
	Langs []mg.Lang

	// Actions is a list of actions on which the reducer is allowed to run.
	// The reducer always runs on the ViewFmt action, even if this list is empty.
	Actions []mg.Action
}

// Reduce implements the FmtCmd reducer.
func (fc FmtCmd) Reduce(mx *mg.Ctx) *mg.State {
	return FmtFunc{Fmt: fc.fmt, Langs: fc.Langs, Actions: fc.Actions}.Reduce(mx)
}

func (fc FmtCmd) fmt(mx *mg.Ctx, src []byte) ([]byte, error) {
	stdin := bytes.NewReader(src)
	stdout := bytes.NewBuffer(nil)
	stderr := bytes.NewBuffer(nil)
	cmd := exec.Command(fc.Name, fc.Args...)
	cmd.Env = mx.Env.Merge(fc.Env).Environ()
	cmd.Stdin = stdin
	cmd.Stdout = stdout
	cmd.Stderr = stderr
	if err := cmd.Run(); err != nil {
		return nil, err
	}
	if stderr.Len() != 0 {
		return nil, fmt.Errorf("fmt completed successfully, but has output on stderr: %s", stderr.Bytes())
	}
	if stdout.Len() == 0 {
		return nil, fmt.Errorf("fmt completed successfully, but has no output on stdout")
	}
	return stdout.Bytes(), nil
}
<commit_msg>sync margo<commit_after>package format

import (
	"bytes"
	"fmt"
	"margo.sh/mg"
	"os/exec"
)

// FmtFunc is a reducer for generic fmt functions
//
// it takes care of reading the view src and properly reporting any errors to the editor
type FmtFunc struct {
	// Fmt receives a copy of the view src and returns the fmt'ed src.
	//
	// Fmt should ideally fail in the face of any uncertainty
	// e.g. if running a command to do the formatting and it prints anything to stderr;
	// it should return an error because commands do not reliably return an error status.
	Fmt func(mx *mg.Ctx, src []byte) ([]byte, error)

	// Langs is the list of languages in which the reducer should run
	Langs []mg.Lang

	// Actions is a list of actions on which the reducer is allowed to run.
	// The reducer always runs on the ViewFmt action, even if this list is empty.
	Actions []mg.Action
}

// Reduce implements the FmtFunc reducer.
func (ff FmtFunc) Reduce(mx *mg.Ctx) *mg.State {
	if !mx.ActionIs(ff.Actions...) || !mx.LangIs(ff.Langs...) {
		return mx.State
	}

	fn := mx.View.Filename()
	src, err := mx.View.ReadAll()
	if err != nil {
		return mx.AddErrorf("failed to read %s: %s\n", fn, err)
	}
	if len(src) == 0 {
		return mx.State
	}

	src, err = ff.Fmt(mx, src)
	if err != nil {
		return mx.AddErrorf("failed to fmt %s: %s\n", fn, err)
	}
	return mx.SetViewSrc(src)
}

// FmtCmd is wrapper around FmtFunc for generic fmt commands.
//
// The view src is passed to the command's stdin.
// It takes care of handling command failure e.g. output on stderr or no output on stdout.
type FmtCmd struct {
	// Name is the command name or path
	Name string

	// Args is a list of args to pass to the command.
	Args []string

	// Env is a map of additional env vars to pass to the command.
	Env mg.EnvMap

	// Langs is the list of languages in which the reducer should run
	Langs []mg.Lang

	// Actions is a list of actions on which the reducer is allowed to run.
	// The reducer always runs on the ViewFmt action, even if this list is empty.
	Actions []mg.Action
}

// Reduce implements the FmtCmd reducer.
func (fc FmtCmd) Reduce(mx *mg.Ctx) *mg.State {
	return FmtFunc{Fmt: fc.fmt, Langs: fc.Langs, Actions: fc.Actions}.Reduce(mx)
}

func (fc FmtCmd) fmt(mx *mg.Ctx, src []byte) ([]byte, error) {
	stdin := bytes.NewReader(src)
	stdout := bytes.NewBuffer(nil)
	stderr := bytes.NewBuffer(nil)
	cmd := exec.Command(fc.Name, fc.Args...)
	cmd.Env = mx.Env.Merge(fc.Env).Environ()
	cmd.Stdin = stdin
	cmd.Stdout = stdout
	cmd.Stderr = stderr
	if err := cmd.Run(); err != nil {
		return nil, err
	}
	if stderr.Len() != 0 {
		return nil, fmt.Errorf("fmt completed successfully, but has output on stderr: %s", stderr.Bytes())
	}
	if stdout.Len() == 0 {
		return nil, fmt.Errorf("fmt completed successfully, but has no output on stdout")
	}
	return stdout.Bytes(), nil
}
<|endoftext|>


<commit_before>package integration_test

import (
	"path/filepath"

	"github.com/cloudfoundry/libbuildpack/cutlass"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("override yml", func() {
	var app *cutlass.App
	var buildpackName string
	AfterEach(func() {
		if buildpackName != "" {
			cutlass.DeleteBuildpack(buildpackName)
		}
		if app != nil {
			app.Destroy()
		}
		app = nil
	})

	BeforeEach(func() {
		if !ApiHasMultiBuildpack() {
			Skip("Multi buildpack support is required")
		}

		buildpackName = "override_yml_" + cutlass.RandStringRunes(5)
		Expect(cutlass.CreateOrUpdateBuildpack(buildpackName, filepath.Join(bpDir, "fixtures", "overrideyml_bp"), "")).To(Succeed())

		app = cutlass.New(filepath.Join(bpDir, "fixtures", "native_vendoring"))
		app.Buildpacks = []string{buildpackName + "_buildpack", "go_buildpack"}
	})

	It("Forces go from override buildpack", func() {
		Expect(app.Push()).ToNot(Succeed())
		Expect(app.Stdout.String()).To(ContainSubstring("-----> OverrideYML Buildpack"))
		Expect(app.ConfirmBuildpack(buildpackVersion)).To(Succeed())

		Eventually(app.Stdout.String).Should(ContainSubstring("-----> Installing go"))
		Eventually(app.Stdout.String).Should(MatchRegexp("Copy .*/go.tgz"))
		Eventually(app.Stdout.String).Should(ContainSubstring("Error installing Go: dependency sha256 mismatch: expected sha256 062d906c87839d03b243e2821e10653c89b4c92878bfe2bf995dec231e117bfc, actual sha256 b56b58ac21f9f42d032e1e4b8bf8b8823e69af5411caa15aee2b140bc756962f"))
	})
})
<commit_msg>Use V3Push in override test<commit_after>package integration_test

import (
	"path/filepath"

	"github.com/cloudfoundry/libbuildpack/cutlass"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("override yml", func() {
	var app *cutlass.App
	var buildpackName string
	AfterEach(func() {
		if buildpackName != "" {
			cutlass.DeleteBuildpack(buildpackName)
		}
		if app != nil {
			app.Destroy()
		}
		app = nil
	})

	BeforeEach(func() {
		if !ApiHasMultiBuildpack() {
			Skip("Multi buildpack support is required")
		}

		buildpackName = "override_yml_" + cutlass.RandStringRunes(5)
		Expect(cutlass.CreateOrUpdateBuildpack(buildpackName, filepath.Join(bpDir, "fixtures", "overrideyml_bp"), "")).To(Succeed())

		app = cutlass.New(filepath.Join(bpDir, "fixtures", "native_vendoring"))
		app.Buildpacks = []string{buildpackName + "_buildpack", "go_buildpack"}
	})

	It("Forces go from override buildpack", func() {
		Expect(app.V3Push()).ToNot(Succeed())
		Expect(app.Stdout.String()).To(ContainSubstring("-----> OverrideYML Buildpack"))
		Expect(app.ConfirmBuildpack(buildpackVersion)).To(Succeed())

		Eventually(app.Stdout.String).Should(ContainSubstring("-----> Installing go"))
		Eventually(app.Stdout.String).Should(MatchRegexp("Copy .*/go.tgz"))
		Eventually(app.Stdout.String).Should(ContainSubstring("Error installing Go: dependency sha256 mismatch: expected sha256 062d906c87839d03b243e2821e10653c89b4c92878bfe2bf995dec231e117bfc, actual sha256 b56b58ac21f9f42d032e1e4b8bf8b8823e69af5411caa15aee2b140bc756962f"))
	})
})
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"
)

func getRemoteJSON(source string) (body []byte, err error) {
	response, err := http.Get(source)
	//Getting our nice http response.

	//This error check may be given it's own function, later. Not sure of best way to do it.
	if err != nil {
		return nil, err

	}

	defer func() {
		err = response.Body.Close() //and not forgetting to close it when it's done. And before we panic and die horribly.
		if err != nil {
			lFatal("Could  not close server response")
		}
	}()

	if !okHTTPStatus(response) { //Checking that we weren't given crap instead of candy
		return nil, fmt.Errorf("Incorrect server response")
	}

	body, err = ioutil.ReadAll(response.Body) //stolen from official documentation
	if err != nil {
		return nil, err
	}

	return body, nil
}

func okHTTPStatus(chk *http.Response) bool {
	switch chk.StatusCode {
	case http.StatusOK, http.StatusNotModified:
		return true
	case http.StatusGatewayTimeout,
		http.StatusInternalServerError,
		http.StatusNotImplemented,
		http.StatusBadGateway,
		http.StatusServiceUnavailable,
		http.StatusHTTPVersionNotSupported:
		lErr("Server error: ", chk.Status)
		return false
	case http.StatusBadRequest,
		http.StatusTeapot,
		http.StatusUnauthorized,
		http.StatusForbidden,
		http.StatusNotFound,
		http.StatusRequestURITooLong,
		http.StatusExpectationFailed:
		lErr("Incorrect request to server, error ", chk.Status)
		lErr("Possible API changes")
		return false
	default:
		lWarn("Got something weird from server: ", chk.Status)
		lWarn("Continuing anyway")
		return true
	}
}

func makeHTTPSUnsafe() {
	lWarn("Disabling HTTPS trust check is unsafe and may lead to your data being monitored, stolen or falsified by third party")
	lWarn("Are you really want to continue? [Yes/No]")
	if !checkUserResponse() {
		lInfo("Continuing in safe mode")
		return
	}
	lWarn("Continuing in unsafe mode")
	http.DefaultClient.Transport = &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		Dial: (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
		}).Dial,
		TLSHandshakeTimeout: 10 * time.Second,
		TLSClientConfig:     &tls.Config{InsecureSkipVerify: true},
	}
}

func checkUserResponse() bool {
	var response string
	for i := 0; i < 10; i++ {
		time.Sleep(3 * time.Second)
		n, _ := fmt.Scanln(&response)
		r := strings.ToLower(response)
		if n != 1 {
			continue
		}
		if r == "y" || r == "yes" {
			return true
		}
		if r == "n" || r == "no" {
			return false
		}

	}
	return false

}
func (imgdata Image) saveImage(opts *Config) (size int64) { // To not hold all the files open when there is no need. All pointers to files are in the scope of this function.

	filepath := opts.ImageDir + string(os.PathSeparator) + imgdata.Filename

	fsize := getFileSize(filepath)

	start := time.Now() //timing download time. We can't begin it sooner, not sure if we can begin it later

	response, err := http.Get(imgdata.URL)

	if err != nil {
		lErr("Error when getting image: ", strconv.Itoa(imgdata.Imgid))
		lErr(err)
		return
	}

	defer func() {
		err = response.Body.Close() //Same, we shall not listen to the void when we finished getting image
		if err != nil {
			lFatal("Could  not close server response")
		}
	}()

	if !okHTTPStatus(response) {
		return
	}

	sizestring, prs := response.Header["Content-Length"]
	if !prs {
		lErr("Filesize not provided")
	}
	expsize, err := strconv.ParseInt(sizestring[0], 10, 64)
	if err != nil {
		lErr("Unable to get expected filesize")
	}

	if expsize == fsize {
		lInfo("Skipping: no-clobber")
		return 0
	}

	if err != nil {
		lErr("Error when getting image: ", strconv.Itoa(imgdata.Imgid))
		lErr(err)
		return
	}

	output, err := os.Create(filepath) //And now, THE FILE! New, truncated, ready to write
	if err != nil {
		lErr("Error when creating file for image: ", strconv.Itoa(imgdata.Imgid))
		lErr(err) //Either we got no permisson or no space, end of line
		return
	}
	defer func() {
		err = output.Close() //Not forgetting to deal with it after completing download
		if err != nil {
			lFatal("Could  not close downloaded file")
		}
	}()

	size, err = io.Copy(output, response.Body) //Preventing creation of temporary buffer in memory
	if err != nil {
		lErr("Unable to write image on disk, id: ", strconv.Itoa(imgdata.Imgid))
		lErr(err)
		return
	}
	timed := time.Since(start).Seconds()

	lInfof("Downloaded %d bytes in %.2fs, speed %s/s\n", size, timed, fmtbytes(float64(size)/timed))

	if expsize != size {
		lErr("Unable to download full image")
		return 0
	}
	return
}

func getFileSize(path string) int64 {
	fstat, err := os.Stat(path)
	if err != nil {
		return 0
	}
	return fstat.Size()

}
<commit_msg>Extracting more extractable parts<commit_after>package main

import (
	"crypto/tls"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"
)

func getRemoteJSON(source string) (body []byte, err error) {
	response, err := http.Get(source)
	//Getting our nice http response.

	//This error check may be given it's own function, later. Not sure of best way to do it.
	if err != nil {
		return nil, err

	}

	defer func() {
		err = response.Body.Close() //and not forgetting to close it when it's done. And before we panic and die horribly.
		if err != nil {
			lFatal("Could  not close server response")
		}
	}()

	if !okHTTPStatus(response) { //Checking that we weren't given crap instead of candy
		return nil, fmt.Errorf("Incorrect server response")
	}

	body, err = ioutil.ReadAll(response.Body) //stolen from official documentation
	if err != nil {
		return nil, err
	}

	return body, nil
}

func okHTTPStatus(chk *http.Response) bool {
	switch chk.StatusCode {
	case http.StatusOK, http.StatusNotModified:
		return true
	case http.StatusGatewayTimeout,
		http.StatusInternalServerError,
		http.StatusNotImplemented,
		http.StatusBadGateway,
		http.StatusServiceUnavailable,
		http.StatusHTTPVersionNotSupported:
		lErr("Server error: ", chk.Status)
		return false
	case http.StatusBadRequest,
		http.StatusTeapot,
		http.StatusUnauthorized,
		http.StatusForbidden,
		http.StatusNotFound,
		http.StatusRequestURITooLong,
		http.StatusExpectationFailed:
		lErr("Incorrect request to server, error ", chk.Status)
		lErr("Possible API changes")
		return false
	default:
		lWarn("Got something weird from server: ", chk.Status)
		lWarn("Continuing anyway")
		return true
	}
}

func makeHTTPSUnsafe() {
	lWarn("Disabling HTTPS trust check is unsafe and may lead to your data being monitored, stolen or falsified by third party")
	lWarn("Are you really want to continue? [Yes/No]")
	if !checkUserResponse() {
		lInfo("Continuing in safe mode")
		return
	}
	lWarn("Continuing in unsafe mode")
	http.DefaultClient.Transport = &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		Dial: (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
		}).Dial,
		TLSHandshakeTimeout: 10 * time.Second,
		TLSClientConfig:     &tls.Config{InsecureSkipVerify: true},
	}
}

func checkUserResponse() bool {
	var response string
	for i := 0; i < 10; i++ {
		time.Sleep(3 * time.Second)
		n, _ := fmt.Scanln(&response)
		r := strings.ToLower(response)
		if n != 1 {
			continue
		}
		if r == "y" || r == "yes" {
			return true
		}
		if r == "n" || r == "no" {
			return false
		}

	}
	return false

}
func (imgdata Image) saveImage(opts *Config) (size int64) { // To not hold all the files open when there is no need. All pointers to files are in the scope of this function.

	filepath := opts.ImageDir + string(os.PathSeparator) + imgdata.Filename

	fsize := getFileSize(filepath)

	start := time.Now() //timing download time. We can't begin it sooner, not sure if we can begin it later

	response, err := http.Get(imgdata.URL)

	if err != nil {
		lErr("Error when getting image: ", strconv.Itoa(imgdata.Imgid))
		lErr(err)
		return
	}

	defer func() {
		err = response.Body.Close() //Same, we shall not listen to the void when we finished getting image
		if err != nil {
			lFatal("Could  not close server response")
		}
	}()

	if !okHTTPStatus(response) {
		return
	}

	expsize := getRemoteSize(response.Header)

	if expsize == fsize {
		lInfo("Skipping: no-clobber")
		return 0
	}

	if err != nil {
		lErr("Error when getting image: ", strconv.Itoa(imgdata.Imgid))
		lErr(err)
		return
	}

	output, err := os.Create(filepath) //And now, THE FILE! New, truncated, ready to write
	if err != nil {
		lErr("Error when creating file for image: ", strconv.Itoa(imgdata.Imgid))
		lErr(err) //Either we got no permisson or no space, end of line
		return
	}
	defer func() {
		err = output.Close() //Not forgetting to deal with it after completing download
		if err != nil {
			lFatal("Could  not close downloaded file")
		}
	}()

	size, err = io.Copy(output, response.Body) //Preventing creation of temporary buffer in memory
	if err != nil {
		lErr("Unable to write image on disk, id: ", strconv.Itoa(imgdata.Imgid))
		lErr(err)
		return
	}
	timed := time.Since(start).Seconds()

	lInfof("Downloaded %d bytes in %.2fs, speed %s/s\n", size, timed, fmtbytes(float64(size)/timed))

	if expsize != size {
		lErr("Unable to download full image")
		return 0
	}
	return
}

func getFileSize(path string) int64 {
	fstat, err := os.Stat(path)
	if err != nil {
		return 0
	}
	return fstat.Size()

}

func getRemoteSize(head http.Header) (expsize int64) {

	sizestring, ok := head["Content-Length"]
	if !ok {
		lErr("Filesize not provided")
	}
	expsize, err := strconv.ParseInt(sizestring[0], 10, 64)
	if err != nil {
		lErr("Unable to get expected filesize")
	}
	return
}
<|endoftext|>


<commit_before>package scenario

import (
	"time"
	"ihmc.us/anglova/pubsub"
	log "github.com/Sirupsen/logrus"
	"ihmc.us/anglova/config"
	"ihmc.us/anglova/msg"
	"ihmc.us/anglova/protocol"
)

type Anglova struct {
	Cfg  config.Manager
	Quit chan bool
}

// sizes in bytes, intervals in seconds
const (
	BlueForceTopic         = "blueforce"
	BlueForceTrackSize     = 512
	BlueForceTrackInterval = 5
	SensorDataTopic        = "sensordata"
	SensorDataSize         = 128000
	SensorDataInterval     = 300 // 5 minutes
	HQReportTopic          = "hqreport"
	HQReportSize           = 1000000
	HQReportInterval       = 600 // 10 minutes
)

func NewAnglova(cfg config.Manager) (*Anglova) {

	return &Anglova{Cfg: cfg,
		Quit:        make(chan bool)}
}

func (a *Anglova) Run() {

	var quitChannels []chan bool

	//for statistics, use nats
	pubStats, err := pubsub.NewPub(protocol.NATS, a.Cfg.StatsAddress, a.Cfg.StatsPort, "", a.Cfg.StatsAddress, a.Cfg.StatsPort)
	if err != nil {
		log.Error(err)
	}
	//TODO verify if using one shared publisher for stats is goroutine safe

	if a.Cfg.EnableBlueForce {
		quitBF := make(chan bool)
		quitChannels = append(quitChannels, quitBF)
		pubBF, err := pubsub.NewPub(a.Cfg.Protocol, a.Cfg.BrokerAddress, a.Cfg.BrokerPort, BlueForceTopic, a.Cfg.StatsAddress, a.Cfg.StatsPort)
		if err != nil {
			log.Error(err)
		}
		blueForcePublishRoutine(pubBF, pubStats, quitBF)
	}

	if a.Cfg.EnableSensorData {
		quitSD := make(chan bool)
		quitChannels = append(quitChannels, quitSD)
		pubSD,err  := pubsub.NewPub(a.Cfg.Protocol, a.Cfg.BrokerAddress, a.Cfg.BrokerPort, SensorDataTopic, a.Cfg.StatsAddress, a.Cfg.StatsPort)
		if err != nil {
			log.Error(err)
		}
		sensorDataPublishRoutine(pubSD, pubStats, quitSD)
	}

	if a.Cfg.EnableHQReport {
		quitHQ := make(chan bool)
		quitChannels = append(quitChannels, quitHQ)
		pubHQ, err := pubsub.NewPub(a.Cfg.Protocol, a.Cfg.BrokerAddress, a.Cfg.BrokerPort, HQReportTopic, a.Cfg.StatsAddress, a.Cfg.StatsPort)
		if err != nil {
			log.Error(err)
		}
		hqReportPublishRoutine(pubHQ, pubStats, quitHQ)

	}

	// wait for the scenario to quit, quit go routines
	select {
	case <-a.Quit:
		{
			for _, quit := range quitChannels {
				quit <- true
			}
		}
	}

}

func blueForcePublishRoutine(pub *pubsub.Pub, pubStats *pubsub.Pub, quit chan bool) {
	ticker := time.NewTicker(time.Second * BlueForceTrackInterval)
	var msgCount int32
	go func() {
		for {
			select {
			case <-ticker.C:
				message, err := msg.New(pub.ID, msgCount, time.Now().UnixNano(), BlueForceTrackSize)
				if err != nil {
					log.Error(err)
				}
				pub.Publish(BlueForceTopic, message.Bytes())
				pub.PublishStats(msgCount)
				msgCount++
			case <-quit:
				ticker.Stop()
				return
			}
		}
	}()
}

func sensorDataPublishRoutine(pub *pubsub.Pub, pubStats *pubsub.Pub, quit chan bool) {
	ticker := time.NewTicker(time.Second * SensorDataInterval)
	var msgCount int32
	go func() {
		for {
			select {
			case <-ticker.C:
				message, err := msg.New(pub.ID, msgCount, time.Now().UnixNano(), SensorDataSize)
				if err != nil {
					log.Error(err)
				}
				pub.Publish(SensorDataTopic, message.Bytes())
				pub.PublishStats(msgCount)
				msgCount++
			case <-quit:
				ticker.Stop()
				return
			}
		}
	}()
}

func hqReportPublishRoutine(pub *pubsub.Pub, pubStats *pubsub.Pub, quit chan bool) {

	ticker := time.NewTicker(time.Second * HQReportInterval)
	var msgCount int32
	go func() {
		for {
			select {
			case <-ticker.C:
				message, err := msg.New(pub.ID, msgCount, time.Now().UnixNano(), HQReportSize)
				if err != nil {
					log.Error(err)
				}
				pub.Publish(HQReportTopic, message.Bytes())
				pub.PublishStats(msgCount)
				msgCount++
			case <-quit:
				ticker.Stop()
				return
			}
		}
	}()
}
<commit_msg>changes in the scenario configurations<commit_after>package scenario

import (
	"time"
	"ihmc.us/anglova/pubsub"
	log "github.com/Sirupsen/logrus"
	"ihmc.us/anglova/config"
	"ihmc.us/anglova/msg"
	"ihmc.us/anglova/protocol"
)

type Anglova struct {
	Cfg  config.Manager
	Quit chan bool
}

// sizes in bytes, intervals in seconds
const (
	BlueForceTopic         = "blueforce"
	BlueForceTrackSize     = 512
	BlueForceTrackInterval = 5
	SensorDataTopic        = "sensordata"
	SensorDataSize         = 128000
	SensorDataInterval     = 300 // 5 minutes
	HQReportTopic          = "hqreport"
	HQReportSize           = 1000000
	HQReportInterval       = 600 // 10 minutes
)

func NewAnglova(cfg config.Manager) (*Anglova) {

	return &Anglova{Cfg: cfg,
		Quit:        make(chan bool)}
}

func (a *Anglova) Run() {

	var quitChannels []chan bool

	//TODO verify if using one shared publisher for stats is goroutine safe

	if a.Cfg.EnableBlueForce {
		quitBF := make(chan bool)
		quitChannels = append(quitChannels, quitBF)
		pubBF, err := pubsub.NewPub(a.Cfg.Protocol, a.Cfg.BrokerAddress, a.Cfg.BrokerPort, BlueForceTopic, a.Cfg.StatsAddress, a.Cfg.StatsPort)
		if err != nil {
			log.Error(err)
		}
		blueForcePublishRoutine(pubBF, quitBF)
	}

	if a.Cfg.EnableSensorData {
		quitSD := make(chan bool)
		quitChannels = append(quitChannels, quitSD)
		pubSD,err  := pubsub.NewPub(a.Cfg.Protocol, a.Cfg.BrokerAddress, a.Cfg.BrokerPort, SensorDataTopic, a.Cfg.StatsAddress, a.Cfg.StatsPort)
		if err != nil {
			log.Error(err)
		}
		sensorDataPublishRoutine(pubSD, quitSD)
	}

	if a.Cfg.EnableHQReport {
		quitHQ := make(chan bool)
		quitChannels = append(quitChannels, quitHQ)
		pubHQ, err := pubsub.NewPub(a.Cfg.Protocol, a.Cfg.BrokerAddress, a.Cfg.BrokerPort, HQReportTopic, a.Cfg.StatsAddress, a.Cfg.StatsPort)
		if err != nil {
			log.Error(err)
		}
		hqReportPublishRoutine(pubHQ, quitHQ)

	}

	// wait for the scenario to quit, quit go routines
	select {
	case <-a.Quit:
		{
			for _, quit := range quitChannels {
				quit <- true
			}
		}
	}

}

func blueForcePublishRoutine(pub *pubsub.Pub, quit chan bool) {
	ticker := time.NewTicker(time.Second * BlueForceTrackInterval)
	var msgCount int32
	go func() {
		for {
			select {
			case <-ticker.C:
				message, err := msg.New(pub.ID, msgCount, time.Now().UnixNano(), BlueForceTrackSize)
				if err != nil {
					log.Error(err)
				}
				pub.Publish(BlueForceTopic, message.Bytes())
				pub.PublishStats(msgCount)
				msgCount++
			case <-quit:
				ticker.Stop()
				return
			}
		}
	}()
}

func sensorDataPublishRoutine(pub *pubsub.Pub, quit chan bool) {
	ticker := time.NewTicker(time.Second * SensorDataInterval)
	var msgCount int32
	go func() {
		for {
			select {
			case <-ticker.C:
				message, err := msg.New(pub.ID, msgCount, time.Now().UnixNano(), SensorDataSize)
				if err != nil {
					log.Error(err)
				}
				pub.Publish(SensorDataTopic, message.Bytes())
				pub.PublishStats(msgCount)
				msgCount++
			case <-quit:
				ticker.Stop()
				return
			}
		}
	}()
}

func hqReportPublishRoutine(pub *pubsub.Pub, quit chan bool) {

	ticker := time.NewTicker(time.Second * HQReportInterval)
	var msgCount int32
	go func() {
		for {
			select {
			case <-ticker.C:
				message, err := msg.New(pub.ID, msgCount, time.Now().UnixNano(), HQReportSize)
				if err != nil {
					log.Error(err)
				}
				pub.Publish(HQReportTopic, message.Bytes())
				pub.PublishStats(msgCount)
				msgCount++
			case <-quit:
				ticker.Stop()
				return
			}
		}
	}()
}
<|endoftext|>


<commit_before>// Copyright 2015 Apcera Inc. All rights reserved.

package test

import (
	"fmt"
	"strings"
	"testing"
	"time"

	"github.com/nats-io/nats"
)

func TestMaxPayload(t *testing.T) {
	srv, opts := RunServerWithConfig("./configs/override.conf")
	defer srv.Shutdown()

	nc, err := nats.Connect(fmt.Sprintf("nats://%s:%d/", opts.Host, opts.Port))
	if err != nil {
		t.Fatalf("Could not connect to server: %v", err)
	}
	defer nc.Close()

	big := sizedBytes(4 * 1024 * 1024)
	nc.Publish("foo", big)
	err = nc.FlushTimeout(1 * time.Second)
	if err == nil {
		t.Fatalf("Expected an error from flush")
	}
	if strings.Contains(err.Error(), "Maximum Payload Violation") != true {
		t.Fatalf("Received wrong error message (%v)\n", err)
	}
	if !nc.IsClosed() {
		t.Fatalf("Expected connection to be closed")
	}
}
<commit_msg>Update max payload test to reflect behavior change in go client.<commit_after>// Copyright 2015 Apcera Inc. All rights reserved.

package test

import (
	"fmt"
	"net"
	"strings"
	"testing"

	"github.com/nats-io/nats"
)

func TestMaxPayload(t *testing.T) {
	srv, opts := RunServerWithConfig("./configs/override.conf")
	defer srv.Shutdown()

	endpoint := fmt.Sprintf("%s:%d", opts.Host, opts.Port)
	nc, err := nats.Connect(fmt.Sprintf("nats://%s/", endpoint))
	if err != nil {
		t.Fatalf("Could not connect to server: %v", err)
	}
	defer nc.Close()

	size := 4 * 1024 * 1024
	big := sizedBytes(size)
	err = nc.Publish("foo", big)

	if err != nats.ErrMaxPayload {
		t.Fatalf("Expected a Max Payload error")
	}

	conn, err := net.DialTimeout("tcp", endpoint, nc.Opts.Timeout)
	if err != nil {
		t.Fatalf("Could not make a raw connection to the server: %v", err)
	}
	info := make([]byte, 200)
	_, err = conn.Read(info)
	if err != nil {
		t.Fatalf("Expected an info message to be sent by the server: %s", err)
	}

	pub := fmt.Sprintf("PUB bar %d\r\n", size)
	conn.Write([]byte(pub))
	if err != nil {
		t.Fatalf("Could not publish event to the server: %s", err)
	}

	errMsg := make([]byte, 35)
	_, err = conn.Read(errMsg)
	if err != nil {
		t.Fatalf("Expected an error message to be sent by the server: %s", err)
	}

	if strings.Contains(string(errMsg), "Maximum Payload Violation") != true {
		t.Errorf("Received wrong error message (%v)\n", string(errMsg))
	}

	// Client proactively omits sending the message so server
	// does not close the connection.
	if nc.IsClosed() {
		t.Errorf("Expected connection to not be closed.")
	}

	// On the other hand client which did not proactively omitted
	// publishing the bytes following what is suggested by server
	// in the info message has its connection closed.
	_, err = conn.Write(big)
	if err == nil {
		t.Errorf("Expected error due to maximum payload transgression.")
	}
}
<|endoftext|>


<commit_before>package main
import (
    // "os_"
 	"log"
 	// "go/build"
    // "github.com/MSOpenTech/azure-sdk-for-go/management"
    // server "github.com/MSOpenTech/azure-sdk-for-go/management/core/http"
	"github.com/MSOpenTech/azure-sdk-for-go/storage"
	// "github.com/MSOpenTech/azure-sdk-for-go/core/http"
	// ""
	// "./azure-sdk-for-go/storage"
	"fmt"
	"strings"
	// "bytes"
	"crypto/rand"
	// "io/ioutil"
	// "net/http"
	// "time"
)

func randString(n int) string {
	if n <= 0 {
		panic("negative number")
	}
	const alphanum = "0123456789abcdefghijklmnopqrstuvwxyz"
	var bytes = make([]byte, n)
	rand.Read(bytes)
	for i, b := range bytes {
		bytes[i] = alphanum[b%byte(len(alphanum))]
	}
	return string(bytes)
}

func getStorageServer() (*storage.BlobStorageClient) {
	client, err := storage.NewClient("spluto", "b7yy+C33a//uLE62Og9CkKDHRLNErMrbX40nKUxiTimgOvkP3MhEbjObmRxumda9grCwY8zqL6nLNcKCAS40Iw==", "core.chinacloudapi.cn", "2014-02-14", true)
	if err != nil {
		log.Fatal(err)
	}
	cli := client.GetBlobService()
	return cli
}

// func (c storage.StorageClient) GetBlobService() *storage.BlobStorageClient {
// 	return &storage.BlobStorageClient{c}
// }

func connectStorageSever() {
	// client, err := storage.NewBasicClient("spluto", "b7yy+C33a//uLE62Og9CkKDHRLNErMrbX40nKUxiTimgOvkP3MhEbjObmRxumda9grCwY8zqL6nLNcKCAS40Iw==")

	// fmt.Println(type(client))
	// fmt.Println(client.accountName)
	cli := getStorageServer()
	cnt := randString(32)
	blob := randString(20)
	// body := []byte(randString(100))
	// expiry := time.Now().UTC().Add(time.Hour)
	// permissions := "r"

	// err = cli.CreateContainer(cnt, storage.ContainerAccessTypePrivate)
	// if err != nil {
	// 	log.Fatal(err)
	// }
	err := cli.CreateContainer(cnt, storage.ContainerAccessTypeBlob)
	if err != nil {
		// fmt.Println(err)
		log.Fatal(err)
	}
	fmt.Println("!!!!df")
	// defer cli.DeleteContainer(cnt)
	err = cli.PutBlockBlob(cnt, blob, strings.NewReader("Hello!"))
	if err != nil {
		log.Fatal(err)
	}
	defer cli.DeleteBlob(cnt, blob)

	ok, err := cli.BlobExists(cnt, blob+".foo")
	if err != nil {
		log.Fatal(err)
	}
	if ok {
		log.Fatal("Non-existing blob returned as existing: %s/%s", cnt, blob)
	}

	ok, err = cli.BlobExists(cnt, blob)
	if err != nil {
		log.Fatal(err)
	}
	if !ok {
		log.Fatal("Existing blob returned as non-existing: %s/%s", cnt, blob)
	}

	// defer cli.DeleteContainer(cnt)
	// log.Printf("3333")
	// err = cli.PutBlockBlob(cnt, blob, bytes.NewReader(body))
	// if err != nil {
	// 	log.Fatal(err)
	// }
	// sasUri, err := cli.GetBlobSASURI(cnt, blob, expiry, permissions)
	// if err != nil {
	// 	log.Fatal(err)
	// }

	// resp, err := http.Get(sasUri)
	// if err != nil {
	// 	// t.Logf("SAS URI: %s", sasUri)
	// 	log.Fatal(err)
	// }

	// blobResp, err := ioutil.ReadAll(resp.Body)
	// defer resp.Body.Close()
	// if err != nil {
	// 	log.Fatal(err)
	// }
	// if resp.StatusCode != http.StatusOK {
	// 	log.Fatalf("Non-ok status code: %s", resp.Status)
	// }

	// if len(blobResp) != len(body) {
	// 	log.Fatalf("Wrong blob size on SAS URI. Expected: %d, Got: %d", len(body), len(blobResp))
	// }
}

// cli, err := getBlobClient()
	// if err != nil {
	// 	log.Fatal(err)
	// }

// func getBlobClient() (*BlobStorageClient, error) {
// 	name := os.Getenv("ACCOUNT_NAME")
// 	if name == "" {
// 		return nil, errors.New("ACCOUNT_NAME not set, need an empty storage account to test")
// 	}
// 	key := os.Getenv("ACCOUNT_KEY")
// 	if key == "" {
// 		return nil, errors.New("ACCOUNT_KEY not set")
// 	}
// 	cli, err := NewBasicClient(name, key)
// 	if err != nil {
// 		return nil, err
// 	}
// 	return cli.GetBlobService(), nil
// }

func main() {

	connectStorageSever()

}<commit_msg>Azure Exist Function Finished<commit_after>package main
import (
    // "os_"
 	// "log"
 	// "strings"
 	// "go/build"
    // "github.com/MSOpenTech/azure-sdk-for-go/management"
    // server "github.com/MSOpenTech/azure-sdk-for-go/management/core/http"
	"github.com/MSOpenTech/azure-sdk-for-go/storage"
	// "github.com/MSOpenTech/azure-sdk-for-go/core/http"
	// ""
	// "./azure-sdk-for-go/storage"
	"fmt"
	"strings"
	// "bytes"
	"crypto/rand"
	// "io/ioutil"
	// "net/http"
	// "time"
)

// type azureConfig struct {
// 	accountName    string
// 	accountKey        string
// 	baseUrl string
// }

type AzureClient struct {
	client *storage.StorageClient
	blobClient *storage.BlobStorageClient
}

func convertToAzurePath(name string) (string, string, error) {
	afterSplit := strings.Split(name, "/")
	if len(afterSplit) != 2 || len(afterSplit[0]) != 32 {
		return "","",fmt.Errorf("Azure : Need Correct Path Name")
	}
	return afterSplit[0], afterSplit[1], nil
}

// func (c *Glob) OpenReadCloser(name string) (io.ReadCloser, error) {
// 	return c.client.Open(name)
// }

// func (c *storage.BlobStorageClient) OpenWriteCloser(name string) (io.WriteCloser, error) {
// 	exist, err := c.Exists(name)
// 	if err != nil {
// 		return nil, err
// 	}
// 	if !exist {
// 		err := c.client.CreateEmptyFile(name)
// 		if err != nil {
// 			return nil, err
// 		}
// 	}
// 	return &HdfsFile{
// 		path:       name,
// 		logger:     log.New(os.Stdout, "", log.Lshortfile|log.LstdFlags),
// 		hdfsConfig: c.hdfsConfig,
// 	}, nil
// }

func (c *AzureClient) Exists(name string) (bool, error) {
	containerName, blobName, err := convertToAzurePath(name)
	if err != nil {
		return false, err
	}
	return c.blobClient.BlobExists(containerName, blobName)
	// _, err := c.client.Stat(name)
	// return existCommon(err)
}

// func (b BlobStorageClient) BlobExists(container, name string) (bool, error)

// func (c *storage.BlobStorageClient) Rename(oldpath, newpath string) error {
// 	c.Copy()
// }

// only supports '*', '?'
// Syntax:
//    /user/hdfs/etl*/part.*
// func (c *storage.BlobStorageClient) Glob(pattern string) (matches []string, err error) {
// 	if pattern == "" {
// 		return nil, fmt.Errorf("Glob pattern shouldn't be empty")
// 	}
// 	if pattern[len(pattern)-1] == '/' {
// 		return nil, fmt.Errorf("Glob pattern shouldn't be a directory")
// 	}
// 	// names will have all the pathnames of the pattern.
// 	// e.g. "/a/b/c" => [a, b, c]
// 	var names []string
// 	for path.Dir(pattern) != "/" {
// 		names = append(names, path.Base(pattern))
// 		pattern = path.Dir(pattern)
// 	}
// 	names = append(names, pattern[1:len(pattern)])
// 	for i, j := 0, len(names)-1; i < j; i, j = i+1, j-1 {
// 		names[i], names[j] = names[j], names[i]
// 	}
// 	return c.glob("/", names)
// }

// func (c *storage.BlobStorageClient) glob(dir string, names []string) (m []string, err error) {
// 	name := names[0]
// 	var dirs []string
// 	if hasMeta(name) {
// 		fileInfos, err := c.client.ReadDir(dir)
// 		if err != nil {
// 			return nil, err
// 		}
// 		for _, fi := range fileInfos {
// 			matched, err := path.Match(name, fi.Name())
// 			if err != nil {
// 				return nil, err
// 			}
// 			if matched {
// 				dirs = append(dirs, path.Join(dir, fi.Name()))
// 			}
// 		}
// 	} else {
// 		dirs = append(dirs, path.Join(dir, name))
// 	}
// 	for _, pathname := range dirs {
// 		if len(names) == 1 {
// 			exist, err := c.Exists(pathname)
// 			if err != nil {
// 				return nil, err
// 			}
// 			if exist {
// 				m = append(m, pathname)
// 			}
// 		} else {
// 			return c.glob(pathname, names[1:len(names)])

// 		}
// 	}
// 	return
// }

func randString(n int) string {
	if n <= 0 {
		panic("negative number")
	}
	const alphanum = "0123456789abcdefghijklmnopqrstuvwxyz"
	var bytes = make([]byte, n)
	rand.Read(bytes)
	for i, b := range bytes {
		bytes[i] = alphanum[b%byte(len(alphanum))]
	}
	return string(bytes)
}

func NewAzureClient(accountName, accountKey, blobServiceBaseUrl, apiVersion string, useHttps bool) (*AzureClient, error) {
	cli, err := storage.NewClient(accountName, accountKey, blobServiceBaseUrl, apiVersion, useHttps)
	if err != nil {
		return nil, err
	}
	return &AzureClient{
		client : &cli,
		blobClient : cli.GetBlobService(),
	}, nil
}

// func NewHdfsClient(namenodeAddr, webHdfsAddr, user string) (Client, error) {
// 	client, err := hdfs.NewForUser(namenodeAddr, user)
// 	if err != nil {
// 		return nil, err
// 	}
// 	return &HdfsClient{
// 		client: client,
// 		hdfsConfig: hdfsConfig{
// 			namenodeAddr: namenodeAddr,
// 			webHdfsAddr:  webHdfsAddr,
// 			user:         user,
// 		},
// 	}, nil
// }

// func (c storage.StorageClient) GetBlobService() *storage.BlobStorageClient {
// 	return &storage.BlobStorageClient{c}
// }

// func connectStorageSever() {
// 	// client, err := storage.NewBasicClient("spluto", "b7yy+C33a//uLE62Og9CkKDHRLNErMrbX40nKUxiTimgOvkP3MhEbjObmRxumda9grCwY8zqL6nLNcKCAS40Iw==")

// 	// fmt.Println(type(client))
// 	// fmt.Println(client.accountName)
// 	cli, err := newStorageServer()
// 	if err != nil {
// 		log.Fatal(err)
// 	}
// 	cnt := randString(32)
// 	blob := randString(20)
// 	// body := []byte(randString(100))
// 	// expiry := time.Now().UTC().Add(time.Hour)
// 	// permissions := "r"

// 	// err = cli.CreateContainer(cnt, storage.ContainerAccessTypePrivate)
// 	// if err != nil {
// 	// 	log.Fatal(err)
// 	// }
// 	err = cli.CreateContainer(cnt, storage.ContainerAccessTypeBlob)
// 	if err != nil {
// 		// fmt.Println(err)
// 		log.Fatal(err)
// 	}
// 	fmt.Println("!!!!df")
// 	// defer cli.DeleteContainer(cnt)
// 	err = cli.PutBlockBlob(cnt, blob, strings.NewReader("Hello!"))
// 	if err != nil {
// 		log.Fatal(err)
// 	}
// 	defer cli.DeleteBlob(cnt, blob)

// 	ok, err := cli.BlobExists(cnt, blob+".foo")
// 	if err != nil {
// 		log.Fatal(err)
// 	}
// 	if ok {
// 		log.Fatal("Non-existing blob returned as existing: %s/%s", cnt, blob)
// 	}

// 	ok, err = cli.BlobExists(cnt, blob)
// 	if err != nil {
// 		log.Fatal(err)
// 	}
// 	if !ok {
// 		log.Fatal("Existing blob returned as non-existing: %s/%s", cnt, blob)
// 	}

// 	// defer cli.DeleteContainer(cnt)
// 	// log.Printf("3333")
// 	// err = cli.PutBlockBlob(cnt, blob, bytes.NewReader(body))
// 	// if err != nil {
// 	// 	log.Fatal(err)
// 	// }
// 	// sasUri, err := cli.GetBlobSASURI(cnt, blob, expiry, permissions)
// 	// if err != nil {
// 	// 	log.Fatal(err)
// 	// }

// 	// resp, err := http.Get(sasUri)
// 	// if err != nil {
// 	// 	// t.Logf("SAS URI: %s", sasUri)
// 	// 	log.Fatal(err)
// 	// }

// 	// blobResp, err := ioutil.ReadAll(resp.Body)
// 	// defer resp.Body.Close()
// 	// if err != nil {
// 	// 	log.Fatal(err)
// 	// }
// 	// if resp.StatusCode != http.StatusOK {
// 	// 	log.Fatalf("Non-ok status code: %s", resp.Status)
// 	// }

// 	// if len(blobResp) != len(body) {
// 	// 	log.Fatalf("Wrong blob size on SAS URI. Expected: %d, Got: %d", len(body), len(blobResp))
// 	// }
// }

// cli, err := getBlobClient()
	// if err != nil {
	// 	log.Fatal(err)
	// }

// func getBlobClient() (*BlobStorageClient, error) {
// 	name := os.Getenv("ACCOUNT_NAME")
// 	if name == "" {
// 		return nil, errors.New("ACCOUNT_NAME not set, need an empty storage account to test")
// 	}
// 	key := os.Getenv("ACCOUNT_KEY")
// 	if key == "" {
// 		return nil, errors.New("ACCOUNT_KEY not set")
// 	}
// 	cli, err := NewBasicClient(name, key)
// 	if err != nil {
// 		return nil, err
// 	}
// 	return cli.GetBlobService(), nil
// }

func main() {
	// AzureClient

	// client, err := storage.NewClient("spluto", "b7yy+C33a//uLE62Og9CkKDHRLNErMrbX40nKUxiTimgOvkP3MhEbjObmRxumda9grCwY8zqL6nLNcKCAS40Iw==", "core.chinacloudapi.cn", "2014-02-14", true)
	// cli, err := newStorageServer()
	cli, err := NewAzureClient("spluto", "b7yy+C33a//uLE62Og9CkKDHRLNErMrbX40nKUxiTimgOvkP3MhEbjObmRxumda9grCwY8zqL6nLNcKCAS40Iw==", "core.chinacloudapi.cn", "2014-02-14", true)
	exist, err := cli.Exists("11111111111111111111111111111111/sa")
	if (err != nil) {
		fmt.Println(err)
	}
	fmt.Println(exist)

	// connectStorageSever()

}<|endoftext|>


<commit_before>package gomol

import (
	"runtime"
	"sync"
)

type queue struct {
	running     bool
	queueCtl    chan int
	senderCtl   chan int
	workersDone sync.WaitGroup

	queueChan chan *message

	queue    []*message
	queueMut sync.RWMutex
}

var curQueue *queue

func init() {
	startQueueWorkers()
}

func newQueue() *queue {
	return &queue{
		running:   false,
		queueChan: make(chan *message, 1000),
		queueCtl:  make(chan int),
		senderCtl: make(chan int),
		queue:     make([]*message, 0),
	}
}

func startQueueWorkers() {
	if curQueue != nil && curQueue.running {
		return
	}
	curQueue = newQueue()
	curQueue.running = true
	go curQueue.queueWorker()
	go curQueue.senderWorker()
}

func (queue *queue) queueWorker() {
	queue.workersDone.Add(1)
	exiting := false
	for {
		select {
		case msg := <-queue.queueChan:
			queue.queueMut.Lock()
			queue.queue = append(queue.queue, msg)
			queue.queueMut.Unlock()
		case <-queue.queueCtl:
			exiting = true
		}

		if exiting && len(queue.queueChan) == 0 {
			break
		}
	}
	queue.workersDone.Done()
}

func (queue *queue) senderWorker() {
	queue.workersDone.Add(1)
	exiting := false
	for {
		select {
		case <-queue.senderCtl:
			exiting = true
		default:
		}

		if exiting && len(queue.queue) == 0 {
			break
		} else if len(queue.queue) == 0 {
			runtime.Gosched()
			continue
		}

		msg := queue.NextMessage()

		if msg == nil {
			continue
		}

		for _, l := range msg.Base.loggers {
			switch msg.Level {
			case levelDbg:
				l.Dbgm(msg.Attrs, msg.Msg)
			case levelInfo:
				l.Infom(msg.Attrs, msg.Msg)
			case levelWarn:
				l.Warnm(msg.Attrs, msg.Msg)
			case levelError:
				l.Errm(msg.Attrs, msg.Msg)
			case levelFatal:
				l.Fatalm(msg.Attrs, msg.Msg)
			}
		}
	}
	queue.workersDone.Done()
}

func (queue *queue) QueueMessage(msg *message) {
	queue.queueChan <- msg
}

func (queue *queue) NextMessage() *message {
	queue.queueMut.Lock()
	q := queue.queue
	msg, q := q[0], q[1:]
	queue.queue = q
	queue.queueMut.Unlock()

	return msg
}

func (queue *queue) Length() int {
	queue.queueMut.RLock()
	defer queue.queueMut.RUnlock()
	return len(queue.queue) + len(queue.queueChan)
}

func queueLen() int {
	return curQueue.Length()
}

func queueMessage(msg *message) {
	curQueue.QueueMessage(msg)
}

/*
Blocks until all messages in the queue have been processed, then returns.
*/
func FlushMessages() {
	if curQueue.running {
		curQueue.queueCtl <- 1
		curQueue.senderCtl <- 1

		curQueue.workersDone.Wait()
		curQueue.running = false
	}
}
<commit_msg>Improve how the queue workers handle waiting.<commit_after>package gomol

import (
	"sync"
)

type queue struct {
	running     bool
	queueCtl    chan int
	senderCtl   chan int
	workersDone sync.WaitGroup

	queueChan chan *message

	queue        []*message
	msgAddedChan chan int
	queueMut     sync.RWMutex
}

var curQueue *queue

func init() {
	startQueueWorkers()
}

func newQueue() *queue {
	return &queue{
		running:      false,
		queueChan:    make(chan *message, 1000),
		queueCtl:     make(chan int),
		senderCtl:    make(chan int),
		queue:        make([]*message, 0),
		msgAddedChan: make(chan int, 1),
	}
}

func startQueueWorkers() {
	if curQueue != nil && curQueue.running {
		return
	}
	curQueue = newQueue()
	curQueue.running = true
	go curQueue.queueWorker()
	go curQueue.senderWorker()
}

func (queue *queue) queueWorker() {
	queue.workersDone.Add(1)
	exiting := false
	for {
		select {
		case msg := <-queue.queueChan:
			queue.queueMut.Lock()
			queue.queue = append(queue.queue, msg)
			queue.queueMut.Unlock()
			queue.msgAddedChan <- 1
		case <-queue.queueCtl:
			exiting = true
		}

		if exiting && len(queue.queueChan) == 0 {
			break
		}
	}
	queue.workersDone.Done()
}

func (queue *queue) senderWorker() {
	queue.workersDone.Add(1)
	exiting := false
	for {
		if exiting && len(queue.queue) == 0 {
			break
		}

		select {
		case <-queue.senderCtl:
			exiting = true
		case <-queue.msgAddedChan:
		}

		for {
			msg := queue.NextMessage()

			if msg == nil {
				break
			}

			for _, l := range msg.Base.loggers {
				switch msg.Level {
				case levelDbg:
					l.Dbgm(msg.Attrs, msg.Msg)
				case levelInfo:
					l.Infom(msg.Attrs, msg.Msg)
				case levelWarn:
					l.Warnm(msg.Attrs, msg.Msg)
				case levelError:
					l.Errm(msg.Attrs, msg.Msg)
				case levelFatal:
					l.Fatalm(msg.Attrs, msg.Msg)
				}
			}
		}
	}
	queue.workersDone.Done()
}

func (queue *queue) QueueMessage(msg *message) {
	queue.queueChan <- msg
}

func (queue *queue) NextMessage() *message {
	var msg *message
	queue.queueMut.Lock()
	if len(queue.queue) > 0 {
		q := queue.queue
		msg, q = q[0], q[1:]
		queue.queue = q
	} else {
		msg = nil
	}
	queue.queueMut.Unlock()

	return msg
}

func (queue *queue) Length() int {
	queue.queueMut.RLock()
	defer queue.queueMut.RUnlock()
	return len(queue.queue) + len(queue.queueChan)
}

func queueLen() int {
	return curQueue.Length()
}

func queueMessage(msg *message) {
	curQueue.QueueMessage(msg)
}

/*
Blocks until all messages in the queue have been processed, then returns.
*/
func FlushMessages() {
	if curQueue.running {
		curQueue.queueCtl <- 1
		curQueue.senderCtl <- 1

		curQueue.workersDone.Wait()
		curQueue.running = false
	}
}
<|endoftext|>


<commit_before>package helpers

import (
	"errors"
	"fmt"
	"os/exec"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"
)

func FlyLogin(atcURL, concourseAlias, flyBinaryPath string, loginInfo LoginInformation) error {
	switch {
	case loginInfo.NoAuth:
		return flyLogin(flyBinaryPath, []string{
			"-c", atcURL,
			"-t", concourseAlias,
			"-n", "main",
		})
	case loginInfo.BasicAuthCreds.Username != "":
		return flyLogin(flyBinaryPath, []string{
			"-c", atcURL,
			"-t", concourseAlias,
			"-u", loginInfo.BasicAuthCreds.Username,
			"-p", loginInfo.BasicAuthCreds.Password,
			"-n", "main",
		})
	case loginInfo.OauthToken != "":
		return flyLoginOauth(flyBinaryPath, atcURL, loginInfo.OauthToken, []string{
			"-c", atcURL,
			"-t", concourseAlias,
			"-n", "main",
		})
	}

	return errors.New("Unable to determine authentication")
}

func flyLogin(flyBinaryPath string, loginArgs []string) error {
	args := append([]string{"login"}, loginArgs...)
	loginCmd := exec.Command(flyBinaryPath, args...)

	loginProcess, err := gexec.Start(loginCmd, GinkgoWriter, GinkgoWriter)
	if err != nil {
		return err
	}

	Eventually(loginProcess, time.Minute).Should(gexec.Exit(0))

	return nil
}

func flyLoginOauth(flyBinaryPath string, atcURL string, oauthToken string, loginArgs []string) error {
	args := append([]string{"login"}, loginArgs...)
	loginCmd := exec.Command(flyBinaryPath, args...)

	stdin, err := loginCmd.StdinPipe()
	if err != nil {
		return err
	}
	defer stdin.Close()

	loginProcess, err := gexec.Start(loginCmd, GinkgoWriter, GinkgoWriter)
	if err != nil {
		return err
	}

	Eventually(loginProcess.Out).Should(gbytes.Say("enter token"))
	fmt.Fprint(stdin, fmt.Sprintf("%s\n", oauthToken))

	Eventually(loginProcess, time.Minute).Should(gexec.Exit(0))

	return nil
}
<commit_msg>--team-name still defaults to main<commit_after>package helpers

import (
	"errors"
	"fmt"
	"os/exec"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"
)

func FlyLogin(atcURL, concourseAlias, flyBinaryPath string, loginInfo LoginInformation) error {
	switch {
	case loginInfo.NoAuth:
		return flyLogin(flyBinaryPath, []string{
			"-c", atcURL,
			"-t", concourseAlias,
		})
	case loginInfo.BasicAuthCreds.Username != "":
		return flyLogin(flyBinaryPath, []string{
			"-c", atcURL,
			"-t", concourseAlias,
			"-u", loginInfo.BasicAuthCreds.Username,
			"-p", loginInfo.BasicAuthCreds.Password,
		})
	case loginInfo.OauthToken != "":
		return flyLoginOauth(flyBinaryPath, atcURL, loginInfo.OauthToken, []string{
			"-c", atcURL,
			"-t", concourseAlias,
		})
	}

	return errors.New("Unable to determine authentication")
}

func flyLogin(flyBinaryPath string, loginArgs []string) error {
	args := append([]string{"login"}, loginArgs...)
	loginCmd := exec.Command(flyBinaryPath, args...)

	loginProcess, err := gexec.Start(loginCmd, GinkgoWriter, GinkgoWriter)
	if err != nil {
		return err
	}

	Eventually(loginProcess, time.Minute).Should(gexec.Exit(0))

	return nil
}

func flyLoginOauth(flyBinaryPath string, atcURL string, oauthToken string, loginArgs []string) error {
	args := append([]string{"login"}, loginArgs...)
	loginCmd := exec.Command(flyBinaryPath, args...)

	stdin, err := loginCmd.StdinPipe()
	if err != nil {
		return err
	}
	defer stdin.Close()

	loginProcess, err := gexec.Start(loginCmd, GinkgoWriter, GinkgoWriter)
	if err != nil {
		return err
	}

	Eventually(loginProcess.Out).Should(gbytes.Say("enter token"))
	fmt.Fprint(stdin, fmt.Sprintf("%s\n", oauthToken))

	Eventually(loginProcess, time.Minute).Should(gexec.Exit(0))

	return nil
}
<|endoftext|>


<commit_before>package radix

import (
	"sort"
	"strings"
)

// WalkFn is used when walking the tree. Takes a
// key and value, returning if iteration should
// be terminated.
type WalkFn func(s string, v interface{}) bool

// leafNode is used to represent a value
type leafNode struct {
	key string
	val interface{}
}

// edge is used to represent an edge node
type edge struct {
	label byte
	node  *node
}

type node struct {
	// leaf is used to store possible leaf
	leaf *leafNode

	// prefix is the common prefix we ignore
	prefix string

	// Edges should be stored in-order for iteration.
	// We avoid a fully materialized slice to save memory,
	// since in most cases we expect to be sparse
	edges edges
}

func (n *node) isLeaf() bool {
	return n.leaf != nil
}

func (n *node) addEdge(e edge) {
	n.edges = append(n.edges, e)
	n.edges.Sort()
}

func (n *node) replaceEdge(e edge) {
	num := len(n.edges)
	idx := sort.Search(num, func(i int) bool {
		return n.edges[i].label >= e.label
	})
	if idx < num && n.edges[idx].label == e.label {
		n.edges[idx].node = e.node
		return
	}
	panic("replacing missing edge")
}

func (n *node) getEdge(label byte) *node {
	num := len(n.edges)
	idx := sort.Search(num, func(i int) bool {
		return n.edges[i].label >= label
	})
	if idx < num && n.edges[idx].label == label {
		return n.edges[idx].node
	}
	return nil
}

func (n *node) delEdge(label byte) {
	num := len(n.edges)
	idx := sort.Search(num, func(i int) bool {
		return n.edges[i].label >= label
	})
	if idx < num && n.edges[idx].label == label {
		copy(n.edges[idx:], n.edges[idx+1:])
		n.edges[len(n.edges)-1] = edge{}
		n.edges = n.edges[:len(n.edges)-1]
	}
}

type edges []edge

func (e edges) Len() int {
	return len(e)
}

func (e edges) Less(i, j int) bool {
	return e[i].label < e[j].label
}

func (e edges) Swap(i, j int) {
	e[i], e[j] = e[j], e[i]
}

func (e edges) Sort() {
	sort.Sort(e)
}

// Tree implements a radix tree. This can be treated as a
// Dictionary abstract data type. The main advantage over
// a standard hash map is prefix-based lookups and
// ordered iteration,
type Tree struct {
	root *node
	size int
}

// New returns an empty Tree
func New() *Tree {
	return NewFromMap(nil)
}

// NewFromMap returns a new tree containing the keys
// from an existing map
func NewFromMap(m map[string]interface{}) *Tree {
	t := &Tree{root: &node{}}
	for k, v := range m {
		t.Insert(k, v)
	}
	return t
}

// Len is used to return the number of elements in the tree
func (t *Tree) Len() int {
	return t.size
}

// longestPrefix finds the length of the shared prefix
// of two strings
func longestPrefix(k1, k2 string) int {
	max := len(k1)
	if l := len(k2); l < max {
		max = l
	}
	var i int
	for i = 0; i < max; i++ {
		if k1[i] != k2[i] {
			break
		}
	}
	return i
}

// Insert is used to add a newentry or update
// an existing entry. Returns if updated.
func (t *Tree) Insert(s string, v interface{}) (interface{}, bool) {
	var parent *node
	n := t.root
	search := s
	for {
		// Handle key exhaution
		if len(search) == 0 {
			if n.isLeaf() {
				old := n.leaf.val
				n.leaf.val = v
				return old, true
			}

			n.leaf = &leafNode{
				key: s,
				val: v,
			}
			t.size++
			return nil, false
		}

		// Look for the edge
		parent = n
		n = n.getEdge(search[0])

		// No edge, create one
		if n == nil {
			e := edge{
				label: search[0],
				node: &node{
					leaf: &leafNode{
						key: s,
						val: v,
					},
					prefix: search,
				},
			}
			parent.addEdge(e)
			t.size++
			return nil, false
		}

		// Determine longest prefix of the search key on match
		commonPrefix := longestPrefix(search, n.prefix)
		if commonPrefix == len(n.prefix) {
			search = search[commonPrefix:]
			continue
		}

		// Split the node
		t.size++
		child := &node{
			prefix: search[:commonPrefix],
		}
		parent.replaceEdge(edge{
			label: search[0],
			node:  child,
		})

		// Restore the existing node
		child.addEdge(edge{
			label: n.prefix[commonPrefix],
			node:  n,
		})
		n.prefix = n.prefix[commonPrefix:]

		// Create a new leaf node
		leaf := &leafNode{
			key: s,
			val: v,
		}

		// If the new key is a subset, add to to this node
		search = search[commonPrefix:]
		if len(search) == 0 {
			child.leaf = leaf
			return nil, false
		}

		// Create a new edge for the node
		child.addEdge(edge{
			label: search[0],
			node: &node{
				leaf:   leaf,
				prefix: search,
			},
		})
		return nil, false
	}
}

// Delete is used to delete a key, returning the previous
// value and if it was deleted
func (t *Tree) Delete(s string) (interface{}, bool) {
	var parent *node
	var label byte
	n := t.root
	search := s
	for {
		// Check for key exhaution
		if len(search) == 0 {
			if !n.isLeaf() {
				break
			}
			goto DELETE
		}

		// Look for an edge
		parent = n
		label = search[0]
		n = n.getEdge(label)
		if n == nil {
			break
		}

		// Consume the search prefix
		if strings.HasPrefix(search, n.prefix) {
			search = search[len(n.prefix):]
		} else {
			break
		}
	}
	return nil, false

DELETE:
	// Delete the leaf
	leaf := n.leaf
	n.leaf = nil
	t.size--

	// Check if we should delete this node from the parent
	if parent != nil && len(n.edges) == 0 {
		parent.delEdge(label)
	}

	// Check if we should merge this node
	if n != t.root && len(n.edges) == 1 {
		n.mergeChild()
	}

	// Check if we should merge the parent's other child
	if parent != nil && parent != t.root && len(parent.edges) == 1 && !parent.isLeaf() {
		parent.mergeChild()
	}

	return leaf.val, true
}

// DeletePrefix is used to delete the subtree under a prefix
// Returns how many nodes were deleted
// Use this to delete large subtrees efficiently
func (t *Tree) DeletePrefix(s string) int {
	return t.deletePrefix(nil, t.root, s)
}

// delete does a recursive deletion
func (t *Tree) deletePrefix(parent, n *node, prefix string) int {
	// Check for key exhaustion
	if len(prefix) == 0 {
		// Remove the leaf node
		subTreeSize := 0
		//recursively walk from all edges of the node to be deleted
		recursiveWalk(n, func(s string, v interface{}) bool {
			subTreeSize++
			return false
		})
		if n.isLeaf() {
			n.leaf = nil
		}
		n.edges = nil // deletes the entire subtree

		// Check if we should merge the parent's other child
		if parent != nil && parent != t.root && len(parent.edges) == 1 && !parent.isLeaf() {
			parent.mergeChild()
		}
		t.size -= subTreeSize
		return subTreeSize
	}

	// Look for an edge
	label := prefix[0]
	child := n.getEdge(label)
	if child == nil || (!strings.HasPrefix(child.prefix, prefix) && !strings.HasPrefix(prefix, child.prefix)) {
		return 0
	}

	// Consume the search prefix
	if len(child.prefix) > len(prefix) {
		prefix = prefix[len(prefix):]
	} else {
		prefix = prefix[len(child.prefix):]
	}
	return t.deletePrefix(n, child, prefix)
}

func (n *node) mergeChild() {
	e := n.edges[0]
	child := e.node
	n.prefix = n.prefix + child.prefix
	n.leaf = child.leaf
	n.edges = child.edges
}

// Get is used to lookup a specific key, returning
// the value and if it was found
func (t *Tree) Get(s string) (interface{}, bool) {
	n := t.root
	search := s
	for {
		// Check for key exhaution
		if len(search) == 0 {
			if n.isLeaf() {
				return n.leaf.val, true
			}
			break
		}

		// Look for an edge
		n = n.getEdge(search[0])
		if n == nil {
			break
		}

		// Consume the search prefix
		if strings.HasPrefix(search, n.prefix) {
			search = search[len(n.prefix):]
		} else {
			break
		}
	}
	return nil, false
}

// LongestPrefix is like Get, but instead of an
// exact match, it will return the longest prefix match.
func (t *Tree) LongestPrefix(s string) (string, interface{}, bool) {
	var last *leafNode
	n := t.root
	search := s
	for {
		// Look for a leaf node
		if n.isLeaf() {
			last = n.leaf
		}

		// Check for key exhaution
		if len(search) == 0 {
			break
		}

		// Look for an edge
		n = n.getEdge(search[0])
		if n == nil {
			break
		}

		// Consume the search prefix
		if strings.HasPrefix(search, n.prefix) {
			search = search[len(n.prefix):]
		} else {
			break
		}
	}
	if last != nil {
		return last.key, last.val, true
	}
	return "", nil, false
}

// Minimum is used to return the minimum value in the tree
func (t *Tree) Minimum() (string, interface{}, bool) {
	n := t.root
	for {
		if n.isLeaf() {
			return n.leaf.key, n.leaf.val, true
		}
		if len(n.edges) > 0 {
			n = n.edges[0].node
		} else {
			break
		}
	}
	return "", nil, false
}

// Maximum is used to return the maximum value in the tree
func (t *Tree) Maximum() (string, interface{}, bool) {
	n := t.root
	for {
		if num := len(n.edges); num > 0 {
			n = n.edges[num-1].node
			continue
		}
		if n.isLeaf() {
			return n.leaf.key, n.leaf.val, true
		}
		break
	}
	return "", nil, false
}

// Walk is used to walk the tree
func (t *Tree) Walk(fn WalkFn) {
	recursiveWalk(t.root, fn)
}

// WalkPrefix is used to walk the tree under a prefix
func (t *Tree) WalkPrefix(prefix string, fn WalkFn) {
	n := t.root
	search := prefix
	for {
		// Check for key exhaution
		if len(search) == 0 {
			recursiveWalk(n, fn)
			return
		}

		// Look for an edge
		n = n.getEdge(search[0])
		if n == nil {
			break
		}

		// Consume the search prefix
		if strings.HasPrefix(search, n.prefix) {
			search = search[len(n.prefix):]

		} else if strings.HasPrefix(n.prefix, search) {
			// Child may be under our search prefix
			recursiveWalk(n, fn)
			return
		} else {
			break
		}
	}

}

// WalkPath is used to walk the tree, but only visiting nodes
// from the root down to a given leaf. Where WalkPrefix walks
// all the entries *under* the given prefix, this walks the
// entries *above* the given prefix.
func (t *Tree) WalkPath(path string, fn WalkFn) {
	n := t.root
	search := path
	for {
		// Visit the leaf values if any
		if n.leaf != nil && fn(n.leaf.key, n.leaf.val) {
			return
		}

		// Check for key exhaution
		if len(search) == 0 {
			return
		}

		// Look for an edge
		n = n.getEdge(search[0])
		if n == nil {
			return
		}

		// Consume the search prefix
		if strings.HasPrefix(search, n.prefix) {
			search = search[len(n.prefix):]
		} else {
			break
		}
	}
}

// recursiveWalk is used to do a pre-order walk of a node
// recursively. Returns true if the walk should be aborted
func recursiveWalk(n *node, fn WalkFn) bool {
	// Visit the leaf values if any
	if n.leaf != nil && fn(n.leaf.key, n.leaf.val) {
		return true
	}

	// Recurse on the children
	for _, e := range n.edges {
		if recursiveWalk(e.node, fn) {
			return true
		}
	}
	return false
}

// ToMap is used to walk the tree and convert it into a map
func (t *Tree) ToMap() map[string]interface{} {
	out := make(map[string]interface{}, t.size)
	t.Walk(func(k string, v interface{}) bool {
		out[k] = v
		return false
	})
	return out
}
<commit_msg>refactor: Change replaceEdge func to updateEdge<commit_after>package radix

import (
	"sort"
	"strings"
)

// WalkFn is used when walking the tree. Takes a
// key and value, returning if iteration should
// be terminated.
type WalkFn func(s string, v interface{}) bool

// leafNode is used to represent a value
type leafNode struct {
	key string
	val interface{}
}

// edge is used to represent an edge node
type edge struct {
	label byte
	node  *node
}

type node struct {
	// leaf is used to store possible leaf
	leaf *leafNode

	// prefix is the common prefix we ignore
	prefix string

	// Edges should be stored in-order for iteration.
	// We avoid a fully materialized slice to save memory,
	// since in most cases we expect to be sparse
	edges edges
}

func (n *node) isLeaf() bool {
	return n.leaf != nil
}

func (n *node) addEdge(e edge) {
	n.edges = append(n.edges, e)
	n.edges.Sort()
}

func (n *node) updateEdge(label byte, node *node) {
	num := len(n.edges)
	idx := sort.Search(num, func(i int) bool {
		return n.edges[i].label >= label
	})
	if idx < num && n.edges[idx].label == label {
		n.edges[idx].node = node
		return
	}
	panic("replacing missing edge")
}

func (n *node) getEdge(label byte) *node {
	num := len(n.edges)
	idx := sort.Search(num, func(i int) bool {
		return n.edges[i].label >= label
	})
	if idx < num && n.edges[idx].label == label {
		return n.edges[idx].node
	}
	return nil
}

func (n *node) delEdge(label byte) {
	num := len(n.edges)
	idx := sort.Search(num, func(i int) bool {
		return n.edges[i].label >= label
	})
	if idx < num && n.edges[idx].label == label {
		copy(n.edges[idx:], n.edges[idx+1:])
		n.edges[len(n.edges)-1] = edge{}
		n.edges = n.edges[:len(n.edges)-1]
	}
}

type edges []edge

func (e edges) Len() int {
	return len(e)
}

func (e edges) Less(i, j int) bool {
	return e[i].label < e[j].label
}

func (e edges) Swap(i, j int) {
	e[i], e[j] = e[j], e[i]
}

func (e edges) Sort() {
	sort.Sort(e)
}

// Tree implements a radix tree. This can be treated as a
// Dictionary abstract data type. The main advantage over
// a standard hash map is prefix-based lookups and
// ordered iteration,
type Tree struct {
	root *node
	size int
}

// New returns an empty Tree
func New() *Tree {
	return NewFromMap(nil)
}

// NewFromMap returns a new tree containing the keys
// from an existing map
func NewFromMap(m map[string]interface{}) *Tree {
	t := &Tree{root: &node{}}
	for k, v := range m {
		t.Insert(k, v)
	}
	return t
}

// Len is used to return the number of elements in the tree
func (t *Tree) Len() int {
	return t.size
}

// longestPrefix finds the length of the shared prefix
// of two strings
func longestPrefix(k1, k2 string) int {
	max := len(k1)
	if l := len(k2); l < max {
		max = l
	}
	var i int
	for i = 0; i < max; i++ {
		if k1[i] != k2[i] {
			break
		}
	}
	return i
}

// Insert is used to add a newentry or update
// an existing entry. Returns if updated.
func (t *Tree) Insert(s string, v interface{}) (interface{}, bool) {
	var parent *node
	n := t.root
	search := s
	for {
		// Handle key exhaution
		if len(search) == 0 {
			if n.isLeaf() {
				old := n.leaf.val
				n.leaf.val = v
				return old, true
			}

			n.leaf = &leafNode{
				key: s,
				val: v,
			}
			t.size++
			return nil, false
		}

		// Look for the edge
		parent = n
		n = n.getEdge(search[0])

		// No edge, create one
		if n == nil {
			e := edge{
				label: search[0],
				node: &node{
					leaf: &leafNode{
						key: s,
						val: v,
					},
					prefix: search,
				},
			}
			parent.addEdge(e)
			t.size++
			return nil, false
		}

		// Determine longest prefix of the search key on match
		commonPrefix := longestPrefix(search, n.prefix)
		if commonPrefix == len(n.prefix) {
			search = search[commonPrefix:]
			continue
		}

		// Split the node
		t.size++
		child := &node{
			prefix: search[:commonPrefix],
		}
		parent.updateEdge(search[0], child)

		// Restore the existing node
		child.addEdge(edge{
			label: n.prefix[commonPrefix],
			node:  n,
		})
		n.prefix = n.prefix[commonPrefix:]

		// Create a new leaf node
		leaf := &leafNode{
			key: s,
			val: v,
		}

		// If the new key is a subset, add to to this node
		search = search[commonPrefix:]
		if len(search) == 0 {
			child.leaf = leaf
			return nil, false
		}

		// Create a new edge for the node
		child.addEdge(edge{
			label: search[0],
			node: &node{
				leaf:   leaf,
				prefix: search,
			},
		})
		return nil, false
	}
}

// Delete is used to delete a key, returning the previous
// value and if it was deleted
func (t *Tree) Delete(s string) (interface{}, bool) {
	var parent *node
	var label byte
	n := t.root
	search := s
	for {
		// Check for key exhaution
		if len(search) == 0 {
			if !n.isLeaf() {
				break
			}
			goto DELETE
		}

		// Look for an edge
		parent = n
		label = search[0]
		n = n.getEdge(label)
		if n == nil {
			break
		}

		// Consume the search prefix
		if strings.HasPrefix(search, n.prefix) {
			search = search[len(n.prefix):]
		} else {
			break
		}
	}
	return nil, false

DELETE:
	// Delete the leaf
	leaf := n.leaf
	n.leaf = nil
	t.size--

	// Check if we should delete this node from the parent
	if parent != nil && len(n.edges) == 0 {
		parent.delEdge(label)
	}

	// Check if we should merge this node
	if n != t.root && len(n.edges) == 1 {
		n.mergeChild()
	}

	// Check if we should merge the parent's other child
	if parent != nil && parent != t.root && len(parent.edges) == 1 && !parent.isLeaf() {
		parent.mergeChild()
	}

	return leaf.val, true
}

// DeletePrefix is used to delete the subtree under a prefix
// Returns how many nodes were deleted
// Use this to delete large subtrees efficiently
func (t *Tree) DeletePrefix(s string) int {
	return t.deletePrefix(nil, t.root, s)
}

// delete does a recursive deletion
func (t *Tree) deletePrefix(parent, n *node, prefix string) int {
	// Check for key exhaustion
	if len(prefix) == 0 {
		// Remove the leaf node
		subTreeSize := 0
		//recursively walk from all edges of the node to be deleted
		recursiveWalk(n, func(s string, v interface{}) bool {
			subTreeSize++
			return false
		})
		if n.isLeaf() {
			n.leaf = nil
		}
		n.edges = nil // deletes the entire subtree

		// Check if we should merge the parent's other child
		if parent != nil && parent != t.root && len(parent.edges) == 1 && !parent.isLeaf() {
			parent.mergeChild()
		}
		t.size -= subTreeSize
		return subTreeSize
	}

	// Look for an edge
	label := prefix[0]
	child := n.getEdge(label)
	if child == nil || (!strings.HasPrefix(child.prefix, prefix) && !strings.HasPrefix(prefix, child.prefix)) {
		return 0
	}

	// Consume the search prefix
	if len(child.prefix) > len(prefix) {
		prefix = prefix[len(prefix):]
	} else {
		prefix = prefix[len(child.prefix):]
	}
	return t.deletePrefix(n, child, prefix)
}

func (n *node) mergeChild() {
	e := n.edges[0]
	child := e.node
	n.prefix = n.prefix + child.prefix
	n.leaf = child.leaf
	n.edges = child.edges
}

// Get is used to lookup a specific key, returning
// the value and if it was found
func (t *Tree) Get(s string) (interface{}, bool) {
	n := t.root
	search := s
	for {
		// Check for key exhaution
		if len(search) == 0 {
			if n.isLeaf() {
				return n.leaf.val, true
			}
			break
		}

		// Look for an edge
		n = n.getEdge(search[0])
		if n == nil {
			break
		}

		// Consume the search prefix
		if strings.HasPrefix(search, n.prefix) {
			search = search[len(n.prefix):]
		} else {
			break
		}
	}
	return nil, false
}

// LongestPrefix is like Get, but instead of an
// exact match, it will return the longest prefix match.
func (t *Tree) LongestPrefix(s string) (string, interface{}, bool) {
	var last *leafNode
	n := t.root
	search := s
	for {
		// Look for a leaf node
		if n.isLeaf() {
			last = n.leaf
		}

		// Check for key exhaution
		if len(search) == 0 {
			break
		}

		// Look for an edge
		n = n.getEdge(search[0])
		if n == nil {
			break
		}

		// Consume the search prefix
		if strings.HasPrefix(search, n.prefix) {
			search = search[len(n.prefix):]
		} else {
			break
		}
	}
	if last != nil {
		return last.key, last.val, true
	}
	return "", nil, false
}

// Minimum is used to return the minimum value in the tree
func (t *Tree) Minimum() (string, interface{}, bool) {
	n := t.root
	for {
		if n.isLeaf() {
			return n.leaf.key, n.leaf.val, true
		}
		if len(n.edges) > 0 {
			n = n.edges[0].node
		} else {
			break
		}
	}
	return "", nil, false
}

// Maximum is used to return the maximum value in the tree
func (t *Tree) Maximum() (string, interface{}, bool) {
	n := t.root
	for {
		if num := len(n.edges); num > 0 {
			n = n.edges[num-1].node
			continue
		}
		if n.isLeaf() {
			return n.leaf.key, n.leaf.val, true
		}
		break
	}
	return "", nil, false
}

// Walk is used to walk the tree
func (t *Tree) Walk(fn WalkFn) {
	recursiveWalk(t.root, fn)
}

// WalkPrefix is used to walk the tree under a prefix
func (t *Tree) WalkPrefix(prefix string, fn WalkFn) {
	n := t.root
	search := prefix
	for {
		// Check for key exhaution
		if len(search) == 0 {
			recursiveWalk(n, fn)
			return
		}

		// Look for an edge
		n = n.getEdge(search[0])
		if n == nil {
			break
		}

		// Consume the search prefix
		if strings.HasPrefix(search, n.prefix) {
			search = search[len(n.prefix):]

		} else if strings.HasPrefix(n.prefix, search) {
			// Child may be under our search prefix
			recursiveWalk(n, fn)
			return
		} else {
			break
		}
	}

}

// WalkPath is used to walk the tree, but only visiting nodes
// from the root down to a given leaf. Where WalkPrefix walks
// all the entries *under* the given prefix, this walks the
// entries *above* the given prefix.
func (t *Tree) WalkPath(path string, fn WalkFn) {
	n := t.root
	search := path
	for {
		// Visit the leaf values if any
		if n.leaf != nil && fn(n.leaf.key, n.leaf.val) {
			return
		}

		// Check for key exhaution
		if len(search) == 0 {
			return
		}

		// Look for an edge
		n = n.getEdge(search[0])
		if n == nil {
			return
		}

		// Consume the search prefix
		if strings.HasPrefix(search, n.prefix) {
			search = search[len(n.prefix):]
		} else {
			break
		}
	}
}

// recursiveWalk is used to do a pre-order walk of a node
// recursively. Returns true if the walk should be aborted
func recursiveWalk(n *node, fn WalkFn) bool {
	// Visit the leaf values if any
	if n.leaf != nil && fn(n.leaf.key, n.leaf.val) {
		return true
	}

	// Recurse on the children
	for _, e := range n.edges {
		if recursiveWalk(e.node, fn) {
			return true
		}
	}
	return false
}

// ToMap is used to walk the tree and convert it into a map
func (t *Tree) ToMap() map[string]interface{} {
	out := make(map[string]interface{}, t.size)
	t.Walk(func(k string, v interface{}) bool {
		out[k] = v
		return false
	})
	return out
}
<|endoftext|>


<commit_before>// Rate limited io
package ratio

import (
	"io"
	"time"
)

type values struct {
	n   int
	err error
}

type operation struct {
	p       []byte
	values  chan values
	written int
}

type rateLimiter struct {
	stop      chan bool
	worker    chan *operation
	action    func(p []byte) (int, error)
	limit     int
	remaining int
	written   int
	op        *operation
}

func (rl *rateLimiter) reset() {
	rl.remaining = rl.limit
}

func (rl *rateLimiter) run(duration time.Duration) {
	for {
		select {
		case <-rl.stop:
			rl.close()
			return
		case <-time.Tick(duration):
			rl.reset()
			rl.write()
		case op := <-rl.worker:
			rl.op = op
			rl.write()
		}
	}
}

func (rl *rateLimiter) write() {
	if rl.op == nil {
		return
	}

	p := rl.op.p
	if rl.remaining < len(rl.op.p) {
		p = rl.op.p[:rl.remaining]
	}
	rl.remaining -= len(p)

	n, err := rl.action(p)
	rl.op.p = rl.op.p[n:]
	rl.op.written += n

	if len(rl.op.p) == 0 || err != nil {
		rl.op.values <- values{rl.op.written, err}
		rl.op = nil
	}
}

func (rl *rateLimiter) record(p []byte) (int, error) {
	op := &operation{
		p:      p,
		values: make(chan values, 1),
	}
	select {
	case <-rl.stop:
		return 0, io.EOF
	case rl.worker <- op:
		val := <-op.values
		return val.n, val.err
	}
}

func (rl *rateLimiter) close() {
	if rl.op != nil {
		rl.op.values <- values{rl.op.written, io.EOF}
	}
}

// Returns a rate limited io.Writer, allowing to write up to size bytes per duration.
func RateLimitedWriter(w io.Writer, size int, duration time.Duration) io.WriteCloser {
	rl := &rateLimiter{
		action: func(p []byte) (int, error) {
			return w.Write(p)
		},
		limit:     size,
		remaining: size,
		stop:      make(chan bool),
		worker:    make(chan *operation),
	}
	go rl.run(duration)
	return rl
}

// Returns a rate limited io.Reader, allowing to read up to size bytes per duration.
func RateLimitedReader(r io.Reader, size int, duration time.Duration) io.ReadCloser {
	rl := &rateLimiter{
		action: func(p []byte) (int, error) {
			return r.Read(p)
		},
		limit:     size,
		remaining: size,
		stop:      make(chan bool),
		worker:    make(chan *operation),
	}
	go rl.run(duration)
	return rl
}

func (rl *rateLimiter) Read(p []byte) (int, error) {
	return rl.record(p)
}

func (rl *rateLimiter) Write(p []byte) (int, error) {
	return rl.record(p)
}

func (rl *rateLimiter) Close() error {
	select {
	case <-rl.stop:
		return nil
	default:
	}
	close(rl.stop)
	return nil
}
<commit_msg>add more documentation<commit_after>/*
Rate limited io

To create a rate limited io.Writer:

	// Create a new rate-limited writer.
	rw := ratio.RateLimitedWriter(w, 10e6, time.Second)
	io.Copy(rw, reader)

*/
package ratio

import (
	"io"
	"time"
)

type values struct {
	n   int
	err error
}

type operation struct {
	p       []byte
	values  chan values
	written int
}

type rateLimiter struct {
	stop      chan bool
	worker    chan *operation
	action    func(p []byte) (int, error)
	limit     int
	remaining int
	written   int
	op        *operation
}

func (rl *rateLimiter) reset() {
	rl.remaining = rl.limit
}

func (rl *rateLimiter) run(duration time.Duration) {
	for {
		select {
		case <-rl.stop:
			rl.close()
			return
		case <-time.Tick(duration):
			rl.reset()
			rl.write()
		case op := <-rl.worker:
			rl.op = op
			rl.write()
		}
	}
}

func (rl *rateLimiter) write() {
	if rl.op == nil {
		return
	}

	p := rl.op.p
	if rl.remaining < len(rl.op.p) {
		p = rl.op.p[:rl.remaining]
	}
	rl.remaining -= len(p)

	n, err := rl.action(p)
	rl.op.p = rl.op.p[n:]
	rl.op.written += n

	if len(rl.op.p) == 0 || err != nil {
		rl.op.values <- values{rl.op.written, err}
		rl.op = nil
	}
}

func (rl *rateLimiter) record(p []byte) (int, error) {
	op := &operation{
		p:      p,
		values: make(chan values, 1),
	}
	select {
	case <-rl.stop:
		return 0, io.EOF
	case rl.worker <- op:
		val := <-op.values
		return val.n, val.err
	}
}

func (rl *rateLimiter) close() {
	if rl.op != nil {
		rl.op.values <- values{rl.op.written, io.EOF}
	}
}

// Returns a rate limited io.Writer, allowing to write up to size bytes per duration.
func RateLimitedWriter(w io.Writer, size int, duration time.Duration) io.WriteCloser {
	rl := &rateLimiter{
		action: func(p []byte) (int, error) {
			return w.Write(p)
		},
		limit:     size,
		remaining: size,
		stop:      make(chan bool),
		worker:    make(chan *operation),
	}
	go rl.run(duration)
	return rl
}

// Returns a rate limited io.Reader, allowing to read up to size bytes per duration.
func RateLimitedReader(r io.Reader, size int, duration time.Duration) io.ReadCloser {
	rl := &rateLimiter{
		action: func(p []byte) (int, error) {
			return r.Read(p)
		},
		limit:     size,
		remaining: size,
		stop:      make(chan bool),
		worker:    make(chan *operation),
	}
	go rl.run(duration)
	return rl
}

func (rl *rateLimiter) Read(p []byte) (int, error) {
	return rl.record(p)
}

func (rl *rateLimiter) Write(p []byte) (int, error) {
	return rl.record(p)
}

func (rl *rateLimiter) Close() error {
	select {
	case <-rl.stop:
		return nil
	default:
	}
	close(rl.stop)
	return nil
}
<|endoftext|>


<commit_before>// Copyright (C) 2010, Kyle Lemons <kyle@kylelemons.net>.  All rights reserved.

package log4go

import (
	//"io"
	"fmt"
    "strings"
    "os"
    "github.com/rjohnsondev/raven-go/raven"
)

var (
    // number of simultaneous requests to send through to sentry
    RavenConcurrencyLevel = 32
)

// This is the standard writer that prints to standard output.
type RavenLogWriter chan *LogRecord

func NewRavenLogWriter(dsn string) RavenLogWriter {

    hostname, err := os.Hostname()
    if err != nil {
		fmt.Fprintf(os.Stderr, "Unable to determine hostname, logging to root\n")
    }

	w := RavenLogWriter(make(chan *LogRecord, LogBufferLength))

    for x := 0; x < RavenConcurrencyLevel; x++ {
        go ravenLogWorker(w, dsn, hostname)
    }

	return w
}

func ravenLogWorker(w RavenLogWriter, dsn string, hostname string) {
    client, err := raven.NewClient(dsn, hostname)
    if err != nil {
        fmt.Fprintf(os.Stderr, "NewRavenLogWriter(%q): %s\n", dsn, err)
        return
    }

    var timestr string
    var timestrAt int64

    for rec := range w {
        if at := rec.Created.UnixNano() / 1e9; at != timestrAt {
            timestr, timestrAt = rec.Created.Format("01/02/06 15:04:05"), at
        }

        extra := make(map[string]interface{})
        extra["StackTrace"] = strings.Split(strings.Replace(string(rec.Stack), "\t", "    ",-1), "\n")
        extra["Time"] = timestr
        extra["Level"] = rec.Level
        extra["Source"] = rec.Source

        var err error = nil
        if rec.Level <= DEBUG {
            err = client.Debug(rec.Message, extra)
        } else if rec.Level <= INFO {
            err = client.Info(rec.Message, extra)
        } else if rec.Level <= WARNING {
            err = client.Warning(rec.Message, extra)
        } else if rec.Level <= ERROR {
            err = client.Error(rec.Message, extra)
        } else if rec.Level <= CRITICAL {
            err = client.Fatal(rec.Message, extra)
        }

        if err != nil {
            fmt.Printf("RavenLogWriter(%q): %s \n %s", dsn, err, rec.Message)
            return
        }
    }
}

// This is the ConsoleLogWriter's output method.  This will block if the output
// buffer is full.
func (w RavenLogWriter) LogWrite(rec *LogRecord) {
	w <- rec
}

// Close stops the logger from sending messages to standard output.  Attempts to
// send log messages to this logger after a Close have undefined behavior.
func (w RavenLogWriter) Close() {
	close(w)
}
<commit_msg>Allow raven logs to back up massively, and then stop adding to the buffer if above a certian threshold instead of blocking<commit_after>// Copyright (C) 2010, Kyle Lemons <kyle@kylelemons.net>.  All rights reserved.

package log4go

import (
	//"io"
	"fmt"
    "strings"
    "os"
    "runtime"
    "github.com/rjohnsondev/raven-go/raven"
)

var (
    // number of simultaneous requests to send through to sentry
    RavenConcurrencyLevel = 32
    // Max size of buffer
    RavenLogBufferLength = 10000
    // When the buffer gets to this size, don't attempt to log any more
    // for safety this should be set to RavenLogBufferLength - number of threads
    // logging.  TODO: something more elegant
    RavenLogBufferThreshold = RavenLogBufferLength - (runtime.NumCPU() * 2)
)

// This is the standard writer that prints to standard output.
type RavenLogWriter chan *LogRecord

func NewRavenLogWriter(dsn string) RavenLogWriter {

    hostname, err := os.Hostname()
    if err != nil {
		fmt.Fprintf(os.Stderr, "Unable to determine hostname, logging to root\n")
    }

	w := RavenLogWriter(make(chan *LogRecord, RavenLogBufferLength))

    for x := 0; x < RavenConcurrencyLevel; x++ {
        go ravenLogWorker(w, dsn, hostname)
    }

	return w
}

func ravenLogWorker(w RavenLogWriter, dsn string, hostname string) {
    client, err := raven.NewClient(dsn, hostname)
    if err != nil {
        fmt.Fprintf(os.Stderr, "NewRavenLogWriter(%q): %s\n", dsn, err)
        return
    }

    var timestr string
    var timestrAt int64

    for rec := range w {
        if at := rec.Created.UnixNano() / 1e9; at != timestrAt {
            timestr, timestrAt = rec.Created.Format("01/02/06 15:04:05"), at
        }

        extra := make(map[string]interface{})
        extra["StackTrace"] = strings.Split(strings.Replace(string(rec.Stack), "\t", "    ",-1), "\n")
        extra["Time"] = timestr
        extra["Level"] = rec.Level
        extra["Source"] = rec.Source

        var err error = nil
        if rec.Level <= DEBUG {
            err = client.Debug(rec.Message, extra)
        } else if rec.Level <= INFO {
            err = client.Info(rec.Message, extra)
        } else if rec.Level <= WARNING {
            err = client.Warning(rec.Message, extra)
        } else if rec.Level <= ERROR {
            err = client.Error(rec.Message, extra)
        } else if rec.Level <= CRITICAL {
            err = client.Fatal(rec.Message, extra)
        }

        if err != nil {
            fmt.Printf("RavenLogWriter(%q): %s \n %s", dsn, err, rec.Message)
            return
        }
    }
}

// This is the ConsoleLogWriter's output method.  This will block if the output
// buffer is full.
func (w RavenLogWriter) LogWrite(rec *LogRecord) {
    l := len(w)
    if l > RavenLogBufferThreshold {
        fmt.Fprintf(os.Stderr, "Raven buffer size now %v, refusing to log to sentry: %v\n", l, rec.Message)
        return
    }
	w <- rec
}

// Close stops the logger from sending messages to standard output.  Attempts to
// send log messages to this logger after a Close have undefined behavior.
func (w RavenLogWriter) Close() {
	close(w)
}
<|endoftext|>


<commit_before>package main

import (
	"archive/zip"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
)

var (
	g_train = flag.Bool("train", false, "run the training")

	trainfile = "training.csv"
	testfile  = "test.csv"
)

func printf(format string, args ...interface{}) (int, error) {
	return fmt.Fprintf(os.Stderr, format, args...)
}

func main() {
	printf("::: higgsml-validate...\n")
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "$ %s zipfile [test.csv [training.csv]]\n", os.Args[0])
		fmt.Fprintf(os.Stderr,
			"  - test.csv is a test file (taken from $PWD if not given.)\n"+
				"  - training.csv is a training file (taken from $PWD if not given.)\n"+
				"    training.csv is needed iif -train is enabled",
		)
		flag.PrintDefaults()
	}

	flag.Parse()

	if flag.NArg() < 1 {
		printf("**error** hml-validate needs the path to a zip file\n")
		flag.Usage()
		os.Exit(1)
	}

	if !*g_train {
		trainfile = ""
	}

	if flag.NArg() > 1 {
		testfile = flag.Arg(1)
	}
	if flag.NArg() > 2 {
		testfile = flag.Arg(1)
		trainfile = flag.Arg(2)
	}

	for _, file := range []*string{&testfile, &trainfile} {
		// printf("file[%d]=%q\n", i, *file)
		if *file == "" {
			continue
		}
		if !filepath.IsAbs(*file) {
			name, err := filepath.Abs(*file)
			if err == nil {
				*file = name
			}
		}
		_, err := os.Lstat(*file)
		if err != nil {
			printf("**error** no such file [%s]\n", *file)
			flag.Usage()
			os.Exit(1)
		}
		// printf("file[%d]=%q\n", i, *file)
	}

	rc := run()
	os.Exit(rc)
}

func run() int {
	fname := flag.Arg(0)
	printf("Archive: %s\n", fname)

	r, err := zip.OpenReader(fname)
	if err != nil {
		printf("**error** %v\n", err)
		return 1
	}
	defer r.Close()

	// printf("comment: %q\n", r.Comment)

	tmpdir, err := ioutil.TempDir("", "higgsml-validate-")
	if err != nil {
		printf("**error** creating tmpdir: %v\n", err)
		return 1
	}
	defer os.RemoveAll(tmpdir)

	err = unzip(tmpdir, r)
	if err != nil {
		printf("**error** unzipping [%s] under [%s]: %v\n", fname, tmpdir, err)
		return 1
	}

	v, err := NewValidator(dir, *g_train)
	if err != nil {
		printf("**error** validating: %v\n", err)
		return 1
	}

	// printf("validate: %#v\n", v)
	err = v.Run()
	if err != nil {
		printf("**error** running validation: %v\n", err)
		return 1
	}

	return 0
}

func unzip(tmpdir string, r *zip.ReadCloser) error {

	var err error

	for _, f := range r.File {
		printf("  inflating: %s\n", f.Name)
		rc, err := f.Open()
		if err != nil {
			printf("**error** opening [%s]: %v\n", f.Name, err)
			return err
		}
		defer rc.Close()

		ofname := filepath.Join(tmpdir, f.Name)
		// printf("extracting into: [%s] (%v)\n", ofname, f.Mode())
		odir := filepath.Dir(ofname)
		err = os.MkdirAll(odir, 0755)
		if err != nil {
			printf("**error** creating output dir [%s]: %v\n", odir, err)
			return err
		}
		w, err := os.Create(ofname)
		if err != nil {
			printf("**error** creating output file [%s]: %v\n", ofname, err)
			return err
		}
		defer func(w *os.File) {
			err = w.Close()
			if err != nil {
				printf("**error** closing output file [%s]: %v\n", ofname, err)
				os.Exit(1)
			}
		}(w)
		_, err = io.Copy(w, rc)
		if err != nil {
			printf("**error** copying to [%s]: %v\n", ofname, err)
			return err
		}

		err = w.Chmod(f.Mode())
		if err != nil {
			printf("**error** setting chmod [%v] for file [%s]: %v\n", f.Mode(), ofname, err)
			return err
		}
	}

	return err
}
<commit_msg>hml-validate: better help message<commit_after>package main

import (
	"archive/zip"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
)

var (
	g_train = flag.Bool("train", false, "run the training")

	trainfile = "training.csv"
	testfile  = "test.csv"
)

func printf(format string, args ...interface{}) (int, error) {
	return fmt.Fprintf(os.Stderr, format, args...)
}

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, `Usage: hml-validate zipfile-or-directory [test.csv [training.csv]]

where:
  - test.csv is a test file (taken from $PWD if not given.)
  - training.csv is a training file (taken from $PWD if not given.)
    training.csv is needed iif -train is enabled.

ex:
 $ hml-validate my-team
 $ hml-validate my-team/
 $ hml-validate my-team.zip
 $ hml-validate my-team.tar.gz
 $ hml-validate my-team /path/to/test.csv
 $ hml-validate my-team /path/to/test.csv /path/to/training.csv

options:
`,
		)
		flag.PrintDefaults()
	}

	flag.Parse()

	printf("::: higgsml-validate...\n")
	if flag.NArg() < 1 {
		printf("**error** hml-validate needs the path to a zip file\n")
		flag.Usage()
		os.Exit(1)
	}

	if !*g_train {
		trainfile = ""
	}

	if flag.NArg() > 1 {
		testfile = flag.Arg(1)
	}
	if flag.NArg() > 2 {
		testfile = flag.Arg(1)
		trainfile = flag.Arg(2)
	}

	for _, file := range []*string{&testfile, &trainfile} {
		// printf("file[%d]=%q\n", i, *file)
		if *file == "" {
			continue
		}
		if !filepath.IsAbs(*file) {
			name, err := filepath.Abs(*file)
			if err == nil {
				*file = name
			}
		}
		_, err := os.Lstat(*file)
		if err != nil {
			printf("**error** no such file [%s]\n", *file)
			flag.Usage()
			os.Exit(1)
		}
		// printf("file[%d]=%q\n", i, *file)
	}

	rc := run()
	os.Exit(rc)
}

func run() int {
	fname := flag.Arg(0)
	printf("Archive: %s\n", fname)

	r, err := zip.OpenReader(fname)
	if err != nil {
		printf("**error** %v\n", err)
		return 1
	}
	defer r.Close()

	// printf("comment: %q\n", r.Comment)

	tmpdir, err := ioutil.TempDir("", "higgsml-validate-")
	if err != nil {
		printf("**error** creating tmpdir: %v\n", err)
		return 1
	}
	defer os.RemoveAll(tmpdir)

	err = unzip(tmpdir, r)
	if err != nil {
		printf("**error** unzipping [%s] under [%s]: %v\n", fname, tmpdir, err)
		return 1
	}

	v, err := NewValidator(dir, *g_train)
	if err != nil {
		printf("**error** validating: %v\n", err)
		return 1
	}

	// printf("validate: %#v\n", v)
	err = v.Run()
	if err != nil {
		printf("**error** running validation: %v\n", err)
		return 1
	}

	return 0
}

func unzip(tmpdir string, r *zip.ReadCloser) error {

	var err error

	for _, f := range r.File {
		printf("  inflating: %s\n", f.Name)
		rc, err := f.Open()
		if err != nil {
			printf("**error** opening [%s]: %v\n", f.Name, err)
			return err
		}
		defer rc.Close()

		ofname := filepath.Join(tmpdir, f.Name)
		// printf("extracting into: [%s] (%v)\n", ofname, f.Mode())
		odir := filepath.Dir(ofname)
		err = os.MkdirAll(odir, 0755)
		if err != nil {
			printf("**error** creating output dir [%s]: %v\n", odir, err)
			return err
		}
		w, err := os.Create(ofname)
		if err != nil {
			printf("**error** creating output file [%s]: %v\n", ofname, err)
			return err
		}
		defer func(w *os.File) {
			err = w.Close()
			if err != nil {
				printf("**error** closing output file [%s]: %v\n", ofname, err)
				os.Exit(1)
			}
		}(w)
		_, err = io.Copy(w, rc)
		if err != nil {
			printf("**error** copying to [%s]: %v\n", ofname, err)
			return err
		}

		err = w.Chmod(f.Mode())
		if err != nil {
			printf("**error** setting chmod [%v] for file [%s]: %v\n", f.Mode(), ofname, err)
			return err
		}
	}

	return err
}
<|endoftext|>