<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"strconv"
	"strings"
	"syscall"

	"github.com/datawire/teleproxy/pkg/k8s"
	"github.com/datawire/teleproxy/pkg/tpu"

	"github.com/datawire/teleproxy/internal/pkg/api"
	"github.com/datawire/teleproxy/internal/pkg/dns"
	"github.com/datawire/teleproxy/internal/pkg/docker"
	"github.com/datawire/teleproxy/internal/pkg/interceptor"
	"github.com/datawire/teleproxy/internal/pkg/proxy"
	"github.com/datawire/teleproxy/internal/pkg/route"
)

func dnsListeners(port string) (listeners []string) {
	// turns out you need to listen on localhost for nat to work
	// properly for udp, otherwise you get an "unexpected source
	// blah thingy" because the dns reply packets look like they
	// are coming from the wrong place
	listeners = append(listeners, "127.0.0.1:"+port)

	if runtime.GOOS == "linux" {
		// This is the default docker bridge. We should
		// probably figure out how to query this out of docker
		// instead of hardcoding it. We need to listen here
		// because the nat logic we use to intercept dns
		// packets will divert the packet to the interface it
		// originates from, which in the case of containers is
		// the docker bridge. Without this dns won't work from
		// inside containers.
		listeners = append(listeners, "172.17.0.1:"+port)
	}

	return
}

const (
	DEFAULT   = ""
	INTERCEPT = "intercept"
	BRIDGE    = "bridge"
)

func main() {
	var mode = flag.String("mode", "", "mode of operation")
	var kubeconfig = flag.String("kubeconfig", "", "absolute path to the kubeconfig file")
	var context = flag.String("context", "", "context to use (default: the current context)")
	var namespace = flag.String("namespace", "", "namespace to use (default: the current namespace for the context")
	var dnsIP = flag.String("dns", "", "dns ip address")
	var fallbackIP = flag.String("fallback", "", "dns fallback")

	flag.Parse()

	checkKubectl()

	switch *mode {
	case "":
	case INTERCEPT:
	case BRIDGE:
		break
	default:
		log.Fatalf("TPY: unrecognized mode: %v", *mode)
	}

	// do this up front so we don't miss out on cleanup if someone
	// Control-C's just after starting us
	signalChan := make(chan os.Signal)
	signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)

	if *mode == DEFAULT || *mode == INTERCEPT {
		shutdown := intercept(dnsIP, fallbackIP)
		defer shutdown()
	}
	if *mode == DEFAULT || *mode == BRIDGE {
		kubeinfo, err := k8s.NewKubeInfo(*kubeconfig, *context, *namespace)
		if err != nil {
			log.Fatalln("KubeInfo failed:", err)
		}
		shutdown := bridges(kubeinfo)
		defer shutdown()
	}

	log.Printf("TPY: %v", <-signalChan)
}

func kubeDie(err error) {
	if err != nil {
		log.Println(err)
	}
	log.Fatal("kubectl version 1.10 or greater is required")
}

func checkKubectl() {
	output, err := tpu.Shell("kubectl version --client -o json")
	if err != nil {
		kubeDie(err)
	}

	var info struct {
		ClientVersion struct {
			Major string
			Minor string
		}
	}

	err = json.Unmarshal([]byte(output), &info)
	if err != nil {
		kubeDie(err)
	}

	major, err := strconv.Atoi(info.ClientVersion.Major)
	if err != nil {
		kubeDie(err)
	}
	minor, err := strconv.Atoi(info.ClientVersion.Minor)
	if err != nil {
		kubeDie(err)
	}

	if major != 1 || minor < 10 {
		kubeDie(err)
	}
}

func intercept(dnsIP, fallbackIP *string) func() {
	if *dnsIP == "" {
		dat, err := ioutil.ReadFile("/etc/resolv.conf")
		if err != nil {
			panic(err)
		}
		for _, line := range strings.Split(string(dat), "\n") {
			if strings.Contains(line, "nameserver") {
				fields := strings.Fields(line)
				*dnsIP = fields[1]
				log.Printf("TPY: Automatically set -dns to %v", *dnsIP)
				break
			}
		}
	}

	if *dnsIP == "" {
		panic("couldn't determine dns ip from /etc/resolv.conf")
	}

	if *fallbackIP == "" {
		if *dnsIP == "8.8.8.8" {
			*fallbackIP = "8.8.4.4"
		} else {
			*fallbackIP = "8.8.8.8"
		}
	}

	if *fallbackIP == *dnsIP {
		panic("if your fallbackIP and your dnsIP are the same, you will have a dns loop")
	}

	iceptor := interceptor.NewInterceptor("teleproxy")

	apis, err := api.NewAPIServer(iceptor)
	if err != nil {
		panic(fmt.Sprintf("API Server: %v", err))
	}

	srv := dns.Server{
		Listeners: dnsListeners("1233"),
		Fallback:  *fallbackIP + ":53",
		Resolve: func(domain string) string {
			route := iceptor.Resolve(domain)
			if route != nil {
				return route.Ip
			} else {
				return ""
			}
		},
	}

	// hmm, we may not actually need to get the original
	// destination, we could just forward each ip to a unique port
	// and either listen on that port or run port-forward
	proxy, err := proxy.NewProxy(":1234", iceptor.Destination)
	if err != nil {
		panic(err)
	}

	bootstrap := route.Table{Name: "bootstrap"}
	bootstrap.Add(route.Route{
		Ip:     *dnsIP,
		Target: "1233",
		Proto:  "udp",
	})
	bootstrap.Add(route.Route{
		Name:   "teleproxy",
		Ip:     "127.254.254.254",
		Target: apis.Port(),
		Proto:  "tcp",
	})

	apis.Start()
	srv.Start()
	proxy.Start(10000)
	restore := dns.OverrideSearchDomains(".")

	iceptor.Start()
	iceptor.Update(bootstrap)

	return func() {
		// stop the api server first since it makes calls into
		// the interceptor
		apis.Stop()
		iceptor.Stop()
		restore()
		dns.Flush()
	}
}

func bridges(kubeinfo *k8s.KubeInfo) func() {
	disconnect := connect(kubeinfo)

	// setup kubernetes bridge
	log.Printf("BRG: kubernetes ctx=%s ns=%s", kubeinfo.Context, kubeinfo.Namespace)
	w := k8s.NewClient(kubeinfo).Watcher()
	w.Watch("services", func(w *k8s.Watcher) {
		table := route.Table{Name: "kubernetes"}
		for _, svc := range w.List("services") {
			ip, ok := svc.Spec()["clusterIP"]
			// for headless services the IP is None, we
			// should properly handle these by listening
			// for endpoints and returning multiple A
			// records at some point
			if ok && ip != "None" {
				qualName := svc.Name() + "." + svc.Namespace() + ".svc.cluster.local"
				table.Add(route.Route{
					Name:   qualName,
					Ip:     ip.(string),
					Proto:  "tcp",
					Target: "1234",
				})
			}
		}
		post(table)
	})

	// Set up DNS search path based on current Kubernetes namespace
	paths := []string{
		kubeinfo.Namespace + ".svc.cluster.local.",
		"svc.cluster.local.",
		"cluster.local.",
		"",
	}
	log.Println("BRG: Setting DNS search path:", paths[0])
	body, err := json.Marshal(paths)
	if err != nil {
		panic(err)
	}
	_, err = http.Post("http://teleproxy/api/search", "application/json", bytes.NewReader(body))
	if err != nil {
		log.Printf("BRG: error setting up search path: %v", err)
		panic(err) // Because this will fail if we win the startup race
	}

	// setup docker bridge
	dw := docker.NewWatcher()
	dw.Start(func(w *docker.Watcher) {
		table := route.Table{Name: "docker"}
		for name, ip := range w.Containers {
			table.Add(route.Route{Name: name, Ip: ip, Proto: "tcp"})
		}
		post(table)
	})

	return func() {
		dw.Stop()
		w.Stop()
		post(route.Table{Name: "kubernetes"}, route.Table{Name: "docker"})
		disconnect()
	}
}

func post(tables ...route.Table) {
	names := make([]string, len(tables))
	for i, t := range tables {
		names[i] = t.Name
	}
	jnames := strings.Join(names, ", ")

	body, err := json.Marshal(tables)
	if err != nil {
		panic(err)
	}
	resp, err := http.Post("http://teleproxy/api/tables/", "application/json", bytes.NewReader(body))
	if err != nil {
		log.Printf("BRG: error posting update to %s: %v", jnames, err)
	} else {
		log.Printf("BRG: posted update to %s: %v", jnames, resp.StatusCode)
	}
}

const TELEPROXY_POD = `
---
apiVersion: v1
kind: Pod
metadata:
  name: teleproxy
  labels:
    name: teleproxy
spec:
  containers:
  - name: proxy
    image: datawire/telepresence-k8s:0.75
    ports:
    - protocol: TCP
      containerPort: 8022
`

func connect(kubeinfo *k8s.KubeInfo) func() {
	// setup remote teleproxy pod
	apply := tpu.NewKeeper("KAP", "kubectl "+kubeinfo.GetKubectl("apply -f -"))
	apply.Input = TELEPROXY_POD
	apply.Limit = 1
	apply.Start()
	apply.Wait()

	pf := tpu.NewKeeper("KPF", "kubectl "+kubeinfo.GetKubectl("port-forward pod/teleproxy 8022"))
	pf.Inspect = "kubectl " + kubeinfo.GetKubectl("get pod/teleproxy")

	// XXX: probably need some kind of keepalive check for ssh, first
	// curl after wakeup seems to trigger detection of death
	ssh := tpu.NewKeeper("SSH", "ssh -D localhost:1080 -C -N -oConnectTimeout=5 -oExitOnForwardFailure=yes "+
		"-oStrictHostKeyChecking=no -oUserKnownHostsFile=/dev/null telepresence@localhost -p 8022")

	pf.Start()
	ssh.Start()

	return func() {
		ssh.Stop()
		pf.Stop()
	}
}
<commit_msg>teleproxy: intercept(): Clean up flow control a bit<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"strconv"
	"strings"
	"syscall"

	"github.com/pkg/errors"

	"github.com/datawire/teleproxy/pkg/k8s"
	"github.com/datawire/teleproxy/pkg/tpu"

	"github.com/datawire/teleproxy/internal/pkg/api"
	"github.com/datawire/teleproxy/internal/pkg/dns"
	"github.com/datawire/teleproxy/internal/pkg/docker"
	"github.com/datawire/teleproxy/internal/pkg/interceptor"
	"github.com/datawire/teleproxy/internal/pkg/proxy"
	"github.com/datawire/teleproxy/internal/pkg/route"
)

func dnsListeners(port string) (listeners []string) {
	// turns out you need to listen on localhost for nat to work
	// properly for udp, otherwise you get an "unexpected source
	// blah thingy" because the dns reply packets look like they
	// are coming from the wrong place
	listeners = append(listeners, "127.0.0.1:"+port)

	if runtime.GOOS == "linux" {
		// This is the default docker bridge. We should
		// probably figure out how to query this out of docker
		// instead of hardcoding it. We need to listen here
		// because the nat logic we use to intercept dns
		// packets will divert the packet to the interface it
		// originates from, which in the case of containers is
		// the docker bridge. Without this dns won't work from
		// inside containers.
		listeners = append(listeners, "172.17.0.1:"+port)
	}

	return
}

const (
	DEFAULT   = ""
	INTERCEPT = "intercept"
	BRIDGE    = "bridge"
)

func main() {
	var mode = flag.String("mode", "", "mode of operation")
	var kubeconfig = flag.String("kubeconfig", "", "absolute path to the kubeconfig file")
	var context = flag.String("context", "", "context to use (default: the current context)")
	var namespace = flag.String("namespace", "", "namespace to use (default: the current namespace for the context")
	var dnsIP = flag.String("dns", "", "dns ip address")
	var fallbackIP = flag.String("fallback", "", "dns fallback")

	flag.Parse()

	checkKubectl()

	switch *mode {
	case "":
	case INTERCEPT:
	case BRIDGE:
		break
	default:
		log.Fatalf("TPY: unrecognized mode: %v", *mode)
	}

	// do this up front so we don't miss out on cleanup if someone
	// Control-C's just after starting us
	signalChan := make(chan os.Signal)
	signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)

	if *mode == DEFAULT || *mode == INTERCEPT {
		shutdown, err := intercept(*dnsIP, *fallbackIP)
		if err != nil {
			log.Printf("TPY: Error: %v", err)
		}
		defer shutdown()
	}
	if *mode == DEFAULT || *mode == BRIDGE {
		kubeinfo, err := k8s.NewKubeInfo(*kubeconfig, *context, *namespace)
		if err != nil {
			log.Fatalln("KubeInfo failed:", err)
		}
		shutdown := bridges(kubeinfo)
		defer shutdown()
	}

	log.Printf("TPY: %v", <-signalChan)
}

func kubeDie(err error) {
	if err != nil {
		log.Println(err)
	}
	log.Fatal("kubectl version 1.10 or greater is required")
}

func checkKubectl() {
	output, err := tpu.Shell("kubectl version --client -o json")
	if err != nil {
		kubeDie(err)
	}

	var info struct {
		ClientVersion struct {
			Major string
			Minor string
		}
	}

	err = json.Unmarshal([]byte(output), &info)
	if err != nil {
		kubeDie(err)
	}

	major, err := strconv.Atoi(info.ClientVersion.Major)
	if err != nil {
		kubeDie(err)
	}
	minor, err := strconv.Atoi(info.ClientVersion.Minor)
	if err != nil {
		kubeDie(err)
	}

	if major != 1 || minor < 10 {
		kubeDie(err)
	}
}

// intercept starts the interceptor, and only returns once the
// interceptor is successfully running in another goroutine.  It
// returns a function to call to shut down that goroutine.
//
// If dnsIP is empty, it will be detected from /etc/resolv.conf
//
// If fallbackIP is empty, it will default to Google DNS.
func intercept(dnsIP string, fallbackIP string) (func(), error) {
	if dnsIP == "" {
		dat, err := ioutil.ReadFile("/etc/resolv.conf")
		if err != nil {
			return nil, err
		}
		for _, line := range strings.Split(string(dat), "\n") {
			if strings.Contains(line, "nameserver") {
				fields := strings.Fields(line)
				dnsIP = fields[1]
				log.Printf("TPY: Automatically set -dns=%v", dnsIP)
				break
			}
		}
	}
	if dnsIP == "" {
		return nil, errors.New("couldn't determine dns ip from /etc/resolv.conf")
	}

	if fallbackIP == "" {
		if dnsIP == "8.8.8.8" {
			fallbackIP = "8.8.4.4"
		} else {
			fallbackIP = "8.8.8.8"
		}
		log.Printf("TPY: Automatically set -fallback=%v", dnsIP)
	}
	if fallbackIP == dnsIP {
		return nil, errors.New("if your fallbackIP and your dnsIP are the same, you will have a dns loop")
	}

	iceptor := interceptor.NewInterceptor("teleproxy")

	apis, err := api.NewAPIServer(iceptor)
	if err != nil {
		return nil, errors.Wrap(err, "API Server")
	}

	srv := dns.Server{
		Listeners: dnsListeners("1233"),
		Fallback:  fallbackIP + ":53",
		Resolve: func(domain string) string {
			route := iceptor.Resolve(domain)
			if route != nil {
				return route.Ip
			} else {
				return ""
			}
		},
	}

	// hmm, we may not actually need to get the original
	// destination, we could just forward each ip to a unique port
	// and either listen on that port or run port-forward
	proxy, err := proxy.NewProxy(":1234", iceptor.Destination)
	if err != nil {
		return nil, errors.Wrap(err, "Proxy")
	}

	bootstrap := route.Table{Name: "bootstrap"}
	bootstrap.Add(route.Route{
		Ip:     dnsIP,
		Target: "1233",
		Proto:  "udp",
	})
	bootstrap.Add(route.Route{
		Name:   "teleproxy",
		Ip:     "127.254.254.254",
		Target: apis.Port(),
		Proto:  "tcp",
	})

	apis.Start()
	srv.Start()
	proxy.Start(10000)
	restore := dns.OverrideSearchDomains(".")

	iceptor.Start()
	iceptor.Update(bootstrap)

	return func() {
		// stop the api server first since it makes calls into
		// the interceptor
		apis.Stop()
		iceptor.Stop()
		restore()
		dns.Flush()
	}, nil
}

func bridges(kubeinfo *k8s.KubeInfo) func() {
	disconnect := connect(kubeinfo)

	// setup kubernetes bridge
	log.Printf("BRG: kubernetes ctx=%s ns=%s", kubeinfo.Context, kubeinfo.Namespace)
	w := k8s.NewClient(kubeinfo).Watcher()
	w.Watch("services", func(w *k8s.Watcher) {
		table := route.Table{Name: "kubernetes"}
		for _, svc := range w.List("services") {
			ip, ok := svc.Spec()["clusterIP"]
			// for headless services the IP is None, we
			// should properly handle these by listening
			// for endpoints and returning multiple A
			// records at some point
			if ok && ip != "None" {
				qualName := svc.Name() + "." + svc.Namespace() + ".svc.cluster.local"
				table.Add(route.Route{
					Name:   qualName,
					Ip:     ip.(string),
					Proto:  "tcp",
					Target: "1234",
				})
			}
		}
		post(table)
	})

	// Set up DNS search path based on current Kubernetes namespace
	paths := []string{
		kubeinfo.Namespace + ".svc.cluster.local.",
		"svc.cluster.local.",
		"cluster.local.",
		"",
	}
	log.Println("BRG: Setting DNS search path:", paths[0])
	body, err := json.Marshal(paths)
	if err != nil {
		panic(err)
	}
	_, err = http.Post("http://teleproxy/api/search", "application/json", bytes.NewReader(body))
	if err != nil {
		log.Printf("BRG: error setting up search path: %v", err)
		panic(err) // Because this will fail if we win the startup race
	}

	// setup docker bridge
	dw := docker.NewWatcher()
	dw.Start(func(w *docker.Watcher) {
		table := route.Table{Name: "docker"}
		for name, ip := range w.Containers {
			table.Add(route.Route{Name: name, Ip: ip, Proto: "tcp"})
		}
		post(table)
	})

	return func() {
		dw.Stop()
		w.Stop()
		post(route.Table{Name: "kubernetes"}, route.Table{Name: "docker"})
		disconnect()
	}
}

func post(tables ...route.Table) {
	names := make([]string, len(tables))
	for i, t := range tables {
		names[i] = t.Name
	}
	jnames := strings.Join(names, ", ")

	body, err := json.Marshal(tables)
	if err != nil {
		panic(err)
	}
	resp, err := http.Post("http://teleproxy/api/tables/", "application/json", bytes.NewReader(body))
	if err != nil {
		log.Printf("BRG: error posting update to %s: %v", jnames, err)
	} else {
		log.Printf("BRG: posted update to %s: %v", jnames, resp.StatusCode)
	}
}

const TELEPROXY_POD = `
---
apiVersion: v1
kind: Pod
metadata:
  name: teleproxy
  labels:
    name: teleproxy
spec:
  containers:
  - name: proxy
    image: datawire/telepresence-k8s:0.75
    ports:
    - protocol: TCP
      containerPort: 8022
`

func connect(kubeinfo *k8s.KubeInfo) func() {
	// setup remote teleproxy pod
	apply := tpu.NewKeeper("KAP", "kubectl "+kubeinfo.GetKubectl("apply -f -"))
	apply.Input = TELEPROXY_POD
	apply.Limit = 1
	apply.Start()
	apply.Wait()

	pf := tpu.NewKeeper("KPF", "kubectl "+kubeinfo.GetKubectl("port-forward pod/teleproxy 8022"))
	pf.Inspect = "kubectl " + kubeinfo.GetKubectl("get pod/teleproxy")

	// XXX: probably need some kind of keepalive check for ssh, first
	// curl after wakeup seems to trigger detection of death
	ssh := tpu.NewKeeper("SSH", "ssh -D localhost:1080 -C -N -oConnectTimeout=5 -oExitOnForwardFailure=yes "+
		"-oStrictHostKeyChecking=no -oUserKnownHostsFile=/dev/null telepresence@localhost -p 8022")

	pf.Start()
	ssh.Start()

	return func() {
		ssh.Stop()
		pf.Stop()
	}
}
<|endoftext|>


<commit_before>package main_test

import (
	"fmt"

	"github.com/cloudfoundry-incubator/tps/cmd/tps/testrunner"
	"github.com/cloudfoundry/gunk/diegonats"
	"github.com/cloudfoundry/gunk/timeprovider/faketimeprovider"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
	"github.com/onsi/gomega/ghttp"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"

	"testing"
	"time"
)

var timeProvider *faketimeprovider.FakeTimeProvider

var tpsAddr string
var tps ifrit.Process
var runner *ginkgomon.Runner

var natsPort int
var gnatsdRunner ifrit.Process
var natsClient diegonats.NATSClient
var receptorServer *ghttp.Server

var heartbeatInterval = 50 * time.Millisecond
var tpsBinPath string

var _ = SynchronizedBeforeSuite(func() []byte {
	synchronizedTpsBinPath, err := gexec.Build("github.com/cloudfoundry-incubator/tps/cmd/tps", "-race")
	Ω(err).ShouldNot(HaveOccurred())
	return []byte(synchronizedTpsBinPath)
}, func(synchronizedTpsBinPath []byte) {
	tpsBinPath = string(synchronizedTpsBinPath)
})

func TestTPS(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "TPS Suite")
}

var _ = BeforeEach(func() {
	tpsAddr = fmt.Sprintf("127.0.0.1:%d", uint16(1518+GinkgoParallelNode()))
	natsPort = 4001 + GinkgoParallelNode()

	timeProvider = faketimeprovider.New(time.Unix(0, 1138))
	receptorServer = ghttp.NewServer()

	runner = testrunner.New(
		string(tpsBinPath),
		tpsAddr,
		receptorServer.URL(),
		[]string{fmt.Sprintf("127.0.0.1:%d", natsPort)},
		heartbeatInterval,
	)

	startAll()
})

var _ = AfterEach(func() {
	stopAll()
})

var _ = SynchronizedAfterSuite(func() {
	stopAll()
}, func() {
	gexec.CleanupBuildArtifacts()
})

func startAll() {
	gnatsdRunner, natsClient = diegonats.StartGnatsd(natsPort)
}

func stopAll() {
	ginkgomon.Kill(gnatsdRunner)
}
<commit_msg>timeprovider -> clock<commit_after>package main_test

import (
	"fmt"

	"github.com/cloudfoundry-incubator/tps/cmd/tps/testrunner"
	"github.com/cloudfoundry/gunk/diegonats"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
	"github.com/onsi/gomega/ghttp"
	"github.com/pivotal-golang/clock/fakeclock"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"

	"testing"
	"time"
)

var clock *fakeclock.FakeClock

var tpsAddr string
var tps ifrit.Process
var runner *ginkgomon.Runner

var natsPort int
var gnatsdRunner ifrit.Process
var natsClient diegonats.NATSClient
var receptorServer *ghttp.Server

var heartbeatInterval = 50 * time.Millisecond
var tpsBinPath string

var _ = SynchronizedBeforeSuite(func() []byte {
	synchronizedTpsBinPath, err := gexec.Build("github.com/cloudfoundry-incubator/tps/cmd/tps", "-race")
	Ω(err).ShouldNot(HaveOccurred())
	return []byte(synchronizedTpsBinPath)
}, func(synchronizedTpsBinPath []byte) {
	tpsBinPath = string(synchronizedTpsBinPath)
})

func TestTPS(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "TPS Suite")
}

var _ = BeforeEach(func() {
	tpsAddr = fmt.Sprintf("127.0.0.1:%d", uint16(1518+GinkgoParallelNode()))
	natsPort = 4001 + GinkgoParallelNode()

	clock = fakeclock.NewFakeClock(time.Unix(0, 1138))
	receptorServer = ghttp.NewServer()

	runner = testrunner.New(
		string(tpsBinPath),
		tpsAddr,
		receptorServer.URL(),
		[]string{fmt.Sprintf("127.0.0.1:%d", natsPort)},
		heartbeatInterval,
	)

	startAll()
})

var _ = AfterEach(func() {
	stopAll()
})

var _ = SynchronizedAfterSuite(func() {
	stopAll()
}, func() {
	gexec.CleanupBuildArtifacts()
})

func startAll() {
	gnatsdRunner, natsClient = diegonats.StartGnatsd(natsPort)
}

func stopAll() {
	ginkgomon.Kill(gnatsdRunner)
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package vfs

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"strings"
)

// VFSContext is a 'context' for VFS, that is normally a singleton
// but allows us to configure S3 credentials, for example
type VFSContext struct {
	s3Context    *S3Context
	memfsContext *MemFSContext
}

var Context = VFSContext{
	s3Context: NewS3Context(),
}

// ReadLocation reads a file from a vfs URL
// It supports additional schemes which don't (yet) have full VFS implementations:
//   metadata: reads from instance metadata on GCE/AWS
//   http / https: reads from HTTP
func (c *VFSContext) ReadFile(location string) ([]byte, error) {
	if strings.Contains(location, "://") {
		// Handle our special case schemas
		u, err := url.Parse(location)
		if err != nil {
			return nil, fmt.Errorf("error parsing location %q - not a valid URI", location)
		}

		switch u.Scheme {
		case "metadata":
			switch u.Host {
			case "gce":
				httpURL := "http://169.254.169.254/computeMetadata/v1/instance/attributes/" + u.Path
				httpHeaders := make(map[string]string)
				httpHeaders["Metadata-Flavor"] = "Google"
				return c.readHttpLocation(httpURL, httpHeaders)
			case "aws":
				httpURL := "http://169.254.169.254/latest/" + u.Path
				return c.readHttpLocation(httpURL, nil)

			default:
				return nil, fmt.Errorf("unknown metadata type: %q in %q", u.Host, location)
			}

		case "http", "https":
			return c.readHttpLocation(location, nil)
		}

	}

	p, err := c.BuildVfsPath(location)
	if err != nil {
		return nil, err
	}
	return p.ReadFile()
}

func (c *VFSContext) BuildVfsPath(p string) (Path, error) {
	if !strings.Contains(p, "://") {
		return NewFSPath(p), nil
	}

	if strings.HasPrefix(p, "s3://") {
		return c.buildS3Path(p)
	}

	if strings.HasPrefix(p, "memfs://") {
		return c.buildMemFSPath(p)
	}

	return nil, fmt.Errorf("unknown / unhandled path type: %q", p)
}

// readHttpLocation reads an http (or https) url.
// It returns the contents, or an error on any non-200 response.  On a 404, it will return os.ErrNotExist
func (c *VFSContext) readHttpLocation(httpURL string, httpHeaders map[string]string) ([]byte, error) {
	req, err := http.NewRequest("GET", httpURL, nil)
	if err != nil {
		return nil, err
	}
	for k, v := range httpHeaders {
		req.Header.Add(k, v)
	}
	response, err := http.DefaultClient.Do(req)
	if response != nil {
		defer response.Body.Close()
	}
	if err != nil {
		return nil, fmt.Errorf("error fetching %q: %v", httpURL, err)
	}
	body, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, fmt.Errorf("error reading response for %q: %v", httpURL, err)
	}
	if response.StatusCode == 404 {
		return nil, os.ErrNotExist
	}
	if response.StatusCode != 200 {
		return nil, fmt.Errorf("unexpected response code %q for %q: %v", response.Status, httpURL, string(body))
	}
	return body, nil
}

func (c *VFSContext) buildS3Path(p string) (*S3Path, error) {
	u, err := url.Parse(p)
	if err != nil {
		return nil, fmt.Errorf("invalid s3 path: %q", err)
	}

	bucket := strings.TrimSuffix(u.Host, "/")
	if bucket == "" {
		return nil, fmt.Errorf("invalid s3 path: %q", err)
	}

	s3path := newS3Path(c.s3Context, bucket, u.Path)
	return s3path, nil
}

func (c *VFSContext) buildMemFSPath(p string) (*MemFSPath, error) {
	if !strings.HasPrefix(p, "memfs://") {
		return nil, fmt.Errorf("memfs path not recognized: %q", p)
	}
	location := strings.TrimPrefix(p, "memfs://")
	if c.memfsContext == nil {
		// We only initialize this in unit tests etc
		return nil, fmt.Errorf("memfs context not initialized")
	}
	fspath := NewMemFSPath(c.memfsContext, location)
	return fspath, nil
}

func (c *VFSContext) ResetMemfsContext(clusterReadable bool) {
	c.memfsContext = NewMemFSContext()
	if clusterReadable {
		c.memfsContext.MarkClusterReadable()
	}
}
<commit_msg>Retry readHttpLocation on 500 errors<commit_after>/*
Copyright 2016 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package vfs

import (
	"fmt"
	"github.com/golang/glog"
	"io/ioutil"
	"k8s.io/kubernetes/pkg/util/wait"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
)

// VFSContext is a 'context' for VFS, that is normally a singleton
// but allows us to configure S3 credentials, for example
type VFSContext struct {
	s3Context    *S3Context
	memfsContext *MemFSContext
}

var Context = VFSContext{
	s3Context: NewS3Context(),
}

// ReadLocation reads a file from a vfs URL
// It supports additional schemes which don't (yet) have full VFS implementations:
//   metadata: reads from instance metadata on GCE/AWS
//   http / https: reads from HTTP
func (c *VFSContext) ReadFile(location string) ([]byte, error) {
	if strings.Contains(location, "://") {
		// Handle our special case schemas
		u, err := url.Parse(location)
		if err != nil {
			return nil, fmt.Errorf("error parsing location %q - not a valid URI", location)
		}

		switch u.Scheme {
		case "metadata":
			switch u.Host {
			case "gce":
				httpURL := "http://169.254.169.254/computeMetadata/v1/instance/attributes/" + u.Path
				httpHeaders := make(map[string]string)
				httpHeaders["Metadata-Flavor"] = "Google"
				return c.readHttpLocation(httpURL, httpHeaders)
			case "aws":
				httpURL := "http://169.254.169.254/latest/" + u.Path
				return c.readHttpLocation(httpURL, nil)

			default:
				return nil, fmt.Errorf("unknown metadata type: %q in %q", u.Host, location)
			}

		case "http", "https":
			return c.readHttpLocation(location, nil)
		}

	}

	p, err := c.BuildVfsPath(location)
	if err != nil {
		return nil, err
	}
	return p.ReadFile()
}

func (c *VFSContext) BuildVfsPath(p string) (Path, error) {
	if !strings.Contains(p, "://") {
		return NewFSPath(p), nil
	}

	if strings.HasPrefix(p, "s3://") {
		return c.buildS3Path(p)
	}

	if strings.HasPrefix(p, "memfs://") {
		return c.buildMemFSPath(p)
	}

	return nil, fmt.Errorf("unknown / unhandled path type: %q", p)
}

// readHttpLocation reads an http (or https) url.
// It returns the contents, or an error on any non-200 response.  On a 404, it will return os.ErrNotExist
// It will retry a few times on a 500 class error
func (c *VFSContext) readHttpLocation(httpURL string, httpHeaders map[string]string) ([]byte, error) {
	// Exponential backoff, starting with 500 milliseconds, doubling each time, 5 steps
	backoff := wait.Backoff{
		Duration: 500 * time.Millisecond,
		Factor:   2,
		Steps:    5,
	}

	var body []byte

	done, err := RetryWithBackoff(backoff, func() (bool, error) {
		glog.V(4).Infof("Performing HTTP request: GET %s", httpURL)
		req, err := http.NewRequest("GET", httpURL, nil)
		if err != nil {
			return false, err
		}
		for k, v := range httpHeaders {
			req.Header.Add(k, v)
		}
		response, err := http.DefaultClient.Do(req)
		if response != nil {
			defer response.Body.Close()
		}
		if err != nil {
			return false, fmt.Errorf("error fetching %q: %v", httpURL, err)
		}
		body, err = ioutil.ReadAll(response.Body)
		if err != nil {
			return false, fmt.Errorf("error reading response for %q: %v", httpURL, err)
		}
		if response.StatusCode == 404 {
			// We retry on 404s in case of eventual consistency
			return true, os.ErrNotExist
		}
		if response.StatusCode >= 500 && response.StatusCode <= 599 {
			// Retry on 5XX errors
			return false, fmt.Errorf("unexpected response code %q for %q: %v", response.Status, httpURL, string(body))
		}

		if response.StatusCode == 200 {
			return true, nil
		}

		// Don't retry on other errors
		return true, fmt.Errorf("unexpected response code %q for %q: %v", response.Status, httpURL, string(body))
	})
	if err != nil {
		return nil, err
	} else if done {
		return body, nil
	} else {
		// Shouldn't happen - we always return a non-nil error with false
		return nil, wait.ErrWaitTimeout
	}
}

// RetryWithBackoff runs until a condition function returns true, or until Steps attempts have been taken
// As compared to wait.ExponentialBackoff, this function returns the results from the function on the final attempt
func RetryWithBackoff(backoff wait.Backoff, condition func() (bool, error)) (bool, error) {
	duration := backoff.Duration
	i := 0
	for {
		if i != 0 {
			adjusted := duration
			if backoff.Jitter > 0.0 {
				adjusted = wait.Jitter(duration, backoff.Jitter)
			}
			time.Sleep(adjusted)
			duration = time.Duration(float64(duration) * backoff.Factor)
		}

		i++

		ok, err := condition()
		if ok {
			return ok, err
		}
		noMoreRetries := (i + 1) >= backoff.Steps
		if !noMoreRetries && err != nil {
			glog.V(2).Infof("retrying after error %v", err)
		}

		if noMoreRetries {
			return ok, err
		}
	}
}

func (c *VFSContext) buildS3Path(p string) (*S3Path, error) {
	u, err := url.Parse(p)
	if err != nil {
		return nil, fmt.Errorf("invalid s3 path: %q", err)
	}

	bucket := strings.TrimSuffix(u.Host, "/")
	if bucket == "" {
		return nil, fmt.Errorf("invalid s3 path: %q", err)
	}

	s3path := newS3Path(c.s3Context, bucket, u.Path)
	return s3path, nil
}

func (c *VFSContext) buildMemFSPath(p string) (*MemFSPath, error) {
	if !strings.HasPrefix(p, "memfs://") {
		return nil, fmt.Errorf("memfs path not recognized: %q", p)
	}
	location := strings.TrimPrefix(p, "memfs://")
	if c.memfsContext == nil {
		// We only initialize this in unit tests etc
		return nil, fmt.Errorf("memfs context not initialized")
	}
	fspath := NewMemFSPath(c.memfsContext, location)
	return fspath, nil
}

func (c *VFSContext) ResetMemfsContext(clusterReadable bool) {
	c.memfsContext = NewMemFSContext()
	if clusterReadable {
		c.memfsContext.MarkClusterReadable()
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 The roc Author. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package rocserv

import (
	"flag"
	"fmt"
	"github.com/shawnfeng/roc/util/service/sla"
	"github.com/shawnfeng/sutil/trace"
	"reflect"

	"github.com/julienschmidt/httprouter"

	"git.apache.org/thrift.git/lib/go/thrift"

	"github.com/shawnfeng/sutil/slog"
)

const (
	PROCESSOR_HTTP   = "http"
	PROCESSOR_THRIFT = "thrift"
	PROCESSOR_GRPC   = "gprc"
)

type Service struct {
}

var service Service

//func NewService

type cmdArgs struct {
	servLoc string
	logDir  string
	sessKey string
}

func (m *Service) parseFlag() (*cmdArgs, error) {
	var serv, logDir, skey string
	flag.StringVar(&serv, "serv", "", "servic name")
	flag.StringVar(&logDir, "logdir", "", "serice log dir")
	flag.StringVar(&skey, "skey", "", "service session key")

	flag.Parse()

	if len(serv) == 0 {
		return nil, fmt.Errorf("serv args need!")
	}

	if len(skey) == 0 {
		return nil, fmt.Errorf("skey args need!")
	}

	return &cmdArgs{
		servLoc: serv,
		logDir:  logDir,
		sessKey: skey,
	}, nil

}

func (m *Service) loadDriver(sb ServBase, procs map[string]Processor) (map[string]*ServInfo, error) {
	fun := "Service.loadDriver -->"

	infos := make(map[string]*ServInfo)

	// load processor's driver
	for n, p := range procs {
		addr, driver := p.Driver()
		if driver == nil {
			slog.Infof("%s processor:%s no driver", fun, n)
			continue
		}

		slog.Infof("%s processor:%s type:%s addr:%s", fun, n, reflect.TypeOf(driver), addr)

		switch d := driver.(type) {
		case *httprouter.Router:
			sa, err := powerHttp(addr, d)
			if err != nil {
				return nil, err
			}

			slog.Infof("%s load ok processor:%s serv addr:%s", fun, n, sa)
			infos[n] = &ServInfo{
				Type: PROCESSOR_HTTP,
				Addr: sa,
			}

		case thrift.TProcessor:
			sa, err := powerThrift(addr, d)
			if err != nil {
				return nil, err
			}

			slog.Infof("%s load ok processor:%s serv addr:%s", fun, n, sa)
			infos[n] = &ServInfo{
				Type: PROCESSOR_THRIFT,
				Addr: sa,
			}
		case *GrpcServer:
			sa, err := powerGrpc(addr, d)
			if err != nil {
				return nil, err
			}

			slog.Infof("%s load ok processor:%s serv addr:%s", fun, n, sa)
			infos[n] = &ServInfo{
				Type: PROCESSOR_GRPC,
				Addr: sa,
			}
		default:
			return nil, fmt.Errorf("processor:%s driver not recognition", n)

		}

	}

	return infos, nil

}

func (m *Service) Serve(confEtcd configEtcd, initfn func(ServBase) error, procs map[string]Processor) error {
	fun := "Service.Serve -->"

	args, err := m.parseFlag()
	if err != nil {
		slog.Panicf("%s parse arg err:%s", fun, err)
		return err
	}

	return m.Init(confEtcd, args.servLoc, args.sessKey, args.logDir, initfn, procs)
}

func (m *Service) Init(confEtcd configEtcd, servLoc, sessKey, logDir string, initfn func(ServBase) error, procs map[string]Processor) error {
	fun := "Service.Init -->"
	// Init ServBase

	sb, err := NewServBaseV2(confEtcd, servLoc, sessKey)
	if err != nil {
		slog.Panicf("%s init servbase loc:%s key:%s err:%s", fun, servLoc, sessKey, err)
		return err
	}

	// Init slog
	var logConfig struct {
		Log struct {
			Level string
			Dir   string
		}
	}
	// default use level INFO
	logConfig.Log.Level = "INFO"

	err = sb.ServConfig(&logConfig)
	if err != nil {
		slog.Panicf("%s serv config err:%s", fun, err)
		fmt.Sprintf("%s serv config err:%s", fun, err)
		return err
	}

	var logdir string
	if len(logConfig.Log.Dir) > 0 {
		logdir = fmt.Sprintf("%s/%s", logConfig.Log.Dir, sb.Copyname())
	}

	if len(logDir) > 0 {
		logdir = fmt.Sprintf("%s/%s", logDir, sb.Copyname())
	}

	if logDir == "console" {
		logdir = ""
	}

	fmt.Sprintf("%s init log dir:%s name:%s level:%s", fun, logdir, servLoc, logConfig.Log.Level)
	slog.Infof("%s init log dir:%s name:%s level:%s", fun, logdir, servLoc, logConfig.Log.Level)

	slog.Init(logdir, "serv.log", logConfig.Log.Level)
	defer slog.Sync()

	// init tracer
	err = trace.InitDefaultTracer(servLoc)
	if err != nil {
		slog.Warnf("%s init tracer fail:%v", err)
	}

	// sla metric埋点 ==================
	//init metric
	//user defualt metric opts
	metrics := rocserv.NewMetricsprocessor()
	if err != nil {
		slog.Warnf("init metrics fail:%v", err)
	}
	err = metrics.Init()
	if err != nil {
		slog.Warnf("%s init metrics err:%s", fun, err)
	}

	minfos, err := m.loadDriver(sb, map[string]Processor{"_PROC_METRICS": metrics})
	if err == nil {
		err = sb.RegisterMetrics(minfos)
		if err != nil {
			slog.Warnf("%s regist backdoor err:%s", fun, err)
		}

	} else {
		slog.Warnf("%s load metrics driver err:%s", fun, err)
	}
	//==============================

	// init callback
	err = initfn(sb)
	if err != nil {
		slog.Panicf("%s serv init err:%s", fun, err)
		return err
	}

	// init processor
	for n, p := range procs {
		if len(n) == 0 {
			slog.Panicf("%s processor name empty", fun)
			return err
		}

		if n[0] == '_' {
			slog.Panicf("%s processor name can not prefix '_'", fun)
			return err
		}

		if p == nil {
			slog.Panicf("%s processor:%s is nil", fun, n)
			return fmt.Errorf("processor:%s is nil", n)
		} else {
			err := p.Init()
			if err != nil {
				slog.Panicf("%s processor:%s init err:%s", fun, err)
				return err
			}
		}
	}

	infos, err := m.loadDriver(sb, procs)
	if err != nil {
		slog.Panicf("%s load driver err:%s", fun, err)
		return err
	}

	err = sb.RegisterService(infos)
	if err != nil {
		slog.Panicf("%s regist service err:%s", fun, err)
		return err
	}

	// 后门接口 ==================
	backdoor := &backDoorHttp{}
	err = backdoor.Init()
	if err != nil {
		slog.Panicf("%s init backdoor err:%s", fun, err)
		return err
	}

	binfos, err := m.loadDriver(sb, map[string]Processor{"_PROC_BACKDOOR": backdoor})
	if err == nil {
		err = sb.RegisterBackDoor(binfos)
		if err != nil {
			slog.Errorf("%s regist backdoor err:%s", fun, err)
		}

	} else {
		slog.Warnf("%s load backdoor driver err:%s", fun, err)
	}
	//==============================

	// pause here
	var pause chan bool
	pause <- true

	return nil

}
func (m *Service) getMetricOps(sb *ServBaseV2) *rocserv.MetricsOpts {
	fun := "Service.getMetricOps -->"
	var metricConfig struct {
		metric *rocserv.MetricsOpts
	}
	err := sb.ServConfig(&metricConfig)
	if err != nil {
		slog.Panicf("%s serv config err:%s", fun, err)
		fmt.Sprintf("%s serv config err:%s", fun, err)
		return nil
	}
	return metricConfig.metric
}

func Serve(etcds []string, baseLoc string, initfn func(ServBase) error, procs map[string]Processor) error {
	return service.Serve(configEtcd{etcds, baseLoc}, initfn, procs)
}

func Init(etcds []string, baseLoc string, servLoc, servKey string, initfn func(ServBase) error, procs map[string]Processor) error {
	return service.Init(configEtcd{etcds, baseLoc}, servLoc, servKey, "", initfn, procs)
}

func Test(etcds []string, baseLoc string, initfn func(ServBase) error) error {
	return service.Init(configEtcd{etcds, baseLoc}, "test/test", "test", "console", initfn, nil)
}
<commit_msg>move metric after initlogic<commit_after>// Copyright 2014 The roc Author. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package rocserv

import (
	"flag"
	"fmt"
	"github.com/shawnfeng/roc/util/service/sla"
	"github.com/shawnfeng/sutil/trace"
	"reflect"

	"github.com/julienschmidt/httprouter"

	"git.apache.org/thrift.git/lib/go/thrift"

	"github.com/shawnfeng/sutil/slog"
)

const (
	PROCESSOR_HTTP   = "http"
	PROCESSOR_THRIFT = "thrift"
	PROCESSOR_GRPC   = "gprc"
)

type Service struct {
}

var service Service

//func NewService

type cmdArgs struct {
	servLoc string
	logDir  string
	sessKey string
}

func (m *Service) parseFlag() (*cmdArgs, error) {
	var serv, logDir, skey string
	flag.StringVar(&serv, "serv", "", "servic name")
	flag.StringVar(&logDir, "logdir", "", "serice log dir")
	flag.StringVar(&skey, "skey", "", "service session key")

	flag.Parse()

	if len(serv) == 0 {
		return nil, fmt.Errorf("serv args need!")
	}

	if len(skey) == 0 {
		return nil, fmt.Errorf("skey args need!")
	}

	return &cmdArgs{
		servLoc: serv,
		logDir:  logDir,
		sessKey: skey,
	}, nil

}

func (m *Service) loadDriver(sb ServBase, procs map[string]Processor) (map[string]*ServInfo, error) {
	fun := "Service.loadDriver -->"

	infos := make(map[string]*ServInfo)

	// load processor's driver
	for n, p := range procs {
		addr, driver := p.Driver()
		if driver == nil {
			slog.Infof("%s processor:%s no driver", fun, n)
			continue
		}

		slog.Infof("%s processor:%s type:%s addr:%s", fun, n, reflect.TypeOf(driver), addr)

		switch d := driver.(type) {
		case *httprouter.Router:
			sa, err := powerHttp(addr, d)
			if err != nil {
				return nil, err
			}

			slog.Infof("%s load ok processor:%s serv addr:%s", fun, n, sa)
			infos[n] = &ServInfo{
				Type: PROCESSOR_HTTP,
				Addr: sa,
			}

		case thrift.TProcessor:
			sa, err := powerThrift(addr, d)
			if err != nil {
				return nil, err
			}

			slog.Infof("%s load ok processor:%s serv addr:%s", fun, n, sa)
			infos[n] = &ServInfo{
				Type: PROCESSOR_THRIFT,
				Addr: sa,
			}
		case *GrpcServer:
			sa, err := powerGrpc(addr, d)
			if err != nil {
				return nil, err
			}

			slog.Infof("%s load ok processor:%s serv addr:%s", fun, n, sa)
			infos[n] = &ServInfo{
				Type: PROCESSOR_GRPC,
				Addr: sa,
			}
		default:
			return nil, fmt.Errorf("processor:%s driver not recognition", n)

		}

	}

	return infos, nil

}

func (m *Service) Serve(confEtcd configEtcd, initfn func(ServBase) error, procs map[string]Processor) error {
	fun := "Service.Serve -->"

	args, err := m.parseFlag()
	if err != nil {
		slog.Panicf("%s parse arg err:%s", fun, err)
		return err
	}

	return m.Init(confEtcd, args.servLoc, args.sessKey, args.logDir, initfn, procs)
}

func (m *Service) Init(confEtcd configEtcd, servLoc, sessKey, logDir string, initfn func(ServBase) error, procs map[string]Processor) error {
	fun := "Service.Init -->"
	// Init ServBase

	sb, err := NewServBaseV2(confEtcd, servLoc, sessKey)
	if err != nil {
		slog.Panicf("%s init servbase loc:%s key:%s err:%s", fun, servLoc, sessKey, err)
		return err
	}

	// Init slog
	var logConfig struct {
		Log struct {
			Level string
			Dir   string
		}
	}
	// default use level INFO
	logConfig.Log.Level = "INFO"

	err = sb.ServConfig(&logConfig)
	if err != nil {
		slog.Panicf("%s serv config err:%s", fun, err)
		fmt.Sprintf("%s serv config err:%s", fun, err)
		return err
	}

	var logdir string
	if len(logConfig.Log.Dir) > 0 {
		logdir = fmt.Sprintf("%s/%s", logConfig.Log.Dir, sb.Copyname())
	}

	if len(logDir) > 0 {
		logdir = fmt.Sprintf("%s/%s", logDir, sb.Copyname())
	}

	if logDir == "console" {
		logdir = ""
	}

	fmt.Sprintf("%s init log dir:%s name:%s level:%s", fun, logdir, servLoc, logConfig.Log.Level)
	slog.Infof("%s init log dir:%s name:%s level:%s", fun, logdir, servLoc, logConfig.Log.Level)

	slog.Init(logdir, "serv.log", logConfig.Log.Level)
	defer slog.Sync()

	// init tracer
	err = trace.InitDefaultTracer(servLoc)
	if err != nil {
		slog.Warnf("%s init tracer fail:%v", err)
	}

	// init callback
	err = initfn(sb)
	if err != nil {
		slog.Panicf("%s serv init err:%s", fun, err)
		return err
	}

	// init processor
	for n, p := range procs {
		if len(n) == 0 {
			slog.Panicf("%s processor name empty", fun)
			return err
		}

		if n[0] == '_' {
			slog.Panicf("%s processor name can not prefix '_'", fun)
			return err
		}

		if p == nil {
			slog.Panicf("%s processor:%s is nil", fun, n)
			return fmt.Errorf("processor:%s is nil", n)
		} else {
			err := p.Init()
			if err != nil {
				slog.Panicf("%s processor:%s init err:%s", fun, err)
				return err
			}
		}
	}

	infos, err := m.loadDriver(sb, procs)
	if err != nil {
		slog.Panicf("%s load driver err:%s", fun, err)
		return err
	}

	err = sb.RegisterService(infos)
	if err != nil {
		slog.Panicf("%s regist service err:%s", fun, err)
		return err
	}

	// 后门接口 ==================
	backdoor := &backDoorHttp{}
	err = backdoor.Init()
	if err != nil {
		slog.Panicf("%s init backdoor err:%s", fun, err)
		return err
	}

	binfos, err := m.loadDriver(sb, map[string]Processor{"_PROC_BACKDOOR": backdoor})
	if err == nil {
		err = sb.RegisterBackDoor(binfos)
		if err != nil {
			slog.Errorf("%s regist backdoor err:%s", fun, err)
		}

	} else {
		slog.Warnf("%s load backdoor driver err:%s", fun, err)
	}
	//==============================

	// sla metric埋点 ==================
	//init metric
	//user defualt metric opts
	metrics := rocserv.NewMetricsprocessor()
	if err != nil {
		slog.Warnf("init metrics fail:%v", err)
	}
	err = metrics.Init()
	if err != nil {
		slog.Warnf("%s init metrics err:%s", fun, err)
	}

	minfos, err := m.loadDriver(sb, map[string]Processor{"_PROC_METRICS": metrics})
	if err == nil {
		err = sb.RegisterMetrics(minfos)
		if err != nil {
			slog.Warnf("%s regist backdoor err:%s", fun, err)
		}

	} else {
		slog.Warnf("%s load metrics driver err:%s", fun, err)
	}
	//==============================

	// pause here
	var pause chan bool
	pause <- true

	return nil

}
func (m *Service) getMetricOps(sb *ServBaseV2) *rocserv.MetricsOpts {
	fun := "Service.getMetricOps -->"
	var metricConfig struct {
		metric *rocserv.MetricsOpts
	}
	err := sb.ServConfig(&metricConfig)
	if err != nil {
		slog.Panicf("%s serv config err:%s", fun, err)
		fmt.Sprintf("%s serv config err:%s", fun, err)
		return nil
	}
	return metricConfig.metric
}

func Serve(etcds []string, baseLoc string, initfn func(ServBase) error, procs map[string]Processor) error {
	return service.Serve(configEtcd{etcds, baseLoc}, initfn, procs)
}

func Init(etcds []string, baseLoc string, servLoc, servKey string, initfn func(ServBase) error, procs map[string]Processor) error {
	return service.Init(configEtcd{etcds, baseLoc}, servLoc, servKey, "", initfn, procs)
}

func Test(etcds []string, baseLoc string, initfn func(ServBase) error) error {
	return service.Init(configEtcd{etcds, baseLoc}, "test/test", "test", "console", initfn, nil)
}
<|endoftext|>


<commit_before>package task

import (
	"bytes"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/spf13/cobra"
	assert "github.com/stretchr/testify/require"
	"github.com/stretchr/testify/suite"
	tcclient "github.com/taskcluster/taskcluster-client-go"
)

const fakeTaskID = "ANnmjMocTymeTID0tlNJAw"
const fakeRunID = "0"

type FakeServerSuite struct {
	suite.Suite
	testServer *httptest.Server
}

func (suite *FakeServerSuite) SetupSuite() {
	// set up a fake server that knows how to answer the `task()` method
	handler := http.NewServeMux()
	handler.HandleFunc("/v1/task/"+fakeTaskID, taskHandler)
	handler.HandleFunc("/v1/task/"+fakeTaskID+"/status", manifestHandler)
	handler.HandleFunc("/v1/task/"+fakeTaskID+"/runs/"+fakeRunID+"/artifacts", artifactsHandler)
	handler.HandleFunc("/v1/task/"+fakeTaskID+"/cancel", cancelHandler)
	handler.HandleFunc("/v1/task/"+fakeTaskID+"/rerun", reRunHandler)
	handler.HandleFunc("/v1/task/"+fakeTaskID+"/runs/"+fakeRunID+"/claim", claimTaskHandler)
	handler.HandleFunc("/v1/task/"+fakeTaskID+"/runs/"+fakeRunID+"/completed", manifestHandler)

	suite.testServer = httptest.NewServer(handler)

	// set the base URL the subcommands use to point to the fake server
	queueBaseURL = suite.testServer.URL + "/v1"
}

func (suite *FakeServerSuite) TearDownSuite() {
	suite.testServer.Close()
	queueBaseURL = ""
}

func TestFakeServerSuite(t *testing.T) {
	suite.Run(t, new(FakeServerSuite))
}

// returns the test task on request
func taskHandler(w http.ResponseWriter, _ *http.Request) {
	metadata := `{"metadata": {"name": "my-test"}, "taskGroupId": "my-test"}`
	io.WriteString(w, metadata)
}

// returns the test status on request
func manifestHandler(w http.ResponseWriter, _ *http.Request) {
	status := `{
				  "status": {
				    "state": "completed",
				    "runs": [
				      {
				        "runId": 0,
				        "state": "completed",
				        "reasonCreated": "scheduled",
				        "reasonResolved": "completed"
				      }
				    ]
				  }
				}`
	io.WriteString(w, status)
}

func artifactsHandler(w http.ResponseWriter, _ *http.Request) {
	artifacts := `{
				  	"artifacts": [
				    {
				      	"storageType": "reference",
				      	"name": "fake_live.log",
				      	"expires": "2318-02-02T21:58:38.425Z",
				      	"contentType": "text/plain"
				    },
				    {
				      	"storageType": "s3",
				      	"name": "fake_live_backing.log",
				      	"expires": "2318-02-02T21:58:37.584Z",
				      	"contentType": "text/plain"
				    }
				  ]
				}`
	io.WriteString(w, artifacts)
}

func setUpCommand() (*bytes.Buffer, *cobra.Command) {
	buf := &bytes.Buffer{}
	cmd := &cobra.Command{}
	cmd.SetOutput(buf)

	return buf, cmd
}

func (suite *FakeServerSuite) TestNameCommand() {
	// set up to run a command and capture output
	buf, cmd := setUpCommand()

	// run the command
	args := []string{fakeTaskID}
	runName(&tcclient.Credentials{}, args, cmd.OutOrStdout(), cmd.Flags())

	suite.Equal("my-test\n", buf.String())
}

func (suite *FakeServerSuite) TestDefCommand() {
	// set up to run a command and capture output
	buf, cmd := setUpCommand()

	// run the command
	args := []string{fakeTaskID}
	runDef(&tcclient.Credentials{}, args, cmd.OutOrStdout(), cmd.Flags())

	var f interface{}
	json.Unmarshal(buf.Bytes(), &f)
	m := f.(map[string]interface{})
	m = m["metadata"].(map[string]interface{})
	suite.Equal("my-test", m["name"])
}

// Test the `task log` subcommand against a real task, since it does its own
// HTTP handling
func TestLogCommand(t *testing.T) {
	assert := assert.New(t)

	buf, cmd := setUpCommand()

	args := []string{"TtAsnXdCS1-tAQxvMO4rHQ"}
	runLog(&tcclient.Credentials{}, args, cmd.OutOrStdout(), cmd.Flags())

	// This is the output of a static log
	s := "[taskcluster 2017-03-03 21:18:34.946Z] Task ID: TtAsnXdCS1-tAQxvMO4rHQ\n" +
		"[taskcluster 2017-03-03 21:18:34.946Z] Worker ID: i-035dd1bd8da876f13\n" +
		"[taskcluster 2017-03-03 21:18:34.946Z] Worker Group: us-west-1b\n" +
		"[taskcluster 2017-03-03 21:18:34.946Z] Worker Node Type: m3.large\n" +
		"[taskcluster 2017-03-03 21:18:34.946Z] Worker Type: tutorial\n" +
		"[taskcluster 2017-03-03 21:18:34.946Z] Public IP: 54.153.13.193\n" +
		"\n" +
		"[taskcluster 2017-03-03 21:18:48.518Z] === Task Starting ===\n" +
		"hello World\n" +
		"[taskcluster 2017-03-03 21:18:48.945Z] === Task Finished ===\n" +
		"[taskcluster 2017-03-03 21:18:48.946Z] Successful task run with exit code: 0 completed in 14.001 seconds\n"

	assert.Equal(s, buf.String(), "Log's are not equal.")
}

func (suite *FakeServerSuite) TestArtifactsCommand() {
	// set up to run a command and capture output
	buf, cmd := setUpCommand()

	// run the command
	args := []string{fakeTaskID}

	runArtifacts(&tcclient.Credentials{}, args, cmd.OutOrStdout(), cmd.Flags())
	suite.Equal("fake_live.log\nfake_live_backing.log\n", buf.String())

}

func (suite *FakeServerSuite) TestGroupCommand() {
	// set up to run a command and capture output
	buf, cmd := setUpCommand()

	// run the command
	args := []string{fakeTaskID}
	runGroup(&tcclient.Credentials{}, args, cmd.OutOrStdout(), cmd.Flags())

	suite.Equal("my-test\n", buf.String())
}

func (suite *FakeServerSuite) TestStatusCommand() {
	// set up to run a command and capture output
	buf, cmd := setUpCommand()

	args := []string{fakeTaskID}

	// Test run flag
	cmd.Flags().IntP("run", "r", 0, "Specifies which run to consider.")
	runStatus(&tcclient.Credentials{}, args, cmd.OutOrStdout(), cmd.Flags())

	suite.Equal("completed 'completed'\n", buf.String())

	// Test all-runs flag
	buf.Reset()

	cmd.Flags().Set("run", "-1")
	cmd.Flags().BoolP("all-runs", "a", true, "Specifies which run to consider.")

	runStatus(&tcclient.Credentials{}, args, cmd.OutOrStdout(), cmd.Flags())

	suite.Equal("Run #0: completed 'completed'\n", buf.String())
}
<commit_msg>Fix TestLogCommand by using new taskId (previous task log had expired)<commit_after>package task

import (
	"bytes"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/spf13/cobra"
	assert "github.com/stretchr/testify/require"
	"github.com/stretchr/testify/suite"
	tcclient "github.com/taskcluster/taskcluster-client-go"
)

const fakeTaskID = "ANnmjMocTymeTID0tlNJAw"
const fakeRunID = "0"

type FakeServerSuite struct {
	suite.Suite
	testServer *httptest.Server
}

func (suite *FakeServerSuite) SetupSuite() {
	// set up a fake server that knows how to answer the `task()` method
	handler := http.NewServeMux()
	handler.HandleFunc("/v1/task/"+fakeTaskID, taskHandler)
	handler.HandleFunc("/v1/task/"+fakeTaskID+"/status", manifestHandler)
	handler.HandleFunc("/v1/task/"+fakeTaskID+"/runs/"+fakeRunID+"/artifacts", artifactsHandler)
	handler.HandleFunc("/v1/task/"+fakeTaskID+"/cancel", cancelHandler)
	handler.HandleFunc("/v1/task/"+fakeTaskID+"/rerun", reRunHandler)
	handler.HandleFunc("/v1/task/"+fakeTaskID+"/runs/"+fakeRunID+"/claim", claimTaskHandler)
	handler.HandleFunc("/v1/task/"+fakeTaskID+"/runs/"+fakeRunID+"/completed", manifestHandler)

	suite.testServer = httptest.NewServer(handler)

	// set the base URL the subcommands use to point to the fake server
	queueBaseURL = suite.testServer.URL + "/v1"
}

func (suite *FakeServerSuite) TearDownSuite() {
	suite.testServer.Close()
	queueBaseURL = ""
}

func TestFakeServerSuite(t *testing.T) {
	suite.Run(t, new(FakeServerSuite))
}

// returns the test task on request
func taskHandler(w http.ResponseWriter, _ *http.Request) {
	metadata := `{"metadata": {"name": "my-test"}, "taskGroupId": "my-test"}`
	io.WriteString(w, metadata)
}

// returns the test status on request
func manifestHandler(w http.ResponseWriter, _ *http.Request) {
	status := `{
				  "status": {
				    "state": "completed",
				    "runs": [
				      {
				        "runId": 0,
				        "state": "completed",
				        "reasonCreated": "scheduled",
				        "reasonResolved": "completed"
				      }
				    ]
				  }
				}`
	io.WriteString(w, status)
}

func artifactsHandler(w http.ResponseWriter, _ *http.Request) {
	artifacts := `{
				  	"artifacts": [
				    {
				      	"storageType": "reference",
				      	"name": "fake_live.log",
				      	"expires": "2318-02-02T21:58:38.425Z",
				      	"contentType": "text/plain"
				    },
				    {
				      	"storageType": "s3",
				      	"name": "fake_live_backing.log",
				      	"expires": "2318-02-02T21:58:37.584Z",
				      	"contentType": "text/plain"
				    }
				  ]
				}`
	io.WriteString(w, artifacts)
}

func setUpCommand() (*bytes.Buffer, *cobra.Command) {
	buf := &bytes.Buffer{}
	cmd := &cobra.Command{}
	cmd.SetOutput(buf)

	return buf, cmd
}

func (suite *FakeServerSuite) TestNameCommand() {
	// set up to run a command and capture output
	buf, cmd := setUpCommand()

	// run the command
	args := []string{fakeTaskID}
	runName(&tcclient.Credentials{}, args, cmd.OutOrStdout(), cmd.Flags())

	suite.Equal("my-test\n", buf.String())
}

func (suite *FakeServerSuite) TestDefCommand() {
	// set up to run a command and capture output
	buf, cmd := setUpCommand()

	// run the command
	args := []string{fakeTaskID}
	runDef(&tcclient.Credentials{}, args, cmd.OutOrStdout(), cmd.Flags())

	var f interface{}
	json.Unmarshal(buf.Bytes(), &f)
	m := f.(map[string]interface{})
	m = m["metadata"].(map[string]interface{})
	suite.Equal("my-test", m["name"])
}

// Test the `task log` subcommand against a real task, since it does its own
// HTTP handling
func TestLogCommand(t *testing.T) {
	assert := assert.New(t)

	buf, cmd := setUpCommand()

	args := []string{"YwxadSlPQm-d9jV8V0F9Ig"}
	runLog(&tcclient.Credentials{}, args, cmd.OutOrStdout(), cmd.Flags())

	// This is the output of a static log
	s := "[taskcluster 2017-03-03 21:18:34.946Z] Task ID: YwxadSlPQm-d9jV8V0F9Ig\n" +
		"[taskcluster 2017-03-03 21:18:34.946Z] Worker ID: i-035dd1bd8da876f13\n" +
		"[taskcluster 2017-03-03 21:18:34.946Z] Worker Group: us-west-1b\n" +
		"[taskcluster 2017-03-03 21:18:34.946Z] Worker Node Type: m3.large\n" +
		"[taskcluster 2017-03-03 21:18:34.946Z] Worker Type: tutorial\n" +
		"[taskcluster 2017-03-03 21:18:34.946Z] Public IP: 54.153.13.193\n" +
		"\n" +
		"[taskcluster 2017-03-03 21:18:48.518Z] === Task Starting ===\n" +
		"hello World\n" +
		"[taskcluster 2017-03-03 21:18:48.945Z] === Task Finished ===\n" +
		"[taskcluster 2017-03-03 21:18:48.946Z] Successful task run with exit code: 0 completed in 14.001 seconds\n"

	assert.Equal(s, buf.String(), "Log's are not equal.")
}

func (suite *FakeServerSuite) TestArtifactsCommand() {
	// set up to run a command and capture output
	buf, cmd := setUpCommand()

	// run the command
	args := []string{fakeTaskID}

	runArtifacts(&tcclient.Credentials{}, args, cmd.OutOrStdout(), cmd.Flags())
	suite.Equal("fake_live.log\nfake_live_backing.log\n", buf.String())

}

func (suite *FakeServerSuite) TestGroupCommand() {
	// set up to run a command and capture output
	buf, cmd := setUpCommand()

	// run the command
	args := []string{fakeTaskID}
	runGroup(&tcclient.Credentials{}, args, cmd.OutOrStdout(), cmd.Flags())

	suite.Equal("my-test\n", buf.String())
}

func (suite *FakeServerSuite) TestStatusCommand() {
	// set up to run a command and capture output
	buf, cmd := setUpCommand()

	args := []string{fakeTaskID}

	// Test run flag
	cmd.Flags().IntP("run", "r", 0, "Specifies which run to consider.")
	runStatus(&tcclient.Credentials{}, args, cmd.OutOrStdout(), cmd.Flags())

	suite.Equal("completed 'completed'\n", buf.String())

	// Test all-runs flag
	buf.Reset()

	cmd.Flags().Set("run", "-1")
	cmd.Flags().BoolP("all-runs", "a", true, "Specifies which run to consider.")

	runStatus(&tcclient.Credentials{}, args, cmd.OutOrStdout(), cmd.Flags())

	suite.Equal("Run #0: completed 'completed'\n", buf.String())
}
<|endoftext|>


<commit_before>//
package vagrant

import (
	"bufio"
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/nanobox-io/nanobox/config"
)

//
var err error

// Exists ensure vagrant is installed
func Exists() (exists bool) {
	var err error

	// check if vagrant is installed
	if _, err = exec.LookPath("vagrant"); err == nil {

		// initilize Vagrant incase it hasn't been; there is a chance that Vagrant has
		// never been used meaning there won't be a .vagrant.d folder, so we initialize
		// vagrant just to ensure it's ready to be used with nanobox by running any
		// vagrant command (in this case "vagrant -v").
		if b, err := exec.Command("vagrant", "-v").CombinedOutput(); err != nil {
			config.Fatal("[util/vagrant/vagrant] exec.Command() failed", string(b))
		}

		// read setup_version to determine if the version of vagrant is too old
		// (< 1.5.0) and needs to be migrated
		b, err := ioutil.ReadFile(filepath.Join(config.Home, ".vagrant.d", "setup_version"))
		if err != nil {
			config.Fatal("[util/vagrant/vagrant] ioutil.ReadFile() failed", err.Error())
		}

		// convert the []byte value from the file into a float 'version'
		version, err := strconv.ParseFloat(string(b), 64)
		if err != nil {
			config.Fatal("[util/vagrant/vagrant] strconv.ParseFloat() failed", err.Error())
		}

		// if the current version of vagrant is less than a 'working version' (1.5)
		// give instructions on how to update
		if version < 1.5 {
			fmt.Println(`
Nanobox has detected that you are using an old version of Vagrant (<1.5). Before
you can continue you'll need to run "vagrant update" and follow the instructions
to update Vagrant.
			`)

			// exit here to allow for upgrade
			os.Exit(0)
		}

		// if all checks pass
		exists = true
	}

	return
}

// run runs a vagrant command
func run(cmd *exec.Cmd) error {

	//
	handleCMDout(cmd)

	// start the command; we need this to 'fire and forget' so that we can manually
	// capture and modify the commands output above
	if err := cmd.Start(); err != nil {
		return err
	}

	// wait for the command to complete/fail and exit, giving us a chance to scan
	// the output
	if err := cmd.Wait(); err != nil {
		return err
	}

	return nil
}

// runInContext runs a command in the context of a Vagrantfile (from the same dir)
func runInContext(cmd *exec.Cmd) error {

	// run the command from ~/.nanobox/apps/<config.App>. Running the command from
	// the directory that contains the Vagratfile ensure that the command can
	// atleast run (especially in cases like 'create' where a VM hadn't been created
	// yet, and a UUID isn't available)
	setContext(config.AppDir)

	//
	handleCMDout(cmd)

	// start the command; we need this to 'fire and forget' so that we can manually
	// capture and modify the commands output above
	if err := cmd.Start(); err != nil {
		return err
	}

	// wait for the command to complete/fail and exit, giving us a chance to scan
	// the output
	if err := cmd.Wait(); err != nil {
		return err
	}

	// switch back to project dir
	setContext(config.CWDir)

	return nil
}

// setContext changes the working directory to the designated context
func setContext(context string) {
	if err := os.Chdir(context); err != nil {
		fmt.Printf("No app found at %s. Exiting...\n", config.AppDir)
		os.Exit(1)
	}
}

func customScanner(data []byte, atEOF bool) (advance int, token []byte, err error) {

	//
	if atEOF && len(data) == 0 {
		return 0, nil, nil
	}

	if i := bytes.IndexByte(data, '\n'); i >= 0 {
		// We have a full newline-terminated line.
		return i + 1, dropCR(data[0:i]), nil
	}

	if i := bytes.IndexByte(data, '\r'); i >= 0 {
		return i + 1, dropCR(data[0:i]), nil
	}

	// If we're at EOF, we have a final, non-terminated line. Return it.
	if atEOF {
		return len(data), dropCR(data), nil
	}

	// Request more data.
	return 0, nil, nil
}

// dropCR drops a terminal \r from the data.
func dropCR(data []byte) []byte {
	if len(data) > 0 && data[len(data)-1] == '\r' {
		return data[0 : len(data)-1]
	}
	return data
}

// handleCMDout
func handleCMDout(cmd *exec.Cmd) {

	// create a stderr pipe that will write any error messages to the log
	stderr, err := cmd.StderrPipe()
	if err != nil {
		Fatal("[util/vagrant/vagrant] cmd.StderrPipe() failed", err.Error())
	}

	// log any command errors to the log
	stderrScanner := bufio.NewScanner(stderr)
	go func() {

		//
		var once sync.Once
		for stderrScanner.Scan() {

			// only display the error message once, but log every error
			once.Do(func() { Error("A vagrant error occured", "") })
			Log.Error(stderrScanner.Text())
		}
	}()

	// create a stdout pipe that will allow for scanning the output line-by-line;
	// if needed a stderr pipe could also be created at some point
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		Fatal("[util/vagrant/vagrant] cmd.StdoutPipe() failed", err.Error())
	}

	//
	output := make(chan string)

	// start a goroutine that will act as an 'outputer' allowing us to add 'dots'
	// to the end of each line (as these lines are a reduced version of the actual
	// output there will be some delay between output)
	go func() {

		tick := time.Second

		// by default, don't print dots until we've received at least one message
		messaged := false

		// begin a loop to read off the channel until it's closed
		for {
			select {

			// print any messages and reset ticker
			case msg, ok := <-output:

				// once an one message is received, indicate that dots can now be printed
				messaged = true

				// once the channel closes print the final newline and close the goroutine
				if !ok {
					fmt.Printf("\n")
					return
				}

				fmt.Printf("\n   - %s", msg)

				tick = time.Second

			// after every tick print a '.' until we get another message one the channel
			// (at which point ticker is reset and it starts all over again)
			case <-time.After(tick):
				if messaged {
					fmt.Print(".")

					// increase the wait time by 1/4 of the total previous time; this should
					// provide a good 'loading' effect w/o printing too many dots
					tick += tick / 4
				}
			}
		}
	}()

	// scan the command output intercepting only 'important' lines of vagrant output'
	// and tailoring their message so as to not flood the output.
	// styled according to: http://nanodocs.gopagoda.io/engines/style-guide
	stdoutScanner := bufio.NewScanner(stdout)
	stdoutScanner.Split(customScanner)
	go func() {
		for stdoutScanner.Scan() {

			txt := strings.TrimSpace(stdoutScanner.Text())
			app := config.Nanofile.Name

			// log all vagrant output (might as well)
			Log.Info(txt)

			// handle generic cases
			switch {

			// show the progress bar when trying to download nanobox/boot2docker
			case strings.Contains(txt, "box: Progress:"):
				subMatch := regexp.MustCompile(`box: Progress: (\d{1,3})% \(Rate: (.*), Estimated time remaining: (\d*:\d*:\d*)`).FindStringSubmatch(txt)

				// ensure we have all the submatches needed before trying to use them
				if len(subMatch) >= 4 {
					i, err := strconv.Atoi(subMatch[1])
					if err != nil {

					}

					// show download progress: [*** progress *** 0.0%] 00:00:00 remaining
					fmt.Printf("\r   [%-41s %s%%] %s (%s remaining)", strings.Repeat("*", int(float64(i)/2.5)), subMatch[1], subMatch[2], subMatch[3])
				}
			}

			// handle specific cases
			switch txt {

			// nanobox vm has not yet been created
			case fmt.Sprintf("==> %v: VM not created. Moving on...", app):
				output <- "Nanobox not yet created, use 'nanobox dev' or 'nanobox run' to create it."

			// nanobox is already running
			case fmt.Sprintf("==> %v: VirtualBox VM is already running.", app):
				continue

			case fmt.Sprintf("==> %v: Importing base box 'nanobox/boot2docker'...", app):
				output <- "Importing nanobox base image"
			case fmt.Sprintf("==> %v: Booting VM...", app):
				output <- "Booting virtual machine"
			case fmt.Sprintf("==> %v: Configuring and enabling network interfaces...", app):
				output <- "Configuring virtual network"
			case fmt.Sprintf("==> %v: Mounting shared folders...", app):
				output <- fmt.Sprintf("Mounting source code (%s)", config.CWDir)
			case fmt.Sprintf("==> %v: Waiting for nanobox server...", app):
				output <- "Starting nanobox server"
			case fmt.Sprintf("==> %v: Attempting graceful shutdown of VM...", app):
				output <- "Shutting down virtual machine"
			case fmt.Sprintf("==> %v: Destroying VM and associated drives...", app):
				// output <- "Destroying virtual machine"
			case fmt.Sprintf("==> %v: Forcing shutdown of VM...", app):
				output <- "Shutting down virtual machine"
			case fmt.Sprintf("==> %v: Saving VM state and suspending execution...", app):
				output <- "Saving virtual machine"
			case fmt.Sprintf("==> %v: Resuming suspended VM...", app):
				// output <- "Resuming virtual machine"
			}
		}

		// close the output channel once all lines of command output have been read
		close(output)
	}()
}
<commit_msg>fix the blank new line from vagrant output<commit_after>//
package vagrant

import (
	"bufio"
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/nanobox-io/nanobox/config"
)

//
var err error

// Exists ensure vagrant is installed
func Exists() (exists bool) {
	var err error

	// check if vagrant is installed
	if _, err = exec.LookPath("vagrant"); err == nil {

		// initilize Vagrant incase it hasn't been; there is a chance that Vagrant has
		// never been used meaning there won't be a .vagrant.d folder, so we initialize
		// vagrant just to ensure it's ready to be used with nanobox by running any
		// vagrant command (in this case "vagrant -v").
		if b, err := exec.Command("vagrant", "-v").CombinedOutput(); err != nil {
			config.Fatal("[util/vagrant/vagrant] exec.Command() failed", string(b))
		}

		// read setup_version to determine if the version of vagrant is too old
		// (< 1.5.0) and needs to be migrated
		b, err := ioutil.ReadFile(filepath.Join(config.Home, ".vagrant.d", "setup_version"))
		if err != nil {
			config.Fatal("[util/vagrant/vagrant] ioutil.ReadFile() failed", err.Error())
		}

		// convert the []byte value from the file into a float 'version'
		version, err := strconv.ParseFloat(string(b), 64)
		if err != nil {
			config.Fatal("[util/vagrant/vagrant] strconv.ParseFloat() failed", err.Error())
		}

		// if the current version of vagrant is less than a 'working version' (1.5)
		// give instructions on how to update
		if version < 1.5 {
			fmt.Println(`
Nanobox has detected that you are using an old version of Vagrant (<1.5). Before
you can continue you'll need to run "vagrant update" and follow the instructions
to update Vagrant.
			`)

			// exit here to allow for upgrade
			os.Exit(0)
		}

		// if all checks pass
		exists = true
	}

	return
}

// run runs a vagrant command
func run(cmd *exec.Cmd) error {

	//
	handleCMDout(cmd)

	// start the command; we need this to 'fire and forget' so that we can manually
	// capture and modify the commands output above
	if err := cmd.Start(); err != nil {
		return err
	}

	// wait for the command to complete/fail and exit, giving us a chance to scan
	// the output
	if err := cmd.Wait(); err != nil {
		return err
	}

	return nil
}

// runInContext runs a command in the context of a Vagrantfile (from the same dir)
func runInContext(cmd *exec.Cmd) error {

	// run the command from ~/.nanobox/apps/<config.App>. Running the command from
	// the directory that contains the Vagratfile ensure that the command can
	// atleast run (especially in cases like 'create' where a VM hadn't been created
	// yet, and a UUID isn't available)
	setContext(config.AppDir)

	//
	handleCMDout(cmd)

	// start the command; we need this to 'fire and forget' so that we can manually
	// capture and modify the commands output above
	if err := cmd.Start(); err != nil {
		return err
	}

	// wait for the command to complete/fail and exit, giving us a chance to scan
	// the output
	if err := cmd.Wait(); err != nil {
		return err
	}

	// switch back to project dir
	setContext(config.CWDir)

	return nil
}

// setContext changes the working directory to the designated context
func setContext(context string) {
	if err := os.Chdir(context); err != nil {
		fmt.Printf("No app found at %s. Exiting...\n", config.AppDir)
		os.Exit(1)
	}
}

func customScanner(data []byte, atEOF bool) (advance int, token []byte, err error) {

	//
	if atEOF && len(data) == 0 {
		return 0, nil, nil
	}

	if i := bytes.IndexByte(data, '\n'); i >= 0 {
		// We have a full newline-terminated line.
		return i + 1, dropCR(data[0:i]), nil
	}

	if i := bytes.IndexByte(data, '\r'); i >= 0 {
		return i + 1, dropCR(data[0:i]), nil
	}

	// If we're at EOF, we have a final, non-terminated line. Return it.
	if atEOF {
		return len(data), dropCR(data), nil
	}

	// Request more data.
	return 0, nil, nil
}

// dropCR drops a terminal \r from the data.
func dropCR(data []byte) []byte {
	if len(data) > 0 && data[len(data)-1] == '\r' {
		return data[0 : len(data)-1]
	}
	return data
}

// handleCMDout
func handleCMDout(cmd *exec.Cmd) {

	// create a stderr pipe that will write any error messages to the log
	stderr, err := cmd.StderrPipe()
	if err != nil {
		Fatal("[util/vagrant/vagrant] cmd.StderrPipe() failed", err.Error())
	}

	// log any command errors to the log
	stderrScanner := bufio.NewScanner(stderr)
	go func() {

		//
		var once sync.Once
		for stderrScanner.Scan() {

			// only display the error message once, but log every error
			once.Do(func() { Error("A vagrant error occured", "") })
			Log.Error(stderrScanner.Text())
		}
	}()

	// create a stdout pipe that will allow for scanning the output line-by-line;
	// if needed a stderr pipe could also be created at some point
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		Fatal("[util/vagrant/vagrant] cmd.StdoutPipe() failed", err.Error())
	}

	//
	output := make(chan string)

	// start a goroutine that will act as an 'outputer' allowing us to add 'dots'
	// to the end of each line (as these lines are a reduced version of the actual
	// output there will be some delay between output)
	go func() {

		tick := time.Second

		// by default, don't print dots until we've received at least one message
		messaged := false

		// empty previous message by default
		prevMsg := ""

		// begin a loop to read off the channel until it's closed
		for {
			select {

			// print any messages and reset ticker
			case msg, ok := <-output:

				// if the previous message isn't empty and doesn't equal the "new" message
				// print a new line; I dont really like this, but it works for now.
				if prevMsg != "" && prevMsg != msg {
					fmt.Printf("\n")
				}

				// set previous message = to current message
				prevMsg = msg

				// once any, one, message is received indicate that dots should now be
				// printed
				messaged = true

				// once the channel closes print the final newline and close the goroutine
				if !ok {
					fmt.Printf("\n")
					return
				}

				fmt.Printf("   - %s", msg)

				tick = time.Second

			// after every tick print a '.' until we get another message one the channel
			// (at which point ticker is reset and it starts all over again)
			case <-time.After(tick):
				if messaged {
					fmt.Print(".")

					// increase the wait time by 1/4 of the total previous time; this should
					// provide a good 'loading' effect w/o printing too many dots
					tick += tick / 4
				}
			}
		}
	}()

	// scan the command output intercepting only 'important' lines of vagrant output'
	// and tailoring their message so as to not flood the output.
	// styled according to: http://nanodocs.gopagoda.io/engines/style-guide
	stdoutScanner := bufio.NewScanner(stdout)
	stdoutScanner.Split(customScanner)
	go func() {
		for stdoutScanner.Scan() {

			txt := strings.TrimSpace(stdoutScanner.Text())
			app := config.Nanofile.Name

			// log all vagrant output (might as well)
			Log.Info(txt)

			// handle generic cases
			switch {

			// show the progress bar when trying to download nanobox/boot2docker
			case strings.Contains(txt, "box: Progress:"):
				subMatch := regexp.MustCompile(`box: Progress: (\d{1,3})% \(Rate: (.*), Estimated time remaining: (\d*:\d*:\d*)`).FindStringSubmatch(txt)

				// ensure we have all the submatches needed before trying to use them
				if len(subMatch) >= 4 {
					i, err := strconv.Atoi(subMatch[1])
					if err != nil {

					}

					// show download progress: [*** progress *** 0.0%] 00:00:00 remaining
					fmt.Printf("\r   [%-41s %s%%] %s (%s remaining)", strings.Repeat("*", int(float64(i)/2.5)), subMatch[1], subMatch[2], subMatch[3])
				}
			}

			// handle specific cases
			switch txt {

			// nanobox vm has not yet been created
			case fmt.Sprintf("==> %v: VM not created. Moving on...", app):
				output <- "Nanobox not yet created, use 'nanobox dev' or 'nanobox run' to create it."

			// nanobox is already running
			case fmt.Sprintf("==> %v: VirtualBox VM is already running.", app):
				continue

			case fmt.Sprintf("==> %v: Importing base box 'nanobox/boot2docker'...", app):
				output <- "Importing nanobox base image"
			case fmt.Sprintf("==> %v: Booting VM...", app):
				output <- "Booting virtual machine"
			case fmt.Sprintf("==> %v: Configuring and enabling network interfaces...", app):
				output <- "Configuring virtual network"
			case fmt.Sprintf("==> %v: Mounting shared folders...", app):
				output <- fmt.Sprintf("Mounting source code (%s)", config.CWDir)
			case fmt.Sprintf("==> %v: Waiting for nanobox server...", app):
				output <- "Starting nanobox server"
			case fmt.Sprintf("==> %v: Attempting graceful shutdown of VM...", app):
				output <- "Shutting down virtual machine"
			case fmt.Sprintf("==> %v: Destroying VM and associated drives...", app):
				// output <- "Destroying virtual machine"
			case fmt.Sprintf("==> %v: Forcing shutdown of VM...", app):
				output <- "Shutting down virtual machine"
			case fmt.Sprintf("==> %v: Saving VM state and suspending execution...", app):
				output <- "Saving virtual machine"
			case fmt.Sprintf("==> %v: Resuming suspended VM...", app):
				// output <- "Resuming virtual machine"
			}
		}

		// close the output channel once all lines of command output have been read
		close(output)
	}()
}
<|endoftext|>


<commit_before>package cc

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

	"github.com/omakoto/zenlog-go/zenlog/builtins/history"
	"github.com/omakoto/zenlog-go/zenlog/config"
	"github.com/omakoto/zenlog-go/zenlog/shell"
	"github.com/omakoto/zenlog-go/zenlog/util"
)

// CC represents "command line context" to detect consecutive calls of the same command.
type CC struct {
	FirstCommandLine string
	FirstCursorPos   int

	BeforeCommandLine string
	BeforeCursorPos   int

	AfterCommandLine string
	AfterCursorPos   int

	LastLog string

	Operation string
	NumRepeat int

	config *config.Config
}

func filenameForConfig(c *config.Config) string {
	return fmt.Sprintf("%s/zenlog_lastcommand_%d.json", c.TempDir, c.ZenlogPid)
}

func FromEnvironment(operation string, proxy shell.Proxy) *CC {
	config := config.InitConfigForCommands()

	ret := CC{}
	ret.config = config
	ret.Operation = operation
	ret.BeforeCommandLine, ret.BeforeCursorPos = proxy.GetCommandLine()
	ret.LastLog = history.NthLastLog(config, 0, 1, history.LogTypeSan)

	// Load the last cc.
	last := FromLastFile()

	if last.Operation == ret.Operation &&
		last.AfterCommandLine == ret.BeforeCommandLine &&
		last.AfterCursorPos == ret.BeforeCursorPos &&
		last.LastLog == ret.LastLog {

		ret.FirstCommandLine = last.FirstCommandLine
		ret.FirstCursorPos = last.FirstCursorPos
		ret.NumRepeat = last.NumRepeat + 1

	} else {
		ret.FirstCommandLine = ret.BeforeCommandLine
		ret.FirstCursorPos = ret.BeforeCursorPos
	}

	return &ret
}

func FromLastFile() *CC {
	config := config.InitConfigForCommands()
	file := filenameForConfig(config)

	ret := CC{config: config}
	if !util.FileExists(file) {
		return &ret
	}
	dat, err := ioutil.ReadFile(file)
	util.Check(err, "ReadFile failed")

	util.MustUnmarshal(string(dat), &ret)
	return &ret
}

func (cc *CC) Save() {
	file := filenameForConfig(cc.config)

	dat := []byte(util.MustMarshal(&cc))

	err := ioutil.WriteFile(file, dat, 0600)
	util.Warn(err, "WriteFile failed")
}

func (cc *CC) ClearSaved() {
	os.Remove(filenameForConfig(cc.config))
}

func (cc *CC) Config() *config.Config {
	return cc.config
}
<commit_msg>Rename CC to make it understandable<commit_after>package cc

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

	"github.com/omakoto/zenlog-go/zenlog/builtins/history"
	"github.com/omakoto/zenlog-go/zenlog/config"
	"github.com/omakoto/zenlog-go/zenlog/shell"
	"github.com/omakoto/zenlog-go/zenlog/util"
)

// CommandLineContext represents "command line context" to detect consecutive calls of the same command.
type CommandLineContext struct {
	FirstCommandLine string
	FirstCursorPos   int

	BeforeCommandLine string
	BeforeCursorPos   int

	AfterCommandLine string
	AfterCursorPos   int

	LastLog string

	Operation string
	NumRepeat int

	config *config.Config
}

func filenameForConfig(c *config.Config) string {
	return fmt.Sprintf("%s/zenlog_lastcommand_%d.json", c.TempDir, c.ZenlogPid)
}

func FromEnvironment(operation string, proxy shell.Proxy) *CommandLineContext {
	config := config.InitConfigForCommands()

	ret := CommandLineContext{}
	ret.config = config
	ret.Operation = operation
	ret.BeforeCommandLine, ret.BeforeCursorPos = proxy.GetCommandLine()
	ret.LastLog = history.NthLastLog(config, 0, 1, history.LogTypeSan)

	// Load the last cc.
	last := FromLastFile()

	if last.Operation == ret.Operation &&
		last.AfterCommandLine == ret.BeforeCommandLine &&
		last.AfterCursorPos == ret.BeforeCursorPos &&
		last.LastLog == ret.LastLog {

		ret.FirstCommandLine = last.FirstCommandLine
		ret.FirstCursorPos = last.FirstCursorPos
		ret.NumRepeat = last.NumRepeat + 1

	} else {
		ret.FirstCommandLine = ret.BeforeCommandLine
		ret.FirstCursorPos = ret.BeforeCursorPos
	}

	return &ret
}

func FromLastFile() *CommandLineContext {
	config := config.InitConfigForCommands()
	file := filenameForConfig(config)

	ret := CommandLineContext{config: config}
	if !util.FileExists(file) {
		return &ret
	}
	dat, err := ioutil.ReadFile(file)
	util.Check(err, "ReadFile failed")

	util.MustUnmarshal(string(dat), &ret)
	return &ret
}

func (cc *CommandLineContext) Save() {
	file := filenameForConfig(cc.config)

	dat := []byte(util.MustMarshal(&cc))

	err := ioutil.WriteFile(file, dat, 0600)
	util.Warn(err, "WriteFile failed")
}

func (cc *CommandLineContext) ClearSaved() {
	os.Remove(filenameForConfig(cc.config))
}

func (cc *CommandLineContext) Config() *config.Config {
	return cc.config
}
<|endoftext|>