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

package main

import (
	"bytes"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"runtime/pprof"
	"strings"
	"text/template"
	"time"

	"github.com/golang/glog"
	"github.com/google/kati"
)

const shellDateTimeformat = time.RFC3339

var (
	makefileFlag string
	jobsFlag     int

	loadJSON string
	saveJSON string
	loadGOB  string
	saveGOB  string
	useCache bool

	m2n  bool
	goma bool

	cpuprofile          string
	heapprofile         string
	memstats            string
	traceEventFile      string
	syntaxCheckOnlyFlag bool
	queryFlag           string
	eagerCmdEvalFlag    bool
	generateNinja       bool
	gomaDir             string
	findCachePrunes     string
	findCacheLeafNames  string
	shellDate           string
)

func init() {
	// TODO: Make this default and replace this by -d flag.
	flag.StringVar(&makefileFlag, "f", "", "Use it as a makefile")
	flag.IntVar(&jobsFlag, "j", 1, "Allow N jobs at once.")

	flag.StringVar(&loadGOB, "load", "", "")
	flag.StringVar(&saveGOB, "save", "", "")
	flag.StringVar(&loadJSON, "load_json", "", "")
	flag.StringVar(&saveJSON, "save_json", "", "")
	flag.BoolVar(&useCache, "use_cache", false, "Use cache.")

	flag.BoolVar(&m2n, "m2n", false, "m2n mode")
	flag.BoolVar(&goma, "goma", false, "ensure goma start")

	flag.StringVar(&cpuprofile, "kati_cpuprofile", "", "write cpu profile to `file`")
	flag.StringVar(&heapprofile, "kati_heapprofile", "", "write heap profile to `file`")
	flag.StringVar(&memstats, "kati_memstats", "", "Show memstats with given templates")
	flag.StringVar(&traceEventFile, "kati_trace_event", "", "write trace event to `file`")
	flag.BoolVar(&syntaxCheckOnlyFlag, "c", false, "Syntax check only.")
	flag.StringVar(&queryFlag, "query", "", "Show the target info")
	flag.BoolVar(&eagerCmdEvalFlag, "eager_cmd_eval", false, "Eval commands first.")
	flag.BoolVar(&generateNinja, "ninja", false, "Generate build.ninja.")
	flag.StringVar(&gomaDir, "goma_dir", "", "If specified, use goma to build C/C++ files.")

	flag.StringVar(&findCachePrunes, "find_cache_prunes", "",
		"space separated prune directories for find cache.")
	flag.StringVar(&findCacheLeafNames, "find_cache_leaf_names", "",
		"space separated leaf names for find cache.")
	flag.StringVar(&shellDate, "shell_date", "", "specify $(shell date) time as "+shellDateTimeformat)

	flag.BoolVar(&kati.StatsFlag, "kati_stats", false, "Show a bunch of statistics")
	flag.BoolVar(&kati.PeriodicStatsFlag, "kati_periodic_stats", false, "Show a bunch of periodic statistics")
	flag.BoolVar(&kati.EvalStatsFlag, "kati_eval_stats", false, "Show eval statistics")

	flag.BoolVar(&kati.DryRunFlag, "n", false, "Only print the commands that would be executed")

	// TODO: Make this default.
	flag.BoolVar(&kati.UseFindCache, "use_find_cache", false, "Use find cache.")
	flag.BoolVar(&kati.UseWildcardCache, "use_wildcard_cache", true, "Use wildcard cache.")
	flag.BoolVar(&kati.UseShellBuiltins, "use_shell_builtins", true, "Use shell builtins")
	flag.StringVar(&kati.IgnoreOptionalInclude, "ignore_optional_include", "", "If specified, skip reading -include directives start with the specified path.")
}

func writeHeapProfile() {
	f, err := os.Create(heapprofile)
	if err != nil {
		panic(err)
	}
	pprof.WriteHeapProfile(f)
	f.Close()
}

type memStatsDumper struct {
	*template.Template
}

func (t memStatsDumper) dump() {
	var ms runtime.MemStats
	runtime.ReadMemStats(&ms)
	var buf bytes.Buffer
	err := t.Template.Execute(&buf, ms)
	fmt.Println(buf.String())
	if err != nil {
		panic(err)
	}
}

func load(req kati.LoadReq) (*kati.DepGraph, error) {
	if loadGOB != "" {
		g, err := kati.GOB.Load(loadGOB)
		return g, err
	}
	if loadJSON != "" {
		g, err := kati.JSON.Load(loadJSON)
		return g, err
	}
	g, err := kati.Load(req)
	return g, err
}

func save(g *kati.DepGraph, targets []string) error {
	var err error
	if saveGOB != "" {
		err = kati.GOB.Save(g, saveGOB, targets)
	}
	if saveJSON != "" {
		serr := kati.JSON.Save(g, saveJSON, targets)
		if err == nil {
			err = serr
		}
	}
	return err
}

func m2nsetup() {
	fmt.Println("kati: m2n mode")
	generateNinja = true
	kati.IgnoreOptionalInclude = "out/%.P"
	kati.UseFindCache = true
	if findCachePrunes == "" {
		findCachePrunes = ".git .repo out"
	}
}

func gomasetup() {
	if gomaDir == "" {
		gomaDir = os.Getenv("GOMA_DIR")
		if gomaDir == "" {
			gomaDir = os.ExpandEnv("${HOME}/goma")
		}
	}
	fmt.Printf("kati: setup goma: %s\n", gomaDir)
	cmd := exec.Command(filepath.Join(gomaDir, "goma_ctl.py"), "ensure_start")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err := cmd.Run()
	if err != nil {
		fmt.Printf("goma failed to start: %v", err)
		os.Exit(1)
	}
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	m2ncmd := false
	if filepath.Base(os.Args[0]) == "m2n" {
		m2nsetup()
		m2ncmd = true
	}
	flag.Parse()
	args := flag.Args()
	if m2n {
		generateNinja = true
		if !m2ncmd {
			m2nsetup()
		}
		if len(args) > 1 {
			fmt.Println("use only first argument as ONE_SHOT_MAKEFILE. ignore rest")
		}
		if len(args) > 0 {
			err := os.Setenv("ONE_SHOT_MAKEFILE", filepath.Join(args[0], "Android.mk"))
			if err != nil {
				fmt.Println(err)
				os.Exit(1)
			}
			fmt.Printf("ONE_SHOT_MAKEFILE=%s\n", os.ExpandEnv("${ONE_SHOT_MAKEFILE}"))
		}
		args = args[:0]
	}
	if goma {
		gomasetup()
	}
	err := katiMain(args)
	if err != nil {
		fmt.Println(err)
		// http://www.gnu.org/software/make/manual/html_node/Running.html
		os.Exit(2)
	}
}

func katiMain(args []string) error {
	defer glog.Flush()
	if cpuprofile != "" {
		f, err := os.Create(cpuprofile)
		if err != nil {
			return err
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}
	if heapprofile != "" {
		defer writeHeapProfile()
	}
	defer kati.DumpStats()
	if memstats != "" {
		ms := memStatsDumper{
			Template: template.Must(template.New("memstats").Parse(memstats)),
		}
		ms.dump()
		defer ms.dump()
	}
	if traceEventFile != "" {
		f, err := os.Create(traceEventFile)
		if err != nil {
			panic(err)
		}
		kati.TraceEventStart(f)
		defer kati.TraceEventStop()
	}

	if shellDate != "" {
		if shellDate == "ref" {
			shellDate = shellDateTimeformat[:20] // until Z, drop 07:00
		}
		t, err := time.Parse(shellDateTimeformat, shellDate)
		if err != nil {
			panic(err)
		}
		kati.ShellDateTimestamp = t
	}

	var leafNames []string
	if findCacheLeafNames != "" {
		leafNames = strings.Fields(findCacheLeafNames)
	}
	if findCachePrunes != "" {
		kati.UseFindCache = true
		kati.AndroidFindCacheInit(strings.Fields(findCachePrunes), leafNames)
	}

	req := kati.FromCommandLine(args)
	if makefileFlag != "" {
		req.Makefile = makefileFlag
	}
	req.EnvironmentVars = os.Environ()
	req.UseCache = useCache
	req.EagerEvalCommand = eagerCmdEvalFlag

	g, err := load(req)
	if err != nil {
		return err
	}

	err = save(g, req.Targets)
	if err != nil {
		return err
	}

	if generateNinja {
		return kati.GenerateNinja(g, gomaDir)
	}

	if syntaxCheckOnlyFlag {
		return nil
	}

	if queryFlag != "" {
		kati.Query(os.Stdout, queryFlag, g)
		return nil
	}

	execOpt := &kati.ExecutorOpt{
		NumJobs: jobsFlag,
	}
	ex, err := kati.NewExecutor(execOpt)
	if err != nil {
		return err
	}
	err = ex.Exec(g)
	if err != nil {
		return err
	}
	return nil
}
<commit_msg>[go] backport [C++] Remove *_WRAPPER automatically<commit_after>// Copyright 2015 Google Inc. All rights reserved
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"bytes"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"runtime/pprof"
	"strings"
	"text/template"
	"time"

	"github.com/golang/glog"
	"github.com/google/kati"
)

const shellDateTimeformat = time.RFC3339

var (
	makefileFlag string
	jobsFlag     int

	loadJSON string
	saveJSON string
	loadGOB  string
	saveGOB  string
	useCache bool

	m2n  bool
	goma bool

	cpuprofile          string
	heapprofile         string
	memstats            string
	traceEventFile      string
	syntaxCheckOnlyFlag bool
	queryFlag           string
	eagerCmdEvalFlag    bool
	generateNinja       bool
	gomaDir             string
	findCachePrunes     string
	findCacheLeafNames  string
	shellDate           string
)

func init() {
	// TODO: Make this default and replace this by -d flag.
	flag.StringVar(&makefileFlag, "f", "", "Use it as a makefile")
	flag.IntVar(&jobsFlag, "j", 1, "Allow N jobs at once.")

	flag.StringVar(&loadGOB, "load", "", "")
	flag.StringVar(&saveGOB, "save", "", "")
	flag.StringVar(&loadJSON, "load_json", "", "")
	flag.StringVar(&saveJSON, "save_json", "", "")
	flag.BoolVar(&useCache, "use_cache", false, "Use cache.")

	flag.BoolVar(&m2n, "m2n", false, "m2n mode")
	flag.BoolVar(&goma, "goma", false, "ensure goma start")

	flag.StringVar(&cpuprofile, "kati_cpuprofile", "", "write cpu profile to `file`")
	flag.StringVar(&heapprofile, "kati_heapprofile", "", "write heap profile to `file`")
	flag.StringVar(&memstats, "kati_memstats", "", "Show memstats with given templates")
	flag.StringVar(&traceEventFile, "kati_trace_event", "", "write trace event to `file`")
	flag.BoolVar(&syntaxCheckOnlyFlag, "c", false, "Syntax check only.")
	flag.StringVar(&queryFlag, "query", "", "Show the target info")
	flag.BoolVar(&eagerCmdEvalFlag, "eager_cmd_eval", false, "Eval commands first.")
	flag.BoolVar(&generateNinja, "ninja", false, "Generate build.ninja.")
	flag.StringVar(&gomaDir, "goma_dir", "", "If specified, use goma to build C/C++ files.")

	flag.StringVar(&findCachePrunes, "find_cache_prunes", "",
		"space separated prune directories for find cache.")
	flag.StringVar(&findCacheLeafNames, "find_cache_leaf_names", "",
		"space separated leaf names for find cache.")
	flag.StringVar(&shellDate, "shell_date", "", "specify $(shell date) time as "+shellDateTimeformat)

	flag.BoolVar(&kati.StatsFlag, "kati_stats", false, "Show a bunch of statistics")
	flag.BoolVar(&kati.PeriodicStatsFlag, "kati_periodic_stats", false, "Show a bunch of periodic statistics")
	flag.BoolVar(&kati.EvalStatsFlag, "kati_eval_stats", false, "Show eval statistics")

	flag.BoolVar(&kati.DryRunFlag, "n", false, "Only print the commands that would be executed")

	// TODO: Make this default.
	flag.BoolVar(&kati.UseFindCache, "use_find_cache", false, "Use find cache.")
	flag.BoolVar(&kati.UseWildcardCache, "use_wildcard_cache", true, "Use wildcard cache.")
	flag.BoolVar(&kati.UseShellBuiltins, "use_shell_builtins", true, "Use shell builtins")
	flag.StringVar(&kati.IgnoreOptionalInclude, "ignore_optional_include", "", "If specified, skip reading -include directives start with the specified path.")
}

func writeHeapProfile() {
	f, err := os.Create(heapprofile)
	if err != nil {
		panic(err)
	}
	pprof.WriteHeapProfile(f)
	f.Close()
}

type memStatsDumper struct {
	*template.Template
}

func (t memStatsDumper) dump() {
	var ms runtime.MemStats
	runtime.ReadMemStats(&ms)
	var buf bytes.Buffer
	err := t.Template.Execute(&buf, ms)
	fmt.Println(buf.String())
	if err != nil {
		panic(err)
	}
}

func load(req kati.LoadReq) (*kati.DepGraph, error) {
	if loadGOB != "" {
		g, err := kati.GOB.Load(loadGOB)
		return g, err
	}
	if loadJSON != "" {
		g, err := kati.JSON.Load(loadJSON)
		return g, err
	}
	g, err := kati.Load(req)
	return g, err
}

func save(g *kati.DepGraph, targets []string) error {
	var err error
	if saveGOB != "" {
		err = kati.GOB.Save(g, saveGOB, targets)
	}
	if saveJSON != "" {
		serr := kati.JSON.Save(g, saveJSON, targets)
		if err == nil {
			err = serr
		}
	}
	return err
}

func m2nsetup() {
	fmt.Println("kati: m2n mode")
	generateNinja = true
	kati.IgnoreOptionalInclude = "out/%.P"
	kati.UseFindCache = true
	if findCachePrunes == "" {
		findCachePrunes = ".git .repo out"
	}
}

func gomasetup() {
	for _, k := range []string{"CC_WRAPPER", "CXX_WRAPPER", "JAVAC_WRAPPER"} {
		v := os.Getenv(k)
		if v != "" {
			fmt.Printf("Note: %s=%s may confuse m2n --goma, unsetting", k, v)
			os.Unsetenv(k)
		}
	}

	if gomaDir == "" {
		gomaDir = os.Getenv("GOMA_DIR")
		if gomaDir == "" {
			gomaDir = os.ExpandEnv("${HOME}/goma")
		}
	}
	fmt.Printf("kati: setup goma: %s\n", gomaDir)
	cmd := exec.Command(filepath.Join(gomaDir, "goma_ctl.py"), "ensure_start")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err := cmd.Run()
	if err != nil {
		fmt.Printf("goma failed to start: %v", err)
		os.Exit(1)
	}
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	m2ncmd := false
	if filepath.Base(os.Args[0]) == "m2n" {
		m2nsetup()
		m2ncmd = true
	}
	flag.Parse()
	args := flag.Args()
	if m2n {
		generateNinja = true
		if !m2ncmd {
			m2nsetup()
		}
		if len(args) > 1 {
			fmt.Println("use only first argument as ONE_SHOT_MAKEFILE. ignore rest")
		}
		if len(args) > 0 {
			err := os.Setenv("ONE_SHOT_MAKEFILE", filepath.Join(args[0], "Android.mk"))
			if err != nil {
				fmt.Println(err)
				os.Exit(1)
			}
			fmt.Printf("ONE_SHOT_MAKEFILE=%s\n", os.ExpandEnv("${ONE_SHOT_MAKEFILE}"))
		}
		args = args[:0]
	}
	if goma {
		gomasetup()
	}
	err := katiMain(args)
	if err != nil {
		fmt.Println(err)
		// http://www.gnu.org/software/make/manual/html_node/Running.html
		os.Exit(2)
	}
}

func katiMain(args []string) error {
	defer glog.Flush()
	if cpuprofile != "" {
		f, err := os.Create(cpuprofile)
		if err != nil {
			return err
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}
	if heapprofile != "" {
		defer writeHeapProfile()
	}
	defer kati.DumpStats()
	if memstats != "" {
		ms := memStatsDumper{
			Template: template.Must(template.New("memstats").Parse(memstats)),
		}
		ms.dump()
		defer ms.dump()
	}
	if traceEventFile != "" {
		f, err := os.Create(traceEventFile)
		if err != nil {
			panic(err)
		}
		kati.TraceEventStart(f)
		defer kati.TraceEventStop()
	}

	if shellDate != "" {
		if shellDate == "ref" {
			shellDate = shellDateTimeformat[:20] // until Z, drop 07:00
		}
		t, err := time.Parse(shellDateTimeformat, shellDate)
		if err != nil {
			panic(err)
		}
		kati.ShellDateTimestamp = t
	}

	var leafNames []string
	if findCacheLeafNames != "" {
		leafNames = strings.Fields(findCacheLeafNames)
	}
	if findCachePrunes != "" {
		kati.UseFindCache = true
		kati.AndroidFindCacheInit(strings.Fields(findCachePrunes), leafNames)
	}

	req := kati.FromCommandLine(args)
	if makefileFlag != "" {
		req.Makefile = makefileFlag
	}
	req.EnvironmentVars = os.Environ()
	req.UseCache = useCache
	req.EagerEvalCommand = eagerCmdEvalFlag

	g, err := load(req)
	if err != nil {
		return err
	}

	err = save(g, req.Targets)
	if err != nil {
		return err
	}

	if generateNinja {
		return kati.GenerateNinja(g, gomaDir)
	}

	if syntaxCheckOnlyFlag {
		return nil
	}

	if queryFlag != "" {
		kati.Query(os.Stdout, queryFlag, g)
		return nil
	}

	execOpt := &kati.ExecutorOpt{
		NumJobs: jobsFlag,
	}
	ex, err := kati.NewExecutor(execOpt)
	if err != nil {
		return err
	}
	err = ex.Exec(g)
	if err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"log"
	"math"
	"math/rand"
	"sort"
	"time"
	"runtime"

	"github.com/phil-mansfield/shellfish/los/geom"
	"github.com/phil-mansfield/shellfish/cmd/catalog"
	"github.com/phil-mansfield/shellfish/cmd/env"
	"github.com/phil-mansfield/shellfish/logging"
	"github.com/phil-mansfield/shellfish/parse"
	"github.com/phil-mansfield/shellfish/io"
	"github.com/phil-mansfield/shellfish/cmd/memo"
)

type PotentialConfig struct {
	ncells int64
	rGridMult float64
	rMinMult, rMaxMult float64
	frac float64
}

var _ Mode = &PotentialConfig{}

func (config *PotentialConfig) ExampleConfig() string {
	return `[potential.config]
NCells = 64
GridRMult = 8
# RMinMult is the minimum radius inside which particles are used to calculate
# the potential.
RMinMult = 1
# RMaxMult is the maximum radius inside which particles are used to calculate
# the potential.
RMaxMult = 50
# Percentage of particles that will be used to compute the potential.
ParticleFraction = 0.01
`
}

func (config *PotentialConfig) ReadConfig(fname string, flags []string) error {
	vars := parse.NewConfigVars("prof.config")

	vars.Int(&config.ncells, "RBins", 64)
	vars.Float(&config.rGridMult, "GridRMult", 8)
	vars.Float(&config.rMaxMult, "RMaxMult", 1.0)
	vars.Float(&config.rMinMult, "RMinMult", 1.0)
	vars.Float(&config.frac, "ParticleFraction", 1.0)

	if fname == "" {
		if len(flags) == 0 { return nil }

		err := parse.ReadFlags(flags, vars)
		if err != nil { return err }
	} else {
		if err := parse.ReadConfig(fname, vars); err != nil { return err }
		if err := parse.ReadFlags(flags, vars); err != nil { return err }
	}

	return config.validate()
}

func (config *PotentialConfig) validate() error {
	if config.ncells < 0 {
		return fmt.Errorf("The variable '%s' was set to %d.",
			"NCells", config.ncells)
	} else if config.rGridMult < 0 {
		return fmt.Errorf("The variable '%s' was set to %d.",
			"RGridMult", config.rGridMult)
	} else if config.rMaxMult < 0 {
		return fmt.Errorf("The variable '%s' was set to %d.",
			"RMaxMult", config.rMaxMult)
	} else if config.rMinMult < 0 {
		return fmt.Errorf("The variable '%s' was set to %d.",
			"RMinMult", config.rMinMult)
	}

	return nil
}

func (config *PotentialConfig) Run(
	gConfig *GlobalConfig, e *env.Environment, stdin []byte,
) ([]string, error) {
	if logging.Mode != logging.Nil {
		log.Println(`
####################
## shellfish prof ##
####################`,
		)
	}

	var t time.Time
	if logging.Mode == logging.Performance {
		t = time.Now()
	}

	icols, fcols, err := catalog.Parse(
		stdin, []int{0, 1}, []int{2, 3, 4, 5, 6},
	)
	if err != nil { return nil, err }

	ids, snaps := icols[0], icols[1]
	hx := [3][]float64{ fcols[0], fcols[1], fcols[2] }
	hr, hm := fcols[3], fcols[4]

	if len(ids) == 0 { return nil, fmt.Errorf("No input halos.") }

	// Initialize phase profiles
	rSets := make([][]float64, len(ids))
	phiSets := [3][][]float64{
		make([][]float64, len(ids)),
		make([][]float64, len(ids)),
		make([][]float64, len(ids)),
	}

	for i := range rSets {
		rSets[i] = make([]float64, config.ncells)
		phiSets[0][i] = make([]float64, config.ncells*config.ncells)
		phiSets[1][i] = make([]float64, config.ncells*config.ncells)
		phiSets[2][i] = make([]float64, config.ncells*config.ncells)
	}

	snapBins, idxBins := binBySnap(snaps, ids)

	sortedSnaps := []int{}
	for snap := range snapBins {
		sortedSnaps = append(sortedSnaps, snap)
	}
	sort.Ints(sortedSnaps)

	buf, err := getVectorBuffer(
		e.ParticleCatalog(snaps[0], 0), gConfig,
	)
	if err != nil {
		return nil, err
	}

	// Count number of workers

	workers := runtime.NumCPU()
	if gConfig.Threads > 0 { workers = int(gConfig.Threads) }
	runtime.GOMAXPROCS(workers)

	for _, snap := range sortedSnaps {
		if snap == -1 {
			continue
		}

		idxs := idxBins[snap]
		snapCoords := [][]float64{
			make([]float64, len(idxs)), make([]float64, len(idxs)),
			make([]float64, len(idxs)), make([]float64, len(idxs)),
		}

		for i, idx := range idxs {
			snapCoords[0][i] = hx[0][idx]
			snapCoords[1][i] = hx[1][idx]
			snapCoords[2][i] = hx[2][idx]
			snapCoords[3][i] = hr[idx]*config.rMaxMult
		}

		hds, files, err := memo.ReadHeaders(snap, buf, e)
		if err != nil {
			return nil, err
		}
		hxBounds, err := boundingSpheres(snapCoords, &hds[0], e)
		if err != nil {
			return nil, err
		}
		_, intrIdxs := binSphereIntersections(hds, hxBounds)

		for i := range hds {
			if len(intrIdxs[i]) == 0 {
				continue
			}

			xs, _, ms, _, err := buf.Read(files[i])
			if err != nil {
				return nil, err
			}

			// Waarrrgggble
			for _, j := range intrIdxs[i] {
				phisXY := phiSets[0][idxs[j]]
				phisYZ := phiSets[1][idxs[j]]
				phisXZ := phiSets[2][idxs[j]]

				lg := NewLockGroup(workers)

				for k := 0; k < workers; k++ {
					insertPotentialPoints(
						phisXY, phisYZ, phisXZ,
						hxBounds[j],
						xs, ms,
						config, &hds[i],
						lg.Lock(k),
					)
				}

				lg.Synchronize()
			}

			buf.Close()
		}
	}

	for i := range rSets {
		processPotential(
			rSets[i],
			phiSets[0][i], phiSets[1][i], phiSets[2][i],
			hr[i], hm[i], config,
		)
	}

	rSets = transpose(rSets)
	phiSets[0] = transpose(phiSets[0])
	phiSets[1] = transpose(phiSets[1])
	phiSets[2] = transpose(phiSets[2])

	order := make([]int, len(rSets) + len(phiSets[0])*3 + 2)
	for i := range order { order[i] = i }
	lines := catalog.FormatCols(
		[][]int{ids, snaps},
		append(append(append(
			rSets, phiSets[0]...),
			phiSets[1]...),
			phiSets[2]...),
		order,
	)

	cString := catalog.CommentString(
		[]string{"ID", "Snapshot", "R [cMpc/h]",
			"Phi_xy/(G Mvir / Rvir)",
			"Phi_yz/(G Mvir / Rvir)",
			"Phi_xz/(G Mvir / Rvir)"},
		[]string{}, []int{0, 1, 2, 3, 4, 5},
		[]int{1, 1, int(config.ncells),
			int(config.ncells*config.ncells),
			int(config.ncells*config.ncells),
			int(config.ncells*config.ncells)},
	)

	if logging.Mode == logging.Performance {
		log.Printf("Time: %s", time.Since(t).String())
		log.Printf("Memory:\n%s", logging.MemString())
	}

	return append([]string{cString}, lines...), nil
}

func insertPotentialPoints(
	phisXY, phisYZ, phisXZ []float64,
	hx geom.Sphere,
	xs [][3]float32,
	ms []float32,
	config *PotentialConfig, hd *io.Header,
	lock *Lock,
) {
	rmax2 := (config.rMaxMult*config.rMaxMult) * float64(hx.R*hx.R)
	rmin2 := (config.rMinMult*config.rMinMult) * float64(hx.R*hx.R)

	gridR := float64(hx.R) * config.rGridMult
	delta := gridR * 2 / float64(config.ncells)
	nc := int(config.ncells)

	for i := range xs {
		// Think carefully about this bit. (although it actually makes
		// convergence tests easier to evaluate.)
		if rand.Float64() > config.frac { continue }

		dx0 := float64(xs[i][0] - hx.C[0])
		dy0 := float64(xs[i][1] - hx.C[1])
		dz0 := float64(xs[i][2] - hx.C[2])

		dr02 := dx0*dx0 + dy0*dy0 + dz0*dz0
		pm := float64(ms[i])

		if dr02 > rmax2 || dr02 < rmin2 { continue }

		for i := lock.Idx; i < len(phisXY); i += lock.Workers {
			var (
				ix, iy, iz int
				dx, dy, dz, dr float64
			)

			// phiXY
			ix, iy = i % nc, i / nc
			dx = dx0 - gridR + delta*float64(ix)
			dy = dy0 - gridR + delta*float64(iy)
			dr = math.Sqrt(dx*dx + dy*dy + dz0*dz0)
			phisXY[i] -= pm/dr

			// phiXZ
			ix, iz = i % nc, i / nc
			dx = dx0 - gridR + delta*float64(ix)
			dz = dz0 - gridR + delta*float64(iz)
			dr = math.Sqrt(dx*dx + dy0*dy0 + dz*dz)
			phisXZ[i] -= pm/dr

			// phiYZ
			iy, iz = i % nc, i / nc
			dy = dy0 - gridR + delta*float64(iy)
			dz = dz0 - gridR + delta*float64(iz)
			dr = math.Sqrt(dx0*dx0 + dy*dy + dz*dz)
			phisYZ[i] -= pm/dr
		}
	}

	lock.Unlock()
}

func processPotential(
	rs, phisXY, phisYZ, phisXZ []float64,
	hr, hm float64, config *PotentialConfig,
) {
	for i := range phisXY {
		phisXY[i] /= config.frac * hm/hr
		phisXZ[i] /= config.frac * hm/hr
		phisYZ[i] /= config.frac * hm/hr
	}

	delta := hr * config.rGridMult * 2 / float64(config.ncells)
	for i := range rs {
		rs[i] = delta*float64(i) - hr*config.rGridMult
	}
}
<commit_msg>Fxed bug with halo-bbox interseciton checks.<commit_after>package cmd

import (
	"fmt"
	"log"
	"math"
	"math/rand"
	"sort"
	"time"
	"runtime"

	"github.com/phil-mansfield/shellfish/los/geom"
	"github.com/phil-mansfield/shellfish/cmd/catalog"
	"github.com/phil-mansfield/shellfish/cmd/env"
	"github.com/phil-mansfield/shellfish/logging"
	"github.com/phil-mansfield/shellfish/parse"
	"github.com/phil-mansfield/shellfish/io"
	"github.com/phil-mansfield/shellfish/cmd/memo"
)

type PotentialConfig struct {
	ncells int64
	rGridMult float64
	rMinMult, rMaxMult float64
	frac float64
}

var _ Mode = &PotentialConfig{}

func (config *PotentialConfig) ExampleConfig() string {
	return `[potential.config]
NCells = 64
GridRMult = 8
# RMinMult is the minimum radius inside which particles are used to calculate
# the potential.
RMinMult = 1
# RMaxMult is the maximum radius inside which particles are used to calculate
# the potential.
RMaxMult = 50
# Percentage of particles that will be used to compute the potential.
ParticleFraction = 0.01
`
}

func (config *PotentialConfig) ReadConfig(fname string, flags []string) error {
	vars := parse.NewConfigVars("prof.config")

	vars.Int(&config.ncells, "RBins", 64)
	vars.Float(&config.rGridMult, "GridRMult", 8)
	vars.Float(&config.rMaxMult, "RMaxMult", 1.0)
	vars.Float(&config.rMinMult, "RMinMult", 1.0)
	vars.Float(&config.frac, "ParticleFraction", 1.0)

	if fname == "" {
		if len(flags) == 0 { return nil }

		err := parse.ReadFlags(flags, vars)
		if err != nil { return err }
	} else {
		if err := parse.ReadConfig(fname, vars); err != nil { return err }
		if err := parse.ReadFlags(flags, vars); err != nil { return err }
	}

	return config.validate()
}

func (config *PotentialConfig) validate() error {
	if config.ncells < 0 {
		return fmt.Errorf("The variable '%s' was set to %d.",
			"NCells", config.ncells)
	} else if config.rGridMult < 0 {
		return fmt.Errorf("The variable '%s' was set to %d.",
			"RGridMult", config.rGridMult)
	} else if config.rMaxMult < 0 {
		return fmt.Errorf("The variable '%s' was set to %d.",
			"RMaxMult", config.rMaxMult)
	} else if config.rMinMult < 0 {
		return fmt.Errorf("The variable '%s' was set to %d.",
			"RMinMult", config.rMinMult)
	}

	return nil
}

func (config *PotentialConfig) Run(
	gConfig *GlobalConfig, e *env.Environment, stdin []byte,
) ([]string, error) {
	if logging.Mode != logging.Nil {
		log.Println(`
####################
## shellfish prof ##
####################`,
		)
	}

	var t time.Time
	if logging.Mode == logging.Performance {
		t = time.Now()
	}

	icols, fcols, err := catalog.Parse(
		stdin, []int{0, 1}, []int{2, 3, 4, 5, 6},
	)
	if err != nil { return nil, err }

	ids, snaps := icols[0], icols[1]
	hx := [3][]float64{ fcols[0], fcols[1], fcols[2] }
	hr, hm := fcols[3], fcols[4]

	if len(ids) == 0 { return nil, fmt.Errorf("No input halos.") }

	// Initialize phase profiles
	rSets := make([][]float64, len(ids))
	phiSets := [3][][]float64{
		make([][]float64, len(ids)),
		make([][]float64, len(ids)),
		make([][]float64, len(ids)),
	}

	for i := range rSets {
		rSets[i] = make([]float64, config.ncells)
		phiSets[0][i] = make([]float64, config.ncells*config.ncells)
		phiSets[1][i] = make([]float64, config.ncells*config.ncells)
		phiSets[2][i] = make([]float64, config.ncells*config.ncells)
	}

	snapBins, idxBins := binBySnap(snaps, ids)

	sortedSnaps := []int{}
	for snap := range snapBins {
		sortedSnaps = append(sortedSnaps, snap)
	}
	sort.Ints(sortedSnaps)

	buf, err := getVectorBuffer(
		e.ParticleCatalog(snaps[0], 0), gConfig,
	)
	if err != nil {
		return nil, err
	}

	// Count number of workers

	workers := runtime.NumCPU()
	if gConfig.Threads > 0 { workers = int(gConfig.Threads) }
	runtime.GOMAXPROCS(workers)

	for _, snap := range sortedSnaps {
		if snap == -1 {
			continue
		}

		idxs := idxBins[snap]
		snapCoords := [][]float64{
			make([]float64, len(idxs)), make([]float64, len(idxs)),
			make([]float64, len(idxs)), make([]float64, len(idxs)),
		}

		for i, idx := range idxs {
			snapCoords[0][i] = hx[0][idx]
			snapCoords[1][i] = hx[1][idx]
			snapCoords[2][i] = hx[2][idx]
			snapCoords[3][i] = hr[idx]*config.rMaxMult
		}

		hds, files, err := memo.ReadHeaders(snap, buf, e)
		if err != nil { return nil, err }
		hxBounds, err := boundingSpheres(snapCoords, &hds[0], e)
		if err != nil { return nil, err }
		_, intrIdxs := binSphereIntersections(hds, hxBounds)

		for i := range hxBounds { hxBounds[i].R /= float32(config.rMaxMult) }

		for i := range hds {
			if len(intrIdxs[i]) == 0 { continue }

			xs, _, ms, _, err := buf.Read(files[i])
			if err != nil {
				return nil, err
			}

			table := make([]bool, len(xs))

			// Waarrrgggble
			for _, j := range intrIdxs[i] {
				phisXY := phiSets[0][idxs[j]]
				phisYZ := phiSets[1][idxs[j]]
				phisXZ := phiSets[2][idxs[j]]

				lg := NewLockGroup(workers)
				for i := range table { table[i] = rand.Float64() <= config.frac }

				for k := 0; k < workers; k++ {
					go insertPotentialPoints(
						phisXY, phisYZ, phisXZ,
						hxBounds[j],
						xs, ms, table,
						config, &hds[i],
						lg.Lock(k),
					)
				}

				lg.Synchronize()
			}

			buf.Close()
		}
	}

	for i := range rSets {
		processPotential(
			rSets[i],
			phiSets[0][i], phiSets[1][i], phiSets[2][i],
			hr[i], hm[i], config,
		)
	}

	rSets = transpose(rSets)
	phiSets[0] = transpose(phiSets[0])
	phiSets[1] = transpose(phiSets[1])
	phiSets[2] = transpose(phiSets[2])

	order := make([]int, len(rSets) + len(phiSets[0])*3 + 2)
	for i := range order { order[i] = i }
	lines := catalog.FormatCols(
		[][]int{ids, snaps},
		append(append(append(
			rSets, phiSets[0]...),
			phiSets[1]...),
			phiSets[2]...),
		order,
	)

	cString := catalog.CommentString(
		[]string{"ID", "Snapshot", "R [cMpc/h]",
			"Phi_xy/(G Mvir / Rvir)",
			"Phi_yz/(G Mvir / Rvir)",
			"Phi_xz/(G Mvir / Rvir)"},
		[]string{}, []int{0, 1, 2, 3, 4, 5},
		[]int{1, 1, int(config.ncells),
			int(config.ncells*config.ncells),
			int(config.ncells*config.ncells),
			int(config.ncells*config.ncells)},
	)

	if logging.Mode == logging.Performance {
		log.Printf("Time: %s", time.Since(t).String())
		log.Printf("Memory:\n%s", logging.MemString())
	}

	return append([]string{cString}, lines...), nil
}

func insertPotentialPoints(
	phisXY, phisYZ, phisXZ []float64,
	hx geom.Sphere,
	xs [][3]float32,
	ms []float32,
	skipTable []bool,
	config *PotentialConfig, hd *io.Header,
	lock *Lock,
) {
	rmax2 := (config.rMaxMult*config.rMaxMult) * float64(hx.R*hx.R)
	rmin2 := (config.rMinMult*config.rMinMult) * float64(hx.R*hx.R)

	gridR := float64(hx.R) * config.rGridMult
	delta := gridR * 2 / float64(config.ncells - 1)
	nc := int(config.ncells)

	for i := 0; i < len(xs); i++ {
		if !skipTable[i] { continue }

		dx0 := float64(xs[i][0] - hx.C[0])
		dy0 := float64(xs[i][1] - hx.C[1])
		dz0 := float64(xs[i][2] - hx.C[2])

		dr02 := dx0*dx0 + dy0*dy0 + dz0*dz0
		pm := float64(ms[i])

		if dr02 > rmax2 || dr02 < rmin2 { continue }

		for i := lock.Idx; i < len(phisXY); i += lock.Workers {
			var (
				ix, iy, iz int
				dx, dy, dz, dr float64
			)

			// phiXY
			ix, iy = i % nc, i / nc
			dx = dx0 - gridR + delta*float64(ix)
			dy = dy0 - gridR + delta*float64(iy)
			dr = math.Sqrt(dx*dx + dy*dy + dz0*dz0)
			phisXY[i] -= pm/dr

			// phiXZ
			ix, iz = i % nc, i / nc
			dx = dx0 - gridR + delta*float64(ix)
			dz = dz0 - gridR + delta*float64(iz)
			dr = math.Sqrt(dx*dx + dy0*dy0 + dz*dz)
			phisXZ[i] -= pm/dr

			// phiYZ
			iy, iz = i % nc, i / nc
			dy = dy0 - gridR + delta*float64(iy)
			dz = dz0 - gridR + delta*float64(iz)
			dr = math.Sqrt(dx0*dx0 + dy*dy + dz*dz)
			phisYZ[i] -= pm/dr
		}
	}

	lock.Unlock()
}

func processPotential(
	rs, phisXY, phisYZ, phisXZ []float64,
	hr, hm float64, config *PotentialConfig,
) {
	for i := range phisXY {
		phisXY[i] /= config.frac * hm/hr
		phisXZ[i] /= config.frac * hm/hr
		phisYZ[i] /= config.frac * hm/hr
	}

	delta := hr * config.rGridMult * 2 / float64(config.ncells)
	for i := range rs {
		rs[i] = delta*float64(i) - hr*config.rGridMult
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 The go-ethereum Authors
// This file is part of go-ethereum.
//
// go-ethereum is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// go-ethereum is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with go-ethereum. If not, see <http://www.gnu.org/licenses/>.

// Package utils contains internal helper functions for go-ethereum commands.
package utils

import (
	"fmt"
	"io"
	"os"
	"os/signal"
	"regexp"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core"
	"github.com/ethereum/go-ethereum/core/types"
	"github.com/ethereum/go-ethereum/internal/debug"
	"github.com/ethereum/go-ethereum/logger"
	"github.com/ethereum/go-ethereum/logger/glog"
	"github.com/ethereum/go-ethereum/node"
	"github.com/ethereum/go-ethereum/rlp"
)

const (
	importBatchSize = 2500
)

func openLogFile(Datadir string, filename string) *os.File {
	path := common.AbsolutePath(Datadir, filename)
	file, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
	if err != nil {
		panic(fmt.Sprintf("error opening log file '%s': %v", filename, err))
	}
	return file
}

// Fatalf formats a message to standard error and exits the program.
// The message is also printed to standard output if standard error
// is redirected to a different file.
func Fatalf(format string, args ...interface{}) {
	w := io.MultiWriter(os.Stdout, os.Stderr)
	outf, _ := os.Stdout.Stat()
	errf, _ := os.Stderr.Stat()
	if outf != nil && errf != nil && os.SameFile(outf, errf) {
		w = os.Stderr
	}
	fmt.Fprintf(w, "Fatal: "+format+"\n", args...)
	logger.Flush()
	os.Exit(1)
}

func StartNode(stack *node.Node) {
	if err := stack.Start(); err != nil {
		Fatalf("Error starting protocol stack: %v", err)
	}
	go func() {
		sigc := make(chan os.Signal, 1)
		signal.Notify(sigc, os.Interrupt)
		defer signal.Stop(sigc)
		<-sigc
		glog.V(logger.Info).Infoln("Got interrupt, shutting down...")
		go stack.Stop()
		logger.Flush()
		for i := 10; i > 0; i-- {
			<-sigc
			if i > 1 {
				glog.V(logger.Info).Infoln("Already shutting down, please be patient.")
				glog.V(logger.Info).Infoln("Interrupt", i-1, "more times to induce panic.")
			}
		}
		glog.V(logger.Error).Infof("Force quitting: this might not end so well.")
		debug.LoudPanic("boom")
	}()
}

func FormatTransactionData(data string) []byte {
	d := common.StringToByteFunc(data, func(s string) (ret []byte) {
		slice := regexp.MustCompile("\\n|\\s").Split(s, 1000000000)
		for _, dataItem := range slice {
			d := common.FormatData(dataItem)
			ret = append(ret, d...)
		}
		return
	})

	return d
}

func ImportChain(chain *core.BlockChain, fn string) error {
	// Watch for Ctrl-C while the import is running.
	// If a signal is received, the import will stop at the next batch.
	interrupt := make(chan os.Signal, 1)
	stop := make(chan struct{})
	signal.Notify(interrupt, os.Interrupt)
	defer signal.Stop(interrupt)
	defer close(interrupt)
	go func() {
		if _, ok := <-interrupt; ok {
			glog.Info("caught interrupt during import, will stop at next batch")
		}
		close(stop)
	}()
	checkInterrupt := func() bool {
		select {
		case <-stop:
			return true
		default:
			return false
		}
	}

	glog.Infoln("Importing blockchain", fn)
	fh, err := os.Open(fn)
	if err != nil {
		return err
	}
	defer fh.Close()
	stream := rlp.NewStream(fh, 0)

	// Run actual the import.
	blocks := make(types.Blocks, importBatchSize)
	n := 0
	for batch := 0; ; batch++ {
		// Load a batch of RLP blocks.
		if checkInterrupt() {
			return fmt.Errorf("interrupted")
		}
		i := 0
		for ; i < importBatchSize; i++ {
			var b types.Block
			if err := stream.Decode(&b); err == io.EOF {
				break
			} else if err != nil {
				return fmt.Errorf("at block %d: %v", n, err)
			}
			// don't import first block
			if b.NumberU64() == 0 {
				i--
				continue
			}
			blocks[i] = &b
			n++
		}
		if i == 0 {
			break
		}
		// Import the batch.
		if checkInterrupt() {
			return fmt.Errorf("interrupted")
		}
		if hasAllBlocks(chain, blocks[:i]) {
			glog.Infof("skipping batch %d, all blocks present [%x / %x]",
				batch, blocks[0].Hash().Bytes()[:4], blocks[i-1].Hash().Bytes()[:4])
			continue
		}

		if _, err := chain.InsertChain(blocks[:i]); err != nil {
			return fmt.Errorf("invalid block %d: %v", n, err)
		}
	}
	return nil
}

func hasAllBlocks(chain *core.BlockChain, bs []*types.Block) bool {
	for _, b := range bs {
		if !chain.HasBlock(b.Hash()) {
			return false
		}
	}
	return true
}

func ExportChain(blockchain *core.BlockChain, fn string) error {
	glog.Infoln("Exporting blockchain to", fn)
	fh, err := os.OpenFile(fn, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, os.ModePerm)
	if err != nil {
		return err
	}
	defer fh.Close()
	if err := blockchain.Export(fh); err != nil {
		return err
	}
	glog.Infoln("Exported blockchain to", fn)
	return nil
}

func ExportAppendChain(blockchain *core.BlockChain, fn string, first uint64, last uint64) error {
	glog.Infoln("Exporting blockchain to", fn)
	// TODO verify mode perms
	fh, err := os.OpenFile(fn, os.O_CREATE|os.O_APPEND|os.O_WRONLY, os.ModePerm)
	if err != nil {
		return err
	}
	defer fh.Close()
	if err := blockchain.ExportN(fh, first, last); err != nil {
		return err
	}
	glog.Infoln("Exported blockchain to", fn)
	return nil
}
<commit_msg>cmd/utils: flush trace and CPU profile data when force-qutting<commit_after>// Copyright 2014 The go-ethereum Authors
// This file is part of go-ethereum.
//
// go-ethereum is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// go-ethereum is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with go-ethereum. If not, see <http://www.gnu.org/licenses/>.

// Package utils contains internal helper functions for go-ethereum commands.
package utils

import (
	"fmt"
	"io"
	"os"
	"os/signal"
	"regexp"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core"
	"github.com/ethereum/go-ethereum/core/types"
	"github.com/ethereum/go-ethereum/internal/debug"
	"github.com/ethereum/go-ethereum/logger"
	"github.com/ethereum/go-ethereum/logger/glog"
	"github.com/ethereum/go-ethereum/node"
	"github.com/ethereum/go-ethereum/rlp"
)

const (
	importBatchSize = 2500
)

func openLogFile(Datadir string, filename string) *os.File {
	path := common.AbsolutePath(Datadir, filename)
	file, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
	if err != nil {
		panic(fmt.Sprintf("error opening log file '%s': %v", filename, err))
	}
	return file
}

// Fatalf formats a message to standard error and exits the program.
// The message is also printed to standard output if standard error
// is redirected to a different file.
func Fatalf(format string, args ...interface{}) {
	w := io.MultiWriter(os.Stdout, os.Stderr)
	outf, _ := os.Stdout.Stat()
	errf, _ := os.Stderr.Stat()
	if outf != nil && errf != nil && os.SameFile(outf, errf) {
		w = os.Stderr
	}
	fmt.Fprintf(w, "Fatal: "+format+"\n", args...)
	logger.Flush()
	os.Exit(1)
}

func StartNode(stack *node.Node) {
	if err := stack.Start(); err != nil {
		Fatalf("Error starting protocol stack: %v", err)
	}
	go func() {
		sigc := make(chan os.Signal, 1)
		signal.Notify(sigc, os.Interrupt)
		defer signal.Stop(sigc)
		<-sigc
		glog.V(logger.Info).Infoln("Got interrupt, shutting down...")
		go stack.Stop()
		for i := 10; i > 0; i-- {
			<-sigc
			if i > 1 {
				glog.V(logger.Info).Infof("Already shutting down, interrupt %d more times for panic.", i-1)
			}
		}
		debug.Exit() // ensure trace and CPU profile data is flushed.
		debug.LoudPanic("boom")
	}()
}

func FormatTransactionData(data string) []byte {
	d := common.StringToByteFunc(data, func(s string) (ret []byte) {
		slice := regexp.MustCompile("\\n|\\s").Split(s, 1000000000)
		for _, dataItem := range slice {
			d := common.FormatData(dataItem)
			ret = append(ret, d...)
		}
		return
	})

	return d
}

func ImportChain(chain *core.BlockChain, fn string) error {
	// Watch for Ctrl-C while the import is running.
	// If a signal is received, the import will stop at the next batch.
	interrupt := make(chan os.Signal, 1)
	stop := make(chan struct{})
	signal.Notify(interrupt, os.Interrupt)
	defer signal.Stop(interrupt)
	defer close(interrupt)
	go func() {
		if _, ok := <-interrupt; ok {
			glog.Info("caught interrupt during import, will stop at next batch")
		}
		close(stop)
	}()
	checkInterrupt := func() bool {
		select {
		case <-stop:
			return true
		default:
			return false
		}
	}

	glog.Infoln("Importing blockchain", fn)
	fh, err := os.Open(fn)
	if err != nil {
		return err
	}
	defer fh.Close()
	stream := rlp.NewStream(fh, 0)

	// Run actual the import.
	blocks := make(types.Blocks, importBatchSize)
	n := 0
	for batch := 0; ; batch++ {
		// Load a batch of RLP blocks.
		if checkInterrupt() {
			return fmt.Errorf("interrupted")
		}
		i := 0
		for ; i < importBatchSize; i++ {
			var b types.Block
			if err := stream.Decode(&b); err == io.EOF {
				break
			} else if err != nil {
				return fmt.Errorf("at block %d: %v", n, err)
			}
			// don't import first block
			if b.NumberU64() == 0 {
				i--
				continue
			}
			blocks[i] = &b
			n++
		}
		if i == 0 {
			break
		}
		// Import the batch.
		if checkInterrupt() {
			return fmt.Errorf("interrupted")
		}
		if hasAllBlocks(chain, blocks[:i]) {
			glog.Infof("skipping batch %d, all blocks present [%x / %x]",
				batch, blocks[0].Hash().Bytes()[:4], blocks[i-1].Hash().Bytes()[:4])
			continue
		}

		if _, err := chain.InsertChain(blocks[:i]); err != nil {
			return fmt.Errorf("invalid block %d: %v", n, err)
		}
	}
	return nil
}

func hasAllBlocks(chain *core.BlockChain, bs []*types.Block) bool {
	for _, b := range bs {
		if !chain.HasBlock(b.Hash()) {
			return false
		}
	}
	return true
}

func ExportChain(blockchain *core.BlockChain, fn string) error {
	glog.Infoln("Exporting blockchain to", fn)
	fh, err := os.OpenFile(fn, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, os.ModePerm)
	if err != nil {
		return err
	}
	defer fh.Close()
	if err := blockchain.Export(fh); err != nil {
		return err
	}
	glog.Infoln("Exported blockchain to", fn)
	return nil
}

func ExportAppendChain(blockchain *core.BlockChain, fn string, first uint64, last uint64) error {
	glog.Infoln("Exporting blockchain to", fn)
	// TODO verify mode perms
	fh, err := os.OpenFile(fn, os.O_CREATE|os.O_APPEND|os.O_WRONLY, os.ModePerm)
	if err != nil {
		return err
	}
	defer fh.Close()
	if err := blockchain.ExportN(fh, first, last); err != nil {
		return err
	}
	glog.Infoln("Exported blockchain to", fn)
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"log"
	"os"
	"time"

	"github.com/datawire/teleproxy/pkg/k8s"
	"github.com/datawire/teleproxy/pkg/limiter"
	"github.com/datawire/teleproxy/pkg/supervisor"
	"github.com/spf13/cobra"
)

var kubernetesNamespace string
var initialSources = make([]string, 0)
var watchHooks = make([]string, 0)
var notifyReceivers = make([]string, 0)
var port int
var interval time.Duration

var rootCmd = &cobra.Command{
	Use:              "watt",
	Short:            "watt",
	Long:             "watt - watch all the things",
	PersistentPreRun: func(cmd *cobra.Command, args []string) {},
	Run:              runWatt,
}

func init() {
	rootCmd.Flags().StringVarP(&kubernetesNamespace, "namespace", "n", "", "namespace to watch (default: all)")
	rootCmd.Flags().StringSliceVarP(&initialSources, "source", "s", []string{}, "configure an initial static source")
	rootCmd.Flags().StringSliceVarP(&watchHooks, "watch", "w", []string{}, "configure watch hook(s)")
	rootCmd.Flags().StringSliceVar(&notifyReceivers, "notify", []string{},
		"invoke the program with the given arguments as a receiver")
	rootCmd.Flags().IntVarP(&port, "port", "p", 7000, "configure the snapshot server port")
	rootCmd.Flags().DurationVarP(&interval, "interval", "i", 250*time.Millisecond,
		"configure the rate limit interval")
}

func runWatt(cmd *cobra.Command, args []string) {
	if len(initialSources) == 0 {
		log.Fatalln("no initial sources configured")
	}

	client := k8s.NewClient(nil)
	kubeAPIWatcher := client.Watcher()
	for idx := range initialSources {
		initialSources[idx] = kubeAPIWatcher.Canonical(initialSources[idx])
	}

	log.Printf("starting watt...")

	// The aggregator sends the current consul resolver set to the
	// consul watch manager.
	aggregatorToConsulwatchmanCh := make(chan []ConsulWatchSpec)

	// The aggregator sends the current k8s watch set to the
	// kubernetes watch manager.
	aggregatorToKubewatchmanCh := make(chan []KubernetesWatchSpec)

	invoker := NewInvoker(port, notifyReceivers)
	limiter := limiter.NewComposite(limiter.NewUnlimited(), limiter.NewInterval(interval), interval)
	aggregator := NewAggregator(invoker.Snapshots, aggregatorToKubewatchmanCh, aggregatorToConsulwatchmanCh,
		initialSources, ExecWatchHook(watchHooks), limiter)

	kubebootstrap := kubebootstrap{
		namespace:      kubernetesNamespace,
		kinds:          initialSources,
		kubeAPIWatcher: kubeAPIWatcher,
		notify:         []chan<- k8sEvent{aggregator.KubernetesEvents},
	}

	consulwatchman := consulwatchman{
		WatchMaker: &ConsulWatchMaker{aggregatorCh: aggregator.ConsulEvents},
		watchesCh:  aggregatorToConsulwatchmanCh,
		watched:    make(map[string]*supervisor.Worker),
	}

	kubewatchman := kubewatchman{
		WatchMaker: &KubernetesWatchMaker{kubeAPI: client, notify: aggregator.KubernetesEvents},
		in:         aggregatorToKubewatchmanCh,
	}

	apiServer := &apiServer{
		port:    port,
		invoker: invoker,
	}

	ctx := context.Background()
	s := supervisor.WithContext(ctx)

	s.Supervise(&supervisor.Worker{
		Name: "kubebootstrap",
		Work: kubebootstrap.Work,
	})

	s.Supervise(&supervisor.Worker{
		Name: "consulwatchman",
		Work: consulwatchman.Work,
	})

	s.Supervise(&supervisor.Worker{
		Name: "kubewatchman",
		Work: kubewatchman.Work,
	})

	s.Supervise(&supervisor.Worker{
		Name: "aggregator",
		Work: aggregator.Work,
	})

	s.Supervise(&supervisor.Worker{
		Name: "invoker",
		Work: invoker.Work,
	})

	s.Supervise(&supervisor.Worker{
		Name: "api",
		Work: apiServer.Work,
	})

	if errs := s.Run(); len(errs) > 0 {
		for _, err := range errs {
			log.Println(err)
		}
		os.Exit(1)
	}
}

func main() {
	if err := rootCmd.Execute(); err != nil {
		log.Fatalln(err)
	}
}
<commit_msg>use kinds as passed in<commit_after>package main

import (
	"context"
	"log"
	"os"
	"time"

	"github.com/datawire/teleproxy/pkg/k8s"
	"github.com/datawire/teleproxy/pkg/limiter"
	"github.com/datawire/teleproxy/pkg/supervisor"
	"github.com/spf13/cobra"
)

var kubernetesNamespace string
var initialSources = make([]string, 0)
var watchHooks = make([]string, 0)
var notifyReceivers = make([]string, 0)
var port int
var interval time.Duration

var rootCmd = &cobra.Command{
	Use:              "watt",
	Short:            "watt",
	Long:             "watt - watch all the things",
	PersistentPreRun: func(cmd *cobra.Command, args []string) {},
	Run:              runWatt,
}

func init() {
	rootCmd.Flags().StringVarP(&kubernetesNamespace, "namespace", "n", "", "namespace to watch (default: all)")
	rootCmd.Flags().StringSliceVarP(&initialSources, "source", "s", []string{}, "configure an initial static source")
	rootCmd.Flags().StringSliceVarP(&watchHooks, "watch", "w", []string{}, "configure watch hook(s)")
	rootCmd.Flags().StringSliceVar(&notifyReceivers, "notify", []string{},
		"invoke the program with the given arguments as a receiver")
	rootCmd.Flags().IntVarP(&port, "port", "p", 7000, "configure the snapshot server port")
	rootCmd.Flags().DurationVarP(&interval, "interval", "i", 250*time.Millisecond,
		"configure the rate limit interval")
}

func runWatt(cmd *cobra.Command, args []string) {
	if len(initialSources) == 0 {
		log.Fatalln("no initial sources configured")
	}

	// XXX: we don't need to create this here anymore
	client := k8s.NewClient(nil)
	kubeAPIWatcher := client.Watcher()
	/*for idx := range initialSources {
		initialSources[idx] = kubeAPIWatcher.Canonical(initialSources[idx])
	}*/

	log.Printf("starting watt...")

	// The aggregator sends the current consul resolver set to the
	// consul watch manager.
	aggregatorToConsulwatchmanCh := make(chan []ConsulWatchSpec)

	// The aggregator sends the current k8s watch set to the
	// kubernetes watch manager.
	aggregatorToKubewatchmanCh := make(chan []KubernetesWatchSpec)

	invoker := NewInvoker(port, notifyReceivers)
	limiter := limiter.NewComposite(limiter.NewUnlimited(), limiter.NewInterval(interval), interval)
	aggregator := NewAggregator(invoker.Snapshots, aggregatorToKubewatchmanCh, aggregatorToConsulwatchmanCh,
		initialSources, ExecWatchHook(watchHooks), limiter)

	kubebootstrap := kubebootstrap{
		namespace:      kubernetesNamespace,
		kinds:          initialSources,
		kubeAPIWatcher: kubeAPIWatcher,
		notify:         []chan<- k8sEvent{aggregator.KubernetesEvents},
	}

	consulwatchman := consulwatchman{
		WatchMaker: &ConsulWatchMaker{aggregatorCh: aggregator.ConsulEvents},
		watchesCh:  aggregatorToConsulwatchmanCh,
		watched:    make(map[string]*supervisor.Worker),
	}

	kubewatchman := kubewatchman{
		WatchMaker: &KubernetesWatchMaker{kubeAPI: client, notify: aggregator.KubernetesEvents},
		in:         aggregatorToKubewatchmanCh,
	}

	apiServer := &apiServer{
		port:    port,
		invoker: invoker,
	}

	ctx := context.Background()
	s := supervisor.WithContext(ctx)

	s.Supervise(&supervisor.Worker{
		Name: "kubebootstrap",
		Work: kubebootstrap.Work,
	})

	s.Supervise(&supervisor.Worker{
		Name: "consulwatchman",
		Work: consulwatchman.Work,
	})

	s.Supervise(&supervisor.Worker{
		Name: "kubewatchman",
		Work: kubewatchman.Work,
	})

	s.Supervise(&supervisor.Worker{
		Name: "aggregator",
		Work: aggregator.Work,
	})

	s.Supervise(&supervisor.Worker{
		Name: "invoker",
		Work: invoker.Work,
	})

	s.Supervise(&supervisor.Worker{
		Name: "api",
		Work: apiServer.Work,
	})

	if errs := s.Run(); len(errs) > 0 {
		for _, err := range errs {
			log.Println(err)
		}
		os.Exit(1)
	}
}

func main() {
	if err := rootCmd.Execute(); err != nil {
		log.Fatalln(err)
	}
}
<|endoftext|>


<commit_before>package task

import (
	"encoding/json"
	"errors"
	"fmt"
	"regexp"
	"strings"
	"time"

	"github.com/spf13/cobra"
	"github.com/taskcluster/slugid-go/slugid"
	"github.com/taskcluster/taskcluster-cli/config"
	tcclient "github.com/taskcluster/taskcluster-client-go"
	"github.com/taskcluster/taskcluster-client-go/queue"
)

var (
	runCmd = &cobra.Command{
		Use:     "run <image> <command>",
		Short:   "creates and schedules a task through a 'docker run'-like interface.",
		PreRunE: checkRunFlags,
		RunE:    runRunTask,
	}

	now = time.Now().UTC()

	runPayload = &queue.TaskDefinitionRequest{
		SchedulerID: "taskcluster-cli",

		Created:  tcclient.Time(now),
		Deadline: tcclient.Time(now.Add(24 * time.Hour)),
		Expires:  tcclient.Time(now.Add(24*time.Hour).AddDate(1, 0, 0)),
	}

	requiredFlags = []string{
		"provisioner",
		"worker-type",
	}
)

func init() {
	fs := runCmd.Flags()

	fs.StringVar(&runPayload.ProvisionerID, "provisioner", "", "ID of the provisioner to use")
	fs.StringVar(&runPayload.WorkerType, "worker-type", "", "worker-type to use within the provisioner")
	fs.StringSliceP("env", "e", []string{}, "Environment variable to add to the task's environment (repeatable) (format: VARIABLE=VALUE)")
	fs.StringVar(&runPayload.Metadata.Name, "name", "Taskcluster-cli Task", "Human readable name of the task")
	fs.StringVar(&runPayload.Metadata.Description, "description", "Created by Taskcluster-cli", "Human readable description of the task")
	fs.StringVar(&runPayload.Metadata.Owner, "owner", "name@example.com", "Email of the task's owner")
	fs.StringVar(&runPayload.Metadata.Source, "source", "http://taskcluster-cli/task/run", "URL pointing to the source of the task")
	fs.StringSliceVar(&runPayload.Dependencies, "dependency", []string{}, "TaskID of a dependency (repeatable)")
	fs.IntVar(&runPayload.Retries, "retries", 5, "Number of retries due to infrastructure issues")

	for _, f := range requiredFlags {
		runCmd.MarkFlagRequired(f)
	}

	Command.AddCommand(runCmd)
}

// checkRunFlags checks that the required flags were specified.
func checkRunFlags(cmd *cobra.Command, args []string) error {
	for _, f := range requiredFlags {
		if !cmd.Flag(f).Changed {
			return fmt.Errorf("flag '%s' is required", f)
		}
	}
	if owner := stringFlagHelper(cmd.Flags(), "owner"); !regexp.MustCompile(".+@.+").MatchString(owner) {
		return errors.New("owner must be an email-like string")
	}
	return nil
}

// runRunTask takes the task creation payload and runs the task.
func runRunTask(cmd *cobra.Command, args []string) error {
	if len(args) < 2 {
		return errors.New("run requires at least 2 arguments: image and command")
	}

	var creds *tcclient.Credentials
	if config.Credentials != nil {
		creds = config.Credentials.ToClientCredentials()
	}

	// Generate a new taskID
	taskID := slugid.V4()
	runPayload.TaskGroupID = taskID

	// Build the environment variables.
	env := make(map[string]string)
	envs, err := cmd.Flags().GetStringSlice("env")
	if err != nil {
		return err
	}
	for _, e := range envs {
		p := strings.SplitN(e, "=", 2)
		switch len(p) {
		case 2:
			env[p[0]] = p[1]
		case 1:
			env[p[0]] = ""
		default:
			return fmt.Errorf("invalid environment option: %s", e)
		}
	}

	// Build the task payload.
	runPayload.Payload, err = json.Marshal(struct {
		Image       string            `json:"image"`
		Command     []string          `json:"command"`
		Environment map[string]string `json:"env"`
		MaxRunTime  int               `json:"maxRunTime"`
	}{
		Image:       args[0],
		Command:     args[1:],
		Environment: env,
		MaxRunTime:  7200, // 2 hours
	})
	if err != nil {
		return fmt.Errorf("could not marshal execution payload: %v", err)
	}

	q := queue.New(creds)
	resp, err := q.CreateTask(taskID, runPayload)
	if err != nil {
		return fmt.Errorf("could not create task: %v", err)
	}

	// If we got no error, that means the task was successfully rund
	fmt.Fprintf(cmd.OutOrStdout(), "Task %s created\n", resp.Status.TaskID)

	return nil
}
<commit_msg>Update one flag help string<commit_after>package task

import (
	"encoding/json"
	"errors"
	"fmt"
	"regexp"
	"strings"
	"time"

	"github.com/spf13/cobra"
	"github.com/taskcluster/slugid-go/slugid"
	"github.com/taskcluster/taskcluster-cli/config"
	tcclient "github.com/taskcluster/taskcluster-client-go"
	"github.com/taskcluster/taskcluster-client-go/queue"
)

var (
	runCmd = &cobra.Command{
		Use:     "run <image> <command>",
		Short:   "creates and schedules a task through a 'docker run'-like interface.",
		PreRunE: checkRunFlags,
		RunE:    runRunTask,
	}

	now = time.Now().UTC()

	runPayload = &queue.TaskDefinitionRequest{
		SchedulerID: "taskcluster-cli",

		Created:  tcclient.Time(now),
		Deadline: tcclient.Time(now.Add(24 * time.Hour)),
		Expires:  tcclient.Time(now.Add(24*time.Hour).AddDate(1, 0, 0)),
	}

	requiredFlags = []string{
		"provisioner",
		"worker-type",
	}
)

func init() {
	fs := runCmd.Flags()

	fs.StringVar(&runPayload.ProvisionerID, "provisioner", "", "ID of the provisioner to use")
	fs.StringVar(&runPayload.WorkerType, "worker-type", "", "Type of worker to use within the provisioner")
	fs.StringSliceP("env", "e", []string{}, "Environment variable to add to the task's environment (repeatable) (format: VARIABLE=VALUE)")
	fs.StringVar(&runPayload.Metadata.Name, "name", "Taskcluster-cli Task", "Human readable name of the task")
	fs.StringVar(&runPayload.Metadata.Description, "description", "Created by Taskcluster-cli", "Human readable description of the task")
	fs.StringVar(&runPayload.Metadata.Owner, "owner", "name@example.com", "Email of the task's owner")
	fs.StringVar(&runPayload.Metadata.Source, "source", "http://taskcluster-cli/task/run", "URL pointing to the source of the task")
	fs.StringSliceVar(&runPayload.Dependencies, "dependency", []string{}, "TaskID of a dependency (repeatable)")
	fs.IntVar(&runPayload.Retries, "retries", 5, "Number of retries due to infrastructure issues")

	for _, f := range requiredFlags {
		runCmd.MarkFlagRequired(f)
	}

	Command.AddCommand(runCmd)
}

// checkRunFlags checks that the required flags were specified.
func checkRunFlags(cmd *cobra.Command, args []string) error {
	for _, f := range requiredFlags {
		if !cmd.Flag(f).Changed {
			return fmt.Errorf("flag '%s' is required", f)
		}
	}
	if owner := stringFlagHelper(cmd.Flags(), "owner"); !regexp.MustCompile(".+@.+").MatchString(owner) {
		return errors.New("owner must be an email-like string")
	}
	return nil
}

// runRunTask takes the task creation payload and runs the task.
func runRunTask(cmd *cobra.Command, args []string) error {
	if len(args) < 2 {
		return errors.New("run requires at least 2 arguments: image and command")
	}

	var creds *tcclient.Credentials
	if config.Credentials != nil {
		creds = config.Credentials.ToClientCredentials()
	}

	// Generate a new taskID
	taskID := slugid.V4()
	runPayload.TaskGroupID = taskID

	// Build the environment variables.
	env := make(map[string]string)
	envs, err := cmd.Flags().GetStringSlice("env")
	if err != nil {
		return err
	}
	for _, e := range envs {
		p := strings.SplitN(e, "=", 2)
		switch len(p) {
		case 2:
			env[p[0]] = p[1]
		case 1:
			env[p[0]] = ""
		default:
			return fmt.Errorf("invalid environment option: %s", e)
		}
	}

	// Build the task payload.
	runPayload.Payload, err = json.Marshal(struct {
		Image       string            `json:"image"`
		Command     []string          `json:"command"`
		Environment map[string]string `json:"env"`
		MaxRunTime  int               `json:"maxRunTime"`
	}{
		Image:       args[0],
		Command:     args[1:],
		Environment: env,
		MaxRunTime:  7200, // 2 hours
	})
	if err != nil {
		return fmt.Errorf("could not marshal execution payload: %v", err)
	}

	q := queue.New(creds)
	resp, err := q.CreateTask(taskID, runPayload)
	if err != nil {
		return fmt.Errorf("could not create task: %v", err)
	}

	// If we got no error, that means the task was successfully rund
	fmt.Fprintf(cmd.OutOrStdout(), "Task %s created\n", resp.Status.TaskID)

	return nil
}
<|endoftext|>


<commit_before>package qb

import (
	"database/sql"
	"fmt"
	"sort"
	"strings"
)

// New generates a new Session given engine and returns session pointer
func New(driver string, dsn string) (*Session, error) {

	engine, err := NewEngine(driver, dsn)
	if err != nil {
		return nil, err
	}

	return &Session{
		queries:  []*Query{},
		mapper:   NewMapper(engine.Driver()),
		metadata: NewMetaData(engine),
		builder:  NewBuilder(engine.Driver()),
	}, nil
}

// Session is the composition of engine connection & orm mappings
type Session struct {
	queries  []*Query
	mapper   *Mapper
	metadata *MetaData
	tx       *sql.Tx
	builder  *Builder
}

func (s *Session) add(query *Query) {
	var err error
	if s.tx == nil {
		s.queries = []*Query{}
		s.tx, err = s.metadata.Engine().DB().Begin()
		if err != nil {
			panic(err)
		}
	}
	s.queries = append(s.queries, query)
}

// Close closes engine db (sqlx) connection
func (s *Session) Close() {
	s.metadata.Engine().DB().Close()
}

// Builder returns query builder
func (s *Session) Builder() *Builder{
	return s.builder
}

// AddQuery adds a query given the query pointer retrieved from Query() function
func (s *Session) AddQuery(query *Query) {
	s.add(query)
}

// Query returns the active query built by session
func (s *Session) Query() *Query {
	return s.builder.Query()
}

// Metadata returns the metadata of session
func (s *Session) Metadata() *MetaData {
	return s.metadata
}

// Delete adds a single delete query to the session
func (s *Session) Delete(model interface{}) {

	kv := s.mapper.ToMap(model)

	tName := s.mapper.ModelName(model)

	d := s.metadata.Table(tName).Delete()
	ands := []string{}
	bindings := []interface{}{}
	for k, v := range kv {
		ands = append(ands, fmt.Sprintf("%s = %s", s.mapper.ColName(k), s.builder.Dialect().Placeholder()))
		bindings = append(bindings, v)
	}

	del := d.Where(d.And(ands...), bindings...).Query()
	s.add(del)
}

// Add adds a single model to the session. The query must be insert or update
func (s *Session) Add(model interface{}) {

	rawMap := s.mapper.ToMap(model)

	kv := map[string]interface{}{}

	for k, v := range rawMap {
		kv[s.mapper.ColName(k)] = v
	}

	q := s.metadata.Table(s.mapper.ModelName(model)).Insert(kv).Query()
	s.add(q)
}

// AddAll adds multiple models an adds an insert statement to current queries
func (s *Session) AddAll(models ...interface{}) {
	for _, m := range models {
		s.Add(m)
	}
}

// Commit commits the current transaction with queries
func (s *Session) Commit() error {

	for _, q := range s.queries {
		_, err := s.tx.Exec(q.SQL(), q.Bindings()...)
		if err != nil {
			return err
		}
	}

	err := s.tx.Commit()
	s.tx = nil
	return err
}

// Find returns a row given model properties
func (s *Session) Find(model interface{}) *Session {

	tName := s.mapper.ModelName(model)
	rModelMap := s.mapper.ToRawMap(model)

	sqlColNames := []string{}
	for k := range rModelMap {
		sqlColNames = append(sqlColNames, s.mapper.ColName(k))
	}

	sort.Strings(sqlColNames)

	s.builder = NewBuilder(s.metadata.Engine().Driver())
	s.builder.Select(sqlColNames...).From(tName)

	modelMap := s.mapper.ToMap(model)

	ands := []string{}
	bindings := []interface{}{}

	for k, v := range modelMap {
		ands = append(ands, fmt.Sprintf("%s = %s", s.mapper.ColName(k), s.builder.Dialect().Placeholder()))
		bindings = append(bindings, v)
	}

	s.builder.Where(s.builder.And(ands...), bindings...)
	return s
}

// First returns the first record mapped as a model
// The interface should be struct pointer instead of struct
func (s *Session) First(model interface{}) error {
	query := s.builder.Query()
	return s.metadata.Engine().Get(query, model)
}

// All returns all the records mapped as a model slice
// The interface should be struct pointer instead of struct
func (s *Session) All(models interface{}) error {
	query := s.builder.Query()
	return s.metadata.Engine().Select(query, models)
}

// builder overrides for session

// Update generates "update %s" statement
func (s *Session) Update(table string) *Session {
	s.builder.Update(table)
	return s
}

// Set generates "set a = placeholder" statement for each key a and add bindings for map value
func (s *Session) Set(m map[string]interface{}) *Session {
	s.builder.Set(m)
	return s
}

// Select generates "select %s" statement
func (s *Session) Select(columns ...string) *Session {
	s.builder.Select(columns...)
	return s
}

// From generates "from %s" statement for each table name
func (s *Session) From(tables ...string) *Session {
	s.builder.From(tables...)
	return s
}

// InnerJoin generates "inner join %s on %s" statement for each expression
func (s *Session) InnerJoin(table string, expressions ...string) *Session {
	s.builder.InnerJoin(table, expressions...)
	return s
}

// CrossJoin generates "cross join %s" statement for table
func (s *Session) CrossJoin(table string) *Session {
	s.builder.CrossJoin(table)
	return s
}

// LeftOuterJoin generates "left outer join %s on %s" statement for each expression
func (s *Session) LeftOuterJoin(table string, expressions ...string) *Session {
	s.builder.LeftOuterJoin(table, expressions...)
	return s
}

// RightOuterJoin generates "right outer join %s on %s" statement for each expression
func (s *Session) RightOuterJoin(table string, expressions ...string) *Session {
	s.builder.RightOuterJoin(table, expressions...)
	return s
}

// FullOuterJoin generates "full outer join %s on %s" for each expression
func (s *Session) FullOuterJoin(table string, expressions ...string) *Session {
	s.builder.FullOuterJoin(table, expressions...)
	return s
}

// Where generates "where %s" for the expression and adds bindings for each value
func (s *Session) Where(expression string, bindings ...interface{}) *Session {
	expression = strings.Replace(expression, "?", s.builder.Dialect().Placeholder(), -1)
	s.builder.Where(expression, bindings...)
	return s
}

// OrderBy generates "order by %s" for each expression
func (s *Session) OrderBy(expressions ...string) *Session {
	s.builder.OrderBy(expressions...)
	return s
}

// GroupBy generates "group by %s" for each column
func (s *Session) GroupBy(columns ...string) *Session {
	s.builder.GroupBy(columns...)
	return s
}

// Having generates "having %s" for each expression
func (s *Session) Having(expressions ...string) *Session {
	s.builder.Having(expressions...)
	return s
}

// Limit generates limit %d offset %d for offset and count
func (s *Session) Limit(offset int, count int) *Session {
	s.builder.Limit(offset, count)
	return s
}

// aggregates

// Avg function generates "avg(%s)" statement for column
func (s *Session) Avg(column string) string {
	return s.builder.Avg(column)
}

// Count function generates "count(%s)" statement for column
func (s *Session) Count(column string) string {
	return s.builder.Count(column)
}

// Sum function generates "sum(%s)" statement for column
func (s *Session) Sum(column string) string {
	return s.builder.Sum(column)
}

// Min function generates "min(%s)" statement for column
func (s *Session) Min(column string) string {
	return s.builder.Min(column)
}

// Max function generates "max(%s)" statement for column
func (s *Session) Max(column string) string {
	return s.builder.Max(column)
}

// expressions

// NotIn function generates "%s not in (%s)" for key and adds bindings for each value
func (s *Session) NotIn(key string, values ...interface{}) string {
	return s.builder.NotIn(key, values...)
}

// In function generates "%s in (%s)" for key and adds bindings for each value
func (s *Session) In(key string, values ...interface{}) string {
	return s.builder.In(key, values...)
}

// NotEq function generates "%s != placeholder" for key and adds binding for value
func (s *Session) NotEq(key string, value interface{}) string {
	return s.builder.NotEq(key, value)
}

// Eq function generates "%s = placeholder" for key and adds binding for value
func (s *Session) Eq(key string, value interface{}) string {
	return s.builder.Eq(key, value)
}

// Gt function generates "%s > placeholder" for key and adds binding for value
func (s *Session) Gt(key string, value interface{}) string {
	return s.builder.Gt(key, value)
}

// Gte function generates "%s >= placeholder" for key and adds binding for value
func (s *Session) Gte(key string, value interface{}) string {
	return s.builder.Gte(key, value)
}

// St function generates "%s < placeholder" for key and adds binding for value
func (s *Session) St(key string, value interface{}) string {
	return s.builder.St(key, value)
}

// Ste function generates "%s <= placeholder" for key and adds binding for value
func (s *Session) Ste(key string, value interface{}) string {
	return s.builder.Ste(key, value)
}

// And function generates " AND " between any number of expressions
func (s *Session) And(expressions ...string) string {
	return s.builder.And(expressions...)
}

// Or function generates " OR " between any number of expressions
func (s *Session) Or(expressions ...string) string {
	return s.builder.Or(expressions...)
}
<commit_msg>escape column names in building select of Find(), cleanup<commit_after>package qb

import (
	"database/sql"
	"fmt"
	"strings"
)

// New generates a new Session given engine and returns session pointer
func New(driver string, dsn string) (*Session, error) {

	engine, err := NewEngine(driver, dsn)
	if err != nil {
		return nil, err
	}

	return &Session{
		queries:  []*Query{},
		mapper:   NewMapper(engine.Driver()),
		metadata: NewMetaData(engine),
		builder:  NewBuilder(engine.Driver()),
	}, nil
}

// Session is the composition of engine connection & orm mappings
type Session struct {
	queries  []*Query
	mapper   *Mapper
	metadata *MetaData
	tx       *sql.Tx
	builder  *Builder
}

func (s *Session) add(query *Query) {
	var err error
	if s.tx == nil {
		s.queries = []*Query{}
		s.tx, err = s.metadata.Engine().DB().Begin()
		if err != nil {
			panic(err)
		}
	}
	s.queries = append(s.queries, query)
}

// Close closes engine db (sqlx) connection
func (s *Session) Close() {
	s.metadata.Engine().DB().Close()
}

// Builder returns query builder
func (s *Session) Builder() *Builder{
	return s.builder
}

// AddQuery adds a query given the query pointer retrieved from Query() function
func (s *Session) AddQuery(query *Query) {
	s.add(query)
}

// Query returns the active query built by session
func (s *Session) Query() *Query {
	return s.builder.Query()
}

// Metadata returns the metadata of session
func (s *Session) Metadata() *MetaData {
	return s.metadata
}

// Delete adds a single delete query to the session
func (s *Session) Delete(model interface{}) {

	kv := s.mapper.ToMap(model)

	tName := s.mapper.ModelName(model)

	d := s.metadata.Table(tName).Delete()
	ands := []string{}
	bindings := []interface{}{}
	for k, v := range kv {
		ands = append(ands, fmt.Sprintf("%s = %s", s.mapper.ColName(k), s.builder.Dialect().Placeholder()))
		bindings = append(bindings, v)
	}

	del := d.Where(d.And(ands...), bindings...).Query()
	s.add(del)
}

// Add adds a single model to the session. The query must be insert or update
func (s *Session) Add(model interface{}) {

	rawMap := s.mapper.ToMap(model)

	kv := map[string]interface{}{}

	for k, v := range rawMap {
		kv[s.mapper.ColName(k)] = v
	}

	q := s.metadata.Table(s.mapper.ModelName(model)).Insert(kv).Query()
	s.add(q)
}

// AddAll adds multiple models an adds an insert statement to current queries
func (s *Session) AddAll(models ...interface{}) {
	for _, m := range models {
		s.Add(m)
	}
}

// Commit commits the current transaction with queries
func (s *Session) Commit() error {

	for _, q := range s.queries {
		_, err := s.tx.Exec(q.SQL(), q.Bindings()...)
		if err != nil {
			return err
		}
	}

	err := s.tx.Commit()
	s.tx = nil
	return err
}

// Find returns a row given model properties
func (s *Session) Find(model interface{}) *Session {

	tName := s.mapper.ModelName(model)
	rModelMap := s.mapper.ToRawMap(model)

	sqlColNames := []string{}
	for k := range rModelMap {
		sqlColNames = append(sqlColNames, s.mapper.ColName(k))
	}

	//sort.Strings(sqlColNames)

	s.builder = NewBuilder(s.metadata.Engine().Driver())
	s.builder.Select(s.builder.Dialect().EscapeAll(sqlColNames)...).From(tName)

	modelMap := s.mapper.ToMap(model)

	ands := []string{}
	bindings := []interface{}{}

	for k, v := range modelMap {
		ands = append(ands, fmt.Sprintf("%s = %s", s.mapper.ColName(k), s.builder.Dialect().Placeholder()))
		bindings = append(bindings, v)
	}

	s.builder.Where(s.builder.And(ands...), bindings...)
	return s
}

// First returns the first record mapped as a model
// The interface should be struct pointer instead of struct
func (s *Session) First(model interface{}) error {
	query := s.builder.Query()
	return s.metadata.Engine().Get(query, model)
}

// All returns all the records mapped as a model slice
// The interface should be struct pointer instead of struct
func (s *Session) All(models interface{}) error {
	query := s.builder.Query()
	return s.metadata.Engine().Select(query, models)
}

// builder overrides for session

// Update generates "update %s" statement
func (s *Session) Update(table string) *Session {
	s.builder.Update(table)
	return s
}

// Set generates "set a = placeholder" statement for each key a and add bindings for map value
func (s *Session) Set(m map[string]interface{}) *Session {
	s.builder.Set(m)
	return s
}

// Select generates "select %s" statement
func (s *Session) Select(columns ...string) *Session {
	s.builder.Select(columns...)
	return s
}

// From generates "from %s" statement for each table name
func (s *Session) From(tables ...string) *Session {
	s.builder.From(tables...)
	return s
}

// InnerJoin generates "inner join %s on %s" statement for each expression
func (s *Session) InnerJoin(table string, expressions ...string) *Session {
	s.builder.InnerJoin(table, expressions...)
	return s
}

// CrossJoin generates "cross join %s" statement for table
func (s *Session) CrossJoin(table string) *Session {
	s.builder.CrossJoin(table)
	return s
}

// LeftOuterJoin generates "left outer join %s on %s" statement for each expression
func (s *Session) LeftOuterJoin(table string, expressions ...string) *Session {
	s.builder.LeftOuterJoin(table, expressions...)
	return s
}

// RightOuterJoin generates "right outer join %s on %s" statement for each expression
func (s *Session) RightOuterJoin(table string, expressions ...string) *Session {
	s.builder.RightOuterJoin(table, expressions...)
	return s
}

// FullOuterJoin generates "full outer join %s on %s" for each expression
func (s *Session) FullOuterJoin(table string, expressions ...string) *Session {
	s.builder.FullOuterJoin(table, expressions...)
	return s
}

// Where generates "where %s" for the expression and adds bindings for each value
func (s *Session) Where(expression string, bindings ...interface{}) *Session {
	expression = strings.Replace(expression, "?", s.builder.Dialect().Placeholder(), -1)
	s.builder.Where(expression, bindings...)
	return s
}

// OrderBy generates "order by %s" for each expression
func (s *Session) OrderBy(expressions ...string) *Session {
	s.builder.OrderBy(expressions...)
	return s
}

// GroupBy generates "group by %s" for each column
func (s *Session) GroupBy(columns ...string) *Session {
	s.builder.GroupBy(columns...)
	return s
}

// Having generates "having %s" for each expression
func (s *Session) Having(expressions ...string) *Session {
	s.builder.Having(expressions...)
	return s
}

// Limit generates limit %d offset %d for offset and count
func (s *Session) Limit(offset int, count int) *Session {
	s.builder.Limit(offset, count)
	return s
}

// aggregates

// Avg function generates "avg(%s)" statement for column
func (s *Session) Avg(column string) string {
	return s.builder.Avg(column)
}

// Count function generates "count(%s)" statement for column
func (s *Session) Count(column string) string {
	return s.builder.Count(column)
}

// Sum function generates "sum(%s)" statement for column
func (s *Session) Sum(column string) string {
	return s.builder.Sum(column)
}

// Min function generates "min(%s)" statement for column
func (s *Session) Min(column string) string {
	return s.builder.Min(column)
}

// Max function generates "max(%s)" statement for column
func (s *Session) Max(column string) string {
	return s.builder.Max(column)
}

// expressions

// NotIn function generates "%s not in (%s)" for key and adds bindings for each value
func (s *Session) NotIn(key string, values ...interface{}) string {
	return s.builder.NotIn(key, values...)
}

// In function generates "%s in (%s)" for key and adds bindings for each value
func (s *Session) In(key string, values ...interface{}) string {
	return s.builder.In(key, values...)
}

// NotEq function generates "%s != placeholder" for key and adds binding for value
func (s *Session) NotEq(key string, value interface{}) string {
	return s.builder.NotEq(key, value)
}

// Eq function generates "%s = placeholder" for key and adds binding for value
func (s *Session) Eq(key string, value interface{}) string {
	return s.builder.Eq(key, value)
}

// Gt function generates "%s > placeholder" for key and adds binding for value
func (s *Session) Gt(key string, value interface{}) string {
	return s.builder.Gt(key, value)
}

// Gte function generates "%s >= placeholder" for key and adds binding for value
func (s *Session) Gte(key string, value interface{}) string {
	return s.builder.Gte(key, value)
}

// St function generates "%s < placeholder" for key and adds binding for value
func (s *Session) St(key string, value interface{}) string {
	return s.builder.St(key, value)
}

// Ste function generates "%s <= placeholder" for key and adds binding for value
func (s *Session) Ste(key string, value interface{}) string {
	return s.builder.Ste(key, value)
}

// And function generates " AND " between any number of expressions
func (s *Session) And(expressions ...string) string {
	return s.builder.And(expressions...)
}

// Or function generates " OR " between any number of expressions
func (s *Session) Or(expressions ...string) string {
	return s.builder.Or(expressions...)
}
<|endoftext|>


<commit_before>package smux

import (
	"encoding/binary"
	"io"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pkg/errors"
)

const (
	defaultAcceptBacklog = 1024
)

const (
	errBrokenPipe      = "broken pipe"
	errInvalidProtocol = "invalid protocol version"
	errGoAway          = "stream id overflows, should start a new connection"
)

type writeRequest struct {
	frame  Frame
	result chan writeResult
}

type writeResult struct {
	n   int
	err error
}

type buffersWriter interface {
	WriteBuffers(v [][]byte) (n int, err error)
}

// Session defines a multiplexed connection for streams
type Session struct {
	conn io.ReadWriteCloser

	config           *Config
	nextStreamID     uint32 // next stream identifier
	nextStreamIDLock sync.Mutex

	bucket       int32         // token bucket
	bucketNotify chan struct{} // used for waiting for tokens

	streams    map[uint32]*Stream // all streams in this session
	streamLock sync.Mutex         // locks streams

	die       chan struct{} // flag session has died
	dieLock   sync.Mutex
	chAccepts chan *Stream

	dataReady int32 // flag data has arrived

	goAway int32 // flag id exhausted

	deadline atomic.Value

	writes chan writeRequest
}

func newSession(config *Config, conn io.ReadWriteCloser, client bool) *Session {
	s := new(Session)
	s.die = make(chan struct{})
	s.conn = conn
	s.config = config
	s.streams = make(map[uint32]*Stream)
	s.chAccepts = make(chan *Stream, defaultAcceptBacklog)
	s.bucket = int32(config.MaxReceiveBuffer)
	s.bucketNotify = make(chan struct{}, 1)
	s.writes = make(chan writeRequest)

	if client {
		s.nextStreamID = 1
	} else {
		s.nextStreamID = 0
	}
	go s.recvLoop()
	go s.sendLoop()
	go s.keepalive()
	return s
}

// OpenStream is used to create a new stream
func (s *Session) OpenStream() (*Stream, error) {
	if s.IsClosed() {
		return nil, errors.New(errBrokenPipe)
	}

	// generate stream id
	s.nextStreamIDLock.Lock()
	if s.goAway > 0 {
		s.nextStreamIDLock.Unlock()
		return nil, errors.New(errGoAway)
	}

	s.nextStreamID += 2
	sid := s.nextStreamID
	if sid == sid%2 { // stream-id overflows
		s.goAway = 1
		s.nextStreamIDLock.Unlock()
		return nil, errors.New(errGoAway)
	}
	s.nextStreamIDLock.Unlock()

	stream := newStream(sid, s.config.MaxFrameSize, s)

	if _, err := s.writeFrame(newFrame(cmdSYN, sid)); err != nil {
		return nil, errors.Wrap(err, "writeFrame")
	}

	s.streamLock.Lock()
	s.streams[sid] = stream
	s.streamLock.Unlock()
	return stream, nil
}

// AcceptStream is used to block until the next available stream
// is ready to be accepted.
func (s *Session) AcceptStream() (*Stream, error) {
	var deadline <-chan time.Time
	if d, ok := s.deadline.Load().(time.Time); ok && !d.IsZero() {
		timer := time.NewTimer(time.Until(d))
		defer timer.Stop()
		deadline = timer.C
	}
	select {
	case stream := <-s.chAccepts:
		return stream, nil
	case <-deadline:
		return nil, errTimeout
	case <-s.die:
		return nil, errors.New(errBrokenPipe)
	}
}

// Close is used to close the session and all streams.
func (s *Session) Close() (err error) {
	s.dieLock.Lock()

	select {
	case <-s.die:
		s.dieLock.Unlock()
		return errors.New(errBrokenPipe)
	default:
		close(s.die)
		s.dieLock.Unlock()
		s.streamLock.Lock()
		for k := range s.streams {
			s.streams[k].sessionClose()
		}
		s.streamLock.Unlock()
		s.notifyBucket()
		return s.conn.Close()
	}
}

// notifyBucket notifies recvLoop that bucket is available
func (s *Session) notifyBucket() {
	select {
	case s.bucketNotify <- struct{}{}:
	default:
	}
}

// IsClosed does a safe check to see if we have shutdown
func (s *Session) IsClosed() bool {
	select {
	case <-s.die:
		return true
	default:
		return false
	}
}

// NumStreams returns the number of currently open streams
func (s *Session) NumStreams() int {
	if s.IsClosed() {
		return 0
	}
	s.streamLock.Lock()
	defer s.streamLock.Unlock()
	return len(s.streams)
}

// SetDeadline sets a deadline used by Accept* calls.
// A zero time value disables the deadline.
func (s *Session) SetDeadline(t time.Time) error {
	s.deadline.Store(t)
	return nil
}

// notify the session that a stream has closed
func (s *Session) streamClosed(sid uint32) {
	s.streamLock.Lock()
	if n := s.streams[sid].recycleTokens(); n > 0 { // return remaining tokens to the bucket
		if atomic.AddInt32(&s.bucket, int32(n)) > 0 {
			s.notifyBucket()
		}
	}
	delete(s.streams, sid)
	s.streamLock.Unlock()
}

// returnTokens is called by stream to return token after read
func (s *Session) returnTokens(n int) {
	if atomic.AddInt32(&s.bucket, int32(n)) > 0 {
		s.notifyBucket()
	}
}

// recvLoop keeps on reading from underlying connection if tokens are available
func (s *Session) recvLoop() {
	var hdr rawHeader

	for {
		for atomic.LoadInt32(&s.bucket) <= 0 && !s.IsClosed() {
			<-s.bucketNotify
		}

		// read header first
		if _, err := io.ReadFull(s.conn, hdr[:]); err == nil {
			if hdr.Version() != version { // just ignore
				continue
			}
			atomic.StoreInt32(&s.dataReady, 1)
			sid := hdr.StreamID()
			switch hdr.Cmd() {
			case cmdNOP:
			case cmdSYN:
				s.streamLock.Lock()
				if _, ok := s.streams[sid]; !ok {
					stream := newStream(sid, s.config.MaxFrameSize, s)
					s.streams[sid] = stream
					select {
					case s.chAccepts <- stream:
					case <-s.die:
					}
				}
				s.streamLock.Unlock()
			case cmdFIN:
				s.streamLock.Lock()
				if stream, ok := s.streams[sid]; ok {
					stream.markRST()
					stream.notifyReadEvent()
				}
				s.streamLock.Unlock()
			case cmdPSH:
				s.streamLock.Lock()
				stream := s.streams[sid]
				s.streamLock.Unlock()

				if stream != nil {
					written, err := stream.receiveBytes(s.conn, int64(hdr.Length()))
					atomic.AddInt32(&s.bucket, -int32(written))
					stream.notifyReadEvent()
					if err != nil {
						s.Close()
						return
					}
				}
			default:
				s.Close()
				return
			}
		} else {
			s.Close()
			return
		}
	}
}

func (s *Session) keepalive() {
	tickerPing := time.NewTicker(s.config.KeepAliveInterval)
	tickerTimeout := time.NewTicker(s.config.KeepAliveTimeout)
	defer tickerPing.Stop()
	defer tickerTimeout.Stop()
	for {
		select {
		case <-tickerPing.C:
			s.writeFrameInternal(newFrame(cmdNOP, 0), tickerPing.C)
			s.notifyBucket() // force a signal to the recvLoop
		case <-tickerTimeout.C:
			if !atomic.CompareAndSwapInt32(&s.dataReady, 1, 0) {
				s.Close()
				return
			}
		case <-s.die:
			return
		}
	}
}

func (s *Session) sendLoop() {
	buf := make([]byte, (1<<16)+headerSize)
	var n int
	var err error
	v := make([][]byte, 2) // vector for writing

	for {
		select {
		case <-s.die:
			return
		case request := <-s.writes:
			buf[0] = request.frame.ver
			buf[1] = request.frame.cmd
			binary.LittleEndian.PutUint16(buf[2:], uint16(len(request.frame.data)))
			binary.LittleEndian.PutUint32(buf[4:], request.frame.sid)

			if bw, ok := s.conn.(buffersWriter); ok {
				v[0] = buf[:headerSize]
				v[1] = request.frame.data
				n, err = bw.WriteBuffers(v)
			} else {
				copy(buf[headerSize:], request.frame.data)
				n, err = s.conn.Write(buf[:headerSize+len(request.frame.data)])
			}

			n -= headerSize
			if n < 0 {
				n = 0
			}

			result := writeResult{
				n:   n,
				err: err,
			}

			request.result <- result
			close(request.result)
		}
	}
}

// writeFrame writes the frame to the underlying connection
// and returns the number of bytes written if successful
func (s *Session) writeFrame(f Frame) (n int, err error) {
	return s.writeFrameInternal(f, nil)
}

// internal writeFrame version to support deadline used in keepalive
func (s *Session) writeFrameInternal(f Frame, deadline <-chan time.Time) (int, error) {
	req := writeRequest{
		frame:  f,
		result: make(chan writeResult, 1),
	}
	select {
	case <-s.die:
		return 0, errors.New(errBrokenPipe)
	case s.writes <- req:
	case <-deadline:
		return 0, errTimeout
	}

	select {
	case result := <-req.result:
		return result.n, result.err
	case <-deadline:
		return 0, errTimeout
	case <-s.die:
		return 0, errors.New(errBrokenPipe)
	}
}
<commit_msg>Revert "make sure streamLock is non-blocking due to receiveBytes"<commit_after>package smux

import (
	"encoding/binary"
	"io"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pkg/errors"
)

const (
	defaultAcceptBacklog = 1024
)

const (
	errBrokenPipe      = "broken pipe"
	errInvalidProtocol = "invalid protocol version"
	errGoAway          = "stream id overflows, should start a new connection"
)

type writeRequest struct {
	frame  Frame
	result chan writeResult
}

type writeResult struct {
	n   int
	err error
}

type buffersWriter interface {
	WriteBuffers(v [][]byte) (n int, err error)
}

// Session defines a multiplexed connection for streams
type Session struct {
	conn io.ReadWriteCloser

	config           *Config
	nextStreamID     uint32 // next stream identifier
	nextStreamIDLock sync.Mutex

	bucket       int32         // token bucket
	bucketNotify chan struct{} // used for waiting for tokens

	streams    map[uint32]*Stream // all streams in this session
	streamLock sync.Mutex         // locks streams

	die       chan struct{} // flag session has died
	dieLock   sync.Mutex
	chAccepts chan *Stream

	dataReady int32 // flag data has arrived

	goAway int32 // flag id exhausted

	deadline atomic.Value

	writes chan writeRequest
}

func newSession(config *Config, conn io.ReadWriteCloser, client bool) *Session {
	s := new(Session)
	s.die = make(chan struct{})
	s.conn = conn
	s.config = config
	s.streams = make(map[uint32]*Stream)
	s.chAccepts = make(chan *Stream, defaultAcceptBacklog)
	s.bucket = int32(config.MaxReceiveBuffer)
	s.bucketNotify = make(chan struct{}, 1)
	s.writes = make(chan writeRequest)

	if client {
		s.nextStreamID = 1
	} else {
		s.nextStreamID = 0
	}
	go s.recvLoop()
	go s.sendLoop()
	go s.keepalive()
	return s
}

// OpenStream is used to create a new stream
func (s *Session) OpenStream() (*Stream, error) {
	if s.IsClosed() {
		return nil, errors.New(errBrokenPipe)
	}

	// generate stream id
	s.nextStreamIDLock.Lock()
	if s.goAway > 0 {
		s.nextStreamIDLock.Unlock()
		return nil, errors.New(errGoAway)
	}

	s.nextStreamID += 2
	sid := s.nextStreamID
	if sid == sid%2 { // stream-id overflows
		s.goAway = 1
		s.nextStreamIDLock.Unlock()
		return nil, errors.New(errGoAway)
	}
	s.nextStreamIDLock.Unlock()

	stream := newStream(sid, s.config.MaxFrameSize, s)

	if _, err := s.writeFrame(newFrame(cmdSYN, sid)); err != nil {
		return nil, errors.Wrap(err, "writeFrame")
	}

	s.streamLock.Lock()
	s.streams[sid] = stream
	s.streamLock.Unlock()
	return stream, nil
}

// AcceptStream is used to block until the next available stream
// is ready to be accepted.
func (s *Session) AcceptStream() (*Stream, error) {
	var deadline <-chan time.Time
	if d, ok := s.deadline.Load().(time.Time); ok && !d.IsZero() {
		timer := time.NewTimer(time.Until(d))
		defer timer.Stop()
		deadline = timer.C
	}
	select {
	case stream := <-s.chAccepts:
		return stream, nil
	case <-deadline:
		return nil, errTimeout
	case <-s.die:
		return nil, errors.New(errBrokenPipe)
	}
}

// Close is used to close the session and all streams.
func (s *Session) Close() (err error) {
	s.dieLock.Lock()

	select {
	case <-s.die:
		s.dieLock.Unlock()
		return errors.New(errBrokenPipe)
	default:
		close(s.die)
		s.dieLock.Unlock()
		s.streamLock.Lock()
		for k := range s.streams {
			s.streams[k].sessionClose()
		}
		s.streamLock.Unlock()
		s.notifyBucket()
		return s.conn.Close()
	}
}

// notifyBucket notifies recvLoop that bucket is available
func (s *Session) notifyBucket() {
	select {
	case s.bucketNotify <- struct{}{}:
	default:
	}
}

// IsClosed does a safe check to see if we have shutdown
func (s *Session) IsClosed() bool {
	select {
	case <-s.die:
		return true
	default:
		return false
	}
}

// NumStreams returns the number of currently open streams
func (s *Session) NumStreams() int {
	if s.IsClosed() {
		return 0
	}
	s.streamLock.Lock()
	defer s.streamLock.Unlock()
	return len(s.streams)
}

// SetDeadline sets a deadline used by Accept* calls.
// A zero time value disables the deadline.
func (s *Session) SetDeadline(t time.Time) error {
	s.deadline.Store(t)
	return nil
}

// notify the session that a stream has closed
func (s *Session) streamClosed(sid uint32) {
	s.streamLock.Lock()
	if n := s.streams[sid].recycleTokens(); n > 0 { // return remaining tokens to the bucket
		if atomic.AddInt32(&s.bucket, int32(n)) > 0 {
			s.notifyBucket()
		}
	}
	delete(s.streams, sid)
	s.streamLock.Unlock()
}

// returnTokens is called by stream to return token after read
func (s *Session) returnTokens(n int) {
	if atomic.AddInt32(&s.bucket, int32(n)) > 0 {
		s.notifyBucket()
	}
}

// recvLoop keeps on reading from underlying connection if tokens are available
func (s *Session) recvLoop() {
	var hdr rawHeader

	for {
		for atomic.LoadInt32(&s.bucket) <= 0 && !s.IsClosed() {
			<-s.bucketNotify
		}

		// read header first
		if _, err := io.ReadFull(s.conn, hdr[:]); err == nil {
			if hdr.Version() != version { // just ignore
				continue
			}
			atomic.StoreInt32(&s.dataReady, 1)
			sid := hdr.StreamID()
			switch hdr.Cmd() {
			case cmdNOP:
			case cmdSYN:
				s.streamLock.Lock()
				if _, ok := s.streams[sid]; !ok {
					stream := newStream(sid, s.config.MaxFrameSize, s)
					s.streams[sid] = stream
					select {
					case s.chAccepts <- stream:
					case <-s.die:
					}
				}
				s.streamLock.Unlock()
			case cmdFIN:
				s.streamLock.Lock()
				if stream, ok := s.streams[sid]; ok {
					stream.markRST()
					stream.notifyReadEvent()
				}
				s.streamLock.Unlock()
			case cmdPSH:
				var written int64
				var err error
				s.streamLock.Lock()
				if stream, ok := s.streams[sid]; ok {
					written, err = stream.receiveBytes(s.conn, int64(hdr.Length()))
					atomic.AddInt32(&s.bucket, -int32(written))
					stream.notifyReadEvent()
				}
				s.streamLock.Unlock()

				// read data error
				if err != nil {
					s.Close()
					return
				}
			default:
				s.Close()
				return
			}
		} else {
			s.Close()
			return
		}
	}
}

func (s *Session) keepalive() {
	tickerPing := time.NewTicker(s.config.KeepAliveInterval)
	tickerTimeout := time.NewTicker(s.config.KeepAliveTimeout)
	defer tickerPing.Stop()
	defer tickerTimeout.Stop()
	for {
		select {
		case <-tickerPing.C:
			s.writeFrameInternal(newFrame(cmdNOP, 0), tickerPing.C)
			s.notifyBucket() // force a signal to the recvLoop
		case <-tickerTimeout.C:
			if !atomic.CompareAndSwapInt32(&s.dataReady, 1, 0) {
				s.Close()
				return
			}
		case <-s.die:
			return
		}
	}
}

func (s *Session) sendLoop() {
	buf := make([]byte, (1<<16)+headerSize)
	var n int
	var err error
	v := make([][]byte, 2) // vector for writing

	for {
		select {
		case <-s.die:
			return
		case request := <-s.writes:
			buf[0] = request.frame.ver
			buf[1] = request.frame.cmd
			binary.LittleEndian.PutUint16(buf[2:], uint16(len(request.frame.data)))
			binary.LittleEndian.PutUint32(buf[4:], request.frame.sid)

			if bw, ok := s.conn.(buffersWriter); ok {
				v[0] = buf[:headerSize]
				v[1] = request.frame.data
				n, err = bw.WriteBuffers(v)
			} else {
				copy(buf[headerSize:], request.frame.data)
				n, err = s.conn.Write(buf[:headerSize+len(request.frame.data)])
			}

			n -= headerSize
			if n < 0 {
				n = 0
			}

			result := writeResult{
				n:   n,
				err: err,
			}

			request.result <- result
			close(request.result)
		}
	}
}

// writeFrame writes the frame to the underlying connection
// and returns the number of bytes written if successful
func (s *Session) writeFrame(f Frame) (n int, err error) {
	return s.writeFrameInternal(f, nil)
}

// internal writeFrame version to support deadline used in keepalive
func (s *Session) writeFrameInternal(f Frame, deadline <-chan time.Time) (int, error) {
	req := writeRequest{
		frame:  f,
		result: make(chan writeResult, 1),
	}
	select {
	case <-s.die:
		return 0, errors.New(errBrokenPipe)
	case s.writes <- req:
	case <-deadline:
		return 0, errTimeout
	}

	select {
	case result := <-req.result:
		return result.n, result.err
	case <-deadline:
		return 0, errTimeout
	case <-s.die:
		return 0, errors.New(errBrokenPipe)
	}
}
<|endoftext|>


<commit_before>/*
Package secsurf wraps https handlers setting various security related headers.

Please read the documentation carefully as the headers have long term security implications.

The following headers are set:
	X-XSS-Protection: 1; mode=block
	X-Frame-Options: deny
	X-Content-Type-Options: nosniff
	Strict-Transport-Security: max-age=31536000; includeSubDomains

X-XSS-Protection

Force XSS protection in some versions of IE.

X-Frame-Options

The pages cannot be used inside frames.

X-Content-Type-Options

Content type is not guessed by the browser.

Strict-Transport-Security

After the page has been succesfully loaded with a valid certificate chain the browser will REFUSE to load the page in the future without a valid https (tls) connection. Strict-Transport-Security is only set on https responses.

*/
package secsurf

import (
	"net/http"
)

// Wrap a HTTP handler.
func New(h http.Handler) http.Handler {
	return wrap{h}
}

type wrap struct {
	http.Handler
}

func (h wrap)ServeHTTP(w http.ResponseWriter, r *http.Request) {
	hdrs := w.Header()
	hdrs.Set(`X-XSS-Protection`, `1; mode=block`)
	hdrs.Set(`X-Frame-Options`, `deny`)
	hdrs.Set(`X-Content-Type-Options`, `nosniff`)
	if r.TLS != nil {
		hdrs.Set(`Strict-Transport-Security`, `max-age=31536000; includeSubDomains`)
	}
	h.Handler.ServeHTTP(w, r)
}

// Wrap a HTTP handler, adds a STS header even on HTTP.
func NewAlwaysSTS(h http.Handler) http.Handler {
	return swrap{h}
}

type swrap struct {
	http.Handler
}

func (h swrap)ServeHTTP(w http.ResponseWriter, r *http.Request) {
	hdrs := w.Header()
	hdrs.Set(`X-XSS-Protection`, `1; mode=block`)
	hdrs.Set(`X-Frame-Options`, `deny`)
	hdrs.Set(`X-Content-Type-Options`, `nosniff`)
	hdrs.Set(`Strict-Transport-Security`, `max-age=31536000; includeSubDomains`)
	h.Handler.ServeHTTP(w, r)
}
<commit_msg>gofmt<commit_after>/*
Package secsurf wraps https handlers setting various security related headers.

Please read the documentation carefully as the headers have long term security implications.

The following headers are set:
	X-XSS-Protection: 1; mode=block
	X-Frame-Options: deny
	X-Content-Type-Options: nosniff
	Strict-Transport-Security: max-age=31536000; includeSubDomains

X-XSS-Protection

Force XSS protection in some versions of IE.

X-Frame-Options

The pages cannot be used inside frames.

X-Content-Type-Options

Content type is not guessed by the browser.

Strict-Transport-Security

After the page has been succesfully loaded with a valid certificate chain the browser will REFUSE to load the page in the future without a valid https (tls) connection. Strict-Transport-Security is only set on https responses.

*/
package secsurf

import (
	"net/http"
)

// Wrap a HTTP handler.
func New(h http.Handler) http.Handler {
	return wrap{h}
}

type wrap struct {
	http.Handler
}

func (h wrap) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	hdrs := w.Header()
	hdrs.Set(`X-XSS-Protection`, `1; mode=block`)
	hdrs.Set(`X-Frame-Options`, `deny`)
	hdrs.Set(`X-Content-Type-Options`, `nosniff`)
	if r.TLS != nil {
		hdrs.Set(`Strict-Transport-Security`, `max-age=31536000; includeSubDomains`)
	}
	h.Handler.ServeHTTP(w, r)
}

// Wrap a HTTP handler, adds a STS header even on HTTP.
func NewAlwaysSTS(h http.Handler) http.Handler {
	return swrap{h}
}

type swrap struct {
	http.Handler
}

func (h swrap) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	hdrs := w.Header()
	hdrs.Set(`X-XSS-Protection`, `1; mode=block`)
	hdrs.Set(`X-Frame-Options`, `deny`)
	hdrs.Set(`X-Content-Type-Options`, `nosniff`)
	hdrs.Set(`Strict-Transport-Security`, `max-age=31536000; includeSubDomains`)
	h.Handler.ServeHTTP(w, r)
}
<|endoftext|>