<commit_before>package main

import (
	"bufio"
	"crypto/md5"
	"fmt"
	"os"
	"regexp"
	"strings"
)

const version = "0.1.0+Go"

type ansiColorer func(color int) string

type options struct {
	foreground ansiColorer
	background ansiColorer
	pattern    string
}

type colorMode int

const (
	mode16Color colorMode = iota
	mode256color
)

func main() {
	interactLinewise(colorLineMaker(parseArgs(os.Args[1:])))
}

func parseArgs(args []string) options {
	colorForeground, colorBackground := true, false
	regexps := make([]string, 0)
	colorDepth := mode256color

	acceptingFlags := true
	for _, arg := range args {
		if !acceptingFlags {
			regexps = append(regexps, arg)
			continue
		}

		if arg == "-h" || arg == "--help" {
			printUsage()
			os.Exit(0)
		} else if arg == "-v" || arg == "--version" {
			fmt.Printf("Synesthesia version %v\n", version)
			os.Exit(0)
		} else if arg == "-f" || arg == "--foreground" {
			colorForeground = true
		} else if arg == "--no-foreground" {
			colorForeground = false
		} else if arg == "-b" || arg == "--background" {
			colorBackground = true
		} else if arg == "--no-background" {
			colorBackground = false
		} else if arg == "--16" {
			colorDepth = mode16Color
		} else if arg == "--256" {
			colorDepth = mode256color
		} else if arg == "--" {
			acceptingFlags = false
		} else if strings.HasPrefix(arg, "-") {
			fmt.Printf("Unrecognized flag: %s\n", arg)
			printUsage()
			os.Exit(1)
		} else {
			regexps = append(regexps, arg)
		}
	}

	foregroundColorer, backgroundColorer := colorNoop, colorNoop
	if colorDepth == mode16Color {
		if colorForeground {
			foregroundColorer = colorAnsi16Foreground
		}
		if colorBackground {
			backgroundColorer = colorAnsi16Background
		}
	} else if colorDepth == mode256color {
		if colorForeground {
			foregroundColorer = colorAnsi256Foreground
		}
		if colorBackground {
			backgroundColorer = colorAnsi256Background
		}
	}

	validateRegexps(regexps)
	for i, pattern := range regexps {
		regexps[i] = fmt.Sprintf("(?:%v)", pattern)
	}

	return options{foregroundColorer, backgroundColorer, strings.Join(regexps, "|")}
}

func printUsage() {
	fmt.Fprintf(os.Stderr, `Usage: %v ([flags]|[patterns])... [-- [patterns]...]
Color standard input based on the values matched by the provided regular
expressions.
Before --, any parameters starting with a dash will be interpreted as flags.
After it, all parameters are interpreted as patterns.
A pattern is any valid Go regular expression.
Multiples of flags controlling the same option are allowed; the last flag
specified will take precedence over the ones that come before it.
Flags:
-b --background     Turn on or off background colorization. (default off)
   --no-background
-f --foreground     Turn on or off foreground colorization. (default on)
   --no-foreground
-h --help           Print usage information.
`, os.Args[0])
}

func validateRegexps(regexps []string) {
	valid := true
	for _, pattern := range regexps {
		if _, err := regexp.Compile(pattern); err != nil {
			valid = false
			fmt.Fprintf(os.Stderr, "Invalid regex: %v\n", pattern)
		}
	}
	if !valid {
		os.Exit(1)
	}
}

func interactLinewise(transform func([]byte) []byte) {
	input := bufio.NewReader(os.Stdin)
	line, err := input.ReadBytes('\n')
	for err == nil {
		fmt.Print(string(transform(line)))
		line, err = input.ReadBytes('\n')
	}
	fmt.Print(string(transform(line)))
}

func colorLineMaker(options options) func([]byte) []byte {
	pattern := regexp.MustCompile(options.pattern)
	return func(line []byte) []byte {
		return colorPatternInString(pattern, line, options.foreground, options.background)
	}
}

func colorPatternInString(pattern *regexp.Regexp, buf []byte, foreground, background ansiColorer) []byte {
	return pattern.ReplaceAllFunc(buf, colorMatchMaker(foreground, background))
}

func colorMatchMaker(foreground, background ansiColorer) func([]byte) []byte {
	return func(match []byte) []byte {
		return color(match, foreground, background)
	}
}

func color(buf []byte, foreground, background ansiColorer) []byte {
	sum := md5.Sum(buf)
	color := int(sum[md5.Size-3])<<16 | int(sum[md5.Size-2])<<8 | int(sum[md5.Size-1])
	return []byte(fmt.Sprintf("%v%v%v\033[m", foreground(color), background(color), string(buf)))
}

func colorNoop(color int) string {
	return ""
}

func colorAnsi16Foreground(color int) string {
	color = color % 16
	index := color%8 + 30
	if color < 8 {
		return fmt.Sprintf("\033[%vm", index)
	}
	return fmt.Sprintf("\033[%v;1m", index)
}

func colorAnsi16Background(color int) string {
	color = color % 16
	index := color%8 + 40
	if color < 8 {
		return fmt.Sprintf("\033[%vm", index)
	}
	return fmt.Sprintf("\033[%v;1m", index)
}

func colorAnsi256Foreground(color int) string {
	r, g, b := color>>(16+3), color>>(8+3)&0x1F, color>>3&0x1F
	return fmt.Sprintf("\033[38;5;%vm", rgbTo216Color(r, g, b)+16)
}

func colorAnsi256Background(color int) string {
	r, g, b := color>>(16+3), color>>(8+3)&0x1F, color>>3&0x1F
	return fmt.Sprintf("\033[48;5;%vm", rgbTo216Color(r, g, b)+16)
}

func rgbFrom216Color(color int) (int, int, int) {
	return color / 36, color / 6 % 6, color % 6
}

func rgbTo216Color(r, g, b int) int {
	return r*36 + g*6 + b
}
<commit_msg>Invert background colors<commit_after>package main

import (
	"bufio"
	"crypto/md5"
	"fmt"
	"os"
	"regexp"
	"strings"
)

const version = "0.1.0+Go"

type ansiColorer func(color int) string

type options struct {
	foreground ansiColorer
	background ansiColorer
	pattern    string
}

type colorMode int

const (
	mode16Color colorMode = iota
	mode256color
)

func main() {
	interactLinewise(colorLineMaker(parseArgs(os.Args[1:])))
}

func parseArgs(args []string) options {
	colorForeground, colorBackground := true, false
	regexps := make([]string, 0)
	colorDepth := mode256color

	acceptingFlags := true
	for _, arg := range args {
		if !acceptingFlags {
			regexps = append(regexps, arg)
			continue
		}

		if arg == "-h" || arg == "--help" {
			printUsage()
			os.Exit(0)
		} else if arg == "-v" || arg == "--version" {
			fmt.Printf("Synesthesia version %v\n", version)
			os.Exit(0)
		} else if arg == "-f" || arg == "--foreground" {
			colorForeground = true
		} else if arg == "--no-foreground" {
			colorForeground = false
		} else if arg == "-b" || arg == "--background" {
			colorBackground = true
		} else if arg == "--no-background" {
			colorBackground = false
		} else if arg == "--16" {
			colorDepth = mode16Color
		} else if arg == "--256" {
			colorDepth = mode256color
		} else if arg == "--" {
			acceptingFlags = false
		} else if strings.HasPrefix(arg, "-") {
			fmt.Printf("Unrecognized flag: %s\n", arg)
			printUsage()
			os.Exit(1)
		} else {
			regexps = append(regexps, arg)
		}
	}

	foregroundColorer, backgroundColorer := colorNoop, colorNoop
	if colorDepth == mode16Color {
		if colorForeground {
			foregroundColorer = colorAnsi16Foreground
		}
		if colorBackground {
			backgroundColorer = colorAnsi16Background
		}
	} else if colorDepth == mode256color {
		if colorForeground {
			foregroundColorer = colorAnsi256Foreground
		}
		if colorBackground {
			backgroundColorer = colorAnsi256Background
		}
	}

	validateRegexps(regexps)
	for i, pattern := range regexps {
		regexps[i] = fmt.Sprintf("(?:%v)", pattern)
	}

	return options{foregroundColorer, backgroundColorer, strings.Join(regexps, "|")}
}

func printUsage() {
	fmt.Fprintf(os.Stderr, `Usage: %v ([flags]|[patterns])... [-- [patterns]...]
Color standard input based on the values matched by the provided regular
expressions.
Before --, any parameters starting with a dash will be interpreted as flags.
After it, all parameters are interpreted as patterns.
A pattern is any valid Go regular expression.
Multiples of flags controlling the same option are allowed; the last flag
specified will take precedence over the ones that come before it.
Flags:
-b --background     Turn on or off background colorization. (default off)
   --no-background
-f --foreground     Turn on or off foreground colorization. (default on)
   --no-foreground
-h --help           Print usage information.
`, os.Args[0])
}

func validateRegexps(regexps []string) {
	valid := true
	for _, pattern := range regexps {
		if _, err := regexp.Compile(pattern); err != nil {
			valid = false
			fmt.Fprintf(os.Stderr, "Invalid regex: %v\n", pattern)
		}
	}
	if !valid {
		os.Exit(1)
	}
}

func interactLinewise(transform func([]byte) []byte) {
	input := bufio.NewReader(os.Stdin)
	line, err := input.ReadBytes('\n')
	for err == nil {
		fmt.Print(string(transform(line)))
		line, err = input.ReadBytes('\n')
	}
	fmt.Print(string(transform(line)))
}

func colorLineMaker(options options) func([]byte) []byte {
	pattern := regexp.MustCompile(options.pattern)
	return func(line []byte) []byte {
		return colorPatternInString(pattern, line, options.foreground, options.background)
	}
}

func colorPatternInString(pattern *regexp.Regexp, buf []byte, foreground, background ansiColorer) []byte {
	return pattern.ReplaceAllFunc(buf, colorMatchMaker(foreground, background))
}

func colorMatchMaker(foreground, background ansiColorer) func([]byte) []byte {
	return func(match []byte) []byte {
		return color(match, foreground, background)
	}
}

func color(buf []byte, foreground, background ansiColorer) []byte {
	sum := md5.Sum(buf)
	color := int(sum[md5.Size-3])<<16 | int(sum[md5.Size-2])<<8 | int(sum[md5.Size-1])
	return []byte(fmt.Sprintf("%v%v%v\033[m", foreground(color), background(0xFFFFFF-color), string(buf)))
}

func colorNoop(color int) string {
	return ""
}

func colorAnsi16Foreground(color int) string {
	color = color % 16
	index := color%8 + 30
	if color < 8 {
		return fmt.Sprintf("\033[%vm", index)
	}
	return fmt.Sprintf("\033[%v;1m", index)
}

func colorAnsi16Background(color int) string {
	color = color % 16
	index := color%8 + 40
	if color < 8 {
		return fmt.Sprintf("\033[%vm", index)
	}
	return fmt.Sprintf("\033[%v;1m", index)
}

func colorAnsi256Foreground(color int) string {
	r, g, b := color>>(16+3), color>>(8+3)&0x1F, color>>3&0x1F
	return fmt.Sprintf("\033[38;5;%vm", rgbTo216Color(r, g, b)+16)
}

func colorAnsi256Background(color int) string {
	r, g, b := color>>(16+3), color>>(8+3)&0x1F, color>>3&0x1F
	return fmt.Sprintf("\033[48;5;%vm", rgbTo216Color(r, g, b)+16)
}

func rgbFrom216Color(color int) (int, int, int) {
	return color / 36, color / 6 % 6, color % 6
}

func rgbTo216Color(r, g, b int) int {
	return r*36 + g*6 + b
}
<|endoftext|>


<commit_before>// +build windows

package overseer

import (
	"bytes"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
)

var (
	supported = true
	uid       = syscall.Getuid()
	gid       = syscall.Getgid()
	SIGUSR1   = syscall.SIGTERM
	SIGUSR2   = syscall.SIGTERM
	SIGTERM   = syscall.SIGTERM
)

func move(dst, src string) error {
	os.MkdirAll(filepath.Dir(dst), 0755)
	if err := os.Rename(src, dst); err == nil {
		return nil
	}
	//HACK: we're shelling out to move because windows
	//throws errors when crossing device boundaryes.
	// https://www.microsoft.com/resources/documentation/windows/xp/all/proddocs/en-us/move.mspx?mfr=true
	cmd := exec.Command("cmd", "/c", `move /y "`+src+`" "`+dst+`"`)
	if b, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("%v: %q: %v", cmd.Args, bytes.TrimSpace(b), err)
	}
	return nil
}

func chmod(f *os.File, perms os.FileMode) error {
	if err := f.Chmod(perms); err != nil && !strings.Contains(err.Error(), "not supported") {
		return err
	}
	return nil
}

func chown(f *os.File, uid, gid int) error {
	if err := f.Chown(uid, gid); err != nil && !strings.Contains(err.Error(), "not supported") {
		return err
	}
	return nil
}
<commit_msg>windows: escape cmd args properly<commit_after>// +build windows

package overseer

import (
	"bytes"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
)

var (
	supported = true
	uid       = syscall.Getuid()
	gid       = syscall.Getgid()
	SIGUSR1   = syscall.SIGTERM
	SIGUSR2   = syscall.SIGTERM
	SIGTERM   = syscall.SIGTERM
)

func move(dst, src string) error {
	os.MkdirAll(filepath.Dir(dst), 0755)
	if err := os.Rename(src, dst); err == nil {
		return nil
	}
	//HACK: we're shelling out to move because windows
	//throws errors when crossing device boundaryes.
	// https://www.microsoft.com/resources/documentation/windows/xp/all/proddocs/en-us/move.mspx?mfr=true

	// https://blogs.msdn.microsoft.com/twistylittlepassagesallalike/2011/04/23/everyone-quotes-command-line-arguments-the-wrong-way/
	R := func(s string) string { return replShellMeta.Replace(syscall.EscapeArg(s)) }
	cmd := exec.Command("cmd", "/c", `move /y `+R(src)+` `+R(dst))
	if b, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("%v: %q: %v", cmd.Args, bytes.TrimSpace(b), err)
	}
	return nil
}

func chmod(f *os.File, perms os.FileMode) error {
	if err := f.Chmod(perms); err != nil && !strings.Contains(err.Error(), "not supported") {
		return err
	}
	return nil
}

func chown(f *os.File, uid, gid int) error {
	if err := f.Chown(uid, gid); err != nil && !strings.Contains(err.Error(), "not supported") {
		return err
	}
	return nil
}

// https://blogs.msdn.microsoft.com/twistylittlepassagesallalike/2011/04/23/everyone-quotes-command-line-arguments-the-wrong-way/
var replShellMeta = strings.NewReplacer(
	`(`, `^(`,
	`)`, `^)`,
	`%`, `^%`,
	`!`, `^!`,
	`^`, `^^`,
	`"`, `^"`,
	`<`, `^<`,
	`>`, `^>`,
	`&`, `^&`,
	`|`, `^|`,
)
<|endoftext|>


<commit_before>//
// Copyright (C) 2014 Sebastian 'tokkee' Harl <sh@tokkee.org>
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
// 1. Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
// TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
// OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
// OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
// ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package sysdb

import (
	"fmt"
	"time"
)

// The SysDB JSON time format.
const jsonTime = `"2006-01-02 15:04:05 -0700"`

// A Duration represents the elapsed time between two instants as a
// nanoseconds count.
//
// It supports marshaling to and unmarshaling from the SysDB JSON format (a
// sequence of decimal numbers with a unit suffix).
type Duration time.Duration

// Common durations. All values greater than or equal to a day are not exact
// values but subject to daylight savings time changes, leap years, etc. They
// are available mostly for providing human readable display formats.
const (
	Second = Duration(1000000000)
	Minute = 60 * Second
	Hour   = 60 * Minute
	Day    = 24 * Hour
	Month  = Duration(30436875 * 24 * 60 * 60 * 1000)
	Year   = Duration(3652425 * 24 * 60 * 60 * 100000)
)

// MarshalJSON implements the json.Marshaler interface. The duration is a
// quoted string in the SysDB JSON format.
func (d Duration) MarshalJSON() ([]byte, error) {
	if d == 0 {
		return []byte(`"0s"`), nil
	}

	s := `"`
	secs := false
	for _, spec := range []struct {
		interval Duration
		suffix   string
	}{{Year, "Y"}, {Month, "M"}, {Day, "D"}, {Hour, "h"}, {Minute, "m"}, {Second, ""}} {
		if d >= spec.interval {
			s += fmt.Sprintf("%d%s", d/spec.interval, spec.suffix)
			d %= spec.interval
			if spec.interval == Second {
				secs = true
			}
		}
	}

	if d > 0 {
		s += fmt.Sprintf(".%09d", d)
		for i := len(s) - 1; i > 0; i-- {
			if s[i] != '0' {
				break
			}
			s = s[:i]
		}
		secs = true
	}
	if secs {
		s += "s"
	}
	s += `"`
	return []byte(s), nil
}

// UnmarshalJSON implements the json.Unmarshaler interface. The duration is
// expected to be a quoted string in the SysDB JSON format.
func (d *Duration) UnmarshalJSON(data []byte) error {
	m := map[string]Duration{
		"Y": Year,
		"M": Month,
		"D": Day,
		"h": Hour,
		"m": Minute,
		"s": Second,
	}

	if data[0] != '"' || data[len(data)-1] != '"' {
		return fmt.Errorf("unquoted duration %q", string(data))
	}
	data = data[1 : len(data)-1]

	orig := string(data)
	var res Duration
	for len(data) != 0 {
		// consume digits
		n := 0
		dec := 0
		frac := false
		for n < len(data) && '0' <= data[n] && data[n] <= '9' {
			dec = dec*10 + int(data[n]-'0')
			n++
		}
		if n < len(data) && data[n] == '.' {
			frac = true
			n++

			// consume fraction
			m := 1000000000
			for n < len(data) && '0' <= data[n] && data[n] <= '9' {
				if m > 1 { // cut of to nanoseconds
					dec = dec*10 + int(data[n]-'0')
					m /= 10
				}
				n++
			}
			dec *= m
		}
		if n >= len(data) {
			return fmt.Errorf("missing unit in duration %q", orig)
		}
		if n == 0 {
			// we found something which is not a number
			return fmt.Errorf("invalid duration %q", orig)
		}

		// consume unit
		u := n
		for u < len(data) && data[u] != '.' && (data[u] < '0' || '9' < data[u]) {
			u++
		}

		unit := string(data[n:u])
		data = data[u:]

		// convert to Duration
		d, ok := m[unit]
		if !ok {
			return fmt.Errorf("invalid unit %q in duration %q", unit, orig)
		}

		if d == Second {
			if frac {
				d = 1
			}
		} else if frac {
			return fmt.Errorf("invalid fraction %q%s in duration %q", dec, unit, orig)
		}

		res += Duration(dec) * d
	}
	*d = res
	return nil
}

// String returns the duration formatted using a predefined format string.
func (d Duration) String() string { return time.Duration(d).String() }

// A Time represents an instant in time with nanosecond precision.
//
// It supports marshaling to and unmarshaling from the SysDB JSON format
// (YYYY-MM-DD hh:mm:ss +-zzzz).
type Time time.Time

// MarshalJSON implements the json.Marshaler interface. The time is a quoted
// string in the SysDB JSON format.
func (t Time) MarshalJSON() ([]byte, error) {
	return []byte(time.Time(t).Format(jsonTime)), nil
}

// UnmarshalJSON implements the json.Unmarshaler interface. The time is
// expected to be a quoted string in the SysDB JSON format.
func (t *Time) UnmarshalJSON(data []byte) error {
	parsed, err := time.Parse(jsonTime, string(data))
	if err == nil {
		*t = Time(parsed)
	}
	return err
}

// Equal reports whether t and u represent the same time instant.
func (t Time) Equal(u Time) bool {
	return time.Time(t).Equal(time.Time(u))
}

// String returns the time formatted using a predefined format string.
func (t Time) String() string { return time.Time(t).String() }

// An Attribute describes a host, metric, or service attribute.
type Attribute struct {
	Name           string   `json:"name"`
	Value          string   `json:"value"`
	LastUpdate     Time     `json:"last_update"`
	UpdateInterval Duration `json:"update_interval"`
	Backends       []string `json:"backends"`
}

// A Metric describes a metric known to SysDB.
type Metric struct {
	Name           string      `json:"name"`
	LastUpdate     Time        `json:"last_update"`
	UpdateInterval Duration    `json:"update_interval"`
	Backends       []string    `json:"backends"`
	Attributes     []Attribute `json:"attributes"`
}

// A Service describes a service object stored in the SysDB store.
type Service struct {
	Name           string      `json:"name"`
	LastUpdate     Time        `json:"last_update"`
	UpdateInterval Duration    `json:"update_interval"`
	Backends       []string    `json:"backends"`
	Attributes     []Attribute `json:"attributes"`
}

// A Host describes a host object stored in the SysDB store.
type Host struct {
	Name           string      `json:"name"`
	LastUpdate     Time        `json:"last_update"`
	UpdateInterval Duration    `json:"update_interval"`
	Backends       []string    `json:"backends"`
	Attributes     []Attribute `json:"attributes"`
	Metrics        []Metric    `json:"metrics"`
	Services       []Service   `json:"services"`
}

// A DataPoint describes a datum at a certain point of time.
type DataPoint struct {
	Timestamp Time    `json:"timestamp"`
	Value     float64 `json:"value,string"`
}

// A Timeseries describes a sequence of data-points.
type Timeseries struct {
	Start Time                   `json:"start"`
	End   Time                   `json:"end"`
	Data  map[string][]DataPoint `json:"data"`
}

// vim: set tw=78 sw=4 sw=4 noexpandtab :
<commit_msg>Parse a metric's "timeseries" field.<commit_after>//
// Copyright (C) 2014 Sebastian 'tokkee' Harl <sh@tokkee.org>
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
// 1. Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
// TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
// OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
// OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
// ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package sysdb

import (
	"fmt"
	"time"
)

// The SysDB JSON time format.
const jsonTime = `"2006-01-02 15:04:05 -0700"`

// A Duration represents the elapsed time between two instants as a
// nanoseconds count.
//
// It supports marshaling to and unmarshaling from the SysDB JSON format (a
// sequence of decimal numbers with a unit suffix).
type Duration time.Duration

// Common durations. All values greater than or equal to a day are not exact
// values but subject to daylight savings time changes, leap years, etc. They
// are available mostly for providing human readable display formats.
const (
	Second = Duration(1000000000)
	Minute = 60 * Second
	Hour   = 60 * Minute
	Day    = 24 * Hour
	Month  = Duration(30436875 * 24 * 60 * 60 * 1000)
	Year   = Duration(3652425 * 24 * 60 * 60 * 100000)
)

// MarshalJSON implements the json.Marshaler interface. The duration is a
// quoted string in the SysDB JSON format.
func (d Duration) MarshalJSON() ([]byte, error) {
	if d == 0 {
		return []byte(`"0s"`), nil
	}

	s := `"`
	secs := false
	for _, spec := range []struct {
		interval Duration
		suffix   string
	}{{Year, "Y"}, {Month, "M"}, {Day, "D"}, {Hour, "h"}, {Minute, "m"}, {Second, ""}} {
		if d >= spec.interval {
			s += fmt.Sprintf("%d%s", d/spec.interval, spec.suffix)
			d %= spec.interval
			if spec.interval == Second {
				secs = true
			}
		}
	}

	if d > 0 {
		s += fmt.Sprintf(".%09d", d)
		for i := len(s) - 1; i > 0; i-- {
			if s[i] != '0' {
				break
			}
			s = s[:i]
		}
		secs = true
	}
	if secs {
		s += "s"
	}
	s += `"`
	return []byte(s), nil
}

// UnmarshalJSON implements the json.Unmarshaler interface. The duration is
// expected to be a quoted string in the SysDB JSON format.
func (d *Duration) UnmarshalJSON(data []byte) error {
	m := map[string]Duration{
		"Y": Year,
		"M": Month,
		"D": Day,
		"h": Hour,
		"m": Minute,
		"s": Second,
	}

	if data[0] != '"' || data[len(data)-1] != '"' {
		return fmt.Errorf("unquoted duration %q", string(data))
	}
	data = data[1 : len(data)-1]

	orig := string(data)
	var res Duration
	for len(data) != 0 {
		// consume digits
		n := 0
		dec := 0
		frac := false
		for n < len(data) && '0' <= data[n] && data[n] <= '9' {
			dec = dec*10 + int(data[n]-'0')
			n++
		}
		if n < len(data) && data[n] == '.' {
			frac = true
			n++

			// consume fraction
			m := 1000000000
			for n < len(data) && '0' <= data[n] && data[n] <= '9' {
				if m > 1 { // cut of to nanoseconds
					dec = dec*10 + int(data[n]-'0')
					m /= 10
				}
				n++
			}
			dec *= m
		}
		if n >= len(data) {
			return fmt.Errorf("missing unit in duration %q", orig)
		}
		if n == 0 {
			// we found something which is not a number
			return fmt.Errorf("invalid duration %q", orig)
		}

		// consume unit
		u := n
		for u < len(data) && data[u] != '.' && (data[u] < '0' || '9' < data[u]) {
			u++
		}

		unit := string(data[n:u])
		data = data[u:]

		// convert to Duration
		d, ok := m[unit]
		if !ok {
			return fmt.Errorf("invalid unit %q in duration %q", unit, orig)
		}

		if d == Second {
			if frac {
				d = 1
			}
		} else if frac {
			return fmt.Errorf("invalid fraction %q%s in duration %q", dec, unit, orig)
		}

		res += Duration(dec) * d
	}
	*d = res
	return nil
}

// String returns the duration formatted using a predefined format string.
func (d Duration) String() string { return time.Duration(d).String() }

// A Time represents an instant in time with nanosecond precision.
//
// It supports marshaling to and unmarshaling from the SysDB JSON format
// (YYYY-MM-DD hh:mm:ss +-zzzz).
type Time time.Time

// MarshalJSON implements the json.Marshaler interface. The time is a quoted
// string in the SysDB JSON format.
func (t Time) MarshalJSON() ([]byte, error) {
	return []byte(time.Time(t).Format(jsonTime)), nil
}

// UnmarshalJSON implements the json.Unmarshaler interface. The time is
// expected to be a quoted string in the SysDB JSON format.
func (t *Time) UnmarshalJSON(data []byte) error {
	parsed, err := time.Parse(jsonTime, string(data))
	if err == nil {
		*t = Time(parsed)
	}
	return err
}

// Equal reports whether t and u represent the same time instant.
func (t Time) Equal(u Time) bool {
	return time.Time(t).Equal(time.Time(u))
}

// String returns the time formatted using a predefined format string.
func (t Time) String() string { return time.Time(t).String() }

// An Attribute describes a host, metric, or service attribute.
type Attribute struct {
	Name           string   `json:"name"`
	Value          string   `json:"value"`
	LastUpdate     Time     `json:"last_update"`
	UpdateInterval Duration `json:"update_interval"`
	Backends       []string `json:"backends"`
}

// A Metric describes a metric known to SysDB.
type Metric struct {
	Name           string      `json:"name"`
	Timeseries     bool        `json:"timeseries"`
	LastUpdate     Time        `json:"last_update"`
	UpdateInterval Duration    `json:"update_interval"`
	Backends       []string    `json:"backends"`
	Attributes     []Attribute `json:"attributes"`
}

// A Service describes a service object stored in the SysDB store.
type Service struct {
	Name           string      `json:"name"`
	LastUpdate     Time        `json:"last_update"`
	UpdateInterval Duration    `json:"update_interval"`
	Backends       []string    `json:"backends"`
	Attributes     []Attribute `json:"attributes"`
}

// A Host describes a host object stored in the SysDB store.
type Host struct {
	Name           string      `json:"name"`
	LastUpdate     Time        `json:"last_update"`
	UpdateInterval Duration    `json:"update_interval"`
	Backends       []string    `json:"backends"`
	Attributes     []Attribute `json:"attributes"`
	Metrics        []Metric    `json:"metrics"`
	Services       []Service   `json:"services"`
}

// A DataPoint describes a datum at a certain point of time.
type DataPoint struct {
	Timestamp Time    `json:"timestamp"`
	Value     float64 `json:"value,string"`
}

// A Timeseries describes a sequence of data-points.
type Timeseries struct {
	Start Time                   `json:"start"`
	End   Time                   `json:"end"`
	Data  map[string][]DataPoint `json:"data"`
}

// vim: set tw=78 sw=4 sw=4 noexpandtab :
<|endoftext|>


<commit_before>package system

import (
	"errors"
	"strconv"
	"strings"

	"github.com/ready-steady/format/tgff"
)

func loadTGFF(path string) (*Platform, *Application, error) {
	result, err := tgff.ParseFile(path)
	if err != nil {
		return nil, nil, err
	}

	platform, err := loadPlatform(result.Tables)
	if err != nil {
		return nil, nil, err
	}

	application, err := loadApplication(result.Graphs)
	if err != nil {
		return nil, nil, err
	}

	return platform, application, nil
}

func loadPlatform(tables []tgff.Table) (*Platform, error) {
	size := uint(len(tables))

	if size == 0 {
		return nil, errors.New("need at least one table")
	}

	platform := &Platform{
		Cores: make([]Core, size),
	}

	var err error

	for _, table := range tables {
		i := table.ID

		if i >= size {
			return nil, errors.New("encountered an unknown table indexing scheme")
		}

		platform.Cores[i], err = loadCore(table)

		if err != nil {
			return nil, err
		}
	}

	rows := len(platform.Cores[0].Time)

	for i := uint(1); i < size; i++ {
		if rows != len(platform.Cores[i].Time) {
			return nil, errors.New("the table data are inconsistent")
		}
	}

	return platform, nil
}

func loadApplication(graphs []tgff.Graph) (*Application, error) {
	if len(graphs) != 1 {
		return nil, errors.New("need exactly one task graph")
	}

	size := uint(len(graphs[0].Tasks))

	application := &Application{
		Tasks: make([]Task, size),
	}

	for _, task := range graphs[0].Tasks {
		i, err := extractTaskID(task.Name, size)

		if err != nil {
			return nil, err
		}

		application.Tasks[i].ID = i
		application.Tasks[i].Type = task.Type
	}

	for _, arc := range graphs[0].Arcs {
		i, err := extractTaskID(arc.From, size)

		if err != nil {
			return nil, err
		}

		j, err := extractTaskID(arc.To, size)

		if err != nil {
			return nil, err
		}

		application.Tasks[i].Children = append(application.Tasks[i].Children, j)
		application.Tasks[j].Parents = append(application.Tasks[j].Parents, i)
	}

	return application, nil
}

func loadCore(table tgff.Table) (Core, error) {
	core := Core{
		ID: table.ID,
	}

	var tycol, tmcol, pwcol *tgff.Column

	for i := range table.Columns {
		col := &table.Columns[i]

		name := strings.ToLower(col.Name)

		if strings.Index(name, "type") >= 0 {
			tycol = col
		} else if strings.Index(name, "time") >= 0 {
			tmcol = col
		} else if strings.Index(name, "power") >= 0 {
			pwcol = col
		}

		if tycol != nil && tmcol != nil && pwcol != nil {
			break
		}
	}

	if tycol == nil || tmcol == nil || pwcol == nil {
		return core, errors.New("need columns named type, time, and power")
	}

	size := len(tycol.Data)

	for i := 0; i < size; i++ {
		if int(tycol.Data[i]) != i {
			return core, errors.New("data should be sorted by type")
		}
	}

	core.Time = make([]float64, size)
	copy(core.Time, tmcol.Data)

	core.Power = make([]float64, size)
	copy(core.Power, pwcol.Data)

	return core, nil
}

func extractTaskID(name string, total uint) (uint, error) {
	if !strings.HasPrefix(name, "t0_") {
		return 0, errors.New("encountered an unknown task naming scheme")
	}

	id, err := strconv.ParseUint(name[3:], 10, 0)
	if err != nil || id < 0 || uint(id) >= total {
		return 0, errors.New("encountered an unknown task indexing scheme")
	}

	return uint(id), nil
}
<commit_msg>system: fix the path to tgff<commit_after>package system

import (
	"errors"
	"strconv"
	"strings"

	"github.com/ready-steady/tgff"
)

func loadTGFF(path string) (*Platform, *Application, error) {
	result, err := tgff.ParseFile(path)
	if err != nil {
		return nil, nil, err
	}

	platform, err := loadPlatform(result.Tables)
	if err != nil {
		return nil, nil, err
	}

	application, err := loadApplication(result.Graphs)
	if err != nil {
		return nil, nil, err
	}

	return platform, application, nil
}

func loadPlatform(tables []tgff.Table) (*Platform, error) {
	size := uint(len(tables))

	if size == 0 {
		return nil, errors.New("need at least one table")
	}

	platform := &Platform{
		Cores: make([]Core, size),
	}

	var err error

	for _, table := range tables {
		i := table.ID

		if i >= size {
			return nil, errors.New("encountered an unknown table indexing scheme")
		}

		platform.Cores[i], err = loadCore(table)

		if err != nil {
			return nil, err
		}
	}

	rows := len(platform.Cores[0].Time)

	for i := uint(1); i < size; i++ {
		if rows != len(platform.Cores[i].Time) {
			return nil, errors.New("the table data are inconsistent")
		}
	}

	return platform, nil
}

func loadApplication(graphs []tgff.Graph) (*Application, error) {
	if len(graphs) != 1 {
		return nil, errors.New("need exactly one task graph")
	}

	size := uint(len(graphs[0].Tasks))

	application := &Application{
		Tasks: make([]Task, size),
	}

	for _, task := range graphs[0].Tasks {
		i, err := extractTaskID(task.Name, size)

		if err != nil {
			return nil, err
		}

		application.Tasks[i].ID = i
		application.Tasks[i].Type = task.Type
	}

	for _, arc := range graphs[0].Arcs {
		i, err := extractTaskID(arc.From, size)

		if err != nil {
			return nil, err
		}

		j, err := extractTaskID(arc.To, size)

		if err != nil {
			return nil, err
		}

		application.Tasks[i].Children = append(application.Tasks[i].Children, j)
		application.Tasks[j].Parents = append(application.Tasks[j].Parents, i)
	}

	return application, nil
}

func loadCore(table tgff.Table) (Core, error) {
	core := Core{
		ID: table.ID,
	}

	var tycol, tmcol, pwcol *tgff.Column

	for i := range table.Columns {
		col := &table.Columns[i]

		name := strings.ToLower(col.Name)

		if strings.Index(name, "type") >= 0 {
			tycol = col
		} else if strings.Index(name, "time") >= 0 {
			tmcol = col
		} else if strings.Index(name, "power") >= 0 {
			pwcol = col
		}

		if tycol != nil && tmcol != nil && pwcol != nil {
			break
		}
	}

	if tycol == nil || tmcol == nil || pwcol == nil {
		return core, errors.New("need columns named type, time, and power")
	}

	size := len(tycol.Data)

	for i := 0; i < size; i++ {
		if int(tycol.Data[i]) != i {
			return core, errors.New("data should be sorted by type")
		}
	}

	core.Time = make([]float64, size)
	copy(core.Time, tmcol.Data)

	core.Power = make([]float64, size)
	copy(core.Power, pwcol.Data)

	return core, nil
}

func extractTaskID(name string, total uint) (uint, error) {
	if !strings.HasPrefix(name, "t0_") {
		return 0, errors.New("encountered an unknown task naming scheme")
	}

	id, err := strconv.ParseUint(name[3:], 10, 0)
	if err != nil || id < 0 || uint(id) >= total {
		return 0, errors.New("encountered an unknown task indexing scheme")
	}

	return uint(id), nil
}
<|endoftext|>


<commit_before>package taipei

import (
	"io"
	"net"
	"sync"
	"time"
)

const MAX_OUR_REQUESTS = 2
const MAX_PEER_REQUESTS = 10
const STANDARD_BLOCK_LENGTH = 16 * 1024

type peerMessage struct {
	peer    *peerState
	message []byte // nil means an error occurred
}

type peerState struct {
	address         string
	id              string
	writeChan       chan []byte
	writeChan2      chan []byte
	lastWriteTime   time.Time
	lastReadTime    time.Time
	have            *Bitset // What the peer has told us it has
	conn            net.Conn
	am_choking      bool // this client is choking the peer
	am_interested   bool // this client is interested in the peer
	peer_choking    bool // peer is choking this client
	peer_interested bool // peer is interested in this client
	peer_requests   map[uint64]bool
	our_requests    map[uint64]time.Time // What we requested, when we requested it
	end             sync.Once
}

func queueingWriter(in, out chan []byte) {
	queue := make(map[int][]byte)
	head, tail := 0, 0
L:
	for {
		if head == tail {
			select {
			case m, ok := <-in:
				if !ok {
					break L
				}
				queue[head] = m
				head++
			}
		} else {
			select {
			case m, ok := <-in:
				if !ok {
					break L
				}
				queue[head] = m
				head++
			case out <- queue[tail]:
				delete(queue, tail)
				tail++
			}
		}
	}
	// We throw away any messages waiting to be sent, including the
	// nil message that is automatically sent when the in channel is closed
	close(out)
}

func NewPeerState(conn net.Conn) *peerState {
	writeChan := make(chan []byte)
	writeChan2 := make(chan []byte)
	go queueingWriter(writeChan, writeChan2)
	return &peerState{writeChan: writeChan, writeChan2: writeChan2, conn: conn,
		am_choking: true, peer_choking: true,
		peer_requests: make(map[uint64]bool, MAX_PEER_REQUESTS),
		our_requests:  make(map[uint64]time.Time, MAX_OUR_REQUESTS)}
}

func (p *peerState) Close() {
	p.end.Do(func() {
		p.conn.Close()
		close(p.writeChan)
	})
}

func (p *peerState) AddRequest(index, begin, length uint32) {
	if !p.am_choking && len(p.peer_requests) < MAX_PEER_REQUESTS {
		offset := (uint64(index) << 32) | uint64(begin)
		p.peer_requests[offset] = true
	}
}

func (p *peerState) CancelRequest(index, begin, length uint32) {
	offset := (uint64(index) << 32) | uint64(begin)
	if _, ok := p.peer_requests[offset]; ok {
		delete(p.peer_requests, offset)
	}
}

func (p *peerState) RemoveRequest() (index, begin, length uint32, ok bool) {
	for k, _ := range p.peer_requests {
		index, begin = uint32(k>>32), uint32(k)
		length = STANDARD_BLOCK_LENGTH
		ok = true
		return
	}
	return
}

func (p *peerState) SetChoke(choke bool) {
	if choke != p.am_choking {
		p.am_choking = choke
		b := byte(1)
		if choke {
			b = 0
			p.peer_requests = make(map[uint64]bool, MAX_PEER_REQUESTS)
		}
		p.sendOneCharMessage(b)
	}
}

func (p *peerState) SetInterested(interested bool) {
	if interested != p.am_interested {
		// log.Println("SetInterested", interested, p.address)
		p.am_interested = interested
		b := byte(3)
		if interested {
			b = 2
		}
		p.sendOneCharMessage(b)
	}
}

func (p *peerState) sendOneCharMessage(b byte) {
	// log.Println("ocm", b, p.address)
	p.sendMessage([]byte{b})
}

func (p *peerState) sendMessage(b []byte) {
	p.writeChan <- b
	p.lastWriteTime = time.Now()
}

func (p *peerState) keepAlive(now time.Time) {
	if now.Sub(p.lastWriteTime) >= 2*time.Minute {
		// log.Stderr("Sending keep alive", p)
		p.sendMessage([]byte{})
	}
}

// There's two goroutines per peer, one to read data from the peer, the other to
// send data to the peer.

func uint32ToBytes(buf []byte, n uint32) {
	buf[0] = byte(n >> 24)
	buf[1] = byte(n >> 16)
	buf[2] = byte(n >> 8)
	buf[3] = byte(n)
}

func writeNBOUint32(conn net.Conn, n uint32) (err error) {
	var buf []byte = make([]byte, 4)
	uint32ToBytes(buf, n)
	_, err = conn.Write(buf[0:])
	return
}

func bytesToUint32(buf []byte) uint32 {
	return (uint32(buf[0]) << 24) |
		(uint32(buf[1]) << 16) |
		(uint32(buf[2]) << 8) | uint32(buf[3])
}

func readNBOUint32(conn net.Conn) (n uint32, err error) {
	var buf [4]byte
	_, err = conn.Read(buf[0:])
	if err != nil {
		return
	}
	n = bytesToUint32(buf[0:])
	return
}

// This func is designed to be run as a goroutine. It
// listens for messages on a channel and sends them to a peer.

func (p *peerState) peerWriter(errorChan chan peerMessage, header []byte) {
	// log.Println("Writing header.")
	_, err := p.conn.Write(header)
	if err != nil {
		goto exit
	}
	// log.Println("Writing messages")
	for msg := range p.writeChan2 {
		// log.Println("Writing", len(msg), conn.RemoteAddr())
		err = writeNBOUint32(p.conn, uint32(len(msg)))
		if err != nil {
			goto exit
		}
		_, err = p.conn.Write(msg)
		if err != nil {
			// log.Println("Failed to write a message", p.address, len(msg), msg, err)
			goto exit
		}
	}
exit:
	// log.Println("peerWriter exiting")
	errorChan <- peerMessage{p, nil}
}

// This func is designed to be run as a goroutine. It
// listens for messages from the peer and forwards them to a channel.

func (p *peerState) peerReader(msgChan chan peerMessage) {
	// log.Println("Reading header.")
	var header [68]byte
	_, err := p.conn.Read(header[0:1])
	if err != nil {
		goto exit
	}
	if header[0] != 19 {
		goto exit
	}
	_, err = p.conn.Read(header[1:20])
	if err != nil {
		goto exit
	}
	if string(header[1:20]) != "BitTorrent protocol" {
		goto exit
	}
	// Read rest of header
	_, err = p.conn.Read(header[20:])
	if err != nil {
		goto exit
	}
	msgChan <- peerMessage{p, header[20:]}
	// log.Println("Reading messages")
	for {
		var n uint32
		n, err = readNBOUint32(p.conn)
		if err != nil {
			goto exit
		}
		if n > 130*1024 {
			// log.Println("Message size too large: ", n)
			goto exit
		}
		buf := make([]byte, n)
		_, err = io.ReadFull(p.conn, buf)
		if err != nil {
			goto exit
		}
		msgChan <- peerMessage{p, buf}
	}

exit:
	msgChan <- peerMessage{p, nil}
	// log.Println("peerReader exiting")
}
<commit_msg>dont close channels. Just throw away any messages that cames after a peer connection is closed<commit_after>package taipei

import (
	"io"
	"net"
	"time"
)

const MAX_OUR_REQUESTS = 2
const MAX_PEER_REQUESTS = 10
const STANDARD_BLOCK_LENGTH = 16 * 1024

type peerMessage struct {
	peer    *peerState
	message []byte // nil means an error occurred
}

type peerState struct {
	address         string
	id              string
	writeChan       chan []byte
	writeChan2      chan []byte
	lastWriteTime   time.Time
	lastReadTime    time.Time
	have            *Bitset // What the peer has told us it has
	conn            net.Conn
	am_choking      bool // this client is choking the peer
	am_interested   bool // this client is interested in the peer
	peer_choking    bool // peer is choking this client
	peer_interested bool // peer is interested in this client
	peer_requests   map[uint64]bool
	our_requests    map[uint64]time.Time // What we requested, when we requested it
}

func queueingWriter(in, out chan []byte) {
	queue := make(map[int][]byte)
	head, tail := 0, 0
L:
	for {
		if head == tail {
			select {
			case m, ok := <-in:
				if !ok {
					break L
				}
				queue[head] = m
				head++
			}
		} else {
			select {
			case m, ok := <-in:
				if !ok {
					break L
				}
				queue[head] = m
				head++
			case out <- queue[tail]:
				delete(queue, tail)
				tail++
			}
		}
	}
	// We throw away any messages waiting to be sent, including the
	// nil message that is automatically sent when the in channel is closed
	// close(out) // No need to close channel.
}

func NewPeerState(conn net.Conn) *peerState {
	writeChan := make(chan []byte)
	writeChan2 := make(chan []byte)
	go queueingWriter(writeChan, writeChan2)
	return &peerState{writeChan: writeChan, writeChan2: writeChan2, conn: conn,
		am_choking: true, peer_choking: true,
		peer_requests: make(map[uint64]bool, MAX_PEER_REQUESTS),
		our_requests:  make(map[uint64]time.Time, MAX_OUR_REQUESTS)}
}

func (p *peerState) Close() {
	p.conn.Close()
	// No need to close p.writeChan. Further writes to p.conn will just fail.
}

func (p *peerState) AddRequest(index, begin, length uint32) {
	if !p.am_choking && len(p.peer_requests) < MAX_PEER_REQUESTS {
		offset := (uint64(index) << 32) | uint64(begin)
		p.peer_requests[offset] = true
	}
}

func (p *peerState) CancelRequest(index, begin, length uint32) {
	offset := (uint64(index) << 32) | uint64(begin)
	if _, ok := p.peer_requests[offset]; ok {
		delete(p.peer_requests, offset)
	}
}

func (p *peerState) RemoveRequest() (index, begin, length uint32, ok bool) {
	for k, _ := range p.peer_requests {
		index, begin = uint32(k>>32), uint32(k)
		length = STANDARD_BLOCK_LENGTH
		ok = true
		return
	}
	return
}

func (p *peerState) SetChoke(choke bool) {
	if choke != p.am_choking {
		p.am_choking = choke
		b := byte(1)
		if choke {
			b = 0
			p.peer_requests = make(map[uint64]bool, MAX_PEER_REQUESTS)
		}
		p.sendOneCharMessage(b)
	}
}

func (p *peerState) SetInterested(interested bool) {
	if interested != p.am_interested {
		// log.Println("SetInterested", interested, p.address)
		p.am_interested = interested
		b := byte(3)
		if interested {
			b = 2
		}
		p.sendOneCharMessage(b)
	}
}

func (p *peerState) sendOneCharMessage(b byte) {
	// log.Println("ocm", b, p.address)
	p.sendMessage([]byte{b})
}

func (p *peerState) sendMessage(b []byte) {
	p.writeChan <- b
	p.lastWriteTime = time.Now()
}

func (p *peerState) keepAlive(now time.Time) {
	if now.Sub(p.lastWriteTime) >= 2*time.Minute {
		// log.Stderr("Sending keep alive", p)
		p.sendMessage([]byte{})
	}
}

// There's two goroutines per peer, one to read data from the peer, the other to
// send data to the peer.

func uint32ToBytes(buf []byte, n uint32) {
	buf[0] = byte(n >> 24)
	buf[1] = byte(n >> 16)
	buf[2] = byte(n >> 8)
	buf[3] = byte(n)
}

func writeNBOUint32(conn net.Conn, n uint32) (err error) {
	var buf []byte = make([]byte, 4)
	uint32ToBytes(buf, n)
	_, err = conn.Write(buf[0:])
	return
}

func bytesToUint32(buf []byte) uint32 {
	return (uint32(buf[0]) << 24) |
		(uint32(buf[1]) << 16) |
		(uint32(buf[2]) << 8) | uint32(buf[3])
}

func readNBOUint32(conn net.Conn) (n uint32, err error) {
	var buf [4]byte
	_, err = conn.Read(buf[0:])
	if err != nil {
		return
	}
	n = bytesToUint32(buf[0:])
	return
}

// This func is designed to be run as a goroutine. It
// listens for messages on a channel and sends them to a peer.

func (p *peerState) peerWriter(errorChan chan peerMessage, header []byte) {
	// log.Println("Writing header.")
	_, err := p.conn.Write(header)
	if err != nil {
		goto exit
	}
	// log.Println("Writing messages")
	for msg := range p.writeChan2 {
		// log.Println("Writing", len(msg), conn.RemoteAddr())
		err = writeNBOUint32(p.conn, uint32(len(msg)))
		if err != nil {
			goto exit
		}
		_, err = p.conn.Write(msg)
		if err != nil {
			// log.Println("Failed to write a message", p.address, len(msg), msg, err)
			goto exit
		}
	}
exit:
	// log.Println("peerWriter exiting")
	errorChan <- peerMessage{p, nil}
}

// This func is designed to be run as a goroutine. It
// listens for messages from the peer and forwards them to a channel.

func (p *peerState) peerReader(msgChan chan peerMessage) {
	// log.Println("Reading header.")
	var header [68]byte
	_, err := p.conn.Read(header[0:1])
	if err != nil {
		goto exit
	}
	if header[0] != 19 {
		goto exit
	}
	_, err = p.conn.Read(header[1:20])
	if err != nil {
		goto exit
	}
	if string(header[1:20]) != "BitTorrent protocol" {
		goto exit
	}
	// Read rest of header
	_, err = p.conn.Read(header[20:])
	if err != nil {
		goto exit
	}
	msgChan <- peerMessage{p, header[20:]}
	// log.Println("Reading messages")
	for {
		var n uint32
		n, err = readNBOUint32(p.conn)
		if err != nil {
			goto exit
		}
		if n > 130*1024 {
			// log.Println("Message size too large: ", n)
			goto exit
		}
		buf := make([]byte, n)
		_, err = io.ReadFull(p.conn, buf)
		if err != nil {
			goto exit
		}
		msgChan <- peerMessage{p, buf}
	}

exit:
	msgChan <- peerMessage{p, nil}
	// log.Println("peerReader exiting")
}
<|endoftext|>


<commit_before>// Package rehttp implements an HTTP transport that handles retries.
//
// An HTTP Client can be created with a *rehttp.Transport as RoundTripper
// and it will apply the retry strategy to its requests, e.g.:
//
//     tr, err := rehttp.NewTransport(
//         nil,                            // will use http.DefaultTransport
//         rehttp.RetryTemporaryErr(3),    // max 3 retries for Temporary errors
//         rehttp.ConstDelay(time.Second), // wait 1s between retries
//     )
//     if err != nil {
//         // handle err
//     }
//     client := &http.Client{
//         Transport: tr,
//         Timeout: 30 * time.Second, // timeout applies to all retries as a whole
//     }
//
// The retry strategy is provided by the Transport, which holds
// a function that returns whether or not the request should be retried,
// and if so, what delay to apply before retrying, based on the ShouldRetryFn
// and DelayFn functions passed to NewTransport.
//
// The package offers common delay strategies as ready-made functions that
// return a DelayFn:
//     - ConstDelay(delay time.Duration) DelayFn
//     - ExponentialDelay(base, max time.Duration) DelayFn
//
// It also provides common retry predicates that return a ShouldRetryFn:
//     - RetryTemporaryErr(maxRetries int) ShouldRetryFn
//     - RetryStatus500(maxRetries int) ShouldRetryFn
//     - RetryHTTPMethods(maxRetries int, methods ...string) ShouldRetryFn
//
// Those can be combined with RetryAny or RetryAll as needed. RetryAny
// enables retries if any of the ShouldRetryFn return true, while RetryAll
// enables retries if all ShouldRetryFn return true.
//
// The Transport will buffer the request's body in order to be able to
// retry the request, as a request attempt will consume and close the
// existing body. Sometimes this is not desirable, so it can be prevented
// by setting PreventRetryWithBody to true on the Transport. Doing so
// will disable retries when a request has a non-nil body.
//
package rehttp

import (
	"bytes"
	"errors"
	"io"
	"io/ioutil"
	"math"
	"math/rand"
	"net/http"
	"strings"
	"sync"
	"time"
)

// PRNG is the *math.Rand value to use to add jitter to the backoff
// algorithm used in ExponentialDelay. By default it uses a *math.Rand
// initialized with a source based on the current time in nanoseconds.
var PRNG = rand.New(rand.NewSource(time.Now().UnixNano()))

// terribly named interface to detect errors that support Temporary.
type temporaryer interface {
	Temporary() bool
}

// CancelRoundTripper is a RoundTripper that supports CancelRequest.
// The *http.Transport type implements this interface.
type CancelRoundTripper interface {
	http.RoundTripper
	CancelRequest(*http.Request)
}

// Attempt holds the data for a RoundTrip attempt.
type Attempt struct {
	// Index is the attempt index starting at 0.
	Index int

	// Request is the request for this attempt. If a non-nil Response
	// is present, this is the same as Response.Request, but since a
	// Response may not be available, it is guaranteed to be set on this
	// field.
	Request *http.Request

	// Response is the response for this attempt. It may be nil if an
	// error occurred an no response was received.
	Response *http.Response

	// Error is the error returned by the attempt, if any.
	Error error
}

// retryFn is the signature for functions that implement retry strategies.
type retryFn func(attempt Attempt) (bool, time.Duration)

// DelayFn is the signature for functions that return the delay to apply
// before the next retry.
type DelayFn func(attempt Attempt) time.Duration

// ShouldRetryFn is the signature for functions that return whether a
// retry should be done for the request.
type ShouldRetryFn func(attempt Attempt) bool

// NewTransport creates a Transport with a retry strategy based on
// shouldRetry and delay to control the retry logic. It uses the provided
// CancelRoundTripper to execute the requests. If rt is nil,
// http.DefaultTransport is used. An error is returned if http.DefaultTransport
// is not a CancelRoundTripper (which it is by default).
func NewTransport(rt CancelRoundTripper, shouldRetry ShouldRetryFn, delay DelayFn) (*Transport, error) {
	if rt == nil {
		var ok bool
		rt, ok = http.DefaultTransport.(CancelRoundTripper)
		if !ok {
			return nil, errors.New("http.DefaultTransport is not a CancelRoundTripper")
		}
	}
	return &Transport{
		CancelRoundTripper: rt,
		retry:              toRetryFn(shouldRetry, delay),
	}, nil
}

// toRetryFn combines shouldRetry and delay into a retryFn.
func toRetryFn(shouldRetry ShouldRetryFn, delay DelayFn) retryFn {
	return func(attempt Attempt) (bool, time.Duration) {
		retry := shouldRetry(attempt)
		if !retry {
			return false, 0
		}
		return true, delay(attempt)
	}
}

// RetryAny returns a ShouldRetryFn that allows a retry as long as one of
// the retryFns returns true. If retryFns is empty, it always returns false.
func RetryAny(retryFns ...ShouldRetryFn) ShouldRetryFn {
	return func(attempt Attempt) bool {
		for _, fn := range retryFns {
			if fn(attempt) {
				return true
			}
		}
		return false
	}
}

// RetryAll returns a ShouldRetryFn that allows a retry if all retryFns
// return true. If retryFns is empty, it always returns true.
func RetryAll(retryFns ...ShouldRetryFn) ShouldRetryFn {
	return func(attempt Attempt) bool {
		for _, fn := range retryFns {
			if !fn(attempt) {
				return false
			}
		}
		return true
	}
}

// RetryTemporaryErr returns a ShouldRetryFn that retries up to maxRetries
// times for a temporary error. A temporary error is one that implements
// the Temporary() bool method. Most errors from the net package implement
// this.
func RetryTemporaryErr(maxRetries int) ShouldRetryFn {
	return func(attempt Attempt) bool {
		if attempt.Index >= maxRetries {
			return false
		}
		if terr, ok := attempt.Error.(temporaryer); ok {
			return terr.Temporary()
		}
		return false
	}
}

// RetryStatus500 returns a ShouldRetryFn that retries up to maxRetries times
// for a status code 5xx.
func RetryStatus500(maxRetries int) ShouldRetryFn {
	return func(attempt Attempt) bool {
		if attempt.Index >= maxRetries {
			return false
		}
		return attempt.Response != nil &&
			attempt.Response.StatusCode >= 500 &&
			attempt.Response.StatusCode < 600 // who knows
	}
}

// RetryHTTPMethods returns a ShouldRetryFn that retries up to maxRetries
// times if the request's HTTP method is one of the provided methods.
// It is meant to be used in conjunction with another ShouldRetryFn such
// as RetryTemporaryErr combined using RetryAll, otherwise this function
// will retry any successful request made with one of the provided methods.
func RetryHTTPMethods(maxRetries int, methods ...string) ShouldRetryFn {
	for i, m := range methods {
		methods[i] = strings.ToUpper(m)
	}

	return func(attempt Attempt) bool {
		if attempt.Index >= maxRetries {
			return false
		}
		curMeth := strings.ToUpper(attempt.Request.Method)
		for _, m := range methods {
			if curMeth == m {
				return true
			}
		}
		return false
	}
}

// ConstDelay returns a DelayFn that always returns the same delay.
func ConstDelay(delay time.Duration) DelayFn {
	return func(attempt Attempt) time.Duration {
		return delay
	}
}

// ExponentialDelay returns a DelayFn that returns a delay between 0 and
// base * 2^attempt, capped at max.
//
// See the full jitter algorithm in:
// http://www.awsarchitectureblog.com/2015/03/backoff.html
func ExponentialDelay(base, max time.Duration) DelayFn {
	return func(attempt Attempt) time.Duration {
		exp := math.Pow(2, float64(attempt.Index))
		top := float64(base) * exp
		return time.Duration(
			PRNG.Int63n(int64(math.Min(float64(max), top))),
		)
	}
}

// Transport wraps a CancelRoundTripper such as *http.Transport and adds
// retry logic.
type Transport struct {
	CancelRoundTripper

	// PreventRetryWithBody prevents retrying if the request has a body. Since
	// the body is consumed on a request attempt, in order to retry a request
	// with a body, the body has to be buffered in memory. Setting this
	// to true avoids this buffering: the retry logic is bypassed if a body
	// is present.
	PreventRetryWithBody bool

	// retry is a function that determines if the request should be retried.
	// Unless a retry is prevented based on PreventRetryWithBody, all requests
	// go through that function, even those that are typically considered
	// successful.
	//
	// If it returns false, no retry is attempted, otherwise a retry is
	// attempted after the specified duration.
	retry retryFn

	mu    sync.Mutex
	reqCh map[*http.Request]chan struct{}
}

// RoundTrip implements http.RoundTripper for the Transport type.
// It calls its underlying RoundTripper to execute the request, and
// adds retry logic as per its configuration.
func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
	var attempt int
	preventRetry := req.Body != nil && t.PreventRetryWithBody

	ch := make(chan struct{})
	t.mu.Lock()
	if t.reqCh == nil {
		t.reqCh = make(map[*http.Request]chan struct{})
	}
	t.reqCh[req] = ch
	t.mu.Unlock()

	defer func() {
		t.mu.Lock()
		delete(t.reqCh, req)
		t.mu.Unlock()
	}()

	// buffer the body if needed
	var br *bytes.Reader
	if req.Body != nil && !preventRetry {
		var buf bytes.Buffer
		if _, err := io.Copy(&buf, req.Body); err != nil {
			// cannot even try the first attempt, body has been consumed
			req.Body.Close()
			return nil, err
		}
		req.Body.Close()

		br = bytes.NewReader(buf.Bytes())
		req.Body = ioutil.NopCloser(br)
	}

	for {
		res, err := t.CancelRoundTripper.RoundTrip(req)
		if preventRetry {
			return res, err
		}

		retry, delay := t.retry(Attempt{
			Request:  req,
			Response: res,
			Index:    attempt,
			Error:    err,
		})
		if !retry {
			return res, err
		}

		if br != nil {
			// Per Go's doc: "RoundTrip should not modify the request,
			// except for consuming and closing the Body", so the only thing
			// to reset on the request is the body, if any.
			if _, serr := br.Seek(0, 0); serr != nil {
				// failed to retry, return the results
				return res, err
			}
			req.Body = ioutil.NopCloser(br)
		}
		// close the disposed response's body, if any
		if res != nil {
			res.Body.Close()
		}

		select {
		case <-time.After(delay):
			attempt++
		case <-req.Cancel:
			// request canceled by caller, don't retry
			return nil, errors.New("net/http: request canceled")
		case <-ch:
			// request canceled by call to CancelRequest, don't retry
			return nil, errors.New("net/http: request canceled")
		}
	}
}

// CancelRequest cancels the specified request, preventing any pending
// retry.
func (t *Transport) CancelRequest(req *http.Request) {
	var ch chan struct{}
	t.mu.Lock()
	if t.reqCh != nil {
		ch = t.reqCh[req]
		delete(t.reqCh, req)
	}
	t.mu.Unlock()

	if ch != nil {
		close(ch)
	}

	t.CancelRoundTripper.CancelRequest(req)
}
<commit_msg>jsonreq,rehttp: drain response body before closing<commit_after>// Package rehttp implements an HTTP transport that handles retries.
//
// An HTTP Client can be created with a *rehttp.Transport as RoundTripper
// and it will apply the retry strategy to its requests, e.g.:
//
//     tr, err := rehttp.NewTransport(
//         nil,                            // will use http.DefaultTransport
//         rehttp.RetryTemporaryErr(3),    // max 3 retries for Temporary errors
//         rehttp.ConstDelay(time.Second), // wait 1s between retries
//     )
//     if err != nil {
//         // handle err
//     }
//     client := &http.Client{
//         Transport: tr,
//         Timeout: 30 * time.Second, // timeout applies to all retries as a whole
//     }
//
// The retry strategy is provided by the Transport, which holds
// a function that returns whether or not the request should be retried,
// and if so, what delay to apply before retrying, based on the ShouldRetryFn
// and DelayFn functions passed to NewTransport.
//
// The package offers common delay strategies as ready-made functions that
// return a DelayFn:
//     - ConstDelay(delay time.Duration) DelayFn
//     - ExponentialDelay(base, max time.Duration) DelayFn
//
// It also provides common retry predicates that return a ShouldRetryFn:
//     - RetryTemporaryErr(maxRetries int) ShouldRetryFn
//     - RetryStatus500(maxRetries int) ShouldRetryFn
//     - RetryHTTPMethods(maxRetries int, methods ...string) ShouldRetryFn
//
// Those can be combined with RetryAny or RetryAll as needed. RetryAny
// enables retries if any of the ShouldRetryFn return true, while RetryAll
// enables retries if all ShouldRetryFn return true.
//
// The Transport will buffer the request's body in order to be able to
// retry the request, as a request attempt will consume and close the
// existing body. Sometimes this is not desirable, so it can be prevented
// by setting PreventRetryWithBody to true on the Transport. Doing so
// will disable retries when a request has a non-nil body.
//
package rehttp

import (
	"bytes"
	"errors"
	"io"
	"io/ioutil"
	"math"
	"math/rand"
	"net/http"
	"strings"
	"sync"
	"time"
)

// PRNG is the *math.Rand value to use to add jitter to the backoff
// algorithm used in ExponentialDelay. By default it uses a *math.Rand
// initialized with a source based on the current time in nanoseconds.
var PRNG = rand.New(rand.NewSource(time.Now().UnixNano()))

// terribly named interface to detect errors that support Temporary.
type temporaryer interface {
	Temporary() bool
}

// CancelRoundTripper is a RoundTripper that supports CancelRequest.
// The *http.Transport type implements this interface.
type CancelRoundTripper interface {
	http.RoundTripper
	CancelRequest(*http.Request)
}

// Attempt holds the data for a RoundTrip attempt.
type Attempt struct {
	// Index is the attempt index starting at 0.
	Index int

	// Request is the request for this attempt. If a non-nil Response
	// is present, this is the same as Response.Request, but since a
	// Response may not be available, it is guaranteed to be set on this
	// field.
	Request *http.Request

	// Response is the response for this attempt. It may be nil if an
	// error occurred an no response was received.
	Response *http.Response

	// Error is the error returned by the attempt, if any.
	Error error
}

// retryFn is the signature for functions that implement retry strategies.
type retryFn func(attempt Attempt) (bool, time.Duration)

// DelayFn is the signature for functions that return the delay to apply
// before the next retry.
type DelayFn func(attempt Attempt) time.Duration

// ShouldRetryFn is the signature for functions that return whether a
// retry should be done for the request.
type ShouldRetryFn func(attempt Attempt) bool

// NewTransport creates a Transport with a retry strategy based on
// shouldRetry and delay to control the retry logic. It uses the provided
// CancelRoundTripper to execute the requests. If rt is nil,
// http.DefaultTransport is used. An error is returned if http.DefaultTransport
// is not a CancelRoundTripper (which it is by default).
func NewTransport(rt CancelRoundTripper, shouldRetry ShouldRetryFn, delay DelayFn) (*Transport, error) {
	if rt == nil {
		var ok bool
		rt, ok = http.DefaultTransport.(CancelRoundTripper)
		if !ok {
			return nil, errors.New("http.DefaultTransport is not a CancelRoundTripper")
		}
	}
	return &Transport{
		CancelRoundTripper: rt,
		retry:              toRetryFn(shouldRetry, delay),
	}, nil
}

// toRetryFn combines shouldRetry and delay into a retryFn.
func toRetryFn(shouldRetry ShouldRetryFn, delay DelayFn) retryFn {
	return func(attempt Attempt) (bool, time.Duration) {
		retry := shouldRetry(attempt)
		if !retry {
			return false, 0
		}
		return true, delay(attempt)
	}
}

// RetryAny returns a ShouldRetryFn that allows a retry as long as one of
// the retryFns returns true. If retryFns is empty, it always returns false.
func RetryAny(retryFns ...ShouldRetryFn) ShouldRetryFn {
	return func(attempt Attempt) bool {
		for _, fn := range retryFns {
			if fn(attempt) {
				return true
			}
		}
		return false
	}
}

// RetryAll returns a ShouldRetryFn that allows a retry if all retryFns
// return true. If retryFns is empty, it always returns true.
func RetryAll(retryFns ...ShouldRetryFn) ShouldRetryFn {
	return func(attempt Attempt) bool {
		for _, fn := range retryFns {
			if !fn(attempt) {
				return false
			}
		}
		return true
	}
}

// RetryTemporaryErr returns a ShouldRetryFn that retries up to maxRetries
// times for a temporary error. A temporary error is one that implements
// the Temporary() bool method. Most errors from the net package implement
// this.
func RetryTemporaryErr(maxRetries int) ShouldRetryFn {
	return func(attempt Attempt) bool {
		if attempt.Index >= maxRetries {
			return false
		}
		if terr, ok := attempt.Error.(temporaryer); ok {
			return terr.Temporary()
		}
		return false
	}
}

// RetryStatus500 returns a ShouldRetryFn that retries up to maxRetries times
// for a status code 5xx.
func RetryStatus500(maxRetries int) ShouldRetryFn {
	return func(attempt Attempt) bool {
		if attempt.Index >= maxRetries {
			return false
		}
		return attempt.Response != nil &&
			attempt.Response.StatusCode >= 500 &&
			attempt.Response.StatusCode < 600 // who knows
	}
}

// RetryHTTPMethods returns a ShouldRetryFn that retries up to maxRetries
// times if the request's HTTP method is one of the provided methods.
// It is meant to be used in conjunction with another ShouldRetryFn such
// as RetryTemporaryErr combined using RetryAll, otherwise this function
// will retry any successful request made with one of the provided methods.
func RetryHTTPMethods(maxRetries int, methods ...string) ShouldRetryFn {
	for i, m := range methods {
		methods[i] = strings.ToUpper(m)
	}

	return func(attempt Attempt) bool {
		if attempt.Index >= maxRetries {
			return false
		}
		curMeth := strings.ToUpper(attempt.Request.Method)
		for _, m := range methods {
			if curMeth == m {
				return true
			}
		}
		return false
	}
}

// ConstDelay returns a DelayFn that always returns the same delay.
func ConstDelay(delay time.Duration) DelayFn {
	return func(attempt Attempt) time.Duration {
		return delay
	}
}

// ExponentialDelay returns a DelayFn that returns a delay between 0 and
// base * 2^attempt, capped at max.
//
// See the full jitter algorithm in:
// http://www.awsarchitectureblog.com/2015/03/backoff.html
func ExponentialDelay(base, max time.Duration) DelayFn {
	return func(attempt Attempt) time.Duration {
		exp := math.Pow(2, float64(attempt.Index))
		top := float64(base) * exp
		return time.Duration(
			PRNG.Int63n(int64(math.Min(float64(max), top))),
		)
	}
}

// Transport wraps a CancelRoundTripper such as *http.Transport and adds
// retry logic.
type Transport struct {
	CancelRoundTripper

	// PreventRetryWithBody prevents retrying if the request has a body. Since
	// the body is consumed on a request attempt, in order to retry a request
	// with a body, the body has to be buffered in memory. Setting this
	// to true avoids this buffering: the retry logic is bypassed if a body
	// is present.
	PreventRetryWithBody bool

	// retry is a function that determines if the request should be retried.
	// Unless a retry is prevented based on PreventRetryWithBody, all requests
	// go through that function, even those that are typically considered
	// successful.
	//
	// If it returns false, no retry is attempted, otherwise a retry is
	// attempted after the specified duration.
	retry retryFn

	mu    sync.Mutex
	reqCh map[*http.Request]chan struct{}
}

// RoundTrip implements http.RoundTripper for the Transport type.
// It calls its underlying RoundTripper to execute the request, and
// adds retry logic as per its configuration.
func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
	var attempt int
	preventRetry := req.Body != nil && t.PreventRetryWithBody

	ch := make(chan struct{})
	t.mu.Lock()
	if t.reqCh == nil {
		t.reqCh = make(map[*http.Request]chan struct{})
	}
	t.reqCh[req] = ch
	t.mu.Unlock()

	defer func() {
		t.mu.Lock()
		delete(t.reqCh, req)
		t.mu.Unlock()
	}()

	// buffer the body if needed
	var br *bytes.Reader
	if req.Body != nil && !preventRetry {
		var buf bytes.Buffer
		if _, err := io.Copy(&buf, req.Body); err != nil {
			// cannot even try the first attempt, body has been consumed
			req.Body.Close()
			return nil, err
		}
		req.Body.Close()

		br = bytes.NewReader(buf.Bytes())
		req.Body = ioutil.NopCloser(br)
	}

	for {
		res, err := t.CancelRoundTripper.RoundTrip(req)
		if preventRetry {
			return res, err
		}

		retry, delay := t.retry(Attempt{
			Request:  req,
			Response: res,
			Index:    attempt,
			Error:    err,
		})
		if !retry {
			return res, err
		}

		if br != nil {
			// Per Go's doc: "RoundTrip should not modify the request,
			// except for consuming and closing the Body", so the only thing
			// to reset on the request is the body, if any.
			if _, serr := br.Seek(0, 0); serr != nil {
				// failed to retry, return the results
				return res, err
			}
			req.Body = ioutil.NopCloser(br)
		}
		// close the disposed response's body, if any
		if res != nil {
			io.Copy(ioutil.Discard, res.Body)
			res.Body.Close()
		}

		select {
		case <-time.After(delay):
			attempt++
		case <-req.Cancel:
			// request canceled by caller, don't retry
			return nil, errors.New("net/http: request canceled")
		case <-ch:
			// request canceled by call to CancelRequest, don't retry
			return nil, errors.New("net/http: request canceled")
		}
	}
}

// CancelRequest cancels the specified request, preventing any pending
// retry.
func (t *Transport) CancelRequest(req *http.Request) {
	var ch chan struct{}
	t.mu.Lock()
	if t.reqCh != nil {
		ch = t.reqCh[req]
		delete(t.reqCh, req)
	}
	t.mu.Unlock()

	if ch != nil {
		close(ch)
	}

	t.CancelRoundTripper.CancelRequest(req)
}
<|endoftext|>


<commit_before>package vsolver

import (
	"fmt"
	"net/url"
	"regexp"
	"strings"
)

// A remoteRepo represents a potential remote repository resource.
//
// RemoteRepos are based purely on lexical analysis; successfully constructing
// one is not a guarantee that the resource it identifies actually exists or is
// accessible.
type remoteRepo struct {
	Base     string
	RelPkg   string
	CloneURL *url.URL
	Schemes  []string
	VCS      []string
}

//type remoteResult struct {
//r   remoteRepo
//err error
//}

// TODO sync access to this map
//var remoteCache = make(map[string]remoteResult)

// Regexes for the different known import path flavors
var (
	ghRegex      = regexp.MustCompile(`^(?P<root>github\.com/([A-Za-z0-9_.\-]+/[A-Za-z0-9_.\-]+))(/[A-Za-z0-9_.\-]+)*$`)
	gpinNewRegex = regexp.MustCompile(`^(?P<root>gopkg\.in/(?:([a-zA-Z0-9][-a-zA-Z0-9]+)/)?([a-zA-Z][-.a-zA-Z0-9]*)\.((?:v0|v[1-9][0-9]*)(?:\.0|\.[1-9][0-9]*){0,2}(-unstable)?)(?:\.git))?((?:/[a-zA-Z0-9][-.a-zA-Z0-9]*)*)$`)
	//gpinOldRegex = regexp.MustCompile(`^(?P<root>gopkg\.in/(?:([a-z0-9][-a-z0-9]+)/)?((?:v0|v[1-9][0-9]*)(?:\.0|\.[1-9][0-9]*){0,2}(-unstable)?))/([a-zA-Z][-a-zA-Z0-9]*)(?:\.git)?((?:/[a-zA-Z][-a-zA-Z0-9]*)*)$`)
	bbRegex = regexp.MustCompile(`^(?P<root>bitbucket\.org/(?P<bitname>[A-Za-z0-9_.\-]+/[A-Za-z0-9_.\-]+))(/[A-Za-z0-9_.\-]+)*$`)
	lpRegex = regexp.MustCompile(`^(?P<root>launchpad.net/([A-Za-z0-9-._]+)(/[A-Za-z0-9-._]+)?)(/.+)?`)
	//glpRegex = regexp.MustCompile(`^(?P<root>git\.launchpad\.net/(([A-Za-z0-9_.\-]+)|~[A-Za-z0-9_.\-]+/(\+git|[A-Za-z0-9_.\-]+)/[A-Za-z0-9_.\-]+))$`)
	//gcRegex      = regexp.MustCompile(`^(?P<root>code\.google\.com/[pr]/(?P<project>[a-z0-9\-]+)(\.(?P<subrepo>[a-z0-9\-]+))?)(/[A-Za-z0-9_.\-]+)*$`)
	jazzRegex    = regexp.MustCompile(`^(?P<root>hub\.jazz\.net/git/[a-z0-9]+/[A-Za-z0-9_.\-]+)(/[A-Za-z0-9_.\-]+)*$`)
	genericRegex = regexp.MustCompile(`^(?P<root>(?P<repo>([a-z0-9.\-]+\.)+[a-z0-9.\-]+(:[0-9]+)?/[A-Za-z0-9_.\-/~]*?)\.(?P<vcs>bzr|git|hg|svn))([/A-Za-z0-9_.\-]+)*$`)
)

// Other helper regexes
var (
	scpSyntaxRe = regexp.MustCompile(`^([a-zA-Z0-9_]+)@([a-zA-Z0-9._-]+):(.*)$`)
	pathvld     = regexp.MustCompile(`^([A-Za-z0-9-]+)(\.[A-Za-z0-9-]+)+(/[A-Za-z0-9-_.~]+)*$`)
)

// deduceRemoteRepo takes a potential import path and returns a RemoteRepo
// representing the remote location of the source of an import path. Remote
// repositories can be bare import paths, or urls including a checkout scheme.
func deduceRemoteRepo(path string) (rr remoteRepo, err error) {
	if m := scpSyntaxRe.FindStringSubmatch(path); m != nil {
		// Match SCP-like syntax and convert it to a URL.
		// Eg, "git@github.com:user/repo" becomes
		// "ssh://git@github.com/user/repo".
		rr.CloneURL = &url.URL{
			Scheme:  "ssh",
			User:    url.User(m[1]),
			Host:    m[2],
			RawPath: m[3],
		}
	} else {
		rr.CloneURL, err = url.Parse(path)
		if err != nil {
			return remoteRepo{}, fmt.Errorf("%q is not a valid import path", path)
		}
	}

	path = rr.CloneURL.Host + rr.CloneURL.Path
	if !pathvld.MatchString(path) {
		return remoteRepo{}, fmt.Errorf("%q is not a valid import path", path)
	}

	if rr.CloneURL.Scheme != "" {
		rr.Schemes = []string{rr.CloneURL.Scheme}
	}

	switch {
	case ghRegex.MatchString(path):
		v := ghRegex.FindStringSubmatch(path)

		rr.CloneURL.Host = "github.com"
		rr.CloneURL.Path = v[2]
		rr.Base = v[1]
		rr.RelPkg = strings.TrimPrefix(v[3], "/")
		rr.VCS = []string{"git"}

		return

	case gpinNewRegex.MatchString(path):
		v := gpinNewRegex.FindStringSubmatch(path)

		// Duplicate some logic from the gopkg.in server in order to validate
		// the import path string without having to hit the server
		if strings.Contains(v[4], ".") {
			return remoteRepo{}, fmt.Errorf("%q is not a valid import path; gopkg.in only allows major versions (%q instead of %q)",
				path, v[4][:strings.Index(v[4], ".")], v[4])
		}

		// If the third position is empty, it's the shortened form that expands
		// to the go-pkg github user
		if v[3] != "" {
			rr.CloneURL.Path = "go-pkg/" + v[4]
		} else {
			rr.CloneURL.Path = v[2] + v[4]
		}
		rr.CloneURL.Host = "github.com"
		rr.Base = v[1]
		rr.RelPkg = strings.TrimPrefix(v[6], "/")
		rr.VCS = []string{"git"}

		return
	//case gpinOldRegex.MatchString(path):

	case bbRegex.MatchString(path):
		v := bbRegex.FindStringSubmatch(path)

		rr.CloneURL.Host = "bitbucket.org"
		rr.CloneURL.Path = v[2]
		rr.Base = v[1]
		rr.RelPkg = strings.TrimPrefix(v[5], "/")
		rr.VCS = []string{"git", "hg"}

		return

	//case gcRegex.MatchString(path):
	//v := gcRegex.FindStringSubmatch(path)

	//rr.CloneURL.Host = "code.google.com"
	//rr.CloneURL.Path = "p/" + v[2]
	//rr.Base = v[1]
	//rr.RelPkg = strings.TrimPrefix(v[5], "/")
	//rr.VCS = []string{"hg", "git"}

	//return

	case lpRegex.MatchString(path):
		v := lpRegex.FindStringSubmatch(path)
		v = append(v, "", "")

		rr.CloneURL.Host = "launchpad.net"
		rr.Base = v[1]
		rr.RelPkg = strings.TrimPrefix(v[4], "/")
		rr.VCS = []string{"bzr"}

		if v[3] == "" {
			// launchpad.net/project"
			rr.Base = fmt.Sprintf("https://launchpad.net/%v", v[2])
		} else {
			// launchpad.net/project/series"
			rr.Base = fmt.Sprintf("https://launchpad.net/%s/%s", v[2], v[3])
		}
		return

	//case glpRegex.MatchString(path):
	//// TODO too many rules for this, commenting out for now
	//v := lpRegex.FindStringSubmatch(path)

	//rr.CloneURL.Host = "launchpad.net"
	//rr.RelPkg = strings.TrimPrefix(v[3], "/")
	//rr.VCS = []string{"git"}

	//v = append(v, "", "")
	//if v[2] == "" {
	//// launchpad.net/project"
	//rr.Base = fmt.Sprintf("https://launchpad.net/%v", v[1])
	//} else {
	//// launchpad.net/project/series"
	//rr.Base = fmt.Sprintf("https://launchpad.net/%s/%s", v[1], v[2])
	//}
	//return

	// try the general syntax
	case genericRegex.MatchString(path):
		v := genericRegex.FindStringSubmatch(path)
		switch v[5] {
		case "git", "hg", "bzr":
			x := strings.SplitN(v[1], "/", 2)
			// TODO is this actually correct for bzr?
			rr.CloneURL.Host = x[0]
			rr.CloneURL.Path = x[1]
			rr.VCS = []string{v[5]}
			rr.Base = v[1]
			rr.RelPkg = strings.TrimPrefix(v[6], "/")
			return
		default:
			return remoteRepo{}, fmt.Errorf("unknown repository type: %q", v[5])
		}
	}

	// TODO use HTTP metadata to resolve vanity imports
	return remoteRepo{}, fmt.Errorf("unable to deduct repository and source type for: %q", path)
}
<commit_msg>Can't forget jazz! (and apache)<commit_after>package vsolver

import (
	"fmt"
	"net/url"
	"regexp"
	"strings"
)

// A remoteRepo represents a potential remote repository resource.
//
// RemoteRepos are based purely on lexical analysis; successfully constructing
// one is not a guarantee that the resource it identifies actually exists or is
// accessible.
type remoteRepo struct {
	Base     string
	RelPkg   string
	CloneURL *url.URL
	Schemes  []string
	VCS      []string
}

//type remoteResult struct {
//r   remoteRepo
//err error
//}

// TODO sync access to this map
//var remoteCache = make(map[string]remoteResult)

// Regexes for the different known import path flavors
var (
	ghRegex      = regexp.MustCompile(`^(?P<root>github\.com/([A-Za-z0-9_.\-]+/[A-Za-z0-9_.\-]+))(/[A-Za-z0-9_.\-]+)*$`)
	gpinNewRegex = regexp.MustCompile(`^(?P<root>gopkg\.in/(?:([a-zA-Z0-9][-a-zA-Z0-9]+)/)?([a-zA-Z][-.a-zA-Z0-9]*)\.((?:v0|v[1-9][0-9]*)(?:\.0|\.[1-9][0-9]*){0,2}(-unstable)?)(?:\.git))?((?:/[a-zA-Z0-9][-.a-zA-Z0-9]*)*)$`)
	//gpinOldRegex = regexp.MustCompile(`^(?P<root>gopkg\.in/(?:([a-z0-9][-a-z0-9]+)/)?((?:v0|v[1-9][0-9]*)(?:\.0|\.[1-9][0-9]*){0,2}(-unstable)?))/([a-zA-Z][-a-zA-Z0-9]*)(?:\.git)?((?:/[a-zA-Z][-a-zA-Z0-9]*)*)$`)
	bbRegex = regexp.MustCompile(`^(?P<root>bitbucket\.org/(?P<bitname>[A-Za-z0-9_.\-]+/[A-Za-z0-9_.\-]+))(/[A-Za-z0-9_.\-]+)*$`)
	lpRegex = regexp.MustCompile(`^(?P<root>launchpad.net/([A-Za-z0-9-._]+)(/[A-Za-z0-9-._]+)?)(/.+)?`)
	//glpRegex = regexp.MustCompile(`^(?P<root>git\.launchpad\.net/(([A-Za-z0-9_.\-]+)|~[A-Za-z0-9_.\-]+/(\+git|[A-Za-z0-9_.\-]+)/[A-Za-z0-9_.\-]+))$`)
	//gcRegex      = regexp.MustCompile(`^(?P<root>code\.google\.com/[pr]/(?P<project>[a-z0-9\-]+)(\.(?P<subrepo>[a-z0-9\-]+))?)(/[A-Za-z0-9_.\-]+)*$`)
	jazzRegex    = regexp.MustCompile(`^(?P<root>hub\.jazz\.net/git/[a-z0-9]+/[A-Za-z0-9_.\-]+)(/[A-Za-z0-9_.\-]+)*$`)
	apacheRegex  = regexp.MustCompile(`^(?P<root>git.apache.org/[a-z0-9_.\-]+\.git)(/[A-Za-z0-9_.\-]+)*$`)
	genericRegex = regexp.MustCompile(`^(?P<root>(?P<repo>([a-z0-9.\-]+\.)+[a-z0-9.\-]+(:[0-9]+)?/[A-Za-z0-9_.\-/~]*?)\.(?P<vcs>bzr|git|hg|svn))([/A-Za-z0-9_.\-]+)*$`)
)

// Other helper regexes
var (
	scpSyntaxRe = regexp.MustCompile(`^([a-zA-Z0-9_]+)@([a-zA-Z0-9._-]+):(.*)$`)
	pathvld     = regexp.MustCompile(`^([A-Za-z0-9-]+)(\.[A-Za-z0-9-]+)+(/[A-Za-z0-9-_.~]+)*$`)
)

// deduceRemoteRepo takes a potential import path and returns a RemoteRepo
// representing the remote location of the source of an import path. Remote
// repositories can be bare import paths, or urls including a checkout scheme.
func deduceRemoteRepo(path string) (rr remoteRepo, err error) {
	if m := scpSyntaxRe.FindStringSubmatch(path); m != nil {
		// Match SCP-like syntax and convert it to a URL.
		// Eg, "git@github.com:user/repo" becomes
		// "ssh://git@github.com/user/repo".
		rr.CloneURL = &url.URL{
			Scheme:  "ssh",
			User:    url.User(m[1]),
			Host:    m[2],
			RawPath: m[3],
		}
	} else {
		rr.CloneURL, err = url.Parse(path)
		if err != nil {
			return remoteRepo{}, fmt.Errorf("%q is not a valid import path", path)
		}
	}

	path = rr.CloneURL.Host + rr.CloneURL.Path
	if !pathvld.MatchString(path) {
		return remoteRepo{}, fmt.Errorf("%q is not a valid import path", path)
	}

	if rr.CloneURL.Scheme != "" {
		rr.Schemes = []string{rr.CloneURL.Scheme}
	}

	switch {
	case ghRegex.MatchString(path):
		v := ghRegex.FindStringSubmatch(path)

		rr.CloneURL.Host = "github.com"
		rr.CloneURL.Path = v[2]
		rr.Base = v[1]
		rr.RelPkg = strings.TrimPrefix(v[3], "/")
		rr.VCS = []string{"git"}

		return

	case gpinNewRegex.MatchString(path):
		v := gpinNewRegex.FindStringSubmatch(path)

		// Duplicate some logic from the gopkg.in server in order to validate
		// the import path string without having to hit the server
		if strings.Contains(v[4], ".") {
			return remoteRepo{}, fmt.Errorf("%q is not a valid import path; gopkg.in only allows major versions (%q instead of %q)",
				path, v[4][:strings.Index(v[4], ".")], v[4])
		}

		// If the third position is empty, it's the shortened form that expands
		// to the go-pkg github user
		if v[3] != "" {
			rr.CloneURL.Path = "go-pkg/" + v[4]
		} else {
			rr.CloneURL.Path = v[2] + v[4]
		}
		rr.CloneURL.Host = "github.com"
		rr.Base = v[1]
		rr.RelPkg = strings.TrimPrefix(v[6], "/")
		rr.VCS = []string{"git"}

		return
	//case gpinOldRegex.MatchString(path):

	case bbRegex.MatchString(path):
		v := bbRegex.FindStringSubmatch(path)

		rr.CloneURL.Host = "bitbucket.org"
		rr.CloneURL.Path = v[2]
		rr.Base = v[1]
		rr.RelPkg = strings.TrimPrefix(v[5], "/")
		rr.VCS = []string{"git", "hg"}

		return

	//case gcRegex.MatchString(path):
	//v := gcRegex.FindStringSubmatch(path)

	//rr.CloneURL.Host = "code.google.com"
	//rr.CloneURL.Path = "p/" + v[2]
	//rr.Base = v[1]
	//rr.RelPkg = strings.TrimPrefix(v[5], "/")
	//rr.VCS = []string{"hg", "git"}

	//return

	case lpRegex.MatchString(path):
		v := lpRegex.FindStringSubmatch(path)
		v = append(v, "", "")

		rr.CloneURL.Host = "launchpad.net"
		rr.Base = v[1]
		rr.RelPkg = strings.TrimPrefix(v[4], "/")
		rr.VCS = []string{"bzr"}

		if v[3] == "" {
			// launchpad.net/project"
			rr.Base = fmt.Sprintf("https://launchpad.net/%v", v[2])
		} else {
			// launchpad.net/project/series"
			rr.Base = fmt.Sprintf("https://launchpad.net/%s/%s", v[2], v[3])
		}
		return

	case jazzRegex.MatchString(path):
		v := jazzRegex.FindStringSubmatch(path)

		rr.CloneURL.Host = "hub.jazz.net"
		rr.CloneURL.Path = "git" + v[2]
		rr.Base = v[1]
		rr.RelPkg = strings.TrimPrefix(v[2], "/")
		rr.VCS = []string{"git"}

		return

	case apacheRegex.MatchString(path):
		v := apacheRegex.FindStringSubmatch(path)

		rr.CloneURL.Host = "git.apache.org"
		rr.Base = v[1]
		rr.RelPkg = strings.TrimPrefix(v[2], "/")
		rr.VCS = []string{"git"}

		return

	//case glpRegex.MatchString(path):
	//// TODO too many rules for this, commenting out for now
	//v := lpRegex.FindStringSubmatch(path)

	//rr.CloneURL.Host = "launchpad.net"
	//rr.RelPkg = strings.TrimPrefix(v[3], "/")
	//rr.VCS = []string{"git"}

	//v = append(v, "", "")
	//if v[2] == "" {
	//// launchpad.net/project"
	//rr.Base = fmt.Sprintf("https://launchpad.net/%v", v[1])
	//} else {
	//// launchpad.net/project/series"
	//rr.Base = fmt.Sprintf("https://launchpad.net/%s/%s", v[1], v[2])
	//}
	//return

	// try the general syntax
	case genericRegex.MatchString(path):
		v := genericRegex.FindStringSubmatch(path)
		switch v[5] {
		case "git", "hg", "bzr":
			x := strings.SplitN(v[1], "/", 2)
			// TODO is this actually correct for bzr?
			rr.CloneURL.Host = x[0]
			rr.CloneURL.Path = x[1]
			rr.VCS = []string{v[5]}
			rr.Base = v[1]
			rr.RelPkg = strings.TrimPrefix(v[6], "/")
			return
		default:
			return remoteRepo{}, fmt.Errorf("unknown repository type: %q", v[5])
		}
	}

	// TODO use HTTP metadata to resolve vanity imports
	return remoteRepo{}, fmt.Errorf("unable to deduct repository and source type for: %q", path)
}
<|endoftext|>


<commit_before>// Package rndrec is used to randomly select records from a pool based on their
// relative weight. For example, if the relative weight of one record is 50, on
// average it will be selected five times more often than a record with a
// relative weight of 10. This is useful for generating plausible data sets for
// testing purposes, for example names based on frequency or regions based on
// population.
package rndrec

import (
	"bytes"
	"encoding/csv"
	"fmt"
	"io"
	"math/rand"
	"os"
	"regexp"
	"sort"
	"strconv"
)

var reIntDelimiter = regexp.MustCompile("[,_]")

type recType struct {
	// cumulative frequency
	cf float64
	// list of record fields
	fields []string
}

// SrcType is used to generate plausible random records based on a list of
// weighted records.
type SrcType struct {
	// list of records with ascending cumulative frequency
	list []recType
	// maximum cumulative frequency
	cfMax float64
	// random number generator
	rand *rand.Rand
}

// String implements the fmt.Stringer interface
func (r SrcType) String() string {
	var b bytes.Buffer
	fmt.Fprintf(&b, "Cumulative frequency maximum: %.2f\n", r.cfMax)
	for j, rec := range r.list {
		fmt.Fprintf(&b, "%2d: [%10.2f] %v\n", j, rec.cf, rec.fields)
	}
	return b.String()
}

// NewRandomRecordSource processes a list of multi-field records in which each
// field is a string. With one exception, one column must be an integer weight.
// In this column, specified by weightColPos, each occurrence of an underscore,
// comma or period is removed and the remaining string is parsed as an integer.
// The values in this column are relative weights; that is, a record that has a
// weight twice that of some other record will be selected by Record() on
// average twice as often. The sum of these weights does not have to be any
// special value. The exception to the requirement that one field be a weight
// is when all records are weighted equally. In this case, weightColPos can be
// set to -1 and records do not need to have a weight column. Records returned
// by the Record() method depend on a local pseudo-random number generator;
// seed is used to seed this generator. If any value in the column specified by
// weightColPos can not be parsed as an integer, or the cumulative value of
// weights is zero, or the number of records is zero, an error is returned.
// Otherwise, err is nil and src may be used to retrieve records that are
// distributed according to their relative weights.
func NewRandomRecordSource(recs [][]string, weightColPos int, seed int64) (src *SrcType, err error) {
	var rec recType
	var weight string
	src = new(SrcType)
	for _, fields := range recs {
		if err == nil {
			if weightColPos == -1 {
				rec.cf = 1
			} else if weightColPos >= 0 && weightColPos < len(fields) {
				weight = fields[weightColPos]
				rec.cf, err = strconv.ParseFloat(reIntDelimiter.ReplaceAllString(weight, ""), 64)
				// rec.cf, err = strconv.ParseInt(reIntDelimiter.ReplaceAllString(weight, ""), 10, 64)
			} else {
				err = fmt.Errorf("specified weight column (%d) is out of range", weightColPos)
			}
			if err == nil {
				rec.fields = fields
				src.list = append(src.list, rec)
			}
		}
	}
	if err == nil {
		if len(src.list) > 0 {
			for j, rec := range src.list {
				src.cfMax += rec.cf
				src.list[j].cf = src.cfMax
			}
			if src.cfMax > 0 {
				src.rand = rand.New(rand.NewSource(seed))
			} else {
				err = fmt.Errorf("cumulative frequency must be greater than zero")
			}
		} else {
			err = fmt.Errorf("number of records must be greater than zero")
		}
	}
	if err != nil {
		src = nil
	}
	return
}

// NewRandomRecordSourceFromReader processes a list of multi-field records in
// the form of a comma-separated-value buffer that can be read with the
// io.Reader r. Each record must be separated by a newline. Each field is
// separated by the value specified by fieldSep. For more information on the
// return value and the other arguments, see NewRandomRecordSource().
func NewRandomRecordSourceFromReader(r io.Reader, weightColPos int, fieldSep rune, seed int64) (src *SrcType, err error) {
	var rdr *csv.Reader
	var recs [][]string
	rdr = csv.NewReader(r)
	rdr.Comma = fieldSep
	recs, err = rdr.ReadAll()
	if err == nil {
		src, err = NewRandomRecordSource(recs, weightColPos, seed)
	}
	return
}

// NewRandomRecordSourceFromFile processes a list of multi-field records in the
// form of a comma-separated-value file with the filename specified by fileStr.
// Each record must be separated by a newline. Each field is separated by the
// value specified by fieldSep. For more information on the return value and
// the other arguments, see NewRandomRecordSource().
func NewRandomRecordSourceFromFile(fileStr string, weightColPos int, fieldSep rune, seed int64) (src *SrcType, err error) {
	var f *os.File
	f, err = os.Open(fileStr)
	if err == nil {
		src, err = NewRandomRecordSourceFromReader(f, weightColPos, fieldSep, seed)
		f.Close()
	}
	return
}

// Record returns a random record based on its relative weight. For example, a
// record with a relative weight of 40 will be returned, on average, four times
// as often as a record with the relative weight of 10. The returned record
// will be in the form of a slice of strings taken directly from the original
// list used to initialize the SrcType instance.
func (r *SrcType) Record() []string {
	var cf float64
	var pos int
	cf = r.rand.Float64() * r.cfMax // 0 <= cf < r.cfMax
	pos = sort.Search(len(r.list), func(j int) bool {
		return r.list[j].cf > cf
	})
	return r.list[pos].fields
}
<commit_msg>Correct documentation regarding ignored characters in weight value<commit_after>// Package rndrec is used to randomly select records from a pool based on their
// relative weight. For example, if the relative weight of one record is 50, on
// average it will be selected five times more often than a record with a
// relative weight of 10. This is useful for generating plausible data sets for
// testing purposes, for example names based on frequency or regions based on
// population.
package rndrec

import (
	"bytes"
	"encoding/csv"
	"fmt"
	"io"
	"math/rand"
	"os"
	"regexp"
	"sort"
	"strconv"
)

var reIntDelimiter = regexp.MustCompile("[,_]")

type recType struct {
	// cumulative frequency
	cf float64
	// list of record fields
	fields []string
}

// SrcType is used to generate plausible random records based on a list of
// weighted records.
type SrcType struct {
	// list of records with ascending cumulative frequency
	list []recType
	// maximum cumulative frequency
	cfMax float64
	// random number generator
	rand *rand.Rand
}

// String implements the fmt.Stringer interface
func (r SrcType) String() string {
	var b bytes.Buffer
	fmt.Fprintf(&b, "Cumulative frequency maximum: %.2f\n", r.cfMax)
	for j, rec := range r.list {
		fmt.Fprintf(&b, "%2d: [%10.2f] %v\n", j, rec.cf, rec.fields)
	}
	return b.String()
}

// NewRandomRecordSource processes a list of multi-field records in which each
// field is a string. With one exception, one column must be an integer weight.
// In this column, specified by weightColPos, each occurrence of an underscore
// or comma is removed and the remaining string is parsed as an integer. The
// values in this column are relative weights; that is, a record that has a
// weight twice that of some other record will be selected by Record() on
// average twice as often. The sum of these weights does not have to be any
// special value. The exception to the requirement that one field be a weight
// is when all records are weighted equally. In this case, weightColPos can be
// set to -1 and records do not need to have a weight column. Records returned
// by the Record() method depend on a local pseudo-random number generator;
// seed is used to seed this generator. If any value in the column specified by
// weightColPos can not be parsed as an integer, or the cumulative value of
// weights is zero, or the number of records is zero, an error is returned.
// Otherwise, err is nil and src may be used to retrieve records that are
// distributed according to their relative weights.
func NewRandomRecordSource(recs [][]string, weightColPos int, seed int64) (src *SrcType, err error) {
	var rec recType
	var weight string
	src = new(SrcType)
	for _, fields := range recs {
		if err == nil {
			if weightColPos == -1 {
				rec.cf = 1
			} else if weightColPos >= 0 && weightColPos < len(fields) {
				weight = fields[weightColPos]
				rec.cf, err = strconv.ParseFloat(reIntDelimiter.ReplaceAllString(weight, ""), 64)
				// rec.cf, err = strconv.ParseInt(reIntDelimiter.ReplaceAllString(weight, ""), 10, 64)
			} else {
				err = fmt.Errorf("specified weight column (%d) is out of range", weightColPos)
			}
			if err == nil {
				rec.fields = fields
				src.list = append(src.list, rec)
			}
		}
	}
	if err == nil {
		if len(src.list) > 0 {
			for j, rec := range src.list {
				src.cfMax += rec.cf
				src.list[j].cf = src.cfMax
			}
			if src.cfMax > 0 {
				src.rand = rand.New(rand.NewSource(seed))
			} else {
				err = fmt.Errorf("cumulative frequency must be greater than zero")
			}
		} else {
			err = fmt.Errorf("number of records must be greater than zero")
		}
	}
	if err != nil {
		src = nil
	}
	return
}

// NewRandomRecordSourceFromReader processes a list of multi-field records in
// the form of a comma-separated-value buffer that can be read with the
// io.Reader r. Each record must be separated by a newline. Each field is
// separated by the value specified by fieldSep. For more information on the
// return value and the other arguments, see NewRandomRecordSource().
func NewRandomRecordSourceFromReader(r io.Reader, weightColPos int, fieldSep rune, seed int64) (src *SrcType, err error) {
	var rdr *csv.Reader
	var recs [][]string
	rdr = csv.NewReader(r)
	rdr.Comma = fieldSep
	recs, err = rdr.ReadAll()
	if err == nil {
		src, err = NewRandomRecordSource(recs, weightColPos, seed)
	}
	return
}

// NewRandomRecordSourceFromFile processes a list of multi-field records in the
// form of a comma-separated-value file with the filename specified by fileStr.
// Each record must be separated by a newline. Each field is separated by the
// value specified by fieldSep. For more information on the return value and
// the other arguments, see NewRandomRecordSource().
func NewRandomRecordSourceFromFile(fileStr string, weightColPos int, fieldSep rune, seed int64) (src *SrcType, err error) {
	var f *os.File
	f, err = os.Open(fileStr)
	if err == nil {
		src, err = NewRandomRecordSourceFromReader(f, weightColPos, fieldSep, seed)
		f.Close()
	}
	return
}

// Record returns a random record based on its relative weight. For example, a
// record with a relative weight of 40 will be returned, on average, four times
// as often as a record with the relative weight of 10. The returned record
// will be in the form of a slice of strings taken directly from the original
// list used to initialize the SrcType instance.
func (r *SrcType) Record() []string {
	var cf float64
	var pos int
	cf = r.rand.Float64() * r.cfMax // 0 <= cf < r.cfMax
	pos = sort.Search(len(r.list), func(j int) bool {
		return r.list[j].cf > cf
	})
	return r.list[pos].fields
}
<|endoftext|>