<commit_before>// Copyright 2017, Kerby Shedden and the Muscato contributors.

// muscato_screen is an initial screening step used by Muscato to
// identify candidate matches of a set of reads into a set of target
// gene sequences.  The results of the screen may contain false
// positives, but will not contain any false negatives.
//
// The approach is to use a Bloom filter to sketch the reads based on
// the subsequences that appear at defined offsets within the reads.
// For example, if position 10 is an offset and we are looking at
// subequences of width 15, then the read subsequences from position
// 10 through position 25 are entered into a Bloom filter.  Then, we
// scan through every target gene looking for matches to the Bloom
// filter.  When a match occurs, the match position (in the target)
// and flanking sequences are saved for subequent checking against the
// full read sequence.
//
// A simple entropy check is used to avoid considering subsequences
// that could match large numbers of reads or genes (and hence would
// be uninformative).  Currently, this check is based on the number of
// distinct dinucleotide subsequences in the window (e.g. in the
// 15-mer in the example above).
//
// The results are saved in files named bmatch*.txt.sz, where * is the
// window number.
//
// The format of the bmatch files is:
//
// (window sequence) (left tail) (right tail) (gene id) (position)

package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"log"
	"math/rand"
	"os"
	"path"
	"strings"
	"sync"

	"github.com/chmduquesne/rollinghash"
	"github.com/chmduquesne/rollinghash/buzhash32"
	"github.com/golang-collections/go-datastructures/bitarray"
	"github.com/golang/snappy"
	"github.com/kshedden/seqmatch/utils"
)

const (
	// Number of goroutines, should probably scale with the number
	// of aailable cores.
	concurrency int = 100
)

var (
	// A log
	logger *log.Logger

	// Configuration information
	config *utils.Config

	// All working files are stored here
	tmpdir string

	// Bitarrays that back the Bloom filters
	smp []bitarray.BitArray

	// Tables to produce independent running hashes
	tables [][256]uint32

	// Communicate results back to driver
	hitchan chan rec

	// Semaphore for limiting goroutines
	limit chan bool

	// Line length for output
	bufsize int
)

// genTables generates base hash functions for a collection of rolling hashes.
func genTables() {
	tables = make([][256]uint32, config.NumHash)
	for j := 0; j < config.NumHash; j++ {
		mp := make(map[uint32]bool)
		for i := 0; i < 256; i++ {
			for {
				x := uint32(rand.Int63())
				if !mp[x] {
					tables[j][i] = x
					mp[x] = true
					break
				}
			}
		}
	}
}

// buildBloom constructs bloom filters for each window
func buildBloom() {

	logger.Printf("Building Bloom sketch of read collection...")

	hashes := make([]rollinghash.Hash32, config.NumHash)
	for j := range hashes {
		hashes[j] = buzhash32.NewFromUint32Array(tables[j])
	}

	fname := path.Join(tmpdir, "reads_sorted.txt.sz")
	fid, err := os.Open(fname)
	if err != nil {
		logger.Print(err)
		panic(err)
	}
	defer fid.Close()
	snr := snappy.NewReader(fid)
	scanner := bufio.NewScanner(snr)
	scanner.Buffer(make([]byte, 1024*1024), 1024*1024)

	// Workspace for sequence diversity checker
	wk := make([]int, 25)

	var j int
	for ; scanner.Scan(); j++ {

		if j%1000000 == 0 {
			logger.Printf("%d\n", j)
		}

		line := scanner.Bytes()
		seq := bytes.Fields(line)[0]

		for k := 0; k < len(config.Windows); k++ {
			q1 := config.Windows[k]
			q2 := q1 + config.WindowWidth
			if q2 > len(seq) {
				continue
			}
			seqw := seq[q1:q2]

			// Check entropy
			if utils.CountDinuc(seqw, wk) < config.MinDinuc {
				continue
			}

			// Update the Bloom filter for this sequence
			for _, ha := range hashes {
				ha.Reset()
				_, err = ha.Write(seqw)
				if err != nil {
					panic(err)
				}
				x := uint64(ha.Sum32()) % config.BloomSize
				err := smp[k].SetBit(x)
				if err != nil {
					logger.Print(err)
					panic(err)
				}
			}
		}
	}

	if err := scanner.Err(); err != nil {
		msg := fmt.Sprintf("Problem reading reads_sorted.txt.sz on line %d\n", j)
		os.Stderr.WriteString(msg)
		logger.Print(err)
		panic(err)
	}

	logger.Printf("Done constructing Bloom filters")
}

type rec struct {
	mseq  string
	left  string
	right string
	win   int
	tnum  int
	pos   uint32
}

// checkWin returns the indices of the Bloom filters that match the
// current state of the hashes.  iw is workspace and hashses contains
// the hashes that define the Bloom filters.
func checkWin(ix []int, iw []uint64, hashes []rollinghash.Hash32) []int {

	// Get the hash states
	for j, ha := range hashes {
		iw[j] = uint64(ha.Sum32()) % config.BloomSize
	}

	ix = ix[0:0]

	// Loop over Bloom filters
	for k, ba := range smp {

		// Determine if the Bloom filter matches
		g := true
		for j := range hashes {
			f, err := ba.GetBit(iw[j])
			if err != nil {
				logger.Print(err)
				panic(err)
			}
			if !f {
				g = false
				break
			}
		}
		if g {
			ix = append(ix, k)
		}
	}

	return ix
}

// process one target sequence, runs concurrently with main loop.
func processseq(seq []byte, genenum int) {

	defer func() { <-limit }()

	hashes := make([]rollinghash.Hash32, config.NumHash)
	for j := range hashes {
		hashes[j] = buzhash32.NewFromUint32Array(tables[j])
	}

	// Initialize the hashes with the first window.
	hlen := config.WindowWidth
	if len(seq) < hlen {
		// Not long enough to fit even one window.
		return
	}
	for j := range hashes {
		_, err := hashes[j].Write(seq[0:hlen])
		if err != nil {
			panic(err)
		}
	}

	ix := make([]int, len(smp))
	iw := make([]uint64, config.NumHash)

	// Check if the initial window is a match
	ix = checkWin(ix, iw, hashes)
	for _, i := range ix {

		q1 := config.Windows[i]
		if q1 != 0 {
			// The only way the read can fit is if the
			// window starts at the beginning of the read.
			continue
		}
		q2 := q1 + config.WindowWidth

		jz := 100 - q2
		if jz > len(seq) {
			jz = len(seq)
		}
		hitchan <- rec{
			mseq:  string(seq[0:hlen]),
			left:  "",
			right: string(seq[hlen:jz]),
			tnum:  genenum,
			win:   i,
			pos:   0,
		}
	}

	// Check the rest of the windows
	for j := hlen; j < len(seq); j++ {

		for _, ha := range hashes {
			ha.Roll(seq[j])
		}
		ix = checkWin(ix, iw, hashes)

		// Process a match
		for _, i := range ix {

			q1 := config.Windows[i]
			q2 := q1 + config.WindowWidth
			if j < q2-1 {
				// The read would not fit
				continue
			}

			// Matching sequence is jx:jy
			jx := j - hlen + 1
			jy := j + 1

			// Left tail is jw:jx
			jw := jx - q1

			// Right tail is jy:jz
			jz := jy + config.MaxReadLength - q2
			if jz > len(seq) {
				// May not be long enough to fit, but
				// we don't know until we merge.
				jz = len(seq)
			}

			if jw >= 0 {
				hitchan <- rec{
					mseq:  string(seq[jx:jy]),
					left:  string(seq[jw:jx]),
					right: string(seq[jy:jz]),
					tnum:  genenum,
					win:   i,
					pos:   uint32(j - hlen + 1),
				}
			}
		}
	}
}

// Retrieve the results and write to disk
func harvest(wg *sync.WaitGroup) {

	var wtrs []io.Writer
	var allwtrs []io.Closer
	for k := 0; k < len(config.Windows); k++ {
		f := fmt.Sprintf("bmatch_%d.txt.sz", k)
		outname := path.Join(tmpdir, f)
		out, err := os.Create(outname)
		if err != nil {
			logger.Print(err)
			panic(err)
		}
		wtr := snappy.NewBufferedWriter(out)
		wtrs = append(wtrs, wtr)
		allwtrs = append(allwtrs, wtr, out)
	}

	bb := bytes.Repeat([]byte(" "), bufsize)
	bb[bufsize-1] = byte('\n')

	for r := range hitchan {

		wtr := wtrs[r.win]

		n1, err1 := wtr.Write([]byte(fmt.Sprintf("%s\t", r.mseq)))
		n2, err2 := wtr.Write([]byte(fmt.Sprintf("%s\t", r.left)))
		n3, err3 := wtr.Write([]byte(fmt.Sprintf("%s\t", r.right)))
		n4, err4 := wtr.Write([]byte(fmt.Sprintf("%011d\t", r.tnum)))
		n5, err5 := wtr.Write([]byte(fmt.Sprintf("%d", r.pos)))

		for _, err := range []error{err1, err2, err3, err4, err5} {
			if err != nil {
				logger.Print(err)
				panic("writing error")
			}
		}

		n := n1 + n2 + n3 + n4 + n5
		if n > bufsize {
			panic("output line is too long")
		}

		// The rest of the line is spaces, then newline.
		_, err := wtr.Write(bb[n:bufsize])
		if err != nil {
			logger.Print(err)
			panic(err)
		}
	}

	for _, wtr := range allwtrs {
		wtr.Close()
	}
	wg.Done()
	logger.Printf("Exiting harvest")
}

// search loops through the target sequences, checking each window
// within each target gene for possible matches to the read
// collection.
func search() {

	fid, err := os.Open(config.GeneFileName)
	if err != nil {
		logger.Print(err)
		panic(err)
	}
	defer fid.Close()
	snr := snappy.NewReader(fid)

	// Target file contains some very long lines
	scanner := bufio.NewScanner(snr)
	sbuf := make([]byte, 1024*1024)
	scanner.Buffer(sbuf, 1024*1024)

	hitchan = make(chan rec)
	limit = make(chan bool, concurrency)

	var wg sync.WaitGroup
	wg.Add(1)
	go harvest(&wg)

	var i int
	for ; scanner.Scan(); i++ {

		if i%1000000 == 0 {
			logger.Printf("%d\n", i)
		}

		line := scanner.Text() // need a copy here

		toks := strings.Split(line, "\t")
		seq := toks[0] // The sequence

		limit <- true
		go processseq([]byte(seq), i)
	}

	if err := scanner.Err(); err != nil {
		msg := fmt.Sprintf("Problem reading %s on line %d\n", config.GeneFileName, i)
		os.Stderr.WriteString(msg)
		logger.Print(err)
		panic(err)
	}

	for k := 0; k < concurrency; k++ {
		limit <- true
	}

	close(hitchan)
	wg.Wait()
	logger.Printf("done with search")
}

func setupLogger() {
	logname := path.Join(tmpdir, "muscato_screen.log")
	logfid, err := os.Create(logname)
	if err != nil {
		panic(err)
	}
	logger = log.New(logfid, "", log.Ltime)
}

func estimateFullness() {

	n := 1000
	logger.Printf("Bloom filter fill rates:\n")

	for j, ba := range smp {
		c := 0
		for k := 0; k < n; k++ {
			i := uint64(rand.Int63()) % config.BloomSize
			f, err := ba.GetBit(i)
			if err != nil {
				panic(err)
			}
			if f {
				c++
			}
		}
		logger.Printf("%3d %.3f\n", j, float64(c)/float64(n))
	}
}

func main() {

	if len(os.Args) != 3 {
		panic("wrong number of arguments")
	}

	config = utils.ReadConfig(os.Args[1])

	if config.TempDir == "" {
		tmpdir = os.Args[2]
	} else {
		tmpdir = config.TempDir
	}

	bufsize = config.MaxReadLength + 50

	setupLogger()
	genTables()

	smp = make([]bitarray.BitArray, len(config.Windows))
	for k := range smp {
		smp[k] = bitarray.NewBitArray(config.BloomSize)
	}

	buildBloom()
	estimateFullness()
	search()
}
<commit_msg>update<commit_after>// Copyright 2017, Kerby Shedden and the Muscato contributors.

// muscato_screen is an initial screening step used by Muscato to
// identify candidate matches of a set of reads into a set of target
// gene sequences.  The results of the screen may contain false
// positives, but will not contain any false negatives.
//
// The approach is to use a Bloom filter to sketch the reads based on
// the subsequences that appear at defined offsets within the reads.
// For example, if position 10 is an offset and we are looking at
// subequences of width 15, then the read subsequences from position
// 10 through position 25 are entered into a Bloom filter.  Then, we
// scan through every target gene looking for matches to the Bloom
// filter.  When a match occurs, the match position (in the target)
// and flanking sequences are saved for subequent checking against the
// full read sequence.
//
// A simple entropy check is used to avoid considering subsequences
// that could match large numbers of reads or genes (and hence would
// be uninformative).  Currently, this check is based on the number of
// distinct dinucleotide subsequences in the window (e.g. in the
// 15-mer in the example above).
//
// The results are saved in files named bmatch*.txt.sz, where * is the
// window number.
//
// The format of the bmatch files is:
//
// (window sequence) (left tail) (right tail) (gene id) (position)

package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"log"
	"math/rand"
	"os"
	"path"
	"strings"
	"sync"

	"github.com/chmduquesne/rollinghash"
	"github.com/chmduquesne/rollinghash/buzhash32"
	"github.com/golang-collections/go-datastructures/bitarray"
	"github.com/golang/snappy"
	"github.com/kshedden/seqmatch/utils"
)

const (
	// Number of goroutines, should probably scale with the number
	// of aailable cores.
	concurrency int = 100
)

var (
	// A log
	logger *log.Logger

	// Configuration information
	config *utils.Config

	// All working files are stored here
	tmpdir string

	// Bitarrays that back the Bloom filters
	smp []bitarray.BitArray

	// Tables to produce independent running hashes
	tables [][256]uint32

	// Communicate results back to driver
	hitchan chan rec

	// Semaphore for limiting goroutines
	limit chan bool

	// Line length for output
	bufsize int
)

// genTables generates base hash functions for a collection of rolling hashes.
func genTables() {
	tables = make([][256]uint32, config.NumHash)
	for j := 0; j < config.NumHash; j++ {
		mp := make(map[uint32]bool)
		for i := 0; i < 256; i++ {
			for {
				x := uint32(rand.Int63())
				if !mp[x] {
					tables[j][i] = x
					mp[x] = true
					break
				}
			}
		}
	}
}

// buildBloom constructs bloom filters for each window
func buildBloom() {

	logger.Printf("Building Bloom sketch of read collection...")

	hashes := make([]rollinghash.Hash32, config.NumHash)
	for j := range hashes {
		hashes[j] = buzhash32.NewFromUint32Array(tables[j])
	}

	fname := path.Join(tmpdir, "reads_sorted.txt.sz")
	fid, err := os.Open(fname)
	if err != nil {
		logger.Print(err)
		panic(err)
	}
	defer fid.Close()
	snr := snappy.NewReader(fid)
	scanner := bufio.NewScanner(snr)
	scanner.Buffer(make([]byte, 1024*1024), 1024*1024)

	// Workspace for sequence diversity checker
	wk := make([]int, 25)

	var j int
	for ; scanner.Scan(); j++ {

		if j%1000000 == 0 {
			logger.Printf("%d\n", j)
		}

		line := scanner.Bytes()
		seq := bytes.Fields(line)[0]

		for k := 0; k < len(config.Windows); k++ {
			q1 := config.Windows[k]
			q2 := q1 + config.WindowWidth
			if q2 > len(seq) {
				continue
			}
			seqw := seq[q1:q2]

			// Check entropy
			if utils.CountDinuc(seqw, wk) < config.MinDinuc {
				continue
			}

			// Update the Bloom filter for this sequence
			for _, ha := range hashes {
				ha.Reset()
				_, err = ha.Write(seqw)
				if err != nil {
					panic(err)
				}
				x := uint64(ha.Sum32()) % config.BloomSize
				err := smp[k].SetBit(x)
				if err != nil {
					logger.Print(err)
					panic(err)
				}
			}
		}
	}

	if err := scanner.Err(); err != nil {
		msg := fmt.Sprintf("Problem reading reads_sorted.txt.sz on line %d\n", j)
		os.Stderr.WriteString(msg)
		logger.Print(err)
		panic(err)
	}

	logger.Printf("Done constructing Bloom filters")
}

type rec struct {
	mseq  string
	left  string
	right string
	win   int
	tnum  int
	pos   uint32
}

// checkWin returns the indices of the Bloom filters that match the
// current state of the hashes.  iw is workspace and hashses contains
// the hashes that define the Bloom filters.
func checkWin(ix []int, iw []uint64, hashes []rollinghash.Hash32) []int {

	// Get the hash states
	for j, ha := range hashes {
		iw[j] = uint64(ha.Sum32()) % config.BloomSize
	}

	ix = ix[0:0]

	// Loop over Bloom filters
	for k, ba := range smp {

		// Determine if the Bloom filter matches
		g := true
		for j := range hashes {
			f, err := ba.GetBit(iw[j])
			if err != nil {
				logger.Print(err)
				panic(err)
			}
			if !f {
				g = false
				break
			}
		}
		if g {
			ix = append(ix, k)
		}
	}

	return ix
}

// process one target sequence, runs concurrently with main loop.
func processseq(seq []byte, genenum int) {

	defer func() { <-limit }()

	hashes := make([]rollinghash.Hash32, config.NumHash)
	for j := range hashes {
		hashes[j] = buzhash32.NewFromUint32Array(tables[j])
	}

	// Initialize the hashes with the first window.
	hlen := config.WindowWidth
	if len(seq) < hlen {
		// Not long enough to fit even one window.
		return
	}
	for j := range hashes {
		_, err := hashes[j].Write(seq[0:hlen])
		if err != nil {
			panic(err)
		}
	}

	// Will contain the indices of the matchine windows
	ix := make([]int, len(smp))

	// Workspace
	iw := make([]uint64, config.NumHash)

	// Check if the initial window is a match
	ix = checkWin(ix, iw, hashes)
	for _, i := range ix {

		q1 := config.Windows[i]
		if q1 != 0 {
			// The only way the read can fit is if the
			// window starts at the beginning of the read.
			continue
		}
		q2 := q1 + config.WindowWidth

		jz := 100 - q2
		if jz > len(seq) {
			jz = len(seq)
		}
		hitchan <- rec{
			mseq:  string(seq[0:hlen]),
			left:  "",
			right: string(seq[hlen:jz]),
			tnum:  genenum,
			win:   i,
			pos:   0,
		}
	}

	// Check the rest of the windows
	for j := hlen; j < len(seq); j++ {

		for _, ha := range hashes {
			ha.Roll(seq[j])
		}
		ix = checkWin(ix, iw, hashes)

		// Process a match
		for _, i := range ix {

			q1 := config.Windows[i]
			q2 := q1 + config.WindowWidth
			if j < q2-1 {
				// The read would not fit
				continue
			}

			// Matching sequence is jx:jy
			jx := j - hlen + 1
			jy := j + 1

			// Left tail is jw:jx
			jw := jx - q1

			// Right tail is jy:jz
			jz := jy + config.MaxReadLength - q2
			if jz > len(seq) {
				// May not be long enough to fit, but
				// we don't know until we merge.
				jz = len(seq)
			}

			if jw >= 0 {
				hitchan <- rec{
					mseq:  string(seq[jx:jy]),
					left:  string(seq[jw:jx]),
					right: string(seq[jy:jz]),
					tnum:  genenum,
					win:   i,
					pos:   uint32(j - hlen + 1),
				}
			}
		}
	}
}

// Retrieve the results and write to disk
func harvest(wg *sync.WaitGroup) {

	var wtrs []io.Writer
	var allwtrs []io.Closer
	for k := 0; k < len(config.Windows); k++ {
		f := fmt.Sprintf("bmatch_%d.txt.sz", k)
		outname := path.Join(tmpdir, f)
		out, err := os.Create(outname)
		if err != nil {
			logger.Print(err)
			panic(err)
		}
		wtr := snappy.NewBufferedWriter(out)
		wtrs = append(wtrs, wtr)
		allwtrs = append(allwtrs, wtr, out)
	}

	bb := bytes.Repeat([]byte(" "), bufsize)
	bb[bufsize-1] = byte('\n')

	for r := range hitchan {

		wtr := wtrs[r.win]

		n1, err1 := wtr.Write([]byte(fmt.Sprintf("%s\t", r.mseq)))
		n2, err2 := wtr.Write([]byte(fmt.Sprintf("%s\t", r.left)))
		n3, err3 := wtr.Write([]byte(fmt.Sprintf("%s\t", r.right)))
		n4, err4 := wtr.Write([]byte(fmt.Sprintf("%011d\t", r.tnum)))
		n5, err5 := wtr.Write([]byte(fmt.Sprintf("%d", r.pos)))

		for _, err := range []error{err1, err2, err3, err4, err5} {
			if err != nil {
				logger.Print(err)
				panic("writing error")
			}
		}

		n := n1 + n2 + n3 + n4 + n5
		if n > bufsize {
			panic("output line is too long")
		}

		// The rest of the line is spaces, then newline.
		_, err := wtr.Write(bb[n:bufsize])
		if err != nil {
			logger.Print(err)
			panic(err)
		}
	}

	for _, wtr := range allwtrs {
		wtr.Close()
	}
	wg.Done()
	logger.Printf("Exiting harvest")
}

// search loops through the target sequences, checking each window
// within each target gene for possible matches to the read
// collection.
func search() {

	fid, err := os.Open(config.GeneFileName)
	if err != nil {
		logger.Print(err)
		panic(err)
	}
	defer fid.Close()
	snr := snappy.NewReader(fid)

	// Target file contains some very long lines
	scanner := bufio.NewScanner(snr)
	sbuf := make([]byte, 1024*1024)
	scanner.Buffer(sbuf, 1024*1024)

	hitchan = make(chan rec)
	limit = make(chan bool, concurrency)

	var wg sync.WaitGroup
	wg.Add(1)
	go harvest(&wg)

	var i int
	for ; scanner.Scan(); i++ {

		if i%1000000 == 0 {
			logger.Printf("%d\n", i)
		}

		line := scanner.Text() // need a copy here

		toks := strings.Split(line, "\t")
		seq := toks[0] // The sequence

		limit <- true
		go processseq([]byte(seq), i)
	}

	if err := scanner.Err(); err != nil {
		msg := fmt.Sprintf("Problem reading %s on line %d\n", config.GeneFileName, i)
		os.Stderr.WriteString(msg)
		logger.Print(err)
		panic(err)
	}

	for k := 0; k < concurrency; k++ {
		limit <- true
	}

	close(hitchan)
	wg.Wait()
	logger.Printf("done with search")
}

func setupLogger() {
	logname := path.Join(tmpdir, "muscato_screen.log")
	logfid, err := os.Create(logname)
	if err != nil {
		panic(err)
	}
	logger = log.New(logfid, "", log.Ltime)
}

func estimateFullness() {

	n := 1000
	logger.Printf("Bloom filter fill rates:\n")

	for j, ba := range smp {
		c := 0
		for k := 0; k < n; k++ {
			i := uint64(rand.Int63()) % config.BloomSize
			f, err := ba.GetBit(i)
			if err != nil {
				panic(err)
			}
			if f {
				c++
			}
		}
		logger.Printf("%3d %.3f\n", j, float64(c)/float64(n))
	}
}

func main() {

	if len(os.Args) != 3 {
		panic("wrong number of arguments")
	}

	config = utils.ReadConfig(os.Args[1])

	if config.TempDir == "" {
		tmpdir = os.Args[2]
	} else {
		tmpdir = config.TempDir
	}

	bufsize = config.MaxReadLength + 50

	setupLogger()
	genTables()

	smp = make([]bitarray.BitArray, len(config.Windows))
	for k := range smp {
		smp[k] = bitarray.NewBitArray(config.BloomSize)
	}

	buildBloom()
	estimateFullness()
	search()
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Jetstack cert-manager contributors.

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

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

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

package flags

import (
	"fmt"
	"os"
	"strings"

	flag "github.com/spf13/pflag"
	"os/exec"

	logf "github.com/jetstack/cert-manager/hack/release/pkg/log"
	"github.com/jetstack/cert-manager/hack/release/pkg/util"
)

var (
	Default = &Global{}

	log = logf.Log.WithName("global")
)

type Global struct {
	// Path to the root of the cert-manager repository
	RepoRoot string

	// DockerRepo is the docker repository used to store release images
	DockerRepo string

	// UpstreamRepoURL is the URL of the git repo used to check for tags
	UpstreamRepoURL string

	// AppVersion is the version tag to use for this release
	AppVersion string

	// GitState contains the state of the git working tree.
	GitState string

	// GitCommitRef is the current git commit hash being built
	GitCommitRef string

	GitPath string

	// Path to the cert-manager Helm chart.
	// This is defined as a global as the manifests plugin also needs
	// access to this flag
	ChartPath string
}

const defaultDockerRepo = "quay.io/jetstack"

func (g *Global) AddFlags(fs *flag.FlagSet) {
	fs.StringVar(&g.RepoRoot, "repo-root", "", "path to the root of the cert-manager repository")
	fs.StringVar(&g.DockerRepo, "docker-repo", defaultDockerRepo, "the docker repository that images will be tagged with")
	fs.StringVar(&g.AppVersion, "app-version", "", "app version to use when building and generating manifests. Defaults to 'git describe --tags --abbrev=0 --exact-match'")
	fs.StringVar(&g.UpstreamRepoURL, "git.upstream-repo-url", "https://github.com/jetstack/cert-manager.git", "the URL of the git repo used to check for tags when generating --app-version")
	fs.StringVar(&g.GitState, "git.state", "", "the state of the git working tree. if set and not 'clean', this will be appended to the app-version during builds")
	fs.StringVar(&g.GitCommitRef, "git.commit-ref", "", "the git commit ref of this build. Defaults to 'git rev-parse --short HEAD'")
	fs.StringVar(&g.GitPath, "git.path", "git", "path to the git binary to use")
	fs.StringVar(&g.ChartPath, "chart.path", "deploy/charts/cert-manager", "the path to the cert-manager helm chart, relative to the repo root")
}

func (g *Global) Validate() []error {
	var errs []error

	if g.UpstreamRepoURL == "" {
		errs = append(errs, fmt.Errorf("--git.upstream-repo-url must be specified"))
	}

	return errs
}

func (g *Global) Complete() error {
	log = log.WithName("default-flags")

	if g.DockerRepo == "" {
		log := log.WithValues("flag", "docker-repo")
		g.DockerRepo = defaultDockerRepo
		log.Info("set default value", "value", g.DockerRepo)
	}

	if g.RepoRoot == "" {
		log := log.WithValues("flag", "repo-root")
		if bwd := os.Getenv("BUILD_WORKSPACE_DIRECTORY"); bwd != "" {
			g.RepoRoot = bwd
		} else {
			dir, err := os.Getwd()
			if err != nil {
				return fmt.Errorf("error determining repo root: %v", err)
			}
			g.RepoRoot = dir
		}

		log.Info("set default value", "value", g.RepoRoot)
	}
	if err := os.Chdir(g.RepoRoot); err != nil {
		return fmt.Errorf("error changing directory to --repo-root=%q: %v", g.RepoRoot, err)
	}

	if g.AppVersion == "" {
		log := log.WithValues("flag", "app-version")

		log.V(logf.LogLevelDebug).Info("fetching upstream git repo tags")
		_, err := g.gitOutput("fetch", "--tags", g.UpstreamRepoURL)
		if err != nil {
			return fmt.Errorf("error fetching tags: %v", err)
		}

		log.V(logf.LogLevelDebug).Info("finding tags that match the current commit ref")
		g.AppVersion, err = g.gitOutput("describe", "--tags", "--abbrev=0", "--exact-match")
		if err != nil {
			log.Error(err, "failed to determine tag for current git ref/HEAD")
			g.AppVersion = ""
		}

		if g.AppVersion == "" {
			// default to 'canary' if no tags point to the current ref
			g.AppVersion = "canary"
		}

		log.WithValues("value", g.AppVersion).Info("set default value")
	}

	if g.GitCommitRef == "" {
		log := log.WithValues("flag", "git.commit-ref")

		log.V(logf.LogLevelDebug).Info("parsing current git commit ref")
		var err error
		g.GitCommitRef, err = g.gitOutput("rev-parse", "--short", "HEAD")
		if err != nil {
			return fmt.Errorf("error getting current commit ref: %v", err)
		}

		log.WithValues("value", g.GitCommitRef).Info("set default value")
	}

	if g.GitState == "" {
		log := log.WithValues("flag", "git.commit-state")

		log.V(logf.LogLevelDebug).Info("evaluating current git working tree dirty status")
		changes, err := g.gitOutput("status", "--porcelain")
		if err != nil {
			return fmt.Errorf("error checking git status: %v", err)
		}

		if len(changes) == 0 {
			g.GitState = "clean"
		} else {
			g.GitState = "dirty"
		}

		log.WithValues("value", g.GitState).Info("set default value")
	}

	return nil
}

func (g *Global) gitOutput(args ...string) (string, error) {
	cmd := exec.Command(g.GitPath, args...)
	b, err := util.RunPrintCombined(log, cmd)
	return strings.TrimSpace(string(b)), err
}
<commit_msg>Set default appVersin to v0.0.0-experimental<commit_after>/*
Copyright 2019 The Jetstack cert-manager contributors.

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

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

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

package flags

import (
	"fmt"
	"os"
	"strings"

	flag "github.com/spf13/pflag"
	"os/exec"

	logf "github.com/jetstack/cert-manager/hack/release/pkg/log"
	"github.com/jetstack/cert-manager/hack/release/pkg/util"
)

var (
	Default = &Global{}

	log = logf.Log.WithName("global")
)

type Global struct {
	// Path to the root of the cert-manager repository
	RepoRoot string

	// DockerRepo is the docker repository used to store release images
	DockerRepo string

	// UpstreamRepoURL is the URL of the git repo used to check for tags
	UpstreamRepoURL string

	// AppVersion is the version tag to use for this release
	AppVersion string

	// GitState contains the state of the git working tree.
	GitState string

	// GitCommitRef is the current git commit hash being built
	GitCommitRef string

	GitPath string

	// Path to the cert-manager Helm chart.
	// This is defined as a global as the manifests plugin also needs
	// access to this flag
	ChartPath string
}

const defaultDockerRepo = "quay.io/jetstack"

func (g *Global) AddFlags(fs *flag.FlagSet) {
	fs.StringVar(&g.RepoRoot, "repo-root", "", "path to the root of the cert-manager repository")
	fs.StringVar(&g.DockerRepo, "docker-repo", defaultDockerRepo, "the docker repository that images will be tagged with")
	fs.StringVar(&g.AppVersion, "app-version", "", "app version to use when building and generating manifests. Defaults to 'git describe --tags --abbrev=0 --exact-match'")
	fs.StringVar(&g.UpstreamRepoURL, "git.upstream-repo-url", "https://github.com/jetstack/cert-manager.git", "the URL of the git repo used to check for tags when generating --app-version")
	fs.StringVar(&g.GitState, "git.state", "", "the state of the git working tree. if set and not 'clean', this will be appended to the app-version during builds")
	fs.StringVar(&g.GitCommitRef, "git.commit-ref", "", "the git commit ref of this build. Defaults to 'git rev-parse --short HEAD'")
	fs.StringVar(&g.GitPath, "git.path", "git", "path to the git binary to use")
	fs.StringVar(&g.ChartPath, "chart.path", "deploy/charts/cert-manager", "the path to the cert-manager helm chart, relative to the repo root")
}

func (g *Global) Validate() []error {
	var errs []error

	if g.UpstreamRepoURL == "" {
		errs = append(errs, fmt.Errorf("--git.upstream-repo-url must be specified"))
	}

	return errs
}

func (g *Global) Complete() error {
	log = log.WithName("default-flags")

	if g.DockerRepo == "" {
		log := log.WithValues("flag", "docker-repo")
		g.DockerRepo = defaultDockerRepo
		log.Info("set default value", "value", g.DockerRepo)
	}

	if g.RepoRoot == "" {
		log := log.WithValues("flag", "repo-root")
		if bwd := os.Getenv("BUILD_WORKSPACE_DIRECTORY"); bwd != "" {
			g.RepoRoot = bwd
		} else {
			dir, err := os.Getwd()
			if err != nil {
				return fmt.Errorf("error determining repo root: %v", err)
			}
			g.RepoRoot = dir
		}

		log.Info("set default value", "value", g.RepoRoot)
	}
	if err := os.Chdir(g.RepoRoot); err != nil {
		return fmt.Errorf("error changing directory to --repo-root=%q: %v", g.RepoRoot, err)
	}

	if g.AppVersion == "" {
		log := log.WithValues("flag", "app-version")

		log.V(logf.LogLevelDebug).Info("fetching upstream git repo tags")
		_, err := g.gitOutput("fetch", "--tags", g.UpstreamRepoURL)
		if err != nil {
			return fmt.Errorf("error fetching tags: %v", err)
		}

		log.V(logf.LogLevelDebug).Info("finding tags that match the current commit ref")
		g.AppVersion, err = g.gitOutput("describe", "--tags", "--abbrev=0", "--exact-match")
		if err != nil {
			log.Error(err, "failed to determine tag for current git ref/HEAD")
			g.AppVersion = ""
		}

		if g.AppVersion == "" {
			// default to 'v0.0.0-experimental' if no tags point to the current ref
			g.AppVersion = "v0.0.0-experimental"
		}

		log.WithValues("value", g.AppVersion).Info("set default value")
	}

	if g.GitCommitRef == "" {
		log := log.WithValues("flag", "git.commit-ref")

		log.V(logf.LogLevelDebug).Info("parsing current git commit ref")
		var err error
		g.GitCommitRef, err = g.gitOutput("rev-parse", "--short", "HEAD")
		if err != nil {
			return fmt.Errorf("error getting current commit ref: %v", err)
		}

		log.WithValues("value", g.GitCommitRef).Info("set default value")
	}

	if g.GitState == "" {
		log := log.WithValues("flag", "git.commit-state")

		log.V(logf.LogLevelDebug).Info("evaluating current git working tree dirty status")
		changes, err := g.gitOutput("status", "--porcelain")
		if err != nil {
			return fmt.Errorf("error checking git status: %v", err)
		}

		if len(changes) == 0 {
			g.GitState = "clean"
		} else {
			g.GitState = "dirty"
		}

		log.WithValues("value", g.GitState).Info("set default value")
	}

	return nil
}

func (g *Global) gitOutput(args ...string) (string, error) {
	cmd := exec.Command(g.GitPath, args...)
	b, err := util.RunPrintCombined(log, cmd)
	return strings.TrimSpace(string(b)), err
}
<|endoftext|>


<commit_before>// Package https is the supplement of the standard library `http`,
// not the protocal `https`.
package https

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"net/http"
	"os"

	"github.com/xgfone/go-tools/lifecycle"
)

var (
	// DefaultHealthURL is the url to register the health HTTP handler.
	DefaultHealthURL = "/health"

	// DefaultServeMux is the default serve multiplexer.
	DefaultServeMux = http.DefaultServeMux
)

// SetHealthHandler registers the health handler with DefaultHealthURL into
// DefaultServeMux.
//
// Notice: The handler may be nil, which will use the default handler that only
// returns the status code 200.
func SetHealthHandler(handler interface{}) {
	if handler == nil {
		handler = http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})
	} else if _, ok := handler.(http.Handler); !ok {
		panic("The handler is not nil or http.Handler")
	}
	DefaultServeMux.Handle(DefaultHealthURL, handler.(http.Handler))
}

// SetHealthHandlerFunc is same as SetHealthHandler, but the handler is
// a function that returns the body content.
//
// Notice: If the handler function panic,it will return the status code 500.
func SetHealthHandlerFunc(handler func() string) {
	SetHealthHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if recover() != nil {
				w.WriteHeader(http.StatusInternalServerError)
			}
		}()

		io.WriteString(w, handler())
	}))
}

// StartHealthServer starts the HTTP server to serve the health check.
//
// The handler may be nil, http.Handler, or a function, the type of which is
// `func() string` that the returned result is the response content.
//
// If you want to start the server on SSL, you can give certFile and keyFile
// as the last two arguments.
func StartHealthServer(addr string, handler interface{}, files ...string) error {
	switch h := handler.(type) {
	case nil, http.Handler:
		SetHealthHandler(h)
	case func() string:
		SetHealthHandlerFunc(h)
	default:
		return fmt.Errorf("Unknown handler type")
	}

	_len := len(files)
	if _len == 0 {
		return ListenAndServe(addr, DefaultServeMux)
	} else if _len == 2 {
		return ListenAndServeTLS(addr, files[0], files[1], DefaultServeMux)
	}
	return fmt.Errorf("the options files is not certFile and keyFile")
}

// HTTPError stands for a HTTP error.
type HTTPError struct {
	// The error information
	Err error

	// The status code
	Code int

	// You can place data into it to carry in an error.
	Data map[string]interface{}
}

// NewHTTPError returns a new HTTPError.
func NewHTTPError(code int, err interface{}) error {
	switch err.(type) {
	case error:
	case []byte:
		err = fmt.Errorf("%s", string(err.([]byte)))
	default:
		err = fmt.Errorf("%v", err)
	}
	return HTTPError{Code: code, Err: err.(error)}
}

func (e HTTPError) Error() string {
	return e.Err.Error()
}

// ErrorLogFunc handles the http error log in ErrorHandler and
// ErrorHandlerWithStatusCode.
//
// Notice: The caller doesn't append the new line, so the function should
// append the new line.
var ErrorLogFunc func(format string, args ...interface{})

func init() {
	ErrorLogFunc = func(format string, args ...interface{}) {
		fmt.Fprintf(os.Stderr, format+"\n", args...)
	}
}

// ErrorHandler handles the error and responds it the client.
func ErrorHandler(f func(http.ResponseWriter, *http.Request) error) http.HandlerFunc {
	return ErrorHandlerWithStatusCode(func(w http.ResponseWriter,
		r *http.Request) (int, error) {
		return 0, f(w, r)
	})
}

// ErrorHandlerWithStatusCode handles the error and responds it the client
// with the status code.
func ErrorHandlerWithStatusCode(f func(http.ResponseWriter, *http.Request) (int, error)) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if code, err := f(w, r); err != nil {
			if code == 0 {
				if _err, ok := err.(HTTPError); ok {
					code = _err.Code
				} else {
					code = http.StatusInternalServerError
				}
			}
			http.Error(w, err.Error(), code)
			ErrorLogFunc("Handling %q: status=%d, err=%v", r.RequestURI, code, err)
		}
	}
}

// HandlerWrapper handles the response result.
func HandlerWrapper(f func(http.ResponseWriter, *http.Request) (int, []byte, error)) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		code, resp, err := f(w, r)

		// Handle the error.
		if err != nil {
			if code == 0 {
				if _err, ok := err.(HTTPError); ok {
					code = _err.Code
				} else {
					code = http.StatusInternalServerError
				}
			}
			http.Error(w, err.Error(), code)
			ErrorLogFunc("Failed to handle %q: %s", r.RequestURI, err)
			return
		}

		// Determine the status code.
		if code == 0 {
			if len(resp) == 0 {
				code = http.StatusNoContent
			} else {
				code = http.StatusOK
			}
		}

		// Send the response result.
		w.WriteHeader(code)
		if _, err = io.CopyN(w, bytes.NewBuffer(resp), int64(len(resp))); err != nil {
			ErrorLogFunc("Failed to send the response of %q: %s", r.RequestURI, err)
		}
	}
}

// ListenAndServe is equal to http.ListenAndServe, but calling the method
// server.Shutdown(context.TODO()) to shutdown the HTTP server gracefully
// when calling lifecycle.Stop().
//
// Notice: It will call lifecycle.Stop() when the server exits.
func ListenAndServe(addr string, handler http.Handler) error {
	server := http.Server{Addr: addr, Handler: handler}
	lifecycle.Register(func() { server.Shutdown(context.TODO()) })
	err := server.ListenAndServe()
	ErrorLogFunc("The server listening on %s has an error: %s", addr, err)
	lifecycle.Stop()
	return err
}

// ListenAndServeTLS is equal to http.ListenAndServeTLS, but calling the method
// server.Shutdown(context.TODO()) to shutdown the HTTP server gracefully
// when calling lifecycle.Stop().
//
// Notice: It will call lifecycle.Stop() when the server exits.
func ListenAndServeTLS(addr, certFile, keyFile string, handler http.Handler) error {
	server := http.Server{Addr: addr, Handler: handler}
	lifecycle.Register(func() { server.Shutdown(context.TODO()) })
	err := server.ListenAndServeTLS(certFile, keyFile)
	ErrorLogFunc("The TLS server listening on %s has an error: %s", addr, err)
	lifecycle.Stop()
	return err
}
<commit_msg>Add the function GetBody<commit_after>// Package https is the supplement of the standard library `http`,
// not the protocal `https`.
package https

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"net/http"
	"os"

	"github.com/xgfone/go-tools/lifecycle"
)

var (
	// DefaultHealthURL is the url to register the health HTTP handler.
	DefaultHealthURL = "/health"

	// DefaultServeMux is the default serve multiplexer.
	DefaultServeMux = http.DefaultServeMux
)

// SetHealthHandler registers the health handler with DefaultHealthURL into
// DefaultServeMux.
//
// Notice: The handler may be nil, which will use the default handler that only
// returns the status code 200.
func SetHealthHandler(handler interface{}) {
	if handler == nil {
		handler = http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})
	} else if _, ok := handler.(http.Handler); !ok {
		panic("The handler is not nil or http.Handler")
	}
	DefaultServeMux.Handle(DefaultHealthURL, handler.(http.Handler))
}

// SetHealthHandlerFunc is same as SetHealthHandler, but the handler is
// a function that returns the body content.
//
// Notice: If the handler function panic,it will return the status code 500.
func SetHealthHandlerFunc(handler func() string) {
	SetHealthHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if recover() != nil {
				w.WriteHeader(http.StatusInternalServerError)
			}
		}()

		io.WriteString(w, handler())
	}))
}

// StartHealthServer starts the HTTP server to serve the health check.
//
// The handler may be nil, http.Handler, or a function, the type of which is
// `func() string` that the returned result is the response content.
//
// If you want to start the server on SSL, you can give certFile and keyFile
// as the last two arguments.
func StartHealthServer(addr string, handler interface{}, files ...string) error {
	switch h := handler.(type) {
	case nil, http.Handler:
		SetHealthHandler(h)
	case func() string:
		SetHealthHandlerFunc(h)
	default:
		return fmt.Errorf("Unknown handler type")
	}

	_len := len(files)
	if _len == 0 {
		return ListenAndServe(addr, DefaultServeMux)
	} else if _len == 2 {
		return ListenAndServeTLS(addr, files[0], files[1], DefaultServeMux)
	}
	return fmt.Errorf("the options files is not certFile and keyFile")
}

// HTTPError stands for a HTTP error.
type HTTPError struct {
	// The error information
	Err error

	// The status code
	Code int

	// You can place data into it to carry in an error.
	Data map[string]interface{}
}

// NewHTTPError returns a new HTTPError.
func NewHTTPError(code int, err interface{}) error {
	switch err.(type) {
	case error:
	case []byte:
		err = fmt.Errorf("%s", string(err.([]byte)))
	default:
		err = fmt.Errorf("%v", err)
	}
	return HTTPError{Code: code, Err: err.(error)}
}

func (e HTTPError) Error() string {
	return e.Err.Error()
}

// ErrorLogFunc handles the http error log in ErrorHandler and
// ErrorHandlerWithStatusCode.
//
// Notice: The caller doesn't append the new line, so the function should
// append the new line.
var ErrorLogFunc func(format string, args ...interface{})

func init() {
	ErrorLogFunc = func(format string, args ...interface{}) {
		fmt.Fprintf(os.Stderr, format+"\n", args...)
	}
}

// ErrorHandler handles the error and responds it the client.
func ErrorHandler(f func(http.ResponseWriter, *http.Request) error) http.HandlerFunc {
	return ErrorHandlerWithStatusCode(func(w http.ResponseWriter,
		r *http.Request) (int, error) {
		return 0, f(w, r)
	})
}

// ErrorHandlerWithStatusCode handles the error and responds it the client
// with the status code.
func ErrorHandlerWithStatusCode(f func(http.ResponseWriter, *http.Request) (int, error)) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if code, err := f(w, r); err != nil {
			if code == 0 {
				if _err, ok := err.(HTTPError); ok {
					code = _err.Code
				} else {
					code = http.StatusInternalServerError
				}
			}
			http.Error(w, err.Error(), code)
			ErrorLogFunc("Handling %q: status=%d, err=%v", r.RequestURI, code, err)
		}
	}
}

// HandlerWrapper handles the response result.
func HandlerWrapper(f func(http.ResponseWriter, *http.Request) (int, []byte, error)) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		code, resp, err := f(w, r)

		// Handle the error.
		if err != nil {
			if code == 0 {
				if _err, ok := err.(HTTPError); ok {
					code = _err.Code
				} else {
					code = http.StatusInternalServerError
				}
			}
			http.Error(w, err.Error(), code)
			ErrorLogFunc("Failed to handle %q: %s", r.RequestURI, err)
			return
		}

		// Determine the status code.
		if code == 0 {
			if len(resp) == 0 {
				code = http.StatusNoContent
			} else {
				code = http.StatusOK
			}
		}

		// Send the response result.
		w.WriteHeader(code)
		if _, err = io.CopyN(w, bytes.NewBuffer(resp), int64(len(resp))); err != nil {
			ErrorLogFunc("Failed to send the response of %q: %s", r.RequestURI, err)
		}
	}
}

// ListenAndServe is equal to http.ListenAndServe, but calling the method
// server.Shutdown(context.TODO()) to shutdown the HTTP server gracefully
// when calling lifecycle.Stop().
//
// Notice: It will call lifecycle.Stop() when the server exits.
func ListenAndServe(addr string, handler http.Handler) error {
	server := http.Server{Addr: addr, Handler: handler}
	lifecycle.Register(func() { server.Shutdown(context.TODO()) })
	err := server.ListenAndServe()
	ErrorLogFunc("The server listening on %s has an error: %s", addr, err)
	lifecycle.Stop()
	return err
}

// ListenAndServeTLS is equal to http.ListenAndServeTLS, but calling the method
// server.Shutdown(context.TODO()) to shutdown the HTTP server gracefully
// when calling lifecycle.Stop().
//
// Notice: It will call lifecycle.Stop() when the server exits.
func ListenAndServeTLS(addr, certFile, keyFile string, handler http.Handler) error {
	server := http.Server{Addr: addr, Handler: handler}
	lifecycle.Register(func() { server.Shutdown(context.TODO()) })
	err := server.ListenAndServeTLS(certFile, keyFile)
	ErrorLogFunc("The TLS server listening on %s has an error: %s", addr, err)
	lifecycle.Stop()
	return err
}

// GetBody returns the body of the HTTP request.
func GetBody(r *http.Request) (body []byte, err error) {
	buf := bytes.NewBuffer(nil)
	if _, err = io.CopyN(buf, r.Body, r.ContentLength); err != nil {
		return
	}
	return buf.Bytes(), nil
}
<|endoftext|>


<commit_before>package network

import (
    "fmt"
    "net"
    "strconv"
    "encoding/json"
    "time"
    "netgrok/buffer"
    . "netgrok/obj"
)

func listener(socket *net.UDPConn) (<-chan Message) {
    out := make(chan Message);
    go func () {

        b := make([]byte, 1024);
        for {
            n, addr, err := socket.ReadFromUDP(b);
            if err != nil {
                return;
            }
            if addr.String() != local_addr.String() && n > 0 {
                msg := Message{};
                err := json.Unmarshal(b[:n], &msg);
                if err != nil {
                    fmt.Println("Could not unmarshal message.");
                } else {
                    out <-msg;
                }
                //fmt.Println("Received message with code", msg.Code, "with body", msg.Body, "from", addr.String());

            }
        }
    }();
    return out;
}

func listening_manager(pop_channel chan<- Message, socket, broadcast_socket *net.UDPConn) (<-chan Message) {
    broadcast_in_channel := listener(broadcast_socket);
    tail_in_channel := listener(socket);

    out := make(chan Message);

    go func() {
        for {
            select {
            case msg := <-broadcast_in_channel:
                out <-msg;
            case msg := <-tail_in_channel:
                if /*len(msg.Signatures) != 0 && */msg.Signatures[0] == local_addr.String() {
                    pop_channel <-msg;
                } else {
                    out <-msg;
                }
            }
        }
    }();
    return out;
}

func send(msg Message, addr *net.UDPAddr) (error) {
    msg.Signatures = append(msg.Signatures, local_addr.String());
    b, err := json.Marshal(msg);
    if err != nil {
        fmt.Println("Could not marshal message.");
    }
    _, err = socket.WriteToUDP(b, addr);
    if err != nil {
        fmt.Println("Could not send!");
    } else {
        //fmt.Println("Sent message with code", msg.Code, "and body", msg.Body, "to:", addr.String());
    }
    return err;
}

func sending_manager(push_channel chan<- Message) (chan<- Message, chan<- Message, chan<- Message, chan<- *net.UDPAddr) {
    send_channel := make(chan Message);
    relay_channel := make(chan Message);
    broadcast_channel := make(chan Message);
    tail_channel := make(chan *net.UDPAddr);

    go func() {
        for {
            select {
            case msg := <-send_channel:
                push_channel <-msg;
                send(msg, head_addr);
            case msg := <-relay_channel:
                send(msg, head_addr);
            case msg := <-broadcast_channel:
                send(msg, broadcast_addr);
            case addr := <-tail_channel:
                b, _ := json.Marshal(local_addr);
                msg := NewMessage(TAIL_REQUEST, b);
                send(*msg, addr);
            }
        }
    }();
    return send_channel, relay_channel, broadcast_channel, tail_channel;
}

var local_addr, head_addr, broadcast_addr *net.UDPAddr;
var socket, broadcast_socket *net.UDPConn;

func Init(in_port, broadcast_in_port string) (chan<- Message, <-chan Message) {
    broadcast_addr, _ = net.ResolveUDPAddr("udp4", "255.255.255.255" + ":" + broadcast_in_port);

    temp_socket, err := net.DialUDP("udp4", nil, broadcast_addr);
    defer temp_socket.Close();
    temp_addr := temp_socket.LocalAddr();
    local_addr, err = net.ResolveUDPAddr("udp4", temp_addr.String());
    local_addr.Port, _ = strconv.Atoi(in_port);

    socket, _ = net.ListenUDP("udp4", local_addr);
    if err != nil {
        fmt.Println("Could not create socket.");
        return nil, nil;
    }
    broadcast_socket, _ := net.ListenUDP("udp", broadcast_addr);
    if err != nil {
        fmt.Println("Could not create broadcast socket.");
        socket.Close();
        return nil, nil;

    } else {
        fmt.Println("Sockets have been created.");
    }

    var tail_timeout_channel/*, cycle_timeout_channel*/ <-chan time.Time;

    push_channel, pop_channel := buffer.Init();
    rcv_channel := listening_manager(pop_channel, socket, broadcast_socket);
    send_channel, relay_channel, broadcast_channel, tail_channel := sending_manager(push_channel);

    to_network_channel := make(chan Message);
    from_network_channel := make(chan Message);
    go func() {
        for {
            if head_addr == nil {
                b, _ := json.Marshal(local_addr)
                broadcast_channel <-*NewMessage(HEAD_REQUEST, b);
                select {
                case <- time.After(4 * time.Second):
                    continue;
                case msg := <-rcv_channel:
                    switch msg.Code {
                    case TAIL_REQUEST:
                        var addr *net.UDPAddr;
                        _ = json.Unmarshal(msg.Body, &addr);
                        head_addr = addr;
                        conn := NewConnection(local_addr, addr);
                        b, _ := json.Marshal(conn);
                        send_channel <-*NewMessage(CONNECTION, b);
                    case HEAD_REQUEST:
                        var addr *net.UDPAddr;
                        _ = json.Unmarshal(msg.Body, &addr);
                        tail_channel <-addr;
                    }
                }
            } else {
                select {
                case msg := <-to_network_channel:
                    send_channel <-msg;
                case msg :=  <-rcv_channel:
                    switch msg.Code {
                    case KEEP_ALIVE:
                        relay_channel <-msg;
                    case CONNECTION:
                        var conn Connection;
                        err := json.Unmarshal(msg.Body, &conn);
                        if err != nil {
                            fmt.Println("Could not unmarshal connection.");
                        } else {
                            if head_addr == nil || head_addr.String() == conn.To.String() {
                               head_addr = conn.From;
                            }
                            relay_channel <-msg;
                        }
                    case HEAD_REQUEST:
                        var addr *net.UDPAddr;
                        err := json.Unmarshal(msg.Body, &addr);
                        if err != nil {
                            fmt.Println("Could not unmarshal message.")
                        } else {
                            tail_channel <-addr;
                        }
                    case TAIL_REQUEST:
                        break;
                    case TAIL_DEAD:
                        time.Sleep(1 * time.Second);
                        relay_channel <-msg;
                        head_addr = nil;
                    default:
                        from_network_channel <-msg;
                        relay_channel <-msg;
                    }
                    tail_timeout_channel = nil;
                case <- time.After(1 * time.Second):
                    send_channel <-*NewMessage(KEEP_ALIVE, []byte{});
                    if tail_timeout_channel == nil {
                        tail_timeout_channel = time.After(2 * time.Second);
                    }
                case <-tail_timeout_channel:
                    send_channel <-*NewMessage(TAIL_DEAD, []byte{});
                    head_addr = nil;
                }
            }
        }
    }();
    return to_network_channel, from_network_channel;
}<commit_msg>Less keep alive spam.<commit_after>package network

import (
    "fmt"
    "net"
    "strconv"
    "encoding/json"
    "time"
    "netgrok/buffer"
    . "netgrok/obj"
)

func listener(socket *net.UDPConn) (<-chan Message) {
    out := make(chan Message);
    go func () {

        b := make([]byte, 1024);
        for {
            n, addr, err := socket.ReadFromUDP(b);
            if err != nil {
                return;
            }
            if addr.String() != local_addr.String() && n > 0 {
                msg := Message{};
                err := json.Unmarshal(b[:n], &msg);
                if err != nil {
                    fmt.Println("Could not unmarshal message.");
                } else {
                    out <-msg;
                }
                //fmt.Println("Received message with code", msg.Code, "with body", msg.Body, "from", addr.String());

            }
        }
    }();
    return out;
}

func listening_manager(pop_channel chan<- Message, socket, broadcast_socket *net.UDPConn) (<-chan Message) {
    broadcast_in_channel := listener(broadcast_socket);
    tail_in_channel := listener(socket);

    out := make(chan Message);

    go func() {
        for {
            select {
            case msg := <-broadcast_in_channel:
                out <-msg;
            case msg := <-tail_in_channel:
                if msg.Code != KEEP_ALIVE && msg.Signatures[0] == local_addr.String() {
                    pop_channel <-msg;
                } else {
                    out <-msg;
                }
            }
        }
    }();
    return out;
}

func send(msg Message, addr *net.UDPAddr) (error) {
    msg.Signatures = append(msg.Signatures, local_addr.String());
    b, err := json.Marshal(msg);
    if err != nil {
        fmt.Println("Could not marshal message.");
    }
    _, err = socket.WriteToUDP(b, addr);
    if err != nil {
        fmt.Println("Could not send!");
    } else {
        //fmt.Println("Sent message with code", msg.Code, "and body", msg.Body, "to:", addr.String());
    }
    return err;
}

func sending_manager(push_channel chan<- Message) (chan<- Message, chan<- Message, chan<- Message, chan<- *net.UDPAddr) {
    send_channel := make(chan Message);
    relay_channel := make(chan Message);
    broadcast_channel := make(chan Message);
    tail_channel := make(chan *net.UDPAddr);

    go func() {
        for {
            select {
            case msg := <-send_channel:
                push_channel <-msg;
                send(msg, head_addr);
            case msg := <-relay_channel:
                send(msg, head_addr);
            case msg := <-broadcast_channel:
                send(msg, broadcast_addr);
            case addr := <-tail_channel:
                b, _ := json.Marshal(local_addr);
                msg := NewMessage(TAIL_REQUEST, b);
                send(*msg, addr);
            }
        }
    }();
    return send_channel, relay_channel, broadcast_channel, tail_channel;
}

var local_addr, head_addr, broadcast_addr *net.UDPAddr;
var socket, broadcast_socket *net.UDPConn;

func Init(in_port, broadcast_in_port string) (chan<- Message, <-chan Message) {
    broadcast_addr, _ = net.ResolveUDPAddr("udp4", "255.255.255.255" + ":" + broadcast_in_port);

    temp_socket, err := net.DialUDP("udp4", nil, broadcast_addr);
    defer temp_socket.Close();
    temp_addr := temp_socket.LocalAddr();
    local_addr, err = net.ResolveUDPAddr("udp4", temp_addr.String());
    local_addr.Port, _ = strconv.Atoi(in_port);

    socket, _ = net.ListenUDP("udp4", local_addr);
    if err != nil {
        fmt.Println("Could not create socket.");
        return nil, nil;
    }
    broadcast_socket, _ := net.ListenUDP("udp", broadcast_addr);
    if err != nil {
        fmt.Println("Could not create broadcast socket.");
        socket.Close();
        return nil, nil;

    } else {
        fmt.Println("Sockets have been created.");
    }

    var tail_timeout_channel/*, cycle_timeout_channel*/ <-chan time.Time;

    push_channel, pop_channel := buffer.Init();
    rcv_channel := listening_manager(pop_channel, socket, broadcast_socket);
    send_channel, relay_channel, broadcast_channel, tail_channel := sending_manager(push_channel);

    to_network_channel := make(chan Message);
    from_network_channel := make(chan Message);
    go func() {
        for {
            if head_addr == nil {
                b, _ := json.Marshal(local_addr)
                broadcast_channel <-*NewMessage(HEAD_REQUEST, b);
                select {
                case <- time.After(4 * time.Second):
                    continue;
                case msg := <-rcv_channel:
                    switch msg.Code {
                    case TAIL_REQUEST:
                        var addr *net.UDPAddr;
                        _ = json.Unmarshal(msg.Body, &addr);
                        head_addr = addr;
                        conn := NewConnection(local_addr, addr);
                        b, _ := json.Marshal(conn);
                        send_channel <-*NewMessage(CONNECTION, b);
                    case HEAD_REQUEST:
                        var addr *net.UDPAddr;
                        _ = json.Unmarshal(msg.Body, &addr);
                        tail_channel <-addr;
                    }
                }
            } else {
                select {
                case msg := <-to_network_channel:
                    send_channel <-msg;
                case msg :=  <-rcv_channel:
                    switch msg.Code {
                    case KEEP_ALIVE:
                        break;//relay_channel <-msg;
                    case CONNECTION:
                        var conn Connection;
                        err := json.Unmarshal(msg.Body, &conn);
                        if err != nil {
                            fmt.Println("Could not unmarshal connection.");
                        } else {
                            if head_addr == nil || head_addr.String() == conn.To.String() {
                               head_addr = conn.From;
                            }
                            relay_channel <-msg;
                        }
                    case HEAD_REQUEST:
                        var addr *net.UDPAddr;
                        err := json.Unmarshal(msg.Body, &addr);
                        if err != nil {
                            fmt.Println("Could not unmarshal message.")
                        } else {
                            tail_channel <-addr;
                        }
                    case TAIL_REQUEST:
                        break;
                    case TAIL_DEAD:
                        time.Sleep(1 * time.Second);
                        relay_channel <-msg;
                        head_addr = nil;
                    default:
                        from_network_channel <-msg;
                        relay_channel <-msg;
                    }
                    tail_timeout_channel = nil;
                case <- time.After(1 * time.Second):
                    send_channel <-*NewMessage(KEEP_ALIVE, []byte{});
                    if tail_timeout_channel == nil {
                        tail_timeout_channel = time.After(2 * time.Second);
                    }
                case <-tail_timeout_channel:
                    send_channel <-*NewMessage(TAIL_DEAD, []byte{});
                    head_addr = nil;
                }
            }
        }
    }();
    return to_network_channel, from_network_channel;
}<|endoftext|>


<commit_before>package sparta

import (
	"fmt"
	"github.com/Sirupsen/logrus"
	"github.com/asaskevich/govalidator"
	"github.com/spf13/cobra"
	"os"
	"path"
	"time"
)

// CommandLineOptions defines the commands available via the Sparta command
// line interface.  Embedding applications can extend existing commands
// and add their own to the `Root` command.  See https://github.com/spf13/cobra
// for more information.
var CommandLineOptions = struct {
	Root      *cobra.Command
	Version   *cobra.Command
	Provision *cobra.Command
	Delete    *cobra.Command
	Execute   *cobra.Command
	Describe  *cobra.Command
	Explore   *cobra.Command
}{}

/******************************************************************************/
// Global options
type optionsGlobalStruct struct {
	Noop     bool           `valid:"-"`
	LogLevel string         `valid:"matches(panic|fatal|error|warn|info|debug])"`
	Logger   *logrus.Logger `valid:"-"`
	Command  string         `valid:"-"`
}

// OptionsGlobal stores the global command line options
var OptionsGlobal optionsGlobalStruct

/******************************************************************************/
// Provision options
// Ref: http://docs.aws.amazon.com/AmazonS3/latest/dev/BucketRestrictions.html
type optionsProvisionStruct struct {
	S3Bucket string `valid:"required,matches(\\w+)"`
}

var optionsProvision optionsProvisionStruct

/******************************************************************************/
// Execute options
type optionsExecuteStruct struct {
	Port            int `valid:"-"`
	SignalParentPID int `valid:"-"`
}

var optionsExecute optionsExecuteStruct

/******************************************************************************/
// Describe options
type optionsDescribeStruct struct {
	OutputFile string `valid:"required"`
}

var optionsDescribe optionsDescribeStruct

/******************************************************************************/
// Explore options
type optionsExploreStruct struct {
	Port int `valid:"-"`
}

var optionsExplore optionsExploreStruct

/******************************************************************************/
// Initialization
// Initialize all the Cobra commands and their associated flags
/******************************************************************************/
func init() {
	// Root
	CommandLineOptions.Root = &cobra.Command{
		Use:   path.Base(os.Args[0]),
		Short: "Sparta-powered AWS Lambda microservice",
	}
	CommandLineOptions.Root.PersistentFlags().BoolVarP(&OptionsGlobal.Noop, "noop",
		"n",
		false,
		"Dry-run behavior only (do not perform mutations)")
	CommandLineOptions.Root.PersistentFlags().StringVarP(&OptionsGlobal.LogLevel,
		"level",
		"l",
		"info",
		"Log level [panic, fatal, error, warn, info, debug]")
	// Version
	CommandLineOptions.Version = &cobra.Command{
		Use:   "version",
		Short: "Sparta framework version",
		Long:  `Displays the Sparta framework version `,
		Run: func(cmd *cobra.Command, args []string) {

		},
	}
	// Provision
	CommandLineOptions.Provision = &cobra.Command{
		Use:   "provision",
		Short: "Provision service",
		Long:  `Provision the service (either create or update) via CloudFormation`,
	}
	CommandLineOptions.Provision.Flags().StringVarP(&optionsProvision.S3Bucket,
		"s3Bucket",
		"s",
		"",
		"S3 Bucket to use for Lambda source")

	// Delete
	CommandLineOptions.Delete = &cobra.Command{
		Use:   "delete",
		Short: "Delete service",
		Long:  `Ensure service is successfully deleted`,
	}

	// Execute
	CommandLineOptions.Execute = &cobra.Command{
		Use:   "execute",
		Short: "Execute",
		Long:  `Startup the localhost HTTP server to handle requests`,
	}
	CommandLineOptions.Execute.Flags().IntVarP(&optionsExecute.Port,
		"port",
		"p",
		9999,
		"Alternative port for HTTP binding (default=9999)")
	CommandLineOptions.Execute.Flags().IntVarP(&optionsExecute.SignalParentPID,
		"signal",
		"s",
		0,
		"Process ID to signal with SIGUSR2 once ready")

	// Describe
	CommandLineOptions.Describe = &cobra.Command{
		Use:   "describe",
		Short: "Describe service",
		Long:  `Produce an HTML report of the service`,
	}

	CommandLineOptions.Describe.Flags().StringVarP(&optionsDescribe.OutputFile,
		"out",
		"o",
		"",
		"Output file for HTML description")

	// Explore
	CommandLineOptions.Explore = &cobra.Command{
		Use:   "explore",
		Short: "Interactively explore service",
		Long:  `Startup a localhost HTTP server to explore the exported Go functions`,
	}
	CommandLineOptions.Explore.Flags().IntVarP(&optionsExplore.Port,
		"port",
		"p",
		9999,
		"Alternative port for HTTP binding (default=9999)")
}

// CommandLineOptionsHook allows embedding applications the ability
// to validate caller-defined command line arguments.  Return an error
// if the command line fails.
type CommandLineOptionsHook func(command *cobra.Command) error

// ParseOptions the command line options
func ParseOptions(handler CommandLineOptionsHook) error {
	var noop bool
	var level string

	// First up, create a dummy Root command for the parse...
	var parseCmdRoot = &cobra.Command{
		Use:           CommandLineOptions.Root.Use,
		Short:         CommandLineOptions.Root.Short,
		SilenceUsage:  true,
		SilenceErrors: false,
		RunE: func(cmd *cobra.Command, args []string) error {
			return nil
		},
	}

	parseCmdRoot.PersistentFlags().BoolVarP(&noop, "noop",
		"n",
		false,
		"Dry-run behavior only (do not perform mutations)")
	parseCmdRoot.PersistentFlags().StringVarP(&level,
		"level",
		"l",
		"info",
		"Log level [panic, fatal, error, warn, info, debug]")

	// Now, for any user-attached commands, add them to the temporary Parse
	// root command.
	for _, eachUserCommand := range CommandLineOptions.Root.Commands() {
		userProxyCmd := &cobra.Command{
			Use:   eachUserCommand.Use,
			Short: eachUserCommand.Short,
		}
		userProxyCmd.PreRunE = func(cmd *cobra.Command, args []string) error {
			_, validateErr := govalidator.ValidateStruct(OptionsGlobal)
			if nil != validateErr {
				return validateErr
			}
			logger, loggerErr := NewLogger(OptionsGlobal.LogLevel)
			if nil != loggerErr {
				return loggerErr
			}
			OptionsGlobal.Logger = logger

			if handler != nil {
				return handler(userProxyCmd)
			}
			return nil
		}
		userProxyCmd.Flags().AddFlagSet(eachUserCommand.Flags())
		parseCmdRoot.AddCommand(userProxyCmd)
	}

	//////////////////////////////////////////////////////////////////////////////
	// Then add the standard Sparta ones...

	spartaCommands := []*cobra.Command{
		CommandLineOptions.Version,
		CommandLineOptions.Provision,
		CommandLineOptions.Delete,
		CommandLineOptions.Execute,
		CommandLineOptions.Describe,
		CommandLineOptions.Explore,
	}
	CommandLineOptions.Version.PreRunE = func(cmd *cobra.Command, args []string) error {
		if handler != nil {
			return handler(CommandLineOptions.Version)
		}
		return nil
	}
	parseCmdRoot.AddCommand(CommandLineOptions.Version)

	CommandLineOptions.Provision.PreRunE = func(cmd *cobra.Command, args []string) error {
		if handler != nil {
			return handler(CommandLineOptions.Provision)
		}
		return nil
	}
	parseCmdRoot.AddCommand(CommandLineOptions.Provision)

	CommandLineOptions.Delete.PreRunE = func(cmd *cobra.Command, args []string) error {
		if handler != nil {
			return handler(CommandLineOptions.Delete)
		}
		return nil
	}
	parseCmdRoot.AddCommand(CommandLineOptions.Delete)

	CommandLineOptions.Execute.PreRunE = func(cmd *cobra.Command, args []string) error {
		if handler != nil {
			return handler(CommandLineOptions.Execute)
		}
		return nil
	}
	parseCmdRoot.AddCommand(CommandLineOptions.Execute)

	CommandLineOptions.Describe.PreRunE = func(cmd *cobra.Command, args []string) error {
		if handler != nil {
			return handler(CommandLineOptions.Describe)
		}
		return nil
	}
	parseCmdRoot.AddCommand(CommandLineOptions.Describe)

	CommandLineOptions.Explore.PreRunE = func(cmd *cobra.Command, args []string) error {
		if handler != nil {
			return handler(CommandLineOptions.Explore)
		}
		return nil
	}
	parseCmdRoot.AddCommand(CommandLineOptions.Explore)

	// Assign each command an empty RunE func s.t.
	// Cobra doesn't print out the command info
	for _, eachCommand := range parseCmdRoot.Commands() {
		eachCommand.RunE = func(cmd *cobra.Command, args []string) error {
			return nil
		}
	}
	// Intercept the usage command - we'll end up showing this later
	// in Main...If there is an error, we will show help there...
	parseCmdRoot.SetHelpFunc(func(*cobra.Command, []string) {
		// Swallow help here
	})

	// Run it...
	executeErr := parseCmdRoot.Execute()

	// Cleanup the Sparta specific ones
	for _, eachCmd := range spartaCommands {
		eachCmd.RunE = nil
		eachCmd.PreRunE = nil
	}

	if nil != executeErr {
		parseCmdRoot.SetHelpFunc(nil)
		parseCmdRoot.Root().Help()
	}
	return executeErr
}

// Main defines the primary handler for transforming an application into a Sparta package.  The
// serviceName is used to uniquely identify your service within a region and will
// be used for subsequent updates.  For provisioning, ensure that you've
// properly configured AWS credentials for the golang SDK.
// See http://docs.aws.amazon.com/sdk-for-go/api/aws/defaults.html#DefaultChainCredentials-constant
// for more information.
func Main(serviceName string, serviceDescription string, lambdaAWSInfos []*LambdaAWSInfo, api *API, site *S3Site) error {
	//////////////////////////////////////////////////////////////////////////////
	// cmdRoot defines the root, non-executable command
	CommandLineOptions.Root.Short = fmt.Sprintf("%s - Sparta powered AWS Lambda Microservice", serviceName)
	CommandLineOptions.Root.Long = serviceDescription
	CommandLineOptions.Root.PersistentPreRunE = func(cmd *cobra.Command, args []string) error {

		_, validateErr := govalidator.ValidateStruct(OptionsGlobal)
		if nil != validateErr {
			return validateErr
		}
		logger, loggerErr := NewLogger(OptionsGlobal.LogLevel)
		if nil != loggerErr {
			return loggerErr
		}
		OptionsGlobal.Logger = logger

		logger.WithFields(logrus.Fields{
			"Option":        cmd.Name(),
			"SpartaVersion": SpartaVersion,
			"UTC":           (time.Now().UTC().Format(time.RFC3339)),
		}).Info("Welcome to " + serviceName)
		return nil
	}

	//////////////////////////////////////////////////////////////////////////////
	// Version
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Version)

	//////////////////////////////////////////////////////////////////////////////
	// Provision
	CommandLineOptions.Provision.PreRunE = func(cmd *cobra.Command, args []string) error {
		validationResults, validateErr := govalidator.ValidateStruct(optionsProvision)

		OptionsGlobal.Logger.WithFields(logrus.Fields{
			"validationResults": validationResults,
			"validateErr":       validateErr,
			"optionsProvision":  optionsProvision,
		}).Debug("Provision Validation")
		return validateErr
	}

	if nil == CommandLineOptions.Provision.RunE {
		CommandLineOptions.Provision.RunE = func(cmd *cobra.Command, args []string) error {
			return Provision(OptionsGlobal.Noop,
				serviceName,
				serviceDescription,
				lambdaAWSInfos,
				api,
				site,
				optionsProvision.S3Bucket,
				nil,
				OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Provision)

	//////////////////////////////////////////////////////////////////////////////
	// Delete
	CommandLineOptions.Delete.RunE = func(cmd *cobra.Command, args []string) error {
		return Delete(serviceName, OptionsGlobal.Logger)
	}

	CommandLineOptions.Root.AddCommand(CommandLineOptions.Delete)

	//////////////////////////////////////////////////////////////////////////////
	// Execute
	if nil == CommandLineOptions.Execute.RunE {
		CommandLineOptions.Execute.RunE = func(cmd *cobra.Command, args []string) error {
			_, validateErr := govalidator.ValidateStruct(optionsExecute)
			if nil != validateErr {
				return validateErr
			}

			OptionsGlobal.Logger.Formatter = new(logrus.JSONFormatter)
			return Execute(lambdaAWSInfos,
				optionsExecute.Port,
				optionsExecute.SignalParentPID,
				OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Execute)

	//////////////////////////////////////////////////////////////////////////////
	// Describe
	if nil == CommandLineOptions.Describe.RunE {
		CommandLineOptions.Describe.RunE = func(cmd *cobra.Command, args []string) error {
			_, validateErr := govalidator.ValidateStruct(optionsDescribe)
			if nil != validateErr {
				return validateErr
			}

			fileWriter, fileWriterErr := os.Create(optionsDescribe.OutputFile)
			if fileWriterErr != nil {
				return fileWriterErr
			}
			defer fileWriter.Close()
			describeErr := Describe(serviceName, serviceDescription, lambdaAWSInfos, api, site, fileWriter, OptionsGlobal.Logger)
			if describeErr == nil {
				describeErr = fileWriter.Sync()
			}
			return describeErr
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Describe)

	//////////////////////////////////////////////////////////////////////////////
	// Explore
	if nil == CommandLineOptions.Explore.RunE {
		CommandLineOptions.Explore.RunE = func(cmd *cobra.Command, args []string) error {
			_, validateErr := govalidator.ValidateStruct(optionsExplore)
			if nil != validateErr {
				return validateErr
			}
			return Explore(lambdaAWSInfos, optionsExplore.Port, OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Explore)

	// Run it!
	executeErr := CommandLineOptions.Root.Execute()
	if nil != OptionsGlobal.Logger && nil != executeErr {
		OptionsGlobal.Logger.Error(executeErr)
	}
	// Cleanup, if for some reason the caller wants to re-execute later...
	CommandLineOptions.Root.PersistentPreRunE = nil
	return executeErr
}
<commit_msg>Include Go version in version output<commit_after>package sparta

import (
	"fmt"
	"github.com/Sirupsen/logrus"
	"github.com/asaskevich/govalidator"
	"github.com/spf13/cobra"
	"os"
	"path"
	"runtime"
	"time"
)

// CommandLineOptions defines the commands available via the Sparta command
// line interface.  Embedding applications can extend existing commands
// and add their own to the `Root` command.  See https://github.com/spf13/cobra
// for more information.
var CommandLineOptions = struct {
	Root      *cobra.Command
	Version   *cobra.Command
	Provision *cobra.Command
	Delete    *cobra.Command
	Execute   *cobra.Command
	Describe  *cobra.Command
	Explore   *cobra.Command
}{}

/******************************************************************************/
// Global options
type optionsGlobalStruct struct {
	Noop     bool           `valid:"-"`
	LogLevel string         `valid:"matches(panic|fatal|error|warn|info|debug])"`
	Logger   *logrus.Logger `valid:"-"`
	Command  string         `valid:"-"`
}

// OptionsGlobal stores the global command line options
var OptionsGlobal optionsGlobalStruct

/******************************************************************************/
// Provision options
// Ref: http://docs.aws.amazon.com/AmazonS3/latest/dev/BucketRestrictions.html
type optionsProvisionStruct struct {
	S3Bucket string `valid:"required,matches(\\w+)"`
}

var optionsProvision optionsProvisionStruct

/******************************************************************************/
// Execute options
type optionsExecuteStruct struct {
	Port            int `valid:"-"`
	SignalParentPID int `valid:"-"`
}

var optionsExecute optionsExecuteStruct

/******************************************************************************/
// Describe options
type optionsDescribeStruct struct {
	OutputFile string `valid:"required"`
}

var optionsDescribe optionsDescribeStruct

/******************************************************************************/
// Explore options
type optionsExploreStruct struct {
	Port int `valid:"-"`
}

var optionsExplore optionsExploreStruct

/******************************************************************************/
// Initialization
// Initialize all the Cobra commands and their associated flags
/******************************************************************************/
func init() {
	// Root
	CommandLineOptions.Root = &cobra.Command{
		Use:   path.Base(os.Args[0]),
		Short: "Sparta-powered AWS Lambda microservice",
	}
	CommandLineOptions.Root.PersistentFlags().BoolVarP(&OptionsGlobal.Noop, "noop",
		"n",
		false,
		"Dry-run behavior only (do not perform mutations)")
	CommandLineOptions.Root.PersistentFlags().StringVarP(&OptionsGlobal.LogLevel,
		"level",
		"l",
		"info",
		"Log level [panic, fatal, error, warn, info, debug]")
	// Version
	CommandLineOptions.Version = &cobra.Command{
		Use:   "version",
		Short: "Sparta framework version",
		Long:  `Displays the Sparta framework version `,
		Run: func(cmd *cobra.Command, args []string) {

		},
	}
	// Provision
	CommandLineOptions.Provision = &cobra.Command{
		Use:   "provision",
		Short: "Provision service",
		Long:  `Provision the service (either create or update) via CloudFormation`,
	}
	CommandLineOptions.Provision.Flags().StringVarP(&optionsProvision.S3Bucket,
		"s3Bucket",
		"s",
		"",
		"S3 Bucket to use for Lambda source")

	// Delete
	CommandLineOptions.Delete = &cobra.Command{
		Use:   "delete",
		Short: "Delete service",
		Long:  `Ensure service is successfully deleted`,
	}

	// Execute
	CommandLineOptions.Execute = &cobra.Command{
		Use:   "execute",
		Short: "Execute",
		Long:  `Startup the localhost HTTP server to handle requests`,
	}
	CommandLineOptions.Execute.Flags().IntVarP(&optionsExecute.Port,
		"port",
		"p",
		9999,
		"Alternative port for HTTP binding (default=9999)")
	CommandLineOptions.Execute.Flags().IntVarP(&optionsExecute.SignalParentPID,
		"signal",
		"s",
		0,
		"Process ID to signal with SIGUSR2 once ready")

	// Describe
	CommandLineOptions.Describe = &cobra.Command{
		Use:   "describe",
		Short: "Describe service",
		Long:  `Produce an HTML report of the service`,
	}

	CommandLineOptions.Describe.Flags().StringVarP(&optionsDescribe.OutputFile,
		"out",
		"o",
		"",
		"Output file for HTML description")

	// Explore
	CommandLineOptions.Explore = &cobra.Command{
		Use:   "explore",
		Short: "Interactively explore service",
		Long:  `Startup a localhost HTTP server to explore the exported Go functions`,
	}
	CommandLineOptions.Explore.Flags().IntVarP(&optionsExplore.Port,
		"port",
		"p",
		9999,
		"Alternative port for HTTP binding (default=9999)")
}

// CommandLineOptionsHook allows embedding applications the ability
// to validate caller-defined command line arguments.  Return an error
// if the command line fails.
type CommandLineOptionsHook func(command *cobra.Command) error

// ParseOptions the command line options
func ParseOptions(handler CommandLineOptionsHook) error {
	var noop bool
	var level string

	// First up, create a dummy Root command for the parse...
	var parseCmdRoot = &cobra.Command{
		Use:           CommandLineOptions.Root.Use,
		Short:         CommandLineOptions.Root.Short,
		SilenceUsage:  true,
		SilenceErrors: false,
		RunE: func(cmd *cobra.Command, args []string) error {
			return nil
		},
	}

	parseCmdRoot.PersistentFlags().BoolVarP(&noop, "noop",
		"n",
		false,
		"Dry-run behavior only (do not perform mutations)")
	parseCmdRoot.PersistentFlags().StringVarP(&level,
		"level",
		"l",
		"info",
		"Log level [panic, fatal, error, warn, info, debug]")

	// Now, for any user-attached commands, add them to the temporary Parse
	// root command.
	for _, eachUserCommand := range CommandLineOptions.Root.Commands() {
		userProxyCmd := &cobra.Command{
			Use:   eachUserCommand.Use,
			Short: eachUserCommand.Short,
		}
		userProxyCmd.PreRunE = func(cmd *cobra.Command, args []string) error {
			_, validateErr := govalidator.ValidateStruct(OptionsGlobal)
			if nil != validateErr {
				return validateErr
			}
			logger, loggerErr := NewLogger(OptionsGlobal.LogLevel)
			if nil != loggerErr {
				return loggerErr
			}
			OptionsGlobal.Logger = logger

			if handler != nil {
				return handler(userProxyCmd)
			}
			return nil
		}
		userProxyCmd.Flags().AddFlagSet(eachUserCommand.Flags())
		parseCmdRoot.AddCommand(userProxyCmd)
	}

	//////////////////////////////////////////////////////////////////////////////
	// Then add the standard Sparta ones...

	spartaCommands := []*cobra.Command{
		CommandLineOptions.Version,
		CommandLineOptions.Provision,
		CommandLineOptions.Delete,
		CommandLineOptions.Execute,
		CommandLineOptions.Describe,
		CommandLineOptions.Explore,
	}
	CommandLineOptions.Version.PreRunE = func(cmd *cobra.Command, args []string) error {
		if handler != nil {
			return handler(CommandLineOptions.Version)
		}
		return nil
	}
	parseCmdRoot.AddCommand(CommandLineOptions.Version)

	CommandLineOptions.Provision.PreRunE = func(cmd *cobra.Command, args []string) error {
		if handler != nil {
			return handler(CommandLineOptions.Provision)
		}
		return nil
	}
	parseCmdRoot.AddCommand(CommandLineOptions.Provision)

	CommandLineOptions.Delete.PreRunE = func(cmd *cobra.Command, args []string) error {
		if handler != nil {
			return handler(CommandLineOptions.Delete)
		}
		return nil
	}
	parseCmdRoot.AddCommand(CommandLineOptions.Delete)

	CommandLineOptions.Execute.PreRunE = func(cmd *cobra.Command, args []string) error {
		if handler != nil {
			return handler(CommandLineOptions.Execute)
		}
		return nil
	}
	parseCmdRoot.AddCommand(CommandLineOptions.Execute)

	CommandLineOptions.Describe.PreRunE = func(cmd *cobra.Command, args []string) error {
		if handler != nil {
			return handler(CommandLineOptions.Describe)
		}
		return nil
	}
	parseCmdRoot.AddCommand(CommandLineOptions.Describe)

	CommandLineOptions.Explore.PreRunE = func(cmd *cobra.Command, args []string) error {
		if handler != nil {
			return handler(CommandLineOptions.Explore)
		}
		return nil
	}
	parseCmdRoot.AddCommand(CommandLineOptions.Explore)

	// Assign each command an empty RunE func s.t.
	// Cobra doesn't print out the command info
	for _, eachCommand := range parseCmdRoot.Commands() {
		eachCommand.RunE = func(cmd *cobra.Command, args []string) error {
			return nil
		}
	}
	// Intercept the usage command - we'll end up showing this later
	// in Main...If there is an error, we will show help there...
	parseCmdRoot.SetHelpFunc(func(*cobra.Command, []string) {
		// Swallow help here
	})

	// Run it...
	executeErr := parseCmdRoot.Execute()

	// Cleanup the Sparta specific ones
	for _, eachCmd := range spartaCommands {
		eachCmd.RunE = nil
		eachCmd.PreRunE = nil
	}

	if nil != executeErr {
		parseCmdRoot.SetHelpFunc(nil)
		parseCmdRoot.Root().Help()
	}
	return executeErr
}

// Main defines the primary handler for transforming an application into a Sparta package.  The
// serviceName is used to uniquely identify your service within a region and will
// be used for subsequent updates.  For provisioning, ensure that you've
// properly configured AWS credentials for the golang SDK.
// See http://docs.aws.amazon.com/sdk-for-go/api/aws/defaults.html#DefaultChainCredentials-constant
// for more information.
func Main(serviceName string, serviceDescription string, lambdaAWSInfos []*LambdaAWSInfo, api *API, site *S3Site) error {
	//////////////////////////////////////////////////////////////////////////////
	// cmdRoot defines the root, non-executable command
	CommandLineOptions.Root.Short = fmt.Sprintf("%s - Sparta powered AWS Lambda Microservice", serviceName)
	CommandLineOptions.Root.Long = serviceDescription
	CommandLineOptions.Root.PersistentPreRunE = func(cmd *cobra.Command, args []string) error {

		_, validateErr := govalidator.ValidateStruct(OptionsGlobal)
		if nil != validateErr {
			return validateErr
		}
		logger, loggerErr := NewLogger(OptionsGlobal.LogLevel)
		if nil != loggerErr {
			return loggerErr
		}
		OptionsGlobal.Logger = logger

		logger.WithFields(logrus.Fields{
			"Option":        cmd.Name(),
			"SpartaVersion": SpartaVersion,
			"Go":            runtime.Version(),
			"UTC":           (time.Now().UTC().Format(time.RFC3339)),
		}).Info("Welcome to " + serviceName)
		return nil
	}

	//////////////////////////////////////////////////////////////////////////////
	// Version
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Version)

	//////////////////////////////////////////////////////////////////////////////
	// Provision
	CommandLineOptions.Provision.PreRunE = func(cmd *cobra.Command, args []string) error {
		validationResults, validateErr := govalidator.ValidateStruct(optionsProvision)

		OptionsGlobal.Logger.WithFields(logrus.Fields{
			"validationResults": validationResults,
			"validateErr":       validateErr,
			"optionsProvision":  optionsProvision,
		}).Debug("Provision Validation")
		return validateErr
	}

	if nil == CommandLineOptions.Provision.RunE {
		CommandLineOptions.Provision.RunE = func(cmd *cobra.Command, args []string) error {
			return Provision(OptionsGlobal.Noop,
				serviceName,
				serviceDescription,
				lambdaAWSInfos,
				api,
				site,
				optionsProvision.S3Bucket,
				nil,
				OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Provision)

	//////////////////////////////////////////////////////////////////////////////
	// Delete
	CommandLineOptions.Delete.RunE = func(cmd *cobra.Command, args []string) error {
		return Delete(serviceName, OptionsGlobal.Logger)
	}

	CommandLineOptions.Root.AddCommand(CommandLineOptions.Delete)

	//////////////////////////////////////////////////////////////////////////////
	// Execute
	if nil == CommandLineOptions.Execute.RunE {
		CommandLineOptions.Execute.RunE = func(cmd *cobra.Command, args []string) error {
			_, validateErr := govalidator.ValidateStruct(optionsExecute)
			if nil != validateErr {
				return validateErr
			}

			OptionsGlobal.Logger.Formatter = new(logrus.JSONFormatter)
			return Execute(lambdaAWSInfos,
				optionsExecute.Port,
				optionsExecute.SignalParentPID,
				OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Execute)

	//////////////////////////////////////////////////////////////////////////////
	// Describe
	if nil == CommandLineOptions.Describe.RunE {
		CommandLineOptions.Describe.RunE = func(cmd *cobra.Command, args []string) error {
			_, validateErr := govalidator.ValidateStruct(optionsDescribe)
			if nil != validateErr {
				return validateErr
			}

			fileWriter, fileWriterErr := os.Create(optionsDescribe.OutputFile)
			if fileWriterErr != nil {
				return fileWriterErr
			}
			defer fileWriter.Close()
			describeErr := Describe(serviceName, serviceDescription, lambdaAWSInfos, api, site, fileWriter, OptionsGlobal.Logger)
			if describeErr == nil {
				describeErr = fileWriter.Sync()
			}
			return describeErr
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Describe)

	//////////////////////////////////////////////////////////////////////////////
	// Explore
	if nil == CommandLineOptions.Explore.RunE {
		CommandLineOptions.Explore.RunE = func(cmd *cobra.Command, args []string) error {
			_, validateErr := govalidator.ValidateStruct(optionsExplore)
			if nil != validateErr {
				return validateErr
			}
			return Explore(lambdaAWSInfos, optionsExplore.Port, OptionsGlobal.Logger)
		}
	}
	CommandLineOptions.Root.AddCommand(CommandLineOptions.Explore)

	// Run it!
	executeErr := CommandLineOptions.Root.Execute()
	if nil != OptionsGlobal.Logger && nil != executeErr {
		OptionsGlobal.Logger.Error(executeErr)
	}
	// Cleanup, if for some reason the caller wants to re-execute later...
	CommandLineOptions.Root.PersistentPreRunE = nil
	return executeErr
}
<|endoftext|>


<commit_before>// package network implements micro network node
package network

import (
	"os"
	"strings"
	"time"

	"github.com/micro/cli"
	"github.com/micro/go-micro"
	"github.com/micro/go-micro/network"
	"github.com/micro/go-micro/network/handler"
	"github.com/micro/go-micro/network/resolver"
	"github.com/micro/go-micro/network/resolver/dns"
	"github.com/micro/go-micro/network/resolver/http"
	"github.com/micro/go-micro/network/resolver/registry"
	"github.com/micro/go-micro/proxy"
	"github.com/micro/go-micro/proxy/mucp"
	"github.com/micro/go-micro/router"
	"github.com/micro/go-micro/server"
	"github.com/micro/go-micro/tunnel"
	"github.com/micro/go-micro/util/log"
	"github.com/micro/go-micro/util/mux"
)

var (
	// Name of the network service
	Name = "go.micro.network"
	// Name of the micro network
	Network = "go.micro"
	// Address is the network address
	Address = ":8085"
	// Resolver is the network resolver
	Resolver = "registry"
)

// run runs the micro server
func run(ctx *cli.Context, srvOpts ...micro.Option) {
	// Init plugins
	for _, p := range Plugins() {
		p.Init(ctx)
	}

	if len(ctx.GlobalString("server_name")) > 0 {
		Name = ctx.GlobalString("server_name")
	}
	if len(ctx.String("address")) > 0 {
		Address = ctx.String("address")
	}
	if len(ctx.String("network")) > 0 {
		Network = ctx.String("network")
	}
	var nodes []string
	if len(ctx.String("node")) > 0 {
		nodes = strings.Split(ctx.String("node"), ",")
	}
	if len(ctx.String("resolver")) > 0 {
		Resolver = ctx.String("resolver")
	}
	var res resolver.Resolver
	switch Resolver {
	case "dns":
		res = &dns.Resolver{}
	case "http":
		res = &http.Resolver{}
	case "registry":
		res = &registry.Resolver{}
	}

	// Initialise service
	service := micro.NewService(
		micro.Name(Name),
		micro.RegisterTTL(time.Duration(ctx.GlobalInt("register_ttl"))*time.Second),
		micro.RegisterInterval(time.Duration(ctx.GlobalInt("register_interval"))*time.Second),
	)

	// create a tunnel
	tun := tunnel.NewTunnel(
		tunnel.Address(Address),
		tunnel.Nodes(nodes...),
	)

	// local tunnel router
	rtr := router.NewRouter(
		router.Network(Network),
		router.Id(service.Server().Options().Id),
		router.Registry(service.Client().Options().Registry),
	)

	// creaate new network
	net := network.NewNetwork(
		network.Id(service.Server().Options().Id),
		network.Name(Network),
		network.Address(Address),
		network.Nodes(nodes...),
		network.Tunnel(tun),
		network.Router(rtr),
		network.Resolver(res),
	)

	// local proxy
	prx := mucp.NewProxy(
		proxy.WithRouter(rtr),
		proxy.WithClient(service.Client()),
		proxy.WithLink("network", net.Client()),
	)

	// create a handler
	h := server.DefaultRouter.NewHandler(
		&handler.Network{
			Network: net,
		},
	)

	// register the handler
	server.DefaultRouter.Handle(h)

	// create a new muxer
	mux := mux.New(Name, prx)

	// init server
	service.Server().Init(
		server.WithRouter(mux),
	)

	// set network server to proxy
	net.Server().Init(
		server.WithRouter(mux),
	)

	// connect network
	if err := net.Connect(); err != nil {
		log.Logf("Network failed to connect: %v", err)
		os.Exit(1)
	}

	log.Logf("Network [%s] listening on %s", Name, Address)

	if err := service.Run(); err != nil {
		log.Logf("Network %s failed: %v", Name, err)
		os.Exit(1)
	}
}

func Commands(options ...micro.Option) []cli.Command {
	command := cli.Command{
		Name:  "network",
		Usage: "Run the micro network node",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:   "address",
				Usage:  "Set the micro network address :8085",
				EnvVar: "MICRO_NETWORK_ADDRESS",
			},
			cli.StringFlag{
				Name:   "network",
				Usage:  "Set the micro network name: go.micro",
				EnvVar: "MICRO_NETWORK",
			},
			cli.StringFlag{
				Name:   "node",
				Usage:  "Set the micro network server node address. This can be a comma separated list.",
				EnvVar: "MICRO_NETWORK_NODE",
			},
			cli.StringFlag{
				Name:   "resolver",
				Usage:  "Set the micro network resolver. This can be a comma separated list.",
				EnvVar: "MICRO_NETWORK_RESOLVER",
			},
		},
		Action: func(ctx *cli.Context) {
			run(ctx, options...)
		},
	}

	for _, p := range Plugins() {
		if cmds := p.Commands(); len(cmds) > 0 {
			command.Subcommands = append(command.Subcommands, cmds...)
		}

		if flags := p.Flags(); len(flags) > 0 {
			command.Flags = append(command.Flags, flags...)
		}
	}

	return []cli.Command{command}
}
<commit_msg>Close network connection on exit<commit_after>// package network implements micro network node
package network

import (
	"os"
	"strings"
	"time"

	"github.com/micro/cli"
	"github.com/micro/go-micro"
	"github.com/micro/go-micro/network"
	"github.com/micro/go-micro/network/handler"
	"github.com/micro/go-micro/network/resolver"
	"github.com/micro/go-micro/network/resolver/dns"
	"github.com/micro/go-micro/network/resolver/http"
	"github.com/micro/go-micro/network/resolver/registry"
	"github.com/micro/go-micro/proxy"
	"github.com/micro/go-micro/proxy/mucp"
	"github.com/micro/go-micro/router"
	"github.com/micro/go-micro/server"
	"github.com/micro/go-micro/tunnel"
	"github.com/micro/go-micro/util/log"
	"github.com/micro/go-micro/util/mux"
)

var (
	// Name of the network service
	Name = "go.micro.network"
	// Name of the micro network
	Network = "go.micro"
	// Address is the network address
	Address = ":8085"
	// Resolver is the network resolver
	Resolver = "registry"
)

// run runs the micro server
func run(ctx *cli.Context, srvOpts ...micro.Option) {
	// Init plugins
	for _, p := range Plugins() {
		p.Init(ctx)
	}

	if len(ctx.GlobalString("server_name")) > 0 {
		Name = ctx.GlobalString("server_name")
	}
	if len(ctx.String("address")) > 0 {
		Address = ctx.String("address")
	}
	if len(ctx.String("network")) > 0 {
		Network = ctx.String("network")
	}
	var nodes []string
	if len(ctx.String("node")) > 0 {
		nodes = strings.Split(ctx.String("node"), ",")
	}
	if len(ctx.String("resolver")) > 0 {
		Resolver = ctx.String("resolver")
	}
	var res resolver.Resolver
	switch Resolver {
	case "dns":
		res = &dns.Resolver{}
	case "http":
		res = &http.Resolver{}
	case "registry":
		res = &registry.Resolver{}
	}

	// Initialise service
	service := micro.NewService(
		micro.Name(Name),
		micro.RegisterTTL(time.Duration(ctx.GlobalInt("register_ttl"))*time.Second),
		micro.RegisterInterval(time.Duration(ctx.GlobalInt("register_interval"))*time.Second),
	)

	// create a tunnel
	tun := tunnel.NewTunnel(
		tunnel.Address(Address),
		tunnel.Nodes(nodes...),
	)

	// local tunnel router
	rtr := router.NewRouter(
		router.Network(Network),
		router.Id(service.Server().Options().Id),
		router.Registry(service.Client().Options().Registry),
	)

	// creaate new network
	net := network.NewNetwork(
		network.Id(service.Server().Options().Id),
		network.Name(Network),
		network.Address(Address),
		network.Nodes(nodes...),
		network.Tunnel(tun),
		network.Router(rtr),
		network.Resolver(res),
	)

	// local proxy
	prx := mucp.NewProxy(
		proxy.WithRouter(rtr),
		proxy.WithClient(service.Client()),
		proxy.WithLink("network", net.Client()),
	)

	// create a handler
	h := server.DefaultRouter.NewHandler(
		&handler.Network{
			Network: net,
		},
	)

	// register the handler
	server.DefaultRouter.Handle(h)

	// create a new muxer
	mux := mux.New(Name, prx)

	// init server
	service.Server().Init(
		server.WithRouter(mux),
	)

	// set network server to proxy
	net.Server().Init(
		server.WithRouter(mux),
	)

	// connect network
	if err := net.Connect(); err != nil {
		log.Logf("Network failed to connect: %v", err)
		os.Exit(1)
	}
	// close on exit
	defer net.Close()

	log.Logf("Network [%s] listening on %s", Name, Address)

	if err := service.Run(); err != nil {
		log.Logf("Network %s failed: %v", Name, err)
		net.Close()
		os.Exit(1)
	}
}

func Commands(options ...micro.Option) []cli.Command {
	command := cli.Command{
		Name:  "network",
		Usage: "Run the micro network node",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:   "address",
				Usage:  "Set the micro network address :8085",
				EnvVar: "MICRO_NETWORK_ADDRESS",
			},
			cli.StringFlag{
				Name:   "network",
				Usage:  "Set the micro network name: go.micro",
				EnvVar: "MICRO_NETWORK",
			},
			cli.StringFlag{
				Name:   "node",
				Usage:  "Set the micro network server node address. This can be a comma separated list.",
				EnvVar: "MICRO_NETWORK_NODE",
			},
			cli.StringFlag{
				Name:   "resolver",
				Usage:  "Set the micro network resolver. This can be a comma separated list.",
				EnvVar: "MICRO_NETWORK_RESOLVER",
			},
		},
		Action: func(ctx *cli.Context) {
			run(ctx, options...)
		},
	}

	for _, p := range Plugins() {
		if cmds := p.Commands(); len(cmds) > 0 {
			command.Subcommands = append(command.Subcommands, cmds...)
		}

		if flags := p.Flags(); len(flags) > 0 {
			command.Flags = append(command.Flags, flags...)
		}
	}

	return []cli.Command{command}
}
<|endoftext|>


<commit_before>package table

import (
	"fmt"
	"strings"
)

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}

func getStringTab(rt *RTable) (*[][][]string, int) {
	s := make([][][]string, len(*rt))
	m := 1

	for i := 0; i < len(s); i++ {
		col := rt.GetColumn(i)
		s[i] = make([][]string, len(*col))
		for ii := 0; ii < len(s[i]); ii++ {
			item := col.GetItem(ii)
			s[i][ii] = make([]string, len(*item))
			for iii := 0; iii < len(s[i][ii]); iii++ {
				s[i][ii][iii] = string((*item)[iii])
				m = max(len(s[i][ii][iii]), m)
			}
		}
	}
	return &s, m
}

func printToken(s string, max int) {
	fmt.Printf(s)
	if l := len(s); l < max {
		fmt.Printf(strings.Repeat(" ", max-l))
	}
}

func PrettyPrint(rt *RTable) string {
	s, m := getStringTab(rt)
	length := len(*rt)

	fmt.Printf("+")
	fmt.Printf(strings.Repeat("-", (m+1)*length-1))
	fmt.Println("+")
	for row := 0; row < length; row++ {
		l := 1
		for tok := 0; tok < l; tok++ {
			fmt.Printf(strings.Repeat(" ", (m+1)*row))
			fmt.Printf("|")
			for col := row; col < length; col++ {
				if len((*s)[col][row]) == 0 {
					printToken("", m)
				} else if tok == 0 {
					l = max(len((*s)[col][row]), l)
					printToken((*s)[col][row][0], m)
				} else if tok < len((*s)[col][row]) {
					printToken((*s)[col][row][tok], m)
				} else {
					printToken("", m)
				}
				fmt.Printf("|")
			}
			fmt.Println()
		}
		fmt.Printf(strings.Repeat(" ", (m+1)*row))
		fmt.Printf("+")
		fmt.Printf(strings.Repeat("-", (m+1)*(length-row)-1))
		if row+1 == length {
			fmt.Println("+")
		} else {
			fmt.Println("|")
		}
	}

	return ""
}
<commit_msg>add comments to table/prettyprint.go<commit_after>package table

import (
	"fmt"
	"strings"
)

// max returns the meximum of two integers.
func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}

// getStringTab returns a table of all the tokens of the recognition
// table converted into strings.
func getStringTab(rt *RTable) (*[][][]string, int) {
	s := make([][][]string, len(*rt))
	m := 1

	for i := 0; i < len(s); i++ {
		col := rt.GetColumn(i)
		s[i] = make([][]string, len(*col))
		for ii := 0; ii < len(s[i]); ii++ {
			item := col.GetItem(ii)
			s[i][ii] = make([]string, len(*item))
			for iii := 0; iii < len(s[i][ii]); iii++ {
				s[i][ii][iii] = string((*item)[iii])
				m = max(len(s[i][ii][iii]), m)
			}
		}
	}
	return &s, m
}

// printToken prints a string followed by spaces corresponding to
// the `max` parameter.
func printToken(s string, max int) {
	fmt.Printf(s)
	if l := len(s); l < max {
		fmt.Printf(strings.Repeat(" ", max-l))
	}
}

// PrettyPrint prints a given recognition table.
func PrettyPrint(rt *RTable) string {
	s, m := getStringTab(rt)
	length := len(*rt)

	fmt.Printf("+")
	fmt.Printf(strings.Repeat("-", (m+1)*length-1))
	fmt.Println("+")
	for row := 0; row < length; row++ {
		l := 1
		for tok := 0; tok < l; tok++ {
			fmt.Printf(strings.Repeat(" ", (m+1)*row))
			fmt.Printf("|")
			for col := row; col < length; col++ {
				if len((*s)[col][row]) == 0 {
					printToken("", m)
				} else if tok == 0 {
					l = max(len((*s)[col][row]), l)
					printToken((*s)[col][row][0], m)
				} else if tok < len((*s)[col][row]) {
					printToken((*s)[col][row][tok], m)
				} else {
					printToken("", m)
				}
				fmt.Printf("|")
			}
			fmt.Println()
		}
		fmt.Printf(strings.Repeat(" ", (m+1)*row))
		fmt.Printf("+")
		fmt.Printf(strings.Repeat("-", (m+1)*(length-row)-1))
		if row+1 == length {
			fmt.Println("+")
		} else {
			fmt.Println("|")
		}
	}

	return ""
}
<|endoftext|>