<commit_before>package main

import (
	"bufio"
	"compress/bzip2"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/CSUNetSec/protoparse"
	mrt "github.com/CSUNetSec/protoparse/protocol/mrt"
	util "github.com/CSUNetSec/protoparse/util"
	radix "github.com/armon/go-radix"
	"io"
	"log"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"
)

//accepts an error (can be nil which turns the func into a noop
//and variadic arguemnts that are fds to be closed before calling
//os.exit in case the error is not nil
func errx(e error, fds ...io.Closer) {
	if e == nil {
		return
	}
	log.Printf("error: %s\n", e)
	for _, fd := range fds {
		fd.Close()
	}
	os.Exit(-1)
}

func getScanner(file *os.File) (scanner *bufio.Scanner) {
	fname := file.Name()
	fext := filepath.Ext(fname)
	if fext == ".bz2" {
		bzreader := bzip2.NewReader(file)
		scanner = bufio.NewScanner(bzreader)
	} else {
		scanner = bufio.NewScanner(file)
	}
	scanner.Split(mrt.SplitMrt)
	scanbuffer := make([]byte, 2<<24) //an internal buffer for the large tokens (1M)
	scanner.Buffer(scanbuffer, cap(scanbuffer))
	return
}

var (
	logout   string
	dumpout  string
	statout  string
	pup      bool
	isJson   bool
	parallel bool
	destAs   int
	srcAs    int
)

func init() {
	flag.StringVar(&logout, "lo", "stdout", "file to dump log output")
	flag.StringVar(&dumpout, "o", "stdout", "file to dump entries")
	flag.StringVar(&statout, "so", "stdout", "file to dump statistics output")
	flag.BoolVar(&parallel, "p", false, "dump files in parallel, may cause out of order output")
	flag.BoolVar(&pup, "pup", false, "print every advertized prefix only once")
	flag.BoolVar(&isJson, "json", false, "print the output as json objects")
	flag.IntVar(&destAs, "destAs", -1, "filter by this destination AS")
	flag.IntVar(&srcAs, "srcAs", -1, "filter by this source AS")
}

func main() {
	flag.Parse()

	args := flag.Args()
	if len(args) == 0 {
		log.Println("mrt file not provided")
		os.Exit(-1)
	}

	if logout != "stdout" {
		lfd, _ := os.Create(logout)
		log.SetOutput(lfd)
		defer lfd.Close()
	}
	var statfd *os.File
	var dumpfd *os.File

	if statout != "stdout" {
		statfd, _ = os.Create(statout)
		defer statfd.Close()
	} else {
		statfd = os.Stdout
	}

	if dumpout != "stdout" {
		dumpfd, _ = os.Create(dumpout)
		defer dumpfd.Close()
	} else {
		dumpfd = os.Stdout
	}

	statstr := fmt.SPrintf("Dumping %d files\n", len(args))
	statfd.WriteString(statstr)

	var tf transformer
	if isJson {
		tf = jsonTransformer{}
	} else if pup {
		upm := NewUniquePrefixMap()
		tf = upm
	} else {
		tf = textTransformer{}
	}
	var vals []validator
	if destAs != -1 {
		av := NewAsValidator(destAs)
		vals = append(vals, av.validateDest)
	}
	if srcAs != -1 {
		av := NewAsValidator(srcAs)
		vals = append(vals, av.validateSrc)
	}

	wg := &sync.WaitGroup{}
	start := time.Now()
	// Each goroutine requires an fd,
	// I should consider adding a struct to
	// manage the number of fds consumed by the program
	for _, name := range args {
		if parallel && false {
			wg.Add(1)
			go dumpFile(name, tf, vals, wg)
		} else {
			dumpFile(name, tf, vals, nil)
		}
	}

	if parallel {
		wg.Wait()
	}
	tf.summarize()
	log.Printf("Total time taken: %s\n", time.Since(start))
}

func dumpFile(fName string, tf transformer, vals []validator, dfd, sfd *os.File, wg *sync.WaitGroup) {
	if wg != nil {
		defer wg.Done()
	}
	mrtfd, err := os.Open(fName)
	errx(err)
	defer mrtfd.Close()

	mrtScanner := getScanner(mrtfd)
	numentries := 0
	unfilteredct := 0
	totsz := 0
	t1 := time.Now()

	for mrtScanner.Scan() {
		ret := ""
		numentries++
		data := mrtScanner.Bytes()
		totsz += len(data)
		mrth := mrt.NewMrtHdrBuf(data)
		bgp4h, bgph, bgpup, err := parseHeaders(mrth, numentries)
		if err != nil {
			log.Printf("[%d] Error:%s\n", numentries, err)
			break
		}
		mbs := &mrt.MrtBufferStack{mrth, bgp4h, bgph, bgpup}
		if validateAll(vals, mbs) {
			unfilteredct++
			ret += tf.transform(numentries, mbs)
			// I'm making transformers responsible for newlines, because
			// the upm doesn't need them
			dfd.WriteString(ret)
		}
	}

	if err := mrtScanner.Err(); err != nil {
		errx(err, mrtfd)
	}
	dt := time.Since(t1)
	statstr := fmt.SPrintf("Scanned %s: %d entries, %d passed filters, total size: %d bytes in %v\n", fName, numentries, unfilteredct, totsz, dt)
	sfd.WriteString(statstr)

}

func validateAll(vals []validator, mbs *mrt.MrtBufferStack) bool {
	for _, v := range vals {
		if v != nil && !v(mbs) {
			return false
		}
	}
	return true
}

// This could maybe not be a pointer
type validator func(*mrt.MrtBufferStack) bool

type AsValidator struct {
	as uint32
}

func NewAsValidator(i int) *AsValidator {
	return &AsValidator{uint32(i)}
}

// This checks the last AS in the ASPath,
// which should be where the update originated from
func (asval *AsValidator) validateSrc(mbs *mrt.MrtBufferStack) bool {
	update := mbs.Bgpupbuf.(protoparse.BGPUpdater).GetUpdate()
	if update == nil || update.Attrs == nil {
		//This happens a lot
		//log.Printf("Error retrieving AS Path\n")
		return false
	}
	pathlen := len(update.Attrs.AsPath)
	if pathlen < 1 {
		// This happens sometimes
		//log.Printf("Error: empty AS Path\n")
		return false
	}

	lastseg := update.Attrs.AsPath[pathlen-1]

	var lastAs uint32 = 0
	if lastseg.AsSeq != nil && len(lastseg.AsSeq) > 0 {
		lastAs = lastseg.AsSeq[len(lastseg.AsSeq)-1]
	} else if lastseg.AsSet != nil && len(lastseg.AsSet) > 0 {
		lastAs = lastseg.AsSet[len(lastseg.AsSet)-1]
	} else {
		// Both are empty, not sure how thats possible, but return false
		return false
	}
	return lastAs == asval.as
}

func (asval *AsValidator) validateDest(mbs *mrt.MrtBufferStack) bool {
	update := mbs.Bgpupbuf.(protoparse.BGPUpdater).GetUpdate()
	if update == nil || update.Attrs == nil {
		//This happens a lot
		//log.Printf("Error retrieving AS Path\n")
		return false
	}
	pathlen := len(update.Attrs.AsPath)
	if pathlen < 1 {
		// This happens sometimes
		//log.Printf("Error: empty AS Path\n")
		return false
	}

	lastseg := update.Attrs.AsPath[0]

	var firstAs uint32 = 0
	if lastseg.AsSeq != nil && len(lastseg.AsSeq) > 0 {
		firstAs = lastseg.AsSeq[0]
	} else if lastseg.AsSet != nil && len(lastseg.AsSet) > 0 {
		firstAs = lastseg.AsSet[0]
	} else {
		// Both are empty, not sure how thats possible, but return false
		return false
	}
	return firstAs == asval.as

}

type transformer interface {
	transform(int, *mrt.MrtBufferStack) string
	summarize()
}

type UniquePrefixMap struct {
	prefixes  map[string]bool
	radixTree *radix.Tree
	maplock   *sync.Mutex
}

func NewUniquePrefixMap() *UniquePrefixMap {
	upm := UniquePrefixMap{}
	upm.prefixes = make(map[string]bool)
	upm.radixTree = radix.New()
	upm.maplock = &sync.Mutex{}
	return &upm
}

type IPWrapper struct {
	ip string
}

func (upm *UniquePrefixMap) transform(msgNum int, mbs *mrt.MrtBufferStack) string {
	update := mbs.Bgpupbuf.(protoparse.BGPUpdater).GetUpdate()

	//If there are no advertized routes
	if update.AdvertizedRoutes == nil || len(update.AdvertizedRoutes.Prefixes) == 0 {
		return ""
	}

	for _, ar := range update.AdvertizedRoutes.Prefixes {
		ipstr := fmt.Sprintf("%s/%d", net.IP(util.GetIP(ar.GetPrefix())), ar.Mask)
		upm.maplock.Lock()
		if !upm.prefixes[ipstr] {
			rkey := util.IpToRadixkey(util.GetIP(ar.GetPrefix()), uint8(ar.Mask))
			upm.radixTree.Insert(rkey, IPWrapper{ipstr})
			upm.prefixes[ipstr] = true
		}
		upm.maplock.Unlock()
	}

	return ""

}

func (upm *UniquePrefixMap) summarize() {
	upm.radixTree.Walk(upm.topWalk)
}

func (upm *UniquePrefixMap) topWalk(s string, v interface{}) bool {
	pref := v.(IPWrapper).ip
	if upm.prefixes[pref] {
		fmt.Printf("%s\n", pref)
		upm.radixTree.WalkPrefix(s, upm.subWalk)
	}
	return false
}

func (upm *UniquePrefixMap) subWalk(s string, v interface{}) bool {
	upm.prefixes[v.(IPWrapper).ip] = false
	return false
}

type jsonTransformer struct{}

func (j jsonTransformer) transform(msgNum int, mbs *mrt.MrtBufferStack) string {
	mbsj, err := json.Marshal(mbs)
	if err != nil {
		log.Printf("Error marshaling to json")
		return ""
	}
	return string(mbsj) + "\n"
}

func (j jsonTransformer) summarize() {}

type textTransformer struct{}

func (t textTransformer) transform(msgNum int, mbs *mrt.MrtBufferStack) string {
	ret := ""
	ret += fmt.Sprintf("[%d] MRT Header: %s\n", msgNum, mbs.MrthBuf)
	ret += fmt.Sprintf("BGP4MP Header:%s\n", mbs.Bgp4mpbuf)
	ret += fmt.Sprintf("BGP Header: %s\n", mbs.Bgphbuf)
	ret += fmt.Sprintf("BGP Update:%s\n", mbs.Bgpupbuf)
	return ret + "\n"
}

func (t textTransformer) summarize() {}

func parseHeaders(mrth protoparse.PbVal, entryCt int) (bgp4h, bgph, bgpup protoparse.PbVal, err error) {
	bgp4h, err = mrth.Parse()
	if err != nil {
		log.Printf("Failed parsing MRT header %d :%s", entryCt, err)
		return
	}

	bgph, err = bgp4h.Parse()
	if err != nil {
		log.Printf("Failed parsing BG4MP header %d :%s", entryCt, err)
		return
	}

	bgpup, err = bgph.Parse()
	if err != nil {
		log.Printf("Failed parsing BGP header %d :%s", entryCt, err)
		return
	}

	_, err = bgpup.Parse()
	if err != nil {
		log.Printf("Failed parsing BGP update %d :%s", entryCt, err)
		return
	}

	return
}
<commit_msg>Added ability to redirect log, stat, and dump output<commit_after>package main

import (
	"bufio"
	"compress/bzip2"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/CSUNetSec/protoparse"
	mrt "github.com/CSUNetSec/protoparse/protocol/mrt"
	util "github.com/CSUNetSec/protoparse/util"
	radix "github.com/armon/go-radix"
	"io"
	"log"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"
)

//accepts an error (can be nil which turns the func into a noop
//and variadic arguemnts that are fds to be closed before calling
//os.exit in case the error is not nil
func errx(e error, fds ...io.Closer) {
	if e == nil {
		return
	}
	log.Printf("error: %s\n", e)
	for _, fd := range fds {
		fd.Close()
	}
	os.Exit(-1)
}

func getScanner(file *os.File) (scanner *bufio.Scanner) {
	fname := file.Name()
	fext := filepath.Ext(fname)
	if fext == ".bz2" {
		bzreader := bzip2.NewReader(file)
		scanner = bufio.NewScanner(bzreader)
	} else {
		scanner = bufio.NewScanner(file)
	}
	scanner.Split(mrt.SplitMrt)
	scanbuffer := make([]byte, 2<<24) //an internal buffer for the large tokens (1M)
	scanner.Buffer(scanbuffer, cap(scanbuffer))
	return
}

var (
	logout   string
	dumpout  string
	statout  string
	pup      bool
	isJson   bool
	parallel bool
	destAs   int
	srcAs    int
)

func init() {
	flag.StringVar(&logout, "lo", "stdout", "file to dump log output")
	flag.StringVar(&dumpout, "o", "stdout", "file to dump entries")
	flag.StringVar(&statout, "so", "stdout", "file to dump statistics output")
	flag.BoolVar(&parallel, "p", false, "dump files in parallel, may cause out of order output")
	flag.BoolVar(&pup, "pup", false, "print every advertized prefix only once")
	flag.BoolVar(&isJson, "json", false, "print the output as json objects")
	flag.IntVar(&destAs, "destAs", -1, "filter by this destination AS")
	flag.IntVar(&srcAs, "srcAs", -1, "filter by this source AS")
}

func main() {
	flag.Parse()

	args := flag.Args()
	if len(args) == 0 {
		log.Println("mrt file not provided")
		os.Exit(-1)
	}

	if logout != "stdout" {
		lfd, _ := os.Create(logout)
		log.SetOutput(lfd)
		defer lfd.Close()
	}
	var statfd *os.File
	var dumpfd *os.File

	if statout != "stdout" {
		statfd, _ = os.Create(statout)
		defer statfd.Close()
	} else {
		statfd = os.Stdout
	}

	if dumpout != "stdout" {
		dumpfd, _ = os.Create(dumpout)
		defer dumpfd.Close()
	} else {
		dumpfd = os.Stdout
	}

	statstr := fmt.Sprintf("Dumping %d files\n", len(args))
	statfd.WriteString(statstr)

	var tf transformer
	if isJson {
		tf = jsonTransformer{}
	} else if pup {
		upm := NewUniquePrefixMap(dumpfd)
		tf = upm
	} else {
		tf = textTransformer{}
	}
	var vals []validator
	if destAs != -1 {
		av := NewAsValidator(destAs)
		vals = append(vals, av.validateDest)
	}
	if srcAs != -1 {
		av := NewAsValidator(srcAs)
		vals = append(vals, av.validateSrc)
	}

	wg := &sync.WaitGroup{}
	start := time.Now()
	// Each goroutine requires an fd,
	// I should consider adding a struct to
	// manage the number of fds consumed by the program
	for _, name := range args {
		if parallel && false {
			wg.Add(1)
			go dumpFile(name, tf, vals, dumpfd, statfd, wg)
		} else {
			dumpFile(name, tf, vals, dumpfd, statfd, nil)
		}
	}

	if parallel {
		wg.Wait()
	}
	tf.summarize()
	str := fmt.Sprintf("Total time taken: %s\n", time.Since(start))
	statfd.WriteString(str)
}

func dumpFile(fName string, tf transformer, vals []validator, dfd, sfd *os.File, wg *sync.WaitGroup) {
	if wg != nil {
		defer wg.Done()
	}
	mrtfd, err := os.Open(fName)
	errx(err)
	defer mrtfd.Close()

	mrtScanner := getScanner(mrtfd)
	numentries := 0
	unfilteredct := 0
	totsz := 0
	t1 := time.Now()

	for mrtScanner.Scan() {
		ret := ""
		numentries++
		data := mrtScanner.Bytes()
		totsz += len(data)
		mrth := mrt.NewMrtHdrBuf(data)
		bgp4h, bgph, bgpup, err := parseHeaders(mrth, numentries)
		if err != nil {
			log.Printf("[%d] Error:%s\n", numentries, err)
			break
		}
		mbs := &mrt.MrtBufferStack{mrth, bgp4h, bgph, bgpup}
		if validateAll(vals, mbs) {
			unfilteredct++
			ret += tf.transform(numentries, mbs)
			// I'm making transformers responsible for newlines, because
			// the upm doesn't need them
			dfd.WriteString(ret)
		}
	}

	if err := mrtScanner.Err(); err != nil {
		errx(err, mrtfd)
	}
	dt := time.Since(t1)
	statstr := fmt.Sprintf("Scanned %s: %d entries, %d passed filters, total size: %d bytes in %v\n", fName, numentries, unfilteredct, totsz, dt)
	sfd.WriteString(statstr)

}

func validateAll(vals []validator, mbs *mrt.MrtBufferStack) bool {
	for _, v := range vals {
		if v != nil && !v(mbs) {
			return false
		}
	}
	return true
}

// This could maybe not be a pointer
type validator func(*mrt.MrtBufferStack) bool

type AsValidator struct {
	as uint32
}

func NewAsValidator(i int) *AsValidator {
	return &AsValidator{uint32(i)}
}

// This checks the last AS in the ASPath,
// which should be where the update originated from
func (asval *AsValidator) validateSrc(mbs *mrt.MrtBufferStack) bool {
	update := mbs.Bgpupbuf.(protoparse.BGPUpdater).GetUpdate()
	if update == nil || update.Attrs == nil {
		//This happens a lot
		//log.Printf("Error retrieving AS Path\n")
		return false
	}
	pathlen := len(update.Attrs.AsPath)
	if pathlen < 1 {
		// This happens sometimes
		//log.Printf("Error: empty AS Path\n")
		return false
	}

	lastseg := update.Attrs.AsPath[pathlen-1]

	var lastAs uint32 = 0
	if lastseg.AsSeq != nil && len(lastseg.AsSeq) > 0 {
		lastAs = lastseg.AsSeq[len(lastseg.AsSeq)-1]
	} else if lastseg.AsSet != nil && len(lastseg.AsSet) > 0 {
		lastAs = lastseg.AsSet[len(lastseg.AsSet)-1]
	} else {
		// Both are empty, not sure how thats possible, but return false
		return false
	}
	return lastAs == asval.as
}

func (asval *AsValidator) validateDest(mbs *mrt.MrtBufferStack) bool {
	update := mbs.Bgpupbuf.(protoparse.BGPUpdater).GetUpdate()
	if update == nil || update.Attrs == nil {
		//This happens a lot
		//log.Printf("Error retrieving AS Path\n")
		return false
	}
	pathlen := len(update.Attrs.AsPath)
	if pathlen < 1 {
		// This happens sometimes
		//log.Printf("Error: empty AS Path\n")
		return false
	}

	lastseg := update.Attrs.AsPath[0]

	var firstAs uint32 = 0
	if lastseg.AsSeq != nil && len(lastseg.AsSeq) > 0 {
		firstAs = lastseg.AsSeq[0]
	} else if lastseg.AsSet != nil && len(lastseg.AsSet) > 0 {
		firstAs = lastseg.AsSet[0]
	} else {
		// Both are empty, not sure how thats possible, but return false
		return false
	}
	return firstAs == asval.as

}

type transformer interface {
	transform(int, *mrt.MrtBufferStack) string
	summarize()
}

type UniquePrefixMap struct {
	output    *os.File
	prefixes  map[string]bool
	radixTree *radix.Tree
	maplock   *sync.Mutex
}

func NewUniquePrefixMap(o *os.File) *UniquePrefixMap {
	upm := UniquePrefixMap{}
	upm.output = o
	upm.prefixes = make(map[string]bool)
	upm.radixTree = radix.New()
	upm.maplock = &sync.Mutex{}
	return &upm
}

type IPWrapper struct {
	ip string
}

func (upm *UniquePrefixMap) transform(msgNum int, mbs *mrt.MrtBufferStack) string {
	update := mbs.Bgpupbuf.(protoparse.BGPUpdater).GetUpdate()

	//If there are no advertized routes
	if update.AdvertizedRoutes == nil || len(update.AdvertizedRoutes.Prefixes) == 0 {
		return ""
	}

	for _, ar := range update.AdvertizedRoutes.Prefixes {
		ipstr := fmt.Sprintf("%s/%d", net.IP(util.GetIP(ar.GetPrefix())), ar.Mask)
		upm.maplock.Lock()
		if !upm.prefixes[ipstr] {
			rkey := util.IpToRadixkey(util.GetIP(ar.GetPrefix()), uint8(ar.Mask))
			upm.radixTree.Insert(rkey, IPWrapper{ipstr})
			upm.prefixes[ipstr] = true
		}
		upm.maplock.Unlock()
	}

	return ""

}

func (upm *UniquePrefixMap) summarize() {
	upm.radixTree.Walk(upm.topWalk)
}

func (upm *UniquePrefixMap) topWalk(s string, v interface{}) bool {
	pref := v.(IPWrapper).ip
	if upm.prefixes[pref] {
		str := fmt.Sprintf("%s\n", pref)
		upm.output.WriteString(str)
		upm.radixTree.WalkPrefix(s, upm.subWalk)
	}
	return false
}

func (upm *UniquePrefixMap) subWalk(s string, v interface{}) bool {
	upm.prefixes[v.(IPWrapper).ip] = false
	return false
}

type jsonTransformer struct{}

func (j jsonTransformer) transform(msgNum int, mbs *mrt.MrtBufferStack) string {
	mbsj, err := json.Marshal(mbs)
	if err != nil {
		log.Printf("Error marshaling to json")
		return ""
	}
	return string(mbsj) + "\n"
}

func (j jsonTransformer) summarize() {}

type textTransformer struct{}

func (t textTransformer) transform(msgNum int, mbs *mrt.MrtBufferStack) string {
	ret := ""
	ret += fmt.Sprintf("[%d] MRT Header: %s\n", msgNum, mbs.MrthBuf)
	ret += fmt.Sprintf("BGP4MP Header:%s\n", mbs.Bgp4mpbuf)
	ret += fmt.Sprintf("BGP Header: %s\n", mbs.Bgphbuf)
	ret += fmt.Sprintf("BGP Update:%s\n", mbs.Bgpupbuf)
	return ret + "\n"
}

func (t textTransformer) summarize() {}

func parseHeaders(mrth protoparse.PbVal, entryCt int) (bgp4h, bgph, bgpup protoparse.PbVal, err error) {
	bgp4h, err = mrth.Parse()
	if err != nil {
		log.Printf("Failed parsing MRT header %d :%s", entryCt, err)
		return
	}

	bgph, err = bgp4h.Parse()
	if err != nil {
		log.Printf("Failed parsing BG4MP header %d :%s", entryCt, err)
		return
	}

	bgpup, err = bgph.Parse()
	if err != nil {
		log.Printf("Failed parsing BGP header %d :%s", entryCt, err)
		return
	}

	_, err = bgpup.Parse()
	if err != nil {
		log.Printf("Failed parsing BGP update %d :%s", entryCt, err)
		return
	}

	return
}
<|endoftext|>


<commit_before>//go:build ignore
// +build ignore

package make_test

import (
	"bytes"
	"errors"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
)

const (
	inWindows = (runtime.GOOS == "windows")
)

var script = "all.bash"
var gopRoot = ""
var gopBinFiles = []string{"gop", "gopfmt"}
var installer = "cmd/make.go"
var versionFile = "VERSION"
var mainVersionFile = "env/version.go"

func checkPathExist(path string, isDir bool) bool {
	stat, err := os.Stat(path)
	isExists := !os.IsNotExist(err)
	if isDir {
		return isExists && stat.IsDir()
	}
	return isExists && !stat.IsDir()
}

func trimRight(s string) string {
	return strings.TrimRight(s, " \t\r\n")
}

func execCommand(command string, arg ...string) (string, error) {
	var stdout, stderr bytes.Buffer
	cmd := exec.Command(command, arg...)
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err := cmd.Run()
	if err != nil {
		if stderr.Len() > 0 {
			err = errors.New(string(stderr.String()))
		}
	}
	return stdout.String(), err
}

func getBranch() string {
	stdout, err := execCommand("git", "rev-parse", "--abbrev-ref", "HEAD")
	if err != nil {
		return ""
	}
	return trimRight(stdout)
}

func detectGoBinPath() string {
	goBin, ok := os.LookupEnv("GOBIN")
	if ok {
		return goBin
	}
	goPath, ok := os.LookupEnv("GOPATH")
	if ok {
		list := filepath.SplitList(goPath)
		if len(list) > 0 {
			// Put in first directory of $GOPATH.
			return filepath.Join(list[0], "bin")
		}
	}
	homeDir, _ := os.UserHomeDir()
	return filepath.Join(homeDir, "go", "bin")
}

func init() {
	pwd, _ := os.Getwd()
	gopRoot = filepath.Join(pwd, "..")
	installer = filepath.Join(gopRoot, installer)
	versionFile = filepath.Join(gopRoot, versionFile)
	mainVersionFile = filepath.Join(gopRoot, mainVersionFile)

	if inWindows {
		script = "all.bat"
		for index, file := range gopBinFiles {
			file += ".exe"
			gopBinFiles[index] = file
		}
	}
}

func cleanGopRunCacheFiles(t *testing.T) {
	homeDir, _ := os.UserHomeDir()
	runCacheDir := filepath.Join(homeDir, ".gop", "run")
	files := []string{"go.mod", "go.sum"}
	for _, file := range files {
		fullPath := filepath.Join(runCacheDir, file)
		if checkPathExist(fullPath, false) {
			t.Fatalf("Failed: %s found in %s directory\n", file, runCacheDir)
		}
	}
}

func TestAllScript(t *testing.T) {
	os.Chdir(gopRoot)
	cmd := exec.Command(filepath.Join(gopRoot, script))
	if output, err := cmd.CombinedOutput(); err != nil {
		t.Fatalf("Failed: %v:\nOut: %s\n", err, output)
	}

	goBinPath := detectGoBinPath()

	for _, file := range gopBinFiles {
		if !checkPathExist(filepath.Join(gopRoot, "bin", file), false) {
			t.Fatalf("Failed: %s not found in ./bin directory\n", file)
		}
		if !checkPathExist(filepath.Join(goBinPath, file), false) {
			t.Fatalf("Failed: %s not found in %s/bin directory\n", file, goBinPath)
		}
	}

	cleanGopRunCacheFiles(t)

	cmd = exec.Command(filepath.Join(gopRoot, "bin", gopBinFiles[0]), "version")
	if output, err := cmd.CombinedOutput(); err != nil {
		t.Fatalf("Failed: %v:\nOut: %s\n", err, output)
	}
}

func TestTagFlagInGitRepo(t *testing.T) {
	os.Chdir(gopRoot)

	// Setup
	tag := "v1.0.90"
	tag2 := "v1.0.91"
	bigtag := "v1.999.12"
	releaseBranch := "v1.0"
	nonExistBranch := "non-exist-branch"
	sourceBranch := getBranch()

	// Teardown
	t.Cleanup(func() {
		gitCmd := exec.Command("git", "tag", "-d", tag)
		gitCmd.CombinedOutput()
		gitCmd = exec.Command("git", "tag", "-d", tag2)
		gitCmd.CombinedOutput()
		execCommand("git", "checkout", sourceBranch)
		execCommand("git", "branch", "-D", nonExistBranch)
		execCommand("git", "branch", "-D", releaseBranch)
	})

	t.Run("release new version with bad tag", func(t *testing.T) {
		cmd := exec.Command("go", "run", installer, "--nopush", "--tag", "xyz")
		if _, err := cmd.CombinedOutput(); err == nil {
			t.Fatal("Failed: release a bad tag should be failed.")
		}
	})

	t.Run("empty tag should failed", func(t *testing.T) {
		cmd := exec.Command("go", "run", installer, "--nopush", "--tag", "")
		if out, err := cmd.CombinedOutput(); err != nil || !strings.Contains(string(out), "Usage") {
			t.Fatalf("Failed: %v, out: %s\n", err, out)
		}
	})

	t.Run("failed when release branch corresponding to tag does not exists", func(t *testing.T) {
		cmd := exec.Command("go", "run", installer, "--nopush", "--tag", bigtag)
		if _, err := cmd.CombinedOutput(); err == nil {
			t.Fatal("Failed: a corresponding release branch does not exists should be failed.")
		}
	})

	t.Run("release new version on release branch", func(t *testing.T) {
		execCommand("git", "checkout", sourceBranch)
		execCommand("git", "branch", "-D", releaseBranch)
		_, err := execCommand("git", "checkout", "-b", releaseBranch)
		if err != nil {
			t.Fatal(err)
		}

		cmd := exec.Command("go", "run", installer, "--nopush", "--tag", tag)
		if out, err := cmd.CombinedOutput(); err != nil {
			t.Log(string(out))
			t.Fatalf("Failed: release tag: %s on branch: %s should not be failed", tag, releaseBranch)
		}

		if !checkPathExist(versionFile, false) {
			t.Fatal("Failed: a VERSION file not found.")
		}

		if data, _ := os.ReadFile(versionFile); string(data) != tag {
			t.Fatalf("Failed: content of VERSION file: '%s' not match tag: %s", data, tag)
		}

		// Make sure tag exists.
		gitCmd := exec.Command("git", "tag")
		if allTags, _ := gitCmd.CombinedOutput(); !strings.Contains(string(allTags), tag) {
			t.Fatalf("Failed: %s tag not found in tag list of this git repo.\n", tag)
		}
	})

	t.Run("release new version on non-release branch", func(t *testing.T) {
		execCommand("git", "checkout", sourceBranch)
		execCommand("git", "branch", "-D", releaseBranch)
		_, err := execCommand("git", "checkout", "-b", releaseBranch)
		if err != nil {
			t.Fatal(err)
		}

		gitCmd := exec.Command("git", "checkout", "-b", nonExistBranch)
		if output, err := gitCmd.CombinedOutput(); err != nil {
			if strings.Contains(string(output), "already exists") {
				gitCmd = exec.Command("git", "checkout", nonExistBranch)
				if output, err := gitCmd.CombinedOutput(); err != nil {
					t.Fatalf("Failed: %v, output: %s\n", err, output)
				}
			} else {
				t.Fatalf("Failed: %v, output: %s\n", err, output)
			}
		}

		cmd := exec.Command("go", "run", installer, "--nopush", "--tag", tag2)
		if out, err := cmd.CombinedOutput(); err != nil {
			t.Log(string(out))
			t.Fatalf("Failed: release tag: %s on branch: %s should not be failed", tag2, nonExistBranch)
		}

		if getBranch() != nonExistBranch {
			t.Fatal("Failed: getBranch() != nonExistBranch")
		}
		execCommand("git", "checkout", releaseBranch)

		if !checkPathExist(versionFile, false) {
			t.Fatal("Failed: a VERSION file not found.")
		}

		if data, _ := os.ReadFile(versionFile); string(data) != tag2 {
			t.Fatalf("Failed: content of VERSION file: '%s' not match tag: %s", data, tag2)
		}

		gitCmd = exec.Command("git", "tag")
		if allTags, _ := gitCmd.CombinedOutput(); !strings.Contains(string(allTags), tag2) {
			t.Fatalf("Failed: %s tag not found in tag list of this git repo.\n", tag)
		}
	})
}

func TestTagFlagInNonGitRepo(t *testing.T) {
	os.Chdir(gopRoot)

	// Setup
	gitDir := filepath.Join(gopRoot, ".git")
	gitBackupDir := filepath.Join(gopRoot, ".gitBackup")
	// Rename .git dir
	if checkPathExist(gitDir, true) {
		os.Rename(gitDir, gitBackupDir)
	}

	// Teardown
	t.Cleanup(func() {
		if checkPathExist(gitBackupDir, true) {
			os.Rename(gitBackupDir, gitDir)
		}
	})

	t.Run("specify new tag", func(t *testing.T) {
		cmd := exec.Command("go", "run", installer, "--tag", "v1.0.98")
		output, err := cmd.CombinedOutput()
		if err == nil || !strings.Contains(string(output), "Error") {
			t.Fatal("Failed: release on a non-git repo should be failed.")
		}
	})
}

func TestInstallInNonGitRepo(t *testing.T) {
	os.Chdir(gopRoot)

	// Setup
	gitDir := filepath.Join(gopRoot, ".git")
	gitBackupDir := filepath.Join(gopRoot, ".gitBackup")
	// Rename .git dir
	if checkPathExist(gitDir, true) {
		os.Rename(gitDir, gitBackupDir)
	}

	// Teardown
	t.Cleanup(func() {
		if checkPathExist(gitBackupDir, true) {
			os.Rename(gitBackupDir, gitDir)
		}
		if checkPathExist(versionFile, false) {
			os.Remove(versionFile)
		}
	})

	installCmd := func() *exec.Cmd {
		return exec.Command("go", "run", installer, "--install")
	}

	t.Run("failed build operation", func(t *testing.T) {
		cmd := installCmd()
		output, err := cmd.CombinedOutput()
		if err == nil || !strings.Contains(string(output), "Error") {
			t.Fatal("Failed: build Go+ in a non-git repo and no VERSION file should be failed.")
		}
	})

	t.Run("install with VERSION file", func(t *testing.T) {
		version := "v1.1.98"
		// Create VERSION file
		if err := os.WriteFile(versionFile, []byte(version), 0644); err != nil {
			t.Fatal(err)
		}

		cmd := installCmd()
		if output, err := cmd.CombinedOutput(); err != nil {
			t.Fatalf("Failed: %v, output: %s\n", err, output)
		}

		cmd = exec.Command(filepath.Join(gopRoot, "bin", gopBinFiles[0]), "version")
		output, err := cmd.CombinedOutput()
		if err != nil || !strings.Contains(string(output), version) {
			t.Fatalf("Failed: %v, output: %s\n", err, output)
		}
	})
}

func TestHandleMultiFlags(t *testing.T) {
	os.Chdir(gopRoot)

	// Setup
	nextVersion := "v1.0.988"

	// Teardown
	t.Cleanup(func() {
		gitCmd := exec.Command("git", "tag", "-d", nextVersion)
		gitCmd.CombinedOutput()
		if checkPathExist(versionFile, false) {
			os.Remove(versionFile)
		}
	})

	cmd := exec.Command("go", "run", installer, "--install", "--test", "--uninstall", "--tag", nextVersion)
	if output, err := cmd.CombinedOutput(); err != nil {
		t.Fatalf("Failed: %v:\nOut: %s\n", err, output)
	}

	goBinPath := detectGoBinPath()

	// Uninstall will be the last action to run, test if all build assets being cleared.
	for _, file := range gopBinFiles {
		if checkPathExist(filepath.Join(gopRoot, "bin", file), false) {
			t.Fatalf("Failed: %s found in ./bin directory\n", file)
		}
		if checkPathExist(filepath.Join(goBinPath, file), false) {
			t.Fatalf("Failed: %s found in %s/bin directory\n", file, goBinPath)
		}
	}

	// Test if git tag list contains nextVersion
	cmd = exec.Command("git", "tag")
	if allTags, err := cmd.CombinedOutput(); err != nil {
		if !strings.Contains(string(allTags), nextVersion) {
			t.Fatalf("Failed: %s tag not found in this git repo.\n", nextVersion)
		}
	}

	if !checkPathExist(versionFile, false) {
		t.Fatal("Failed: a VERSION file not found.")
	}

	if data, _ := os.ReadFile(versionFile); string(data) != nextVersion {
		t.Fatalf("Failed: content of VERSION file: '%s' not match tag: %s", data, nextVersion)
	}
}
<commit_msg>TestHandleMultiFlags<commit_after>//go:build ignore
// +build ignore

package make_test

import (
	"bytes"
	"errors"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
)

const (
	inWindows = (runtime.GOOS == "windows")
)

var script = "all.bash"
var gopRoot = ""
var gopBinFiles = []string{"gop", "gopfmt"}
var installer = "cmd/make.go"
var versionFile = "VERSION"
var mainVersionFile = "env/version.go"

func checkPathExist(path string, isDir bool) bool {
	stat, err := os.Stat(path)
	isExists := !os.IsNotExist(err)
	if isDir {
		return isExists && stat.IsDir()
	}
	return isExists && !stat.IsDir()
}

func trimRight(s string) string {
	return strings.TrimRight(s, " \t\r\n")
}

func execCommand(command string, arg ...string) (string, error) {
	var stdout, stderr bytes.Buffer
	cmd := exec.Command(command, arg...)
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err := cmd.Run()
	if err != nil {
		if stderr.Len() > 0 {
			err = errors.New(string(stderr.String()))
		}
	}
	return stdout.String(), err
}

func getBranch() string {
	stdout, err := execCommand("git", "rev-parse", "--abbrev-ref", "HEAD")
	if err != nil {
		return ""
	}
	return trimRight(stdout)
}

func detectGoBinPath() string {
	goBin, ok := os.LookupEnv("GOBIN")
	if ok {
		return goBin
	}
	goPath, ok := os.LookupEnv("GOPATH")
	if ok {
		list := filepath.SplitList(goPath)
		if len(list) > 0 {
			// Put in first directory of $GOPATH.
			return filepath.Join(list[0], "bin")
		}
	}
	homeDir, _ := os.UserHomeDir()
	return filepath.Join(homeDir, "go", "bin")
}

func init() {
	pwd, _ := os.Getwd()
	gopRoot = filepath.Join(pwd, "..")
	installer = filepath.Join(gopRoot, installer)
	versionFile = filepath.Join(gopRoot, versionFile)
	mainVersionFile = filepath.Join(gopRoot, mainVersionFile)

	if inWindows {
		script = "all.bat"
		for index, file := range gopBinFiles {
			file += ".exe"
			gopBinFiles[index] = file
		}
	}
}

func cleanGopRunCacheFiles(t *testing.T) {
	homeDir, _ := os.UserHomeDir()
	runCacheDir := filepath.Join(homeDir, ".gop", "run")
	files := []string{"go.mod", "go.sum"}
	for _, file := range files {
		fullPath := filepath.Join(runCacheDir, file)
		if checkPathExist(fullPath, false) {
			t.Fatalf("Failed: %s found in %s directory\n", file, runCacheDir)
		}
	}
}

func TestAllScript(t *testing.T) {
	os.Chdir(gopRoot)
	cmd := exec.Command(filepath.Join(gopRoot, script))
	if output, err := cmd.CombinedOutput(); err != nil {
		t.Fatalf("Failed: %v:\nOut: %s\n", err, output)
	}

	goBinPath := detectGoBinPath()

	for _, file := range gopBinFiles {
		if !checkPathExist(filepath.Join(gopRoot, "bin", file), false) {
			t.Fatalf("Failed: %s not found in ./bin directory\n", file)
		}
		if !checkPathExist(filepath.Join(goBinPath, file), false) {
			t.Fatalf("Failed: %s not found in %s/bin directory\n", file, goBinPath)
		}
	}

	cleanGopRunCacheFiles(t)

	cmd = exec.Command(filepath.Join(gopRoot, "bin", gopBinFiles[0]), "version")
	if output, err := cmd.CombinedOutput(); err != nil {
		t.Fatalf("Failed: %v:\nOut: %s\n", err, output)
	}
}

func TestTagFlagInGitRepo(t *testing.T) {
	os.Chdir(gopRoot)

	// Setup
	tag := "v1.0.90"
	tag2 := "v1.0.91"
	bigtag := "v1.999.12"
	releaseBranch := "v1.0"
	nonExistBranch := "non-exist-branch"
	sourceBranch := getBranch()

	// Teardown
	t.Cleanup(func() {
		gitCmd := exec.Command("git", "tag", "-d", tag)
		gitCmd.CombinedOutput()
		gitCmd = exec.Command("git", "tag", "-d", tag2)
		gitCmd.CombinedOutput()
		execCommand("git", "checkout", sourceBranch)
		execCommand("git", "branch", "-D", nonExistBranch)
		execCommand("git", "branch", "-D", releaseBranch)
	})

	t.Run("release new version with bad tag", func(t *testing.T) {
		cmd := exec.Command("go", "run", installer, "--nopush", "--tag", "xyz")
		if _, err := cmd.CombinedOutput(); err == nil {
			t.Fatal("Failed: release a bad tag should be failed.")
		}
	})

	t.Run("empty tag should failed", func(t *testing.T) {
		cmd := exec.Command("go", "run", installer, "--nopush", "--tag", "")
		if out, err := cmd.CombinedOutput(); err != nil || !strings.Contains(string(out), "Usage") {
			t.Fatalf("Failed: %v, out: %s\n", err, out)
		}
	})

	t.Run("failed when release branch corresponding to tag does not exists", func(t *testing.T) {
		cmd := exec.Command("go", "run", installer, "--nopush", "--tag", bigtag)
		if _, err := cmd.CombinedOutput(); err == nil {
			t.Fatal("Failed: a corresponding release branch does not exists should be failed.")
		}
	})

	t.Run("release new version on release branch", func(t *testing.T) {
		execCommand("git", "checkout", sourceBranch)
		execCommand("git", "branch", "-D", releaseBranch)
		_, err := execCommand("git", "checkout", "-b", releaseBranch)
		if err != nil {
			t.Fatal(err)
		}

		cmd := exec.Command("go", "run", installer, "--nopush", "--tag", tag)
		if out, err := cmd.CombinedOutput(); err != nil {
			t.Log(string(out))
			t.Fatalf("Failed: release tag: %s on branch: %s should not be failed", tag, releaseBranch)
		}

		if !checkPathExist(versionFile, false) {
			t.Fatal("Failed: a VERSION file not found.")
		}

		if data, _ := os.ReadFile(versionFile); string(data) != tag {
			t.Fatalf("Failed: content of VERSION file: '%s' not match tag: %s", data, tag)
		}

		// Make sure tag exists.
		gitCmd := exec.Command("git", "tag")
		if allTags, _ := gitCmd.CombinedOutput(); !strings.Contains(string(allTags), tag) {
			t.Fatalf("Failed: %s tag not found in tag list of this git repo.\n", tag)
		}
	})

	t.Run("release new version on non-release branch", func(t *testing.T) {
		execCommand("git", "checkout", sourceBranch)
		execCommand("git", "branch", "-D", releaseBranch)
		_, err := execCommand("git", "checkout", "-b", releaseBranch)
		if err != nil {
			t.Fatal(err)
		}

		gitCmd := exec.Command("git", "checkout", "-b", nonExistBranch)
		if output, err := gitCmd.CombinedOutput(); err != nil {
			if strings.Contains(string(output), "already exists") {
				gitCmd = exec.Command("git", "checkout", nonExistBranch)
				if output, err := gitCmd.CombinedOutput(); err != nil {
					t.Fatalf("Failed: %v, output: %s\n", err, output)
				}
			} else {
				t.Fatalf("Failed: %v, output: %s\n", err, output)
			}
		}

		cmd := exec.Command("go", "run", installer, "--nopush", "--tag", tag2)
		if out, err := cmd.CombinedOutput(); err != nil {
			t.Log(string(out))
			t.Fatalf("Failed: release tag: %s on branch: %s should not be failed", tag2, nonExistBranch)
		}

		if getBranch() != nonExistBranch {
			t.Fatal("Failed: getBranch() != nonExistBranch")
		}
		execCommand("git", "checkout", releaseBranch)

		if !checkPathExist(versionFile, false) {
			t.Fatal("Failed: a VERSION file not found.")
		}

		if data, _ := os.ReadFile(versionFile); string(data) != tag2 {
			t.Fatalf("Failed: content of VERSION file: '%s' not match tag: %s", data, tag2)
		}

		gitCmd = exec.Command("git", "tag")
		if allTags, _ := gitCmd.CombinedOutput(); !strings.Contains(string(allTags), tag2) {
			t.Fatalf("Failed: %s tag not found in tag list of this git repo.\n", tag)
		}
	})
}

func TestTagFlagInNonGitRepo(t *testing.T) {
	os.Chdir(gopRoot)

	// Setup
	gitDir := filepath.Join(gopRoot, ".git")
	gitBackupDir := filepath.Join(gopRoot, ".gitBackup")
	// Rename .git dir
	if checkPathExist(gitDir, true) {
		os.Rename(gitDir, gitBackupDir)
	}

	// Teardown
	t.Cleanup(func() {
		if checkPathExist(gitBackupDir, true) {
			os.Rename(gitBackupDir, gitDir)
		}
	})

	t.Run("specify new tag", func(t *testing.T) {
		cmd := exec.Command("go", "run", installer, "--tag", "v1.0.98")
		output, err := cmd.CombinedOutput()
		if err == nil || !strings.Contains(string(output), "Error") {
			t.Fatal("Failed: release on a non-git repo should be failed.")
		}
	})
}

func TestInstallInNonGitRepo(t *testing.T) {
	os.Chdir(gopRoot)

	// Setup
	gitDir := filepath.Join(gopRoot, ".git")
	gitBackupDir := filepath.Join(gopRoot, ".gitBackup")
	// Rename .git dir
	if checkPathExist(gitDir, true) {
		os.Rename(gitDir, gitBackupDir)
	}

	// Teardown
	t.Cleanup(func() {
		if checkPathExist(gitBackupDir, true) {
			os.Rename(gitBackupDir, gitDir)
		}
		if checkPathExist(versionFile, false) {
			os.Remove(versionFile)
		}
	})

	installCmd := func() *exec.Cmd {
		return exec.Command("go", "run", installer, "--install")
	}

	t.Run("failed build operation", func(t *testing.T) {
		cmd := installCmd()
		output, err := cmd.CombinedOutput()
		if err == nil || !strings.Contains(string(output), "Error") {
			t.Fatal("Failed: build Go+ in a non-git repo and no VERSION file should be failed.")
		}
	})

	t.Run("install with VERSION file", func(t *testing.T) {
		version := "v1.1.98"
		// Create VERSION file
		if err := os.WriteFile(versionFile, []byte(version), 0644); err != nil {
			t.Fatal(err)
		}

		cmd := installCmd()
		if output, err := cmd.CombinedOutput(); err != nil {
			t.Fatalf("Failed: %v, output: %s\n", err, output)
		}

		cmd = exec.Command(filepath.Join(gopRoot, "bin", gopBinFiles[0]), "version")
		output, err := cmd.CombinedOutput()
		if err != nil || !strings.Contains(string(output), version) {
			t.Fatalf("Failed: %v, output: %s\n", err, output)
		}
	})
}

func TestHandleMultiFlags(t *testing.T) {
	os.Chdir(gopRoot)

	cmd := exec.Command("go", "run", installer, "--install", "--test", "--uninstall")
	if output, err := cmd.CombinedOutput(); err != nil {
		t.Fatalf("Failed: %v:\nOut: %s\n", err, output)
	}

	goBinPath := detectGoBinPath()

	// Uninstall will be the last action to run, test if all build assets being cleared.
	for _, file := range gopBinFiles {
		if checkPathExist(filepath.Join(gopRoot, "bin", file), false) {
			t.Fatalf("Failed: %s found in ./bin directory\n", file)
		}
		if checkPathExist(filepath.Join(goBinPath, file), false) {
			t.Fatalf("Failed: %s found in %s/bin directory\n", file, goBinPath)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"reflect"

	"github.com/codegangsta/cli"
	"github.com/dustin/go-humanize"
	"github.com/pivotal-golang/bytefmt"
	"github.com/unigraph/rdb"
)

const (
	kB = 1024
	MB = kB * 1024
	GB = MB * 1024
)

type bSize uint64

func (b *bSize) Set(value string) (err error) {
	val, err := humanize.ParseBytes(value)
	if err != nil {
		return err
	}
	*b = bSize(val)
	return
}

func (b *bSize) String() string { return bytefmt.ByteSize(uint64(*b)) }

// https://github.com/facebook/rocksdb/blob/master/include/rocksdb/options.h
var (
	compression_type = cli.StringFlag{
		Name:  "compression_type",
		Value: DefaultOptions.CompressionType,
		Usage: "(none, lz4, snappy, zlib, bzip"}

	num_levels = cli.IntFlag{
		Name:  "num_levels",
		Value: DefaultOptions.NumLevels,
		Usage: "It is safe for num_levels to be bigger than expected number of levels in the database. Some higher levels may be empty, but this will not impact performance in any way. Only change this option if you expect your number of levels will be greater than 7",
	}

	write_buffer_size = cli.GenericFlag{
		Name:  "write_buffer_size",
		Value: &DefaultOptions.WriteBufferSize,
		Usage: "Size of a single memtable. Once memtable exceeds this size, it is marked immutable and a new one is created."}

	max_write_buffer_number = cli.IntFlag{
		Name:  "max_write_buffer_number",
		Value: DefaultOptions.MaxWriteBufferNumber,
		Usage: "Maximum number of memtables, both active and immutable. If the active memtable fills up and the total number of memtables is larger than max_write_buffer_number we stall further writes. This may happen if the flush process is slower than the write rate.",
	}

	min_write_buffer_number_to_merge = cli.IntFlag{
		Name:  "min_write_buffe_number_to_merge",
		Value: DefaultOptions.MinWriteBufferNumberToMerge,
		Usage: "Minimum number of memtables to be merged before flushing to storage. For example, if this option is set to 2, immutable memtables are only flushed when there are two of them - a single immutable memtable will never be flushed. If multiple memtables are merged together, less data may be written to storage since two updates are merged to a single key. However, every Get() must traverse all immutable memtables linearly to check if the key is there. Setting this option too high may hurt read performance.",
	}

	level0_file_num_compaction_trigger = cli.IntFlag{
		Name:  "level0_file_num_compaction_trigger",
		Value: DefaultOptions.Level0FileNumCompactionTrigger,
		Usage: `Once level 0 reaches this number of files, L0->L1 compaction is triggered. We can therefore estimate level 0 size in stable state as write_buffer_size * min_write_buffer_number_to_merge * level0_file_num_compaction_trigger.`,
	}

	level0_slowdown_writes_trigger = cli.IntFlag{
		Name:  "level0_slowdown_writes_trigger",
		Value: DefaultOptions.Level0SlowdownWritesTrigger,
		Usage: "When the number of level 0 files is greater than the slowdown limit, writes are stalled."}

	level0_stop_writes_trigger = cli.IntFlag{
		Name:  "level0_stop_writes_trigger",
		Value: DefaultOptions.Level0StopWritesTrigger,
		Usage: "When the number is greater than stop limit, writes are fully stopped until compaction is done."}

	target_file_size_base = cli.GenericFlag{
		Name:  "target_file_size_base",
		Value: &DefaultOptions.TargetFileSizeBase,
		Usage: "Files in level 1 will have target_file_size_base bytes. Each next level's file size will be target_file_size_multiplier bigger than previous one. However, by default target_file_size_multiplier is 1, so files in all L1..Lmax levels are equal. Increasing target_file_size_base will reduce total number of database files, which is generally a good thing. We recommend setting target_file_size_base to be max_bytes_for_level_base / 10, so that there are 10 files in level 1.",
	}

	target_file_size_multiplier = cli.IntFlag{
		Name:  "target_file_size_multiplier",
		Value: DefaultOptions.TargetFileSizeMultiplier,
		Usage: "Each next level's file size will be target_file_size_multiplier bigger than previous one. By default target_file_size_multiplier is 1, which means by default files in different levels will have similar size.",
	}

	max_bytes_for_level_base = cli.GenericFlag{
		Name:  "max_bytes_for_level_base",
		Value: &DefaultOptions.MaxBytesForLevelBase,
		Usage: "Total size of level 1. As mentioned, we recommend that this be around the size of level 0. ",
	}

	max_bytes_for_level_multiplier = cli.IntFlag{
		Name:  "max_bytes_for_level_multiplier",
		Value: DefaultOptions.MaxBytesForLevelMultiplier,
		Usage: "Each subsequent level is max_bytes_for_level_multiplier larger than previous one. The default is 10 and we do not recommend changing that."}

	bulk = cli.BoolFlag{
		Name:  "bulk",
		Usage: "Sets options for bulk data load. Modifies level0_file_num_compaction_trigger (1G), level0_slowdown_writes_trigger(1G), level0_stop_writes_trigger(1G)"}

	source_compaction_factor = cli.IntFlag{
		Name:  "source_compaction_factor",
		Value: DefaultOptions.SourceCompactionFactor,
		Usage: "Maximum number of bytes in all source files to be compacted in a single compaction run. We avoid picking too many files in the source level so that we do not exceed the total source bytes for compaction to exceed (source_compaction_factor * targetFileSizeLevel()) many bytes. Default:1, i.e. pick maxfilesize amount of data as the source of a compaction."}

	disable_auto_compactions = cli.BoolFlag{
		Name:  "disable_auto_compactions",
		Usage: "disables auto compactions",
	}

	config = cli.StringFlag{
		Name:  "config",
		Usage: "rdb configuration file with options (alternative to command line parameters)",
	}
)

var DefaultOptions = Options{
	CompressionType:                "lz4",
	NumLevels:                      7,
	WriteBufferSize:                1 * GB,
	MaxWriteBufferNumber:           8,
	MinWriteBufferNumberToMerge:    2,
	Level0FileNumCompactionTrigger: 4,
	Level0SlowdownWritesTrigger:    20,
	Level0StopWritesTrigger:        24,
	TargetFileSizeBase:             1 * GB,
	TargetFileSizeMultiplier:       1,
	MaxBytesForLevelBase:           2 * GB,
	MaxBytesForLevelMultiplier:     10,
	SourceCompactionFactor:         1,
	DisableAutoCompactions:         false,
	Bulk: false,
}

type Options struct {
	CompressionType                string
	NumLevels                      int
	WriteBufferSize                bSize
	MaxWriteBufferNumber           int
	MinWriteBufferNumberToMerge    int
	Level0FileNumCompactionTrigger int
	Level0SlowdownWritesTrigger    int
	Level0StopWritesTrigger        int
	TargetFileSizeBase             bSize
	TargetFileSizeMultiplier       int
	MaxBytesForLevelBase           bSize
	MaxBytesForLevelMultiplier     int
	SourceCompactionFactor         int
	DisableAutoCompactions         bool
	Bulk                           bool
}

type flags []cli.Flag

var defaultFlags = flags{
	compression_type,
	num_levels,
	write_buffer_size,
	max_write_buffer_number,
	min_write_buffer_number_to_merge,
	level0_file_num_compaction_trigger,
	level0_slowdown_writes_trigger,
	level0_stop_writes_trigger,
	max_bytes_for_level_base,
	max_bytes_for_level_multiplier,
	target_file_size_base,
	target_file_size_multiplier,
	source_compaction_factor,
	disable_auto_compactions,
	bulk,
	config,
}

func (o *Options) Update(c *cli.Context) {
	o.CompressionType = c.GlobalString(compression_type.Name)
	o.NumLevels = c.GlobalInt(num_levels.Name)
	o.MaxWriteBufferNumber = c.GlobalInt(max_write_buffer_number.Name)
	o.MinWriteBufferNumberToMerge = c.GlobalInt(min_write_buffer_number_to_merge.Name)
	o.Level0FileNumCompactionTrigger = (c.GlobalInt(level0_file_num_compaction_trigger.Name))
	o.Level0SlowdownWritesTrigger = (c.GlobalInt(level0_slowdown_writes_trigger.Name))
	o.Level0StopWritesTrigger = (c.GlobalInt(level0_stop_writes_trigger.Name))

	o.MaxBytesForLevelBase = *c.GlobalGeneric(max_bytes_for_level_base.Name).(*bSize)
	o.WriteBufferSize = *c.GlobalGeneric(write_buffer_size.Name).(*bSize)
	o.TargetFileSizeBase = *c.GlobalGeneric(target_file_size_base.Name).(*bSize)

	o.MaxBytesForLevelMultiplier = (c.GlobalInt(max_bytes_for_level_multiplier.Name))
	o.TargetFileSizeMultiplier = (c.GlobalInt(target_file_size_multiplier.Name))
	o.SourceCompactionFactor = (c.GlobalInt(source_compaction_factor.Name))
	o.DisableAutoCompactions = (c.GlobalBool(disable_auto_compactions.Name))
	if c.GlobalBool(bulk.Name) {
		o.Bulk = true
	}
}

func (o *Options) SetOptions(dbOptions *rdb.Options) {
	setCompression(dbOptions, o.CompressionType)
	if o.Bulk {
		dbOptions.PrepareForBulkLoad()
		// this is what happens internaly for bulk load
		// see https://github.com/facebook/rocksdb/blob/master/util/options.cc#L625
		/*
			dbOptions.SetLevel0FileNumCompactionTrigger(1 << 30)
			dbOptions.SetLevel0SlowdownWritesTrigger(1 << 30)
			dbOptions.SetLevel0StopWritesTrigger(1 << 30)
			dbOptions.SetDisableAutoCompactions(true)
			dbOptions.SetDisableDataSync(true)
			dbOptions.SetSourceCompactionFactor(1 << 30)
			dbOptions.SetNumLevels(2)
			dbOptions.SetMaxWriteBufferNumber(6)
			dbOptions.SetMinWriteBufferNumberToMerge(1)
			dbOptions.SetMaxBackgroundFlushes(4)
			dbOptions.SetMaxBackgroundCompactions(4)
			dbOptions.SetTargetFileSizeBase(256 * 1024 * 1024)
		*/

	}
	dbOptions.SetNumLevels(o.NumLevels)
	dbOptions.SetWriteBufferSize(int(o.WriteBufferSize))
	dbOptions.SetMaxWriteBufferNumber(o.MaxWriteBufferNumber)
	dbOptions.SetMinWriteBufferNumberToMerge(o.MinWriteBufferNumberToMerge)
	//
	dbOptions.SetLevel0FileNumCompactionTrigger(o.Level0FileNumCompactionTrigger)
	dbOptions.SetLevel0SlowdownWritesTrigger(o.Level0SlowdownWritesTrigger)
	dbOptions.SetLevel0StopWritesTrigger(o.Level0StopWritesTrigger)
	dbOptions.SetMaxBytesForLevelBase(uint64(o.MaxBytesForLevelBase))
	dbOptions.SetMaxBytesForLevelMultiplier(o.MaxBytesForLevelMultiplier)
	dbOptions.SetTargetFileSizeBase(uint64(o.TargetFileSizeBase))
	dbOptions.SetTargetFileSizeMultiplier(o.TargetFileSizeMultiplier)
	dbOptions.SetSourceCompactionFactor(o.SourceCompactionFactor)
	dbOptions.SetDisableAutoCompactions(o.DisableAutoCompactions)

}

func (f flags) setOptions(dbOptions *rdb.Options, c *cli.Context) {
	configOptions := map[string]interface{}{}
	if configFile := c.GlobalString("config"); configFile != "" {
		file, err := os.Open(configFile)
		if err != nil {
			log.Fatal(err)
		}
		dec := json.NewDecoder(file)
		err = dec.Decode(&configOptions)
		if err != nil {
			log.Fatal(err)
		}
	}

	if c.GlobalBool("bulk") {
		configOptions["NumLevels"] = float64(2)
		configOptions["Level0FileNumCompactionTrigger"] = float64(1 * GB)
		configOptions["Level0SlowdownWritesTrigger"] = float64(1 * GB)
		configOptions["Level0StopWritesTrigger"] = float64(1 * GB)
		configOptions["DisableAutoCompactions"] = true
		configOptions["SourceCompactionFactor"] = float64(1 * GB)
	}

	DefaultOptions.Update(c)
	val := reflect.ValueOf(&DefaultOptions).Elem()
	for k, v := range configOptions {
		switch a := v.(type) {
		case float64:
			val.FieldByName(k).SetInt(int64(a))
		case string:
			val.FieldByName(k).SetString(a)
		case bool:
			val.FieldByName(k).SetBool(a)
		default:
			fmt.Println(reflect.TypeOf(a))
		}
	}
	DefaultOptions.SetOptions(dbOptions)
}

func setCompression(dbOptions *rdb.Options, compressionType string) {
	switch compressionType {
	case "none":
		dbOptions.SetCompression(rdb.NoCompression)
	case "snappy":
		dbOptions.SetCompression(rdb.SnappyCompression)
	case "zlib":
		dbOptions.SetCompression(rdb.ZLibCompression)
	case "bzip":
		dbOptions.SetCompression(rdb.Bz2Compression)
	case "lz4":
		dbOptions.SetCompression(rdb.Lz4Compression)
	default:
		dbOptions.SetCompression(rdb.Lz4Compression)
	}
}
<commit_msg>update<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"reflect"

	"github.com/codegangsta/cli"
	"github.com/dustin/go-humanize"
	"github.com/unigraph/rdb"
)

const (
	kB = 1024
	MB = kB * 1024
	GB = MB * 1024
)

type bSize uint64

func (b *bSize) Set(value string) (err error) {
	val, err := humanize.ParseBytes(value)
	if err != nil {
		return err
	}
	*b = bSize(val)
	return
}

func (b *bSize) String() string { return humanize.IBytes(uint64(*b)) }

// https://github.com/facebook/rocksdb/blob/master/include/rocksdb/options.h
var (
	compression_type = cli.StringFlag{
		Name:  "compression_type",
		Value: DefaultOptions.CompressionType,
		Usage: "(none, lz4, snappy, zlib, bzip"}

	num_levels = cli.IntFlag{
		Name:  "num_levels",
		Value: DefaultOptions.NumLevels,
		Usage: "It is safe for num_levels to be bigger than expected number of levels in the database. Some higher levels may be empty, but this will not impact performance in any way. Only change this option if you expect your number of levels will be greater than 7",
	}

	write_buffer_size = cli.GenericFlag{
		Name:  "write_buffer_size",
		Value: &DefaultOptions.WriteBufferSize,
		Usage: "Size of a single memtable. Once memtable exceeds this size, it is marked immutable and a new one is created."}

	max_write_buffer_number = cli.IntFlag{
		Name:  "max_write_buffer_number",
		Value: DefaultOptions.MaxWriteBufferNumber,
		Usage: "Maximum number of memtables, both active and immutable. If the active memtable fills up and the total number of memtables is larger than max_write_buffer_number we stall further writes. This may happen if the flush process is slower than the write rate.",
	}

	min_write_buffer_number_to_merge = cli.IntFlag{
		Name:  "min_write_buffe_number_to_merge",
		Value: DefaultOptions.MinWriteBufferNumberToMerge,
		Usage: "Minimum number of memtables to be merged before flushing to storage. For example, if this option is set to 2, immutable memtables are only flushed when there are two of them - a single immutable memtable will never be flushed. If multiple memtables are merged together, less data may be written to storage since two updates are merged to a single key. However, every Get() must traverse all immutable memtables linearly to check if the key is there. Setting this option too high may hurt read performance.",
	}

	level0_file_num_compaction_trigger = cli.IntFlag{
		Name:  "level0_file_num_compaction_trigger",
		Value: DefaultOptions.Level0FileNumCompactionTrigger,
		Usage: `Once level 0 reaches this number of files, L0->L1 compaction is triggered. We can therefore estimate level 0 size in stable state as write_buffer_size * min_write_buffer_number_to_merge * level0_file_num_compaction_trigger.`,
	}

	level0_slowdown_writes_trigger = cli.IntFlag{
		Name:  "level0_slowdown_writes_trigger",
		Value: DefaultOptions.Level0SlowdownWritesTrigger,
		Usage: "When the number of level 0 files is greater than the slowdown limit, writes are stalled."}

	level0_stop_writes_trigger = cli.IntFlag{
		Name:  "level0_stop_writes_trigger",
		Value: DefaultOptions.Level0StopWritesTrigger,
		Usage: "When the number is greater than stop limit, writes are fully stopped until compaction is done."}

	target_file_size_base = cli.GenericFlag{
		Name:  "target_file_size_base",
		Value: &DefaultOptions.TargetFileSizeBase,
		Usage: "Files in level 1 will have target_file_size_base bytes. Each next level's file size will be target_file_size_multiplier bigger than previous one. However, by default target_file_size_multiplier is 1, so files in all L1..Lmax levels are equal. Increasing target_file_size_base will reduce total number of database files, which is generally a good thing. We recommend setting target_file_size_base to be max_bytes_for_level_base / 10, so that there are 10 files in level 1.",
	}

	target_file_size_multiplier = cli.IntFlag{
		Name:  "target_file_size_multiplier",
		Value: DefaultOptions.TargetFileSizeMultiplier,
		Usage: "Each next level's file size will be target_file_size_multiplier bigger than previous one. By default target_file_size_multiplier is 1, which means by default files in different levels will have similar size.",
	}

	max_bytes_for_level_base = cli.GenericFlag{
		Name:  "max_bytes_for_level_base",
		Value: &DefaultOptions.MaxBytesForLevelBase,
		Usage: "Total size of level 1. As mentioned, we recommend that this be around the size of level 0. ",
	}

	max_bytes_for_level_multiplier = cli.IntFlag{
		Name:  "max_bytes_for_level_multiplier",
		Value: DefaultOptions.MaxBytesForLevelMultiplier,
		Usage: "Each subsequent level is max_bytes_for_level_multiplier larger than previous one. The default is 10 and we do not recommend changing that."}

	bulk = cli.BoolFlag{
		Name:  "bulk",
		Usage: "Sets options for bulk data load. Modifies level0_file_num_compaction_trigger (1G), level0_slowdown_writes_trigger(1G), level0_stop_writes_trigger(1G)"}

	source_compaction_factor = cli.IntFlag{
		Name:  "source_compaction_factor",
		Value: DefaultOptions.SourceCompactionFactor,
		Usage: "Maximum number of bytes in all source files to be compacted in a single compaction run. We avoid picking too many files in the source level so that we do not exceed the total source bytes for compaction to exceed (source_compaction_factor * targetFileSizeLevel()) many bytes. Default:1, i.e. pick maxfilesize amount of data as the source of a compaction."}

	disable_auto_compactions = cli.BoolFlag{
		Name:  "disable_auto_compactions",
		Usage: "disables auto compactions",
	}

	config = cli.StringFlag{
		Name:  "config",
		Usage: "rdb configuration file with options (alternative to command line parameters)",
	}
)

var DefaultOptions = Options{
	CompressionType:                "lz4",
	NumLevels:                      7,
	WriteBufferSize:                1 * GB,
	MaxWriteBufferNumber:           8,
	MinWriteBufferNumberToMerge:    2,
	Level0FileNumCompactionTrigger: 4,
	Level0SlowdownWritesTrigger:    20,
	Level0StopWritesTrigger:        24,
	TargetFileSizeBase:             1 * GB,
	TargetFileSizeMultiplier:       1,
	MaxBytesForLevelBase:           2 * GB,
	MaxBytesForLevelMultiplier:     10,
	SourceCompactionFactor:         1,
	DisableAutoCompactions:         false,
	Bulk: false,
}

type Options struct {
	CompressionType                string
	NumLevels                      int
	WriteBufferSize                bSize
	MaxWriteBufferNumber           int
	MinWriteBufferNumberToMerge    int
	Level0FileNumCompactionTrigger int
	Level0SlowdownWritesTrigger    int
	Level0StopWritesTrigger        int
	TargetFileSizeBase             bSize
	TargetFileSizeMultiplier       int
	MaxBytesForLevelBase           bSize
	MaxBytesForLevelMultiplier     int
	SourceCompactionFactor         int
	DisableAutoCompactions         bool
	Bulk                           bool
}

type flags []cli.Flag

var defaultFlags = flags{
	compression_type,
	num_levels,
	write_buffer_size,
	max_write_buffer_number,
	min_write_buffer_number_to_merge,
	level0_file_num_compaction_trigger,
	level0_slowdown_writes_trigger,
	level0_stop_writes_trigger,
	max_bytes_for_level_base,
	max_bytes_for_level_multiplier,
	target_file_size_base,
	target_file_size_multiplier,
	source_compaction_factor,
	disable_auto_compactions,
	bulk,
	config,
}

func (o *Options) Update(c *cli.Context) {
	o.CompressionType = c.GlobalString(compression_type.Name)
	o.NumLevels = c.GlobalInt(num_levels.Name)
	o.MaxWriteBufferNumber = c.GlobalInt(max_write_buffer_number.Name)
	o.MinWriteBufferNumberToMerge = c.GlobalInt(min_write_buffer_number_to_merge.Name)
	o.Level0FileNumCompactionTrigger = (c.GlobalInt(level0_file_num_compaction_trigger.Name))
	o.Level0SlowdownWritesTrigger = (c.GlobalInt(level0_slowdown_writes_trigger.Name))
	o.Level0StopWritesTrigger = (c.GlobalInt(level0_stop_writes_trigger.Name))

	o.MaxBytesForLevelBase = *c.GlobalGeneric(max_bytes_for_level_base.Name).(*bSize)
	o.WriteBufferSize = *c.GlobalGeneric(write_buffer_size.Name).(*bSize)
	o.TargetFileSizeBase = *c.GlobalGeneric(target_file_size_base.Name).(*bSize)

	o.MaxBytesForLevelMultiplier = (c.GlobalInt(max_bytes_for_level_multiplier.Name))
	o.TargetFileSizeMultiplier = (c.GlobalInt(target_file_size_multiplier.Name))
	o.SourceCompactionFactor = (c.GlobalInt(source_compaction_factor.Name))
	o.DisableAutoCompactions = (c.GlobalBool(disable_auto_compactions.Name))
	if c.GlobalBool(bulk.Name) {
		o.Bulk = true
	}
}

func (o *Options) SetOptions(dbOptions *rdb.Options) {
	setCompression(dbOptions, o.CompressionType)
	if o.Bulk {
		dbOptions.PrepareForBulkLoad()
		// this is what happens internaly for bulk load
		// see https://github.com/facebook/rocksdb/blob/master/util/options.cc#L625
		/*
			dbOptions.SetLevel0FileNumCompactionTrigger(1 << 30)
			dbOptions.SetLevel0SlowdownWritesTrigger(1 << 30)
			dbOptions.SetLevel0StopWritesTrigger(1 << 30)
			dbOptions.SetDisableAutoCompactions(true)
			dbOptions.SetDisableDataSync(true)
			dbOptions.SetSourceCompactionFactor(1 << 30)
			dbOptions.SetNumLevels(2)
			dbOptions.SetMaxWriteBufferNumber(6)
			dbOptions.SetMinWriteBufferNumberToMerge(1)
			dbOptions.SetMaxBackgroundFlushes(4)
			dbOptions.SetMaxBackgroundCompactions(4)
			dbOptions.SetTargetFileSizeBase(256 * 1024 * 1024)
		*/

	}
	dbOptions.SetNumLevels(o.NumLevels)
	dbOptions.SetWriteBufferSize(int(o.WriteBufferSize))
	dbOptions.SetMaxWriteBufferNumber(o.MaxWriteBufferNumber)
	dbOptions.SetMinWriteBufferNumberToMerge(o.MinWriteBufferNumberToMerge)
	//
	dbOptions.SetLevel0FileNumCompactionTrigger(o.Level0FileNumCompactionTrigger)
	dbOptions.SetLevel0SlowdownWritesTrigger(o.Level0SlowdownWritesTrigger)
	dbOptions.SetLevel0StopWritesTrigger(o.Level0StopWritesTrigger)
	dbOptions.SetMaxBytesForLevelBase(uint64(o.MaxBytesForLevelBase))
	dbOptions.SetMaxBytesForLevelMultiplier(o.MaxBytesForLevelMultiplier)
	dbOptions.SetTargetFileSizeBase(uint64(o.TargetFileSizeBase))
	dbOptions.SetTargetFileSizeMultiplier(o.TargetFileSizeMultiplier)
	dbOptions.SetSourceCompactionFactor(o.SourceCompactionFactor)
	dbOptions.SetDisableAutoCompactions(o.DisableAutoCompactions)

}

func (f flags) setOptions(dbOptions *rdb.Options, c *cli.Context) {
	configOptions := map[string]interface{}{}
	if configFile := c.GlobalString("config"); configFile != "" {
		file, err := os.Open(configFile)
		if err != nil {
			log.Fatal(err)
		}
		dec := json.NewDecoder(file)
		err = dec.Decode(&configOptions)
		if err != nil {
			log.Fatal(err)
		}
	}

	if c.GlobalBool("bulk") {
		configOptions["NumLevels"] = float64(2)
		configOptions["Level0FileNumCompactionTrigger"] = float64(1 * GB)
		configOptions["Level0SlowdownWritesTrigger"] = float64(1 * GB)
		configOptions["Level0StopWritesTrigger"] = float64(1 * GB)
		configOptions["DisableAutoCompactions"] = true
		configOptions["SourceCompactionFactor"] = float64(1 * GB)
	}

	DefaultOptions.Update(c)
	val := reflect.ValueOf(&DefaultOptions).Elem()
	for k, v := range configOptions {
		switch a := v.(type) {
		case float64:
			val.FieldByName(k).SetInt(int64(a))
		case string:
			val.FieldByName(k).SetString(a)
		case bool:
			val.FieldByName(k).SetBool(a)
		default:
			fmt.Println(reflect.TypeOf(a))
		}
	}
	DefaultOptions.SetOptions(dbOptions)
}

func setCompression(dbOptions *rdb.Options, compressionType string) {
	switch compressionType {
	case "none":
		dbOptions.SetCompression(rdb.NoCompression)
	case "snappy":
		dbOptions.SetCompression(rdb.SnappyCompression)
	case "zlib":
		dbOptions.SetCompression(rdb.ZLibCompression)
	case "bzip":
		dbOptions.SetCompression(rdb.Bz2Compression)
	case "lz4":
		dbOptions.SetCompression(rdb.Lz4Compression)
	default:
		dbOptions.SetCompression(rdb.Lz4Compression)
	}
}
<|endoftext|>


<commit_before>package flags

import (
	"fmt"
	"io"
	"os"
)

// IniError contains location information on where in the ini file an error
// occured.
type IniError struct {
	// The error message.
	Message string

	// The filename of the file in which the error occurred.
	File string

	// The line number at which the error occurred.
	LineNumber uint
}

// Error provides a "file:line: message" formatted message of the ini error.
func (x *IniError) Error() string {
	return fmt.Sprintf(
		"%s:%d: %s",
		x.File,
		x.LineNumber,
		x.Message,
	)
}

// IniOptions for writing ini files
type IniOptions uint

const (
	// IniNone indicates no options.
	IniNone IniOptions = 0

	// IniIncludeDefaults indicates that default values should be written
	// when writing options to an ini file.
	IniIncludeDefaults = 1 << iota

	// IniCommentDefaults indicates that if IniIncludeDefaults is used
	// options with default values are written but commented out.
	IniCommentDefaults

	// IniIncludeComments indicates that comments containing the description
	// of an option should be written when writing options to an ini file.
	IniIncludeComments

	// IniDefault provides a default set of options.
	IniDefault = IniIncludeComments
)

// IniParser is a utility to read and write flags options from and to ini
// files.
type IniParser struct {
	parser *Parser
}

// NewIniParser creates a new ini parser for a given Parser.
func NewIniParser(p *Parser) *IniParser {
	return &IniParser{
		parser: p,
	}
}

// IniParse is a convenience function to parse command line options with default
// settings from an ini file. The provided data is a pointer to a struct
// representing the default option group (named "Application Options"). For
// more control, use flags.NewParser.
func IniParse(filename string, data interface{}) error {
	p := NewParser(data, Default)

	return NewIniParser(p).ParseFile(filename)
}

// ParseFile parses flags from an ini formatted file. See Parse for more
// information on the ini file format. The returned errors can be of the type
// flags.Error or flags.IniError.
func (i *IniParser) ParseFile(filename string) error {
	i.parser.storeDefaults()

	ini, err := readIniFromFile(filename)

	if err != nil {
		return err
	}

	return i.parse(ini)
}

// Parse parses flags from an ini format. You can use ParseFile as a
// convenience function to parse from a filename instead of a general
// io.Reader.
//
// The format of the ini file is as follows:
//
//     [Option group name]
//     option = value
//
// Each section in the ini file represents an option group or command in the
// flags parser. The default flags parser option group (i.e. when using
// flags.Parse) is named 'Application Options'. The ini option name is matched
// in the following order:
//
//     1. Compared to the ini-name tag on the option struct field (if present)
//     2. Compared to the struct field name
//     3. Compared to the option long name (if present)
//     4. Compared to the option short name (if present)
//
// Sections for nested groups and commands can be addressed using a dot `.'
// namespacing notation (i.e [subcommand.Options]). Group section names are
// matched case insensitive.
//
// The returned errors can be of the type flags.Error or
// flags.IniError.
func (i *IniParser) Parse(reader io.Reader) error {
	i.parser.storeDefaults()

	ini, err := readIni(reader, "")

	if err != nil {
		return err
	}

	return i.parse(ini)
}

// WriteFile writes the flags as ini format into a file. See WriteIni
// for more information. The returned error occurs when the specified file
// could not be opened for writing.
func (i *IniParser) WriteFile(filename string, options IniOptions) error {
	file, err := os.Create(filename)

	if err != nil {
		return err
	}

	defer file.Close()

	i.Write(file, options)

	return nil
}

// Write writes the current values of all the flags to an ini format.
// See Parse for more information on the ini file format. You typically
// call this only after settings have been parsed since the default values of each
// option are stored just before parsing the flags (this is only relevant when
// IniIncludeDefaults is _not_ set in options).
func (i *IniParser) Write(writer io.Writer, options IniOptions) {
	writeIni(i, writer, options)
}
<commit_msg>redo some comments<commit_after>package flags

import (
	"fmt"
	"io"
	"os"
)

// IniError contains location information on where an error occured.
type IniError struct {
	// The error message.
	Message string

	// The filename of the file in which the error occurred.
	File string

	// The line number at which the error occurred.
	LineNumber uint
}

// Error provides a "file:line: message" formatted message of the ini error.
func (x *IniError) Error() string {
	return fmt.Sprintf(
		"%s:%d: %s",
		x.File,
		x.LineNumber,
		x.Message,
	)
}

// IniOptions for writing
type IniOptions uint

const (
	// IniNone indicates no options.
	IniNone IniOptions = 0

	// IniIncludeDefaults indicates that default values should be written.
	IniIncludeDefaults = 1 << iota

	// IniCommentDefaults indicates that if IniIncludeDefaults is used
	// options with default values are written but commented out.
	IniCommentDefaults

	// IniIncludeComments indicates that comments containing the description
	// of an option should be written.
	IniIncludeComments

	// IniDefault provides a default set of options.
	IniDefault = IniIncludeComments
)

// IniParser is a utility to read and write flags options from and to ini
// formatted strings.
type IniParser struct {
	parser *Parser
}

// NewIniParser creates a new ini parser for a given Parser.
func NewIniParser(p *Parser) *IniParser {
	return &IniParser{
		parser: p,
	}
}

// IniParse is a convenience function to parse command line options with default
// settings from an ini formatted file. The provided data is a pointer to a struct
// representing the default option group (named "Application Options"). For
// more control, use flags.NewParser.
func IniParse(filename string, data interface{}) error {
	p := NewParser(data, Default)

	return NewIniParser(p).ParseFile(filename)
}

// ParseFile parses flags from an ini formatted file. See Parse for more
// information on the ini file format. The returned errors can be of the type
// flags.Error or flags.IniError.
func (i *IniParser) ParseFile(filename string) error {
	i.parser.storeDefaults()

	ini, err := readIniFromFile(filename)

	if err != nil {
		return err
	}

	return i.parse(ini)
}

// Parse parses flags from an ini format. You can use ParseFile as a
// convenience function to parse from a filename instead of a general
// io.Reader.
//
// The format of the ini file is as follows:
//
//     [Option group name]
//     option = value
//
// Each section in the ini file represents an option group or command in the
// flags parser. The default flags parser option group (i.e. when using
// flags.Parse) is named 'Application Options'. The ini option name is matched
// in the following order:
//
//     1. Compared to the ini-name tag on the option struct field (if present)
//     2. Compared to the struct field name
//     3. Compared to the option long name (if present)
//     4. Compared to the option short name (if present)
//
// Sections for nested groups and commands can be addressed using a dot `.'
// namespacing notation (i.e [subcommand.Options]). Group section names are
// matched case insensitive.
//
// The returned errors can be of the type flags.Error or flags.IniError.
func (i *IniParser) Parse(reader io.Reader) error {
	i.parser.storeDefaults()

	ini, err := readIni(reader, "")

	if err != nil {
		return err
	}

	return i.parse(ini)
}

// WriteFile writes the flags as ini format into a file. See WriteIni
// for more information. The returned error occurs when the specified file
// could not be opened for writing.
func (i *IniParser) WriteFile(filename string, options IniOptions) error {
	file, err := os.Create(filename)

	if err != nil {
		return err
	}

	defer file.Close()

	i.Write(file, options)

	return nil
}

// Write writes the current values of all the flags to an ini format.
// See Parse for more information on the ini file format. You typically
// call this only after settings have been parsed since the default values of each
// option are stored just before parsing the flags (this is only relevant when
// IniIncludeDefaults is _not_ set in options).
func (i *IniParser) Write(writer io.Writer, options IniOptions) {
	writeIni(i, writer, options)
}
<|endoftext|>


<commit_before>//  Copyright (c) 2015 Couchbase, Inc.

package gson

import "strconv"
import "fmt"
import "encoding/json"

func gson2collate(obj interface{}, code []byte, config *Config) int {
	if obj == nil {
		code[0], code[1] = TypeNull, Terminator
		return 2
	}

	switch value := obj.(type) {
	case bool:
		if value {
			code[0] = TypeTrue
		} else {
			code[0] = TypeFalse
		}
		code[1] = Terminator
		return 2

	case float64:
		n := 0
		code[n] = TypeNumber
		n++
		n += normalizeFloat(obj, code[n:], config.nk)
		code[n] = Terminator
		n++
		return n

	case int64:
		n := 0
		code[n] = TypeNumber
		n++
		n += normalizeFloat(obj, code[n:], config.nk)
		code[n] = Terminator
		n++
		return n

	case int:
		n := 0
		code[n] = TypeNumber
		n++
		n += normalizeFloat(int64(value), code[n:], config.nk)
		code[n] = Terminator
		n++
		return n

	case json.Number:
		n := 0
		code[n] = TypeNumber
		n++
		f, err := strconv.ParseFloat(string(value), 64)
		if err != nil {
			panic(err)
		}
		n += normalizeFloat(f, code[n:], config.nk)
		code[n] = Terminator
		n++
		return n

	case Length:
		n := 0
		code[n] = TypeLength
		n++
		n += normalizeFloat(int64(value), code[n:], IntNumber)
		code[n] = Terminator
		n++
		return n

	case Missing:
		if config.doMissing && MissingLiteral.Equal(string(value)) {
			code[0], code[1] = TypeMissing, Terminator
			return 2
		}
		panic("collate missing not configured")

	case string:
		if config.doMissing && MissingLiteral.Equal(value) {
			code[0], code[1] = TypeMissing, Terminator
			return 2
		}
		n := 0
		code[n] = TypeString
		n++
		n += suffixEncodeString(str2bytes(value), code[n:])
		code[n] = Terminator
		n++
		return n

	case []byte:
		n := 0
		code[n] = TypeBinary
		n++
		m := copy(code[n:], value)
		n += m
		code[n] = Terminator
		n++
		return n

	case []interface{}:
		n := 0
		code[n] = TypeArray
		n++
		if config.arrayLenPrefix {
			n += gson2collate(Length(len(value)), code[n:], config)
		}
		for _, val := range value {
			n += gson2collate(val, code[n:], config)
		}
		code[n] = Terminator
		n++
		return n

	case map[string]interface{}:
		n := 0
		code[n] = TypeObj
		n++
		if config.propertyLenPrefix {
			n += gson2collate(Length(len(value)), code[n:], config)
		}

		keys := config.pools.keysPool.Get().([]string)
		defer config.pools.keysPool.Put(keys[:0])
		for _, key := range sortProps(value, keys) {
			n += gson2collate(key, code[n:], config)        // encode key
			n += gson2collate(value[key], code[n:], config) // encode value
		}
		code[n] = Terminator
		n++
		return n
	}
	panic(fmt.Errorf("collate invalid golang type %T", obj))
}

func collate2gson(code []byte, config *Config) (interface{}, int) {
	if len(code) == 0 {
		return nil, 0
	}

	var scratch [64]byte
	n := 1
	switch code[0] {
	case TypeMissing:
		return MissingLiteral, 2

	case TypeNull:
		return nil, 2

	case TypeTrue:
		return true, 2

	case TypeFalse:
		return false, 2

	case TypeLength:
		m := getDatum(code[n:])
		_, y := collated2Int(code[n:n+m-1], scratch[:]) // -1 to skip terminator
		length, err := strconv.Atoi(bytes2str(scratch[:y]))
		if err != nil {
			panic(err)
		}
		return Length(length), n + m

	case TypeNumber:
		m := getDatum(code[n:])
		f := denormalizeFloat(code[n:n+m-1], config.nk) // -1 to skip terminator
		return f, n + m

	case TypeString:
		s := make([]byte, encodedStringSize(code[n:]))
		x, y := suffixDecodeString(code[n:], s)
		return bytes2str(s[:y]), n + x

	case TypeBinary:
		m := getDatum(code[n:])
		bs := make([]byte, m-1)
		copy(bs, code[n:n+m-1])
		return bs, n + m

	case TypeArray:
		var arr []interface{}
		if config.arrayLenPrefix {
			if code[n] != TypeLength {
				panic("collate decode expected array length prefix")
			}
			n++
			m := getDatum(code[n:])
			_, y := collated2Int(code[n:n+m], scratch[:])
			ln, err := strconv.Atoi(bytes2str(scratch[:y]))
			if err != nil {
				panic(err)
			}
			arr = make([]interface{}, 0, ln)
			n += m
		} else {
			arr = make([]interface{}, 0, 8)
		}
		for code[n] != Terminator {
			item, y := collate2gson(code[n:], config)
			arr = append(arr, item)
			n += y
		}
		return arr, n + 1 // +1 to skip terminator

	case TypeObj:
		obj := make(map[string]interface{})
		if config.propertyLenPrefix {
			if code[n] != TypeLength {
				panic("collate decode expected object length prefix")
			}
			n++
			m := getDatum(code[n:])
			_, y := collated2Int(code[n:n+m], scratch[:])
			_, err := strconv.Atoi(bytes2str(scratch[:y])) // just skip
			if err != nil {
				panic(err)
			}
			n += m
		}
		for code[n] != Terminator {
			key, m := collate2gson(code[n:], config)
			n += m
			value, m := collate2gson(code[n:], config)
			obj[key.(string)] = value
			n += m
		}
		return obj, n + 1 // +1 to skip terminator
	}
	panic("collate decode invalid binary")
}

// get the collated datum based on Terminator and return the length
// of the datum.
func getDatum(code []byte) int {
	var i int
	var b byte
	for i, b = range code {
		if b == Terminator {
			return i + 1
		}
	}
	panic("collate decode terminator not found")
}
<commit_msg>bug fix: gson2collate() to support uint64 and float32.<commit_after>//  Copyright (c) 2015 Couchbase, Inc.

package gson

import "strconv"
import "fmt"
import "encoding/json"

func gson2collate(obj interface{}, code []byte, config *Config) int {
	if obj == nil {
		code[0], code[1] = TypeNull, Terminator
		return 2
	}

	switch value := obj.(type) {
	case bool:
		if value {
			code[0] = TypeTrue
		} else {
			code[0] = TypeFalse
		}
		code[1] = Terminator
		return 2

	case float64:
		n := 0
		code[n] = TypeNumber
		n++
		n += normalizeFloat(obj, code[n:], config.nk)
		code[n] = Terminator
		n++
		return n

	case float32:
		n := 0
		code[n] = TypeNumber
		n++
		n += normalizeFloat(float64(value), code[n:], config.nk)
		code[n] = Terminator
		n++
		return n

	case int64:
		n := 0
		code[n] = TypeNumber
		n++
		n += normalizeFloat(obj, code[n:], config.nk)
		code[n] = Terminator
		n++
		return n

	case uint64:
		n := 0
		code[n] = TypeNumber
		n++
		n += normalizeFloat(int64(value), code[n:], config.nk)
		code[n] = Terminator
		n++
		return n

	case int:
		n := 0
		code[n] = TypeNumber
		n++
		n += normalizeFloat(int64(value), code[n:], config.nk)
		code[n] = Terminator
		n++
		return n

	case json.Number:
		n := 0
		code[n] = TypeNumber
		n++
		f, err := strconv.ParseFloat(string(value), 64)
		if err != nil {
			panic(err)
		}
		n += normalizeFloat(f, code[n:], config.nk)
		code[n] = Terminator
		n++
		return n

	case Length:
		n := 0
		code[n] = TypeLength
		n++
		n += normalizeFloat(int64(value), code[n:], IntNumber)
		code[n] = Terminator
		n++
		return n

	case Missing:
		if config.doMissing && MissingLiteral.Equal(string(value)) {
			code[0], code[1] = TypeMissing, Terminator
			return 2
		}
		panic("collate missing not configured")

	case string:
		if config.doMissing && MissingLiteral.Equal(value) {
			code[0], code[1] = TypeMissing, Terminator
			return 2
		}
		n := 0
		code[n] = TypeString
		n++
		n += suffixEncodeString(str2bytes(value), code[n:])
		code[n] = Terminator
		n++
		return n

	case []byte:
		n := 0
		code[n] = TypeBinary
		n++
		m := copy(code[n:], value)
		n += m
		code[n] = Terminator
		n++
		return n

	case []interface{}:
		n := 0
		code[n] = TypeArray
		n++
		if config.arrayLenPrefix {
			n += gson2collate(Length(len(value)), code[n:], config)
		}
		for _, val := range value {
			n += gson2collate(val, code[n:], config)
		}
		code[n] = Terminator
		n++
		return n

	case map[string]interface{}:
		n := 0
		code[n] = TypeObj
		n++
		if config.propertyLenPrefix {
			n += gson2collate(Length(len(value)), code[n:], config)
		}

		keys := config.pools.keysPool.Get().([]string)
		defer config.pools.keysPool.Put(keys[:0])
		for _, key := range sortProps(value, keys) {
			n += gson2collate(key, code[n:], config)        // encode key
			n += gson2collate(value[key], code[n:], config) // encode value
		}
		code[n] = Terminator
		n++
		return n
	}
	panic(fmt.Errorf("collate invalid golang type %T", obj))
}

func collate2gson(code []byte, config *Config) (interface{}, int) {
	if len(code) == 0 {
		return nil, 0
	}

	var scratch [64]byte
	n := 1
	switch code[0] {
	case TypeMissing:
		return MissingLiteral, 2

	case TypeNull:
		return nil, 2

	case TypeTrue:
		return true, 2

	case TypeFalse:
		return false, 2

	case TypeLength:
		m := getDatum(code[n:])
		_, y := collated2Int(code[n:n+m-1], scratch[:]) // -1 to skip terminator
		length, err := strconv.Atoi(bytes2str(scratch[:y]))
		if err != nil {
			panic(err)
		}
		return Length(length), n + m

	case TypeNumber:
		m := getDatum(code[n:])
		f := denormalizeFloat(code[n:n+m-1], config.nk) // -1 to skip terminator
		return f, n + m

	case TypeString:
		s := make([]byte, encodedStringSize(code[n:]))
		x, y := suffixDecodeString(code[n:], s)
		return bytes2str(s[:y]), n + x

	case TypeBinary:
		m := getDatum(code[n:])
		bs := make([]byte, m-1)
		copy(bs, code[n:n+m-1])
		return bs, n + m

	case TypeArray:
		var arr []interface{}
		if config.arrayLenPrefix {
			if code[n] != TypeLength {
				panic("collate decode expected array length prefix")
			}
			n++
			m := getDatum(code[n:])
			_, y := collated2Int(code[n:n+m], scratch[:])
			ln, err := strconv.Atoi(bytes2str(scratch[:y]))
			if err != nil {
				panic(err)
			}
			arr = make([]interface{}, 0, ln)
			n += m
		} else {
			arr = make([]interface{}, 0, 8)
		}
		for code[n] != Terminator {
			item, y := collate2gson(code[n:], config)
			arr = append(arr, item)
			n += y
		}
		return arr, n + 1 // +1 to skip terminator

	case TypeObj:
		obj := make(map[string]interface{})
		if config.propertyLenPrefix {
			if code[n] != TypeLength {
				panic("collate decode expected object length prefix")
			}
			n++
			m := getDatum(code[n:])
			_, y := collated2Int(code[n:n+m], scratch[:])
			_, err := strconv.Atoi(bytes2str(scratch[:y])) // just skip
			if err != nil {
				panic(err)
			}
			n += m
		}
		for code[n] != Terminator {
			key, m := collate2gson(code[n:], config)
			n += m
			value, m := collate2gson(code[n:], config)
			obj[key.(string)] = value
			n += m
		}
		return obj, n + 1 // +1 to skip terminator
	}
	panic("collate decode invalid binary")
}

// get the collated datum based on Terminator and return the length
// of the datum.
func getDatum(code []byte) int {
	var i int
	var b byte
	for i, b = range code {
		if b == Terminator {
			return i + 1
		}
	}
	panic("collate decode terminator not found")
}
<|endoftext|>


<commit_before>package collect

import (
	"bytes"
	"compress/gzip"
	"encoding/json"
	"io/ioutil"
	"net/http"
	"time"

	"github.com/StackExchange/slog"
	"github.com/bosun-monitor/opentsdb"
)

func queuer() {
	for dp := range tchan {
		qlock.Lock()
		for {
			if len(queue) > MaxQueueLen {
				slock.Lock()
				dropped++
				slock.Unlock()
				break
			}
			m, err := json.Marshal(dp)
			if err != nil {
				slog.Error(err)
			} else {
				queue = append(queue, m)
			}
			select {
			case dp = <-tchan:
				continue
			default:
			}
			break
		}
		qlock.Unlock()
	}
}

func send() {
	for {
		qlock.Lock()
		if i := len(queue); i > 0 {
			if i > BatchSize {
				i = BatchSize
			}
			sending := queue[:i]
			queue = queue[i:]
			if Debug {
				slog.Infof("sending: %d, remaining: %d", i, len(queue))
			}
			qlock.Unlock()
			sendBatch(sending)
		} else {
			qlock.Unlock()
			time.Sleep(time.Second)
		}
	}
}

func sendBatch(batch []json.RawMessage) {
	if Print {
		for _, d := range batch {
			slog.Info(string(d))
		}
		recordSent(len(batch))
		return
	}
	var buf bytes.Buffer
	g := gzip.NewWriter(&buf)
	if err := json.NewEncoder(g).Encode(batch); err != nil {
		slog.Error(err)
		return
	}
	if err := g.Close(); err != nil {
		slog.Error(err)
		return
	}
	req, err := http.NewRequest("POST", tsdbURL, &buf)
	if err != nil {
		slog.Error(err)
		return
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Content-Encoding", "gzip")
	resp, err := client.Do(req)
	if err == nil {
		defer resp.Body.Close()
	}
	// Some problem with connecting to the server; retry later.
	if err != nil || resp.StatusCode != http.StatusNoContent {
		if err != nil {
			slog.Error(err)
		} else if resp.StatusCode != http.StatusNoContent {
			slog.Errorln(resp.Status)
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				slog.Error(err)
			}
			if len(body) > 0 {
				slog.Error(string(body))
			}
		}
		restored := 0
		for _, msg := range batch {
			var dp opentsdb.DataPoint
			if err := json.Unmarshal(msg, &dp); err != nil {
				slog.Error(err)
				continue
			}
			restored++
			tchan <- &dp
		}
		d := time.Second * 5
		slog.Infof("restored %d, sleeping %s", restored, d)
		time.Sleep(d)
		return
	}
	recordSent(len(batch))
}

func recordSent(num int) {
	if Debug {
		slog.Infoln("sent", num)
	}
	slock.Lock()
	sent += int64(num)
	slock.Unlock()
}
<commit_msg>collect: Merge pull request #1 from bosun-monitor/inst_post<commit_after>package collect

import (
	"bytes"
	"compress/gzip"
	"encoding/json"
	"io/ioutil"
	"net/http"
	"time"

	"github.com/StackExchange/slog"
	"github.com/bosun-monitor/opentsdb"
)

func queuer() {
	for dp := range tchan {
		qlock.Lock()
		for {
			if len(queue) > MaxQueueLen {
				slock.Lock()
				dropped++
				slock.Unlock()
				break
			}
			m, err := json.Marshal(dp)
			if err != nil {
				slog.Error(err)
			} else {
				queue = append(queue, m)
			}
			select {
			case dp = <-tchan:
				continue
			default:
			}
			break
		}
		qlock.Unlock()
	}
}

func send() {
	for {
		qlock.Lock()
		if i := len(queue); i > 0 {
			if i > BatchSize {
				i = BatchSize
			}
			sending := queue[:i]
			queue = queue[i:]
			if Debug {
				slog.Infof("sending: %d, remaining: %d", i, len(queue))
			}
			qlock.Unlock()
			sendBatch(sending)
		} else {
			qlock.Unlock()
			time.Sleep(time.Second)
		}
	}
}

func sendBatch(batch []json.RawMessage) {
	if Print {
		for _, d := range batch {
			slog.Info(string(d))
		}
		recordSent(len(batch))
		return
	}
	var buf bytes.Buffer
	g := gzip.NewWriter(&buf)
	if err := json.NewEncoder(g).Encode(batch); err != nil {
		slog.Error(err)
		return
	}
	if err := g.Close(); err != nil {
		slog.Error(err)
		return
	}
	req, err := http.NewRequest("POST", tsdbURL, &buf)
	if err != nil {
		slog.Error(err)
		return
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Content-Encoding", "gzip")
	now := time.Now()
	resp, err := client.Do(req)
	d := time.Since(now).Nanoseconds() / 1e6
	if err == nil {
		defer resp.Body.Close()
	}
	Add("collect.post.total_duration", nil, d)
	Add("collect.post.count", nil, 1)
	// Some problem with connecting to the server; retry later.
	if err != nil || resp.StatusCode != http.StatusNoContent {
		if err != nil {
			slog.Error(err)
		} else if resp.StatusCode != http.StatusNoContent {
			slog.Errorln(resp.Status)
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				slog.Error(err)
			}
			if len(body) > 0 {
				slog.Error(string(body))
			}
		}
		restored := 0
		for _, msg := range batch {
			var dp opentsdb.DataPoint
			if err := json.Unmarshal(msg, &dp); err != nil {
				slog.Error(err)
				continue
			}
			restored++
			tchan <- &dp
		}
		d := time.Second * 5
		slog.Infof("restored %d, sleeping %s", restored, d)
		time.Sleep(d)
		return
	}
	recordSent(len(batch))
}

func recordSent(num int) {
	if Debug {
		slog.Infoln("sent", num)
	}
	slock.Lock()
	sent += int64(num)
	slock.Unlock()
}
<|endoftext|>


<commit_before>package commands

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/url"
	"os"

	"github.com/codegangsta/cli"

	"github.com/yuuki1/droot/aws"
	"github.com/yuuki1/droot/log"
	"github.com/yuuki1/droot/osutil"
)

var CommandArgPull = "--dest DESTINATION_DIRECTORY --src S3_ENDPOINT"
var CommandPull = cli.Command{
	Name:   "pull",
	Usage:  "Pull an extracted docker image from s3",
	Action: fatalOnError(doPull),
	Flags: []cli.Flag{
		cli.StringFlag{Name: "dest, d", Usage: "Local filesystem path (ex. /var/containers/app)"},
		cli.StringFlag{Name: "src, s", Usage: "Amazon S3 endpoint (ex. s3://example.com/containers/app.tar.gz)"},
	},
}

func doPull(c *cli.Context) error {
	destDir := c.String("dest")
	srcURL := c.String("src")
	if destDir == "" || srcURL == "" {
		cli.ShowCommandHelp(c, "pull")
		return errors.New("--src and --dest option required ")
	}

	s3URL, err := url.Parse(srcURL)
	if err != nil {
		return err
	}
	if s3URL.Scheme != "s3" {
		return fmt.Errorf("Not s3 scheme %s", srcURL)
	}

	if !osutil.ExistsDir(destDir) {
		return fmt.Errorf("No such directory %s", destDir)
	}

	tmp, err := ioutil.TempFile(os.TempDir(), "droot_gzip")
	if err != nil {
		return err
	}
	defer tmp.Close()
	defer os.Remove(tmp.Name())

	nBytes, err := aws.NewS3Client().Download(s3URL, tmp)
	if err != nil {
		return err
	}
	log.Info("downloaded", "from", s3URL, "to", tmp.Name(), nBytes, "bytes")

	rawDir, err := ioutil.TempDir(os.TempDir(), "droot_raw")
	if err != nil {
		return err
	}
	defer os.RemoveAll(rawDir)

	cwd, _ := os.Getwd()
	if err := os.Chdir(rawDir); err != nil {
		return err
	}

	log.Info("Extract archive:", tmp.Name(), "to", rawDir)
	if err := osutil.ExtractTarGz(tmp.Name()); err != nil {
		return err
	}

	log.Info("Sync:", "from", rawDir, "to", destDir)
	if err := osutil.Rsync(rawDir, destDir); err != nil {
		return err
	}

	if err = os.Chdir(cwd); err != nil {
		return err
	}

	return nil
}
<commit_msg>Fix defer process because rm calls close<commit_after>package commands

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/url"
	"os"

	"github.com/codegangsta/cli"

	"github.com/yuuki1/droot/aws"
	"github.com/yuuki1/droot/log"
	"github.com/yuuki1/droot/osutil"
)

var CommandArgPull = "--dest DESTINATION_DIRECTORY --src S3_ENDPOINT"
var CommandPull = cli.Command{
	Name:   "pull",
	Usage:  "Pull an extracted docker image from s3",
	Action: fatalOnError(doPull),
	Flags: []cli.Flag{
		cli.StringFlag{Name: "dest, d", Usage: "Local filesystem path (ex. /var/containers/app)"},
		cli.StringFlag{Name: "src, s", Usage: "Amazon S3 endpoint (ex. s3://example.com/containers/app.tar.gz)"},
	},
}

func doPull(c *cli.Context) error {
	destDir := c.String("dest")
	srcURL := c.String("src")
	if destDir == "" || srcURL == "" {
		cli.ShowCommandHelp(c, "pull")
		return errors.New("--src and --dest option required ")
	}

	s3URL, err := url.Parse(srcURL)
	if err != nil {
		return err
	}
	if s3URL.Scheme != "s3" {
		return fmt.Errorf("Not s3 scheme %s", srcURL)
	}

	if !osutil.ExistsDir(destDir) {
		return fmt.Errorf("No such directory %s", destDir)
	}

	tmp, err := ioutil.TempFile(os.TempDir(), "droot_gzip")
	if err != nil {
		return err
	}
	defer func(f *os.File){
		f.Close()
		os.Remove(f.Name())
	}(tmp)

	nBytes, err := aws.NewS3Client().Download(s3URL, tmp)
	if err != nil {
		return err
	}
	log.Info("downloaded", "from", s3URL, "to", tmp.Name(), nBytes, "bytes")

	rawDir, err := ioutil.TempDir(os.TempDir(), "droot_raw")
	if err != nil {
		return err
	}
	defer os.RemoveAll(rawDir)

	cwd, _ := os.Getwd()
	if err := os.Chdir(rawDir); err != nil {
		return err
	}

	log.Info("Extract archive:", tmp.Name(), "to", rawDir)
	if err := osutil.ExtractTarGz(tmp.Name()); err != nil {
		return err
	}

	log.Info("Sync:", "from", rawDir, "to", destDir)
	if err := osutil.Rsync(rawDir, destDir); err != nil {
		return err
	}

	if err = os.Chdir(cwd); err != nil {
		return err
	}

	return nil
}
<|endoftext|>