<commit_before>package link

import (
	"errors"
	"os"
	"testing"

	qt "github.com/frankban/quicktest"

	"github.com/cilium/ebpf"
	"github.com/cilium/ebpf/asm"
	"github.com/cilium/ebpf/internal/testutils"
)

var (
	kprobeSpec = ebpf.ProgramSpec{
		Type:    ebpf.Kprobe,
		License: "MIT",
		Instructions: asm.Instructions{
			// set exit code to 0
			asm.Mov.Imm(asm.R0, 0),
			asm.Return(),
		},
	}
)

func TestKprobe(t *testing.T) {
	c := qt.New(t)

	prog, err := ebpf.NewProgram(&kprobeSpec)
	if err != nil {
		t.Fatal(err)
	}
	defer prog.Close()

	k, err := Kprobe("printk", prog)
	c.Assert(err, qt.IsNil)
	defer k.Close()

	testLink(t, k, testLinkOptions{
		prog: prog,
	})

	k, err = Kprobe("bogus", prog)
	c.Assert(errors.Is(err, os.ErrNotExist), qt.IsTrue, qt.Commentf("got error: %s", err))
	if k != nil {
		k.Close()
	}
}

func TestKretprobe(t *testing.T) {
	c := qt.New(t)

	prog, err := ebpf.NewProgram(&kprobeSpec)
	if err != nil {
		t.Fatal(err)
	}
	defer prog.Close()

	k, err := Kretprobe("printk", prog)
	c.Assert(err, qt.IsNil)
	defer k.Close()

	testLink(t, k, testLinkOptions{
		prog: prog,
	})

	k, err = Kretprobe("bogus", prog)
	c.Assert(errors.Is(err, os.ErrNotExist), qt.IsTrue, qt.Commentf("got error: %s", err))
	if k != nil {
		k.Close()
	}
}

func TestKprobeErrors(t *testing.T) {
	c := qt.New(t)

	// Invalid Kprobe incantations. Kretprobe uses the same code paths
	// with a different ret flag.
	_, err := Kprobe("", nil) // empty symbol
	c.Assert(errors.Is(err, errInvalidInput), qt.IsTrue)

	_, err = Kprobe("_", nil) // empty prog
	c.Assert(errors.Is(err, errInvalidInput), qt.IsTrue)

	_, err = Kprobe(".", &ebpf.Program{}) // illegal chars in symbol
	c.Assert(errors.Is(err, errInvalidInput), qt.IsTrue)

	_, err = Kprobe("foo", &ebpf.Program{}) // wrong prog type
	c.Assert(errors.Is(err, errInvalidInput), qt.IsTrue)
}

// Test k(ret)probe creation using perf_kprobe PMU.
func TestKprobeCreatePMU(t *testing.T) {

	// Requires at least 4.17 (e12f03d7031a "perf/core: Implement the 'perf_kprobe' PMU")
	testutils.SkipOnOldKernel(t, "4.17", "perf_kprobe PMU")

	c := qt.New(t)

	// kprobe happy path. printk is always present.
	pk, err := pmuKprobe("printk", false)
	c.Assert(err, qt.IsNil)
	defer pk.Close()

	c.Assert(pk.typ, qt.Equals, kprobeEvent)

	// kretprobe happy path.
	pr, err := pmuKprobe("printk", true)
	c.Assert(err, qt.IsNil)
	defer pr.Close()

	c.Assert(pr.typ, qt.Equals, kretprobeEvent)

	// Expect os.ErrNotExist when specifying a non-existent kernel symbol
	// on kernels 4.17 and up.
	_, err = pmuKprobe("bogus", false)
	c.Assert(errors.Is(err, os.ErrNotExist), qt.IsTrue, qt.Commentf("got error: %s", err))

	// A kernel bug was fixed in 97c753e62e6c where EINVAL was returned instead
	// of ENOENT, but only for kretprobes.
	_, err = pmuKprobe("bogus", true)
	c.Assert(errors.Is(err, os.ErrNotExist), qt.IsTrue, qt.Commentf("got error: %s", err))
}

// Test fallback behaviour on kernels without perf_kprobe PMU available.
func TestKprobePMUUnavailable(t *testing.T) {
	c := qt.New(t)

	pk, err := pmuKprobe("printk", false)
	if err == nil {
		pk.Close()
		t.Skipf("Kernel supports perf_kprobe PMU, not asserting error.")
	}

	// Only allow a PMU creation with a valid kernel symbol to fail with ErrNotSupported.
	c.Assert(errors.Is(err, ErrNotSupported), qt.IsTrue, qt.Commentf("got error: %s", err))
}

func BenchmarkKprobeCreatePMU(b *testing.B) {
	for n := 0; n < b.N; n++ {
		pr, err := pmuKprobe("printk", false)
		if err != nil {
			b.Error("error creating perf_kprobe PMU:", err)
		}

		if err := pr.Close(); err != nil {
			b.Error("error closing perf_kprobe PMU:", err)
		}
	}
}

// Test tracefs k(ret)probe creation on all kernel versions.
func TestKprobeTraceFS(t *testing.T) {
	c := qt.New(t)

	symbol := "printk"

	// Open and close tracefs k(ret)probes, checking all errors.
	kp, err := tracefsKprobe(symbol, false)
	c.Assert(err, qt.IsNil)
	c.Assert(kp.Close(), qt.IsNil)
	c.Assert(kp.typ, qt.Equals, kprobeEvent)

	kp, err = tracefsKprobe(symbol, true)
	c.Assert(err, qt.IsNil)
	c.Assert(kp.Close(), qt.IsNil)
	c.Assert(kp.typ, qt.Equals, kretprobeEvent)

	// Create two identical trace events, ensure their IDs differ.
	k1, err := tracefsKprobe(symbol, false)
	c.Assert(err, qt.IsNil)
	defer k1.Close()
	c.Assert(k1.tracefsID, qt.Not(qt.Equals), 0)

	k2, err := tracefsKprobe(symbol, false)
	c.Assert(err, qt.IsNil)
	defer k2.Close()
	c.Assert(k2.tracefsID, qt.Not(qt.Equals), 0)

	// Compare the kprobes' tracefs IDs.
	c.Assert(k1.tracefsID, qt.Not(qt.CmpEquals()), k2.tracefsID)

	// Write a k(ret)probe event for a non-existing symbol.
	err = createTraceFSProbeEvent(kprobeType, "testgroup", "bogus", "", 0, false)
	c.Assert(errors.Is(err, os.ErrNotExist), qt.IsTrue, qt.Commentf("got error: %s", err))

	// A kernel bug was fixed in 97c753e62e6c where EINVAL was returned instead
	// of ENOENT, but only for kretprobes.
	err = createTraceFSProbeEvent(kprobeType, "testgroup", "bogus", "", 0, true)
	c.Assert(errors.Is(err, os.ErrNotExist), qt.IsTrue, qt.Commentf("got error: %s", err))
}

func BenchmarkKprobeCreateTraceFS(b *testing.B) {
	for n := 0; n < b.N; n++ {
		// Include <tracefs>/kprobe_events operations in the benchmark loop
		// because we create one per perf event.
		pr, err := tracefsKprobe("printk", false)
		if err != nil {
			b.Error("error creating tracefs perf event:", err)
		}

		if err := pr.Close(); err != nil {
			b.Error("error closing tracefs perf event:", err)
		}
	}
}

// Test k(ret)probe creation writing directly to <tracefs>/kprobe_events.
// Only runs on 5.0 and over. Earlier versions ignored writes of duplicate
// events, while 5.0 started returning -EEXIST when a kprobe event already
// exists.
func TestKprobeCreateTraceFS(t *testing.T) {
	testutils.SkipOnOldKernel(t, "5.0", "<tracefs>/kprobe_events doesn't reject duplicate events")

	c := qt.New(t)

	symbol := "printk"
	pg, _ := randomGroup("ebpftest")
	rg, _ := randomGroup("ebpftest")

	// Tee up cleanups in case any of the Asserts abort the function.
	defer func() {
		_ = closeTraceFSProbeEvent(kprobeType, pg, symbol)
		_ = closeTraceFSProbeEvent(kprobeType, rg, symbol)
	}()

	// Create a kprobe.
	err := createTraceFSProbeEvent(kprobeType, pg, symbol, "", 0, false)
	c.Assert(err, qt.IsNil)

	// Attempt to create an identical kprobe using tracefs,
	// expect it to fail with os.ErrExist.
	err = createTraceFSProbeEvent(kprobeType, pg, symbol, "", 0, false)
	c.Assert(errors.Is(err, os.ErrExist), qt.IsTrue,
		qt.Commentf("expected consecutive kprobe creation to contain os.ErrExist, got: %v", err))

	// Expect a successful close of the kprobe.
	c.Assert(closeTraceFSProbeEvent(kprobeType, pg, symbol), qt.IsNil)

	// Same test for a kretprobe.
	err = createTraceFSProbeEvent(kprobeType, rg, symbol, "", 0, true)
	c.Assert(err, qt.IsNil)

	err = createTraceFSProbeEvent(kprobeType, rg, symbol, "", 0, true)
	c.Assert(os.IsExist(err), qt.IsFalse,
		qt.Commentf("expected consecutive kretprobe creation to contain os.ErrExist, got: %v", err))

	// Expect a successful close of the kretprobe.
	c.Assert(closeTraceFSProbeEvent(kprobeType, rg, symbol), qt.IsNil)
}

func TestKprobeTraceFSGroup(t *testing.T) {
	c := qt.New(t)

	// Expect <prefix>_<16 random hex chars>.
	g, err := randomGroup("ebpftest")
	c.Assert(err, qt.IsNil)
	c.Assert(g, qt.Matches, `ebpftest_[a-f0-9]{16}`)

	// Expect error when the generator's output exceeds 63 characters.
	p := make([]byte, 47) // 63 - 17 (length of the random suffix and underscore) + 1
	for i := range p {
		p[i] = byte('a')
	}
	_, err = randomGroup(string(p))
	c.Assert(err, qt.Not(qt.IsNil))

	// Reject non-alphanumeric characters.
	_, err = randomGroup("/")
	c.Assert(err, qt.Not(qt.IsNil))
}

func TestDetermineRetprobeBit(t *testing.T) {
	testutils.SkipOnOldKernel(t, "4.17", "perf_kprobe PMU")
	c := qt.New(t)

	rpk, err := kretprobeBit()
	c.Assert(err, qt.IsNil)
	c.Assert(rpk, qt.Equals, uint64(0))

	rpu, err := uretprobeBit()
	c.Assert(err, qt.IsNil)
	c.Assert(rpu, qt.Equals, uint64(0))
}
<commit_msg>link: test kprobe program call<commit_after>package link

import (
	"errors"
	"os"
	"testing"

	qt "github.com/frankban/quicktest"

	"github.com/cilium/ebpf"
	"github.com/cilium/ebpf/asm"
	"github.com/cilium/ebpf/internal/testutils"
	"github.com/cilium/ebpf/internal/unix"
)

var (
	kprobeSpec = ebpf.ProgramSpec{
		Type:    ebpf.Kprobe,
		License: "MIT",
		Instructions: asm.Instructions{
			// set exit code to 0
			asm.Mov.Imm(asm.R0, 0),
			asm.Return(),
		},
	}
)

func TestKprobe(t *testing.T) {
	c := qt.New(t)

	prog, err := ebpf.NewProgram(&kprobeSpec)
	if err != nil {
		t.Fatal(err)
	}
	defer prog.Close()

	k, err := Kprobe("printk", prog)
	c.Assert(err, qt.IsNil)
	defer k.Close()

	testLink(t, k, testLinkOptions{
		prog: prog,
	})

	k, err = Kprobe("bogus", prog)
	c.Assert(errors.Is(err, os.ErrNotExist), qt.IsTrue, qt.Commentf("got error: %s", err))
	if k != nil {
		k.Close()
	}
}

func TestKretprobe(t *testing.T) {
	c := qt.New(t)

	prog, err := ebpf.NewProgram(&kprobeSpec)
	if err != nil {
		t.Fatal(err)
	}
	defer prog.Close()

	k, err := Kretprobe("printk", prog)
	c.Assert(err, qt.IsNil)
	defer k.Close()

	testLink(t, k, testLinkOptions{
		prog: prog,
	})

	k, err = Kretprobe("bogus", prog)
	c.Assert(errors.Is(err, os.ErrNotExist), qt.IsTrue, qt.Commentf("got error: %s", err))
	if k != nil {
		k.Close()
	}
}

func TestKprobeErrors(t *testing.T) {
	c := qt.New(t)

	// Invalid Kprobe incantations. Kretprobe uses the same code paths
	// with a different ret flag.
	_, err := Kprobe("", nil) // empty symbol
	c.Assert(errors.Is(err, errInvalidInput), qt.IsTrue)

	_, err = Kprobe("_", nil) // empty prog
	c.Assert(errors.Is(err, errInvalidInput), qt.IsTrue)

	_, err = Kprobe(".", &ebpf.Program{}) // illegal chars in symbol
	c.Assert(errors.Is(err, errInvalidInput), qt.IsTrue)

	_, err = Kprobe("foo", &ebpf.Program{}) // wrong prog type
	c.Assert(errors.Is(err, errInvalidInput), qt.IsTrue)
}

// Test k(ret)probe creation using perf_kprobe PMU.
func TestKprobeCreatePMU(t *testing.T) {

	// Requires at least 4.17 (e12f03d7031a "perf/core: Implement the 'perf_kprobe' PMU")
	testutils.SkipOnOldKernel(t, "4.17", "perf_kprobe PMU")

	c := qt.New(t)

	// kprobe happy path. printk is always present.
	pk, err := pmuKprobe("printk", false)
	c.Assert(err, qt.IsNil)
	defer pk.Close()

	c.Assert(pk.typ, qt.Equals, kprobeEvent)

	// kretprobe happy path.
	pr, err := pmuKprobe("printk", true)
	c.Assert(err, qt.IsNil)
	defer pr.Close()

	c.Assert(pr.typ, qt.Equals, kretprobeEvent)

	// Expect os.ErrNotExist when specifying a non-existent kernel symbol
	// on kernels 4.17 and up.
	_, err = pmuKprobe("bogus", false)
	c.Assert(errors.Is(err, os.ErrNotExist), qt.IsTrue, qt.Commentf("got error: %s", err))

	// A kernel bug was fixed in 97c753e62e6c where EINVAL was returned instead
	// of ENOENT, but only for kretprobes.
	_, err = pmuKprobe("bogus", true)
	c.Assert(errors.Is(err, os.ErrNotExist), qt.IsTrue, qt.Commentf("got error: %s", err))
}

// Test fallback behaviour on kernels without perf_kprobe PMU available.
func TestKprobePMUUnavailable(t *testing.T) {
	c := qt.New(t)

	pk, err := pmuKprobe("printk", false)
	if err == nil {
		pk.Close()
		t.Skipf("Kernel supports perf_kprobe PMU, not asserting error.")
	}

	// Only allow a PMU creation with a valid kernel symbol to fail with ErrNotSupported.
	c.Assert(errors.Is(err, ErrNotSupported), qt.IsTrue, qt.Commentf("got error: %s", err))
}

func BenchmarkKprobeCreatePMU(b *testing.B) {
	for n := 0; n < b.N; n++ {
		pr, err := pmuKprobe("printk", false)
		if err != nil {
			b.Error("error creating perf_kprobe PMU:", err)
		}

		if err := pr.Close(); err != nil {
			b.Error("error closing perf_kprobe PMU:", err)
		}
	}
}

// Test tracefs k(ret)probe creation on all kernel versions.
func TestKprobeTraceFS(t *testing.T) {
	c := qt.New(t)

	symbol := "printk"

	// Open and close tracefs k(ret)probes, checking all errors.
	kp, err := tracefsKprobe(symbol, false)
	c.Assert(err, qt.IsNil)
	c.Assert(kp.Close(), qt.IsNil)
	c.Assert(kp.typ, qt.Equals, kprobeEvent)

	kp, err = tracefsKprobe(symbol, true)
	c.Assert(err, qt.IsNil)
	c.Assert(kp.Close(), qt.IsNil)
	c.Assert(kp.typ, qt.Equals, kretprobeEvent)

	// Create two identical trace events, ensure their IDs differ.
	k1, err := tracefsKprobe(symbol, false)
	c.Assert(err, qt.IsNil)
	defer k1.Close()
	c.Assert(k1.tracefsID, qt.Not(qt.Equals), 0)

	k2, err := tracefsKprobe(symbol, false)
	c.Assert(err, qt.IsNil)
	defer k2.Close()
	c.Assert(k2.tracefsID, qt.Not(qt.Equals), 0)

	// Compare the kprobes' tracefs IDs.
	c.Assert(k1.tracefsID, qt.Not(qt.CmpEquals()), k2.tracefsID)

	// Write a k(ret)probe event for a non-existing symbol.
	err = createTraceFSProbeEvent(kprobeType, "testgroup", "bogus", "", 0, false)
	c.Assert(errors.Is(err, os.ErrNotExist), qt.IsTrue, qt.Commentf("got error: %s", err))

	// A kernel bug was fixed in 97c753e62e6c where EINVAL was returned instead
	// of ENOENT, but only for kretprobes.
	err = createTraceFSProbeEvent(kprobeType, "testgroup", "bogus", "", 0, true)
	c.Assert(errors.Is(err, os.ErrNotExist), qt.IsTrue, qt.Commentf("got error: %s", err))
}

func BenchmarkKprobeCreateTraceFS(b *testing.B) {
	for n := 0; n < b.N; n++ {
		// Include <tracefs>/kprobe_events operations in the benchmark loop
		// because we create one per perf event.
		pr, err := tracefsKprobe("printk", false)
		if err != nil {
			b.Error("error creating tracefs perf event:", err)
		}

		if err := pr.Close(); err != nil {
			b.Error("error closing tracefs perf event:", err)
		}
	}
}

// Test k(ret)probe creation writing directly to <tracefs>/kprobe_events.
// Only runs on 5.0 and over. Earlier versions ignored writes of duplicate
// events, while 5.0 started returning -EEXIST when a kprobe event already
// exists.
func TestKprobeCreateTraceFS(t *testing.T) {
	testutils.SkipOnOldKernel(t, "5.0", "<tracefs>/kprobe_events doesn't reject duplicate events")

	c := qt.New(t)

	symbol := "printk"
	pg, _ := randomGroup("ebpftest")
	rg, _ := randomGroup("ebpftest")

	// Tee up cleanups in case any of the Asserts abort the function.
	defer func() {
		_ = closeTraceFSProbeEvent(kprobeType, pg, symbol)
		_ = closeTraceFSProbeEvent(kprobeType, rg, symbol)
	}()

	// Create a kprobe.
	err := createTraceFSProbeEvent(kprobeType, pg, symbol, "", 0, false)
	c.Assert(err, qt.IsNil)

	// Attempt to create an identical kprobe using tracefs,
	// expect it to fail with os.ErrExist.
	err = createTraceFSProbeEvent(kprobeType, pg, symbol, "", 0, false)
	c.Assert(errors.Is(err, os.ErrExist), qt.IsTrue,
		qt.Commentf("expected consecutive kprobe creation to contain os.ErrExist, got: %v", err))

	// Expect a successful close of the kprobe.
	c.Assert(closeTraceFSProbeEvent(kprobeType, pg, symbol), qt.IsNil)

	// Same test for a kretprobe.
	err = createTraceFSProbeEvent(kprobeType, rg, symbol, "", 0, true)
	c.Assert(err, qt.IsNil)

	err = createTraceFSProbeEvent(kprobeType, rg, symbol, "", 0, true)
	c.Assert(os.IsExist(err), qt.IsFalse,
		qt.Commentf("expected consecutive kretprobe creation to contain os.ErrExist, got: %v", err))

	// Expect a successful close of the kretprobe.
	c.Assert(closeTraceFSProbeEvent(kprobeType, rg, symbol), qt.IsNil)
}

func TestKprobeTraceFSGroup(t *testing.T) {
	c := qt.New(t)

	// Expect <prefix>_<16 random hex chars>.
	g, err := randomGroup("ebpftest")
	c.Assert(err, qt.IsNil)
	c.Assert(g, qt.Matches, `ebpftest_[a-f0-9]{16}`)

	// Expect error when the generator's output exceeds 63 characters.
	p := make([]byte, 47) // 63 - 17 (length of the random suffix and underscore) + 1
	for i := range p {
		p[i] = byte('a')
	}
	_, err = randomGroup(string(p))
	c.Assert(err, qt.Not(qt.IsNil))

	// Reject non-alphanumeric characters.
	_, err = randomGroup("/")
	c.Assert(err, qt.Not(qt.IsNil))
}

func TestDetermineRetprobeBit(t *testing.T) {
	testutils.SkipOnOldKernel(t, "4.17", "perf_kprobe PMU")
	c := qt.New(t)

	rpk, err := kretprobeBit()
	c.Assert(err, qt.IsNil)
	c.Assert(rpk, qt.Equals, uint64(0))

	rpu, err := uretprobeBit()
	c.Assert(err, qt.IsNil)
	c.Assert(rpu, qt.Equals, uint64(0))
}

func TestKprobeProgramCall(t *testing.T) {
	// Create ebpf map. Will contain only one key with initial value 0.
	m, err := ebpf.NewMap(&ebpf.MapSpec{
		Name:       "kprobetestmap",
		Type:       ebpf.Array,
		KeySize:    4,
		ValueSize:  4,
		MaxEntries: 1,
	})
	if err != nil {
		t.Fatal(err)
	}

	// Create ebpf program. When called, will set the value of key 0 in
	// the map created above to 1.
	p, err := ebpf.NewProgram(&ebpf.ProgramSpec{
		Name: "kprobetestprog",
		Type: ebpf.Kprobe,
		Instructions: asm.Instructions{
			// u32 key = 0
			asm.Mov.Imm(asm.R1, 0),
			asm.StoreMem(asm.RFP, -4, asm.R1, asm.Word),

			// u32 val = 1
			asm.Mov.Imm(asm.R1, 1),
			asm.StoreMem(asm.RFP, -8, asm.R1, asm.Word),

			// bpf_map_update_elem(...)
			asm.Mov.Reg(asm.R2, asm.RFP),
			asm.Add.Imm(asm.R2, -4),
			asm.Mov.Reg(asm.R3, asm.RFP),
			asm.Add.Imm(asm.R3, -8),
			asm.LoadMapPtr(asm.R1, m.FD()),
			asm.Mov.Imm(asm.R4, 0),
			asm.FnMapUpdateElem.Call(),

			// exit 0
			asm.Mov.Imm(asm.R0, 0),
			asm.Return(),
		},
		License: "Dual MIT/GPL",
	})
	if err != nil {
		t.Fatal(err)
	}

	// Open Kprobe on `__x64_sys_getpid` and attach it
	// to the ebpf program created above.
	k, err := Kprobe("__x64_sys_getpid", p)
	if errors.Is(err, os.ErrNotExist) {
		// Use the correct symbol based on the kernel version.
		// Since 4.17, syscalls symbols are generated with the `__x64_` prefix.
		// https://github.com/torvalds/linux/commit/d5a00528b58cdb2c71206e18bd021e34c4eab878
		k, err = Kprobe("sys_getpid", p)
		if err != nil {
			t.Fatal(err)
		}
	}
	if err != nil {
		t.Fatal(err)
	}
	defer k.Close()

	// Trigger ebpf program call.
	unix.Getpid()

	// Assert that the value has been updated to 1.
	var val uint32
	if err := m.Lookup(uint32(0), &val); err != nil {
		t.Fatal(err)
	}
	if val != 1 {
		t.Fatalf("unexpected value: want '1', got %d", val)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"net/http"
)

func main() {
	http.HandleFunc("/", handler)
	log.Fatal(http.ListenAndServe("localhost:8000", nil))
}

func handler(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "URL.Path = %q\n", r.URL.Path)
}
<commit_msg>Add counter, mutexes, and request echo to server.<commit_after>package main

import (
	"fmt"
	"log"
	"net/http"
	"sync"
)

var mu sync.Mutex
var count int

func main() {
	log.Print("Server running...")
	http.HandleFunc("/", handler)
	http.HandleFunc("/count", counter)
	log.Fatal(http.ListenAndServe("localhost:8000", nil))
}

func handler(w http.ResponseWriter, r *http.Request) {
	mu.Lock()
	count++
	mu.Unlock()
	fmt.Fprintf(w, "%s %s %s\n", r.Method, r.URL, r.Proto)
	for k, v := range r.Header {
		fmt.Fprintf(w, "Header[%q]: %q\n", k, v)
	}
	fmt.Fprintf(w, "Host: %q\n", r.Host)
	fmt.Fprintf(w, "RemoteAddr: %q\n", r.RemoteAddr)
	if err := r.ParseForm(); err != nil {
		log.Print(err)
	}
	for k, v := range r.Form {
		fmt.Fprintf(w, "Form[%q]: %q\n", k, v)
	}
}

func counter(w http.ResponseWriter, r *http.Request) {
	mu.Lock()
	fmt.Fprintf(w, "Count: %d\n", count)
	mu.Unlock()
}
<|endoftext|>


<commit_before>package oauth

import (
	"encoding/xml"
	"github.com/masci/flickr.go/flickr"
	flickErr "github.com/masci/flickr.go/flickr/error"
)

// TODO docs
type CheckTokenResponse struct {
	flickr.FlickrResponse
	OAuth struct {
		XMLName xml.Name `xml:"oauth"`
		Token   string   `xml:"token"`
		Perms   string   `xml:"perms"`
		User    struct {
			XMLName xml.Name `xml:"user"`
			// Flickr ID
			ID string `xml:"nsid,attr"`
			// Flickr Username
			Username string `xml:"username,attr"`
			// Flickr full name
			Fullname string `xml:"fullname,attr"`
		}
	}
}

// TODO docs
func CheckToken(client *flickr.FlickrClient, oauthToken string) (*CheckTokenResponse, error) {
	client.EndpointUrl = flickr.API_ENDPOINT
	client.ClearArgs()
	client.Args.Set("method", "flickr.auth.oauth.checkToken")
	client.Args.Set("oauth_token", oauthToken)
	client.Args.Set("api_key", client.ApiKey)
	client.ApiSign(client.ApiSecret)

	response := &CheckTokenResponse{}
	err := flickr.GetResponse(client, response)

	if err != nil {
		return nil, err
	}

	if response.HasErrors() {
		return response, flickErr.NewError(10)
	}

	return response, nil
}

func getAccessToken() {

}
<commit_msg>added docs<commit_after>package oauth

import (
	"encoding/xml"
	"github.com/masci/flickr.go/flickr"
	flickErr "github.com/masci/flickr.go/flickr/error"
)

// Response type representing data returned by CheckToken
type CheckTokenResponse struct {
	flickr.FlickrResponse
	OAuth struct {
		XMLName xml.Name `xml:"oauth"`
		// OAuth token
		Token string `xml:"token"`
		// String containing permissions bonded to this token
		Perms string `xml:"perms"`
		// The owner of this token
		User struct {
			XMLName xml.Name `xml:"user"`
			// Flickr ID
			ID string `xml:"nsid,attr"`
			// Flickr Username
			Username string `xml:"username,attr"`
			// Flickr full name
			Fullname string `xml:"fullname,attr"`
		}
	}
}

// Returns the credentials attached to an OAuth authentication token.
// This method does not require user authentication, but the request must be api-signed.
func CheckToken(client *flickr.FlickrClient, oauthToken string) (*CheckTokenResponse, error) {
	client.EndpointUrl = flickr.API_ENDPOINT
	client.ClearArgs()
	client.Args.Set("method", "flickr.auth.oauth.checkToken")
	client.Args.Set("oauth_token", oauthToken)
	client.Args.Set("api_key", client.ApiKey)
	client.ApiSign(client.ApiSecret)

	response := &CheckTokenResponse{}
	err := flickr.GetResponse(client, response)

	if err != nil {
		return nil, err
	}

	if response.HasErrors() {
		return response, flickErr.NewError(10)
	}

	return response, nil
}

func getAccessToken() {

}
<|endoftext|>


<commit_before>package consumergroup

import (
	"errors"
	"sync"
	"time"

	"github.com/Shopify/sarama"
)

type ConsumerGroupConfig struct {
	// The Zookeeper read timeout
	ZookeeperTimeout time.Duration

	// Zookeeper chroot to use. Should not include a trailing slash.
	// Leave this empty for to not set a chroot.
	ZookeeperChroot string

	KafkaClientConfig   *sarama.ClientConfig   // This will be passed to Sarama when creating a new sarama.Client
	KafkaConsumerConfig *sarama.ConsumerConfig // This will be passed to Sarama when creating a new sarama.Consumer

	ChannelBufferSize int
	CommitInterval    time.Duration
}

func NewConsumerGroupConfig() *ConsumerGroupConfig {
	return &ConsumerGroupConfig{
		ZookeeperTimeout:    1 * time.Second,
		KafkaClientConfig:   sarama.NewClientConfig(),
		KafkaConsumerConfig: sarama.NewConsumerConfig(),
		ChannelBufferSize:   10,
		CommitInterval:      10 * time.Second,
	}
}

func (cgc *ConsumerGroupConfig) Validate() error {
	if cgc.ZookeeperTimeout <= 0 {
		return errors.New("ZookeeperTimeout should have a duration > 0")
	}

	if cgc.KafkaClientConfig == nil {
		return errors.New("KafkaClientConfig is not set!")
	} else if err := cgc.KafkaClientConfig.Validate(); err != nil {
		return err
	}

	if cgc.KafkaConsumerConfig == nil {
		return errors.New("KafkaConsumerConfig is not set!")
	} else if err := cgc.KafkaConsumerConfig.Validate(); err != nil {
		return err
	}

	return nil
}

type ConsumerGroup struct {
	id, name string

	config *ConsumerGroupConfig

	client *sarama.Client
	zk     *ZK

	wg sync.WaitGroup

	events  chan *sarama.ConsumerEvent
	stopper chan struct{}

	brokers   map[int]string
	consumers []string
	offsets   map[string]map[int32]int64
}

// Connects to a consumer group, using Zookeeper for auto-discovery
func JoinConsumerGroup(name string, topics []string, zookeeper []string, config *ConsumerGroupConfig) (cg *ConsumerGroup, err error) {

	if name == "" {
		return nil, sarama.ConfigurationError("Empty consumergroup name")
	}

	if len(topics) == 0 {
		return nil, sarama.ConfigurationError("No topics provided")
	}

	if len(zookeeper) == 0 {
		return nil, errors.New("You need to provide at least one zookeeper node address!")
	}

	if config == nil {
		config = NewConsumerGroupConfig()
	}

	// Validate configuration
	if err = config.Validate(); err != nil {
		return
	}

	var zk *ZK
	if zk, err = NewZK(zookeeper, config.ZookeeperChroot, config.ZookeeperTimeout); err != nil {
		return
	}

	brokers, err := zk.Brokers()
	if err != nil {
		return
	}

	brokerList := make([]string, 0, len(brokers))
	for _, broker := range brokers {
		brokerList = append(brokerList, broker)
	}

	var client *sarama.Client
	if client, err = sarama.NewClient(name, brokerList, config.KafkaClientConfig); err != nil {
		return
	}

	var consumerID string
	consumerID, err = generateConsumerID()
	if err != nil {
		return
	}

	// Register consumer group
	if err = zk.RegisterGroup(name); err != nil {
		return
	}

	// Register itself with zookeeper
	if err = zk.RegisterConsumer(name, consumerID, topics); err != nil {
		sarama.Logger.Printf("[%s] FAILED to register consumer %s!\n", name, consumerID)
		return
	} else {
		sarama.Logger.Printf("[%s] Registered consumer %s.\n", name, consumerID)
	}

	group := &ConsumerGroup{
		id:      consumerID,
		name:    name,
		config:  config,
		brokers: brokers,
		client:  client,
		zk:      zk,
		events:  make(chan *sarama.ConsumerEvent, config.ChannelBufferSize),
		stopper: make(chan struct{}),
		offsets: make(map[string]map[int32]int64),
	}

	go group.topicListConsumer(topics)

	return group, nil
}

func (cg *ConsumerGroup) Events() <-chan *sarama.ConsumerEvent {
	return cg.events
}

func (cg *ConsumerGroup) Close() (err error) {
	defer cg.zk.Close()
	close(cg.stopper)
	cg.wg.Wait()

	cg.client.Close()

	if err = cg.zk.DeregisterConsumer(cg.name, cg.id); err != nil {
		sarama.Logger.Printf("[%s] FAILED deregistering consumer %s!\n", cg.name, cg.id)
		return err
	} else {
		sarama.Logger.Printf("[%s] Deregistered consumer %s.\n", cg.name, cg.id)
	}

	close(cg.events)
	return
}

func (cg *ConsumerGroup) topicListConsumer(topics []string) {
	for {

		consumers, consumerChanges, err := cg.zk.Consumers(cg.name)
		if err != nil {
			panic(err)
		}

		cg.consumers = consumers
		sarama.Logger.Printf("[%s] Currently registered consumers: %d\n", cg.name, len(cg.consumers))

		stopper := make(chan struct{})

		for _, topic := range topics {
			cg.wg.Add(1)
			go cg.topicConsumer(topic, cg.events, stopper)
		}

		select {
		case <-cg.stopper:
			close(stopper)
			return

		case <-consumerChanges:
			sarama.Logger.Printf("[%s] Triggering rebalance due to consumer list change.\n", cg.name)
			close(stopper)
			cg.wg.Wait()
		}
	}
}

func (cg *ConsumerGroup) topicConsumer(topic string, events chan<- *sarama.ConsumerEvent, stopper <-chan struct{}) {
	defer cg.wg.Done()

	sarama.Logger.Printf("[%s] Started topic consumer for %s\n", cg.name, topic)

	// Fetch a list of partition IDs
	partitions, err := cg.zk.Partitions(topic)
	if err != nil {
		panic(err)
	}

	dividedPartitions := dividePartitionsBetweenConsumers(cg.consumers, partitions)
	myPartitions := dividedPartitions[cg.id]
	sarama.Logger.Printf("[%s] Claiming %d of %d partitions for topic %s.", cg.name, len(myPartitions), len(partitions), topic)

	// Consume all the assigned partitions
	var wg sync.WaitGroup
	for _, pid := range myPartitions {

		wg.Add(1)
		go cg.partitionConsumer(topic, pid.id, events, &wg, stopper)
	}

	wg.Wait()
	sarama.Logger.Printf("[%s] Stopped topic consumer for %s\n", cg.name, topic)
}

// Consumes a partition
func (cg *ConsumerGroup) partitionConsumer(topic string, partition int32, events chan<- *sarama.ConsumerEvent, wg *sync.WaitGroup, stopper <-chan struct{}) {
	defer wg.Done()

	err := cg.zk.Claim(cg.name, topic, partition, cg.id)
	if err != nil {
		panic(err)
	}
	defer cg.zk.Release(cg.name, topic, partition, cg.id)

	lastOffset, err := cg.zk.Offset(cg.name, topic, partition)
	if err != nil {
		panic(err)
	}

	if lastOffset > 0 {
		cg.config.KafkaConsumerConfig.OffsetMethod = sarama.OffsetMethodManual
		cg.config.KafkaConsumerConfig.OffsetValue = lastOffset + 1
	} else {
		cg.config.KafkaConsumerConfig.OffsetMethod = sarama.OffsetMethodOldest
	}

	consumer, err := sarama.NewConsumer(cg.client, topic, partition, cg.name, cg.config.KafkaConsumerConfig)
	if err != nil {
		panic(err)
	}
	defer consumer.Close()

	sarama.Logger.Printf("[%s] Started partition consumer for %s:%d at offset %d.\n", cg.name, topic, partition, lastOffset)

	partitionEvents := consumer.Events()
	commitInterval := time.After(cg.config.CommitInterval)
partitionConsumerLoop:
	for {
		select {
		case event := <-partitionEvents:
			if event.Err != nil {
				panic(event.Err)
			}

			lastOffset = event.Offset
			events <- event

		case <-commitInterval:
			if err := cg.zk.Commit(cg.name, topic, partition, lastOffset); err != nil {
				sarama.Logger.Printf("[%s] Failed to commit offset for %s:%d\n", cg.name, topic, partition)
			}

		case <-stopper:
			break partitionConsumerLoop
		}
	}

	if err := cg.zk.Commit(cg.name, topic, partition, lastOffset); err != nil {
		sarama.Logger.Printf("[%s] Failed to commit offset for %s:%d\n", cg.name, topic, partition)
	}

	sarama.Logger.Printf("[%s] Stopping partition consumer for %s:%d at offset %d.\n", cg.name, topic, partition, lastOffset)
}
<commit_msg>Fix race condition when rebelance and close would happen simultaneously<commit_after>package consumergroup

import (
	"errors"
	"sync"
	"time"

	"github.com/Shopify/sarama"
)

type ConsumerGroupConfig struct {
	// The Zookeeper read timeout
	ZookeeperTimeout time.Duration

	// Zookeeper chroot to use. Should not include a trailing slash.
	// Leave this empty for to not set a chroot.
	ZookeeperChroot string

	KafkaClientConfig   *sarama.ClientConfig   // This will be passed to Sarama when creating a new sarama.Client
	KafkaConsumerConfig *sarama.ConsumerConfig // This will be passed to Sarama when creating a new sarama.Consumer

	ChannelBufferSize int
	CommitInterval    time.Duration
}

func NewConsumerGroupConfig() *ConsumerGroupConfig {
	return &ConsumerGroupConfig{
		ZookeeperTimeout:    1 * time.Second,
		KafkaClientConfig:   sarama.NewClientConfig(),
		KafkaConsumerConfig: sarama.NewConsumerConfig(),
		ChannelBufferSize:   10,
		CommitInterval:      10 * time.Second,
	}
}

func (cgc *ConsumerGroupConfig) Validate() error {
	if cgc.ZookeeperTimeout <= 0 {
		return errors.New("ZookeeperTimeout should have a duration > 0")
	}

	if cgc.KafkaClientConfig == nil {
		return errors.New("KafkaClientConfig is not set!")
	} else if err := cgc.KafkaClientConfig.Validate(); err != nil {
		return err
	}

	if cgc.KafkaConsumerConfig == nil {
		return errors.New("KafkaConsumerConfig is not set!")
	} else if err := cgc.KafkaConsumerConfig.Validate(); err != nil {
		return err
	}

	return nil
}

type ConsumerGroup struct {
	id, name string

	config *ConsumerGroupConfig

	client *sarama.Client
	zk     *ZK

	wg sync.WaitGroup

	events  chan *sarama.ConsumerEvent
	stopper chan struct{}

	brokers   map[int]string
	consumers []string
	offsets   map[string]map[int32]int64
}

// Connects to a consumer group, using Zookeeper for auto-discovery
func JoinConsumerGroup(name string, topics []string, zookeeper []string, config *ConsumerGroupConfig) (cg *ConsumerGroup, err error) {

	if name == "" {
		return nil, sarama.ConfigurationError("Empty consumergroup name")
	}

	if len(topics) == 0 {
		return nil, sarama.ConfigurationError("No topics provided")
	}

	if len(zookeeper) == 0 {
		return nil, errors.New("You need to provide at least one zookeeper node address!")
	}

	if config == nil {
		config = NewConsumerGroupConfig()
	}

	// Validate configuration
	if err = config.Validate(); err != nil {
		return
	}

	var zk *ZK
	if zk, err = NewZK(zookeeper, config.ZookeeperChroot, config.ZookeeperTimeout); err != nil {
		return
	}

	brokers, err := zk.Brokers()
	if err != nil {
		return
	}

	brokerList := make([]string, 0, len(brokers))
	for _, broker := range brokers {
		brokerList = append(brokerList, broker)
	}

	var client *sarama.Client
	if client, err = sarama.NewClient(name, brokerList, config.KafkaClientConfig); err != nil {
		return
	}

	var consumerID string
	consumerID, err = generateConsumerID()
	if err != nil {
		return
	}

	// Register consumer group
	if err = zk.RegisterGroup(name); err != nil {
		return
	}

	// Register itself with zookeeper
	if err = zk.RegisterConsumer(name, consumerID, topics); err != nil {
		sarama.Logger.Printf("[%s] FAILED to register consumer %s!\n", name, consumerID)
		return
	} else {
		sarama.Logger.Printf("[%s] Registered consumer %s.\n", name, consumerID)
	}

	group := &ConsumerGroup{
		id:      consumerID,
		name:    name,
		config:  config,
		brokers: brokers,
		client:  client,
		zk:      zk,
		events:  make(chan *sarama.ConsumerEvent, config.ChannelBufferSize),
		stopper: make(chan struct{}),
		offsets: make(map[string]map[int32]int64),
	}

	go group.topicListConsumer(topics)

	return group, nil
}

func (cg *ConsumerGroup) Events() <-chan *sarama.ConsumerEvent {
	return cg.events
}

func (cg *ConsumerGroup) Close() (err error) {
	defer cg.zk.Close()
	close(cg.stopper)
	cg.wg.Wait()

	cg.client.Close()

	if err = cg.zk.DeregisterConsumer(cg.name, cg.id); err != nil {
		sarama.Logger.Printf("[%s] FAILED deregistering consumer %s!\n", cg.name, cg.id)
		return err
	} else {
		sarama.Logger.Printf("[%s] Deregistered consumer %s.\n", cg.name, cg.id)
	}

	close(cg.events)
	return
}

func (cg *ConsumerGroup) topicListConsumer(topics []string) {
	for {
		select {
		case <-cg.stopper:
			return
		default:
		}

		consumers, consumerChanges, err := cg.zk.Consumers(cg.name)
		if err != nil {
			panic(err)
		}

		cg.consumers = consumers
		sarama.Logger.Printf("[%s] Currently registered consumers: %d\n", cg.name, len(cg.consumers))

		stopper := make(chan struct{})

		for _, topic := range topics {
			cg.wg.Add(1)
			go cg.topicConsumer(topic, cg.events, stopper)
		}

		select {
		case <-cg.stopper:
			close(stopper)
			return

		case <-consumerChanges:
			sarama.Logger.Printf("[%s] Triggering rebalance due to consumer list change.\n", cg.name)
			close(stopper)
			cg.wg.Wait()
		}
	}
}

func (cg *ConsumerGroup) topicConsumer(topic string, events chan<- *sarama.ConsumerEvent, stopper <-chan struct{}) {
	defer cg.wg.Done()

	sarama.Logger.Printf("[%s] Started topic consumer for %s\n", cg.name, topic)

	// Fetch a list of partition IDs
	partitions, err := cg.zk.Partitions(topic)
	if err != nil {
		panic(err)
	}

	dividedPartitions := dividePartitionsBetweenConsumers(cg.consumers, partitions)
	myPartitions := dividedPartitions[cg.id]
	sarama.Logger.Printf("[%s] Claiming %d of %d partitions for topic %s.", cg.name, len(myPartitions), len(partitions), topic)

	// Consume all the assigned partitions
	var wg sync.WaitGroup
	for _, pid := range myPartitions {

		wg.Add(1)
		go cg.partitionConsumer(topic, pid.id, events, &wg, stopper)
	}

	wg.Wait()
	sarama.Logger.Printf("[%s] Stopped topic consumer for %s\n", cg.name, topic)
}

// Consumes a partition
func (cg *ConsumerGroup) partitionConsumer(topic string, partition int32, events chan<- *sarama.ConsumerEvent, wg *sync.WaitGroup, stopper <-chan struct{}) {
	defer wg.Done()

	err := cg.zk.Claim(cg.name, topic, partition, cg.id)
	if err != nil {
		panic(err)
	}
	defer cg.zk.Release(cg.name, topic, partition, cg.id)

	lastOffset, err := cg.zk.Offset(cg.name, topic, partition)
	if err != nil {
		panic(err)
	}

	if lastOffset > 0 {
		cg.config.KafkaConsumerConfig.OffsetMethod = sarama.OffsetMethodManual
		cg.config.KafkaConsumerConfig.OffsetValue = lastOffset + 1
	} else {
		cg.config.KafkaConsumerConfig.OffsetMethod = sarama.OffsetMethodOldest
	}

	consumer, err := sarama.NewConsumer(cg.client, topic, partition, cg.name, cg.config.KafkaConsumerConfig)
	if err != nil {
		panic(err)
	}
	defer consumer.Close()

	sarama.Logger.Printf("[%s] Started partition consumer for %s:%d at offset %d.\n", cg.name, topic, partition, lastOffset)

	partitionEvents := consumer.Events()
	commitInterval := time.After(cg.config.CommitInterval)
partitionConsumerLoop:
	for {
		select {
		case event := <-partitionEvents:
			if event.Err != nil {
				panic(event.Err)
			}

			lastOffset = event.Offset
			events <- event

		case <-commitInterval:
			if err := cg.zk.Commit(cg.name, topic, partition, lastOffset); err != nil {
				sarama.Logger.Printf("[%s] Failed to commit offset for %s:%d\n", cg.name, topic, partition)
			}

		case <-stopper:
			break partitionConsumerLoop
		}
	}

	if err := cg.zk.Commit(cg.name, topic, partition, lastOffset); err != nil {
		sarama.Logger.Printf("[%s] Failed to commit offset for %s:%d\n", cg.name, topic, partition)
	}

	sarama.Logger.Printf("[%s] Stopping partition consumer for %s:%d at offset %d.\n", cg.name, topic, partition, lastOffset)
}
<|endoftext|>


<commit_before>package main

import "fmt"
import "os"
import "log"

import "github.com/moovweb/gokogiri/xpath"
import "github.com/moovweb/gokogiri/xml"
import "github.com/codegangsta/cli"
import "github.com/gosuri/uitable"

func multiple(vrouter string, vrf_name string, count bool) {
	url := "http://" + vrouter + ":8085" + "/Snh_PageReq?x=begin:-1,end:-1,table:" + vrf_name + ".uc.route.0,"

	var doc = load(url, false)
	defer doc.Free()
	xps := xpath.Compile("//route_list/list/RouteUcSandeshData/path_list/list/PathSandeshData/nh/NhSandeshData/mc_list/../../../../../../src_ip/text()")
	ss, _ := doc.Root().Search(xps)
	if count {
		fmt.Printf("%d\n", len(ss))
	} else {
		for _, s := range ss {
			fmt.Printf("%s\n", s)
		}
	}
}

func DescPeering() DescCollection {
	return DescCollection{
		PageArgs:  []string{"vrouter-fqdn"},
		BaseXpath: "AgentXmppConnectionStatus/peer/list",
		PageBuilder: func(args []string) Sourcer {
			return Webui{Path: "Snh_AgentXmppConnectionStatusReq", VrouterUrl: args[0], Port: 8085}
		},
		DescElt: DescElement{
			ShortDetailXpath: "controller_ip/text()",
			LongDetail:       LongFormatXpaths([]string{"controller_ip", "state", "flap_count", "cfg_controller"}),
		},
		PrimaryField: "name",
	}
}

func DescItf() DescCollection {
	return DescCollection{
		BaseXpath: "__ItfResp_list/ItfResp/itf_list/list",
		DescElt: DescElement{
			ShortDetailXpath: "name/text()",
			LongDetail:       LongFormatXpaths([]string{"uuid", "name", "vrf_name", "vm_uuid"}),
		},
		PageArgs: []string{"vrouter-fqdn"},
		PageBuilder: func(args []string) Sourcer {
			return Remote{Table: "db.interface.0", VrouterUrl: args[0], Port: 8085}
		},
		PrimaryField: "name",
	}
}
func DescRoute() DescCollection {
	return DescCollection{
		PageArgs: []string{"vrouter-fqdn", "vrf-name"},
		PageBuilder: func(args []string) Sourcer {
			return Remote{VrouterUrl: args[0], Table: args[1] + ".uc.route.0,", Port: 8085}
		},
		BaseXpath: "__Inet4UcRouteResp_list/Inet4UcRouteResp/route_list/list",
		DescElt: DescElement{
			ShortDetailXpath: "src_ip/text()",
			LongDetail:       LongFormatFn(routeDetail)},
		PrimaryField: "src_ip",
	}
}
func DescVrf() DescCollection {
	return DescCollection{
		PageArgs: []string{"vrouter-fqdn"},
		PageBuilder: func(args []string) Sourcer {
			return Remote{Table: "db.vrf.0", VrouterUrl: args[0], Port: 8085}
		},
		BaseXpath: "__VrfListResp_list/VrfListResp/vrf_list/list",
		DescElt: DescElement{
			ShortDetailXpath: "name/text()",
			LongDetail:       LongFormatXpaths([]string{"name", "uc_index"}),
		},
		PrimaryField: "name",
	}
}
func DescVn() DescCollection {
	return DescCollection{
		PageArgs: []string{"vrouter-fqdn"},
		PageBuilder: func(args []string) Sourcer {
			return Remote{Table: "db.vn.0", VrouterUrl: args[0], Port: 8085}
		},
		BaseXpath: "__VnListResp_list/VnListResp/vn_list/list",
		DescElt: DescElement{
			ShortDetailXpath: "name/text()",
			LongDetail:       LongFormatXpaths([]string{"name", "vrf_name"}),
		},
		PrimaryField: "name",
	}
}

func DescRiSummary() DescCollection {
	return DescCollection{
		PageArgs: []string{"controller-fqdn"},
		PageBuilder: func(args []string) Sourcer {
			return Webui{Path: "Snh_ShowRoutingInstanceSummaryReq?search_string=", VrouterUrl: args[0], Port: 8083}
		},
		BaseXpath: "ShowRoutingInstanceSummaryResp/instances/list",
		DescElt: DescElement{
			ShortDetailXpath: "name/text()",
			LongDetail:       LongFormatXpaths([]string{"name", "virtual_network"}),
		},
		PrimaryField: "name",
	}
}

func DescCtrlRouteSummary() DescCollection {
	return DescCollection{
		PageArgs: []string{"controller-fqdn", "search"},
		PageBuilder: func(args []string) Sourcer {
			path := fmt.Sprintf("Snh_ShowRouteSummaryReq?search_string=%s", args[1])
			return Webui{Path: path, VrouterUrl: args[0], Port: 8083}
		},
		BaseXpath: "ShowRouteSummaryResp/tables/list",
		DescElt: DescElement{
			ShortDetailXpath: "name/text()",
			LongDetail:       LongFormatXpaths([]string{"name", "prefixes", "paths", "primary_paths", "secondary_paths", "pending_updates"}),
		},
		PrimaryField: "name",
	}
}

func DescCtrlRoute() DescCollection {
	return DescCollection{
		PageArgs: []string{"controller-fqdn", "routing-instance"},
		PageBuilder: func(args []string) Sourcer {
			path := fmt.Sprintf("Snh_ShowRouteReq?x=%s.inet.0", args[1])
			return Webui{Path: path, VrouterUrl: args[0], Port: 8083}
		},
		BaseXpath: "ShowRouteResp/tables/list/ShowRouteTable/routes/list",
		DescElt: DescElement{
			ShortDetailXpath: "prefix/text()",
			LongDetail:       LongFormatFn(controllerRoutePath),
		},
		PrimaryField: "prefix",
	}
}

func DescMpls() DescCollection {
	return DescCollection{
		PageArgs: []string{"vrouter-fqdn"},
		PageBuilder: func(args []string) Sourcer {
			return Remote{Table: "db.mpls.0", VrouterUrl: args[0], Port: 8085}
		},
		BaseXpath: "__MplsResp_list/MplsResp/mpls_list/list",
		DescElt: DescElement{
			ShortDetailXpath: "label/text()",
			LongDetail:       LongFormatFn(mplsDetail),
		},
		PrimaryField: "label",
	}
}

func routeSummaryDetail(e Element) {
	table := uitable.New()
	table.MaxColWidth = 80
	table.AddRow("Name", "Prefixes", "Paths", "Primary paths", "Secondary paths", "Pending Updates")
	fields := []string{"name", "prefixes", "paths", "primary_paths",
		"secondary_paths", "pending_updates"}
	paths, _ := e.node.Search(".")
	for _, path := range paths {
		values := [6]string{}
		for i, field := range fields {
			value, _ := path.Search(fmt.Sprintf("%s/text()", field))
			values[i] = Pretty(value)
		}
		table.AddRow(values[0], values[1], values[2], values[3], values[4], values[5])
	}
	fmt.Println(table)
}

func routeDetail(e Element) {
	srcIp, _ := e.node.Search("src_ip/text()")
	fmt.Printf("Src %s\n", srcIp[0])
	paths, _ := e.node.Search("path_list/list/PathSandeshData")

	table := uitable.New()
	table.MaxColWidth = 80
	table.AddRow("    Dst", "Peers", "MPLS label", "Interface", "Dest VN")
	for _, path := range paths {
		nhs, _ := path.Search("nh/NhSandeshData//dip/text()")
		peers, _ := path.Search("peer/text()")
		label, _ := path.Search("label/text()")
		destvn, _ := path.Search("dest_vn/text()")
		itf, _ := path.Search("nh/NhSandeshData/itf/text()")
		table.AddRow("    "+Pretty(nhs), Pretty(peers), Pretty(label), Pretty(itf), Pretty(destvn))
	}
	fmt.Println(table)
}

func mplsDetail(e Element) {
	fmt.Printf("Label: %s\n", e.GetField("label"))
	nexthopDetail(e.node)
}

func nexthopDetail(node xml.Node) {
	table := uitable.New()
	table.MaxColWidth = 80
	table.AddRow("    Type", "Interface", "Nexthop index")
	nhs, _ := node.Search("nh/NhSandeshData/type/text()")
	itf, _ := node.Search("nh/NhSandeshData/itf/text()")
	nhIdx, _ := node.Search("nh/NhSandeshData/nh_index/text()")
	table.AddRow("    "+Pretty(nhs), Pretty(itf), Pretty(nhIdx))
	fmt.Println(table)
}

func controllerRoutePath(e Element) {
	srcIp, _ := e.node.Search("prefix/text()")
	fmt.Printf("Prefix %s\n", srcIp[0])
	paths, _ := e.node.Search("paths/list/ShowRoutePath")

	table := uitable.New()
	table.MaxColWidth = 80
	table.AddRow("    Protocol", "Nexthop", "Local Pref", "Peers", "MPLS label")
	for _, path := range paths {
		protocol, _ := path.Search("protocol/text()")
		nhs, _ := path.Search("next_hop/text()")
		peers, _ := path.Search("source/text()")
		label, _ := path.Search("label/text()")
		localPref, _ := path.Search("local_preference/text()")
		table.AddRow("    "+Pretty(protocol), Pretty(nhs), Pretty(localPref), Pretty(peers), Pretty(label))
	}
	fmt.Println(table)
}

func main() {
	var count bool
	var hosts_file string

	app := cli.NewApp()
	app.Name = "contrail-introspect-cli"
	app.Usage = "CLI on ContraiL Introspects"
	app.Version = "0.0.3"
	app.EnableBashCompletion = true
	app.Before = func(c *cli.Context) error {
		if c.GlobalIsSet("hosts") {
			var err error
			hosts, err = LoadHostsFile(c.GlobalString("hosts"))
			return err
		}
		return nil
	}
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:        "hosts",
			Usage:       "host file to do DNS resolution",
			Destination: &hosts_file,
		}}
	app.Commands = []cli.Command{
		GenCommand(DescRoute(), "agent-route", "Show routes on agent"),
		GenCommand(DescItf(), "agent-itf", "Show interfaces on agent"),
		GenCommand(DescVrf(), "agent-vrf", "Show vrfs on agent "),
		GenCommand(DescPeering(), "agent-peering", "Peering with controller on agent"),
		GenCommand(DescVn(), "agent-vn", "Show virtual networks on agent"),
		GenCommand(DescMpls(), "agent-mpls", "Show mpls on agent"),
		Follow(),
		Path(),
		GenCommand(DescRiSummary(), "controller-ri", "Show routing instances on controller"),
		GenCommand(DescCtrlRoute(), "controller-route", "Show routes on controller"),
		GenCommand(DescCtrlRouteSummary(), "controller-route-summary", "Show routes summary on controller"),
		{
			Name:      "agent-multiple",
			Usage:     "List routes with multiple nexthops",
			ArgsUsage: "vrouter vrf_name",
			Flags: []cli.Flag{
				cli.BoolFlag{
					Name:        "count",
					Destination: &count,
				}},
			Action: func(c *cli.Context) error {
				if c.NArg() != 2 {
					log.Fatal("Wrong argument number!")
				}
				vrouter := c.Args()[0]
				vrf_name := c.Args()[1]
				multiple(vrouter, vrf_name, count)
				return nil
			},
		},
	}
	app.Run(os.Args)
}
<commit_msg>Release 0.0.4<commit_after>package main

import "fmt"
import "os"
import "log"

import "github.com/moovweb/gokogiri/xpath"
import "github.com/moovweb/gokogiri/xml"
import "github.com/codegangsta/cli"
import "github.com/gosuri/uitable"

func multiple(vrouter string, vrf_name string, count bool) {
	url := "http://" + vrouter + ":8085" + "/Snh_PageReq?x=begin:-1,end:-1,table:" + vrf_name + ".uc.route.0,"

	var doc = load(url, false)
	defer doc.Free()
	xps := xpath.Compile("//route_list/list/RouteUcSandeshData/path_list/list/PathSandeshData/nh/NhSandeshData/mc_list/../../../../../../src_ip/text()")
	ss, _ := doc.Root().Search(xps)
	if count {
		fmt.Printf("%d\n", len(ss))
	} else {
		for _, s := range ss {
			fmt.Printf("%s\n", s)
		}
	}
}

func DescPeering() DescCollection {
	return DescCollection{
		PageArgs:  []string{"vrouter-fqdn"},
		BaseXpath: "AgentXmppConnectionStatus/peer/list",
		PageBuilder: func(args []string) Sourcer {
			return Webui{Path: "Snh_AgentXmppConnectionStatusReq", VrouterUrl: args[0], Port: 8085}
		},
		DescElt: DescElement{
			ShortDetailXpath: "controller_ip/text()",
			LongDetail:       LongFormatXpaths([]string{"controller_ip", "state", "flap_count", "cfg_controller"}),
		},
		PrimaryField: "name",
	}
}

func DescItf() DescCollection {
	return DescCollection{
		BaseXpath: "__ItfResp_list/ItfResp/itf_list/list",
		DescElt: DescElement{
			ShortDetailXpath: "name/text()",
			LongDetail:       LongFormatXpaths([]string{"uuid", "name", "vrf_name", "vm_uuid"}),
		},
		PageArgs: []string{"vrouter-fqdn"},
		PageBuilder: func(args []string) Sourcer {
			return Remote{Table: "db.interface.0", VrouterUrl: args[0], Port: 8085}
		},
		PrimaryField: "name",
	}
}
func DescRoute() DescCollection {
	return DescCollection{
		PageArgs: []string{"vrouter-fqdn", "vrf-name"},
		PageBuilder: func(args []string) Sourcer {
			return Remote{VrouterUrl: args[0], Table: args[1] + ".uc.route.0,", Port: 8085}
		},
		BaseXpath: "__Inet4UcRouteResp_list/Inet4UcRouteResp/route_list/list",
		DescElt: DescElement{
			ShortDetailXpath: "src_ip/text()",
			LongDetail:       LongFormatFn(routeDetail)},
		PrimaryField: "src_ip",
	}
}
func DescVrf() DescCollection {
	return DescCollection{
		PageArgs: []string{"vrouter-fqdn"},
		PageBuilder: func(args []string) Sourcer {
			return Remote{Table: "db.vrf.0", VrouterUrl: args[0], Port: 8085}
		},
		BaseXpath: "__VrfListResp_list/VrfListResp/vrf_list/list",
		DescElt: DescElement{
			ShortDetailXpath: "name/text()",
			LongDetail:       LongFormatXpaths([]string{"name", "uc_index"}),
		},
		PrimaryField: "name",
	}
}
func DescVn() DescCollection {
	return DescCollection{
		PageArgs: []string{"vrouter-fqdn"},
		PageBuilder: func(args []string) Sourcer {
			return Remote{Table: "db.vn.0", VrouterUrl: args[0], Port: 8085}
		},
		BaseXpath: "__VnListResp_list/VnListResp/vn_list/list",
		DescElt: DescElement{
			ShortDetailXpath: "name/text()",
			LongDetail:       LongFormatXpaths([]string{"name", "vrf_name"}),
		},
		PrimaryField: "name",
	}
}

func DescRiSummary() DescCollection {
	return DescCollection{
		PageArgs: []string{"controller-fqdn"},
		PageBuilder: func(args []string) Sourcer {
			return Webui{Path: "Snh_ShowRoutingInstanceSummaryReq?search_string=", VrouterUrl: args[0], Port: 8083}
		},
		BaseXpath: "ShowRoutingInstanceSummaryResp/instances/list",
		DescElt: DescElement{
			ShortDetailXpath: "name/text()",
			LongDetail:       LongFormatXpaths([]string{"name", "virtual_network"}),
		},
		PrimaryField: "name",
	}
}

func DescCtrlRouteSummary() DescCollection {
	return DescCollection{
		PageArgs: []string{"controller-fqdn", "search"},
		PageBuilder: func(args []string) Sourcer {
			path := fmt.Sprintf("Snh_ShowRouteSummaryReq?search_string=%s", args[1])
			return Webui{Path: path, VrouterUrl: args[0], Port: 8083}
		},
		BaseXpath: "ShowRouteSummaryResp/tables/list",
		DescElt: DescElement{
			ShortDetailXpath: "name/text()",
			LongDetail:       LongFormatXpaths([]string{"name", "prefixes", "paths", "primary_paths", "secondary_paths", "pending_updates"}),
		},
		PrimaryField: "name",
	}
}

func DescCtrlRoute() DescCollection {
	return DescCollection{
		PageArgs: []string{"controller-fqdn", "routing-instance"},
		PageBuilder: func(args []string) Sourcer {
			path := fmt.Sprintf("Snh_ShowRouteReq?x=%s.inet.0", args[1])
			return Webui{Path: path, VrouterUrl: args[0], Port: 8083}
		},
		BaseXpath: "ShowRouteResp/tables/list/ShowRouteTable/routes/list",
		DescElt: DescElement{
			ShortDetailXpath: "prefix/text()",
			LongDetail:       LongFormatFn(controllerRoutePath),
		},
		PrimaryField: "prefix",
	}
}

func DescMpls() DescCollection {
	return DescCollection{
		PageArgs: []string{"vrouter-fqdn"},
		PageBuilder: func(args []string) Sourcer {
			return Remote{Table: "db.mpls.0", VrouterUrl: args[0], Port: 8085}
		},
		BaseXpath: "__MplsResp_list/MplsResp/mpls_list/list",
		DescElt: DescElement{
			ShortDetailXpath: "label/text()",
			LongDetail:       LongFormatFn(mplsDetail),
		},
		PrimaryField: "label",
	}
}

func routeSummaryDetail(e Element) {
	table := uitable.New()
	table.MaxColWidth = 80
	table.AddRow("Name", "Prefixes", "Paths", "Primary paths", "Secondary paths", "Pending Updates")
	fields := []string{"name", "prefixes", "paths", "primary_paths",
		"secondary_paths", "pending_updates"}
	paths, _ := e.node.Search(".")
	for _, path := range paths {
		values := [6]string{}
		for i, field := range fields {
			value, _ := path.Search(fmt.Sprintf("%s/text()", field))
			values[i] = Pretty(value)
		}
		table.AddRow(values[0], values[1], values[2], values[3], values[4], values[5])
	}
	fmt.Println(table)
}

func routeDetail(e Element) {
	srcIp, _ := e.node.Search("src_ip/text()")
	fmt.Printf("Src %s\n", srcIp[0])
	paths, _ := e.node.Search("path_list/list/PathSandeshData")

	table := uitable.New()
	table.MaxColWidth = 80
	table.AddRow("    Dst", "Peers", "MPLS label", "Interface", "Dest VN")
	for _, path := range paths {
		nhs, _ := path.Search("nh/NhSandeshData//dip/text()")
		peers, _ := path.Search("peer/text()")
		label, _ := path.Search("label/text()")
		destvn, _ := path.Search("dest_vn/text()")
		itf, _ := path.Search("nh/NhSandeshData/itf/text()")
		table.AddRow("    "+Pretty(nhs), Pretty(peers), Pretty(label), Pretty(itf), Pretty(destvn))
	}
	fmt.Println(table)
}

func mplsDetail(e Element) {
	fmt.Printf("Label: %s\n", e.GetField("label"))
	nexthopDetail(e.node)
}

func nexthopDetail(node xml.Node) {
	table := uitable.New()
	table.MaxColWidth = 80
	table.AddRow("    Type", "Interface", "Nexthop index")
	nhs, _ := node.Search("nh/NhSandeshData/type/text()")
	itf, _ := node.Search("nh/NhSandeshData/itf/text()")
	nhIdx, _ := node.Search("nh/NhSandeshData/nh_index/text()")
	table.AddRow("    "+Pretty(nhs), Pretty(itf), Pretty(nhIdx))
	fmt.Println(table)
}

func controllerRoutePath(e Element) {
	srcIp, _ := e.node.Search("prefix/text()")
	fmt.Printf("Prefix %s\n", srcIp[0])
	paths, _ := e.node.Search("paths/list/ShowRoutePath")

	table := uitable.New()
	table.MaxColWidth = 80
	table.AddRow("    Protocol", "Nexthop", "Local Pref", "Peers", "MPLS label")
	for _, path := range paths {
		protocol, _ := path.Search("protocol/text()")
		nhs, _ := path.Search("next_hop/text()")
		peers, _ := path.Search("source/text()")
		label, _ := path.Search("label/text()")
		localPref, _ := path.Search("local_preference/text()")
		table.AddRow("    "+Pretty(protocol), Pretty(nhs), Pretty(localPref), Pretty(peers), Pretty(label))
	}
	fmt.Println(table)
}

func main() {
	var count bool
	var hosts_file string

	app := cli.NewApp()
	app.Name = "contrail-introspect-cli"
	app.Usage = "CLI on ContraiL Introspects"
	app.Version = "0.0.4"
	app.EnableBashCompletion = true
	app.Before = func(c *cli.Context) error {
		if c.GlobalIsSet("hosts") {
			var err error
			hosts, err = LoadHostsFile(c.GlobalString("hosts"))
			return err
		}
		return nil
	}
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:        "hosts",
			Usage:       "host file to do DNS resolution",
			Destination: &hosts_file,
		}}
	app.Commands = []cli.Command{
		GenCommand(DescRoute(), "agent-route", "Show routes on agent"),
		GenCommand(DescItf(), "agent-itf", "Show interfaces on agent"),
		GenCommand(DescVrf(), "agent-vrf", "Show vrfs on agent "),
		GenCommand(DescPeering(), "agent-peering", "Peering with controller on agent"),
		GenCommand(DescVn(), "agent-vn", "Show virtual networks on agent"),
		GenCommand(DescMpls(), "agent-mpls", "Show mpls on agent"),
		Follow(),
		Path(),
		GenCommand(DescRiSummary(), "controller-ri", "Show routing instances on controller"),
		GenCommand(DescCtrlRoute(), "controller-route", "Show routes on controller"),
		GenCommand(DescCtrlRouteSummary(), "controller-route-summary", "Show routes summary on controller"),
		{
			Name:      "agent-multiple",
			Usage:     "List routes with multiple nexthops",
			ArgsUsage: "vrouter vrf_name",
			Flags: []cli.Flag{
				cli.BoolFlag{
					Name:        "count",
					Destination: &count,
				}},
			Action: func(c *cli.Context) error {
				if c.NArg() != 2 {
					log.Fatal("Wrong argument number!")
				}
				vrouter := c.Args()[0]
				vrf_name := c.Args()[1]
				multiple(vrouter, vrf_name, count)
				return nil
			},
		},
	}
	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package scipipe

import (
	"errors"
	"fmt"
	"os"
	"os/exec"
	re "regexp"
	str "strings"
)

type ShellProcess struct {
	process
	InPorts          map[string]chan *FileTarget
	OutPorts         map[string]chan *FileTarget
	OutPortsDoStream map[string]bool
	OutPathFuncs     map[string]func(*ShellTask) string
	ParamPorts       map[string]chan string
	Prepend          string
	CommandPattern   string
	Spawn            bool
}

func NewShellProcess(command string) *ShellProcess {
	return &ShellProcess{
		CommandPattern:   command,
		InPorts:          make(map[string]chan *FileTarget),
		OutPorts:         make(map[string]chan *FileTarget),
		OutPortsDoStream: make(map[string]bool),
		OutPathFuncs:     make(map[string]func(*ShellTask) string),
		ParamPorts:       make(map[string]chan string),
		Spawn:            true,
	}
}

func Shell(cmd string) *ShellProcess {
	if !LogExists {
		InitLogInfo()
	}
	p := NewShellProcess(cmd)
	p.initPortsFromCmdPattern(cmd, nil)
	return p
}

func Sh(cmd string) *ShellProcess {
	return Shell(cmd)
}

func (p *ShellProcess) Run() {
	defer p.closeOutPorts()

	tasks := []*ShellTask{}
	Debug.Printf("[%s] Starting to loop over tasks\n", p.CommandPattern)
	for t := range p.createTasks() {
		tasks = append(tasks, t)
		Debug.Println("Now processing task", t.Command, "...")

		t.createFifos()

		Debug.Println("[%s] Now sending fifos for task\n", t.Command, "...")
		// Sending FIFOs for the task
		for oname, otgt := range t.OutTargets {
			if otgt.doStream {
				Debug.Printf("[%s] Sending FIFO target on outport %s ...\n", p.CommandPattern, oname)
				p.OutPorts[oname] <- otgt
			}
		}

		Debug.Println("[%s] Now starting to run task\n", t.Command, "...")
		// Run the task
		go t.Run()
	}

	// Wait for finish, and send out targets in arrival order
	for _, t := range tasks {
		Debug.Printf("[%s] Waiting for Done from task: %s\n", p.CommandPattern, t.Command)
		<-t.Done
		Debug.Printf("[%s] Receiving Done from task: %s\n", p.CommandPattern, t.Command)
		for oname, otgt := range t.OutTargets {
			if !otgt.doStream {
				Debug.Printf("[%s] Sending target on outport %s ...\n", p.CommandPattern, oname)
				p.OutPorts[oname] <- otgt
			}
		}
	}
}

func (p *ShellProcess) initPortsFromCmdPattern(cmd string, params map[string]string) {
	// Find in/out port names and Params and set up in struct fields
	r := getPlaceHolderRegex()
	ms := r.FindAllStringSubmatch(cmd, -1)
	for _, m := range ms {
		if len(m) < 3 {
			Check(errors.New("Too few matches"))
		}
		typ := m[1]
		name := m[2]
		if typ == "o" || typ == "os" {
			p.OutPorts[name] = make(chan *FileTarget, BUFSIZE)
			if typ == "os" {
				p.OutPortsDoStream[name] = true
			}
		} else if typ == "i" {
			// Set up a channel on the inports, even though this is
			// often replaced by another processes output port channel.
			// It might be nice to have it init'ed with a channel
			// anyways, for use cases when we want to send FileTargets
			// on the inport manually.
			p.InPorts[name] = make(chan *FileTarget, BUFSIZE)
		} else if typ == "p" {
			if params == nil || params[name] == "" {
				p.ParamPorts[name] = make(chan string, BUFSIZE)
			}
		}
	}
}

func (p *ShellProcess) createTasks() (ch chan *ShellTask) {
	ch = make(chan *ShellTask)
	go func() {
		defer close(ch)
		for {
			inTargets, inPortsOpen := p.receiveInputs()
			Debug.Printf("[%s] Got inTargets: %s", p.CommandPattern, inTargets)
			params, paramPortsOpen := p.receiveParams()
			Debug.Printf("[%s] Got params: %s", p.CommandPattern, params)
			if !inPortsOpen && !paramPortsOpen {
				Debug.Printf("[%s] Breaking: Both inPorts and paramPorts closed", p.CommandPattern)
				break
			}
			if len(p.InPorts) == 0 && !paramPortsOpen {
				Debug.Printf("[%s] Breaking: No inports, and params closed", p.CommandPattern)
				break
			}
			if len(p.ParamPorts) == 0 && !inPortsOpen {
				Debug.Printf("[%s] Breaking: No params, and inPorts closed", p.CommandPattern)
				break
			}
			t := NewShellTask(p.CommandPattern, inTargets, p.OutPathFuncs, p.OutPortsDoStream, params, p.Prepend)
			ch <- t
			if len(p.InPorts) == 0 && len(p.ParamPorts) == 0 {
				Debug.Printf("[%s] Breaking: No inports nor params", p.CommandPattern)
				break
			}
		}
	}()
	return ch
}

func (p *ShellProcess) receiveInputs() (inTargets map[string]*FileTarget, inPortsOpen bool) {
	inPortsOpen = true
	inTargets = make(map[string]*FileTarget)
	// Read input targets on in-ports and set up path mappings
	for iname, ichan := range p.InPorts {
		Debug.Printf("[%s] Receieving on inPort %s ...", p.CommandPattern, iname)
		inTarget, open := <-ichan
		if !open {
			inPortsOpen = false
			continue
		}
		Debug.Printf("[%s] Got inTarget %s ...", p.CommandPattern, inTarget.GetPath())
		inTargets[iname] = inTarget
	}
	return
}

func (p *ShellProcess) receiveParams() (params map[string]string, paramPortsOpen bool) {
	paramPortsOpen = true
	params = make(map[string]string)
	// Read input targets on in-ports and set up path mappings
	for pname, pchan := range p.ParamPorts {
		pval, open := <-pchan
		if !open {
			paramPortsOpen = false
			continue
		}
		Debug.Println("Receiving param:", pname, "with value", pval)
		params[pname] = pval
	}
	return
}

func (p *ShellProcess) closeOutPorts() {
	for oname, oport := range p.OutPorts {
		Debug.Printf("[%s] Closing port %s ...\n", p.CommandPattern, oname)
		close(oport)
	}
}

func getPlaceHolderRegex() *re.Regexp {
	r, err := re.Compile("{(o|os|i|is|p):([^{}:]+)}")
	Check(err)
	return r
}

// ------- ShellTask -------

type ShellTask struct {
	InTargets  map[string]*FileTarget
	OutTargets map[string]*FileTarget
	Params     map[string]string
	Command    string
	Done       chan int
}

func NewShellTask(cmdPat string, inTargets map[string]*FileTarget, outPathFuncs map[string]func(*ShellTask) string, outPortsDoStream map[string]bool, params map[string]string, prepend string) *ShellTask {
	t := &ShellTask{
		InTargets:  inTargets,
		OutTargets: make(map[string]*FileTarget),
		Params:     params,
		Command:    "",
		Done:       make(chan int),
	}
	// Create out targets
	Debug.Printf("[%s] Creating outTargets now ...", cmdPat)
	outTargets := make(map[string]*FileTarget)
	for oname, ofun := range outPathFuncs {
		opath := ofun(t)
		otgt := NewFileTarget(opath)
		if outPortsDoStream[oname] {
			otgt.doStream = true
		}
		Debug.Printf("[%s] Creating outTarget with path %s ...", cmdPat, opath)
		outTargets[oname] = otgt
	}
	t.OutTargets = outTargets
	t.Command = formatCommand(cmdPat, inTargets, outTargets, params, prepend)
	Debug.Printf("[%s] Created formatted command: %s", cmdPat, t.Command)
	return t
}

func (t *ShellTask) Run() {
	defer close(t.Done) // TODO: Is this needed?
	if !t.anyOutputExists() && !t.fifosInOutTargetsMissing() {
		t.executeCommand(t.Command)
		t.atomizeTargets()
		t.cleanUpFifos()
	}
	Debug.Printf("[%s] Starting to send Done in t.Run() ...)\n", t.Command)
	t.Done <- 1
	Debug.Printf("[%s] Done sending Done, in t.Run()\n", t.Command)
}

func (t *ShellTask) executeCommand(cmd string) {
	Info.Printf("[%s] Executing command: %s \n", t.Command, cmd)
	_, err := exec.Command("bash", "-c", cmd).Output()
	Check(err)
}

func (t *ShellTask) GetInPath(inPort string) string {
	return t.InTargets[inPort].GetPath()
}

func formatCommand(cmd string, inTargets map[string]*FileTarget, outTargets map[string]*FileTarget, params map[string]string, prepend string) string {

	// Debug.Println("Formatting command with the following data:")
	// Debug.Println("prepend:", prepend)
	// Debug.Println("cmd:", cmd)
	// Debug.Println("inTargets:", inTargets)
	// Debug.Println("outTargets:", outTargets)
	// Debug.Println("params:", params)

	r := getPlaceHolderRegex()
	ms := r.FindAllStringSubmatch(cmd, -1)
	for _, m := range ms {
		whole := m[0]
		typ := m[1]
		name := m[2]
		var newstr string
		if typ == "o" || typ == "os" {
			if outTargets[name] == nil {
				msg := fmt.Sprint("Missing outpath for outport '", name, "' for command '", cmd, "'")
				Check(errors.New(msg))
			} else {
				if typ == "o" {
					newstr = outTargets[name].GetTempPath() // Means important to Atomize afterwards!
				} else if typ == "os" {
					newstr = outTargets[name].GetFifoPath()
				}
			}
		} else if typ == "i" {
			if inTargets[name] == nil {
				msg := fmt.Sprint("Missing intarget for inport '", name, "' for command '", cmd, "'")
				Check(errors.New(msg))
			} else if inTargets[name].GetPath() == "" {
				msg := fmt.Sprint("Missing inpath for inport '", name, "' for command '", cmd, "'")
				Check(errors.New(msg))
			} else {
				if typ == "i" {
					if inTargets[name].doStream {
						newstr = inTargets[name].GetFifoPath()
					} else {
						newstr = inTargets[name].GetPath()
					}
				}
			}
		} else if typ == "p" {
			if params[name] == "" {
				msg := fmt.Sprint("Missing param value param '", name, "' for command '", cmd, "'")
				Check(errors.New(msg))
			} else {
				newstr = params[name]
			}
		}
		if newstr == "" {
			msg := fmt.Sprint("Replace failed for port ", name, " forcommand '", cmd, "'")
			Check(errors.New(msg))
		}
		cmd = str.Replace(cmd, whole, newstr, -1)
	}
	// Add prepend string to the command
	if prepend != "" {
		cmd = fmt.Sprintf("%s %s", prepend, cmd)
	}
	return cmd
}

func (t *ShellTask) atomizeTargets() {
	for _, tgt := range t.OutTargets {
		if !tgt.doStream {
			Debug.Printf("Atomizing file: %s -> %s", tgt.GetTempPath(), tgt.GetPath())
			tgt.Atomize()
			Debug.Printf("Done atomizing file: %s -> %s", tgt.GetTempPath(), tgt.GetPath())
		} else {
			Debug.Printf("Target is streaming, so not atomizing: %s", tgt.GetPath())
		}
	}
}

func (t *ShellTask) cleanUpFifos() {
	for _, tgt := range t.OutTargets {
		if tgt.doStream {
			Debug.Printf("[%s] Cleaning up FIFO for input target: %s\n", t.Command, tgt.GetFifoPath())
			tgt.RemoveFifo()
		} else {
			Debug.Printf("[%s] input target is not FIFO, so not removing any FIFO: %s\n", t.Command, tgt.GetPath())
		}
	}
}

func (t *ShellTask) anyOutputExists() (anyFileExists bool) {
	anyFileExists = false
	for _, tgt := range t.OutTargets {
		opath := tgt.GetPath()
		otmpPath := tgt.GetTempPath()
		if !tgt.doStream {
			if _, err := os.Stat(opath); err == nil {
				Warn.Printf("[%s] Output file already exists: %s. Check your workflow for correctness!\n", t.Command, opath)
				anyFileExists = true
			}
			if _, err := os.Stat(otmpPath); err == nil {
				Warn.Printf("[%s] Temporary Output file already exists: %s. Check your workflow for correctness!\n", t.Command, otmpPath)
				anyFileExists = true
			}
		}
	}
	return
}

func (t *ShellTask) fifosInOutTargetsMissing() (fifosInOutTargetsMissing bool) {
	fifosInOutTargetsMissing = false
	for _, tgt := range t.OutTargets {
		if tgt.doStream {
			ofifoPath := tgt.GetFifoPath()
			if _, err := os.Stat(ofifoPath); err != nil {
				Warn.Printf("[%s] FIFO Output file missing, for streaming output: %s. Check your workflow for correctness!\n", t.Command, ofifoPath)
				fifosInOutTargetsMissing = true
			}
		}
	}
	return
}

func (t *ShellTask) createFifos() {
	Debug.Printf("[%s] Now creating fifos for task\n", t.Command)
	for _, otgt := range t.OutTargets {
		if otgt.doStream {
			otgt.CreateFifo()
		}
	}
}
<commit_msg>Add back convenience methods for creating path formatters<commit_after>package scipipe

import (
	"errors"
	"fmt"
	"os"
	"os/exec"
	re "regexp"
	str "strings"
)

type ShellProcess struct {
	process
	InPorts          map[string]chan *FileTarget
	OutPorts         map[string]chan *FileTarget
	OutPortsDoStream map[string]bool
	OutPathFuncs     map[string]func(*ShellTask) string
	ParamPorts       map[string]chan string
	Prepend          string
	CommandPattern   string
	Spawn            bool
}

func NewShellProcess(command string) *ShellProcess {
	return &ShellProcess{
		CommandPattern:   command,
		InPorts:          make(map[string]chan *FileTarget),
		OutPorts:         make(map[string]chan *FileTarget),
		OutPortsDoStream: make(map[string]bool),
		OutPathFuncs:     make(map[string]func(*ShellTask) string),
		ParamPorts:       make(map[string]chan string),
		Spawn:            true,
	}
}

func Shell(cmd string) *ShellProcess {
	if !LogExists {
		InitLogInfo()
	}
	p := NewShellProcess(cmd)
	p.initPortsFromCmdPattern(cmd, nil)
	return p
}

func Sh(cmd string) *ShellProcess {
	return Shell(cmd)
}

func (p *ShellProcess) Run() {
	defer p.closeOutPorts()

	tasks := []*ShellTask{}
	Debug.Printf("[%s] Starting to loop over tasks\n", p.CommandPattern)
	for t := range p.createTasks() {
		tasks = append(tasks, t)
		Debug.Println("Now processing task", t.Command, "...")

		t.createFifos()

		Debug.Println("[%s] Now sending fifos for task\n", t.Command, "...")
		// Sending FIFOs for the task
		for oname, otgt := range t.OutTargets {
			if otgt.doStream {
				Debug.Printf("[%s] Sending FIFO target on outport %s ...\n", p.CommandPattern, oname)
				p.OutPorts[oname] <- otgt
			}
		}

		Debug.Println("[%s] Now starting to run task\n", t.Command, "...")
		// Run the task
		go t.Run()
	}

	// Wait for finish, and send out targets in arrival order
	for _, t := range tasks {
		Debug.Printf("[%s] Waiting for Done from task: %s\n", p.CommandPattern, t.Command)
		<-t.Done
		Debug.Printf("[%s] Receiving Done from task: %s\n", p.CommandPattern, t.Command)
		for oname, otgt := range t.OutTargets {
			if !otgt.doStream {
				Debug.Printf("[%s] Sending target on outport %s ...\n", p.CommandPattern, oname)
				p.OutPorts[oname] <- otgt
			}
		}
	}
}

func (p *ShellProcess) initPortsFromCmdPattern(cmd string, params map[string]string) {
	// Find in/out port names and Params and set up in struct fields
	r := getPlaceHolderRegex()
	ms := r.FindAllStringSubmatch(cmd, -1)
	for _, m := range ms {
		if len(m) < 3 {
			Check(errors.New("Too few matches"))
		}
		typ := m[1]
		name := m[2]
		if typ == "o" || typ == "os" {
			p.OutPorts[name] = make(chan *FileTarget, BUFSIZE)
			if typ == "os" {
				p.OutPortsDoStream[name] = true
			}
		} else if typ == "i" {
			// Set up a channel on the inports, even though this is
			// often replaced by another processes output port channel.
			// It might be nice to have it init'ed with a channel
			// anyways, for use cases when we want to send FileTargets
			// on the inport manually.
			p.InPorts[name] = make(chan *FileTarget, BUFSIZE)
		} else if typ == "p" {
			if params == nil || params[name] == "" {
				p.ParamPorts[name] = make(chan string, BUFSIZE)
			}
		}
	}
}

func (p *ShellProcess) createTasks() (ch chan *ShellTask) {
	ch = make(chan *ShellTask)
	go func() {
		defer close(ch)
		for {
			inTargets, inPortsOpen := p.receiveInputs()
			Debug.Printf("[%s] Got inTargets: %s", p.CommandPattern, inTargets)
			params, paramPortsOpen := p.receiveParams()
			Debug.Printf("[%s] Got params: %s", p.CommandPattern, params)
			if !inPortsOpen && !paramPortsOpen {
				Debug.Printf("[%s] Breaking: Both inPorts and paramPorts closed", p.CommandPattern)
				break
			}
			if len(p.InPorts) == 0 && !paramPortsOpen {
				Debug.Printf("[%s] Breaking: No inports, and params closed", p.CommandPattern)
				break
			}
			if len(p.ParamPorts) == 0 && !inPortsOpen {
				Debug.Printf("[%s] Breaking: No params, and inPorts closed", p.CommandPattern)
				break
			}
			t := NewShellTask(p.CommandPattern, inTargets, p.OutPathFuncs, p.OutPortsDoStream, params, p.Prepend)
			ch <- t
			if len(p.InPorts) == 0 && len(p.ParamPorts) == 0 {
				Debug.Printf("[%s] Breaking: No inports nor params", p.CommandPattern)
				break
			}
		}
	}()
	return ch
}

func (p *ShellProcess) receiveInputs() (inTargets map[string]*FileTarget, inPortsOpen bool) {
	inPortsOpen = true
	inTargets = make(map[string]*FileTarget)
	// Read input targets on in-ports and set up path mappings
	for iname, ichan := range p.InPorts {
		Debug.Printf("[%s] Receieving on inPort %s ...", p.CommandPattern, iname)
		inTarget, open := <-ichan
		if !open {
			inPortsOpen = false
			continue
		}
		Debug.Printf("[%s] Got inTarget %s ...", p.CommandPattern, inTarget.GetPath())
		inTargets[iname] = inTarget
	}
	return
}

func (p *ShellProcess) receiveParams() (params map[string]string, paramPortsOpen bool) {
	paramPortsOpen = true
	params = make(map[string]string)
	// Read input targets on in-ports and set up path mappings
	for pname, pchan := range p.ParamPorts {
		pval, open := <-pchan
		if !open {
			paramPortsOpen = false
			continue
		}
		Debug.Println("Receiving param:", pname, "with value", pval)
		params[pname] = pval
	}
	return
}

func (p *ShellProcess) closeOutPorts() {
	for oname, oport := range p.OutPorts {
		Debug.Printf("[%s] Closing port %s ...\n", p.CommandPattern, oname)
		close(oport)
	}
}

// Convenience method to create an (output) path formatter returning a static string
func (p *ShellProcess) OutPathGenString(outPort string, path string) {
	p.OutPathFuncs[outPort] = func(t *ShellTask) string {
		return path
	}
}

// Convenience method to create an (output) path formatter that extends the path of
// and input FileTarget
func (p *ShellProcess) OutPathGenExtend(outPort string, inPort string, extension string) {
	p.OutPathFuncs[outPort] = func(t *ShellTask) string {
		return t.InTargets[inPort].GetPath() + extension
	}
}

// Convenience method to create an (output) path formatter that uses an input's path
// but replaces parts of it.
func (p *ShellProcess) OutPathGenReplace(outPort string, inPort string, old string, new string) {
	p.OutPathFuncs[outPort] = func(t *ShellTask) string {
		return str.Replace(t.InTargets[inPort].GetPath(), old, new, -1)
	}
}

func getPlaceHolderRegex() *re.Regexp {
	r, err := re.Compile("{(o|os|i|is|p):([^{}:]+)}")
	Check(err)
	return r
}

// ------- ShellTask -------

type ShellTask struct {
	InTargets  map[string]*FileTarget
	OutTargets map[string]*FileTarget
	Params     map[string]string
	Command    string
	Done       chan int
}

func NewShellTask(cmdPat string, inTargets map[string]*FileTarget, outPathFuncs map[string]func(*ShellTask) string, outPortsDoStream map[string]bool, params map[string]string, prepend string) *ShellTask {
	t := &ShellTask{
		InTargets:  inTargets,
		OutTargets: make(map[string]*FileTarget),
		Params:     params,
		Command:    "",
		Done:       make(chan int),
	}
	// Create out targets
	Debug.Printf("[%s] Creating outTargets now ...", cmdPat)
	outTargets := make(map[string]*FileTarget)
	for oname, ofun := range outPathFuncs {
		opath := ofun(t)
		otgt := NewFileTarget(opath)
		if outPortsDoStream[oname] {
			otgt.doStream = true
		}
		Debug.Printf("[%s] Creating outTarget with path %s ...", cmdPat, opath)
		outTargets[oname] = otgt
	}
	t.OutTargets = outTargets
	t.Command = formatCommand(cmdPat, inTargets, outTargets, params, prepend)
	Debug.Printf("[%s] Created formatted command: %s", cmdPat, t.Command)
	return t
}

func (t *ShellTask) Run() {
	defer close(t.Done) // TODO: Is this needed?
	if !t.anyOutputExists() && !t.fifosInOutTargetsMissing() {
		t.executeCommand(t.Command)
		t.atomizeTargets()
		t.cleanUpFifos()
	}
	Debug.Printf("[%s] Starting to send Done in t.Run() ...)\n", t.Command)
	t.Done <- 1
	Debug.Printf("[%s] Done sending Done, in t.Run()\n", t.Command)
}

func (t *ShellTask) executeCommand(cmd string) {
	Info.Printf("[%s] Executing command: %s \n", t.Command, cmd)
	_, err := exec.Command("bash", "-c", cmd).Output()
	Check(err)
}

func (t *ShellTask) GetInPath(inPort string) string {
	return t.InTargets[inPort].GetPath()
}

func formatCommand(cmd string, inTargets map[string]*FileTarget, outTargets map[string]*FileTarget, params map[string]string, prepend string) string {

	// Debug.Println("Formatting command with the following data:")
	// Debug.Println("prepend:", prepend)
	// Debug.Println("cmd:", cmd)
	// Debug.Println("inTargets:", inTargets)
	// Debug.Println("outTargets:", outTargets)
	// Debug.Println("params:", params)

	r := getPlaceHolderRegex()
	ms := r.FindAllStringSubmatch(cmd, -1)
	for _, m := range ms {
		whole := m[0]
		typ := m[1]
		name := m[2]
		var newstr string
		if typ == "o" || typ == "os" {
			if outTargets[name] == nil {
				msg := fmt.Sprint("Missing outpath for outport '", name, "' for command '", cmd, "'")
				Check(errors.New(msg))
			} else {
				if typ == "o" {
					newstr = outTargets[name].GetTempPath() // Means important to Atomize afterwards!
				} else if typ == "os" {
					newstr = outTargets[name].GetFifoPath()
				}
			}
		} else if typ == "i" {
			if inTargets[name] == nil {
				msg := fmt.Sprint("Missing intarget for inport '", name, "' for command '", cmd, "'")
				Check(errors.New(msg))
			} else if inTargets[name].GetPath() == "" {
				msg := fmt.Sprint("Missing inpath for inport '", name, "' for command '", cmd, "'")
				Check(errors.New(msg))
			} else {
				if typ == "i" {
					if inTargets[name].doStream {
						newstr = inTargets[name].GetFifoPath()
					} else {
						newstr = inTargets[name].GetPath()
					}
				}
			}
		} else if typ == "p" {
			if params[name] == "" {
				msg := fmt.Sprint("Missing param value param '", name, "' for command '", cmd, "'")
				Check(errors.New(msg))
			} else {
				newstr = params[name]
			}
		}
		if newstr == "" {
			msg := fmt.Sprint("Replace failed for port ", name, " forcommand '", cmd, "'")
			Check(errors.New(msg))
		}
		cmd = str.Replace(cmd, whole, newstr, -1)
	}
	// Add prepend string to the command
	if prepend != "" {
		cmd = fmt.Sprintf("%s %s", prepend, cmd)
	}
	return cmd
}

func (t *ShellTask) atomizeTargets() {
	for _, tgt := range t.OutTargets {
		if !tgt.doStream {
			Debug.Printf("Atomizing file: %s -> %s", tgt.GetTempPath(), tgt.GetPath())
			tgt.Atomize()
			Debug.Printf("Done atomizing file: %s -> %s", tgt.GetTempPath(), tgt.GetPath())
		} else {
			Debug.Printf("Target is streaming, so not atomizing: %s", tgt.GetPath())
		}
	}
}

func (t *ShellTask) cleanUpFifos() {
	for _, tgt := range t.OutTargets {
		if tgt.doStream {
			Debug.Printf("[%s] Cleaning up FIFO for input target: %s\n", t.Command, tgt.GetFifoPath())
			tgt.RemoveFifo()
		} else {
			Debug.Printf("[%s] input target is not FIFO, so not removing any FIFO: %s\n", t.Command, tgt.GetPath())
		}
	}
}

func (t *ShellTask) anyOutputExists() (anyFileExists bool) {
	anyFileExists = false
	for _, tgt := range t.OutTargets {
		opath := tgt.GetPath()
		otmpPath := tgt.GetTempPath()
		if !tgt.doStream {
			if _, err := os.Stat(opath); err == nil {
				Warn.Printf("[%s] Output file already exists: %s. Check your workflow for correctness!\n", t.Command, opath)
				anyFileExists = true
			}
			if _, err := os.Stat(otmpPath); err == nil {
				Warn.Printf("[%s] Temporary Output file already exists: %s. Check your workflow for correctness!\n", t.Command, otmpPath)
				anyFileExists = true
			}
		}
	}
	return
}

func (t *ShellTask) fifosInOutTargetsMissing() (fifosInOutTargetsMissing bool) {
	fifosInOutTargetsMissing = false
	for _, tgt := range t.OutTargets {
		if tgt.doStream {
			ofifoPath := tgt.GetFifoPath()
			if _, err := os.Stat(ofifoPath); err != nil {
				Warn.Printf("[%s] FIFO Output file missing, for streaming output: %s. Check your workflow for correctness!\n", t.Command, ofifoPath)
				fifosInOutTargetsMissing = true
			}
		}
	}
	return
}

func (t *ShellTask) createFifos() {
	Debug.Printf("[%s] Now creating fifos for task\n", t.Command)
	for _, otgt := range t.OutTargets {
		if otgt.doStream {
			otgt.CreateFifo()
		}
	}
}
<|endoftext|>


<commit_before>package rel

// A Slice is a concrete implementation of a List interface
// that still affords standard slice mechanics.
type Slice []*Pair

// Implement a List and sort interface.
func (s Slice) Get(i int) interface{}                 { return s[i].Element }
func (s Slice) Set(i int, e interface{}, val float64) { s[i] = &Pair{e, val} }
func (s Slice) Val(i int) float64                     { return s[i].Value }
func (s Slice) Len() int                              { return len(s) }
func (s Slice) Less(i, j int) bool                    { return s.Val(i) < s.Val(j) }
func (s Slice) Swap(i, j int)                         { s[i], s[j] = s[j], s[i] }

func MakeSlice(l List) Slice {
	// Create a new Slice struct that represents the input List.
	s := make(Slice, l.Len())

	// Populate the list struct with the input List's data.
	for i := range s {
		s[i] = &Pair{
			Element: l.Get(i),
			Value:   l.Val(i),
		}
	}

	return s
}
<commit_msg>MakeSlice accepts interfaces.<commit_after>package rel

// A Slice is a concrete implementation of a List interface
// that still affords standard slice mechanics.
type Slice []*Pair

// Implement a List and sort interface.
func (s Slice) Get(i int) interface{}                 { return s[i].Element }
func (s Slice) Set(i int, e interface{}, val float64) { s[i] = NewPair(e, val) }
func (s Slice) Val(i int) float64                     { return s[i].Value }
func (s Slice) Len() int                              { return len(s) }
func (s Slice) Less(i, j int) bool                    { return s.Val(i) < s.Val(j) }
func (s Slice) Swap(i, j int)                         { s[i], s[j] = s[j], s[i] }

// MakeSlice will produce a Slice instance either from an input List
// or []float64 slice.  If the input is type []float64, then the
// underlying Element is an empty struct.
func MakeSlice(l interface{}) Slice {

	var s Slice

	switch t := l.(type) {
	case List:
		s = makeSliceFromList(t)
	case []float64:
		s = makeSliceFromFloat64s(t)
	}

	return s
}

func makeSliceFromList(l List) Slice {
	// Create a new Slice struct that represents the input List.
	s := make(Slice, l.Len())

	// Populate the list struct with the input List's data.
	for i := range s {
		s.Set(i, l.Get(i), l.Val(i))
	}

	return s
}

func makeSliceFromFloat64s(l []float64) Slice {
	s := make(Slice, len(l))
	for i := range l {
		s.Set(i, struct{}{}, l[i])
	}

	return s
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/fraenkel/candiedyaml"
	"github.com/gorilla/websocket"
	"github.com/mgutz/ansi"
	"github.com/pivotal-golang/archiver/compressor"
)

type Build struct {
	Guid   string `json:"guid,omitempty"`
	Image  string `json:"image"`
	Path   string `json:"path"`
	Script string `json:"script"`
}

type BuildResult struct {
	Status string `json:"status"`
}

type BuildConfig struct {
	Image  string `yaml:"image"`
	Path   string `yaml:"path"`
	Script string `yaml:"script"`
}

var buildConfig = flag.String(
	"c",
	"build.yml",
	"build configuration file",
)

var buildDir = flag.String(
	"d",
	".",
	"source directory to build",
)

var redgreenAddr = flag.String(
	"redgreenAddr",
	"127.0.0.1:5637",
	"address denoting the redgreen service",
)

func main() {
	flag.Parse()

	build := create(loadConfig())

	buildLog := fmt.Sprintf("ws://%s/builds/%s/log/output", *redgreenAddr, build.Guid)

	conn, res, err := websocket.DefaultDialer.Dial(buildLog, nil)
	if err != nil {
		log.Println("failed to stream output:", err, res)
		return
	}

	streaming := new(sync.WaitGroup)
	streaming.Add(1)

	go stream(conn, streaming)

	upload(build)

	exitCode := poll(build)

	res.Body.Close()
	conn.Close()

	streaming.Wait()

	os.Exit(exitCode)
}

func loadConfig() BuildConfig {
	configFile, err := os.Open(*buildConfig)
	if err != nil {
		log.Fatalln("could not open config file:", err)
	}

	var config BuildConfig

	err = candiedyaml.NewDecoder(configFile).Decode(&config)
	if err != nil {
		log.Fatalln("could not parse config file:", err)
	}

	return config
}

func create(config BuildConfig) Build {
	buffer := &bytes.Buffer{}

	build := Build{
		Image:  config.Image,
		Path:   config.Path,
		Script: config.Script,
	}

	if build.Path == "" {
		build.Path = "."
	}

	err := json.NewEncoder(buffer).Encode(build)
	if err != nil {
		log.Fatalln("encoding build failed:", err)
	}

	response, err := http.Post(
		"http://"+*redgreenAddr+"/builds",
		"application/json",
		buffer,
	)
	if err != nil {
		log.Fatalln("request failed:", err)
	}

	defer response.Body.Close()

	if response.StatusCode != http.StatusCreated {
		log.Println("bad response when creating build:", response)
		response.Write(os.Stderr)
		os.Exit(1)
	}

	err = json.NewDecoder(response.Body).Decode(&build)
	if err != nil {
		log.Fatalln("response decoding failed:", err)
	}

	return build
}

func stream(conn *websocket.Conn, streaming *sync.WaitGroup) {
	defer streaming.Done()

	for {
		_, data, err := conn.ReadMessage()
		if err != nil {
			break
		}

		fmt.Print(string(data))
	}
}

func upload(build Build) {
	src, err := filepath.Abs(*buildDir)
	if err != nil {
		log.Fatalln("could not locate build config:", err)
	}

	compressor := compressor.NewTgz()

	tmpfile, err := ioutil.TempFile("", "smith")
	if err != nil {
		log.Fatalln("creating tempfile failed:", err)
	}

	tmpfile.Close()

	defer os.Remove(tmpfile.Name())

	err = compressor.Compress(src, tmpfile.Name())
	if err != nil {
		log.Fatalln("creating archive failed:", err)
	}

	archive, err := os.Open(tmpfile.Name())
	if err != nil {
		log.Fatalln("could not open archive:", err)
	}

	response, err := http.Post(
		"http://"+*redgreenAddr+"/builds/"+build.Guid+"/bits",
		"application/octet-stream",
		archive,
	)
	if err != nil {
		log.Fatalln("request failed:", err)
	}

	defer response.Body.Close()

	if response.StatusCode != http.StatusCreated {
		log.Println("bad response when uploading bits:", response)
		response.Write(os.Stderr)
		os.Exit(1)
	}
}

func poll(build Build) int {
	for {
		var result BuildResult

		response, err := http.Get("http://" + *redgreenAddr + "/builds/" + build.Guid + "/result")
		if err != nil {
			log.Fatalln("error polling for result:", err)
		}

		err = json.NewDecoder(response.Body).Decode(&result)
		if err != nil {
			response.Body.Close()
			log.Fatalln("error decoding result:", err)
		}

		response.Body.Close()

		var color string
		var exitCode int

		switch result.Status {
		case "succeeded":
			color = "green"
			exitCode = 0
		case "failed":
			color = "red"
			exitCode = 1
		case "errored":
			color = "magenta"
			exitCode = 2
		default:
			time.Sleep(time.Second)
			continue
		}

		fmt.Println(ansi.Color(result.Status, color))
		return exitCode
	}
}
<commit_msg>report upload progress<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/cheggaaa/pb"
	"github.com/fraenkel/candiedyaml"
	"github.com/gorilla/websocket"
	"github.com/mgutz/ansi"
	"github.com/pivotal-golang/archiver/compressor"
)

type Build struct {
	Guid   string `json:"guid,omitempty"`
	Image  string `json:"image"`
	Path   string `json:"path"`
	Script string `json:"script"`
}

type BuildResult struct {
	Status string `json:"status"`
}

type BuildConfig struct {
	Image  string `yaml:"image"`
	Path   string `yaml:"path"`
	Script string `yaml:"script"`
}

var buildConfig = flag.String(
	"c",
	"build.yml",
	"build configuration file",
)

var buildDir = flag.String(
	"d",
	".",
	"source directory to build",
)

var redgreenAddr = flag.String(
	"redgreenAddr",
	"127.0.0.1:5637",
	"address denoting the redgreen service",
)

func main() {
	flag.Parse()

	build := create(loadConfig())

	buildLog := fmt.Sprintf("ws://%s/builds/%s/log/output", *redgreenAddr, build.Guid)

	conn, res, err := websocket.DefaultDialer.Dial(buildLog, nil)
	if err != nil {
		log.Println("failed to stream output:", err, res)
		return
	}

	streaming := new(sync.WaitGroup)
	streaming.Add(1)

	go stream(conn, streaming)

	upload(build)

	exitCode := poll(build)

	res.Body.Close()
	conn.Close()

	streaming.Wait()

	os.Exit(exitCode)
}

func loadConfig() BuildConfig {
	configFile, err := os.Open(*buildConfig)
	if err != nil {
		log.Fatalln("could not open config file:", err)
	}

	var config BuildConfig

	err = candiedyaml.NewDecoder(configFile).Decode(&config)
	if err != nil {
		log.Fatalln("could not parse config file:", err)
	}

	return config
}

func create(config BuildConfig) Build {
	buffer := &bytes.Buffer{}

	build := Build{
		Image:  config.Image,
		Path:   config.Path,
		Script: config.Script,
	}

	if build.Path == "" {
		build.Path = "."
	}

	err := json.NewEncoder(buffer).Encode(build)
	if err != nil {
		log.Fatalln("encoding build failed:", err)
	}

	response, err := http.Post(
		"http://"+*redgreenAddr+"/builds",
		"application/json",
		buffer,
	)
	if err != nil {
		log.Fatalln("request failed:", err)
	}

	defer response.Body.Close()

	if response.StatusCode != http.StatusCreated {
		log.Println("bad response when creating build:", response)
		response.Write(os.Stderr)
		os.Exit(1)
	}

	err = json.NewDecoder(response.Body).Decode(&build)
	if err != nil {
		log.Fatalln("response decoding failed:", err)
	}

	return build
}

func stream(conn *websocket.Conn, streaming *sync.WaitGroup) {
	defer streaming.Done()

	for {
		_, data, err := conn.ReadMessage()
		if err != nil {
			break
		}

		fmt.Print(string(data))
	}
}

func upload(build Build) {
	src, err := filepath.Abs(*buildDir)
	if err != nil {
		log.Fatalln("could not locate build config:", err)
	}

	compressor := compressor.NewTgz()

	tmpfile, err := ioutil.TempFile("", "smith")
	if err != nil {
		log.Fatalln("creating tempfile failed:", err)
	}

	tmpfile.Close()

	defer os.Remove(tmpfile.Name())

	err = compressor.Compress(src, tmpfile.Name())
	if err != nil {
		log.Fatalln("creating archive failed:", err)
	}

	archive, err := os.Open(tmpfile.Name())
	if err != nil {
		log.Fatalln("could not open archive:", err)
	}

	info, err := archive.Stat()
	if err != nil {
		log.Fatalln("could not stat archive:", err)
	}

	progress := pb.New64(info.Size())
	progress.SetUnits(pb.U_BYTES)

	progress.Start()
	defer progress.Finish()

	response, err := http.Post(
		"http://"+*redgreenAddr+"/builds/"+build.Guid+"/bits",
		"application/octet-stream",
		progress.NewProxyReader(archive),
	)
	if err != nil {
		log.Fatalln("request failed:", err)
	}

	defer response.Body.Close()

	if response.StatusCode != http.StatusCreated {
		log.Println("bad response when uploading bits:", response)
		response.Write(os.Stderr)
		os.Exit(1)
	}
}

func poll(build Build) int {
	for {
		var result BuildResult

		response, err := http.Get("http://" + *redgreenAddr + "/builds/" + build.Guid + "/result")
		if err != nil {
			log.Fatalln("error polling for result:", err)
		}

		err = json.NewDecoder(response.Body).Decode(&result)
		if err != nil {
			response.Body.Close()
			log.Fatalln("error decoding result:", err)
		}

		response.Body.Close()

		var color string
		var exitCode int

		switch result.Status {
		case "succeeded":
			color = "green"
			exitCode = 0
		case "failed":
			color = "red"
			exitCode = 1
		case "errored":
			color = "magenta"
			exitCode = 2
		default:
			time.Sleep(time.Second)
			continue
		}

		fmt.Println(ansi.Color(result.Status, color))
		return exitCode
	}
}
<|endoftext|>


<commit_before>package stats

import (
	"fmt"
	"net/http"
	"os"
	"sync"
	"time"
)

type StatsMiddleware struct {
	Lock                sync.RWMutex
	Start               time.Time
	Pid                 int
	ResponseCounts      map[string]int
	TotalResponseCounts map[string]int
	TotalResponseTime   time.Time
}

func New() *StatsMiddleware {
	stats := &StatsMiddleware{
		Start:               time.Now(),
		Pid:                 os.Getpid(),
		ResponseCounts:      map[string]int{},
		TotalResponseCounts: map[string]int{},
		TotalResponseTime:   time.Time{},
	}

	ticker := time.NewTicker(time.Second)
	go func() {
		for _ = range ticker.C {
			stats.Lock.Lock()

			stats.ResponseCounts = map[string]int{}

			defer stats.Lock.Unlock()
		}
	}()

	return stats
}

type recorderResponseWriter struct {
	http.ResponseWriter
	StatusCode int
}

func (w *recorderResponseWriter) WriteHeader(code int) {
	w.ResponseWriter.WriteHeader(code)
	w.StatusCode = code
}

// MiddlewareFunc makes StatsMiddleware implement the Middleware interface.
func (mw *StatsMiddleware) Handler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()

		writer := &recorderResponseWriter{w, 0}

		h.ServeHTTP(writer, r)

		mw.handleWriter(start, writer)
	})
}

// Negroni compatible interface
func (mw *StatsMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	start := time.Now()

	writer := &recorderResponseWriter{w, 0}

	next(writer, r)

	mw.handleWriter(start, writer)
}

func (mw *StatsMiddleware) handleWriter(start time.Time, writer *recorderResponseWriter) {
	end := time.Now()

	responseTime := end.Sub(start)

	mw.Lock.Lock()

	defer mw.Lock.Unlock()

	statusCode := fmt.Sprintf("%d", writer.StatusCode)

	mw.ResponseCounts[statusCode]++
	mw.TotalResponseCounts[statusCode]++
	mw.TotalResponseTime = mw.TotalResponseTime.Add(responseTime)
}

type Stats struct {
	Pid                    int            `json: "pid"`
	UpTime                 string         `json: "uptime"`
	UpTimeSec              float64        `json: "uptime_sec"`
	Time                   string         `json: "time"`
	TimeUnix               int64          `json: "unixtime"`
	StatusCodeCount        map[string]int `json: "status_code_count"`
	TotalStatusCodeCount   map[string]int `json: "total_status_code_count"`
	Count                  int            `json: "count"`
	TotalCount             int            `json: "total_count"`
	TotalResponseTime      string         `json" "total_response_time`
	TotalResponseTimeSec   float64        `json: "total_response_time_sec"`
	AverageResponseTime    string         `json: "average_response_time"`
	AverageResponseTimeSec float64        `json: "average_response_time_sec"`
}

func (mw *StatsMiddleware) GetStats() *Stats {

	mw.Lock.RLock()

	now := time.Now()

	uptime := now.Sub(mw.Start)

	count := 0
	for _, current := range mw.ResponseCounts {
		count += current
	}

	totalCount := 0
	for _, count := range mw.TotalResponseCounts {
		totalCount += count
	}

	totalResponseTime := mw.TotalResponseTime.Sub(time.Time{})

	averageResponseTime := time.Duration(0)
	if totalCount > 0 {
		avgNs := int64(totalResponseTime) / int64(totalCount)
		averageResponseTime = time.Duration(avgNs)
	}

	stats := &Stats{
		Pid:                    mw.Pid,
		UpTime:                 uptime.String(),
		UpTimeSec:              uptime.Seconds(),
		Time:                   now.String(),
		TimeUnix:               now.Unix(),
		StatusCodeCount:        mw.ResponseCounts,
		TotalStatusCodeCount:   mw.TotalResponseCounts,
		Count:                  count,
		TotalCount:             totalCount,
		TotalResponseTime:      totalResponseTime.String(),
		TotalResponseTimeSec:   totalResponseTime.Seconds(),
		AverageResponseTime:    averageResponseTime.String(),
		AverageResponseTimeSec: averageResponseTime.Seconds(),
	}

	mw.Lock.RUnlock()

	return stats
}
<commit_msg>Add ResetResponseCounts to StatsMiddleware<commit_after>package stats

import (
	"fmt"
	"net/http"
	"os"
	"sync"
	"time"
)

type StatsMiddleware struct {
	Lock                sync.RWMutex
	Start               time.Time
	Pid                 int
	ResponseCounts      map[string]int
	TotalResponseCounts map[string]int
	TotalResponseTime   time.Time
}

func New() *StatsMiddleware {
	stats := &StatsMiddleware{
		Start:               time.Now(),
		Pid:                 os.Getpid(),
		ResponseCounts:      map[string]int{},
		TotalResponseCounts: map[string]int{},
		TotalResponseTime:   time.Time{},
	}

	go func() {
		for {
			stats.ResetResponseCounts()

			time.Sleep(time.Second * 1)
		}
	}()

	return stats
}

func (mw *StatsMiddleware) ResetResponseCounts() {
	mw.Lock.Lock()
	defer mw.Lock.Unlock()
	mw.ResponseCounts = map[string]int{}
}

type recorderResponseWriter struct {
	http.ResponseWriter
	StatusCode int
}

func (w *recorderResponseWriter) WriteHeader(code int) {
	w.ResponseWriter.WriteHeader(code)
	w.StatusCode = code
}

// MiddlewareFunc makes StatsMiddleware implement the Middleware interface.
func (mw *StatsMiddleware) Handler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()

		writer := &recorderResponseWriter{w, 0}

		h.ServeHTTP(writer, r)

		mw.handleWriter(start, writer)
	})
}

// Negroni compatible interface
func (mw *StatsMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	start := time.Now()

	writer := &recorderResponseWriter{w, 200}

	next(writer, r)

	mw.handleWriter(start, writer)
}

func (mw *StatsMiddleware) handleWriter(start time.Time, writer *recorderResponseWriter) {
	end := time.Now()

	responseTime := end.Sub(start)

	mw.Lock.Lock()

	defer mw.Lock.Unlock()

	statusCode := fmt.Sprintf("%d", writer.StatusCode)

	mw.ResponseCounts[statusCode]++
	mw.TotalResponseCounts[statusCode]++
	mw.TotalResponseTime = mw.TotalResponseTime.Add(responseTime)
}

type Stats struct {
	Pid                    int            `json: "pid"`
	UpTime                 string         `json: "uptime"`
	UpTimeSec              float64        `json: "uptime_sec"`
	Time                   string         `json: "time"`
	TimeUnix               int64          `json: "unixtime"`
	StatusCodeCount        map[string]int `json: "status_code_count"`
	TotalStatusCodeCount   map[string]int `json: "total_status_code_count"`
	Count                  int            `json: "count"`
	TotalCount             int            `json: "total_count"`
	TotalResponseTime      string         `json" "total_response_time`
	TotalResponseTimeSec   float64        `json: "total_response_time_sec"`
	AverageResponseTime    string         `json: "average_response_time"`
	AverageResponseTimeSec float64        `json: "average_response_time_sec"`
}

func (mw *StatsMiddleware) GetStats() *Stats {

	mw.Lock.RLock()

	now := time.Now()

	uptime := now.Sub(mw.Start)

	count := 0
	for _, current := range mw.ResponseCounts {
		count += current
	}

	totalCount := 0
	for _, count := range mw.TotalResponseCounts {
		totalCount += count
	}

	totalResponseTime := mw.TotalResponseTime.Sub(time.Time{})

	averageResponseTime := time.Duration(0)
	if totalCount > 0 {
		avgNs := int64(totalResponseTime) / int64(totalCount)
		averageResponseTime = time.Duration(avgNs)
	}

	stats := &Stats{
		Pid:                    mw.Pid,
		UpTime:                 uptime.String(),
		UpTimeSec:              uptime.Seconds(),
		Time:                   now.String(),
		TimeUnix:               now.Unix(),
		StatusCodeCount:        mw.ResponseCounts,
		TotalStatusCodeCount:   mw.TotalResponseCounts,
		Count:                  count,
		TotalCount:             totalCount,
		TotalResponseTime:      totalResponseTime.String(),
		TotalResponseTimeSec:   totalResponseTime.Seconds(),
		AverageResponseTime:    averageResponseTime.String(),
		AverageResponseTimeSec: averageResponseTime.Seconds(),
	}

	mw.Lock.RUnlock()

	return stats
}
<|endoftext|>