<commit_before>// This is a client that writes out to a file, and optionally rolls the file

package main

import (
	"../../nsq"
	"../../util"
	"compress/gzip"
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"path"
	"strings"
	"syscall"
	"time"
)

var (
	datetimeFormat   = flag.String("datetime-format", "%Y-%m-%d_%H", "strftime compatible format for <DATETIME> in filename format")
	filenameFormat   = flag.String("filename-format", "<TOPIC>.<HOST>.<DATETIME>.log", "output filename format (<TOPIC>, <HOST>, <DATETIME> are replaced)")
	showVersion      = flag.Bool("version", false, "print version string")
	hostIdentifier   = flag.String("host-identifier", "", "value to output in log filename in place of hostname. <SHORT_HOST> and <HOSTNAME> are valid replacement tokens")
	outputDir        = flag.String("output-dir", "/tmp", "directory to write output files to")
	topic            = flag.String("topic", "", "nsq topic")
	channel          = flag.String("channel", "nsq_to_file", "nsq channel")
	maxInFlight      = flag.Int("max-in-flight", 1000, "max number of messages to allow in flight")
	gzipFlag         = flag.Bool("gzip", false, "gzip output files.")
	verbose          = flag.Bool("verbose", false, "verbose logging")
	nsqdTCPAddrs     = util.StringArray{}
	lookupdHTTPAddrs = util.StringArray{}
)

func init() {
	flag.Var(&nsqdTCPAddrs, "nsqd-tcp-address", "nsqd TCP address (may be given multiple times)")
	flag.Var(&lookupdHTTPAddrs, "lookupd-http-address", "lookupd HTTP address (may be given multiple times)")
}

type FileLogger struct {
	out        *os.File
	gzipWriter *gzip.Writer
	filename   string
	logChan    chan *Message
}

type Message struct {
	*nsq.Message
	returnChannel chan *nsq.FinishedMessage
}

type SyncMsg struct {
	m             *nsq.FinishedMessage
	returnChannel chan *nsq.FinishedMessage
}

func (l *FileLogger) HandleMessage(m *nsq.Message, responseChannel chan *nsq.FinishedMessage) {
	l.logChan <- &Message{m, responseChannel}
}

func router(r *nsq.Reader, f *FileLogger, termChan chan os.Signal, hupChan chan os.Signal) {
	pos := 0
	output := make([]*Message, *maxInFlight)
	sync := false
	ticker := time.NewTicker(time.Duration(30) * time.Second)
	closing := false

	for {
		select {
		case <-termChan:
			ticker.Stop()
			r.Stop()
			// ensures that we keep flushing whatever is left in the channels
			closing = true
		case <-hupChan:
			f.Close()
			f.updateFile()
			sync = true
		case <-ticker.C:
			f.updateFile()
			sync = true
		case m := <-f.logChan:
			if f.updateFile() {
				sync = true
			}
			_, err := f.Write(m.Body)
			if err != nil {
				log.Fatalf("ERROR: writing message to disk - %s", err.Error())
			}
			_, err = f.Write([]byte("\n"))
			if err != nil {
				log.Fatalf("ERROR: writing newline to disk - %s", err.Error())
			}
			output[pos] = m
			pos++
			if pos == *maxInFlight {
				sync = true
			}
		}

		if closing || sync || r.IsStarved() {
			if pos > 0 {
				log.Printf("syncing %d records to disk", pos)
				err := f.Sync()
				if err != nil {
					log.Fatalf("ERROR: failed syncing messages - %s", err.Error())
				}
				for pos > 0 {
					pos--
					m := output[pos]
					m.returnChannel <- &nsq.FinishedMessage{m.Id, 0, true}
					output[pos] = nil
				}
			}
			sync = false
		}
	}
}

func (f *FileLogger) Close() {
	if f.out != nil {
		if f.gzipWriter != nil {
			f.gzipWriter.Close()
		}
		f.out.Close()
		f.out = nil
	}
}
func (f *FileLogger) Write(p []byte) (n int, err error) {
	if f.gzipWriter != nil {
		return f.gzipWriter.Write(p)
	}
	return f.out.Write(p)
}
func (f *FileLogger) Sync() error {
	var err error
	if f.gzipWriter != nil {
		f.gzipWriter.Close()
		err = f.out.Sync()
		f.gzipWriter, _ = gzip.NewWriterLevel(f.out, gzip.BestSpeed)
	} else {
		err = f.out.Sync()
	}
	return err
}

func (f *FileLogger) updateFile() bool {
	t := time.Now()

	hostname, err := os.Hostname()
	if err != nil {
		log.Fatalf("ERROR: unable to get hostname %s", err.Error())
	}
	shortHostname := strings.Split(hostname, ".")[0]
	identifier := shortHostname
	if len(*hostIdentifier) != 0 {
		identifier = strings.Replace(*hostIdentifier, "<SHORT_HOST>", shortHostname, -1)
		identifier = strings.Replace(identifier, "<HOSTNAME>", hostname, -1)
	}

	filename := strings.Replace(*filenameFormat, "<TOPIC>", *topic, -1)
	filename = strings.Replace(filename, "<HOST>", identifier, -1)
	datetime := strftime(*datetimeFormat, t)
	filename = strings.Replace(filename, "<DATETIME>", datetime, -1)
	if *gzipFlag {
		filename = filename + ".gz"
	}

	if filename != f.filename || f.out == nil {
		log.Printf("old %s new %s", f.filename, filename)
		f.Close()
		os.MkdirAll(*outputDir, 777)
		log.Printf("opening %s/%s", *outputDir, filename)
		newfile, err := os.OpenFile(path.Join(*outputDir, filename), os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
		f.out = newfile
		if err != nil {
			log.Fatal(err)
		}
		f.filename = filename
		if *gzipFlag {
			f.gzipWriter, _ = gzip.NewWriterLevel(newfile, gzip.BestSpeed)
		}
		return true
	}

	return false
}

func main() {
	flag.Parse()

	if *showVersion {
		fmt.Printf("nsq_to_file v%s\n", util.BINARY_VERSION)
		return
	}

	if *topic == "" || *channel == "" {
		log.Fatalf("--topic and --channel are required")
	}

	if *maxInFlight < 0 {
		log.Fatalf("--max-in-flight must be > 0")
	}

	if len(nsqdTCPAddrs) == 0 && len(lookupdHTTPAddrs) == 0 {
		log.Fatalf("--nsqd-tcp-address or --lookupd-http-address required.")
	}
	if len(nsqdTCPAddrs) != 0 && len(lookupdHTTPAddrs) != 0 {
		log.Fatalf("use --nsqd-tcp-address or --lookupd-http-address not both")
	}

	hupChan := make(chan os.Signal, 1)
	termChan := make(chan os.Signal, 1)
	signal.Notify(hupChan, syscall.SIGHUP)
	signal.Notify(termChan, syscall.SIGINT, syscall.SIGTERM)

	f := &FileLogger{
		logChan: make(chan *Message, 1),
	}

	r, err := nsq.NewReader(*topic, *channel)
	if err != nil {
		log.Fatalf(err.Error())
	}
	r.SetMaxInFlight(*maxInFlight)
	r.VerboseLogging = *verbose

	r.AddAsyncHandler(f)
	go router(r, f, termChan, hupChan)

	for _, addrString := range nsqdTCPAddrs {
		err := r.ConnectToNSQ(addrString)
		if err != nil {
			log.Fatalf(err.Error())
		}
	}

	for _, addrString := range lookupdHTTPAddrs {
		log.Printf("lookupd addr %s", addrString)
		err := r.ConnectToLookupd(addrString)
		if err != nil {
			log.Fatalf(err.Error())
		}
	}

	<-r.ExitChan
}
<commit_msg>set gzip compression via command line arg<commit_after>// This is a client that writes out to a file, and optionally rolls the file

package main

import (
	"../../nsq"
	"../../util"
	"compress/gzip"
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"path"
	"strings"
	"syscall"
	"time"
)

var (
	datetimeFormat   = flag.String("datetime-format", "%Y-%m-%d_%H", "strftime compatible format for <DATETIME> in filename format")
	filenameFormat   = flag.String("filename-format", "<TOPIC>.<HOST>.<DATETIME>.log", "output filename format (<TOPIC>, <HOST>, <DATETIME> are replaced)")
	showVersion      = flag.Bool("version", false, "print version string")
	hostIdentifier   = flag.String("host-identifier", "", "value to output in log filename in place of hostname. <SHORT_HOST> and <HOSTNAME> are valid replacement tokens")
	outputDir        = flag.String("output-dir", "/tmp", "directory to write output files to")
	topic            = flag.String("topic", "", "nsq topic")
	channel          = flag.String("channel", "nsq_to_file", "nsq channel")
	maxInFlight      = flag.Int("max-in-flight", 1000, "max number of messages to allow in flight")
	gzipCompression  = flag.Int("gzip-compression", 3, "gzip compression level. 1 BestSpeed, 2 BestCompression, 3 DefaultCompression")
	gzipFlag         = flag.Bool("gzip", false, "gzip output files.")
	verbose          = flag.Bool("verbose", false, "verbose logging")
	nsqdTCPAddrs     = util.StringArray{}
	lookupdHTTPAddrs = util.StringArray{}
)

func init() {
	flag.Var(&nsqdTCPAddrs, "nsqd-tcp-address", "nsqd TCP address (may be given multiple times)")
	flag.Var(&lookupdHTTPAddrs, "lookupd-http-address", "lookupd HTTP address (may be given multiple times)")
}

type FileLogger struct {
	out              *os.File
	gzipWriter       *gzip.Writer
	filename         string
	logChan          chan *Message
	compressionLevel int
}

type Message struct {
	*nsq.Message
	returnChannel chan *nsq.FinishedMessage
}

type SyncMsg struct {
	m             *nsq.FinishedMessage
	returnChannel chan *nsq.FinishedMessage
}

func (l *FileLogger) HandleMessage(m *nsq.Message, responseChannel chan *nsq.FinishedMessage) {
	l.logChan <- &Message{m, responseChannel}
}

func router(r *nsq.Reader, f *FileLogger, termChan chan os.Signal, hupChan chan os.Signal) {
	pos := 0
	output := make([]*Message, *maxInFlight)
	sync := false
	ticker := time.NewTicker(time.Duration(30) * time.Second)
	closing := false

	for {
		select {
		case <-termChan:
			ticker.Stop()
			r.Stop()
			// ensures that we keep flushing whatever is left in the channels
			closing = true
		case <-hupChan:
			f.Close()
			f.updateFile()
			sync = true
		case <-ticker.C:
			f.updateFile()
			sync = true
		case m := <-f.logChan:
			if f.updateFile() {
				sync = true
			}
			_, err := f.Write(m.Body)
			if err != nil {
				log.Fatalf("ERROR: writing message to disk - %s", err.Error())
			}
			_, err = f.Write([]byte("\n"))
			if err != nil {
				log.Fatalf("ERROR: writing newline to disk - %s", err.Error())
			}
			output[pos] = m
			pos++
			if pos == *maxInFlight {
				sync = true
			}
		}

		if closing || sync || r.IsStarved() {
			if pos > 0 {
				log.Printf("syncing %d records to disk", pos)
				err := f.Sync()
				if err != nil {
					log.Fatalf("ERROR: failed syncing messages - %s", err.Error())
				}
				for pos > 0 {
					pos--
					m := output[pos]
					m.returnChannel <- &nsq.FinishedMessage{m.Id, 0, true}
					output[pos] = nil
				}
			}
			sync = false
		}
	}
}

func (f *FileLogger) Close() {
	if f.out != nil {
		if f.gzipWriter != nil {
			f.gzipWriter.Close()
		}
		f.out.Close()
		f.out = nil
	}
}
func (f *FileLogger) Write(p []byte) (n int, err error) {
	if f.gzipWriter != nil {
		return f.gzipWriter.Write(p)
	}
	return f.out.Write(p)
}
func (f *FileLogger) Sync() error {
	var err error
	if f.gzipWriter != nil {
		f.gzipWriter.Close()
		err = f.out.Sync()
		f.gzipWriter, _ = gzip.NewWriterLevel(f.out, f.compressionLevel)
	} else {
		err = f.out.Sync()
	}
	return err
}

func (f *FileLogger) updateFile() bool {
	t := time.Now()

	hostname, err := os.Hostname()
	if err != nil {
		log.Fatalf("ERROR: unable to get hostname %s", err.Error())
	}
	shortHostname := strings.Split(hostname, ".")[0]
	identifier := shortHostname
	if len(*hostIdentifier) != 0 {
		identifier = strings.Replace(*hostIdentifier, "<SHORT_HOST>", shortHostname, -1)
		identifier = strings.Replace(identifier, "<HOSTNAME>", hostname, -1)
	}

	filename := strings.Replace(*filenameFormat, "<TOPIC>", *topic, -1)
	filename = strings.Replace(filename, "<HOST>", identifier, -1)
	datetime := strftime(*datetimeFormat, t)
	filename = strings.Replace(filename, "<DATETIME>", datetime, -1)
	if *gzipFlag {
		filename = filename + ".gz"
	}

	if filename != f.filename || f.out == nil {
		log.Printf("old %s new %s", f.filename, filename)
		f.Close()
		os.MkdirAll(*outputDir, 777)
		log.Printf("opening %s/%s", *outputDir, filename)
		newfile, err := os.OpenFile(path.Join(*outputDir, filename), os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
		f.out = newfile
		if err != nil {
			log.Fatal(err)
		}
		f.filename = filename
		if *gzipFlag {
			f.gzipWriter, _ = gzip.NewWriterLevel(newfile, f.compressionLevel)
		}
		return true
	}

	return false
}

func NewFileLogger(compressionLevel int) *FileLogger {
	var speed int
	switch compressionLevel {
	case 1:
		speed = gzip.BestSpeed
	case 2:
		speed = gzip.BestCompression
	case 3:
		speed = gzip.DefaultCompression
	}
	return &FileLogger{
		logChan:          make(chan *Message, 1),
		compressionLevel: speed,
	}
}

func main() {
	flag.Parse()

	if *showVersion {
		fmt.Printf("nsq_to_file v%s\n", util.BINARY_VERSION)
		return
	}

	if *topic == "" || *channel == "" {
		log.Fatalf("--topic and --channel are required")
	}

	if *maxInFlight < 0 {
		log.Fatalf("--max-in-flight must be > 0")
	}

	if len(nsqdTCPAddrs) == 0 && len(lookupdHTTPAddrs) == 0 {
		log.Fatalf("--nsqd-tcp-address or --lookupd-http-address required.")
	}
	if len(nsqdTCPAddrs) != 0 && len(lookupdHTTPAddrs) != 0 {
		log.Fatalf("use --nsqd-tcp-address or --lookupd-http-address not both")
	}

	if *gzipCompression < 1 || *gzipCompression > 3 {
		log.Fatalf("invalid --gzip-compresion value (%v). should be 1,2 or 3", *gzipCompression)
	}

	hupChan := make(chan os.Signal, 1)
	termChan := make(chan os.Signal, 1)
	signal.Notify(hupChan, syscall.SIGHUP)
	signal.Notify(termChan, syscall.SIGINT, syscall.SIGTERM)

	f := NewFileLogger(*gzipCompression)

	r, err := nsq.NewReader(*topic, *channel)
	if err != nil {
		log.Fatalf(err.Error())
	}
	r.SetMaxInFlight(*maxInFlight)
	r.VerboseLogging = *verbose

	r.AddAsyncHandler(f)
	go router(r, f, termChan, hupChan)

	for _, addrString := range nsqdTCPAddrs {
		err := r.ConnectToNSQ(addrString)
		if err != nil {
			log.Fatalf(err.Error())
		}
	}

	for _, addrString := range lookupdHTTPAddrs {
		log.Printf("lookupd addr %s", addrString)
		err := r.ConnectToLookupd(addrString)
		if err != nil {
			log.Fatalf(err.Error())
		}
	}

	<-r.ExitChan
}
<|endoftext|>


<commit_before>//
// Variable frequency CPU usage sampling on Linux via /proc
// Maybe this will turn into something like prstat on Solaris
//

// TOOD - check for rollover in sched
// TODO - check max field length for assumed present fields
// TODO - tui fill out text area
// TODO - tui y axis scale
// TODO - tui x axis labels
// TODO - tui consistent proc colors
// TODO - tui proc labels
// TODO - tui top procs should not include ctime

package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"runtime/pprof"
	"sort"
	"time"
)

const histMin = 0
const histMax = 100000000
const histSigFigs = 2
const maxProcsToScan = 2048

func main() {
	var interval = flag.Int("i", 1000, "interval (ms) between measurements")
	var samples = flag.Int("s", 60, "sample counts to aggregate for output")
	var topN = flag.Int("n", 10, "show top N processes")
	var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
	var jiffy = flag.Int("jiffy", 100, "length of a jiffy")
	var useTui = flag.Bool("t", false, "use fancy terminal mode")

	flag.Parse()

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		if err = pprof.StartCPUProfile(f); err != nil {
			log.Fatal(err)
		}
	}

	uiQuitChan := make(chan string)
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt)

	go func() {
		select {
		case <-sigChan:
			fmt.Fprintln(os.Stderr, "quitting on signal")
		case msg := <-uiQuitChan:
			fmt.Fprintln(os.Stderr, msg)
		}
		pprof.StopCPUProfile()
		os.Exit(0)
	}()

	if *useTui {
		go tuiInit(uiQuitChan)
	}

	cmdNames := make(cmdlineMap)

	procCur := make(procStatsMap)
	procPrev := make(procStatsMap)
	procSum := make(procStatsMap)
	procHist := make(procStatsHistMap)

	schedCur := make(schedStatsMap)
	schedPrev := make(schedStatsMap)
	schedSum := make(schedStatsMap)

	var sysCur *systemStats
	var sysPrev *systemStats
	var sysSum *systemStats
	var sysHist *systemStatsHist

	var t1, t2 time.Time

	// run all scans one time to establish a baseline
	pids := make(pidlist, 0, maxProcsToScan)
	getPidList(&pids)

	schedPrev = schedReaderPids(pids)
	t1 = time.Now()
	procPrev = statReader(pids, cmdNames)
	sysPrev = statReaderGlobal()
	sysSum = &systemStats{}
	sysHist = &systemStatsHist{}
	t2 = time.Now()

	targetSleep := time.Duration(*interval) * time.Millisecond
	adjustedSleep := targetSleep - t2.Sub(t1)

	topPids := make(pidlist, *topN)
	for {
		for count := 0; count < *samples; count++ {
			time.Sleep(adjustedSleep)

			t1 = time.Now()
			getPidList(&pids)

			procCur = statReader(pids, cmdNames)
			procDelta := statRecord(procCur, procPrev, procSum, procHist)
			procPrev = procCur

			sysCur = statReaderGlobal()
			sysDelta := statsRecordGlobal(sysCur, sysPrev, sysSum, sysHist)
			sysPrev = sysCur

			if *useTui {
				tuiGraphUpdate(sysDelta, procDelta, topPids)
			}

			t2 = time.Now()
			adjustedSleep = targetSleep - t2.Sub(t1)
		}

		topHist := sortList(procHist, *topN)
		for i := 0; i < *topN; i++ {
			topPids[i] = topHist[i].pid
		}

		schedCur = schedReaderPids(topPids)
		schedRecord(schedCur, schedPrev, schedSum)
		schedPrev = schedCur

		if *useTui {
			tuiListUpdate(topPids, procSum, procHist, sysHist)
		} else {
			dumpStats(cmdNames, topPids, procSum, procHist, sysSum, sysHist, schedSum, jiffy, interval, samples)
		}
		procHist = make(procStatsHistMap)
		procSum = make(procStatsMap)
		sysHist = &systemStatsHist{}
		sysSum = &systemStats{}
		t2 = time.Now()
		adjustedSleep = targetSleep - t2.Sub(t1)
	}
}

type bar struct {
	baz int
}

// Wrapper to sort histograms by max but remember which pid they are
type sortHist struct {
	pid  int
	hist *procStatsHist
}

// ByMax sorts stats by max usage
type ByMax []*sortHist

func (m ByMax) Len() int {
	return len(m)
}
func (m ByMax) Swap(i, j int) {
	m[i], m[j] = m[j], m[i]
}
func (m ByMax) Less(i, j int) bool {
	maxI := maxList([]float64{
		float64(m[i].hist.ustime.Max()),
		float64(m[i].hist.custime.Max()),
		float64(m[i].hist.delayacctBlkioTicks.Max()),
	})
	maxJ := maxList([]float64{
		float64(m[j].hist.ustime.Max()),
		float64(m[j].hist.custime.Max()),
		float64(m[j].hist.delayacctBlkioTicks.Max()),
	})
	return maxI > maxJ
}
func maxList(list []float64) float64 {
	ret := list[0]
	for i := 1; i < len(list); i++ {
		if list[i] > ret {
			ret = list[i]
		}
	}
	return ret
}

func sortList(histStats procStatsHistMap, limit int) []*sortHist {
	var list []*sortHist

	for pid, hist := range histStats {
		list = append(list, &sortHist{pid, hist})
	}
	sort.Sort(ByMax(list))

	if len(list) > limit {
		list = list[:limit]
	}

	return list
}

func formatMem(num uint64) string {
	letter := string("K")

	num = num * 4
	if num >= 10000 {
		num = (num + 512) / 1024
		letter = "M"
		if num >= 10000 {
			num = (num + 512) / 1024
			letter = "G"
		}
	}
	return fmt.Sprintf("%d%s", num, letter)
}

func dumpStats(cmdNames cmdlineMap, list pidlist, sumStats procStatsMap, histStats procStatsHistMap,
	sysSum *systemStats, sysHist *systemStatsHist, sumSched schedStatsMap, jiffy, interval, samples *int) {

	scale := func(val float64) float64 {
		return val / float64(*jiffy) / float64(*interval) * 1000 * 100
	}
	scaleSum := func(val float64, count int64) float64 {
		valSec := val / float64(*jiffy)
		sampleSec := float64(*interval) * float64(count) / 1000.0
		ret := (valSec / sampleSec) * 100
		return ret
	}
	scaleSched := func(val float64) float64 {
		return val / float64(*jiffy) / float64((*interval)*(*samples)) * 100
	}

	fmt.Printf("usr:    %4s/%4s/%4s   sys:%4s/%4s/%4s  nice:%4s/%4s/%4s    idle:%4s/%4s/%4s\n",
		trim(scale(float64(sysHist.usr.Min())), 4),
		trim(scale(float64(sysHist.usr.Max())), 4),
		trim(scale(sysHist.usr.Mean()), 4),

		trim(scale(float64(sysHist.sys.Min())), 4),
		trim(scale(float64(sysHist.sys.Max())), 4),
		trim(scale(sysHist.sys.Mean()), 4),

		trim(scale(float64(sysHist.nice.Min())), 4),
		trim(scale(float64(sysHist.nice.Max())), 4),
		trim(scale(sysHist.nice.Mean()), 4),

		trim(scale(float64(sysHist.idle.Min())), 4),
		trim(scale(float64(sysHist.idle.Max())), 4),
		trim(scale(sysHist.idle.Mean()), 4),
	)
	fmt.Printf("iowait: %4s/%4s/%4s  ctxt:%4s/%4s/%4s  prun:%4s/%4s/%4s  pblock:%4s/%4s/%4s  pstart: %4d\n",
		trim(scale(float64(sysHist.iowait.Min())), 4),
		trim(scale(float64(sysHist.iowait.Max())), 4),
		trim(scale(sysHist.iowait.Mean()), 4),

		trim(scale(float64(sysHist.ctxt.Min())), 4),
		trim(scale(float64(sysHist.ctxt.Max())), 4),
		trim(scale(sysHist.ctxt.Mean()), 4),

		trim(float64(sysHist.procsRunning.Min()), 4),
		trim(float64(sysHist.procsRunning.Max()), 4),
		trim(sysHist.procsRunning.Mean(), 4),

		trim(float64(sysHist.procsBlocked.Min()), 4),
		trim(float64(sysHist.procsBlocked.Max()), 4),
		trim(sysHist.procsBlocked.Mean(), 4),

		sysSum.procsTotal,
	)

	fmt.Print("                   comm     pid     min     max     usr     sys    nice   ctime    slat     ctx     icx     rss  iowait thrd  sam\n")
	for _, pid := range list {
		hist := histStats[pid]

		var schedWait, nrSwitches, nrInvoluntarySwitches string
		sched, ok := sumSched[pid]
		if ok == true {
			schedWait = trim(scaleSched(sched.waitSum), 7)
			nrSwitches = fmt.Sprintf("%d", sched.nrSwitches)
			nrInvoluntarySwitches = fmt.Sprintf("%d", sched.nrInvoluntarySwitches)
		} else {
			schedWait = "-"
			nrSwitches = "-"
			nrInvoluntarySwitches = "-"
		}
		sampleCount := hist.utime.TotalCount()
		var friendlyName string
		cmdName, ok := cmdNames[pid]
		if ok == true && len(cmdName.friendly) > 0 {
			friendlyName = cmdName.friendly
		} else {
			// This should not happen as long as the cmdline resolver works
			fmt.Println("using comm for ", cmdName, pid)
			friendlyName = sumStats[pid].comm
		}
		if len(friendlyName) > 23 {
			friendlyName = friendlyName[:23]
		}

		fmt.Printf("%23s %7d %7s %7s %7s %7s %7s %7s %7s %7s %7s %7s %7s %4d %4d\n",
			friendlyName,
			pid,
			trim(scale(float64(hist.ustime.Min())), 7),
			trim(scale(float64(hist.ustime.Max())), 7),
			trim(scaleSum(float64(sumStats[pid].utime), sampleCount), 7),
			trim(scaleSum(float64(sumStats[pid].stime), sampleCount), 7),
			trim(float64(sumStats[pid].nice), 7),
			trim(scaleSum(float64(sumStats[pid].cutime+sumStats[pid].cstime), sampleCount), 7),
			schedWait,
			nrSwitches,
			nrInvoluntarySwitches,
			formatMem(sumStats[pid].rss),
			trim(scaleSum(float64(sumStats[pid].delayacctBlkioTicks), sampleCount), 7),
			sumStats[pid].numThreads,
			sampleCount,
		)
	}
	fmt.Println()
}
<commit_msg>Slightly better formatting.<commit_after>//
// Variable frequency CPU usage sampling on Linux via /proc
// Maybe this will turn into something like prstat on Solaris
//

// easy
// TOOD - check for rollover in sched
// TODO - check max field length for assumed present fields
// TODO - tui fill out text area
// TODO - tui better colors
// TODO - tui y axis labels should go back down after growing, either that
//        or figure out where this giant values are coming from
// TODO - use the actual time of each measurement not the expected time
// TODO - move dumpStats into a separate file like termui is

// hard
// TODO - split into long running backend and multiple frontends

package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"runtime/pprof"
	"sort"
	"time"
)

const histMin = 0
const histMax = 100000000
const histSigFigs = 2
const maxProcsToScan = 2048

func main() {
	var interval = flag.Int("i", 1000, "interval (ms) between measurements")
	var samples = flag.Int("s", 60, "sample counts to aggregate for output")
	var topN = flag.Int("n", 10, "show top N processes")
	var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
	var jiffy = flag.Int("jiffy", 100, "length of a jiffy")
	var useTui = flag.Bool("t", false, "use fancy terminal mode")

	flag.Parse()

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		if err = pprof.StartCPUProfile(f); err != nil {
			log.Fatal(err)
		}
	}

	uiQuitChan := make(chan string)
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt)

	go func() {
		select {
		case <-sigChan:
			fmt.Fprintln(os.Stderr, "quitting on signal")
		case msg := <-uiQuitChan:
			fmt.Fprintln(os.Stderr, msg)
		}
		pprof.StopCPUProfile()
		os.Exit(0)
	}()

	if *useTui {
		go tuiInit(uiQuitChan, *interval)
	}

	cmdNames := make(cmdlineMap)

	procCur := make(procStatsMap)
	procPrev := make(procStatsMap)
	procSum := make(procStatsMap)
	procHist := make(procStatsHistMap)

	schedCur := make(schedStatsMap)
	schedPrev := make(schedStatsMap)
	schedSum := make(schedStatsMap)

	var sysCur *systemStats
	var sysPrev *systemStats
	var sysSum *systemStats
	var sysHist *systemStatsHist

	var t1, t2 time.Time

	// run all scans one time to establish a baseline
	pids := make(pidlist, 0, maxProcsToScan)
	getPidList(&pids)

	schedPrev = schedReaderPids(pids)
	t1 = time.Now()
	procPrev = statReader(pids, cmdNames)
	sysPrev = statReaderGlobal()
	sysSum = &systemStats{}
	sysHist = &systemStatsHist{}
	t2 = time.Now()

	targetSleep := time.Duration(*interval) * time.Millisecond
	adjustedSleep := targetSleep - t2.Sub(t1)

	topPids := make(pidlist, *topN)
	for {
		for count := 0; count < *samples; count++ {
			time.Sleep(adjustedSleep)

			t1 = time.Now()
			getPidList(&pids)

			procCur = statReader(pids, cmdNames)
			procDelta := statRecord(procCur, procPrev, procSum, procHist)
			procPrev = procCur

			sysCur = statReaderGlobal()
			sysDelta := statsRecordGlobal(sysCur, sysPrev, sysSum, sysHist)
			sysPrev = sysCur

			if *useTui {
				tuiGraphUpdate(sysDelta, procDelta, topPids, *jiffy, *interval)
			}

			t2 = time.Now()
			adjustedSleep = targetSleep - t2.Sub(t1)
		}

		topHist := sortList(procHist, *topN)
		for i := 0; i < *topN; i++ {
			topPids[i] = topHist[i].pid
		}

		schedCur = schedReaderPids(topPids)
		schedRecord(schedCur, schedPrev, schedSum)
		schedPrev = schedCur

		if *useTui {
			tuiListUpdate(topPids, procSum, cmdNames, procHist, sysHist)
		} else {
			dumpStats(cmdNames, topPids, procSum, procHist, sysSum, sysHist, schedSum, jiffy, interval, samples)
		}
		procHist = make(procStatsHistMap)
		procSum = make(procStatsMap)
		sysHist = &systemStatsHist{}
		sysSum = &systemStats{}
		t2 = time.Now()
		adjustedSleep = targetSleep - t2.Sub(t1)
	}
}

type bar struct {
	baz int
}

// Wrapper to sort histograms by max but remember which pid they are
type sortHist struct {
	pid  int
	hist *procStatsHist
}

// ByMax sorts stats by max usage
type ByMax []*sortHist

func (m ByMax) Len() int {
	return len(m)
}
func (m ByMax) Swap(i, j int) {
	m[i], m[j] = m[j], m[i]
}
func (m ByMax) Less(i, j int) bool {
	maxI := maxList([]float64{
		float64(m[i].hist.ustime.Max()),
		float64(m[i].hist.delayacctBlkioTicks.Max()),
	})
	maxJ := maxList([]float64{
		float64(m[j].hist.ustime.Max()),
		float64(m[j].hist.delayacctBlkioTicks.Max()),
	})
	return maxI > maxJ
}
func maxList(list []float64) float64 {
	ret := list[0]
	for i := 1; i < len(list); i++ {
		if list[i] > ret {
			ret = list[i]
		}
	}
	return ret
}

func sortList(histStats procStatsHistMap, limit int) []*sortHist {
	var list []*sortHist

	for pid, hist := range histStats {
		list = append(list, &sortHist{pid, hist})
	}
	sort.Sort(ByMax(list))

	if len(list) > limit {
		list = list[:limit]
	}

	return list
}

func formatMem(num uint64) string {
	letter := string("K")

	num = num * 4
	if num >= 10000 {
		num = (num + 512) / 1024
		letter = "M"
		if num >= 10000 {
			num = (num + 512) / 1024
			letter = "G"
		}
	}
	return fmt.Sprintf("%d%s", num, letter)
}

func formatNum(num uint64) string {
	if num > 1000000 {
		return fmt.Sprintf("%dM", num/1000000)
	}
	if num > 1000 {
		return fmt.Sprintf("%dK", num/1000)
	}
	return fmt.Sprintf("%d", num)
}

func dumpStats(cmdNames cmdlineMap, list pidlist, sumStats procStatsMap, histStats procStatsHistMap,
	sysSum *systemStats, sysHist *systemStatsHist, sumSched schedStatsMap, jiffy, interval, samples *int) {

	scale := func(val float64) float64 {
		return val / float64(*jiffy) / float64(*interval) * 1000 * 100
	}
	scaleSum := func(val float64, count int64) float64 {
		valSec := val / float64(*jiffy)
		sampleSec := float64(*interval) * float64(count) / 1000.0
		ret := (valSec / sampleSec) * 100
		return ret
	}
	scaleSched := func(val float64) float64 {
		return val / float64(*jiffy) / float64((*interval)*(*samples)) * 100
	}

	fmt.Printf("usr:    %4s/%4s/%4s   sys:%4s/%4s/%4s  nice:%4s/%4s/%4s    idle:%4s/%4s/%4s\n",
		trim(scale(float64(sysHist.usr.Min())), 4),
		trim(scale(float64(sysHist.usr.Max())), 4),
		trim(scale(sysHist.usr.Mean()), 4),

		trim(scale(float64(sysHist.sys.Min())), 4),
		trim(scale(float64(sysHist.sys.Max())), 4),
		trim(scale(sysHist.sys.Mean()), 4),

		trim(scale(float64(sysHist.nice.Min())), 4),
		trim(scale(float64(sysHist.nice.Max())), 4),
		trim(scale(sysHist.nice.Mean()), 4),

		trim(scale(float64(sysHist.idle.Min())), 4),
		trim(scale(float64(sysHist.idle.Max())), 4),
		trim(scale(sysHist.idle.Mean()), 4),
	)
	fmt.Printf("iowait: %4s/%4s/%4s  ctxt:%4s/%4s/%4s  prun:%4s/%4s/%4s  pblock:%4s/%4s/%4s  pstart: %4d\n",
		trim(scale(float64(sysHist.iowait.Min())), 4),
		trim(scale(float64(sysHist.iowait.Max())), 4),
		trim(scale(sysHist.iowait.Mean()), 4),

		trim(scale(float64(sysHist.ctxt.Min())), 4),
		trim(scale(float64(sysHist.ctxt.Max())), 4),
		trim(scale(sysHist.ctxt.Mean()), 4),

		trim(float64(sysHist.procsRunning.Min()), 4),
		trim(float64(sysHist.procsRunning.Max()), 4),
		trim(sysHist.procsRunning.Mean(), 4),

		trim(float64(sysHist.procsBlocked.Min()), 4),
		trim(float64(sysHist.procsBlocked.Max()), 4),
		trim(sysHist.procsBlocked.Mean(), 4),

		sysSum.procsTotal,
	)

	fmt.Print("                   comm     pid     min     max     usr     sys   nice   ctime    slat   ctx   icx   rss   iow  thrd  sam\n")
	for _, pid := range list {
		hist := histStats[pid]

		var schedWait, nrSwitches, nrInvoluntarySwitches string
		sched, ok := sumSched[pid]
		if ok == true {
			schedWait = trim(scaleSched(sched.waitSum), 7)
			nrSwitches = formatNum(sched.nrSwitches)
			nrInvoluntarySwitches = formatNum(sched.nrInvoluntarySwitches)
		} else {
			schedWait = "-"
			nrSwitches = "-"
			nrInvoluntarySwitches = "-"
		}
		sampleCount := hist.utime.TotalCount()
		var friendlyName string
		cmdName, ok := cmdNames[pid]
		if ok == true && len(cmdName.friendly) > 0 {
			friendlyName = cmdName.friendly
		} else {
			// This should not happen as long as the cmdline resolver works
			fmt.Println("using comm for ", cmdName, pid)
			friendlyName = sumStats[pid].comm
		}
		if len(friendlyName) > 23 {
			friendlyName = friendlyName[:23]
		}

		fmt.Printf("%23s %7d %7s %7s %7s %7s %6s %7s %7s %5s %5s %5s %5s %5d %4d\n",
			friendlyName,
			pid,
			trim(scale(float64(hist.ustime.Min())), 7),
			trim(scale(float64(hist.ustime.Max())), 7),
			trim(scaleSum(float64(sumStats[pid].utime), sampleCount), 7),
			trim(scaleSum(float64(sumStats[pid].stime), sampleCount), 7),
			trim(float64(sumStats[pid].nice), 6),
			trim(scaleSum(float64(sumStats[pid].cutime+sumStats[pid].cstime), sampleCount), 7),
			schedWait,
			nrSwitches,
			nrInvoluntarySwitches,
			formatMem(sumStats[pid].rss),
			trim(scaleSum(float64(sumStats[pid].delayacctBlkioTicks), sampleCount), 7),
			sumStats[pid].numThreads,
			sampleCount,
		)
	}
	fmt.Println()
}
<|endoftext|>


<commit_before>package goose

import (
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/PuerkitoBio/goquery"
)

// Crawler can fetch the target HTML page
type Crawler struct {
	config  Configuration
	Charset string
}

// NewCrawler returns a crawler object initialised with the URL and the [optional] raw HTML body
func NewCrawler(config Configuration) Crawler {
	return Crawler{
		config:  config,
		Charset: "",
	}
}

func getCharsetFromContentType(cs string) string {
	cs = strings.ToLower(strings.Replace(cs, " ", "", -1))
	cs = strings.TrimPrefix(cs, "text/html;charset=")
	cs = strings.TrimPrefix(cs, "text/xhtml;charset=")
	cs = strings.TrimPrefix(cs, "application/xhtml+xml;charset=")
	return NormaliseCharset(cs)
}

// SetCharset can be used to force a charset (e.g. when read from the HTTP headers)
// rather than relying on the detection from the HTML meta tags
func (c *Crawler) SetCharset(cs string) {
	c.Charset = getCharsetFromContentType(cs)
}

// GetContentType returns the Content-Type string extracted from the meta tags
func (c Crawler) GetContentType(document *goquery.Document) string {
	var attr string
	// <meta http-equiv="Content-Type" content="text/html; charset=utf-8" />
	document.Find("meta[http-equiv#=(?i)^Content\\-type$]").Each(func(i int, s *goquery.Selection) {
		attr, _ = s.Attr("content")
	})
	return attr
}

// GetCharset returns a normalised charset string extracted from the meta tags
func (c Crawler) GetCharset(document *goquery.Document) string {
	// manually-provided charset (from HTTP headers?) takes priority
	if "" != c.Charset {
		return c.Charset
	}

	// <meta http-equiv="Content-Type" content="text/html; charset=utf-8" />
	ct := c.GetContentType(document)
	if "" != ct && strings.Contains(strings.ToLower(ct), "charset") {
		return getCharsetFromContentType(ct)
	}

	// <meta charset="utf-8">
	selection := document.Find("meta").EachWithBreak(func(i int, s *goquery.Selection) bool {
		_, exists := s.Attr("charset")
		return !exists
	})

	if selection != nil {
		cs, _ := selection.Attr("charset")
		return NormaliseCharset(cs)
	}

	return ""
}

// Preprocess fetches the HTML page if needed, converts it to UTF-8 and applies
// some text normalisation to guarantee better results when extracting the content
func (c *Crawler) Preprocess(RawHTML string) (*goquery.Document, error) {
	var err error

	if RawHTML == "" {
		return nil, errors.New("cannot process empty HTML content")
	}

	RawHTML = c.addSpacesBetweenTags(RawHTML)

	reader := strings.NewReader(RawHTML)
	document, err := goquery.NewDocumentFromReader(reader)
	if err != nil {
		return nil, err
	}

	cs := c.GetCharset(document)
	//log.Println("-------------------------------------------CHARSET:", cs)
	if "" != cs && "UTF-8" != cs {
		// the net/html parser and goquery require UTF-8 data
		RawHTML = UTF8encode(RawHTML, cs)
		reader = strings.NewReader(RawHTML)
		if document, err = goquery.NewDocumentFromReader(reader); err != nil {
			return nil, err
		}
	}

	return document, nil
}

// Crawl fetches the HTML body and returns an Article
func (c Crawler) Crawl(RawHTML string, url string) (*Article, error) {
	article := new(Article)
	fmt.Println(0)
	document, err := c.Preprocess(RawHTML)
	fmt.Println(1)
	if nil != err {
		return nil, err
	}
	fmt.Println(2)
	if nil == document {
		return article, nil
	}
	fmt.Println(3)
	extractor := NewExtractor(c.config)
	fmt.Println(4)
	startTime := time.Now().UnixNano()

	article.RawHTML, err = document.Html()
	fmt.Println(5)
	if nil != err {
		return nil, err
	}
	fmt.Println(6)
	fmt.Println("len of url is ", string(len(article.RawHTML)))
	article.FinalURL = url
	article.Doc = document

	article.Title = extractor.GetTitle(document)
	article.MetaLang = extractor.GetMetaLanguage(document)
	article.MetaFavicon = extractor.GetFavicon(document)

	article.MetaDescription = extractor.GetMetaContentWithSelector(document, "meta[name#=(?i)^description$]")
	article.MetaKeywords = extractor.GetMetaContentWithSelector(document, "meta[name#=(?i)^keywords$]")
	article.CanonicalLink = extractor.GetCanonicalLink(document)
	if "" == article.CanonicalLink {
		article.CanonicalLink = article.FinalURL
	}
	article.Domain = extractor.GetDomain(article.CanonicalLink)
	article.Tags = extractor.GetTags(document)

	if c.config.extractPublishDate {
		if timestamp := extractor.GetPublishDate(document); timestamp != nil {
			article.PublishDate = timestamp
		}
	}

	cleaner := NewCleaner(c.config)
	article.Doc = cleaner.Clean(article.Doc)

	article.TopImage = OpenGraphResolver(document)
	if article.TopImage == "" {
		article.TopImage = WebPageResolver(article)
	}

	article.TopNode = extractor.CalculateBestNode(document)
	if article.TopNode != nil {
		article.TopNode = extractor.PostCleanup(article.TopNode)

		article.CleanedText, article.Links = extractor.GetCleanTextAndLinks(article.TopNode, article.MetaLang)

		videoExtractor := NewVideoExtractor()
		article.Movies = videoExtractor.GetVideos(document)
	}

	article.Delta = time.Now().UnixNano() - startTime

	return article, nil
}

// In many cases, like at the end of each
 element or between  tags,
// we need to add spaces, otherwise the text on either side will get joined together into one word.
// This method also adds newlines after each
 tag to preserve paragraphs.
func (c Crawler) addSpacesBetweenTags(text string) string {
	text = strings.Replace(text, "><", "> <", -1)
	text = strings.Replace(text, "
", "
\n", -1)
	text = strings.Replace(text, "<img ", "\n<img ", -1)
	text = strings.Replace(text, "
", "
\n", -1)
	return strings.Replace(text, "
", "
\n", -1)
}
<commit_msg>add more debugging<commit_after>package goose

import (
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/PuerkitoBio/goquery"
)

// Crawler can fetch the target HTML page
type Crawler struct {
	config  Configuration
	Charset string
}

// NewCrawler returns a crawler object initialised with the URL and the [optional] raw HTML body
func NewCrawler(config Configuration) Crawler {
	return Crawler{
		config:  config,
		Charset: "",
	}
}

func getCharsetFromContentType(cs string) string {
	cs = strings.ToLower(strings.Replace(cs, " ", "", -1))
	cs = strings.TrimPrefix(cs, "text/html;charset=")
	cs = strings.TrimPrefix(cs, "text/xhtml;charset=")
	cs = strings.TrimPrefix(cs, "application/xhtml+xml;charset=")
	return NormaliseCharset(cs)
}

// SetCharset can be used to force a charset (e.g. when read from the HTTP headers)
// rather than relying on the detection from the HTML meta tags
func (c *Crawler) SetCharset(cs string) {
	c.Charset = getCharsetFromContentType(cs)
}

// GetContentType returns the Content-Type string extracted from the meta tags
func (c Crawler) GetContentType(document *goquery.Document) string {
	var attr string
	// <meta http-equiv="Content-Type" content="text/html; charset=utf-8" />
	document.Find("meta[http-equiv#=(?i)^Content\\-type$]").Each(func(i int, s *goquery.Selection) {
		attr, _ = s.Attr("content")
	})
	return attr
}

// GetCharset returns a normalised charset string extracted from the meta tags
func (c Crawler) GetCharset(document *goquery.Document) string {
	// manually-provided charset (from HTTP headers?) takes priority
	if "" != c.Charset {
		return c.Charset
	}

	// <meta http-equiv="Content-Type" content="text/html; charset=utf-8" />
	ct := c.GetContentType(document)
	if "" != ct && strings.Contains(strings.ToLower(ct), "charset") {
		return getCharsetFromContentType(ct)
	}

	// <meta charset="utf-8">
	selection := document.Find("meta").EachWithBreak(func(i int, s *goquery.Selection) bool {
		_, exists := s.Attr("charset")
		return !exists
	})

	if selection != nil {
		cs, _ := selection.Attr("charset")
		return NormaliseCharset(cs)
	}

	return ""
}

// Preprocess fetches the HTML page if needed, converts it to UTF-8 and applies
// some text normalisation to guarantee better results when extracting the content
func (c *Crawler) Preprocess(RawHTML string) (*goquery.Document, error) {
	var err error

	if RawHTML == "" {
		return nil, errors.New("cannot process empty HTML content")
	}

	RawHTML = c.addSpacesBetweenTags(RawHTML)

	reader := strings.NewReader(RawHTML)
	document, err := goquery.NewDocumentFromReader(reader)
	if err != nil {
		return nil, err
	}

	cs := c.GetCharset(document)
	//log.Println("-------------------------------------------CHARSET:", cs)
	if "" != cs && "UTF-8" != cs {
		// the net/html parser and goquery require UTF-8 data
		RawHTML = UTF8encode(RawHTML, cs)
		reader = strings.NewReader(RawHTML)
		if document, err = goquery.NewDocumentFromReader(reader); err != nil {
			return nil, err
		}
	}

	return document, nil
}

// Crawl fetches the HTML body and returns an Article
func (c Crawler) Crawl(RawHTML string, url string) (*Article, error) {
	article := new(Article)
	fmt.Println(0)
	document, err := c.Preprocess(RawHTML)
	fmt.Println(1)
	if nil != err {
		return nil, err
	}
	fmt.Println(2)
	if nil == document {
		return article, nil
	}
	fmt.Println(3)
	extractor := NewExtractor(c.config)
	fmt.Println(4)
	startTime := time.Now().UnixNano()

	article.RawHTML, err = document.Html()
	fmt.Println(5)
	if nil != err {
		return nil, err
	}
	fmt.Println(6)
	fmt.Println(len(article.RawHTML))
	fmt.Println(article.RawHTML)
	article.FinalURL = url
	article.Doc = document

	article.Title = extractor.GetTitle(document)
	article.MetaLang = extractor.GetMetaLanguage(document)
	article.MetaFavicon = extractor.GetFavicon(document)

	article.MetaDescription = extractor.GetMetaContentWithSelector(document, "meta[name#=(?i)^description$]")
	fmt.Println(7)
	fmt.Println(len(article.MetaDescription))
	article.MetaKeywords = extractor.GetMetaContentWithSelector(document, "meta[name#=(?i)^keywords$]")
	article.CanonicalLink = extractor.GetCanonicalLink(document)
	if "" == article.CanonicalLink {
		article.CanonicalLink = article.FinalURL
	}
	article.Domain = extractor.GetDomain(article.CanonicalLink)
	article.Tags = extractor.GetTags(document)

	if c.config.extractPublishDate {
		if timestamp := extractor.GetPublishDate(document); timestamp != nil {
			article.PublishDate = timestamp
		}
	}

	cleaner := NewCleaner(c.config)
	article.Doc = cleaner.Clean(article.Doc)

	article.TopImage = OpenGraphResolver(document)
	if article.TopImage == "" {
		article.TopImage = WebPageResolver(article)
	}

	article.TopNode = extractor.CalculateBestNode(document)
	if article.TopNode != nil {
		article.TopNode = extractor.PostCleanup(article.TopNode)

		article.CleanedText, article.Links = extractor.GetCleanTextAndLinks(article.TopNode, article.MetaLang)

		videoExtractor := NewVideoExtractor()
		article.Movies = videoExtractor.GetVideos(document)
	}
	fmt.Println(len(article.CleanedText))

	article.Delta = time.Now().UnixNano() - startTime

	return article, nil
}

// In many cases, like at the end of each
 element or between  tags,
// we need to add spaces, otherwise the text on either side will get joined together into one word.
// This method also adds newlines after each
 tag to preserve paragraphs.
func (c Crawler) addSpacesBetweenTags(text string) string {
	text = strings.Replace(text, "><", "> <", -1)
	text = strings.Replace(text, "
", "
\n", -1)
	text = strings.Replace(text, "<img ", "\n<img ", -1)
	text = strings.Replace(text, "
", "
\n", -1)
	return strings.Replace(text, "
", "
\n", -1)
}
<|endoftext|>


<commit_before>package vcsstore

import (
	"crypto/sha1"
	"encoding/base64"
	"errors"
	"fmt"
	"log"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/sourcegraph/go-vcs/vcs"
)

var (
	// RepositoryPath is called to determine the directory, relative to a
	// Config's StorageDir, to which the repository should be cloned to. The
	// default implementation stores repositories in directories of the form
	// "vcs-type/escaped-clone-url".
	RepositoryPath = func(vcsType string, cloneURL *url.URL) string {
		return filepath.Join(vcsType, url.QueryEscape(cloneURL.String()))
	}
)

// HashedRepositoryPath may be assigned to RepositoryPath to use paths
// of the form "xx/yy/zzzzzzzz" where xx and yy are the first 4 characters of
// some hash of vcsType and cloneURL, and zzzzzzzz is the full hash (minus the
// first 4 characters).
func HashedRepositoryPath(vcsType string, cloneURL *url.URL) string {
	h := sha1.New()
	h.Write([]byte(vcsType))
	h.Write([]byte(cloneURL.String()))
	s := base64.URLEncoding.EncodeToString(h.Sum(nil))
	return fmt.Sprintf("%s/%s/%s/%s", vcsType, s[:2], s[2:4], s[4:])
}

type Service interface {
	Open(vcs string, cloneURL *url.URL) (interface{}, error)
}

type Config struct {
	// StorageDir is where cloned repositories are stored. If empty, the current
	// working directory is used.
	StorageDir string

	Log *log.Logger
}

func NewService(c *Config) Service {
	if c == nil {
		c = &Config{
			StorageDir: ".",
			Log:        log.New(os.Stderr, "vcsstore: ", log.LstdFlags),
		}
	}
	return &service{
		Config: *c,
		repoMu: make(map[repoKey]*sync.Mutex),
	}
}

type service struct {
	Config

	// repoMu synchronizes access to repository data on the filesystem.
	repoMu map[repoKey]*sync.Mutex

	// repoMuMu synchronizes access to repoMu.
	repoMuMu sync.Mutex
}

type repoKey struct {
	vcsType  string
	cloneURL string
}

func (s *service) Open(vcsType string, cloneURL *url.URL) (interface{}, error) {
	if !isLowercaseLetter(vcsType) {
		return nil, errors.New("invalid VCS type")
	}
	if cloneURL.Scheme == "" || cloneURL.Host == "" {
		return nil, errors.New("invalid clone URL")
	}

	cloneDir := filepath.Join(s.StorageDir, RepositoryPath(vcsType, cloneURL))

	mu := s.Mutex(vcsType, cloneURL)
	mu.Lock()
	defer mu.Unlock()

	_, err := os.Stat(cloneDir)
	if os.IsNotExist(err) {
		// Repository hasn't been cloned locally. Try cloning and opening it.
		if err := os.MkdirAll(filepath.Join(s.StorageDir, vcsType), 0700); err != nil {
			return nil, err
		}
		start := time.Now()
		msg := fmt.Sprintf("%s %s to %s", vcsType, cloneURL.String(), cloneDir)
		s.Log.Print("Cloning ", msg, "...")
		time.Sleep(time.Second * 2)
		defer s.Log.Print("Finished cloning ", msg, " in ", time.Since(start))
		return vcs.CloneMirror(vcsType, cloneURL.String(), cloneDir)
	} else if err != nil {
		return nil, err
	}

	return vcs.OpenMirror(vcsType, cloneDir)
}

func (s *service) Mutex(vcsType string, cloneURL *url.URL) *sync.Mutex {
	s.repoMuMu.Lock()
	defer s.repoMuMu.Unlock()

	k := repoKey{vcsType, cloneURL.String()}
	if mu, ok := s.repoMu[k]; ok {
		return mu
	}
	s.repoMu[k] = &sync.Mutex{}
	return s.repoMu[k]
}

func isLowercaseLetter(s string) bool {
	return strings.IndexFunc(s, func(c rune) bool {
		return !(c >= 'a' && c <= 'z')
	}) == -1
}
<commit_msg>rm sleep<commit_after>package vcsstore

import (
	"crypto/sha1"
	"encoding/base64"
	"errors"
	"fmt"
	"log"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/sourcegraph/go-vcs/vcs"
)

var (
	// RepositoryPath is called to determine the directory, relative to a
	// Config's StorageDir, to which the repository should be cloned to. The
	// default implementation stores repositories in directories of the form
	// "vcs-type/escaped-clone-url".
	RepositoryPath = func(vcsType string, cloneURL *url.URL) string {
		return filepath.Join(vcsType, url.QueryEscape(cloneURL.String()))
	}
)

// HashedRepositoryPath may be assigned to RepositoryPath to use paths
// of the form "xx/yy/zzzzzzzz" where xx and yy are the first 4 characters of
// some hash of vcsType and cloneURL, and zzzzzzzz is the full hash (minus the
// first 4 characters).
func HashedRepositoryPath(vcsType string, cloneURL *url.URL) string {
	h := sha1.New()
	h.Write([]byte(vcsType))
	h.Write([]byte(cloneURL.String()))
	s := base64.URLEncoding.EncodeToString(h.Sum(nil))
	return fmt.Sprintf("%s/%s/%s/%s", vcsType, s[:2], s[2:4], s[4:])
}

type Service interface {
	Open(vcs string, cloneURL *url.URL) (interface{}, error)
}

type Config struct {
	// StorageDir is where cloned repositories are stored. If empty, the current
	// working directory is used.
	StorageDir string

	Log *log.Logger
}

func NewService(c *Config) Service {
	if c == nil {
		c = &Config{
			StorageDir: ".",
			Log:        log.New(os.Stderr, "vcsstore: ", log.LstdFlags),
		}
	}
	return &service{
		Config: *c,
		repoMu: make(map[repoKey]*sync.Mutex),
	}
}

type service struct {
	Config

	// repoMu synchronizes access to repository data on the filesystem.
	repoMu map[repoKey]*sync.Mutex

	// repoMuMu synchronizes access to repoMu.
	repoMuMu sync.Mutex
}

type repoKey struct {
	vcsType  string
	cloneURL string
}

func (s *service) Open(vcsType string, cloneURL *url.URL) (interface{}, error) {
	if !isLowercaseLetter(vcsType) {
		return nil, errors.New("invalid VCS type")
	}
	if cloneURL.Scheme == "" || cloneURL.Host == "" {
		return nil, errors.New("invalid clone URL")
	}

	cloneDir := filepath.Join(s.StorageDir, RepositoryPath(vcsType, cloneURL))

	mu := s.Mutex(vcsType, cloneURL)
	mu.Lock()
	defer mu.Unlock()

	_, err := os.Stat(cloneDir)
	if os.IsNotExist(err) {
		// Repository hasn't been cloned locally. Try cloning and opening it.
		if err := os.MkdirAll(filepath.Join(s.StorageDir, vcsType), 0700); err != nil {
			return nil, err
		}
		start := time.Now()
		msg := fmt.Sprintf("%s %s to %s", vcsType, cloneURL.String(), cloneDir)
		s.Log.Print("Cloning ", msg, "...")
		defer s.Log.Print("Finished cloning ", msg, " in ", time.Since(start))
		return vcs.CloneMirror(vcsType, cloneURL.String(), cloneDir)
	} else if err != nil {
		return nil, err
	}

	return vcs.OpenMirror(vcsType, cloneDir)
}

func (s *service) Mutex(vcsType string, cloneURL *url.URL) *sync.Mutex {
	s.repoMuMu.Lock()
	defer s.repoMuMu.Unlock()

	k := repoKey{vcsType, cloneURL.String()}
	if mu, ok := s.repoMu[k]; ok {
		return mu
	}
	s.repoMu[k] = &sync.Mutex{}
	return s.repoMu[k]
}

func isLowercaseLetter(s string) bool {
	return strings.IndexFunc(s, func(c rune) bool {
		return !(c >= 'a' && c <= 'z')
	}) == -1
}
<|endoftext|>


<commit_before>package goaspen

import (
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"syscall"

	"github.com/jteeuwen/go-pkg-optarg"
)

type serverWrapper struct {
	PackageName string
	ServerBind  string
	WwwRoot     string
}

func AddCommonServingOptions(serverBind, wwwRoot string, debug bool) {
	optarg.Add("w", "www_root",
		"Filesystem path of the document publishing root", wwwRoot)
	optarg.Add("a", "network_address", "The IPv4 or IPv6 address to which "+
		"the generated server will bind by default", serverBind)
	optarg.Add("x", "debug", "Print debugging output", debug)
}

func RunServerMain(wwwRoot, serverBind, packageName string) {
	debug := false

	AddCommonServingOptions(serverBind, wwwRoot, debug)
	for opt := range optarg.Parse() {
		switch opt.Name {
		case "network_address":
			serverBind = opt.String()
		case "www_root":
			wwwRoot = opt.String()
		case "debug":
			debug = opt.Bool()
		}
	}

	SetDebug(debug)

	server := newServerWrapper(packageName, serverBind, wwwRoot)

	err := server.Run()
	if err != nil {
		log.Fatal(err)
	}
}

func newServerWrapper(packageName, serverBind, wwwRoot string) *serverWrapper {
	return &serverWrapper{
		PackageName: packageName,
		ServerBind:  serverBind,
		WwwRoot:     wwwRoot,
	}
}

func (me *serverWrapper) serverQuitListener() {
	ch := make(chan os.Signal)
	signal.Notify(ch, syscall.SIGQUIT)
	<-ch
	log.Println("Received SIGQUIT; exiting")
	os.Exit(0)
}

func (me *serverWrapper) Run() error {
	go me.serverQuitListener()

	err := expandAllHandlerFuncRegistrations()
	if err != nil {
		return err
	}

	err = registerAllHandlerFuncs()
	if err != nil {
		return err
	}

	fmt.Printf("%s-server serving on %q\n", me.PackageName, me.ServerBind)
	return http.ListenAndServe(me.ServerBind, nil)
}
<commit_msg>TRIVIAL correcting reported server binary name<commit_after>package goaspen

import (
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"syscall"

	"github.com/jteeuwen/go-pkg-optarg"
)

type serverWrapper struct {
	PackageName string
	ServerBind  string
	WwwRoot     string
}

func AddCommonServingOptions(serverBind, wwwRoot string, debug bool) {
	optarg.Add("w", "www_root",
		"Filesystem path of the document publishing root", wwwRoot)
	optarg.Add("a", "network_address", "The IPv4 or IPv6 address to which "+
		"the generated server will bind by default", serverBind)
	optarg.Add("x", "debug", "Print debugging output", debug)
}

func RunServerMain(wwwRoot, serverBind, packageName string) {
	debug := false

	AddCommonServingOptions(serverBind, wwwRoot, debug)
	for opt := range optarg.Parse() {
		switch opt.Name {
		case "network_address":
			serverBind = opt.String()
		case "www_root":
			wwwRoot = opt.String()
		case "debug":
			debug = opt.Bool()
		}
	}

	SetDebug(debug)

	server := newServerWrapper(packageName, serverBind, wwwRoot)

	err := server.Run()
	if err != nil {
		log.Fatal(err)
	}
}

func newServerWrapper(packageName, serverBind, wwwRoot string) *serverWrapper {
	return &serverWrapper{
		PackageName: packageName,
		ServerBind:  serverBind,
		WwwRoot:     wwwRoot,
	}
}

func (me *serverWrapper) serverQuitListener() {
	ch := make(chan os.Signal)
	signal.Notify(ch, syscall.SIGQUIT)
	<-ch
	log.Println("Received SIGQUIT; exiting")
	os.Exit(0)
}

func (me *serverWrapper) Run() error {
	go me.serverQuitListener()

	err := expandAllHandlerFuncRegistrations()
	if err != nil {
		return err
	}

	err = registerAllHandlerFuncs()
	if err != nil {
		return err
	}

	fmt.Printf("%s-http-server serving on %q\n", me.PackageName, me.ServerBind)
	return http.ListenAndServe(me.ServerBind, nil)
}
<|endoftext|>


<commit_before>package cron

import (
	"github.com/Cepave/sender/g"
	"github.com/Cepave/sender/model"
	"github.com/Cepave/sender/proc"
	"github.com/Cepave/sender/redis"
	//"github.com/toolkits/net/httplib"
	"github.com/kardianos/osext"
	"log"
	"os/exec"
	"time"
)

func ConsumeQQ() {
	queue := g.Config().Queue.QQ
	for {
		L := redis.PopAllQQ(queue)
		if len(L) == 0 {
			time.Sleep(time.Millisecond * 200)
			continue
		}
		SendQQList(L)
	}
}

func SendQQList(L []*model.QQ) {
	for _, qq := range L {
		QQWorkerChan <- 1
		go SendQQ(qq)
	}
}

func SendQQ(qq *model.QQ) {
	defer func() {
		<-QQWorkerChan
	}()

	url := g.Config().Api.QQ
	exePath, _ := osext.Executable()
	cmd := exec.Command("/bin/bash", "./qq_sms.sh", url, qq.Subject, qq.Content)
	err := cmd.Run()
	if err != nil {
		log.Println(err)
	}

	proc.IncreQQCount()

	if g.Config().Debug {
		log.Println("==qq==>>>>", qq.Subject)
	}

}
<commit_msg>fixup<commit_after>package cron

import (
	"github.com/Cepave/sender/g"
	"github.com/Cepave/sender/model"
	"github.com/Cepave/sender/proc"
	"github.com/Cepave/sender/redis"
	//"github.com/toolkits/net/httplib"
	"log"
	"os/exec"
	"time"
)

func ConsumeQQ() {
	queue := g.Config().Queue.QQ
	for {
		L := redis.PopAllQQ(queue)
		if len(L) == 0 {
			time.Sleep(time.Millisecond * 200)
			continue
		}
		SendQQList(L)
	}
}

func SendQQList(L []*model.QQ) {
	for _, qq := range L {
		QQWorkerChan <- 1
		go SendQQ(qq)
	}
}

func SendQQ(qq *model.QQ) {
	defer func() {
		<-QQWorkerChan
	}()

	url := g.Config().Api.QQ
	cmd := exec.Command("/bin/bash", "./qq_sms.sh", url, qq.Subject, qq.Content)
	err := cmd.Run()
	if err != nil {
		log.Println(err)
	}

	proc.IncreQQCount()

	if g.Config().Debug {
		log.Println("==qq==>>>>", qq.Subject)
	}

}
<|endoftext|>


<commit_before>package smux

import (
	"io"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pkg/errors"
)

const (
	defaultAcceptBacklog = 1024
	defaultCloseWait     = 1024
)

const (
	errBrokenPipe      = "broken pipe"
	errInvalidProtocol = "invalid protocol version"
)

// Session defines a multiplexed connection for streams
type Session struct {
	conn io.ReadWriteCloser

	nextStreamID uint32                   // next stream identifier
	streams      map[uint32]*Stream       // all streams in this session
	frameSize    uint16                   // max frame size
	rdEvents     map[uint32]chan struct{} // stream read notification

	tbf         chan struct{}      // tokenbuffer
	frameQueues map[uint32][]Frame // stream input frame queue

	die       chan struct{} // flag session has died
	chAccepts chan *Stream
	chClose   chan uint32
	dataReady int32 // flag data has arrived
	mu        sync.Mutex
}

func newSession(conn io.ReadWriteCloser, client bool, maxframes int, framesize uint16) *Session {
	s := new(Session)
	s.conn = conn
	s.frameSize = framesize
	s.streams = make(map[uint32]*Stream)
	s.tbf = make(chan struct{}, maxframes)
	s.frameQueues = make(map[uint32][]Frame)
	s.rdEvents = make(map[uint32]chan struct{})
	s.chAccepts = make(chan *Stream, defaultAcceptBacklog)
	s.chClose = make(chan uint32, defaultCloseWait)
	s.die = make(chan struct{})
	for i := 0; i < maxframes; i++ {
		s.tbf <- struct{}{}
	}
	if client {
		s.nextStreamID = 1
	} else {
		s.nextStreamID = 2
	}
	go s.recvLoop()
	go s.monitor()
	go s.keepalive()
	return s
}

// OpenStream is used to create a new stream
func (s *Session) OpenStream() (*Stream, error) {
	s.mu.Lock()
	sid := s.nextStreamID
	s.nextStreamID += 2
	chNotifyReader := make(chan struct{}, 1)
	stream := newStream(sid, s.frameSize, chNotifyReader, s)
	s.rdEvents[sid] = chNotifyReader
	s.streams[sid] = stream
	s.mu.Unlock()
	s.sendFrame(newFrame(cmdSYN, sid))
	return stream, nil
}

// AcceptStream is used to block until the next available stream
// is ready to be accepted.
func (s *Session) AcceptStream() (*Stream, error) {
	select {
	case stream := <-s.chAccepts:
		return stream, nil
	case <-s.die:
		return nil, errors.New(errBrokenPipe)
	}
}

// Close is used to close the session and all streams.
func (s *Session) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	select {
	case <-s.die:
		return errors.New(errBrokenPipe)
	default:
		close(s.die)
		for k := range s.streams {
			s.streams[k].Close()
		}
		s.sendFrame(newFrame(cmdTerminate, 0))
		s.conn.Close()
	}
	return nil
}

// NumStreams returns the number of currently open streams
func (s *Session) NumStreams() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.streams)
}

// notify the session that a session has closed
func (s *Session) streamClose(sid uint32) {
	s.chClose <- sid
}

// nonblocking read from session pool, for streams
func (s *Session) nioread(sid uint32) *Frame {
	s.mu.Lock()
	frames := s.frameQueues[sid]
	if len(frames) > 0 {
		f := frames[0]
		s.frameQueues[sid] = frames[1:]
		s.tbf <- struct{}{}
		s.mu.Unlock()
		return &f
	}
	s.mu.Unlock()
	return nil
}

// session read a frame from underlying connection
func (s *Session) readFrame(buffer []byte) (f Frame, err error) {
	if _, err := io.ReadFull(s.conn, buffer[:headerSize]); err != nil {
		return f, errors.Wrap(err, "readFrame")
	}

	dec := rawHeader(buffer)
	if dec.Version() != version {
		return f, errors.New(errInvalidProtocol)
	}

	if length := dec.Length(); length > 0 {
		if _, err := io.ReadFull(s.conn, buffer[headerSize:headerSize+length]); err != nil {
			return f, errors.Wrap(err, "readFrame")
		}
		f.UnmarshalBinary(buffer[:headerSize+length])
		return f, nil
	}
	f.UnmarshalBinary(buffer[:headerSize])
	return f, nil
}

// monitors streams
func (s *Session) monitor() {
	for {
		select {
		case sid := <-s.chClose:
			s.mu.Lock()
			delete(s.streams, sid)
			delete(s.rdEvents, sid)
			ntokens := len(s.frameQueues[sid])
			delete(s.frameQueues, sid)
			s.mu.Unlock()
			for i := 0; i < ntokens; i++ { // return remaining tokens to the pool
				s.tbf <- struct{}{}
			}
		case <-s.die:
			return
		}
	}
}

// recvLoop keeps on reading from underlying connection if tokens are available
func (s *Session) recvLoop() {
	buffer := make([]byte, (1<<16)+headerSize)
	for {
		select {
		case <-s.tbf:
			if f, err := s.readFrame(buffer); err == nil {
				s.mu.Lock()
				switch f.cmd {
				case cmdNOP:
					s.tbf <- struct{}{}
				case cmdTerminate:
					s.Close()
					return
				case cmdSYN:
					chNotifyReader := make(chan struct{}, 1)
					s.streams[f.sid] = newStream(f.sid, s.frameSize, chNotifyReader, s)
					s.rdEvents[f.sid] = chNotifyReader
					s.chAccepts <- s.streams[f.sid]
					s.tbf <- struct{}{}
				default:
					if _, ok := s.streams[f.sid]; ok {
						s.frameQueues[f.sid] = append(s.frameQueues[f.sid], f)
						select {
						case s.rdEvents[f.sid] <- struct{}{}:
						default:
						}
					} else {
						s.tbf <- struct{}{}
					}
				}
				s.mu.Unlock()
				atomic.StoreInt32(&s.dataReady, 1)
			} else {
				s.Close()
				return
			}
		case <-s.die:
			return
		}
	}
}

func (s *Session) keepalive() {
	tickerPing := time.NewTicker(10 * time.Second)
	tickerTimeout := time.NewTicker(20 * time.Second)
	defer tickerPing.Stop()
	defer tickerTimeout.Stop()
	for {
		select {
		case <-tickerPing.C:
			s.sendFrame(newFrame(cmdNOP, 0))
		case <-tickerTimeout.C:
			if !atomic.CompareAndSwapInt32(&s.dataReady, 1, 0) {
				s.Close()
				return
			}
		case <-s.die:
			return
		}
	}
}

func (s *Session) sendFrame(f Frame) {
	bts, _ := f.MarshalBinary()
	s.conn.Write(bts)
}
<commit_msg>add IsClosed<commit_after>package smux

import (
	"io"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pkg/errors"
)

const (
	defaultAcceptBacklog = 1024
	defaultCloseWait     = 1024
)

const (
	errBrokenPipe      = "broken pipe"
	errInvalidProtocol = "invalid protocol version"
)

// Session defines a multiplexed connection for streams
type Session struct {
	conn io.ReadWriteCloser

	nextStreamID uint32                   // next stream identifier
	streams      map[uint32]*Stream       // all streams in this session
	frameSize    uint16                   // max frame size
	rdEvents     map[uint32]chan struct{} // stream read notification

	tbf         chan struct{}      // tokenbuffer
	frameQueues map[uint32][]Frame // stream input frame queue

	die       chan struct{} // flag session has died
	chAccepts chan *Stream
	chClose   chan uint32
	dataReady int32 // flag data has arrived
	mu        sync.Mutex
}

func newSession(conn io.ReadWriteCloser, client bool, maxframes int, framesize uint16) *Session {
	s := new(Session)
	s.conn = conn
	s.frameSize = framesize
	s.streams = make(map[uint32]*Stream)
	s.tbf = make(chan struct{}, maxframes)
	s.frameQueues = make(map[uint32][]Frame)
	s.rdEvents = make(map[uint32]chan struct{})
	s.chAccepts = make(chan *Stream, defaultAcceptBacklog)
	s.chClose = make(chan uint32, defaultCloseWait)
	s.die = make(chan struct{})
	for i := 0; i < maxframes; i++ {
		s.tbf <- struct{}{}
	}
	if client {
		s.nextStreamID = 1
	} else {
		s.nextStreamID = 2
	}
	go s.recvLoop()
	go s.monitor()
	go s.keepalive()
	return s
}

// OpenStream is used to create a new stream
func (s *Session) OpenStream() (*Stream, error) {
	if s.IsClosed() {
		return nil, errors.New(errBrokenPipe)
	}

	sid := atomic.AddUint32(&s.nextStreamID, 2)
	chNotifyReader := make(chan struct{}, 1)
	stream := newStream(sid, s.frameSize, chNotifyReader, s)

	s.mu.Lock()
	s.rdEvents[sid] = chNotifyReader
	s.streams[sid] = stream
	s.mu.Unlock()

	s.sendFrame(newFrame(cmdSYN, sid))
	return stream, nil
}

// AcceptStream is used to block until the next available stream
// is ready to be accepted.
func (s *Session) AcceptStream() (*Stream, error) {
	select {
	case stream := <-s.chAccepts:
		return stream, nil
	case <-s.die:
		return nil, errors.New(errBrokenPipe)
	}
}

// Close is used to close the session and all streams.
func (s *Session) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	select {
	case <-s.die:
		return errors.New(errBrokenPipe)
	default:
		close(s.die)
		for k := range s.streams {
			s.streams[k].Close()
		}
		s.sendFrame(newFrame(cmdTerminate, 0))
		s.conn.Close()
	}
	return nil
}

// IsClosed does a safe check to see if we have shutdown
func (s *Session) IsClosed() bool {
	select {
	case <-s.die:
		return true
	default:
		return false
	}
}

// NumStreams returns the number of currently open streams
func (s *Session) NumStreams() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.streams)
}

// notify the session that a session has closed
func (s *Session) streamClose(sid uint32) {
	s.chClose <- sid
}

// nonblocking read from session pool, for streams
func (s *Session) nioread(sid uint32) *Frame {
	s.mu.Lock()
	frames := s.frameQueues[sid]
	if len(frames) > 0 {
		f := frames[0]
		s.frameQueues[sid] = frames[1:]
		s.tbf <- struct{}{}
		s.mu.Unlock()
		return &f
	}
	s.mu.Unlock()
	return nil
}

// session read a frame from underlying connection
func (s *Session) readFrame(buffer []byte) (f Frame, err error) {
	if _, err := io.ReadFull(s.conn, buffer[:headerSize]); err != nil {
		return f, errors.Wrap(err, "readFrame")
	}

	dec := rawHeader(buffer)
	if dec.Version() != version {
		return f, errors.New(errInvalidProtocol)
	}

	if length := dec.Length(); length > 0 {
		if _, err := io.ReadFull(s.conn, buffer[headerSize:headerSize+length]); err != nil {
			return f, errors.Wrap(err, "readFrame")
		}
		f.UnmarshalBinary(buffer[:headerSize+length])
		return f, nil
	}
	f.UnmarshalBinary(buffer[:headerSize])
	return f, nil
}

// monitors streams
func (s *Session) monitor() {
	for {
		select {
		case sid := <-s.chClose:
			s.mu.Lock()
			delete(s.streams, sid)
			delete(s.rdEvents, sid)
			ntokens := len(s.frameQueues[sid])
			delete(s.frameQueues, sid)
			s.mu.Unlock()
			for i := 0; i < ntokens; i++ { // return remaining tokens to the pool
				s.tbf <- struct{}{}
			}
		case <-s.die:
			return
		}
	}
}

// recvLoop keeps on reading from underlying connection if tokens are available
func (s *Session) recvLoop() {
	buffer := make([]byte, (1<<16)+headerSize)
	for {
		select {
		case <-s.tbf:
			if f, err := s.readFrame(buffer); err == nil {
				s.mu.Lock()
				switch f.cmd {
				case cmdNOP:
					s.tbf <- struct{}{}
				case cmdTerminate:
					s.Close()
					return
				case cmdSYN:
					chNotifyReader := make(chan struct{}, 1)
					s.streams[f.sid] = newStream(f.sid, s.frameSize, chNotifyReader, s)
					s.rdEvents[f.sid] = chNotifyReader
					s.chAccepts <- s.streams[f.sid]
					s.tbf <- struct{}{}
				default:
					if _, ok := s.streams[f.sid]; ok {
						s.frameQueues[f.sid] = append(s.frameQueues[f.sid], f)
						select {
						case s.rdEvents[f.sid] <- struct{}{}:
						default:
						}
					} else {
						s.tbf <- struct{}{}
					}
				}
				s.mu.Unlock()
				atomic.StoreInt32(&s.dataReady, 1)
			} else {
				s.Close()
				return
			}
		case <-s.die:
			return
		}
	}
}

func (s *Session) keepalive() {
	tickerPing := time.NewTicker(10 * time.Second)
	tickerTimeout := time.NewTicker(20 * time.Second)
	defer tickerPing.Stop()
	defer tickerTimeout.Stop()
	for {
		select {
		case <-tickerPing.C:
			s.sendFrame(newFrame(cmdNOP, 0))
		case <-tickerTimeout.C:
			if !atomic.CompareAndSwapInt32(&s.dataReady, 1, 0) {
				s.Close()
				return
			}
		case <-s.die:
			return
		}
	}
}

func (s *Session) sendFrame(f Frame) {
	bts, _ := f.MarshalBinary()
	s.conn.Write(bts)
}
<|endoftext|>


<commit_before>package darksky

import (
	"fmt"
	"github.com/google/go-querystring/query"
)

const (
	baseUrl = "https://api.darksky.net/forecast"
)

// DarkSky Api client
type DarkSky struct {
	BaseUrl string
}

// New creates a new DarkSky client
func New() *DarkSky {
	return &DarkSky{baseUrl}
}

// DarkSky.Forecast request a forecast
func (d *DarkSky) Forecast(request ForecastRequest) (ForecastResponse, error) {
	response := ForecastResponse{}

	url := d.buildRequestUrl(request)

	err := get(url, nil, &response)

	return response, err
}

func (d *DarkSky) buildRequestUrl(request ForecastRequest) string {
	url := fmt.Sprintf("%s/%s/%f,%f", d.BaseUrl, request.ApiKey, request.Latitude, request.Longitude)

	if request.Time != nil {
		url = url + fmt.Sprintf(",%d", request.Time)
	}

	values, _ := query.Values(request.Options)
	queryString := values.Encode()

	if len(queryString) > 0 {
		url = url + "?" + queryString
	}

	return url
}
<commit_msg>fix doc<commit_after>package darksky

import (
	"fmt"
	"github.com/google/go-querystring/query"
)

const (
	baseUrl = "https://api.darksky.net/forecast"
)

// DarkSky Api client
type DarkSky struct {
	BaseUrl string
}

// New creates a new DarkSky client
func New() *DarkSky {
	return &DarkSky{baseUrl}
}

// Forecast request a forecast
func (d *DarkSky) Forecast(request ForecastRequest) (ForecastResponse, error) {
	response := ForecastResponse{}

	url := d.buildRequestUrl(request)

	err := get(url, nil, &response)

	return response, err
}

func (d *DarkSky) buildRequestUrl(request ForecastRequest) string {
	url := fmt.Sprintf("%s/%s/%f,%f", d.BaseUrl, request.ApiKey, request.Latitude, request.Longitude)

	if request.Time != nil {
		url = url + fmt.Sprintf(",%d", request.Time)
	}

	values, _ := query.Values(request.Options)
	queryString := values.Encode()

	if len(queryString) > 0 {
		url = url + "?" + queryString
	}

	return url
}
<|endoftext|>


<commit_before>package main

import (
	"testing"
)

func TestPageListBack(t *testing.T) {
	_, artcs := PageListBack(0, "", false, false, 1, 20)
	for _, artc := range artcs {
		t.Log(*artc)
	}
	t.Log("------------------------------------------------------------")
	_, artcs = PageListBack(0, "", false, false, 2, 10)
	for _, artc := range artcs {
		t.Log(*artc)
	}
	t.Log("------------------------------------------------------------")
	_, artcs = PageListBack(3, "", false, false, 1, 20)
	for _, artc := range artcs {
		t.Log(*artc)
	}
	t.Log("------------------------------------------------------------")
	_, artcs = PageListBack(3, "19", false, false, 1, 20)
	for _, artc := range artcs {
		t.Log(*artc)
	}
	t.Log("------------------------------------------------------------")
	_, artcs = PageListBack(0, "", false, true, 1, 20)
	for _, artc := range artcs {
		t.Log(*artc)
	}
	t.Log("------------------------------------------------------------")
	_, artcs = PageListBack(0, "", true, false, 1, 20)
	for _, artc := range artcs {
		t.Log(*artc)
	}
}

func TestAddSerie(t *testing.T) {
	err := AddSerie("测试", "nothing", "这里是描述")
	if err != nil {
		t.Error(err)
	}
}
<commit_msg>update<commit_after>package main

import (
	"html/template"
	"testing"
)

func TestPageListBack(t *testing.T) {
	_, artcs := PageListBack(0, "", false, false, 1, 20)
	for _, artc := range artcs {
		t.Log(*artc)
	}
	t.Log("------------------------------------------------------------")
	_, artcs = PageListBack(0, "", false, false, 2, 10)
	for _, artc := range artcs {
		t.Log(*artc)
	}
	t.Log("------------------------------------------------------------")
	_, artcs = PageListBack(3, "", false, false, 1, 20)
	for _, artc := range artcs {
		t.Log(*artc)
	}
	t.Log("------------------------------------------------------------")
	_, artcs = PageListBack(3, "19", false, false, 1, 20)
	for _, artc := range artcs {
		t.Log(*artc)
	}
	t.Log("------------------------------------------------------------")
	_, artcs = PageListBack(0, "", false, true, 1, 20)
	for _, artc := range artcs {
		t.Log(*artc)
	}
	t.Log("------------------------------------------------------------")
	_, artcs = PageListBack(0, "", true, false, 1, 20)
	for _, artc := range artcs {
		t.Log(*artc)
	}
}

func TestAddSerie(t *testing.T) {
	err := AddSerie("测试", "nothing", "这里是描述")
	if err != nil {
		t.Error(err)
	}
}

func TestRenderPage(t *testing.T) {
	data := []byte(`

一人游走「不错的小伙子」

Leo同学「小伙子，该干活了」

razeen同学「Stay hungry. Stay foolish.」

Mirreal Ellison「kissing the fire」

`)

	t.Log(IgnoreHtmlTag(string(data)))
	data = renderPage(data)

	t.Log(template.HTML(string(data)))
}
<|endoftext|>


<commit_before>// Command deplist prints a detailed list of package imports for a package.
// The output contains the path to the source on the GOPATH for those imported packages.
package main

import (
	"flag"
	"fmt"
	"go/build"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"
	"text/tabwriter"
)

var (
	tags   = flag.String("tags", "", "comma-separated list of build tags to apply")
	goroot = flag.Bool("goroot", false, "include imports in GOROOT")
	tsv    = flag.Bool("tsv", false, "use only a single tab between columns")
)

type flusher interface {
	Flush() error
}

func main() {
	flag.Usage = func() {
		fmt.Fprintln(os.Stderr, "usage: pkglist [-tags] [-goroot] <dirs...>")
		flag.PrintDefaults()
		os.Exit(1)
	}

	flag.Parse()

	if flag.NArg() == 0 {
		flag.Usage()
	}

	buildctx := build.Default
	buildctx.BuildTags = strings.Split(*tags, ",")

	visited := make(map[imp]bool)
	var imports []imp

	for _, dir := range flag.Args() {
		abs, err := filepath.Abs(dir)
		if err != nil {
			log.Fatalf("could not get absolute path for dir %q: %v", dir, err)
		}

		pkg, err := buildctx.ImportDir(abs, 0)
		if err != nil {
			log.Fatalf("could not get package for dir %q: %v", dir, err)
		}

		for _, importPath := range pkg.Imports {
			i := imp{path: importPath, from: abs}
			imports = append(imports, i)
		}
	}

	var w io.Writer = os.Stdout
	if !*tsv {
		w = tabwriter.NewWriter(os.Stdout, 8, 0, 4, ' ', 0)
	}

	for len(imports) != 0 {
		i := imports[0]
		imports = imports[1:] // shift

		if _, ok := visited[i]; ok || i.path == "C" {
			continue
		}
		visited[i] = true

		pkg, err := buildctx.Import(i.path, i.from, 0)
		if err != nil {
			log.Fatalf("could not get package %q, imported from %q: %v", i.path, i.from, err)
		}

		if !*goroot && pkg.Goroot {
			continue
		}

		fmt.Fprintf(w, "%s\t%s\t%s\t%s\n", i.from, i.path, pkg.SrcRoot, pkg.ImportPath)

		for _, i := range pkg.Imports {
			imports = append(imports, imp{path: i, from: pkg.Dir})
		}
	}
	if f, ok := w.(flusher); ok {
		f.Flush()
	}
}

type imp struct {
	path, from string
}
<commit_msg>deplist: clean up<commit_after>// Command deplist prints a detailed list of package imports for a package.
// The output contains the path to the source on the GOPATH for those imported packages.
package main

import (
	"flag"
	"fmt"
	"go/build"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"
	"text/tabwriter"
)

var (
	tags   = flag.String("tags", "", "comma-separated list of build tags to apply")
	goroot = flag.Bool("goroot", false, "include imports in GOROOT")
	tsv    = flag.Bool("tsv", false, "use only a single tab between columns")
)

type flusher interface {
	Flush() error
}

type importFrom struct {
	path    string
	fromDir string
}

func main() {
	flag.Usage = func() {
		fmt.Fprintln(os.Stderr, "usage: deplist [-tags] [-goroot] <dirs...>")
		flag.PrintDefaults()
		os.Exit(1)
	}

	flag.Parse()

	if flag.NArg() == 0 {
		flag.Usage()
	}

	buildctx := build.Default
	buildctx.BuildTags = strings.Split(*tags, ",")

	var w io.Writer = os.Stdout
	if !*tsv {
		w = tabwriter.NewWriter(os.Stdout, 8, 0, 4, ' ', 0)
	}

	visited := make(map[importFrom]bool)
	var imports []importFrom

	for _, dir := range flag.Args() {
		abs, err := filepath.Abs(dir)
		if err != nil {
			log.Fatalf("could not get absolute path for dir %q: %v", dir, err)
		}

		pkg, err := buildctx.ImportDir(abs, 0)
		if err != nil {
			log.Fatalf("could not get package for dir %q: %v", dir, err)
		}

		for _, importPath := range pkg.Imports {
			imports = append(imports, importFrom{
				path:    importPath,
				fromDir: abs,
			})
		}
	}

	for len(imports) != 0 {
		i := imports[0]
		imports = imports[1:] // shift

		if _, ok := visited[i]; ok || i.path == "C" {
			continue
		}
		visited[i] = true

		pkg, err := buildctx.Import(i.path, i.fromDir, 0)
		if err != nil {
			log.Fatalf("could not get package %q, imported from %q: %v", i.path, i.fromDir, err)
		}

		if !*goroot && pkg.Goroot {
			continue
		}

		fmt.Fprintf(w, "%s\t%s\t%s\t%s\n", i.fromDir, i.path, pkg.SrcRoot, pkg.ImportPath)

		for _, importPath := range pkg.Imports {
			imports = append(imports, importFrom{
				path:    importPath,
				fromDir: pkg.Dir,
			})
		}
	}

	if f, ok := w.(flusher); ok {
		f.Flush()
	}
}
<|endoftext|>


<commit_before>package gopus

// #cgo pkg-config: opus
// #include "gopus.h"
import "C"

import (
	"unsafe"
)

type Decoder struct {
	data     []byte
	cDecoder *C.struct_OpusDecoder
	channels int
}

func NewDecoder(sampleRate, channels int) (*Decoder, error) {
	decoder := &Decoder{}
	decoder.data = make([]byte, int(C.opus_decoder_get_size(C.int(channels))))
	decoder.cDecoder = (*C.struct_OpusDecoder)(unsafe.Pointer(&decoder.data[0]))

	ret := C.opus_decoder_init(decoder.cDecoder, C.opus_int32(sampleRate), C.int(channels))
	if err := getErr(ret); err != nil {
		return nil, err
	}
	decoder.channels = channels

	return decoder, nil
}

func (d *Decoder) Decode(data []byte, frameSize int, fec bool) ([]int16, error) {
	dataPtr := (*C.uchar)(unsafe.Pointer(&data[0]))
	dataLen := C.opus_int32(len(data))

	output := make([]int16, d.channels * frameSize)
	outputPtr := (*C.opus_int16)(unsafe.Pointer(&output[0]))

	var cFec C.int
	if fec {
		cFec = 1
	} else {
		cFec = 0
	}

	cRet := C.opus_decode(d.cDecoder, dataPtr, dataLen, outputPtr, C.int(frameSize), cFec)
	ret := int(cRet)

	if ret < 0 {
		return nil, getErr(cRet)
	}
	return output[:ret], nil
}
<commit_msg>add CountFrames<commit_after>package gopus

// #cgo pkg-config: opus
// #include "gopus.h"
import "C"

import (
	"unsafe"
)

type Decoder struct {
	data     []byte
	cDecoder *C.struct_OpusDecoder
	channels int
}

func NewDecoder(sampleRate, channels int) (*Decoder, error) {
	decoder := &Decoder{}
	decoder.data = make([]byte, int(C.opus_decoder_get_size(C.int(channels))))
	decoder.cDecoder = (*C.struct_OpusDecoder)(unsafe.Pointer(&decoder.data[0]))

	ret := C.opus_decoder_init(decoder.cDecoder, C.opus_int32(sampleRate), C.int(channels))
	if err := getErr(ret); err != nil {
		return nil, err
	}
	decoder.channels = channels

	return decoder, nil
}

func (d *Decoder) Decode(data []byte, frameSize int, fec bool) ([]int16, error) {
	dataPtr := (*C.uchar)(unsafe.Pointer(&data[0]))
	dataLen := C.opus_int32(len(data))

	output := make([]int16, d.channels * frameSize)
	outputPtr := (*C.opus_int16)(unsafe.Pointer(&output[0]))

	var cFec C.int
	if fec {
		cFec = 1
	} else {
		cFec = 0
	}

	cRet := C.opus_decode(d.cDecoder, dataPtr, dataLen, outputPtr, C.int(frameSize), cFec)
	ret := int(cRet)

	if ret < 0 {
		return nil, getErr(cRet)
	}
	return output[:ret], nil
}

func CountFrames(data []byte) (int, error) {
	dataPtr := (*C.uchar)(unsafe.Pointer(&data[0]))
	cLen := C.opus_int32(len(data))

	cRet := C.opus_packet_get_nb_frames(dataPtr, cLen)
	if err := getErr(cRet); err != nil {
		return 0, err
	}
	return int(cRet), nil
}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"time"
	"sync"
	"github.com/codegangsta/cli"
	"github.com/fsouza/go-dockerclient"
)

func collectAPIImages(images []docker.APIImages, client *docker.Client, ctx *cli.Context) {
	var imageSync sync.WaitGroup
	for _, image := range images {
		imageSync.Add(1)
		go func(image *docker.APIImages, client *docker.Client, grace time.Duration, output bool) {
			defer imageSync.Done()
			imageDetail, _ := client.InspectImage(image.ID)
			handleImage(imageDetail, client, grace, output)
		}(&image, client, ctx.Duration("grace"), ctx.BoolT("quiet"))
	}
	imageSync.Wait()
}

func handleImage(image *docker.Image, client *docker.Client, grace time.Duration, output bool) {
}

func collectAPIContainers(containers []docker.APIContainers, client *docker.Client, ctx *cli.Context) {
	var containerSync sync.WaitGroup
	for _, container := range containers {
		containerSync.Add(1)
		go func(container *docker.APIContainers, client *docker.Client, grace time.Duration, output bool) {
			defer containerSync.Done()
			containerDetail, _ := client.InspectContainer(container.ID)
			handleContainer(containerDetail, client, grace, output)
		}(&container, client, ctx.Duration("grace"), ctx.BoolT("quiet"))
	}
	containerSync.Wait()
}

func handleContainer(container *docker.Container, client *docker.Client, grace time.Duration, output bool) {
}

func runDgc(ctx *cli.Context) {
	var dgcSync sync.WaitGroup
	client, _ := docker.NewClient(ctx.String("socket"))
	images, _ := client.ListImages(docker.ListImagesOptions{All: false})
	containers, _ := client.ListContainers(docker.ListContainersOptions{All: false})
	dgcSync.Add(2)
	go func(containers []docker.APIContainers, client *docker.Client, ctx *cli.Context) {
		defer dgcSync.Done()
		collectAPIContainers(containers, client, ctx)
	}(containers, client, ctx)
	go func(images []docker.APIImages, client *docker.Client, ctx *cli.Context) {
		defer dgcSync.Done()
		collectAPIImages(images, client, ctx)
	}(images, client, ctx)
	dgcSync.Wait()
}

func main() {
	dgc := cli.NewApp()
	dgc.EnableBashCompletion = true
	dgc.Name = "dgc"
	dgc.Usage = "A minimal docker garbage collector"
	dgc.Version = "0.1.0"
	dgc.Author = "David J Felix <davidjfelix@davidjfelix.com>"
	dgc.Action = runDgc
	dgc.Flags = []cli.Flag {
		cli.DurationFlag {
			Name: "grace, g",
			Value: time.Duration(3600)*time.Second,
			Usage: "the grace period for a container. Accepted compostable time units: [h, m, s, ms, ns us]",
			EnvVar: "GRACE_PERIOD_SECONDS,GRACE_PERIOD",
		},
		cli.StringFlag {
			Name: "socket, s",
			Value: "unix:///var/run/docker.sock",
			Usage: "the docker remote socket",
			EnvVar: "DOCKER_SOCKET",
		},
		cli.StringFlag {
			Name: "exclude, e",
			Value: "/etc/docker-gc-exclude",
			Usage: "the list of containers to exclude from garbage collection, as a file or directory",
			EnvVar: "EXCLUDE_FROM_GC",
		},
		cli.BoolTFlag {
			Name: "quiet, q",
			Usage: "don't print name of garbage-collected containers",
			EnvVar: "QUIET",
		},
	}
	dgc.Run(os.Args)
}
<commit_msg>Switch quiet to off by default<commit_after>package main

import (
	"os"
	"time"
	"sync"
	"github.com/codegangsta/cli"
	"github.com/fsouza/go-dockerclient"
)

func collectAPIImages(images []docker.APIImages, client *docker.Client, ctx *cli.Context) {
	var imageSync sync.WaitGroup
	for _, image := range images {
		imageSync.Add(1)
		go func(image *docker.APIImages, client *docker.Client, grace time.Duration, output bool) {
			defer imageSync.Done()
			imageDetail, _ := client.InspectImage(image.ID)
			handleImage(imageDetail, client, grace, output)
		}(&image, client, ctx.Duration("grace"), ctx.Bool("quiet"))
	}
	imageSync.Wait()
}

func handleImage(image *docker.Image, client *docker.Client, grace time.Duration, output bool) {
}

func collectAPIContainers(containers []docker.APIContainers, client *docker.Client, ctx *cli.Context) {
	var containerSync sync.WaitGroup
	for _, container := range containers {
		containerSync.Add(1)
		go func(container *docker.APIContainers, client *docker.Client, grace time.Duration, output bool) {
			defer containerSync.Done()
			containerDetail, _ := client.InspectContainer(container.ID)
			handleContainer(containerDetail, client, grace, output)
		}(&container, client, ctx.Duration("grace"), ctx.Bool("quiet"))
	}
	containerSync.Wait()
}

func handleContainer(container *docker.Container, client *docker.Client, grace time.Duration, output bool) {
}

func runDgc(ctx *cli.Context) {
	var dgcSync sync.WaitGroup
	client, _ := docker.NewClient(ctx.String("socket"))
	images, _ := client.ListImages(docker.ListImagesOptions{All: false})
	containers, _ := client.ListContainers(docker.ListContainersOptions{All: false})
	dgcSync.Add(2)
	go func(containers []docker.APIContainers, client *docker.Client, ctx *cli.Context) {
		defer dgcSync.Done()
		collectAPIContainers(containers, client, ctx)
	}(containers, client, ctx)
	go func(images []docker.APIImages, client *docker.Client, ctx *cli.Context) {
		defer dgcSync.Done()
		collectAPIImages(images, client, ctx)
	}(images, client, ctx)
	dgcSync.Wait()
}

func main() {
	dgc := cli.NewApp()
	dgc.EnableBashCompletion = true
	dgc.Name = "dgc"
	dgc.Usage = "A minimal docker garbage collector"
	dgc.Version = "0.1.0"
	dgc.Author = "David J Felix <davidjfelix@davidjfelix.com>"
	dgc.Action = runDgc
	dgc.Flags = []cli.Flag {
		cli.DurationFlag {
			Name: "grace, g",
			Value: time.Duration(3600)*time.Second,
			Usage: "the grace period for a container. Accepted compostable time units: [h, m, s, ms, ns us]",
			EnvVar: "GRACE_PERIOD_SECONDS,GRACE_PERIOD",
		},
		cli.StringFlag {
			Name: "socket, s",
			Value: "unix:///var/run/docker.sock",
			Usage: "the docker remote socket",
			EnvVar: "DOCKER_SOCKET",
		},
		cli.StringFlag {
			Name: "exclude, e",
			Value: "/etc/docker-gc-exclude",
			Usage: "the list of containers to exclude from garbage collection, as a file or directory",
			EnvVar: "EXCLUDE_FROM_GC",
		},
		cli.BoolFlag {
			Name: "quiet, q",
			Usage: "don't print name of garbage-collected containers",
			EnvVar: "QUIET",
		},
	}
	dgc.Run(os.Args)
}
<|endoftext|>


<commit_before>package dht

import (
	"crypto"
	_ "crypto/sha1"
	"encoding/binary"
	"errors"
	"fmt"
	"github.com/nsf/libtorgo/bencode"
	"io"
	"log"
	"net"
	"os"
	"sync"
	"time"
)

type Server struct {
	ID               string
	Socket           *net.UDPConn
	transactions     []*transaction
	transactionIDInt uint64
	nodes            map[string]*Node
	mu               sync.Mutex
}

type Node struct {
	addr          *net.UDPAddr
	id            string
	lastHeardFrom time.Time
	lastSentTo    time.Time
}

func (n *Node) Good() bool {
	if len(n.id) != 20 {
		return false
	}
	if time.Now().Sub(n.lastHeardFrom) >= 15*time.Minute {
		return false
	}
	return true
}

type Msg map[string]interface{}

var _ fmt.Stringer = Msg{}

func (m Msg) String() string {
	return fmt.Sprintf("%#v", m)
}

type transaction struct {
	remoteAddr net.Addr
	t          string
	Response   chan Msg
	response   chan Msg
}

func (s *Server) setDefaults() {
	if s.ID == "" {
		var id [20]byte
		h := crypto.SHA1.New()
		ss, err := os.Hostname()
		if err != nil {
			log.Print(err)
		}
		ss += s.Socket.LocalAddr().String()
		h.Write([]byte(ss))
		if b := h.Sum(id[:0:20]); len(b) != 20 {
			panic(len(b))
		}
		if len(id) != 20 {
			panic(len(id))
		}
		s.ID = string(id[:])
	}
}

func (s *Server) Init() {
	s.setDefaults()
}

func (s *Server) Serve() error {
	for {
		var b [1500]byte
		n, addr, err := s.Socket.ReadFromUDP(b[:])
		if err != nil {
			return err
		}
		var d map[string]interface{}
		err = bencode.Unmarshal(b[:n], &d)
		if err != nil {
			log.Printf("bad krpc message: %s", err)
			continue
		}
		s.mu.Lock()
		if d["y"] == "q" {
			s.handleQuery(addr, d)
			s.mu.Unlock()
			continue
		}
		t := s.findResponseTransaction(d["t"].(string), addr)
		if t == nil {
			log.Printf("unexpected message: %#v", d)
			s.mu.Unlock()
			continue
		}
		t.response <- d
		s.removeTransaction(t)
		id := ""
		if d["y"] == "r" {
			id = d["r"].(map[string]interface{})["id"].(string)
		}
		s.heardFromNode(addr, id)
		s.mu.Unlock()
	}
}

func (s *Server) AddNode(ni NodeInfo) {
	if s.nodes == nil {
		s.nodes = make(map[string]*Node)
	}
	n := s.getNode(ni.Addr)
	if n.id == "" {
		n.id = string(ni.ID[:])
	}
}

func (s *Server) handleQuery(source *net.UDPAddr, m Msg) {
	log.Print(m["q"])
	if m["q"] != "ping" {
		return
	}
	s.heardFromNode(source, m["a"].(map[string]interface{})["id"].(string))
	s.reply(source, m["t"].(string))
}

func (s *Server) reply(addr *net.UDPAddr, t string) {
	m := map[string]interface{}{
		"t": t,
		"y": "r",
		"r": map[string]string{
			"id": s.IDString(),
		},
	}
	b, err := bencode.Marshal(m)
	if err != nil {
		panic(err)
	}
	_, err = s.Socket.WriteTo(b, addr)
	if err != nil {
		panic(err)
	}
}

func (s *Server) heardFromNode(addr *net.UDPAddr, id string) {
	n := s.getNode(addr)
	n.id = id
	n.lastHeardFrom = time.Now()
}

func (s *Server) getNode(addr *net.UDPAddr) (n *Node) {
	n = s.nodes[addr.String()]
	if n == nil {
		n = &Node{
			addr: addr,
		}
		s.nodes[addr.String()] = n
	}
	return
}

func (s *Server) sentToNode(addr *net.UDPAddr) {
	n := s.getNode(addr)
	n.lastSentTo = time.Now()
}

func (s *Server) findResponseTransaction(transactionID string, sourceNode net.Addr) *transaction {
	for _, t := range s.transactions {
		if t.t == transactionID && t.remoteAddr.String() == sourceNode.String() {
			return t
		}
	}
	return nil
}

func (s *Server) nextTransactionID() string {
	var b [binary.MaxVarintLen64]byte
	n := binary.PutUvarint(b[:], s.transactionIDInt)
	s.transactionIDInt++
	return string(b[:n])
}

func (s *Server) removeTransaction(t *transaction) {
	for i, tt := range s.transactions {
		if t == tt {
			last := len(s.transactions) - 1
			s.transactions[i] = s.transactions[last]
			s.transactions = s.transactions[:last]
			return
		}
	}
	panic("transaction not found")
}

func (s *Server) addTransaction(t *transaction) {
	s.transactions = append(s.transactions, t)
}

func (s *Server) IDString() string {
	if len(s.ID) != 20 {
		panic("bad node id")
	}
	return s.ID
}

func (s *Server) query(node *net.UDPAddr, q string, a map[string]string) (t *transaction, err error) {
	tid := s.nextTransactionID()
	if a == nil {
		a = make(map[string]string, 1)
	}
	a["id"] = s.IDString()
	d := map[string]interface{}{
		"t": tid,
		"y": "q",
		"q": q,
		"a": a,
	}
	b, err := bencode.Marshal(d)
	if err != nil {
		return
	}
	t = &transaction{
		remoteAddr: node,
		t:          tid,
		Response:   make(chan Msg, 1),
	}
	t.response = t.Response
	s.addTransaction(t)
	n, err := s.Socket.WriteTo(b, node)
	if err != nil {
		s.removeTransaction(t)
		return
	}
	if n != len(b) {
		err = io.ErrShortWrite
		s.removeTransaction(t)
		return
	}
	s.sentToNode(node)
	return
}

const CompactNodeInfoLen = 26

type NodeInfo struct {
	ID   [20]byte
	Addr *net.UDPAddr
}

func (ni *NodeInfo) PutCompact(b []byte) error {
	if n := copy(b[:], ni.ID[:]); n != 20 {
		panic(n)
	}
	ip := ni.Addr.IP.To4()
	if len(ip) != 4 {
		panic(ip)
	}
	if n := copy(b[20:], ip); n != 4 {
		panic(n)
	}
	binary.BigEndian.PutUint16(b[24:], uint16(ni.Addr.Port))
	return nil
}

func (cni *NodeInfo) UnmarshalCompact(b []byte) error {
	if len(b) != 26 {
		return errors.New("expected 26 bytes")
	}
	if 20 != copy(cni.ID[:], b[:20]) {
		panic("impossibru!")
	}
	if cni.Addr == nil {
		cni.Addr = &net.UDPAddr{}
	}
	cni.Addr.IP = net.IPv4(b[20], b[21], b[22], b[23])
	cni.Addr.Port = int(binary.BigEndian.Uint16(b[24:26]))
	return nil
}

func (s *Server) Ping(node *net.UDPAddr) (*transaction, error) {
	return s.query(node, "ping", nil)
}

type findNodeResponse struct {
	Nodes []NodeInfo
}

func getResponseNodes(m Msg) (s string, err error) {
	defer func() {
		r := recover()
		if r == nil {
			return
		}
		err = fmt.Errorf("couldn't get response nodes: %s: %#v", r, m)
	}()
	s = m["r"].(map[string]interface{})["nodes"].(string)
	return
}

func (me *findNodeResponse) UnmarshalKRPCMsg(m Msg) error {
	b, err := getResponseNodes(m)
	if err != nil {
		return err
	}
	for i := 0; i < len(b); i += 26 {
		var n NodeInfo
		err := n.UnmarshalCompact([]byte(b[i : i+26]))
		if err != nil {
			return err
		}
		me.Nodes = append(me.Nodes, n)
	}
	return nil
}

func (s *Server) FindNode(addr *net.UDPAddr, targetID string) (t *transaction, err error) {
	t, err = s.query(addr, "find_node", map[string]string{"target": targetID})
	if err != nil {
		return
	}
	ch := make(chan Msg)
	t.response = ch
	go func() {
		d, ok := <-t.response
		if !ok {
			close(t.Response)
			return
		}
		if d["y"] == "r" {
			var r findNodeResponse
			err = r.UnmarshalKRPCMsg(d)
			if err != nil {
				log.Print(err)
			} else {
				s.mu.Lock()
				for _, cni := range r.Nodes {
					n := s.getNode(cni.Addr)
					n.id = string(cni.ID[:])
				}
				s.mu.Unlock()
			}
		}
		t.Response <- d
	}()
	return
}

func (s *Server) addRootNode() error {
	addr, err := net.ResolveUDPAddr("udp4", "router.bittorrent.com:6881")
	if err != nil {
		return err
	}
	s.nodes[addr.String()] = &Node{
		addr: addr,
	}
	return nil
}

// Populates the node table.
func (s *Server) Bootstrap() (err error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if len(s.nodes) == 0 {
		err = s.addRootNode()
		if err != nil {
			return
		}
	}
	for _, node := range s.nodes {
		var t *transaction
		s.mu.Unlock()
		t, err = s.FindNode(node.addr, s.ID)
		s.mu.Lock()
		if err != nil {
			return
		}
		go func() {
			<-t.Response
		}()
	}
	return
}

func (s *Server) GoodNodes() (nis []NodeInfo) {
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, node := range s.nodes {
		if !node.Good() {
			continue
		}
		ni := NodeInfo{
			Addr: node.addr,
		}
		if n := copy(ni.ID[:], node.id); n != 20 {
			panic(n)
		}
		nis = append(nis, ni)
	}
	return
}

func (s *Server) StopServing() {
	s.Socket.Close()
}
<commit_msg>Wrap server UDP sends to ensure the node last send is updated<commit_after>package dht

import (
	"crypto"
	_ "crypto/sha1"
	"encoding/binary"
	"errors"
	"fmt"
	"github.com/nsf/libtorgo/bencode"
	"io"
	"log"
	"net"
	"os"
	"sync"
	"time"
)

type Server struct {
	ID               string
	Socket           *net.UDPConn
	transactions     []*transaction
	transactionIDInt uint64
	nodes            map[string]*Node
	mu               sync.Mutex
}

type Node struct {
	addr          *net.UDPAddr
	id            string
	lastHeardFrom time.Time
	lastSentTo    time.Time
}

func (n *Node) Good() bool {
	if len(n.id) != 20 {
		return false
	}
	if time.Now().Sub(n.lastHeardFrom) >= 15*time.Minute {
		return false
	}
	return true
}

type Msg map[string]interface{}

var _ fmt.Stringer = Msg{}

func (m Msg) String() string {
	return fmt.Sprintf("%#v", m)
}

type transaction struct {
	remoteAddr net.Addr
	t          string
	Response   chan Msg
	response   chan Msg
}

func (s *Server) setDefaults() {
	if s.ID == "" {
		var id [20]byte
		h := crypto.SHA1.New()
		ss, err := os.Hostname()
		if err != nil {
			log.Print(err)
		}
		ss += s.Socket.LocalAddr().String()
		h.Write([]byte(ss))
		if b := h.Sum(id[:0:20]); len(b) != 20 {
			panic(len(b))
		}
		if len(id) != 20 {
			panic(len(id))
		}
		s.ID = string(id[:])
	}
}

func (s *Server) Init() {
	s.setDefaults()
}

func (s *Server) Serve() error {
	for {
		var b [1500]byte
		n, addr, err := s.Socket.ReadFromUDP(b[:])
		if err != nil {
			return err
		}
		var d map[string]interface{}
		err = bencode.Unmarshal(b[:n], &d)
		if err != nil {
			log.Printf("bad krpc message: %s", err)
			continue
		}
		s.mu.Lock()
		if d["y"] == "q" {
			s.handleQuery(addr, d)
			s.mu.Unlock()
			continue
		}
		t := s.findResponseTransaction(d["t"].(string), addr)
		if t == nil {
			log.Printf("unexpected message: %#v", d)
			s.mu.Unlock()
			continue
		}
		t.response <- d
		s.removeTransaction(t)
		id := ""
		if d["y"] == "r" {
			id = d["r"].(map[string]interface{})["id"].(string)
		}
		s.heardFromNode(addr, id)
		s.mu.Unlock()
	}
}

func (s *Server) AddNode(ni NodeInfo) {
	if s.nodes == nil {
		s.nodes = make(map[string]*Node)
	}
	n := s.getNode(ni.Addr)
	if n.id == "" {
		n.id = string(ni.ID[:])
	}
}

func (s *Server) handleQuery(source *net.UDPAddr, m Msg) {
	log.Print(m["q"])
	if m["q"] != "ping" {
		return
	}
	s.heardFromNode(source, m["a"].(map[string]interface{})["id"].(string))
	s.reply(source, m["t"].(string))
}

func (s *Server) reply(addr *net.UDPAddr, t string) {
	m := map[string]interface{}{
		"t": t,
		"y": "r",
		"r": map[string]string{
			"id": s.IDString(),
		},
	}
	b, err := bencode.Marshal(m)
	if err != nil {
		panic(err)
	}
	err = s.writeToNode(b, addr)
	if err != nil {
		panic(err)
	}
}

func (s *Server) heardFromNode(addr *net.UDPAddr, id string) {
	n := s.getNode(addr)
	n.id = id
	n.lastHeardFrom = time.Now()
}

func (s *Server) getNode(addr *net.UDPAddr) (n *Node) {
	n = s.nodes[addr.String()]
	if n == nil {
		n = &Node{
			addr: addr,
		}
		s.nodes[addr.String()] = n
	}
	return
}

func (s *Server) writeToNode(b []byte, node *net.UDPAddr) (err error) {
	n, err := s.Socket.WriteTo(b, node)
	if err != nil {
		return
	}
	if n != len(b) {
		err = io.ErrShortWrite
		return
	}
	s.sentToNode(node)
	return
}

func (s *Server) sentToNode(addr *net.UDPAddr) {
	n := s.getNode(addr)
	n.lastSentTo = time.Now()
}

func (s *Server) findResponseTransaction(transactionID string, sourceNode net.Addr) *transaction {
	for _, t := range s.transactions {
		if t.t == transactionID && t.remoteAddr.String() == sourceNode.String() {
			return t
		}
	}
	return nil
}

func (s *Server) nextTransactionID() string {
	var b [binary.MaxVarintLen64]byte
	n := binary.PutUvarint(b[:], s.transactionIDInt)
	s.transactionIDInt++
	return string(b[:n])
}

func (s *Server) removeTransaction(t *transaction) {
	for i, tt := range s.transactions {
		if t == tt {
			last := len(s.transactions) - 1
			s.transactions[i] = s.transactions[last]
			s.transactions = s.transactions[:last]
			return
		}
	}
	panic("transaction not found")
}

func (s *Server) addTransaction(t *transaction) {
	s.transactions = append(s.transactions, t)
}

func (s *Server) IDString() string {
	if len(s.ID) != 20 {
		panic("bad node id")
	}
	return s.ID
}

func (s *Server) query(node *net.UDPAddr, q string, a map[string]string) (t *transaction, err error) {
	tid := s.nextTransactionID()
	if a == nil {
		a = make(map[string]string, 1)
	}
	a["id"] = s.IDString()
	d := map[string]interface{}{
		"t": tid,
		"y": "q",
		"q": q,
		"a": a,
	}
	b, err := bencode.Marshal(d)
	if err != nil {
		return
	}
	t = &transaction{
		remoteAddr: node,
		t:          tid,
		Response:   make(chan Msg, 1),
	}
	t.response = t.Response
	s.addTransaction(t)
	err = s.writeToNode(b, node)
	if err != nil {
		s.removeTransaction(t)
	}
	return
}

const CompactNodeInfoLen = 26

type NodeInfo struct {
	ID   [20]byte
	Addr *net.UDPAddr
}

func (ni *NodeInfo) PutCompact(b []byte) error {
	if n := copy(b[:], ni.ID[:]); n != 20 {
		panic(n)
	}
	ip := ni.Addr.IP.To4()
	if len(ip) != 4 {
		panic(ip)
	}
	if n := copy(b[20:], ip); n != 4 {
		panic(n)
	}
	binary.BigEndian.PutUint16(b[24:], uint16(ni.Addr.Port))
	return nil
}

func (cni *NodeInfo) UnmarshalCompact(b []byte) error {
	if len(b) != 26 {
		return errors.New("expected 26 bytes")
	}
	if 20 != copy(cni.ID[:], b[:20]) {
		panic("impossibru!")
	}
	if cni.Addr == nil {
		cni.Addr = &net.UDPAddr{}
	}
	cni.Addr.IP = net.IPv4(b[20], b[21], b[22], b[23])
	cni.Addr.Port = int(binary.BigEndian.Uint16(b[24:26]))
	return nil
}

func (s *Server) Ping(node *net.UDPAddr) (*transaction, error) {
	return s.query(node, "ping", nil)
}

type findNodeResponse struct {
	Nodes []NodeInfo
}

func getResponseNodes(m Msg) (s string, err error) {
	defer func() {
		r := recover()
		if r == nil {
			return
		}
		err = fmt.Errorf("couldn't get response nodes: %s: %#v", r, m)
	}()
	s = m["r"].(map[string]interface{})["nodes"].(string)
	return
}

func (me *findNodeResponse) UnmarshalKRPCMsg(m Msg) error {
	b, err := getResponseNodes(m)
	if err != nil {
		return err
	}
	for i := 0; i < len(b); i += 26 {
		var n NodeInfo
		err := n.UnmarshalCompact([]byte(b[i : i+26]))
		if err != nil {
			return err
		}
		me.Nodes = append(me.Nodes, n)
	}
	return nil
}

func (s *Server) FindNode(addr *net.UDPAddr, targetID string) (t *transaction, err error) {
	t, err = s.query(addr, "find_node", map[string]string{"target": targetID})
	if err != nil {
		return
	}
	ch := make(chan Msg)
	t.response = ch
	go func() {
		d, ok := <-t.response
		if !ok {
			close(t.Response)
			return
		}
		if d["y"] == "r" {
			var r findNodeResponse
			err = r.UnmarshalKRPCMsg(d)
			if err != nil {
				log.Print(err)
			} else {
				s.mu.Lock()
				for _, cni := range r.Nodes {
					n := s.getNode(cni.Addr)
					n.id = string(cni.ID[:])
				}
				s.mu.Unlock()
			}
		}
		t.Response <- d
	}()
	return
}

func (s *Server) addRootNode() error {
	addr, err := net.ResolveUDPAddr("udp4", "router.bittorrent.com:6881")
	if err != nil {
		return err
	}
	s.nodes[addr.String()] = &Node{
		addr: addr,
	}
	return nil
}

// Populates the node table.
func (s *Server) Bootstrap() (err error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if len(s.nodes) == 0 {
		err = s.addRootNode()
		if err != nil {
			return
		}
	}
	for _, node := range s.nodes {
		var t *transaction
		s.mu.Unlock()
		t, err = s.FindNode(node.addr, s.ID)
		s.mu.Lock()
		if err != nil {
			return
		}
		go func() {
			<-t.Response
		}()
	}
	return
}

func (s *Server) GoodNodes() (nis []NodeInfo) {
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, node := range s.nodes {
		if !node.Good() {
			continue
		}
		ni := NodeInfo{
			Addr: node.addr,
		}
		if n := copy(ni.ID[:], node.id); n != 20 {
			panic(n)
		}
		nis = append(nis, ni)
	}
	return
}

func (s *Server) StopServing() {
	s.Socket.Close()
}
<|endoftext|>