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

// assumptions
// we've been booted into a ramfs with all this stuff unpacked and ready.
// we don't need a loop device mount because it's all there.
// So we run /go/bin/go build installcommand
// and then exec /buildbin/sh

package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"syscall"
	"uroot"
)

const PATH = "/bin:/buildbin:/usr/local/bin"

type dir struct {
	name string
	mode   os.FileMode
}

type mount struct {
	source string
	target string
	fstype string
	flags  uintptr
	opts   string
}

var (
	env = map[string]string{
		"LD_LIBRARY_PATH": "/usr/local/lib",
		"GOROOT":          "/go",
		"GOPATH":          "/",
		"CGO_ENABLED":     "0",
	}

	dirs = []dir {
		{name: "/proc", mode: os.FileMode(0555),},
		{name: "/buildbin", mode: os.FileMode(0777),},
		{name: "/bin", mode: os.FileMode(0777),},
		{name: "/tmp", mode: os.FileMode(0777),},
		{name: "/env", mode: os.FileMode(0777),},
		{name: "/etc", mode: os.FileMode(0777),},
		{name: "/tcz", mode: os.FileMode(0777),},
		{name: "/dev", mode: os.FileMode(0777),},
		{name: "/lib", mode: os.FileMode(0777),},
		{name: "/usr/lib", mode: os.FileMode(0777),},
		{name: "/go/pkg/linux_amd64", mode: os.FileMode(0777),},
	}
	namespace = []mount{
		{source: "proc", target: "/proc", fstype: "proc", flags: syscall.MS_MGC_VAL | syscall.MS_RDONLY, opts: "",},
	}
)

func main() {
	log.Printf("Welcome to u-root")
	// Pick some reasonable values in the (unlikely!) even that Uname fails.
	uname := "linux"
	mach := "x86_64"
	// There are two possible places for go:
	// The first is in /go/bin
	// The second is in /go/pkg/tool/$OS_$ARCH
	if u, err := uroot.Uname(); err != nil {
		log.Printf("uroot.Utsname fails: %v, so assume %v_%v\n", uname, mach)
	} else {
		// Sadly, go and the OS disagree on case.
		uname = strings.ToLower(u.Sysname)
	}
	env["PATH"] = fmt.Sprintf("/go/bin:/go/pkg/tool/%s_%s:%v", uname, mach, PATH)
	envs := []string{}
	for k, v := range env {
		os.Setenv(k, v)
		envs = append(envs, k+"="+v)
	}

	for _, m := range dirs {
		if err := os.MkdirAll(m.name, m.mode); err != nil {
			log.Printf("mkdir :%s: mode %o: %v\n", m.name, m.mode)
			continue
		}
	}

	for _, m := range namespace {
		if err := syscall.Mount(m.source, m.target, m.fstype, m.flags, m.opts); err != nil {
			log.Printf("Mount :%s: on :%s: type :%s: flags %x: %v\n", m.source, m.target, m.fstype, m.flags, m.opts, err)
		}

	}
//HACK
	if err := filepath.Walk("/", func(name string, fi os.FileInfo, err error) error {
			if err != nil {
				fmt.Printf(" WALK FAIL%v: %v\n", name, err)
				// That's ok, sometimes things are not there.
				return nil
			}
			fmt.Printf("%v\n", name)
			return nil
	}); err != nil {
		log.Printf("WALK fails %v\n", err)
	}
// HACK

	// populate buildbin

	if commands, err := ioutil.ReadDir("/src/cmds"); err == nil {
		for _, v := range commands {
			name := v.Name()
			if name == "installcommand" || name == "init" {
				continue
			} else {
				destPath := path.Join("/buildbin", name)
				source := "/buildbin/installcommand"
				if err := os.Symlink(source, destPath); err != nil {
					log.Printf("Symlink %v -> %v failed; %v", source, destPath, err)
				}
			}
		}
	} else {
		log.Fatalf("Can't read %v; %v", "/src", err)
	}
	log.Printf("envs %v", envs)
	os.Setenv("GOBIN", "/buildbin")
	cmd := exec.Command("go", "install", "-x", path.Join("cmds", "installcommand"))
	installenvs := envs
	installenvs = append(envs, "GOBIN=/buildbin")
	cmd.Env = installenvs
	cmd.Dir = "/"

	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	log.Printf("Run %v", cmd)
	err := cmd.Run()
	if err != nil {
		log.Printf("%v\n", err)
	}

	// install /env.
	os.Setenv("GOBIN", "/bin")
	envs = append(envs, "GOBIN=/bin")
	for _, e := range envs {
		nv := strings.SplitN(e, "=", 2)
		if len(nv) < 2 {
			nv = append(nv, "")
		}
		n := path.Join("/env", nv[0])
		if err := ioutil.WriteFile(n, []byte(nv[1]), 0666); err != nil {
			log.Printf("%v: %v", n, err)
		}
	}

	cmd = exec.Command("/buildbin/sh")
	cmd.Env = envs
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	// TODO: figure out why we get EPERM when we use this.
	//cmd.SysProcAttr = &syscall.SysProcAttr{Setctty: true, Setsid: true,}
	log.Printf("Run %v", cmd)
	err = cmd.Run()
	if err != nil {
		log.Printf("%v\n", err)
	}
	log.Printf("init: /bin/sh returned!\n")
}
<commit_msg>Create device nodes<commit_after>// Copyright 2012 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// assumptions
// we've been booted into a ramfs with all this stuff unpacked and ready.
// we don't need a loop device mount because it's all there.
// So we run /go/bin/go build installcommand
// and then exec /buildbin/sh

package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"syscall"
	"uroot"
)

const PATH = "/bin:/buildbin:/usr/local/bin"

type dir struct {
	name string
	mode   os.FileMode
}

type dev struct {
	name string
	mode   os.FileMode
	magic  int
}

type mount struct {
	source string
	target string
	fstype string
	flags  uintptr
	opts   string
}

var (
	env = map[string]string{
		"LD_LIBRARY_PATH": "/usr/local/lib",
		"GOROOT":          "/go",
		"GOPATH":          "/",
		"CGO_ENABLED":     "0",
	}

	dirs = []dir {
		{name: "/proc", mode: os.FileMode(0555),},
		{name: "/buildbin", mode: os.FileMode(0777),},
		{name: "/bin", mode: os.FileMode(0777),},
		{name: "/tmp", mode: os.FileMode(0777),},
		{name: "/env", mode: os.FileMode(0777),},
		{name: "/etc", mode: os.FileMode(0777),},
		{name: "/tcz", mode: os.FileMode(0777),},
		{name: "/dev", mode: os.FileMode(0777),},
		{name: "/lib", mode: os.FileMode(0777),},
		{name: "/usr/lib", mode: os.FileMode(0777),},
		{name: "/go/pkg/linux_amd64", mode: os.FileMode(0777),},
	}
	devs = []dev {
		{name: "/dev/null", mode: os.FileMode(0660)|os.ModeDevice|os.ModeCharDevice, magic: 0x0301,},
	}
	namespace = []mount{
		{source: "proc", target: "/proc", fstype: "proc", flags: syscall.MS_MGC_VAL | syscall.MS_RDONLY, opts: "",},
	}
)

func main() {
	log.Printf("Welcome to u-root")
	// Pick some reasonable values in the (unlikely!) even that Uname fails.
	uname := "linux"
	mach := "x86_64"
	// There are two possible places for go:
	// The first is in /go/bin
	// The second is in /go/pkg/tool/$OS_$ARCH
	if u, err := uroot.Uname(); err != nil {
		log.Printf("uroot.Utsname fails: %v, so assume %v_%v\n", uname, mach)
	} else {
		// Sadly, go and the OS disagree on case.
		uname = strings.ToLower(u.Sysname)
	}
	env["PATH"] = fmt.Sprintf("/go/bin:/go/pkg/tool/%s_%s:%v", uname, mach, PATH)
	envs := []string{}
	for k, v := range env {
		os.Setenv(k, v)
		envs = append(envs, k+"="+v)
	}

	for _, m := range dirs {
		if err := os.MkdirAll(m.name, m.mode); err != nil {
			log.Printf("mkdir :%s: mode %o: %v\n", m.name, m.mode, err)
			continue
		}
	}

	for _, d := range devs {
		if err := syscall.Mknod(d.name, uint32(d.mode), d.magic); err != nil {
			log.Printf("mknod :%s: mode %o: magic: %v: %v\n", d.name, d.mode, d.magic, err)
			continue
		}
	}

	for _, m := range namespace {
		if err := syscall.Mount(m.source, m.target, m.fstype, m.flags, m.opts); err != nil {
			log.Printf("Mount :%s: on :%s: type :%s: flags %x: %v\n", m.source, m.target, m.fstype, m.flags, m.opts, err)
		}

	}

	// only in case of emergency.
	if false {
		if err := filepath.Walk("/", func(name string, fi os.FileInfo, err error) error {
			if err != nil {
				fmt.Printf(" WALK FAIL%v: %v\n", name, err)
				// That's ok, sometimes things are not there.
				return nil
			}
			fmt.Printf("%v\n", name)
			return nil
		}); err != nil {
			log.Printf("WALK fails %v\n", err)
		}
	}

	// populate buildbin

	if commands, err := ioutil.ReadDir("/src/cmds"); err == nil {
		for _, v := range commands {
			name := v.Name()
			if name == "installcommand" || name == "init" {
				continue
			} else {
				destPath := path.Join("/buildbin", name)
				source := "/buildbin/installcommand"
				if err := os.Symlink(source, destPath); err != nil {
					log.Printf("Symlink %v -> %v failed; %v", source, destPath, err)
				}
			}
		}
	} else {
		log.Fatalf("Can't read %v; %v", "/src", err)
	}
	log.Printf("envs %v", envs)
	os.Setenv("GOBIN", "/buildbin")
	cmd := exec.Command("go", "install", "-x", path.Join("cmds", "installcommand"))
	installenvs := envs
	installenvs = append(envs, "GOBIN=/buildbin")
	cmd.Env = installenvs
	cmd.Dir = "/"

	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	log.Printf("Run %v", cmd)
	err := cmd.Run()
	if err != nil {
		log.Printf("%v\n", err)
	}

	// install /env.
	os.Setenv("GOBIN", "/bin")
	envs = append(envs, "GOBIN=/bin")
	for _, e := range envs {
		nv := strings.SplitN(e, "=", 2)
		if len(nv) < 2 {
			nv = append(nv, "")
		}
		n := path.Join("/env", nv[0])
		if err := ioutil.WriteFile(n, []byte(nv[1]), 0666); err != nil {
			log.Printf("%v: %v", n, err)
		}
	}

	cmd = exec.Command("/buildbin/sh")
	cmd.Env = envs
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	// TODO: figure out why we get EPERM when we use this.
	//cmd.SysProcAttr = &syscall.SysProcAttr{Setctty: true, Setsid: true,}
	log.Printf("Run %v", cmd)
	err = cmd.Run()
	if err != nil {
		log.Printf("%v\n", err)
	}
	log.Printf("init: /bin/sh returned!\n")
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"math/rand"
	"net"
	"net/http"
	"os"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"
)

var (
	version   = "1.0.0"
	github    = "https://github.com/ekalinin/awsping"
	useragent = fmt.Sprintf("AwsPing/%s (+%s)", version, github)
)

var (
	repeats = flag.Int("repeats", 1, "Number of repeats")
	useHTTP = flag.Bool("http", false, "Use http transport (default is tcp)")
	showVer = flag.Bool("v", false, "Show version")
	verbose = flag.Int("verbose", 0, "Verbosity level")
	service = flag.String("service", "dynamodb", "AWS Service: ec2, sdb, sns, sqs, ...")
)

var letterRunes = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")

// Duration2ms converts time.Duration to ms (float64)
func Duration2ms(d time.Duration) float64 {
	return float64(d.Nanoseconds()) / 1000 / 1000
}

func mkRandoString(n int) string {
	b := make([]rune, n)
	for i := range b {
		b[i] = letterRunes[rand.Intn(len(letterRunes))]
	}
	return string(b)
}

// AWSRegion description of the AWS EC2 region
type AWSRegion struct {
	Name      string
	Code      string
	Service   string
	Latencies []time.Duration
	Error     error
}

// CheckLatencyHTTP Test Latency via HTTP
func (r *AWSRegion) CheckLatencyHTTP(wg *sync.WaitGroup) {
	defer wg.Done()
	url := fmt.Sprintf("http://%s.%s.amazonaws.com/ping?x=%s", r.Service,
		r.Code, mkRandoString(13))
	client := &http.Client{}

	req, err := http.NewRequest("GET", url, nil)
	req.Header.Set("User-Agent", useragent)

	start := time.Now()
	resp, err := client.Do(req)
	r.Latencies = append(r.Latencies, time.Since(start))
	defer resp.Body.Close()

	r.Error = err
}

// CheckLatencyTCP Test Latency via TCP
func (r *AWSRegion) CheckLatencyTCP(wg *sync.WaitGroup) {
	defer wg.Done()
	tcpURI := fmt.Sprintf("%s.%s.amazonaws.com:80", r.Service, r.Code)
	tcpAddr, err := net.ResolveTCPAddr("tcp4", tcpURI)
	if err != nil {
		r.Error = err
		return
	}
	start := time.Now()
	conn, err := net.DialTCP("tcp", nil, tcpAddr)
	if err != nil {
		r.Error = err
		return
	}
	r.Latencies = append(r.Latencies, time.Since(start))
	defer conn.Close()

	r.Error = err
}

// GetLatency returns Latency in ms
func (r *AWSRegion) GetLatency() float64 {
	sum := float64(0)
	for _, l := range r.Latencies {
		sum += Duration2ms(l)
	}
	return sum / float64(len(r.Latencies))
}

// GetLatencyStr returns Latency in string
func (r *AWSRegion) GetLatencyStr() string {
	if r.Error != nil {
		return r.Error.Error()
	}
	return fmt.Sprintf("%.2f ms", r.GetLatency())
}

// AWSRegions slice of the AWSRegion
type AWSRegions []AWSRegion

func (rs AWSRegions) Len() int {
	return len(rs)
}

func (rs AWSRegions) Less(i, j int) bool {
	return rs[i].GetLatency() < rs[j].GetLatency()
}

func (rs AWSRegions) Swap(i, j int) {
	rs[i], rs[j] = rs[j], rs[i]
}

// CalcLatency returns list of aws regions sorted by Latency
func CalcLatency(repeats int, useHTTP bool, service string) *AWSRegions {
	regions := AWSRegions{
		{Service: service, Name: "US-East (Virginia)", Code: "us-east-1"},
		{Service: service, Name: "US-East (Ohio)", Code: "us-east-2"},
		{Service: service, Name: "US-West (California)", Code: "us-west-1"},
		{Service: service, Name: "US-West (Oregon)", Code: "us-west-2"},
		{Service: service, Name: "Canada (Central)", Code: "ca-central-1"},
		{Service: service, Name: "Europe (Ireland)", Code: "eu-west-1"},
		{Service: service, Name: "Europe (Frankfurt)", Code: "eu-central-1"},
		{Service: service, Name: "Europe (London)", Code: "eu-west-2"},
		{Service: service, Name: "Europe (Milan)", Code: "eu-south-1"},
		{Service: service, Name: "Europe (Paris)", Code: "eu-west-3"},
		{Service: service, Name: "Europe (Stockholm)", Code: "eu-north-1"},
		{Service: service, Name: "Africa (Cape Town)", Code: "af-south-1"},
		{Service: service, Name: "Asia Pacific (Osaka-Local)", Code: "ap-northeast-3"},
		{Service: service, Name: "Asia Pacific (Hong Kong)", Code: "ap-east-1"},
		{Service: service, Name: "Asia Pacific (Tokyo)", Code: "ap-northeast-1"},
		{Service: service, Name: "Asia Pacific (Seoul)", Code: "ap-northeast-2"},
		{Service: service, Name: "Asia Pacific (Singapore)", Code: "ap-southeast-1"},
		{Service: service, Name: "Asia Pacific (Mumbai)", Code: "ap-south-1"},
		{Service: service, Name: "Asia Pacific (Sydney)", Code: "ap-southeast-2"},
		{Service: service, Name: "South America (São Paulo)", Code: "sa-east-1"},
		{Service: service, Name: "Middle East (Bahrain)", Code: "me-south-1"},
		{Service: service, Name: "South America (São Paulo)", Code: "sa-east-1"},
	}
	var wg sync.WaitGroup

	for n := 1; n <= repeats; n++ {
		wg.Add(len(regions))
		for i := range regions {
			if useHTTP {
				go regions[i].CheckLatencyHTTP(&wg)
			} else {
				go regions[i].CheckLatencyTCP(&wg)
			}
		}
		wg.Wait()
	}

	sort.Sort(regions)
	return &regions
}

// LatencyOutput prints data into console
type LatencyOutput struct {
	Level int
}

func (lo *LatencyOutput) show0(regions *AWSRegions) {
	for _, r := range *regions {
		fmt.Printf("%-25s %20s\n", r.Name, r.GetLatencyStr())
	}
}

func (lo *LatencyOutput) show1(regions *AWSRegions) {
	outFmt := "%5v %-15s %-30s %20s\n"
	fmt.Printf(outFmt, "", "Code", "Region", "Latency")
	for i, r := range *regions {
		fmt.Printf(outFmt, i, r.Code, r.Name, r.GetLatencyStr())
	}
}

func (lo *LatencyOutput) show2(regions *AWSRegions) {
	// format
	outFmt := "%5v %-15s %-25s"
	outFmt += strings.Repeat(" %15s", *repeats) + " %15s\n"
	// header
	outStr := []interface{}{"", "Code", "Region"}
	for i := 0; i < *repeats; i++ {
		outStr = append(outStr, "Try #"+strconv.Itoa(i+1))
	}
	outStr = append(outStr, "Avg Latency")

	// show header
	fmt.Printf(outFmt, outStr...)

	// each region stats
	for i, r := range *regions {
		outData := []interface{}{strconv.Itoa(i), r.Code, r.Name}
		for n := 0; n < *repeats; n++ {
			outData = append(outData, fmt.Sprintf("%.2f ms",
				Duration2ms(r.Latencies[n])))
		}
		outData = append(outData, fmt.Sprintf("%.2f ms", r.GetLatency()))
		fmt.Printf(outFmt, outData...)
	}
}

// Show print data
func (lo *LatencyOutput) Show(regions *AWSRegions) {
	switch lo.Level {
	case 0:
		lo.show0(regions)
	case 1:
		lo.show1(regions)
	case 2:
		lo.show2(regions)
	}
}

func main() {

	flag.Parse()

	if *showVer {
		fmt.Println(version)
		os.Exit(0)
	}

	regions := CalcLatency(*repeats, *useHTTP, *service)
	lo := LatencyOutput{*verbose}
	lo.Show(regions)
}

func init() {
	rand.Seed(time.Now().UnixNano())
}
<commit_msg>added HTTPS<commit_after>package main

import (
	"flag"
	"fmt"
	"math/rand"
	"net"
	"net/http"
	"os"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"
)

var (
	version   = "1.0.0"
	github    = "https://github.com/ekalinin/awsping"
	useragent = fmt.Sprintf("AwsPing/%s (+%s)", version, github)
)

var (
	repeats  = flag.Int("repeats", 1, "Number of repeats")
	useHTTP  = flag.Bool("http", false, "Use http transport (default is tcp)")
	useHTTPS = flag.Bool("https", false, "Use https transport (default is tcp)")
	showVer  = flag.Bool("v", false, "Show version")
	verbose  = flag.Int("verbose", 0, "Verbosity level")
	service  = flag.String("service", "dynamodb", "AWS Service: ec2, sdb, sns, sqs, ...")
)

var letterRunes = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")

// Duration2ms converts time.Duration to ms (float64)
func Duration2ms(d time.Duration) float64 {
	return float64(d.Nanoseconds()) / 1000 / 1000
}

func mkRandoString(n int) string {
	b := make([]rune, n)
	for i := range b {
		b[i] = letterRunes[rand.Intn(len(letterRunes))]
	}
	return string(b)
}

// AWSRegion description of the AWS EC2 region
type AWSRegion struct {
	Name      string
	Code      string
	Service   string
	Latencies []time.Duration
	Error     error
}

// CheckLatencyHTTP Test Latency via HTTP
func (r *AWSRegion) CheckLatencyHTTP(wg *sync.WaitGroup, https bool) {
	defer wg.Done()
	proto := "http"
	if https {
		proto = "https"
	}

	url := fmt.Sprintf("%s://%s.%s.amazonaws.com/ping?x=%s", proto, r.Service,
		r.Code, mkRandoString(13))
	client := &http.Client{}

	req, err := http.NewRequest("GET", url, nil)
	req.Header.Set("User-Agent", useragent)

	start := time.Now()
	resp, err := client.Do(req)
	r.Latencies = append(r.Latencies, time.Since(start))
	defer resp.Body.Close()

	r.Error = err
}

// CheckLatencyTCP Test Latency via TCP
func (r *AWSRegion) CheckLatencyTCP(wg *sync.WaitGroup) {
	defer wg.Done()
	tcpURI := fmt.Sprintf("%s.%s.amazonaws.com:80", r.Service, r.Code)
	tcpAddr, err := net.ResolveTCPAddr("tcp4", tcpURI)
	if err != nil {
		r.Error = err
		return
	}
	start := time.Now()
	conn, err := net.DialTCP("tcp", nil, tcpAddr)
	if err != nil {
		r.Error = err
		return
	}
	r.Latencies = append(r.Latencies, time.Since(start))
	defer conn.Close()

	r.Error = err
}

// GetLatency returns Latency in ms
func (r *AWSRegion) GetLatency() float64 {
	sum := float64(0)
	for _, l := range r.Latencies {
		sum += Duration2ms(l)
	}
	return sum / float64(len(r.Latencies))
}

// GetLatencyStr returns Latency in string
func (r *AWSRegion) GetLatencyStr() string {
	if r.Error != nil {
		return r.Error.Error()
	}
	return fmt.Sprintf("%.2f ms", r.GetLatency())
}

// AWSRegions slice of the AWSRegion
type AWSRegions []AWSRegion

func (rs AWSRegions) Len() int {
	return len(rs)
}

func (rs AWSRegions) Less(i, j int) bool {
	return rs[i].GetLatency() < rs[j].GetLatency()
}

func (rs AWSRegions) Swap(i, j int) {
	rs[i], rs[j] = rs[j], rs[i]
}

// CalcLatency returns list of aws regions sorted by Latency
func CalcLatency(repeats int, useHTTP bool, useHTTPS bool, service string) *AWSRegions {
	regions := AWSRegions{
		{Service: service, Name: "US-East (Virginia)", Code: "us-east-1"},
		{Service: service, Name: "US-East (Ohio)", Code: "us-east-2"},
		{Service: service, Name: "US-West (California)", Code: "us-west-1"},
		{Service: service, Name: "US-West (Oregon)", Code: "us-west-2"},
		{Service: service, Name: "Canada (Central)", Code: "ca-central-1"},
		{Service: service, Name: "Europe (Ireland)", Code: "eu-west-1"},
		{Service: service, Name: "Europe (Frankfurt)", Code: "eu-central-1"},
		{Service: service, Name: "Europe (London)", Code: "eu-west-2"},
		{Service: service, Name: "Europe (Milan)", Code: "eu-south-1"},
		{Service: service, Name: "Europe (Paris)", Code: "eu-west-3"},
		{Service: service, Name: "Europe (Stockholm)", Code: "eu-north-1"},
		{Service: service, Name: "Africa (Cape Town)", Code: "af-south-1"},
		{Service: service, Name: "Asia Pacific (Osaka-Local)", Code: "ap-northeast-3"},
		{Service: service, Name: "Asia Pacific (Hong Kong)", Code: "ap-east-1"},
		{Service: service, Name: "Asia Pacific (Tokyo)", Code: "ap-northeast-1"},
		{Service: service, Name: "Asia Pacific (Seoul)", Code: "ap-northeast-2"},
		{Service: service, Name: "Asia Pacific (Singapore)", Code: "ap-southeast-1"},
		{Service: service, Name: "Asia Pacific (Mumbai)", Code: "ap-south-1"},
		{Service: service, Name: "Asia Pacific (Sydney)", Code: "ap-southeast-2"},
		{Service: service, Name: "South America (São Paulo)", Code: "sa-east-1"},
		{Service: service, Name: "Middle East (Bahrain)", Code: "me-south-1"},
		{Service: service, Name: "South America (São Paulo)", Code: "sa-east-1"},
	}
	var wg sync.WaitGroup

	for n := 1; n <= repeats; n++ {
		wg.Add(len(regions))
		for i := range regions {
			if useHTTP || useHTTPS {
				go regions[i].CheckLatencyHTTP(&wg, useHTTPS)
			} else {
				go regions[i].CheckLatencyTCP(&wg)
			}
		}
		wg.Wait()
	}

	sort.Sort(regions)
	return &regions
}

// LatencyOutput prints data into console
type LatencyOutput struct {
	Level int
}

func (lo *LatencyOutput) show0(regions *AWSRegions) {
	for _, r := range *regions {
		fmt.Printf("%-25s %20s\n", r.Name, r.GetLatencyStr())
	}
}

func (lo *LatencyOutput) show1(regions *AWSRegions) {
	outFmt := "%5v %-15s %-30s %20s\n"
	fmt.Printf(outFmt, "", "Code", "Region", "Latency")
	for i, r := range *regions {
		fmt.Printf(outFmt, i, r.Code, r.Name, r.GetLatencyStr())
	}
}

func (lo *LatencyOutput) show2(regions *AWSRegions) {
	// format
	outFmt := "%5v %-15s %-25s"
	outFmt += strings.Repeat(" %15s", *repeats) + " %15s\n"
	// header
	outStr := []interface{}{"", "Code", "Region"}
	for i := 0; i < *repeats; i++ {
		outStr = append(outStr, "Try #"+strconv.Itoa(i+1))
	}
	outStr = append(outStr, "Avg Latency")

	// show header
	fmt.Printf(outFmt, outStr...)

	// each region stats
	for i, r := range *regions {
		outData := []interface{}{strconv.Itoa(i), r.Code, r.Name}
		for n := 0; n < *repeats; n++ {
			outData = append(outData, fmt.Sprintf("%.2f ms",
				Duration2ms(r.Latencies[n])))
		}
		outData = append(outData, fmt.Sprintf("%.2f ms", r.GetLatency()))
		fmt.Printf(outFmt, outData...)
	}
}

// Show print data
func (lo *LatencyOutput) Show(regions *AWSRegions) {
	switch lo.Level {
	case 0:
		lo.show0(regions)
	case 1:
		lo.show1(regions)
	case 2:
		lo.show2(regions)
	}
}

func main() {

	flag.Parse()

	if *showVer {
		fmt.Println(version)
		os.Exit(0)
	}

	regions := CalcLatency(*repeats, *useHTTP, *useHTTPS, *service)
	lo := LatencyOutput{*verbose}
	lo.Show(regions)
}

func init() {
	rand.Seed(time.Now().UnixNano())
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"math"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"code.google.com/p/gogoprotobuf/proto"

	pb "github.com/dgryski/carbonzipper/carbonzipperpb"
)

type zipper string

var Zipper zipper

var timeFormats = []string{"15:04 20060102", "20060102", "01/02/06"}

// dateParamToEpoch turns a passed string parameter into an epoch we can send to the Zipper
func dateParamToEpoch(s string, d int64) string {

	if s == "" {
		// return the default if nothing was passed
		return strconv.Itoa(int(d))
	}

	// relative timestamp
	if s[0] == '-' {

		offset, err := intervalString(s[1:])
		if err != nil {
			return strconv.Itoa(int(d))
		}

		return strconv.Itoa(int(timeNow().Add(-time.Duration(offset) * time.Second).Unix()))
	}

	_, err := strconv.Atoi(s)
	if err == nil && len(s) > 8 {
		return s // We got a timestamp so returning it
	}

	if strings.Contains(s, "_") {
		s = strings.Replace(s, "_", " ", 1) // Go can't parse _ in date strings
	}

	for _, format := range timeFormats {
		t, err := time.Parse(format, s)
		if err == nil {
			return strconv.Itoa(int(t.Unix()))
		}
	}
	return strconv.Itoa(int(d))
}

func intervalString(s string) (int32, error) {

	var j int

	for j < len(s) && s[j] >= '0' && s[j] <= '9' {
		j++
	}
	offsetStr, unitStr := s[:j], s[j:]

	var units int
	switch unitStr {
	case "s", "sec", "secs", "second", "seconds":
		units = 1
	case "min", "minute", "minutes":
		units = 60
	case "h", "hour", "hours":
		units = 60 * 60
	case "d", "day", "days":
		units = 24 * 60 * 60
	case "w", "week", "weeks":
		units = 7 * 24 * 60 * 60
	case "mon", "month", "months":
		units = 30 * 24 * 60 * 60
	case "y", "year", "years":
		units = 365 * 24 * 60 * 60
	default:
		return 0, errors.New("unknown time units")
	}

	offset, err := strconv.Atoi(offsetStr)
	if err != nil {
		return 0, err
	}

	return int32(offset * units), nil
}

// FIXME(dgryski): extract the http.Get + unproto code into its own function

func (z zipper) Find(metric string) (pb.GlobResponse, error) {

	u, _ := url.Parse(string(z) + "/metrics/find/")

	u.RawQuery = url.Values{
		"query":  []string{metric},
		"format": []string{"protobuf"},
	}.Encode()

	resp, err := http.Get(u.String())
	if err != nil {
		log.Printf("Find: http.Get: %+v\n", err)
		return pb.GlobResponse{}, err
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Printf("Find: ioutil.ReadAll: %+v\n", err)
		return pb.GlobResponse{}, err
	}

	var pbresp pb.GlobResponse

	err = proto.Unmarshal(body, &pbresp)
	if err != nil {
		log.Printf("Find: proto.Unmarshal: %+v\n", err)
		return pb.GlobResponse{}, err
	}

	return pbresp, nil
}

func (z zipper) Render(metric, from, until string) (pb.FetchResponse, error) {

	u, _ := url.Parse(string(z) + "/render/")

	u.RawQuery = url.Values{
		"target": []string{metric},
		"format": []string{"protobuf"},
		"from":   []string{from},
		"until":  []string{until},
	}.Encode()

	resp, err := http.Get(u.String())
	if err != nil {
		log.Printf("Render: http.Get: %s: %+v\n", metric, err)
		return pb.FetchResponse{}, err
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Printf("Render: ioutil.ReadAll: %s: %+v\n", metric, err)
		return pb.FetchResponse{}, err
	}

	var pbresp pb.FetchResponse

	err = proto.Unmarshal(body, &pbresp)
	if err != nil {
		log.Printf("Render: proto.Unmarshal: %s: %+v\n", metric, err)
		return pb.FetchResponse{}, err
	}

	return pbresp, nil
}

type limiter chan struct{}

func (l limiter) enter() { l <- struct{}{} }
func (l limiter) leave() { <-l }

var Limiter limiter

// for testing
var timeNow = time.Now

type graphitePoint struct {
	value float64
	t     int32
}

func (g graphitePoint) MarshalJSON() ([]byte, error) {
	// TODO(dgryski): fmt.Sprintf() is slow, use strconv.Append{Float,Int}
	if math.IsNaN(g.value) {
		return []byte(fmt.Sprintf("[null,%d]", g.t)), nil
	}
	return []byte(fmt.Sprintf("[%g,%d]", g.value, g.t)), nil
}

type jsonResponse struct {
	Target     string          `json:"target"`
	Datapoints []graphitePoint `json:"datapoints"`
}

func renderHandler(w http.ResponseWriter, r *http.Request) {

	r.ParseForm()
	targets := r.Form["target"]
	from := r.FormValue("from")
	until := r.FormValue("until")

	cacheKey := r.Form.Encode()

	if response, ok := queryCache.get(cacheKey); ok {
		w.Header().Set("Content-Type", "application/json")
		w.Write(response.([]byte))
		return
	}

	// normalize from and until values
	from = dateParamToEpoch(from, timeNow().Add(-24*time.Hour).Unix())
	until = dateParamToEpoch(until, timeNow().Unix())

	var results []*pb.FetchResponse

	for _, target := range targets {

		exp, e, err := parseExpr(target)
		if err != nil || e != "" {
			http.Error(w, http.StatusText(http.StatusBadRequest), http.StatusBadRequest)
			return
		}

		metricMap := make(map[string][]*pb.FetchResponse)

		for _, metric := range exp.metrics() {

			// TODO(dgryski): cache Find results for 1 minute

			// query zipper for find
			glob, err := Zipper.Find(metric)
			if err != nil {
				continue
			}

			// for each server in find response query render
			rch := make(chan *pb.FetchResponse, len(glob.GetMatches()))
			for _, m := range glob.GetMatches() {
				go func(m *pb.GlobMatch) {
					Limiter.enter()
					var rptr *pb.FetchResponse
					if m.GetIsLeaf() {
						r, err := Zipper.Render(m.GetPath(), from, until)
						if err == nil {
							rptr = &r
						}
					}
					rch <- rptr
					Limiter.leave()
				}(m)
			}

			for i := 0; i < len(glob.GetMatches()); i++ {
				r := <-rch
				if r != nil {
					metricMap[metric] = append(metricMap[metric], r)
				}
			}
		}

		exprs := evalExpr(exp, metricMap)
		results = append(results, exprs...)
	}

	var jresults []jsonResponse

	for _, r := range results {
		if r == nil {
			log.Println("skipping nil result")
		}
		datapoints := make([]graphitePoint, 0, len(r.Values))
		t := *r.StartTime
		for i, v := range r.Values {
			if r.IsAbsent[i] {
				v = math.NaN()
			}
			datapoints = append(datapoints, graphitePoint{value: v, t: t})
			t += *r.StepTime
		}
		jresults = append(jresults, jsonResponse{Target: r.GetName(), Datapoints: datapoints})
	}

	jout, err := json.Marshal(jresults)
	if err != nil {
		http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
		return
	}

	queryCache.set(cacheKey, jout)

	w.Header().Set("Content-Type", "application/json")
	w.Write(jout)
}

var queryCache *expireCache

type cacheElement struct {
	validUntil time.Time
	data       interface{}
}

type expireCache struct {
	sync.Mutex
	cache map[string]cacheElement
}

func (ec *expireCache) get(k string) (interface{}, bool) {
	ec.Lock()
	v, ok := ec.cache[k]
	ec.Unlock()
	if !ok || v.validUntil.Before(time.Now()) {
		return nil, false
	}
	return v.data, ok
}

func (ec *expireCache) set(k string, v interface{}) {
	ec.Lock()
	ec.cache[k] = cacheElement{validUntil: time.Now().Add(60 * time.Second), data: v}
	ec.Unlock()
}

func (ec *expireCache) cleaner() {

	var keys []string

	for {
		time.Sleep(5 * time.Minute)

		now := time.Now()
		ec.Lock()

		for k, v := range ec.cache {
			if v.validUntil.Before(now) {
				keys = append(keys, k)
			}
		}

		for _, k := range keys {
			log.Println("cleaning", k)
			delete(ec.cache, k)
		}

		keys = keys[:0]
		ec.Unlock()
	}
}

func main() {

	z := flag.String("z", "", "zipper")
	port := flag.Int("p", 8080, "port")
	l := flag.Int("l", 20, "concurrency limit")

	flag.Parse()

	if *z == "" {
		log.Fatal("no zipper (-z) provided")
	}

	if p := os.Getenv("PORT"); p != "" {
		*port, _ = strconv.Atoi(p)
	}

	Limiter = make(chan struct{}, *l)

	if *z == "" {
		log.Fatal("no zipper provided")
	}

	if _, err := url.Parse(*z); err != nil {
		log.Fatal("unable to parze zipper:", err)
	}

	log.Println("using zipper", *z)

	queryCache = &expireCache{cache: make(map[string]cacheElement)}
	go queryCache.cleaner()

	Zipper = zipper(*z)

	http.HandleFunc("/render/", renderHandler)

	log.Println("listening on port", *port)
	log.Fatalln(http.ListenAndServe(":"+strconv.Itoa(*port), nil))

}
<commit_msg>Cache zipper find requests too<commit_after>package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"math"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"code.google.com/p/gogoprotobuf/proto"

	pb "github.com/dgryski/carbonzipper/carbonzipperpb"
)

type zipper string

var Zipper zipper

var timeFormats = []string{"15:04 20060102", "20060102", "01/02/06"}

// dateParamToEpoch turns a passed string parameter into an epoch we can send to the Zipper
func dateParamToEpoch(s string, d int64) string {

	if s == "" {
		// return the default if nothing was passed
		return strconv.Itoa(int(d))
	}

	// relative timestamp
	if s[0] == '-' {

		offset, err := intervalString(s[1:])
		if err != nil {
			return strconv.Itoa(int(d))
		}

		return strconv.Itoa(int(timeNow().Add(-time.Duration(offset) * time.Second).Unix()))
	}

	_, err := strconv.Atoi(s)
	if err == nil && len(s) > 8 {
		return s // We got a timestamp so returning it
	}

	if strings.Contains(s, "_") {
		s = strings.Replace(s, "_", " ", 1) // Go can't parse _ in date strings
	}

	for _, format := range timeFormats {
		t, err := time.Parse(format, s)
		if err == nil {
			return strconv.Itoa(int(t.Unix()))
		}
	}
	return strconv.Itoa(int(d))
}

func intervalString(s string) (int32, error) {

	var j int

	for j < len(s) && s[j] >= '0' && s[j] <= '9' {
		j++
	}
	offsetStr, unitStr := s[:j], s[j:]

	var units int
	switch unitStr {
	case "s", "sec", "secs", "second", "seconds":
		units = 1
	case "min", "minute", "minutes":
		units = 60
	case "h", "hour", "hours":
		units = 60 * 60
	case "d", "day", "days":
		units = 24 * 60 * 60
	case "w", "week", "weeks":
		units = 7 * 24 * 60 * 60
	case "mon", "month", "months":
		units = 30 * 24 * 60 * 60
	case "y", "year", "years":
		units = 365 * 24 * 60 * 60
	default:
		return 0, errors.New("unknown time units")
	}

	offset, err := strconv.Atoi(offsetStr)
	if err != nil {
		return 0, err
	}

	return int32(offset * units), nil
}

// FIXME(dgryski): extract the http.Get + unproto code into its own function

func (z zipper) Find(metric string) (pb.GlobResponse, error) {

	u, _ := url.Parse(string(z) + "/metrics/find/")

	u.RawQuery = url.Values{
		"query":  []string{metric},
		"format": []string{"protobuf"},
	}.Encode()

	resp, err := http.Get(u.String())
	if err != nil {
		log.Printf("Find: http.Get: %+v\n", err)
		return pb.GlobResponse{}, err
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Printf("Find: ioutil.ReadAll: %+v\n", err)
		return pb.GlobResponse{}, err
	}

	var pbresp pb.GlobResponse

	err = proto.Unmarshal(body, &pbresp)
	if err != nil {
		log.Printf("Find: proto.Unmarshal: %+v\n", err)
		return pb.GlobResponse{}, err
	}

	return pbresp, nil
}

func (z zipper) Render(metric, from, until string) (pb.FetchResponse, error) {

	u, _ := url.Parse(string(z) + "/render/")

	u.RawQuery = url.Values{
		"target": []string{metric},
		"format": []string{"protobuf"},
		"from":   []string{from},
		"until":  []string{until},
	}.Encode()

	resp, err := http.Get(u.String())
	if err != nil {
		log.Printf("Render: http.Get: %s: %+v\n", metric, err)
		return pb.FetchResponse{}, err
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Printf("Render: ioutil.ReadAll: %s: %+v\n", metric, err)
		return pb.FetchResponse{}, err
	}

	var pbresp pb.FetchResponse

	err = proto.Unmarshal(body, &pbresp)
	if err != nil {
		log.Printf("Render: proto.Unmarshal: %s: %+v\n", metric, err)
		return pb.FetchResponse{}, err
	}

	return pbresp, nil
}

type limiter chan struct{}

func (l limiter) enter() { l <- struct{}{} }
func (l limiter) leave() { <-l }

var Limiter limiter

// for testing
var timeNow = time.Now

type graphitePoint struct {
	value float64
	t     int32
}

func (g graphitePoint) MarshalJSON() ([]byte, error) {
	// TODO(dgryski): fmt.Sprintf() is slow, use strconv.Append{Float,Int}
	if math.IsNaN(g.value) {
		return []byte(fmt.Sprintf("[null,%d]", g.t)), nil
	}
	return []byte(fmt.Sprintf("[%g,%d]", g.value, g.t)), nil
}

type jsonResponse struct {
	Target     string          `json:"target"`
	Datapoints []graphitePoint `json:"datapoints"`
}

func renderHandler(w http.ResponseWriter, r *http.Request) {

	r.ParseForm()
	targets := r.Form["target"]
	from := r.FormValue("from")
	until := r.FormValue("until")

	cacheKey := r.Form.Encode()

	if response, ok := queryCache.get(cacheKey); ok {
		w.Header().Set("Content-Type", "application/json")
		w.Write(response.([]byte))
		return
	}

	// normalize from and until values
	from = dateParamToEpoch(from, timeNow().Add(-24*time.Hour).Unix())
	until = dateParamToEpoch(until, timeNow().Unix())

	var results []*pb.FetchResponse

	for _, target := range targets {

		exp, e, err := parseExpr(target)
		if err != nil || e != "" {
			http.Error(w, http.StatusText(http.StatusBadRequest), http.StatusBadRequest)
			return
		}

		metricMap := make(map[string][]*pb.FetchResponse)

		for _, metric := range exp.metrics() {

			var glob pb.GlobResponse

			if response, ok := findCache.get(metric); ok {
				glob = response.(pb.GlobResponse)
			} else {
				var err error
				glob, err = Zipper.Find(metric)
				if err != nil {
					continue
				}
				findCache.set(metric, glob)
			}

			// for each server in find response query render
			rch := make(chan *pb.FetchResponse, len(glob.GetMatches()))
			for _, m := range glob.GetMatches() {
				go func(m *pb.GlobMatch) {
					Limiter.enter()
					var rptr *pb.FetchResponse
					if m.GetIsLeaf() {
						r, err := Zipper.Render(m.GetPath(), from, until)
						if err == nil {
							rptr = &r
						}
					}
					rch <- rptr
					Limiter.leave()
				}(m)
			}

			for i := 0; i < len(glob.GetMatches()); i++ {
				r := <-rch
				if r != nil {
					metricMap[metric] = append(metricMap[metric], r)
				}
			}
		}

		exprs := evalExpr(exp, metricMap)
		results = append(results, exprs...)
	}

	var jresults []jsonResponse

	for _, r := range results {
		if r == nil {
			log.Println("skipping nil result")
		}
		datapoints := make([]graphitePoint, 0, len(r.Values))
		t := *r.StartTime
		for i, v := range r.Values {
			if r.IsAbsent[i] {
				v = math.NaN()
			}
			datapoints = append(datapoints, graphitePoint{value: v, t: t})
			t += *r.StepTime
		}
		jresults = append(jresults, jsonResponse{Target: r.GetName(), Datapoints: datapoints})
	}

	jout, err := json.Marshal(jresults)
	if err != nil {
		http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
		return
	}

	queryCache.set(cacheKey, jout)

	w.Header().Set("Content-Type", "application/json")
	w.Write(jout)
}

var queryCache *expireCache
var findCache *expireCache

type cacheElement struct {
	validUntil time.Time
	data       interface{}
}

type expireCache struct {
	sync.Mutex
	cache map[string]cacheElement
}

func (ec *expireCache) get(k string) (interface{}, bool) {
	ec.Lock()
	v, ok := ec.cache[k]
	ec.Unlock()
	if !ok || v.validUntil.Before(time.Now()) {
		return nil, false
	}
	return v.data, ok
}

func (ec *expireCache) set(k string, v interface{}) {
	ec.Lock()
	ec.cache[k] = cacheElement{validUntil: time.Now().Add(60 * time.Second), data: v}
	ec.Unlock()
}

func (ec *expireCache) cleaner() {

	var keys []string

	for {
		time.Sleep(5 * time.Minute)

		now := time.Now()
		ec.Lock()

		for k, v := range ec.cache {
			if v.validUntil.Before(now) {
				keys = append(keys, k)
			}
		}

		for _, k := range keys {
			log.Println("cleaning", k)
			delete(ec.cache, k)
		}

		keys = keys[:0]
		ec.Unlock()
	}
}

func main() {

	z := flag.String("z", "", "zipper")
	port := flag.Int("p", 8080, "port")
	l := flag.Int("l", 20, "concurrency limit")

	flag.Parse()

	if *z == "" {
		log.Fatal("no zipper (-z) provided")
	}

	if p := os.Getenv("PORT"); p != "" {
		*port, _ = strconv.Atoi(p)
	}

	Limiter = make(chan struct{}, *l)

	if *z == "" {
		log.Fatal("no zipper provided")
	}

	if _, err := url.Parse(*z); err != nil {
		log.Fatal("unable to parze zipper:", err)
	}

	log.Println("using zipper", *z)

	queryCache = &expireCache{cache: make(map[string]cacheElement)}
	go queryCache.cleaner()

	findCache = &expireCache{cache: make(map[string]cacheElement)}
	go findCache.cleaner()

	Zipper = zipper(*z)

	http.HandleFunc("/render/", renderHandler)

	log.Println("listening on port", *port)
	log.Fatalln(http.ListenAndServe(":"+strconv.Itoa(*port), nil))

}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"golang.org/x/net/html"
	"log"
	"net/http"
	"net/url"
	"os"
	"os/signal"
)

type Page struct {
	Url   *url.URL
	Next  *Page
}

func (page *Page) IsMatch(opage *Page) bool {
	return page.Url.String() == opage.Url.String()
}

// Gets the first link of a wikipedia page.
func (page *Page) FollowLink(acceptFunc func(ur *url.URL) bool) (*Page, error) {
	resp, err := http.Get(page.Url.String())
	if err != nil {
		return page, err
	}

	body := resp.Body
	defer body.Close()

	z := html.NewTokenizer(body)
	inBody := false
	depth := 0
	for {
		tt := z.Next()
		switch tt {
		case html.ErrorToken:
			return page, z.Err()
		case html.StartTagToken, html.EndTagToken:
			tn, _ := z.TagName()
			if string(tn) == "div" {
				if tt == html.StartTagToken {
					if inBody {
						// Descend into an inner div
						depth++
					} else {
						// This is a div tag
						// Wikipedia puts the main section of the article
						// within a div tag with the id "bodyContent"
						// Loop through attributes for an id
						more := true
						for more {
							key, val, m := z.TagAttr()
							more = m
							if string(key) == "id" && string(val) == "bodyContent" {
								inBody = true
							}
						}
					}
				} else {
					if depth == 0 {
						inBody = false
					}
				}
			} else if inBody && tt == html.StartTagToken && string(tn) == "a" {
				// This is an anchor tag
				// This is an anchor tag in a div
				// Check if it has an href attribute
				more := true
				for more {
					key, val, m := z.TagAttr()
					more = m
					if string(key) == "href" {
						// Parse URL
						ur, err := page.Url.Parse(string(val))
						if err != nil {
							// If this url is not parseable,
							// skip to the second url
							break
						}

						if acceptFunc(ur) {
							p := &Page{Url: ur}
							page.Next = p
							return p, nil
						}
					}
				}
			}
		}
	}
}

func main() {
	haveVisited := make(map[url.URL]Page)

	var targetPage *Page
	var startPage *Page

	if len(os.Args) == 3 {

		ur, err := url.Parse(os.Args[1])
		if err != nil {
			log.Fatal(err)
		}

		targetPage = &Page{Url: ur}

		ur, err = url.Parse(os.Args[2])
		if err != nil {
			log.Fatal(err)
		}

		// Initial page to start crawler
		startPage = &Page{Url: ur}
	} else {
		fmt.Println("Needs url to start crawler")
		return
	}

	done := make(chan bool)
	go func () {
		page := startPage
		for {
			pg, err := page.FollowLink(func(ur *url.URL) bool {
				p := haveVisited[*ur]
				if p.Url != nil {
					return false
				}
				return true
			})
			page = pg

			if err != nil {
				log.Fatal(fmt.Sprintf("Could find link on %s: %s\n", page.Url, err))
			}

			fmt.Printf("Have Followed %d links\r", len(haveVisited))

			haveVisited[*page.Url] = *page

			if page.IsMatch(targetPage) {
				fmt.Printf("Found match, took %d follows\n", len(haveVisited))
				break
			}
		}
		done <- true
	}()

	sig := make(chan os.Signal)
	signal.Notify(sig, os.Interrupt)

	// Print path
	select {
	case <-done:
	case <-sig:
	}

	page := startPage
	i := 0
	for page != nil {
		fmt.Printf("%d: %s\n", i, page.Url)
		page = page.Next
		i++
	}
}
<commit_msg>update<commit_after>package main

import (
	"fmt"
	"golang.org/x/net/html"
	"log"
	"net/http"
	"net/url"
	"os"
	"os/signal"
)

type Page struct {
	Url   *url.URL
	Next  *Page
}

func (page *Page) IsMatch(opage *Page) bool {
	return page.Url.String() == opage.Url.String()
}

// Gets the first link of a wikipedia page.
func (page *Page) FollowLink(acceptFunc func(ur *url.URL) bool) (*Page, error) {
	resp, err := http.Get(page.Url.String())
	if err != nil {
		return page, err
	}

	body := resp.Body
	defer body.Close()

	z := html.NewTokenizer(body)
	inBody := false
	depth := 0
	for {
		tt := z.Next()
		switch tt {
		case html.ErrorToken:
			return page, z.Err()
		case html.StartTagToken, html.EndTagToken:
			tn, _ := z.TagName()
			if string(tn) == "div" {
				if tt == html.StartTagToken {
					if inBody {
						// Descend into an inner div
						depth++
					} else {
						// This is a div tag
						// Wikipedia puts the main section of the article
						// within a div tag with the id "bodyContent"
						// Loop through attributes for an id
						more := true
						for more {
							key, val, m := z.TagAttr()
							more = m
							if string(key) == "id" && string(val) == "bodyContent" {
								inBody = true
							}
						}
					}
				} else {
					if depth == 0 {
						inBody = false
					}
				}
			} else if inBody && tt == html.StartTagToken && string(tn) == "a" {
				// This is an anchor tag
				// This is an anchor tag in a div
				// Check if it has an href attribute
				more := true
				for more {
					key, val, m := z.TagAttr()
					more = m
					if string(key) == "href" {
						// Parse URL
						ur, err := page.Url.Parse(string(val))
						if err != nil {
							// If this url is not parseable,
							// skip to the second url
							break
						}

						if acceptFunc(ur) {
							p := &Page{Url: ur}
							page.Next = p
							return p, nil
						}
					}
				}
			}
		}
	}
}

func main() {
	haveVisited := make(map[url.URL]Page)
	visits := 0

	var targetPage *Page
	var startPage *Page

	if len(os.Args) == 3 {

		ur, err := url.Parse(os.Args[1])
		if err != nil {
			log.Fatal(err)
		}

		targetPage = &Page{Url: ur}

		ur, err = url.Parse(os.Args[2])
		if err != nil {
			log.Fatal(err)
		}

		// Initial page to start crawler
		startPage = &Page{Url: ur}
	} else {
		fmt.Println("Needs url to start crawler")
		return
	}

	done := make(chan bool)
	go func () {
		page := startPage
		for {
			pg, err := page.FollowLink(func(ur *url.URL) bool {
				p := haveVisited[*ur]
				if p.Url != nil {
					return false
				}
				return true
			})
			if err != nil {
				if err.Error() == "EOF" {
					// Could not find a link on this file,
					// Go back up one page
					lp := startPage
					p := startPage
					for p.Next != nil {
						lp = p
						p = p.Next
					}
					if (lp == startPage) {
						log.Fatal("Cannot find links on provided page")
					}
					page = lp
					visits--
					continue
				}
				log.Fatal(err)
			}
			page = pg

			fmt.Printf("Have Followed %d links\r", visits)

			haveVisited[*page.Url] = *page
			visits++

			if page.IsMatch(targetPage) {
				fmt.Printf("Found match, took %d follows\n", visits)
				break
			}
		}
		done <- true
	}()

	sig := make(chan os.Signal)
	signal.Notify(sig, os.Interrupt)

	// Print path
	select {
	case <-done:
	case <-sig:
	}

	page := startPage
	i := 0
	for page != nil {
		fmt.Printf("%d: %s\n", i, page.Url)
		page = page.Next
		i++
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"fmt"
	"github.com/monochromegane/hoi/server"
	"os"
	"os/user"
	"path/filepath"
)

func main() {

	file := os.Args[1]

	// create hoi public directory
	usr, _ := user.Current()
	homeDir := usr.HomeDir
	publicDir := filepath.Join(homeDir, ".hoi", "public")
	os.MkdirAll(publicDir, 0755)

	// create random directory
	random := randomString(32)
	randomDir := filepath.Join(publicDir, random)
	os.Mkdir(randomDir, 0755)

	// create symblic link
	os.Symlink(file, filepath.Join(randomDir, filepath.Base(file)))

	fmt.Println(filepath.Join(random, filepath.Base(file)))

	// start hoi server
	server.Start(publicDir)
}

func randomString(length int) string {
	alphanum := "0123456789abcdefghijklmnopqrstuvwxyz"
	var bytes = make([]byte, length)
	rand.Read(bytes)
	for i, b := range bytes {
		bytes[i] = alphanum[b%byte(len(alphanum))]
	}
	return string(bytes)
}
<commit_msg>Used go-flags.<commit_after>package main

import (
	"crypto/rand"
	"fmt"
	flags "github.com/jessevdk/go-flags"
	"github.com/monochromegane/hoi/option"
	"github.com/monochromegane/hoi/server"
	"os"
	"os/user"
	"path/filepath"
)

var opts option.Options

func main() {

        args, err := flags.Parse(&opts)
        if err != nil {
                os.Exit(1)
        }

	file := args[0]

	// create hoi public directory
	usr, _ := user.Current()
	homeDir := usr.HomeDir
	publicDir := filepath.Join(homeDir, ".hoi", "public")
	os.MkdirAll(publicDir, 0755)

	// create random directory
	random := randomString(32)
	randomDir := filepath.Join(publicDir, random)
	os.Mkdir(randomDir, 0755)

	// create symblic link
	os.Symlink(file, filepath.Join(randomDir, filepath.Base(file)))

	fmt.Println(filepath.Join(random, filepath.Base(file)))

	// start hoi server
	server.Start(publicDir)
}

func randomString(length int) string {
	alphanum := "0123456789abcdefghijklmnopqrstuvwxyz"
	var bytes = make([]byte, length)
	rand.Read(bytes)
	for i, b := range bytes {
		bytes[i] = alphanum[b%byte(len(alphanum))]
	}
	return string(bytes)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
)

func shortUsage() {
	os.Stderr.WriteString(`
Usage: vub [OPTION]... URI...
Try 'vub --help' for more information.
`[1:])
}

func usage() {
	os.Stderr.WriteString(`
Usage: vub [OPTION]... URI...
Install Vim plugin to under the management of vim-unbundle.

URI:
  sunaku/vim-unbundle                    # short URI
  https://github.com/sunaku/vim-unbundle # full URI

Options:
  -f, --filetype=TYPE       installing under the ftbundle/TYPE
  -l, --list                change the behavior to list packages
  -r, --remove              change the behavior to remove
  -u, --update              change the behavior to update
      --help                show this help message
      --version             output version information
`[1:])
}

func version() {
	os.Stderr.WriteString(`
0.3.1
`[1:])
}

func countTrue(bls ...bool) int {
	cnt := 0
	for _, b := range bls {
		if b {
			cnt++
		}
	}
	return cnt
}

func printError(err error) {
	fmt.Fprintln(os.Stderr, "vub:", err)
}

func main() {
	f := flag.NewFlagSet("vub", flag.ContinueOnError)
	f.SetOutput(ioutil.Discard)

	var filetype string
	f.StringVar(&filetype, "f", "", "")
	f.StringVar(&filetype, "filetype", "", "")

	var listMode, removeMode, updateMode bool
	f.BoolVar(&listMode, "l", false, "")
	f.BoolVar(&listMode, "list", false, "")
	f.BoolVar(&removeMode, "r", false, "")
	f.BoolVar(&removeMode, "remove", false, "")
	f.BoolVar(&updateMode, "u", false, "")
	f.BoolVar(&updateMode, "update", false, "")

	var isHelp, isVersion bool
	f.BoolVar(&isHelp, "h", false, "")
	f.BoolVar(&isHelp, "help", false, "")
	f.BoolVar(&isVersion, "version", false, "")

	if err := f.Parse(os.Args[1:]); err != nil {
		printError(err)
		os.Exit(2)
	}
	switch {
	case isHelp:
		usage()
		os.Exit(0)
	case isVersion:
		version()
		os.Exit(0)
	case !listMode && f.NArg() < 1:
		shortUsage()
		os.Exit(2)
	case countTrue(listMode, removeMode, updateMode) > 1:
		printError(fmt.Errorf("cannot specify multiple mode"))
		os.Exit(2)
	}

	switch {
	case listMode:
		ListPackages(filetype)
	default:
		var err error
		for _, uri := range f.Args() {
			p := NewPackage(uri, filetype)
			switch {
			case removeMode:
				err = p.Remove()
			case updateMode:
				err = p.Update()
			default:
				err = p.Install()
			}
			if err != nil {
				printError(err)
				os.Exit(1)
			}
		}
	}
}
<commit_msg>s/version/printVersion/<commit_after>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
)

func shortUsage() {
	os.Stderr.WriteString(`
Usage: vub [OPTION]... URI...
Try 'vub --help' for more information.
`[1:])
}

func usage() {
	os.Stderr.WriteString(`
Usage: vub [OPTION]... URI...
Install Vim plugin to under the management of vim-unbundle.

URI:
  sunaku/vim-unbundle                    # short URI
  https://github.com/sunaku/vim-unbundle # full URI

Options:
  -f, --filetype=TYPE       installing under the ftbundle/TYPE
  -l, --list                change the behavior to list packages
  -r, --remove              change the behavior to remove
  -u, --update              change the behavior to update
      --help                show this help message
      --version             output version information
`[1:])
}

func printVersion() {
	os.Stderr.WriteString(`
0.3.1
`[1:])
}

func countTrue(bls ...bool) int {
	cnt := 0
	for _, b := range bls {
		if b {
			cnt++
		}
	}
	return cnt
}

func printError(err error) {
	fmt.Fprintln(os.Stderr, "vub:", err)
}

func main() {
	f := flag.NewFlagSet("vub", flag.ContinueOnError)
	f.SetOutput(ioutil.Discard)

	var filetype string
	f.StringVar(&filetype, "f", "", "")
	f.StringVar(&filetype, "filetype", "", "")

	var listMode, removeMode, updateMode bool
	f.BoolVar(&listMode, "l", false, "")
	f.BoolVar(&listMode, "list", false, "")
	f.BoolVar(&removeMode, "r", false, "")
	f.BoolVar(&removeMode, "remove", false, "")
	f.BoolVar(&updateMode, "u", false, "")
	f.BoolVar(&updateMode, "update", false, "")

	var isHelp, isVersion bool
	f.BoolVar(&isHelp, "h", false, "")
	f.BoolVar(&isHelp, "help", false, "")
	f.BoolVar(&isVersion, "version", false, "")

	if err := f.Parse(os.Args[1:]); err != nil {
		printError(err)
		os.Exit(2)
	}
	switch {
	case isHelp:
		usage()
		os.Exit(0)
	case isVersion:
		printVersion()
		os.Exit(0)
	case !listMode && f.NArg() < 1:
		shortUsage()
		os.Exit(2)
	case countTrue(listMode, removeMode, updateMode) > 1:
		printError(fmt.Errorf("cannot specify multiple mode"))
		os.Exit(2)
	}

	switch {
	case listMode:
		ListPackages(filetype)
	default:
		var err error
		for _, uri := range f.Args() {
			p := NewPackage(uri, filetype)
			switch {
			case removeMode:
				err = p.Remove()
			case updateMode:
				err = p.Update()
			default:
				err = p.Install()
			}
			if err != nil {
				printError(err)
				os.Exit(1)
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"flag"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/fsouza/go-dockerclient"
	"github.com/mesos/mesos-go/executor"
	mesos "github.com/mesos/mesos-go/mesosproto"
	"github.com/newrelic/sidecar/service"
	"github.com/nitro/sidecar-executor/container"
	"github.com/relistan/envconfig"
	"github.com/relistan/go-director"
)

const (
	TaskRunning  = 0
	TaskFinished = iota
	TaskFailed   = iota
	TaskKilled   = iota
)

const (
	StatusSleepTime = 1 * time.Second // How long we wait for final status updates to hit Mesos worker
)

var (
	config Config
)

type Config struct {
	KillTaskTimeout     uint          `split_words:"true" default:"5"` // Seconds
	HttpTimeout         time.Duration `split_words:"true" default:"2s"`
	SidecarRetryCount   int           `split_words:"true" default:"5"`
	SidecarRetryDelay   time.Duration `split_words:"true" default:"3s"`
	SidecarUrl          string        `split_words:"true" default:"http://localhost:7777/state.json"`
	SidecarBackoff      time.Duration `split_words:"true" default:"1m"`
	SidecarPollInterval time.Duration `split_words:"true" default:"30s"`
	SidecarMaxFails     int           `split_words:"true" default:"3"`
	DockerRepository    string        `split_words:"true" default:"https://index.docker.io/v1/"`
	LogsSince           time.Duration `split_words:"true" default:"3m"`
}

type sidecarExecutor struct {
	driver      *executor.MesosExecutorDriver
	client      container.DockerClient
	fetcher     SidecarFetcher
	watchLooper director.Looper
	dockerAuth  *docker.AuthConfiguration
	failCount   int
}

type SidecarServices struct {
	Servers map[string]struct {
		Services map[string]service.Service
	}
}

type SidecarFetcher interface {
	Get(url string) (resp *http.Response, err error)
}

func newSidecarExecutor(client container.DockerClient, auth *docker.AuthConfiguration) *sidecarExecutor {
	return &sidecarExecutor{
		client:     client,
		fetcher:    &http.Client{Timeout: config.HttpTimeout},
		dockerAuth: auth,
	}
}

func logConfig() {
	log.Infof("Executor Config -----------------------")
	log.Infof(" * KillTaskTimeout:     %d", config.KillTaskTimeout)
	log.Infof(" * HttpTimeout:         %s", config.HttpTimeout.String())
	log.Infof(" * SidecarRetryCount:   %d", config.SidecarRetryCount)
	log.Infof(" * SidecarRetryDelay:   %s", config.SidecarRetryDelay.String())
	log.Infof(" * SidecarUrl:          %s", config.SidecarUrl)
	log.Infof(" * SidecarBackoff:      %s", config.SidecarBackoff.String())
	log.Infof(" * SidecarPollInterval: %s", config.SidecarPollInterval.String())
	log.Infof(" * SidecarMaxFails:     %d", config.SidecarMaxFails)
	log.Infof(" * DockerRepository:    %s", config.DockerRepository)
	log.Infof(" * LogsSince:           %s", config.LogsSince.String())

	log.Infof("Environment ---------------------------")
	for _, setting := range os.Environ() {
		if (len(setting) >= 5 && setting[:5] == "MESOS") ||
			((len(setting) >= 8) && setting[:8] == "EXECUTOR") ||
			(setting == "HOME") {
			pair := strings.Split(setting, "=")
			log.Infof(" * %-25s: %s", pair[0], pair[1])
		}
	}
	log.Infof("---------------------------------------")
}

func logTaskEnv(taskInfo *mesos.TaskInfo) {
	env := container.EnvForTask(taskInfo)
	if len(env) < 1 {
		log.Info("No Docker environment provided")
		return
	}

	log.Infof("Docker Environment --------------------")
	for _, setting := range env {
		log.Infof(" * %s", setting)
	}
	log.Infof("---------------------------------------")
}

// Send task status updates to Mesos via the executor driver
func (exec *sidecarExecutor) sendStatus(status int64, taskId *mesos.TaskID) {
	var mesosStatus *mesos.TaskState
	switch status {
	case TaskRunning:
		mesosStatus = mesos.TaskState_TASK_RUNNING.Enum()
	case TaskFinished:
		mesosStatus = mesos.TaskState_TASK_FINISHED.Enum()
	case TaskFailed:
		mesosStatus = mesos.TaskState_TASK_FAILED.Enum()
	case TaskKilled:
		mesosStatus = mesos.TaskState_TASK_KILLED.Enum()
	}

	update := &mesos.TaskStatus{
		TaskId: taskId,
		State:  mesosStatus,
	}

	if _, err := exec.driver.SendStatusUpdate(update); err != nil {
		log.Errorf("Error sending status update %s", err.Error())
		panic(err.Error())
	}
}

// Tell Mesos and thus the framework that the task finished. Shutdown driver.
func (exec *sidecarExecutor) finishTask(taskInfo *mesos.TaskInfo) {
	exec.sendStatus(TaskFinished, taskInfo.GetTaskId())
	time.Sleep(StatusSleepTime)
	exec.driver.Stop()
}

// Tell Mesos and thus the framework that the task failed. Shutdown driver.
func (exec *sidecarExecutor) failTask(taskInfo *mesos.TaskInfo) {
	exec.sendStatus(TaskFailed, taskInfo.GetTaskId())

	// Unfortunately the status updates are sent async and we can't
	// get a handle on the channel used to send them. So we wait
	time.Sleep(StatusSleepTime)

	// We're done with this executor, so let's stop now.
	exec.driver.Stop()
}

// Loop on a timed basis and check the health of the process in Sidecar.
// Note that because of the way the retries work, the loop timing is a
// lower bound on the delay.
func (exec *sidecarExecutor) watchContainer(containerId string) {
	time.Sleep(config.SidecarBackoff)

	exec.watchLooper.Loop(func() error {
		containers, err := exec.client.ListContainers(
			docker.ListContainersOptions{
				All: true,
			},
		)
		if err != nil {
			return err
		}

		// Loop through all the containers, looking for a running
		// container with our Id.
		ok := false
		for _, entry := range containers {
			if entry.ID == containerId {
				ok = true
				break
			}
		}
		if !ok {
			return errors.New("Container " + containerId + " not running!")
		}

		// Validate health status with Sidecar
		if err = exec.sidecarStatus(containerId); err != nil {
			return err
		}

		return nil
	})
}

// Lookup a container in a service list
func sidecarLookup(containerId string, services SidecarServices) (*service.Service, bool) {
	hostname := os.Getenv("TASK_HOST") // Mesos supplies this
	if _, ok := services.Servers[hostname]; !ok {
		// Don't even have this host!
		log.Warnf("Host not found in Sidecar! (%s)", hostname)
		return nil, ok
	}

	svc, ok := services.Servers[hostname].Services[containerId[:12]]

	return &svc, ok
}

// We only want to kill things that are definitely unhealthy
func shouldBeKilled(svc *service.Service) bool {
	return svc.Status == service.UNHEALTHY || svc.Status == service.TOMBSTONE
}

func (exec *sidecarExecutor) exceededFailCount() bool {
	return exec.failCount >= config.SidecarMaxFails
}

// Validate the status of this task with Sidecar
func (exec *sidecarExecutor) sidecarStatus(containerId string) error {
	fetch := func() ([]byte, error) {
		resp, err := exec.fetcher.Get(config.SidecarUrl)
		defer func() {
			if resp != nil {
				resp.Body.Close()
			}
		}()

		if err != nil {
			return nil, err
		}

		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return nil, err
		}

		return body, nil
	}

	// Try to connect to Sidecar, with some retries
	var err error
	var data []byte
	for i := 0; i <= config.SidecarRetryCount; i++ {
		data, err = fetch()
		if err != nil {
			log.Warnf("Failed %d health checks!", i+1)
			time.Sleep(config.SidecarRetryDelay)
			continue
		}
	}

	// We really really don't want to shut off all the jobs if Sidecar
	// is down. That would make it impossible to deploy Sidecar, and
	// would make the entire system dependent on it for services to
	// even start.
	if err != nil {
		log.Error("Can't contact Sidecar! Assuming healthy...")
		return nil
	}

	// We got a successful result from Sidecar, so let's parse it!
	var services SidecarServices
	err = json.Unmarshal(data, &services)
	if err != nil {
		log.Error("Can't parse Sidecar results! Assuming healthy...")
		return nil
	}

	svc, ok := sidecarLookup(containerId, services)
	if !ok {
		log.Errorf("Can't find this service in Sidecar yet! Assuming healthy...")
		return nil
	}

	// This is the one and only place where we're going to raise our hand
	// and say something is wrong with this service and it needs to be
	// shot by Mesos.
	if shouldBeKilled(svc) {
		// Only bail out if we've exceed the setting for number of failures
		if !exec.exceededFailCount() {
			exec.failCount += 1
			log.Warnf("Failed Sidecar health check, but below fail limit")
			return nil
		}

		log.Errorf("Health failure count exceeded %d", config.SidecarMaxFails)

		exec.failCount = 0
		return errors.New("Unhealthy container: " + containerId + " failing task!")
	}

	exec.failCount = 0 // Reset because we were healthy!

	return nil
}

// Try to find the Docker registry auth information. This will look in the usual
// locations. Note that an environment variable of DOCKER_CONFIG can override the
// other locations. If it's set, we'll look in $DOCKER_CONFIG/config.json.
// https://godoc.org/github.com/fsouza/go-dockerclient#NewAuthConfigurationsFromDockerCfg
func getDockerAuthConfig() docker.AuthConfiguration {
	lookup := func() (docker.AuthConfiguration, error) {
		// Attempt to fetch and configure Docker auth
		auths, err := docker.NewAuthConfigurationsFromDockerCfg()
		if err == nil {
			if auth, ok := auths.Configs[config.DockerRepository]; ok {
				log.Infof("Found Docker auth configuration for '%s'", config.DockerRepository)
				return auth, nil
			}
		}

		return docker.AuthConfiguration{}, errors.New("No auth match for repository")
	}

	// Try the first time
	auth, err := lookup()
	if err == nil {
		return auth
	}

	// Set the home dir to the likely path, then try one more time
	os.Setenv("HOME", "/root")
	auth, err = lookup()
	if err != nil {
		log.Warnf(
			"No docker auth match for repository '%s'... proceeding anyway",
			config.DockerRepository,
		)
	}

	return auth
}

func init() {
	flag.Parse()
	err := envconfig.Process("executor", &config)
	if err != nil {
		log.Fatal(err.Error())
	}

	log.SetOutput(os.Stdout)
	log.SetLevel(log.DebugLevel)
}

func main() {
	log.Info("Starting Sidecar Executor")
	logConfig()

	// Get a Docker client. Without one, we can't do anything.
	dockerClient, err := docker.NewClientFromEnv()
	if err != nil {
		log.Fatal(err.Error())
	}

	dockerAuth := getDockerAuthConfig()
	scExec := newSidecarExecutor(dockerClient, &dockerAuth)
	dconfig := executor.DriverConfig{Executor: scExec}

	driver, err := executor.NewMesosExecutorDriver(dconfig)
	if err != nil || driver == nil {
		log.Info("Unable to create an ExecutorDriver ", err.Error())
	}

	// Give the executor a reference to the driver
	scExec.driver = driver

	_, err = driver.Start()
	if err != nil {
		log.Info("Got error:", err)
		return
	}

	log.Info("Executor process has started")

	_, err = driver.Join()
	if err != nil {
		log.Info("driver failed:", err)
	}

	log.Info("Sidecar Executor exiting")
}
<commit_msg>Capture some fatal signals and clean up<commit_after>package main

import (
	"encoding/json"
	"errors"
	"flag"
	"io/ioutil"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/fsouza/go-dockerclient"
	"github.com/mesos/mesos-go/executor"
	mesos "github.com/mesos/mesos-go/mesosproto"
	"github.com/newrelic/sidecar/service"
	"github.com/nitro/sidecar-executor/container"
	"github.com/relistan/envconfig"
	"github.com/relistan/go-director"
)

const (
	TaskRunning  = 0
	TaskFinished = iota
	TaskFailed   = iota
	TaskKilled   = iota
)

const (
	StatusSleepTime = 1 * time.Second // How long we wait for final status updates to hit Mesos worker
)

var (
	config Config
)

type Config struct {
	KillTaskTimeout     uint          `split_words:"true" default:"5"` // Seconds
	HttpTimeout         time.Duration `split_words:"true" default:"2s"`
	SidecarRetryCount   int           `split_words:"true" default:"5"`
	SidecarRetryDelay   time.Duration `split_words:"true" default:"3s"`
	SidecarUrl          string        `split_words:"true" default:"http://localhost:7777/state.json"`
	SidecarBackoff      time.Duration `split_words:"true" default:"1m"`
	SidecarPollInterval time.Duration `split_words:"true" default:"30s"`
	SidecarMaxFails     int           `split_words:"true" default:"3"`
	DockerRepository    string        `split_words:"true" default:"https://index.docker.io/v1/"`
	LogsSince           time.Duration `split_words:"true" default:"3m"`
}

type sidecarExecutor struct {
	driver      *executor.MesosExecutorDriver
	client      container.DockerClient
	fetcher     SidecarFetcher
	watchLooper director.Looper
	dockerAuth  *docker.AuthConfiguration
	failCount   int
}

type SidecarServices struct {
	Servers map[string]struct {
		Services map[string]service.Service
	}
}

type SidecarFetcher interface {
	Get(url string) (resp *http.Response, err error)
}

func newSidecarExecutor(client container.DockerClient, auth *docker.AuthConfiguration) *sidecarExecutor {
	return &sidecarExecutor{
		client:     client,
		fetcher:    &http.Client{Timeout: config.HttpTimeout},
		dockerAuth: auth,
	}
}

func logConfig() {
	log.Infof("Executor Config -----------------------")
	log.Infof(" * KillTaskTimeout:     %d", config.KillTaskTimeout)
	log.Infof(" * HttpTimeout:         %s", config.HttpTimeout.String())
	log.Infof(" * SidecarRetryCount:   %d", config.SidecarRetryCount)
	log.Infof(" * SidecarRetryDelay:   %s", config.SidecarRetryDelay.String())
	log.Infof(" * SidecarUrl:          %s", config.SidecarUrl)
	log.Infof(" * SidecarBackoff:      %s", config.SidecarBackoff.String())
	log.Infof(" * SidecarPollInterval: %s", config.SidecarPollInterval.String())
	log.Infof(" * SidecarMaxFails:     %d", config.SidecarMaxFails)
	log.Infof(" * DockerRepository:    %s", config.DockerRepository)
	log.Infof(" * LogsSince:           %s", config.LogsSince.String())

	log.Infof("Environment ---------------------------")
	for _, setting := range os.Environ() {
		if (len(setting) >= 5 && setting[:5] == "MESOS") ||
			((len(setting) >= 8) && setting[:8] == "EXECUTOR") ||
			(setting == "HOME") {
			pair := strings.Split(setting, "=")
			log.Infof(" * %-25s: %s", pair[0], pair[1])
		}
	}
	log.Infof("---------------------------------------")
}

func logTaskEnv(taskInfo *mesos.TaskInfo) {
	env := container.EnvForTask(taskInfo)
	if len(env) < 1 {
		log.Info("No Docker environment provided")
		return
	}

	log.Infof("Docker Environment --------------------")
	for _, setting := range env {
		log.Infof(" * %s", setting)
	}
	log.Infof("---------------------------------------")
}

// Send task status updates to Mesos via the executor driver
func (exec *sidecarExecutor) sendStatus(status int64, taskId *mesos.TaskID) {
	var mesosStatus *mesos.TaskState
	switch status {
	case TaskRunning:
		mesosStatus = mesos.TaskState_TASK_RUNNING.Enum()
	case TaskFinished:
		mesosStatus = mesos.TaskState_TASK_FINISHED.Enum()
	case TaskFailed:
		mesosStatus = mesos.TaskState_TASK_FAILED.Enum()
	case TaskKilled:
		mesosStatus = mesos.TaskState_TASK_KILLED.Enum()
	}

	update := &mesos.TaskStatus{
		TaskId: taskId,
		State:  mesosStatus,
	}

	if _, err := exec.driver.SendStatusUpdate(update); err != nil {
		log.Errorf("Error sending status update %s", err.Error())
		panic(err.Error())
	}
}

// Tell Mesos and thus the framework that the task finished. Shutdown driver.
func (exec *sidecarExecutor) finishTask(taskInfo *mesos.TaskInfo) {
	exec.sendStatus(TaskFinished, taskInfo.GetTaskId())
	time.Sleep(StatusSleepTime)
	exec.driver.Stop()
}

// Tell Mesos and thus the framework that the task failed. Shutdown driver.
func (exec *sidecarExecutor) failTask(taskInfo *mesos.TaskInfo) {
	exec.sendStatus(TaskFailed, taskInfo.GetTaskId())

	// Unfortunately the status updates are sent async and we can't
	// get a handle on the channel used to send them. So we wait
	time.Sleep(StatusSleepTime)

	// We're done with this executor, so let's stop now.
	exec.driver.Stop()
}

// Loop on a timed basis and check the health of the process in Sidecar.
// Note that because of the way the retries work, the loop timing is a
// lower bound on the delay.
func (exec *sidecarExecutor) watchContainer(containerId string) {
	time.Sleep(config.SidecarBackoff)

	exec.watchLooper.Loop(func() error {
		containers, err := exec.client.ListContainers(
			docker.ListContainersOptions{
				All: true,
			},
		)
		if err != nil {
			return err
		}

		// Loop through all the containers, looking for a running
		// container with our Id.
		ok := false
		for _, entry := range containers {
			if entry.ID == containerId {
				ok = true
				break
			}
		}
		if !ok {
			return errors.New("Container " + containerId + " not running!")
		}

		// Validate health status with Sidecar
		if err = exec.sidecarStatus(containerId); err != nil {
			return err
		}

		return nil
	})
}

// Lookup a container in a service list
func sidecarLookup(containerId string, services SidecarServices) (*service.Service, bool) {
	hostname := os.Getenv("TASK_HOST") // Mesos supplies this
	if _, ok := services.Servers[hostname]; !ok {
		// Don't even have this host!
		log.Warnf("Host not found in Sidecar! (%s)", hostname)
		return nil, ok
	}

	svc, ok := services.Servers[hostname].Services[containerId[:12]]

	return &svc, ok
}

// We only want to kill things that are definitely unhealthy
func shouldBeKilled(svc *service.Service) bool {
	return svc.Status == service.UNHEALTHY || svc.Status == service.TOMBSTONE
}

func (exec *sidecarExecutor) exceededFailCount() bool {
	return exec.failCount >= config.SidecarMaxFails
}

// Validate the status of this task with Sidecar
func (exec *sidecarExecutor) sidecarStatus(containerId string) error {
	fetch := func() ([]byte, error) {
		resp, err := exec.fetcher.Get(config.SidecarUrl)
		defer func() {
			if resp != nil {
				resp.Body.Close()
			}
		}()

		if err != nil {
			return nil, err
		}

		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return nil, err
		}

		return body, nil
	}

	// Try to connect to Sidecar, with some retries
	var err error
	var data []byte
	for i := 0; i <= config.SidecarRetryCount; i++ {
		data, err = fetch()
		if err != nil {
			log.Warnf("Failed %d health checks!", i+1)
			time.Sleep(config.SidecarRetryDelay)
			continue
		}
	}

	// We really really don't want to shut off all the jobs if Sidecar
	// is down. That would make it impossible to deploy Sidecar, and
	// would make the entire system dependent on it for services to
	// even start.
	if err != nil {
		log.Error("Can't contact Sidecar! Assuming healthy...")
		return nil
	}

	// We got a successful result from Sidecar, so let's parse it!
	var services SidecarServices
	err = json.Unmarshal(data, &services)
	if err != nil {
		log.Error("Can't parse Sidecar results! Assuming healthy...")
		return nil
	}

	svc, ok := sidecarLookup(containerId, services)
	if !ok {
		log.Errorf("Can't find this service in Sidecar yet! Assuming healthy...")
		return nil
	}

	// This is the one and only place where we're going to raise our hand
	// and say something is wrong with this service and it needs to be
	// shot by Mesos.
	if shouldBeKilled(svc) {
		// Only bail out if we've exceed the setting for number of failures
		if !exec.exceededFailCount() {
			exec.failCount += 1
			log.Warnf("Failed Sidecar health check, but below fail limit")
			return nil
		}

		log.Errorf("Health failure count exceeded %d", config.SidecarMaxFails)

		exec.failCount = 0
		return errors.New("Unhealthy container: " + containerId + " failing task!")
	}

	exec.failCount = 0 // Reset because we were healthy!

	return nil
}

// Try to find the Docker registry auth information. This will look in the usual
// locations. Note that an environment variable of DOCKER_CONFIG can override the
// other locations. If it's set, we'll look in $DOCKER_CONFIG/config.json.
// https://godoc.org/github.com/fsouza/go-dockerclient#NewAuthConfigurationsFromDockerCfg
func getDockerAuthConfig() docker.AuthConfiguration {
	lookup := func() (docker.AuthConfiguration, error) {
		// Attempt to fetch and configure Docker auth
		auths, err := docker.NewAuthConfigurationsFromDockerCfg()
		if err == nil {
			if auth, ok := auths.Configs[config.DockerRepository]; ok {
				log.Infof("Found Docker auth configuration for '%s'", config.DockerRepository)
				return auth, nil
			}
		}

		return docker.AuthConfiguration{}, errors.New("No auth match for repository")
	}

	// Try the first time
	auth, err := lookup()
	if err == nil {
		return auth
	}

	// Set the home dir to the likely path, then try one more time
	os.Setenv("HOME", "/root")
	auth, err = lookup()
	if err != nil {
		log.Warnf(
			"No docker auth match for repository '%s'... proceeding anyway",
			config.DockerRepository,
		)
	}

	return auth
}

func handleSignals(scExec *sidecarExecutor) {
	sigChan := make(chan os.Signal, 1) // Buffered!

	// Grab some signals we want to catch where possible
	signal.Notify(sigChan, os.Interrupt)
	signal.Notify(sigChan, os.Kill)
	signal.Notify(sigChan, syscall.SIGTERM)

	sig := <-sigChan
	log.Warnf("Received signal '%s', attempting clean shutdown", sig)
	if scExec.watchLooper != nil {
		scExec.watchLooper.Done(errors.New("Got " + sig.String() + " signal!"))
	}
	time.Sleep(3*time.Second) // Try to let it quit
	os.Exit(130) // Ctrl-C received or equivalent
}

func init() {
	flag.Parse()
	err := envconfig.Process("executor", &config)
	if err != nil {
		log.Fatal(err.Error())
	}

	log.SetOutput(os.Stdout)
	log.SetLevel(log.DebugLevel)
}

func main() {
	log.Info("Starting Sidecar Executor")
	logConfig()

	// Get a Docker client. Without one, we can't do anything.
	dockerClient, err := docker.NewClientFromEnv()
	if err != nil {
		log.Fatal(err.Error())
	}

	dockerAuth := getDockerAuthConfig()
	scExec := newSidecarExecutor(dockerClient, &dockerAuth)
	dconfig := executor.DriverConfig{Executor: scExec}

	driver, err := executor.NewMesosExecutorDriver(dconfig)
	if err != nil || driver == nil {
		log.Info("Unable to create an ExecutorDriver ", err.Error())
	}

	// Give the executor a reference to the driver
	scExec.driver = driver

	go handleSignals(scExec)

	_, err = driver.Start()
	if err != nil {
		log.Info("Got error:", err)
		return
	}

	log.Info("Executor process has started")

	_, err = driver.Join()
	if err != nil {
		log.Info("driver failed:", err)
	}

	log.Info("Sidecar Executor exiting")
}
<|endoftext|>


<commit_before>/*
RS Ingester loads processed event data into a Redshift database and
migrates schemas when they change. It receives pointers to tsv files from
the Spade processor, loads the files in batches, and migrates tables.

The outer binary is aliased as rsloadmanager, and is responsible for running
loads, retrying loads on failure, and executing migrations on the Redshift instance.
*/
package main

import (
	"flag"
	"fmt"
	"net"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
	"github.com/aws/aws-sdk-go/service/s3/s3manager/s3manageriface"
	"github.com/twitchscience/aws_utils/logger"
	"github.com/twitchscience/rs_ingester/blueprint"
	"github.com/twitchscience/rs_ingester/control"
	"github.com/twitchscience/rs_ingester/migrator"
	"github.com/twitchscience/rs_ingester/versions"

	"github.com/twitchscience/rs_ingester/backend"
	"github.com/twitchscience/rs_ingester/healthcheck"
	"github.com/twitchscience/rs_ingester/loadclient"
	"github.com/twitchscience/rs_ingester/metadata"
	"github.com/twitchscience/rs_ingester/monitoring"
	"github.com/twitchscience/rs_ingester/reporter"
)

const (
	healthCheckPoolSize = 1
)

var (
	poolSize             int
	statsPrefix          string
	manifestBucket       string
	rsURL                string
	rollbarToken         string
	rollbarEnvironment   string
	blueprintHost        string
	pgConfig             metadata.PGConfig
	loadAgeSeconds       int
	workerGroup          sync.WaitGroup
	waitProcessorPeriod  time.Duration
	migratorPollPeriod   time.Duration
	reporterPollPeriod   time.Duration
	offpeakStartHour     int
	offpeakDurationHours int
)

type loadWorker struct {
	MetadataBackend metadata.Backend
	Loader          loadclient.Loader
}

func (i *loadWorker) Work(stats monitoring.SafeStatter) {

	c := i.MetadataBackend.LoadReady()
	for load := range c {
		logfields := logger.WithField("loadUUID", load.UUID).
			WithField("numFiles", len(load.Loads)).
			WithField("table", load.TableName)
		logfields.Info("Loading manifest into table")
		err := i.Loader.LoadManifest(load)
		if err != nil {
			if err.Retryable() {
				i.MetadataBackend.LoadError(load.UUID, err.Error())
			}
			logfields.WithError(err).WithField("retryable", err.Retryable()).
				Error("Error loading files into table.")

			stats.SafeInc("manifest_load.failures", 1, 1.0)
			continue
		}
		logfields.Info("Loaded manifest into table")
		i.MetadataBackend.LoadDone(load.UUID)

		stats.SafeInc("manifest_load.count", 1, 1.0)
		statsdPattern := "tsv_files.%s.loaded"
		stats.SafeInc(fmt.Sprintf(statsdPattern, load.TableName), int64(len(load.Loads)), 1.0)
		stats.SafeInc(fmt.Sprintf(statsdPattern, "total"), int64(len(load.Loads)), 1.0)
	}
	workerGroup.Done()
}

func startWorkers(s3Uploader s3manageriface.UploaderAPI, b metadata.Backend, stats monitoring.SafeStatter, aceBackend backend.Backend) ([]loadWorker, error) {
	workers := make([]loadWorker, poolSize)
	for i := 0; i < poolSize; i++ {
		loadclient, err := loadclient.NewRSLoader(s3Uploader, aceBackend, manifestBucket, stats)
		if err != nil {
			return workers, err
		}
		workers[i] = loadWorker{MetadataBackend: b, Loader: loadclient}
		workerGroup.Add(1)
		index := i
		logger.Go(func() {
			workers[index].Work(stats)
		})
	}
	return workers, nil
}

func init() {
	flag.DurationVar(&migratorPollPeriod, "migratorPollPeriod", time.Minute, "the period betwen each poll the migrator does of ingesterdb for new versions to migrate to")
	flag.DurationVar(&reporterPollPeriod, "reporterPollPeriod", time.Minute, "the period betwen each poll the reporter does of ingesterdb to query current stats")
	flag.DurationVar(&waitProcessorPeriod, "waitProcessorPeriod", time.Minute*3, "the period we wait for processor to process all old version TSVs")
	flag.StringVar(&statsPrefix, "statsPrefix", "ingester", "the prefix to statsd")
	flag.StringVar(&pgConfig.DatabaseURL, "databaseURL", "", "Postgres-scheme url for the RDS instance")
	flag.StringVar(&manifestBucket, "manifestBucket", "", "S3 bucket for manifests.")
	flag.IntVar(&pgConfig.MaxConnections, "maxDBConnections", 5, "Number of database connections to open")
	flag.IntVar(&pgConfig.LoadCountTrigger, "loadCountTrigger", 5, "Number of queued tsvs before a load into redshift is triggered")
	flag.IntVar(&loadAgeSeconds, "loadAgeSeconds", 1800, "Max age of tsvs in queue before a load into redshift is triggered")
	flag.IntVar(&poolSize, "n_workers", 5, "Number of load workers and therefore redshift connections. Set to 0 to turn off ingests (COPYs).")
	flag.StringVar(&blueprintHost, "blueprint_host", "", "Host name (and optionally :port) for communicating with blueprint")
	flag.StringVar(&rsURL, "rsURL", "", "URL for Redshift")
	flag.StringVar(&rollbarToken, "rollbarToken", "", "Rollbar post_server_item token")
	flag.StringVar(&rollbarEnvironment, "rollbarEnvironment", "", "Rollbar environment")
	flag.IntVar(&offpeakStartHour, "offpeakStartHour", 3, "Hour that offpeak period starts and migrations can happen, in UTC")
	flag.IntVar(&offpeakDurationHours, "offpeakDurationHours", 8, "Duration of the offpeak migration period, in hours")
}

func main() {
	flag.Parse()
	pgConfig.LoadAgeTrigger = time.Second * time.Duration(loadAgeSeconds)

	stats, err := monitoring.InitStats(statsPrefix)
	if err != nil {
		logger.WithError(err).Fatal("Failed to setup statter")
	}

	logger.InitWithRollbar("info", rollbarToken, rollbarEnvironment)
	logger.CaptureDefault()
	logger.Info("starting")
	defer logger.LogPanic()

	session, err := session.NewSession()
	if err != nil {
		logger.WithError(err).Fatal("Failed to setup aws session")
	}
	s3Uploader := s3manager.NewUploader(session)
	aceBackend, err := backend.BuildRedshiftBackend(session.Config.Credentials, poolSize+healthCheckPoolSize, rsURL)
	if err != nil {
		logger.WithError(err).Fatal("Failed to setup redshift connection")
	}

	rsConnection, err := loadclient.NewRSLoader(s3Uploader, aceBackend, manifestBucket, stats)
	if err != nil {
		logger.WithError(err).Fatal("Failed to setup Redshift loading client for postgres")
	}

	initVersions, err := aceBackend.TableVersions()
	if err != nil {
		logger.WithError(err).Fatal("Failed initialization of table version cache")
	}
	tableVersions := versions.New(initVersions)

	var metaBackend metadata.Backend

	if poolSize > 0 {
		metaBackend, err = metadata.NewPostgresLoader(&pgConfig, rsConnection, tableVersions)
		if err != nil {
			logger.WithError(err).Fatal("Failed to setup postgres backend")
		}

		_, err = startWorkers(s3Uploader, metaBackend, stats, aceBackend)
		if err != nil {
			logger.WithError(err).Fatal("Failed to start workers")
		}
	}

	metaReader, err := metadata.NewPostgresReader(&pgConfig, tableVersions)
	if err != nil {
		logger.WithError(err).Fatal("Failed to setup postgres reader")
	}

	statsReporter := reporter.New(metaReader, stats, reporterPollPeriod)
	blueprintClient := blueprint.New(blueprintHost)
	versionIncrement := make(chan migrator.VersionIncrement)
	migrator := migrator.New(aceBackend, metaReader, blueprintClient, tableVersions, migratorPollPeriod,
		waitProcessorPeriod, offpeakStartHour, offpeakDurationHours, versionIncrement)

	hcBackend := healthcheck.NewBackend(rsConnection, metaReader)
	hcHandler := healthcheck.NewHandler(hcBackend)

	serveMux := http.NewServeMux()
	serveMux.Handle("/health", healthcheck.NewHealthRouter(hcHandler))

	controlBackend := control.NewControlBackend(metaReader, tableVersions, versionIncrement)
	controlHandler := control.NewControlHandler(controlBackend, stats)

	serveMux.Handle("/control/", control.NewControlRouter(controlHandler))

	logger.Go(func() {
		logger.WithError(http.ListenAndServe(net.JoinHostPort("localhost", "8080"), serveMux)).
			Fatal("Serving health and control failed")
	})

	logger.Go(func() {
		logger.WithError(http.ListenAndServe(":7766", http.DefaultServeMux)).
			Error("Serving pprof failed")
	})

	wait := make(chan struct{})
	sigc := make(chan os.Signal, 1)
	signal.Notify(sigc, syscall.SIGINT)
	logger.Info("Loader is set up")
	logger.Go(func() {
		<-sigc
		logger.Info("Sigint received -- shutting down")
		migrator.Close()
		statsReporter.Close()
		if metaBackend != nil {
			metaBackend.Close()
		}
		workerGroup.Wait()
		// Cause flush
		err = stats.Close()
		if err != nil {
			logger.WithError(err).Error("Error closing statter")
		}
		logger.Info("Exiting main cleanly.")
		logger.Wait()
		close(wait)
	})
	<-wait
}
<commit_msg>Update package docs to begin with Package<commit_after>/*
Package rs_ingester provides servers which load processed event data from
the Spade pipeline into a Redshift database.It receives pointers to tsv files
from the Spade processor, loads the files in batches, and migrates tables when
they are changed in Blueprint.

The outer binary is aliased as rsloadmanager, and is responsible for running
loads, retrying loads on failure, and executing migrations on the Redshift
instance.
*/
package main

import (
	"flag"
	"fmt"
	"net"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
	"github.com/aws/aws-sdk-go/service/s3/s3manager/s3manageriface"
	"github.com/twitchscience/aws_utils/logger"
	"github.com/twitchscience/rs_ingester/blueprint"
	"github.com/twitchscience/rs_ingester/control"
	"github.com/twitchscience/rs_ingester/migrator"
	"github.com/twitchscience/rs_ingester/versions"

	"github.com/twitchscience/rs_ingester/backend"
	"github.com/twitchscience/rs_ingester/healthcheck"
	"github.com/twitchscience/rs_ingester/loadclient"
	"github.com/twitchscience/rs_ingester/metadata"
	"github.com/twitchscience/rs_ingester/monitoring"
	"github.com/twitchscience/rs_ingester/reporter"
)

const (
	healthCheckPoolSize = 1
)

var (
	poolSize             int
	statsPrefix          string
	manifestBucket       string
	rsURL                string
	rollbarToken         string
	rollbarEnvironment   string
	blueprintHost        string
	pgConfig             metadata.PGConfig
	loadAgeSeconds       int
	workerGroup          sync.WaitGroup
	waitProcessorPeriod  time.Duration
	migratorPollPeriod   time.Duration
	reporterPollPeriod   time.Duration
	offpeakStartHour     int
	offpeakDurationHours int
)

type loadWorker struct {
	MetadataBackend metadata.Backend
	Loader          loadclient.Loader
}

func (i *loadWorker) Work(stats monitoring.SafeStatter) {

	c := i.MetadataBackend.LoadReady()
	for load := range c {
		logfields := logger.WithField("loadUUID", load.UUID).
			WithField("numFiles", len(load.Loads)).
			WithField("table", load.TableName)
		logfields.Info("Loading manifest into table")
		err := i.Loader.LoadManifest(load)
		if err != nil {
			if err.Retryable() {
				i.MetadataBackend.LoadError(load.UUID, err.Error())
			}
			logfields.WithError(err).WithField("retryable", err.Retryable()).
				Error("Error loading files into table.")

			stats.SafeInc("manifest_load.failures", 1, 1.0)
			continue
		}
		logfields.Info("Loaded manifest into table")
		i.MetadataBackend.LoadDone(load.UUID)

		stats.SafeInc("manifest_load.count", 1, 1.0)
		statsdPattern := "tsv_files.%s.loaded"
		stats.SafeInc(fmt.Sprintf(statsdPattern, load.TableName), int64(len(load.Loads)), 1.0)
		stats.SafeInc(fmt.Sprintf(statsdPattern, "total"), int64(len(load.Loads)), 1.0)
	}
	workerGroup.Done()
}

func startWorkers(s3Uploader s3manageriface.UploaderAPI, b metadata.Backend, stats monitoring.SafeStatter, aceBackend backend.Backend) ([]loadWorker, error) {
	workers := make([]loadWorker, poolSize)
	for i := 0; i < poolSize; i++ {
		loadclient, err := loadclient.NewRSLoader(s3Uploader, aceBackend, manifestBucket, stats)
		if err != nil {
			return workers, err
		}
		workers[i] = loadWorker{MetadataBackend: b, Loader: loadclient}
		workerGroup.Add(1)
		index := i
		logger.Go(func() {
			workers[index].Work(stats)
		})
	}
	return workers, nil
}

func init() {
	flag.DurationVar(&migratorPollPeriod, "migratorPollPeriod", time.Minute, "the period betwen each poll the migrator does of ingesterdb for new versions to migrate to")
	flag.DurationVar(&reporterPollPeriod, "reporterPollPeriod", time.Minute, "the period betwen each poll the reporter does of ingesterdb to query current stats")
	flag.DurationVar(&waitProcessorPeriod, "waitProcessorPeriod", time.Minute*3, "the period we wait for processor to process all old version TSVs")
	flag.StringVar(&statsPrefix, "statsPrefix", "ingester", "the prefix to statsd")
	flag.StringVar(&pgConfig.DatabaseURL, "databaseURL", "", "Postgres-scheme url for the RDS instance")
	flag.StringVar(&manifestBucket, "manifestBucket", "", "S3 bucket for manifests.")
	flag.IntVar(&pgConfig.MaxConnections, "maxDBConnections", 5, "Number of database connections to open")
	flag.IntVar(&pgConfig.LoadCountTrigger, "loadCountTrigger", 5, "Number of queued tsvs before a load into redshift is triggered")
	flag.IntVar(&loadAgeSeconds, "loadAgeSeconds", 1800, "Max age of tsvs in queue before a load into redshift is triggered")
	flag.IntVar(&poolSize, "n_workers", 5, "Number of load workers and therefore redshift connections. Set to 0 to turn off ingests (COPYs).")
	flag.StringVar(&blueprintHost, "blueprint_host", "", "Host name (and optionally :port) for communicating with blueprint")
	flag.StringVar(&rsURL, "rsURL", "", "URL for Redshift")
	flag.StringVar(&rollbarToken, "rollbarToken", "", "Rollbar post_server_item token")
	flag.StringVar(&rollbarEnvironment, "rollbarEnvironment", "", "Rollbar environment")
	flag.IntVar(&offpeakStartHour, "offpeakStartHour", 3, "Hour that offpeak period starts and migrations can happen, in UTC")
	flag.IntVar(&offpeakDurationHours, "offpeakDurationHours", 8, "Duration of the offpeak migration period, in hours")
}

func main() {
	flag.Parse()
	pgConfig.LoadAgeTrigger = time.Second * time.Duration(loadAgeSeconds)

	stats, err := monitoring.InitStats(statsPrefix)
	if err != nil {
		logger.WithError(err).Fatal("Failed to setup statter")
	}

	logger.InitWithRollbar("info", rollbarToken, rollbarEnvironment)
	logger.CaptureDefault()
	logger.Info("starting")
	defer logger.LogPanic()

	session, err := session.NewSession()
	if err != nil {
		logger.WithError(err).Fatal("Failed to setup aws session")
	}
	s3Uploader := s3manager.NewUploader(session)
	aceBackend, err := backend.BuildRedshiftBackend(session.Config.Credentials, poolSize+healthCheckPoolSize, rsURL)
	if err != nil {
		logger.WithError(err).Fatal("Failed to setup redshift connection")
	}

	rsConnection, err := loadclient.NewRSLoader(s3Uploader, aceBackend, manifestBucket, stats)
	if err != nil {
		logger.WithError(err).Fatal("Failed to setup Redshift loading client for postgres")
	}

	initVersions, err := aceBackend.TableVersions()
	if err != nil {
		logger.WithError(err).Fatal("Failed initialization of table version cache")
	}
	tableVersions := versions.New(initVersions)

	var metaBackend metadata.Backend

	if poolSize > 0 {
		metaBackend, err = metadata.NewPostgresLoader(&pgConfig, rsConnection, tableVersions)
		if err != nil {
			logger.WithError(err).Fatal("Failed to setup postgres backend")
		}

		_, err = startWorkers(s3Uploader, metaBackend, stats, aceBackend)
		if err != nil {
			logger.WithError(err).Fatal("Failed to start workers")
		}
	}

	metaReader, err := metadata.NewPostgresReader(&pgConfig, tableVersions)
	if err != nil {
		logger.WithError(err).Fatal("Failed to setup postgres reader")
	}

	statsReporter := reporter.New(metaReader, stats, reporterPollPeriod)
	blueprintClient := blueprint.New(blueprintHost)
	versionIncrement := make(chan migrator.VersionIncrement)
	migrator := migrator.New(aceBackend, metaReader, blueprintClient, tableVersions, migratorPollPeriod,
		waitProcessorPeriod, offpeakStartHour, offpeakDurationHours, versionIncrement)

	hcBackend := healthcheck.NewBackend(rsConnection, metaReader)
	hcHandler := healthcheck.NewHandler(hcBackend)

	serveMux := http.NewServeMux()
	serveMux.Handle("/health", healthcheck.NewHealthRouter(hcHandler))

	controlBackend := control.NewControlBackend(metaReader, tableVersions, versionIncrement)
	controlHandler := control.NewControlHandler(controlBackend, stats)

	serveMux.Handle("/control/", control.NewControlRouter(controlHandler))

	logger.Go(func() {
		logger.WithError(http.ListenAndServe(net.JoinHostPort("localhost", "8080"), serveMux)).
			Fatal("Serving health and control failed")
	})

	logger.Go(func() {
		logger.WithError(http.ListenAndServe(":7766", http.DefaultServeMux)).
			Error("Serving pprof failed")
	})

	wait := make(chan struct{})
	sigc := make(chan os.Signal, 1)
	signal.Notify(sigc, syscall.SIGINT)
	logger.Info("Loader is set up")
	logger.Go(func() {
		<-sigc
		logger.Info("Sigint received -- shutting down")
		migrator.Close()
		statsReporter.Close()
		if metaBackend != nil {
			metaBackend.Close()
		}
		workerGroup.Wait()
		// Cause flush
		err = stats.Close()
		if err != nil {
			logger.WithError(err).Error("Error closing statter")
		}
		logger.Info("Exiting main cleanly.")
		logger.Wait()
		close(wait)
	})
	<-wait
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"time"
)

// A small programme written to get some understanding of how time zone handling is done in Go.
func main() {

	// Input, perhaps from user via an HTML form or some other means
	fromStr := "2013-07-09"
	toStr := "2013-11-10" // 24 hours will be added to this date below to make it exclusive (eg. include all day!)
	locationStr := "Europe/Copenhagen"

	// Get location, fall back to UTC
	location, err := time.LoadLocation(locationStr)
	if err != nil {
		location = time.UTC
	}

	// Parse dates, fall back to today for both dates if necessary

	from, err := time.ParseInLocation("2006-01-02", fromStr, location)
	if err != nil {
		fmt.Printf("Unable to parse from-date '%v'\n", fromStr)
		from, _ = time.ParseInLocation("2006-01-02", time.Now().Format("2006-01-02"), location)
	}

	to, err := time.ParseInLocation("2006-01-02", toStr, location)
	if err != nil {
		fmt.Printf("Unable to parse to date '%v'\n", toStr)
		to, _ = time.ParseInLocation("2006-01-02", time.Now().Format("2006-01-02"), location)
	}

	// Want all day of "to" - so add 24 hours and use it as exclusive when doing database queries for example.
	to = to.Add(24 * time.Hour)

	fmt.Printf("%11s %v\n%11s %v\n%11s %v.\n", "Period from", from, "to", to, "in", location)
	fmt.Printf("Notice offset. Daylight saving time means that %v can be either %v or %v.\n", location, from.Format("-0700"), to.Format("-0700"))
	fmt.Printf("In UTC this is from '%v' to '%v' which is useful for querying databases that stores dates in UTC.\n", from.UTC(), to.UTC())
}
<commit_msg>Alter UTC comment<commit_after>package main

import (
	"fmt"
	"time"
)

// A small programme written to get some understanding of how time zone handling is done in Go.
func main() {

	// Input, perhaps from user via an HTML form or some other means
	fromStr := "2013-07-09"
	toStr := "2013-11-10" // 24 hours will be added to this date below to make it exclusive (eg. include all day!)
	locationStr := "Europe/Copenhagen"

	// Get location, fall back to UTC
	location, err := time.LoadLocation(locationStr)
	if err != nil {
		location = time.UTC
	}

	// Parse dates, fall back to today for both dates if necessary

	from, err := time.ParseInLocation("2006-01-02", fromStr, location)
	if err != nil {
		fmt.Printf("Unable to parse from-date '%v'\n", fromStr)
		from, _ = time.ParseInLocation("2006-01-02", time.Now().Format("2006-01-02"), location)
	}

	to, err := time.ParseInLocation("2006-01-02", toStr, location)
	if err != nil {
		fmt.Printf("Unable to parse to date '%v'\n", toStr)
		to, _ = time.ParseInLocation("2006-01-02", time.Now().Format("2006-01-02"), location)
	}

	// Want all day of "to" - so add 24 hours and use it as exclusive when doing database queries for example.
	to = to.Add(24 * time.Hour)

	fmt.Printf("%11s %v\n%11s %v\n%11s %v.\n", "Period from", from, "to", to, "in", location)
	fmt.Printf("Notice offset. Daylight saving time means that %v can be either %v or %v.\n", location, from.Format("-0700"), to.Format("-0700"))
	fmt.Printf("In UTC these dates are from '%v' to '%v' which is useful for querying databases that stores dates in UTC.\n", from.UTC(), to.UTC())
}
<|endoftext|>